Skip to content

Commit ff6dcac

Browse files
committed
feat(codecompanion): add format_tool
1 parent c3fc53b commit ff6dcac

3 files changed

Lines changed: 53 additions & 8 deletions

File tree

doc/extensions/codecompanion.md

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,27 @@ require("codecompanion").setup({
3232
show_server_tools_in_chat = true, -- Show individual tools in chat completion (when make_tools=true)
3333
add_mcp_prefix_to_tool_names = false, -- Add mcp__ prefix (e.g `@mcp__github`, `@mcp__neovim__list_issues`)
3434
show_result_in_chat = true, -- Show tool results directly in chat buffer
35+
format_tool = function(display_name, tool)
36+
local args = vim.deepcopy(tool.args)
37+
-- Replace 'use_mcp_tool' with actual tool name and params.
38+
if name == 'use_mcp_tool' then
39+
name = string.format('%s__%s', args.server_name, args.tool_name)
40+
args = args.tool_input
41+
end
42+
43+
if name == 'filesystem__edit_file' then
44+
if type(args.edits) == 'table' then
45+
args.edits = '󰩫'
46+
end
47+
elseif name == 'filesystem__write_file' then
48+
if type(args.content) == 'string' then
49+
args.content = '󰩫'
50+
end
51+
end
52+
53+
local args = vim.inspect(args):gsub('%s+', ' ')
54+
return name .. ' ' .. args
55+
end,
3556
-- MCP Resources
3657
make_vars = true, -- Convert MCP resources to #variables for prompts
3758
-- MCP Prompts

doc/mcphub.txt

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2785,6 +2785,27 @@ Register MCP Hub as an extension in your CodeCompanion configuration:
27852785
show_server_tools_in_chat = true, -- Show individual tools in chat completion (when make_tools=true)
27862786
add_mcp_prefix_to_tool_names = false, -- Add mcp__ prefix (e.g `@mcp__github`, `@mcp__neovim__list_issues`)
27872787
show_result_in_chat = true, -- Show tool results directly in chat buffer
2788+
format_tool = function(display_name, tool)
2789+
local args = vim.deepcopy(tool.args)
2790+
-- Replace 'use_mcp_tool' with actual tool name and params.
2791+
if name == 'use_mcp_tool' then
2792+
name = string.format('%s__%s', args.server_name, args.tool_name)
2793+
args = args.tool_input
2794+
end
2795+
2796+
if name == 'filesystem__edit_file' then
2797+
if type(args.edits) == 'table' then
2798+
args.edits = '󰩫'
2799+
end
2800+
elseif name == 'filesystem__write_file' then
2801+
if type(args.content) == 'string' then
2802+
args.content = '󰩫'
2803+
end
2804+
end
2805+
2806+
local args = vim.inspect(args):gsub('%s+', ' ')
2807+
return name .. ' ' .. args
2808+
end,
27882809
-- MCP Resources
27892810
make_vars = true, -- Convert MCP resources to #variables for prompts
27902811
-- MCP Prompts

lua/mcphub/extensions/codecompanion/core.lua

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ function M.execute_mcp_tool(params, agent, output_handler, context)
8787
end
8888

8989
---@param display_name string
90-
---@param tool table
90+
---@param tool CodeCompanion.Agent.Tool
9191
---@param chat any
9292
---@param llm_msg string
9393
---@param is_error boolean
@@ -111,13 +111,14 @@ local function add_tool_output(
111111
local helpers = require("codecompanion.strategies.chat.helpers")
112112
local show_result_in_chat = opts.show_result_in_chat == true
113113
local text = llm_msg
114+
local formatted_name = opts.format_tool and opts.format_tool(display_name, tool) or display_name
114115

115116
if has_function_calling then
116117
chat:add_tool_output(
117118
tool,
118119
text,
119120
(user_msg or show_result_in_chat or is_error) and (user_msg or text)
120-
or string.format("**`%s` Tool**: Successfully finished", display_name)
121+
or string.format("**`%s` Tool**: Successfully finished", formatted_name)
121122
)
122123
for _, image in ipairs(images) do
123124
helpers.add_image(chat, image)
@@ -135,7 +136,7 @@ local function add_tool_output(
135136
})
136137
chat:add_buf_message({
137138
role = config.constants.USER_ROLE,
138-
content = string.format("I've shared the result of the `%s` tool with you.\n", display_name),
139+
content = string.format("I've shared the result of the `%s` tool with you.\n", formatted_name),
139140
})
140141
end
141142
end
@@ -159,14 +160,15 @@ function M.create_output_handlers(display_name, has_function_calling, opts)
159160
---@diagnostic disable-next-line: cast-local-type
160161
stderr = vim.inspect(stderr)
161162
end
163+
local formatted_name = opts.format_tool and opts.format_tool(display_name, self) or display_name
162164
local err_msg = string.format(
163165
[[**`%s` Tool**: Failed with the following error:
164166
165167
````
166168
%s
167169
````
168170
]],
169-
display_name,
171+
formatted_name,
170172
stderr
171173
)
172174
add_tool_output(display_name, self, agent.chat, err_msg, true, has_function_calling, opts, nil, {})
@@ -182,6 +184,7 @@ function M.create_output_handlers(display_name, has_function_calling, opts)
182184
local result = has_function_calling and stdout[#stdout] or cmd[#cmd]
183185
---@diagnostic disable-next-line: cast-local-type
184186
agent = has_function_calling and agent or self
187+
local formatted_name = opts.format_tool and opts.format_tool(display_name, self) or display_name
185188
local to_llm = nil
186189
local to_user = nil
187190
local images = {}
@@ -193,7 +196,7 @@ function M.create_output_handlers(display_name, has_function_calling, opts)
193196
````
194197
%s
195198
````]],
196-
display_name,
199+
formatted_name,
197200
result.text
198201
)
199202
end
@@ -216,7 +219,7 @@ function M.create_output_handlers(display_name, has_function_calling, opts)
216219
````
217220
%s
218221
````]],
219-
display_name,
222+
formatted_name,
220223
string.format("%d image%s returned", #result.images, #result.images > 1 and "s" or "")
221224
)
222225
end
@@ -233,9 +236,9 @@ function M.create_output_handlers(display_name, has_function_calling, opts)
233236
end
234237
end
235238

236-
local fallback_to_llm = string.format("**`%s` Tool**: Completed with no output", display_name)
239+
local fallback_to_llm = string.format("**`%s` Tool**: Completed with no output", formatted_name)
237240
if opts.show_result_in_chat == false and not to_user then
238-
to_user = string.format("**`%s` Tool**: Successfully finished", display_name)
241+
to_user = string.format("**`%s` Tool**: Successfully finished", formatted_name)
239242
end
240243
add_tool_output(
241244
display_name,

0 commit comments

Comments
 (0)