Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
9 changes: 8 additions & 1 deletion src/core/tools/UseMcpToolTool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -290,7 +290,14 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
return item.text
}
if (item.type === "resource") {
const { blob: _, ...rest } = item.resource
const { blob, ...rest } = item.resource
if (blob && item.resource.mimeType?.startsWith("image")) {
if (blob.startsWith("data:")) {
images.push(blob)
} else {
images.push(`data:${item.resource.mimeType};base64,${blob}`)
}
}
return JSON.stringify(rest, null, 2)
}
if (item.type === "image") {
Expand Down
177 changes: 177 additions & 0 deletions src/core/tools/__tests__/useMcpToolTool.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -859,6 +859,183 @@ describe("useMcpToolTool", () => {
])
})

it("should extract image data from embedded resource blobs", async () => {
const block: ToolUse<"use_mcp_tool"> = {
type: "tool_use",
name: "use_mcp_tool",
params: {
server_name: "godot-server",
tool_name: "game_screenshot",
arguments: "{}",
},
nativeArgs: {
server_name: "godot-server",
tool_name: "game_screenshot",
arguments: {},
},
partial: false,
}

mockAskApproval.mockResolvedValue(true)

const mockToolResult = {
content: [
{ type: "text", text: "Screenshot captured: 1152x648" },
{
type: "resource",
resource: {
uri: "godot://screenshot/latest",
mimeType: "image/png",
blob: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
},
},
],
isError: false,
}

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

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

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

it("should not extract non-image resource blobs", async () => {
const block: ToolUse<"use_mcp_tool"> = {
type: "tool_use",
name: "use_mcp_tool",
params: {
server_name: "godot-server",
tool_name: "read_scene",
arguments: "{}",
},
nativeArgs: {
server_name: "godot-server",
tool_name: "read_scene",
arguments: {},
},
partial: false,
}

mockAskApproval.mockResolvedValue(true)

const mockToolResult = {
content: [
{
type: "resource",
resource: {
uri: "godot://scene/main",
mimeType: "text/plain",
blob: "c2NlbmUgZGF0YQ==",
},
},
],
isError: false,
}

mockProviderRef.deref.mockReturnValue({
getMcpHub: () => ({
callTool: vi.fn().mockResolvedValue(mockToolResult),
getAllServers: vi.fn().mockReturnValue([
{
name: "godot-server",
tools: [{ name: "read_scene", description: "Read scene" }],
},
]),
}),
postMessageToWebview: vi.fn(),
})

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

expect(mockTask.say).toHaveBeenCalledWith(
"mcp_server_response",
expect.stringContaining("godot://scene/main"),
[],
)
})

it("should not double-prefix image resource blobs already formatted as data URLs", async () => {
const block: ToolUse<"use_mcp_tool"> = {
type: "tool_use",
name: "use_mcp_tool",
params: {
server_name: "godot-server",
tool_name: "game_screenshot",
arguments: "{}",
},
nativeArgs: {
server_name: "godot-server",
tool_name: "game_screenshot",
arguments: {},
},
partial: false,
}

mockAskApproval.mockResolvedValue(true)

const mockToolResult = {
content: [
{
type: "resource",
resource: {
uri: "godot://screenshot/latest",
mimeType: "image/png",
blob: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
},
},
],
isError: false,
}

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

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

expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", expect.any(String), [
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ",
])
})

it("should handle multiple images in response", async () => {
const block: ToolUse = {
type: "tool_use",
Expand Down
20 changes: 20 additions & 0 deletions webview-ui/src/components/chat/ChatRow.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -1577,6 +1577,19 @@ export const ChatRowContent = ({
/>
)
}
case "mcp_server_response":
return (
<div style={{ paddingTop: 10 }}>
<Markdown markdown={message.text} partial={message.partial} />
{message.images && message.images.length > 0 && (
<div style={{ marginTop: "10px" }}>
{message.images.map((image, index) => (
<ImageBlock key={index} imageData={image} />
))}
</div>
)}
</div>
)
default:
return (
<>
Expand All @@ -1588,6 +1601,13 @@ export const ChatRowContent = ({
)}
<div style={{ paddingTop: 10 }}>
<Markdown markdown={message.text} partial={message.partial} />
{message.images && message.images.length > 0 && (
<div style={{ marginTop: "10px" }}>
{message.images.map((image, index) => (
<ImageBlock key={index} imageData={image} />
))}
</div>
)}
</div>
</>
)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
import React from "react"

import { render, screen } from "@/utils/test-utils"
import { QueryClient, QueryClientProvider } from "@tanstack/react-query"
import { ExtensionStateContextProvider } from "@src/context/ExtensionStateContext"
import { ChatRowContent } from "../ChatRow"

// Mock i18n
vi.mock("react-i18next", () => ({
useTranslation: () => ({
t: (key: string) => key,
i18n: {
exists: () => false,
},
}),
Trans: ({ children }: { children?: React.ReactNode }) => <>{children}</>,
initReactI18next: { type: "3rdParty", init: () => {} },
}))

const queryClient = new QueryClient()

function renderChatRow(message: any) {
return render(
<ExtensionStateContextProvider>
<QueryClientProvider client={queryClient}>
<ChatRowContent
message={message}
isExpanded={false}
isLast={false}
isStreaming={false}
onToggleExpand={() => {}}
onSuggestionClick={() => {}}
onBatchFileResponse={() => {}}
onFollowUpUnmount={() => {}}
isFollowUpAnswered={false}
/>
</QueryClientProvider>
</ExtensionStateContextProvider>,
)
}

describe("ChatRow - mcp_server_response", () => {
it("renders images attached to the MCP server response", () => {
const message: any = {
type: "say",
say: "mcp_server_response",
ts: Date.now(),
partial: false,
text: "Screenshot captured: 1152x648",
images: ["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"],
}

renderChatRow(message)

const img = screen.getByRole("img")
expect(img).toHaveAttribute("src", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ")
})

it("renders multiple images attached to the MCP server response", () => {
const message: any = {
type: "say",
say: "mcp_server_response",
ts: Date.now(),
partial: false,
text: "[2 image(s) received]",
images: ["data:image/png;base64,image1data", "data:image/png;base64,image2data"],
}

renderChatRow(message)

const imgs = screen.getAllByRole("img")
expect(imgs).toHaveLength(2)
expect(imgs[0]).toHaveAttribute("src", "data:image/png;base64,image1data")
expect(imgs[1]).toHaveAttribute("src", "data:image/png;base64,image2data")
})

it("renders only the text when the MCP server response has no images", () => {
const message: any = {
type: "say",
say: "mcp_server_response",
ts: Date.now(),
partial: false,
text: "Plain text result",
}

renderChatRow(message)

expect(screen.queryByRole("img")).toBeNull()
})

it("renders attached images for say types that use the default renderer", () => {
const message: any = {
type: "say",
say: "command_output",
ts: Date.now(),
partial: false,
text: "Some output",
images: ["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"],
}

renderChatRow(message)

const img = screen.getByRole("img")
expect(img).toHaveAttribute("src", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ")
})
})
Loading