Skip to content

fix: add image content support to MCP tool responses #10874

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
roomote wants to merge 1 commit into
base: main
Choose a base branch
from
+272 −12
Draft

fix: add image content support to MCP tool responses #10874

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 28 additions & 9 deletions src/core/tools/UseMcpToolTool.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,12 +250,14 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
})
}

private processToolContent(toolResult: any): string {
private processToolContent(toolResult: any): { text: string; images: string[] } {
if (!toolResult?.content || toolResult.content.length === 0) {
return ""
return { text: "", images: [] }
}

return toolResult.content
const images: string[] = []

const textContent = toolResult.content
.map((item: any) => {
if (item.type === "text") {
return item.text
Expand All @@ -264,10 +266,23 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
const { blob: _, ...rest } = item.resource
return JSON.stringify(rest, null, 2)
}
if (item.type === "image") {
// Handle image content (MCP image content has mimeType and data properties)
if (item.mimeType && item.data) {
if (item.data.startsWith("data:")) {
images.push(item.data)
} else {
images.push(`data:${item.mimeType};base64,${item.data}`)
}
}
return ""
}
return ""
})
.filter(Boolean)
.join("\n\n")

return { text: textContent, images }
}

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

let toolResultPretty = "(No response)"
let images: string[] = []

if (toolResult) {
const outputText = this.processToolContent(toolResult)
const { text: outputText, images: extractedImages } = this.processToolContent(toolResult)
images = extractedImages

if (outputText) {
if (outputText || images.length > 0) {
await this.sendExecutionStatus(task, {
executionId,
status: "output",
response: outputText,
response: outputText || (images.length > 0 ? `[${images.length} image(s)]` : ""),
})

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

// Send completion status
Expand All @@ -321,8 +340,8 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
})
}

await task.say("mcp_server_response", toolResultPretty)
pushToolResult(formatResponse.toolResult(toolResultPretty))
await task.say("mcp_server_response", toolResultPretty, images)
pushToolResult(formatResponse.toolResult(toolResultPretty, images))
}
}

