Skip to content
This repository was archived by the owner on May 15, 2026. It is now read-only.

Commit bb488fe

Browse files
roomoteroomote
committed
fix: add image content support to MCP tool responses
Fixes #10872 The processToolContent method in UseMcpToolTool.ts now handles image content types from MCP protocol responses (e.g., Figma's get_screenshot). Changes: - Modified processToolContent to return both text and images - Updated executeToolAndProcessResult to pass images to task.say() and pushToolResult() - Added tests for image handling scenarios
1 parent 9ab279a commit bb488fe

2 files changed

Lines changed: 272 additions & 12 deletions

File tree

‎src/core/tools/UseMcpToolTool.ts‎

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -250,12 +250,14 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
250250
})
251251
}
252252

253-
private processToolContent(toolResult: any): string {
253+
private processToolContent(toolResult: any): { text: string; images: string[] } {
254254
if (!toolResult?.content || toolResult.content.length === 0) {
255-
return ""
255+
return { text: "", images: [] }
256256
}
257257

258-
return toolResult.content
258+
const images: string[] = []
259+
260+
const textContent = toolResult.content
259261
.map((item: any) => {
260262
if (item.type === "text") {
261263
return item.text
@@ -264,10 +266,23 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
264266
const { blob: _, ...rest } = item.resource
265267
return JSON.stringify(rest, null, 2)
266268
}
269+
if (item.type === "image") {
270+
// Handle image content (MCP image content has mimeType and data properties)
271+
if (item.mimeType && item.data) {
272+
if (item.data.startsWith("data:")) {
273+
images.push(item.data)
274+
} else {
275+
images.push(`data:${item.mimeType};base64,${item.data}`)
276+
}
277+
}
278+
return ""
279+
}
267280
return ""
268281
})
269282
.filter(Boolean)
270283
.join("\n\n")
284+
285+
return { text: textContent, images }
271286
}
272287

273288
private async executeToolAndProcessResult(
@@ -291,18 +306,22 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
291306
const toolResult = await task.providerRef.deref()?.getMcpHub()?.callTool(serverName, toolName, parsedArguments)
292307

293308
let toolResultPretty = "(No response)"
309+
let images: string[] = []
294310

295311
if (toolResult) {
296-
const outputText = this.processToolContent(toolResult)
312+
const { text: outputText, images: extractedImages } = this.processToolContent(toolResult)
313+
images = extractedImages
297314

298-
if (outputText) {
315+
if (outputText || images.length > 0) {
299316
await this.sendExecutionStatus(task, {
300317
executionId,
301318
status: "output",
302-
response: outputText,
319+
response: outputText || (images.length > 0 ? `[${images.length} image(s)]` : ""),
303320
})
304321

305-
toolResultPretty = (toolResult.isError ? "Error:\n" : "") + outputText
322+
toolResultPretty =
323+
(toolResult.isError ? "Error:\n" : "") +
324+
(outputText || (images.length > 0 ? `[${images.length} image(s) received]` : ""))
306325
}
307326

308327
// Send completion status
@@ -321,8 +340,8 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
321340
})
322341
}
323342

324-
await task.say("mcp_server_response", toolResultPretty)
325-
pushToolResult(formatResponse.toolResult(toolResultPretty))
343+
await task.say("mcp_server_response", toolResultPretty, images)
344+
pushToolResult(formatResponse.toolResult(toolResultPretty, images))
326345
}
327346
}
328347

‎src/core/tools/__tests__/useMcpToolTool.spec.ts‎

