8888--- @param is_error boolean
8989--- @param has_function_calling boolean
9090--- @param opts MCPHubCodeCompanionConfig
91- local function add_tool_output (action_name , tool , chat , llm_msg , is_error , has_function_calling , opts )
91+ --- @param user_msg string ?
92+ --- @param images { id : string , base64 : string , mimetype : string , cached_file_path : string | nil } []
93+ local function add_tool_output (action_name , tool , chat , llm_msg , is_error , has_function_calling , opts , user_msg , images )
9294 local config = require (" codecompanion.config"
95+ local helpers = require (" codecompanion.strategies.chat.helpers"
9396 local show_result_in_chat = opts .show_result_in_chat == true
9497 -- local text = show_result_in_chat and replace_headers(llm_msg) or llm_msg
9598 local text = llm_msg
9699 if has_function_calling then
97100 chat :add_tool_output (
98101 tool ,
99102 text ,
100- (show_result_in_chat or is_error ) and text
103+ (user_msg or show_result_in_chat or is_error ) and ( user_msg or text )
101104 or string.format (" **`%s` Tool**: Successfully finished" action_name )
102105 )
106+ for _ , image in ipairs (images ) do
107+ helpers .add_image (chat , image )
108+ end
103109 else
104110 if show_result_in_chat or is_error then
105111 chat :add_buf_message ({
@@ -143,16 +149,18 @@ function M.create_output_handlers(action_name, has_function_calling, opts)
143149 action_name ,
144150 stderr
145151 )
146- add_tool_output (action_name , self , agent .chat , err_msg , true , has_function_calling , opts )
152+ add_tool_output (action_name , self , agent .chat , err_msg , true , has_function_calling , opts , nil , {} )
147153 end ,
148154 success = function (self , agent , cmd , stdout )
155+ local image_cache = require (" mcphub.utils.image_cache"
149156 --- @type MCPResponseOutput
150157 local result = has_function_calling and stdout [# stdout ] or cmd [# cmd ]
151158 agent = has_function_calling and agent or self
152- -- Show text content if present
153- local tool_call_result_added = false
159+ local to_llm = nil
160+ local to_user = nil
161+ local images = {}
154162 if result .text and result .text ~= " " then
155- local to_llm = string.format (
163+ to_llm = string.format (
156164 [[ **`%s` Tool**: Returned the following:
157165
158166````
@@ -161,16 +169,52 @@ function M.create_output_handlers(action_name, has_function_calling, opts)
161169 action_name ,
162170 result .text
163171 )
164- add_tool_output (action_name , self , agent .chat , to_llm , false , has_function_calling , opts )
165- tool_call_result_added = true
166172 end
167- if not tool_call_result_added then
168- -- When a tool returns no text content, still send a message to
169- -- ensure the tool_call_id protocol is satisfied
170- local to_llm = string.format (" **`%s` Tool**: Completed with no output" action_name )
171- add_tool_output (action_name , self , agent .chat , to_llm , false , has_function_calling , opts )
173+ if result .images and # result .images > 0 then
174+ --- When the mcp call returns just images, we need to add the tool output
175+ for _ , image in ipairs (result .images ) do
176+ local id = string.format (" mcp-%s" os.time ())
177+ table.insert (images , {
178+ id = id ,
179+ base64 = image .data ,
180+ mimetype = image .mimeType ,
181+ cached_file_path = image_cache .save_image (image .data , image .mimeType ),
182+ })
183+ end
184+ --- If there is no text response, add no of images returned
185+ if not to_llm then
186+ to_llm = string.format (
187+ [[ **`%s` Tool**: Returned the following:
188+ ````
189+ %s
190+ ````]] ,
191+ action_name ,
192+ string.format (" %d image%s returned" # result .images , # result .images > 1 and " s" or " "
193+ )
194+ end
195+ to_user = to_llm .. (# images > 0 and string.format (" \n\n > Preview Images\n " or " "
196+ for _ , image in ipairs (images ) do
197+ local file = image .cached_file_path
198+ if file then
199+ local file_name = vim .fn .fnamemodify (file , " :t"
200+ to_user = to_user .. string.format (" \n \n " file_name , vim .fn .fnameescape (file ))
201+ else
202+ to_user = to_user .. string.format (" \n \n " file )
203+ end
204+ end
172205 end
173- -- TODO: Add image support when codecompanion supports it
206+ local fallback_to_llm = string.format (" **`%s` Tool**: Completed with no output" action_name )
207+ add_tool_output (
208+ action_name ,
209+ self ,
210+ agent .chat ,
211+ to_llm or fallback_to_llm ,
212+ false ,
213+ has_function_calling ,
214+ opts ,
215+ to_user ,
216+ images
217+ )
174218 end ,
175219 }
176220end
@@ -214,7 +258,7 @@ function M.setup_codecompanion_variables(opts)
214258 description = description ,
215259 callback = function (self )
216260 -- this is sync and will block the UI (can't use async in variables yet)
217- local response = hub :access_resource (server_name , uri , {
261+ local result = hub :access_resource (server_name , uri , {
218262 caller = {
219263 type = " codecompanion"
220264 codecompanion = self ,
@@ -224,7 +268,22 @@ function M.setup_codecompanion_variables(opts)
224268 },
225269 parse_response = true ,
226270 })
227- return response and response .text
271+ if not result then
272+ return string.format (" Accessing resource failed: %s" uri )
273+ end
274+
275+ if result .images and # result .images > 0 then
276+ local helpers = require (" codecompanion.strategies.chat.helpers"
277+ for _ , image in ipairs (result .images ) do
278+ local id = string.format (" mcp-%s" os.time ())
279+ helpers .add_image (self .Chat , {
280+ id = id ,
281+ base64 = image .data ,
282+ mimetype = image .mimeType ,
283+ })
284+ end
285+ end
286+ return result .text
228287 end ,
229288 }
230289 end
@@ -299,11 +358,10 @@ function M.setup_codecompanion_slash_commands(opts)
299358 local text_messages = 0
300359 for i , message in ipairs (messages ) do
301360 local output = message .output
302- -- TODO: Currently codecompanion only supports text messages
361+ local mapped_role = message .role == " assistant" and config .constants .LLM_ROLE
362+ or message .role == " system" and config .constants .SYSTEM_ROLE
363+ or config .constants .USER_ROLE
303364 if output .text and output .text ~= " " then
304- local mapped_role = message .role == " assistant" and config .constants .LLM_ROLE
305- or message .role == " system" and config .constants .SYSTEM_ROLE
306- or config .constants .USER_ROLE
307365 text_messages = text_messages + 1
308366 -- if last message is from user, add it to the chat buffer
309367 if i == # messages and mapped_role == config .constants .USER_ROLE then
@@ -318,6 +376,17 @@ function M.setup_codecompanion_slash_commands(opts)
318376 })
319377 end
320378 end
379+ if output .images and # output .images > 0 then
380+ local helpers = require (" codecompanion.strategies.chat.helpers"
381+ for _ , image in ipairs (output .images ) do
382+ local id = string.format (" mcp-%s" os.time ())
383+ helpers .add_image (self , {
384+ id = id ,
385+ base64 = image .data ,
386+ mimetype = image .mimeType ,
387+ }, { role = mapped_role })
388+ end
389+ end
321390 end
322391 vim .notify (
323392 string.format (
0 commit comments