Skip to content

Commit c75e886

Browse files
ravitemerravitemer
committed
fix(codecompanion): immediately register capabilities once extension is called
1 parent a358f92 commit c75e886

3 files changed

Lines changed: 326 additions & 301 deletions

File tree

‎lua/mcphub/extensions/codecompanion/slash_commands.lua‎

Lines changed: 113 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -1,129 +1,139 @@
11
local M = {}
2+
local mcphub = require("mcphub")
3+
24
local shared = require("mcphub.extensions.shared")
35

4-
-- Setup MCP prompts as CodeCompanion slash commands
5-
---@param opts MCPHubCodeCompanionConfig
6-
function M.setup(opts)
7-
if not opts.make_slash_commands then
6+
function M.register()
7+
local config = require("codecompanion.config")
8+
local hub = mcphub.get_hub_instance()
9+
if not hub then
810
return
911
end
1012

11-
local mcphub = require("mcphub")
12-
local config = require("codecompanion.config")
13+
local prompts = hub:get_prompts()
14+
local slash_commands = config.strategies.chat.slash_commands
1315

14-
mcphub.on({ "servers_updated", "prompt_list_changed" }, function(_)
15-
local hub = mcphub.get_hub_instance()
16-
if not hub then
17-
return
16+
-- Remove existing MCP slash commands
17+
for key, value in pairs(slash_commands) do
18+
local id = value.id or ""
19+
if id:sub(1, 3) == "mcp" then
20+
slash_commands[key] = nil
1821
end
22+
end
1923

20-
local prompts = hub:get_prompts()
21-
local slash_commands = config.strategies.chat.slash_commands
24+
-- Add current prompts as slash commands
25+
for _, prompt in ipairs(prompts) do
26+
local server_name = prompt.server_name
27+
local prompt_name = prompt.name or ""
28+
local description = prompt.description or ""
29+
description = description:gsub("\n", " ")
30+
description = prompt_name .. " (" .. description .. ")"
2231

23-
-- Remove existing MCP slash commands
24-
for key, value in pairs(slash_commands) do
25-
local id = value.id or ""
26-
if id:sub(1, 3) == "mcp" then
27-
slash_commands[key] = nil
32+
local arguments = prompt.arguments or {}
33+
if type(arguments) == "function" then
34+
local ok, args = pcall(arguments, prompt)
35+
if ok then
36+
arguments = args or {}
37+
else
38+
vim.notify("Error in arguments function: " .. (args or ""), vim.log.levels.ERROR)
39+
arguments = {}
2840
end
2941
end
3042

31-
-- Add current prompts as slash commands
32-
for _, prompt in ipairs(prompts) do
33-
local server_name = prompt.server_name
34-
local prompt_name = prompt.name or ""
35-
local description = prompt.description or ""
36-
description = description:gsub("\n", " ")
37-
description = prompt_name .. " (" .. description .. ")"
38-
39-
local arguments = prompt.arguments or {}
40-
if type(arguments) == "function" then
41-
local ok, args = pcall(arguments, prompt)
42-
if ok then
43-
arguments = args or {}
44-
else
45-
vim.notify("Error in arguments function: " .. (args or ""), vim.log.levels.ERROR)
46-
arguments = {}
47-
end
48-
end
49-
50-
slash_commands["mcp:" .. prompt_name] = {
51-
id = "mcp" .. server_name .. prompt_name,
52-
description = description,
53-
callback = function(self)
54-
shared.collect_arguments(arguments, function(values)
55-
-- Sync call - blocks UI (can't use async in slash_commands yet)
56-
local response, err = hub:get_prompt(server_name, prompt_name, values, {
57-
caller = {
58-
type = "codecompanion",
59-
codecompanion = self,
60-
meta = {
61-
is_within_slash_command = true,
62-
},
43+
slash_commands["mcp:" .. prompt_name] = {
44+
id = "mcp" .. server_name .. prompt_name,
45+
description = description,
46+
callback = function(self)
47+
shared.collect_arguments(arguments, function(values)
48+
-- Sync call - blocks UI (can't use async in slash_commands yet)
49+
local response, err = hub:get_prompt(server_name, prompt_name, values, {
50+
caller = {
51+
type = "codecompanion",
52+
codecompanion = self,
53+
meta = {
54+
is_within_slash_command = true,
6355
},
64-
parse_response = true,
65-
})
56+
},
57+
parse_response = true,
58+
})
6659

67-
if not response then
68-
if err then
69-
vim.notify("Error in slash command: " .. err, vim.log.levels.ERROR)
70-
vim.notify("Prompt cancelled", vim.log.levels.INFO)
71-
end
72-
return
60+
if not response then
61+
if err then
62+
vim.notify("Error in slash command: " .. err, vim.log.levels.ERROR)
63+
vim.notify("Prompt cancelled", vim.log.levels.INFO)
7364
end
65+
return
66+
end
67+
68+
local messages = response.messages or {}
69+
local text_messages = 0
70+
71+
for i, message in ipairs(messages) do
72+
local output = message.output
73+
local mapped_role = message.role == "assistant" and config.constants.LLM_ROLE
74+
or message.role == "system" and config.constants.SYSTEM_ROLE
75+
or config.constants.USER_ROLE
7476

75-
local messages = response.messages or {}
76-
local text_messages = 0
77-
78-
for i, message in ipairs(messages) do
79-
local output = message.output
80-
local mapped_role = message.role == "assistant" and config.constants.LLM_ROLE
81-
or message.role == "system" and config.constants.SYSTEM_ROLE
82-
or config.constants.USER_ROLE
83-
84-
if output.text and output.text ~= "" then
85-
text_messages = text_messages + 1
86-
-- If last message is from user, add it to chat buffer
87-
if i == #messages and mapped_role == config.constants.USER_ROLE then
88-
self:add_buf_message({
89-
role = mapped_role,
90-
content = output.text,
91-
})
92-
else
93-
self:add_message({
94-
role = mapped_role,
95-
content = output.text,
96-
})
97-
end
77+
if output.text and output.text ~= "" then
78+
text_messages = text_messages + 1
79+
-- If last message is from user, add it to chat buffer
80+
if i == #messages and mapped_role == config.constants.USER_ROLE then
81+
self:add_buf_message({
82+
role = mapped_role,
83+
content = output.text,
84+
})
85+
else
86+
self:add_message({
87+
role = mapped_role,
88+
content = output.text,
89+
})
9890
end
91+
end
9992

100-
-- Handle images
101-
if output.images and #output.images > 0 then
102-
local helpers = require("codecompanion.strategies.chat.helpers")
103-
for _, image in ipairs(output.images) do
104-
local id = string.format("mcp-%s", os.time())
105-
helpers.add_image(self, {
106-
id = id,
107-
base64 = image.data,
108-
mimetype = image.mimeType,
109-
}, { role = mapped_role })
110-
end
93+
-- Handle images
94+
if output.images and #output.images > 0 then
95+
local helpers = require("codecompanion.strategies.chat.helpers")
96+
for _, image in ipairs(output.images) do
97+
local id = string.format("mcp-%s", os.time())
98+
helpers.add_image(self, {
99+
id = id,
100+
base64 = image.data,
101+
mimetype = image.mimeType,
102+
}, { role = mapped_role })
111103
end
112104
end
105+
end
113106

114-
vim.notify(
115-
string.format(
116-
"%s message%s added successfully",
117-
text_messages,
118-
text_messages == 1 and "" or "s"
119-
),
120-
vim.log.levels.INFO
121-
)
122-
end)
123-
end,
124-
}
125-
end
107+
vim.notify(
108+
string.format(
109+
"%s message%s added successfully",
110+
text_messages,
111+
text_messages == 1 and "" or "s"
112+
),
113+
vim.log.levels.INFO
114+
)
115+
end)
116+
end,
117+
}
118+
end
119+
end
120+
121+
-- Setup MCP prompts as CodeCompanion slash commands
122+
---@param opts MCPHubCodeCompanionConfig
123+
function M.setup(opts)
124+
if not opts.make_slash_commands then
125+
return
126+
end
127+
128+
vim.schedule(function()
129+
M.register()
126130
end)
131+
mcphub.on(
132+
{ "servers_updated", "prompt_list_changed" },
133+
vim.schedule_wrap(function()
134+
M.register()
135+
end)
136+
)
127137
end
128138

129139
return M

0 commit comments

Comments
 (0)