Expand Down
247 changes: 244 additions & 3 deletions src/core/tools/__tests__/useMcpToolTool.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,12 @@ import { ToolUse } from "../../../shared/tools"
// Mock dependencies
vi.mock("../../prompts/responses", () => ({
formatResponse: {
toolResult: vi.fn((result: string) => `Tool result: ${result}`),
toolResult: vi.fn((result: string, images?: string[]) => {
if (images && images.length > 0) {
return `Tool result: ${result} [with ${images.length} image(s)]`
}
return `Tool result: ${result}`
}),
toolError: vi.fn((error: string) => `Tool error: ${error}`),
invalidMcpToolArgumentError: vi.fn((server: string, tool: string) => `Invalid args for ${server}:${tool}`),
unknownMcpToolError: vi.fn((server: string, tool: string, availableTools: string[]) => {
Expand Down Expand Up @@ -245,7 +250,7 @@ describe("useMcpToolTool", () => {
expect(mockTask.consecutiveMistakeCount).toBe(0)
expect(mockAskApproval).toHaveBeenCalled()
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully")
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully", [])
expect(mockPushToolResult).toHaveBeenCalledWith("Tool result: Tool executed successfully")
})

Expand Down Expand Up @@ -483,7 +488,7 @@ describe("useMcpToolTool", () => {
expect(mockTask.consecutiveMistakeCount).toBe(0)
expect(mockTask.recordToolError).not.toHaveBeenCalled()
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully")
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully", [])
})

it("should reject unknown server names with available servers listed", async () => {
Expand Down Expand Up @@ -578,4 +583,240 @@ describe("useMcpToolTool", () => {
expect(mockAskApproval).not.toHaveBeenCalled()
})
})

describe("image handling", () => {
it("should handle tool response with image content", async () => {
const block: ToolUse = {
type: "tool_use",
name: "use_mcp_tool",
params: {
server_name: "figma-server",
tool_name: "get_screenshot",
arguments: '{"nodeId": "123"}',
},
nativeArgs: {
server_name: "figma-server",
tool_name: "get_screenshot",
arguments: { nodeId: "123" },
},
partial: false,
}

mockAskApproval.mockResolvedValue(true)

const mockToolResult = {
content: [
{
type: "image",
mimeType: "image/png",
data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
},
],
isError: false,
}

mockProviderRef.deref.mockReturnValue({
getMcpHub: () => ({
callTool: vi.fn().mockResolvedValue(mockToolResult),
getAllServers: vi
.fn()
.mockReturnValue([
{
name: "figma-server",
tools: [{ name: "get_screenshot", description: "Get screenshot" }],
},
]),
}),
postMessageToWebview: vi.fn(),
})

await useMcpToolTool.handle(mockTask as Task, block as any, {
askApproval: mockAskApproval,
handleError: mockHandleError,
pushToolResult: mockPushToolResult,
})

expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[1 image(s) received]", [
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
])
expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)"))
})

it("should handle tool response with both text and image content", async () => {
const block: ToolUse = {
type: "tool_use",
name: "use_mcp_tool",
params: {
server_name: "figma-server",
tool_name: "get_node_info",
arguments: '{"nodeId": "123"}',
},
nativeArgs: {
server_name: "figma-server",
tool_name: "get_node_info",
arguments: { nodeId: "123" },
},
partial: false,
}

mockAskApproval.mockResolvedValue(true)

const mockToolResult = {
content: [
{ type: "text", text: "Node name: Button" },
{
type: "image",
mimeType: "image/png",
data: "base64imagedata",
},
],
isError: false,
}

mockProviderRef.deref.mockReturnValue({
getMcpHub: () => ({
callTool: vi.fn().mockResolvedValue(mockToolResult),
getAllServers: vi
.fn()
.mockReturnValue([
{ name: "figma-server", tools: [{ name: "get_node_info", description: "Get node info" }] },
]),
}),
postMessageToWebview: vi.fn(),
})

await useMcpToolTool.handle(mockTask as Task, block as any, {
askApproval: mockAskApproval,
handleError: mockHandleError,
pushToolResult: mockPushToolResult,
})

expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Node name: Button", [
"data:image/png;base64,base64imagedata",
])
expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 1 image(s)"))
})

it("should handle image with data URL already formatted", async () => {
const block: ToolUse = {
type: "tool_use",
name: "use_mcp_tool",
params: {
server_name: "figma-server",
tool_name: "get_screenshot",
arguments: '{"nodeId": "123"}',
},
nativeArgs: {
server_name: "figma-server",
tool_name: "get_screenshot",
arguments: { nodeId: "123" },
},
partial: false,
}

mockAskApproval.mockResolvedValue(true)

const mockToolResult = {
content: [
{
type: "image",
mimeType: "image/jpeg",
data: "data:image/jpeg;base64,/9j/4AAQSkZJRg==",
},
],
isError: false,
}

mockProviderRef.deref.mockReturnValue({
getMcpHub: () => ({
callTool: vi.fn().mockResolvedValue(mockToolResult),
getAllServers: vi
.fn()
.mockReturnValue([
{
name: "figma-server",
tools: [{ name: "get_screenshot", description: "Get screenshot" }],
},
]),
}),
postMessageToWebview: vi.fn(),
})

await useMcpToolTool.handle(mockTask as Task, block as any, {
askApproval: mockAskApproval,
handleError: mockHandleError,
pushToolResult: mockPushToolResult,
})

// Should not double-prefix the data URL
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[1 image(s) received]", [
"data:image/jpeg;base64,/9j/4AAQSkZJRg==",
])
})

it("should handle multiple images in response", async () => {
const block: ToolUse = {
type: "tool_use",
name: "use_mcp_tool",
params: {
server_name: "figma-server",
tool_name: "get_screenshots",
arguments: '{"nodeIds": ["1", "2"]}',
},
nativeArgs: {
server_name: "figma-server",
tool_name: "get_screenshots",
arguments: { nodeIds: ["1", "2"] },
},
partial: false,
}

mockAskApproval.mockResolvedValue(true)

const mockToolResult = {
content: [
{
type: "image",
mimeType: "image/png",
data: "image1data",
},
{
type: "image",
mimeType: "image/png",
data: "image2data",
},
],
isError: false,
}

mockProviderRef.deref.mockReturnValue({
getMcpHub: () => ({
callTool: vi.fn().mockResolvedValue(mockToolResult),
getAllServers: vi
.fn()
.mockReturnValue([
{
name: "figma-server",
tools: [{ name: "get_screenshots", description: "Get screenshots" }],
},
]),
}),
postMessageToWebview: vi.fn(),
})

await useMcpToolTool.handle(mockTask as Task, block as any, {
askApproval: mockAskApproval,
handleError: mockHandleError,
pushToolResult: mockPushToolResult,
})

expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "[2 image(s) received]", [
"data:image/png;base64,image1data",
"data:image/png;base64,image2data",
])
expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("with 2 image(s)"))
})
})
})
Loading