Lines changed: 244 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,12 @@ import { ToolUse } from "../../../shared/tools"
77
// Mock dependencies
88
vi.mock("../../prompts/responses", () => ({
99
formatResponse: {
10-
toolResult: vi.fn((result: string) => `Tool result: ${result}`),
10+
toolResult: vi.fn((result: string, images?: string[]) => {
11+
if (images && images.length > 0) {
12+
return `Tool result: ${result} [with ${images.length} image(s)]`
13+
}
14+
return `Tool result: ${result}`
15+
}),
1116
toolError: vi.fn((error: string) => `Tool error: ${error}`),
1217
invalidMcpToolArgumentError: vi.fn((server: string, tool: string) => `Invalid args for ${server}:${tool}`),
1318
unknownMcpToolError: vi.fn((server: string, tool: string, availableTools: string[]) => {
@@ -245,7 +250,7 @@ describe("useMcpToolTool", () => {
245250
expect(mockTask.consecutiveMistakeCount).toBe(0)
246251
expect(mockAskApproval).toHaveBeenCalled()
247252
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
248-
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully")
253+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully", [])
249254
expect(mockPushToolResult).toHaveBeenCalledWith("Tool result: Tool executed successfully")
250255
})
251256

@@ -483,7 +488,7 @@ describe("useMcpToolTool", () => {
483488
expect(mockTask.consecutiveMistakeCount).toBe(0)
484489
expect(mockTask.recordToolError).not.toHaveBeenCalled()
485490
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
486-
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully")
491+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully", [])
487492
})
488493

489494
it("should reject unknown server names with available servers listed", async () => {
@@ -578,4 +583,240 @@ describe("useMcpToolTool", () => {
578583
expect(mockAskApproval).not.toHaveBeenCalled()
579584
})
580585
})
586+
587+
describe("image handling", () => {
588+
it("should handle tool response with image content", async () => {
589+
const block: ToolUse = {
590+
type: "tool_use",
591+
name: "use_mcp_tool",
592+
params: {
593+
server_name: "figma-server",
594+
tool_name: "get_screenshot",
595+
arguments: '{"nodeId": "123"}',
596+
},
597+
nativeArgs: {
598+
server_name: "figma-server",
599+
tool_name: "get_screenshot",
600+
arguments: { nodeId: "123" },
601+
},
602+
partial: false,
603+
}
604+
605+
mockAskApproval.mockResolvedValue(true)
606+
607+
const mockToolResult = {
608+
content: [
609+
{
610+
type: "image",
611+
mimeType: "image/png",
612+
data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
613+
},
614+
],
615+
isError: false,
616+
}
617+
618+
mockProviderRef.deref.mockReturnValue({
619+
getMcpHub: () => ({
620+
callTool: vi.fn().mockResolvedValue(mockToolResult),
621+
getAllServers: vi
622+
.fn()
623+
.mockReturnValue([
624+
{
625+
name: "figma-server",
626+
tools: [{ name: "get_screenshot", description: "Get screenshot" }],
627+
},
628+
]),
629+
}),
630+
postMessageToWebview: vi.fn(),
631+
})
632+
633+
await useMcpToolTool.handle(mockTask as Task, block as any, {
634+
askApproval: mockAskApproval,
635+
handleError: mockHandleError,
636+
pushToolResult: mockPushToolResult,
637+
})
638+
639+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
640+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[1 image(s) received]", [
641+
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
642+
])
643+
expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)"))
644+
})
645+
646+
it("should handle tool response with both text and image content", async () => {
647+
const block: ToolUse = {
648+
type: "tool_use",
649+
name: "use_mcp_tool",
650+
params: {
651+
server_name: "figma-server",
652+
tool_name: "get_node_info",
653+
arguments: '{"nodeId": "123"}',
654+
},
655+
nativeArgs: {
656+
server_name: "figma-server",
657+
tool_name: "get_node_info",
658+
arguments: { nodeId: "123" },
659+
},
660+
partial: false,
661+
}
662+
663+
mockAskApproval.mockResolvedValue(true)
664+
665+
const mockToolResult = {
666+
content: [
667+
{ type: "text", text: "Node name: Button" },
668+
{
669+
type: "image",
670+
mimeType: "image/png",
671+
data: "base64imagedata",
672+
},
673+
],
674+
isError: false,
675+
}
676+
677+
mockProviderRef.deref.mockReturnValue({
678+
getMcpHub: () => ({
679+
callTool: vi.fn().mockResolvedValue(mockToolResult),
680+
getAllServers: vi
681+
.fn()
682+
.mockReturnValue([
683+
{ name: "figma-server", tools: [{ name: "get_node_info", description: "Get node info" }] },
684+
]),
685+
}),
686+
postMessageToWebview: vi.fn(),
687+
})
688+
689+
await useMcpToolTool.handle(mockTask as Task, block as any, {
690+
askApproval: mockAskApproval,
691+
handleError: mockHandleError,
692+
pushToolResult: mockPushToolResult,
693+
})
694+
695+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
696+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Node name: Button", [
697+
"data:image/png;base64,base64imagedata",
698+
])
699+
expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)"))
700+
})
701+
702+
it("should handle image with data URL already formatted", async () => {
703+
const block: ToolUse = {
704+
type: "tool_use",
705+
name: "use_mcp_tool",
706+
params: {
707+
server_name: "figma-server",
708+
tool_name: "get_screenshot",
709+
arguments: '{"nodeId": "123"}',
710+
},
711+
nativeArgs: {
712+
server_name: "figma-server",
713+
tool_name: "get_screenshot",
714+
arguments: { nodeId: "123" },
715+
},
716+
partial: false,
717+
}
718+
719+
mockAskApproval.mockResolvedValue(true)
720+
721+
const mockToolResult = {
722+
content: [
723+
{
724+
type: "image",
725+
mimeType: "image/jpeg",
726+
data: "data:image/jpeg;base64,/9j/4AAQSkZJRg==",
727+
},
728+
],
729+
isError: false,
730+
}
731+
732+
mockProviderRef.deref.mockReturnValue({
733+
getMcpHub: () => ({
734+
callTool: vi.fn().mockResolvedValue(mockToolResult),
735+
getAllServers: vi
736+
.fn()
737+
.mockReturnValue([
738+
{
739+
name: "figma-server",
740+
tools: [{ name: "get_screenshot", description: "Get screenshot" }],
741+
},
742+
]),
743+
}),
744+
postMessageToWebview: vi.fn(),
745+
})
746+
747+
await useMcpToolTool.handle(mockTask as Task, block as any, {
748+
askApproval: mockAskApproval,
749+
handleError: mockHandleError,
750+
pushToolResult: mockPushToolResult,
751+
})
752+
753+
// Should not double-prefix the data URL
754+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[1 image(s) received]", [
755+
"data:image/jpeg;base64,/9j/4AAQSkZJRg==",
756+
])
757+
})
758+
759+
it("should handle multiple images in response", async () => {
760+
const block: ToolUse = {
761+
type: "tool_use",
762+
name: "use_mcp_tool",
763+
params: {
764+
server_name: "figma-server",
765+
tool_name: "get_screenshots",
766+
arguments: '{"nodeIds": ["1", "2"]}',
767+
},
768+
nativeArgs: {
769+
server_name: "figma-server",
770+
tool_name: "get_screenshots",
771+
arguments: { nodeIds: ["1", "2"] },
772+
},
773+
partial: false,
774+
}
775+
776+
mockAskApproval.mockResolvedValue(true)
777+
778+
const mockToolResult = {
779+
content: [
780+
{
781+
type: "image",
782+
mimeType: "image/png",
783+
data: "image1data",
784+
},
785+
{
786+
type: "image",
787+
mimeType: "image/png",
788+
data: "image2data",
789+
},
790+
],
791+
isError: false,
792+
}
793+
794+
mockProviderRef.deref.mockReturnValue({
795+
getMcpHub: () => ({
796+
callTool: vi.fn().mockResolvedValue(mockToolResult),
797+
getAllServers: vi
798+
.fn()
799+
.mockReturnValue([
800+
{
801+
name: "figma-server",
802+
tools: [{ name: "get_screenshots", description: "Get screenshots" }],
803+
},
804+
]),
805+
}),
806+
postMessageToWebview: vi.fn(),
807+
})
808+
809+
await useMcpToolTool.handle(mockTask as Task, block as any, {
810+
askApproval: mockAskApproval,
811+
handleError: mockHandleError,
812+
pushToolResult: mockPushToolResult,
813+
})
814+
815+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[2 image(s) received]", [
816+
"data:image/png;base64,image1data",
817+
"data:image/png;base64,image2data",
818+
])
819+
expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 2 image(s)"))
820+
})
821+
})
581822
})

0 commit comments

Comments
 (0)