From b4f4e43feb3cf79e501867465e466c718e853c64 Mon Sep 17 00:00:00 2001 From: anntnzrb Date: Sun, 4 Oct 2026 22:47:53 -0500 Subject: [PATCH 1/2] fix(server): pass MCP image blocks through to Pi - Preserve image blocks from MCP tool results alongside text content instead of dropping them or flattening them into text JSON. - Add unit tests verifying Pi tool results include image blocks and ignore malformed images. --- .../Adapters/piT3McpExtensionSource.test.ts | 85 +++++++++++++++++++ .../Adapters/piT3McpExtensionSource.ts | 39 +++++++-- 2 files changed, 115 insertions(+), 9 deletions(-) diff --git a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts index 9dc9fa4526e6..cc0cdb89aa63 100644 --- a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts @@ -57,3 +57,88 @@ describe("Pi upstream output-budget workaround", () => { ); }); }); + +type RegisteredTool = { + readonly name: string; + readonly execute: ( + toolCallId: string, + params: unknown, + signal?: AbortSignal, + ) => Promise<{ readonly content: ReadonlyArray; readonly isError?: boolean }>; +}; + +// Loads the extension against an MCP endpoint that answers tools/call with `callResult`. +async function loadMcpTool(callResult: unknown): Promise { + const tools: RegisteredTool[] = []; + const source = NodeModule.stripTypeScriptTypes( + PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace( + "export default async function", + "async function", + ), + ); + const fetch = async (_url: string, init: { body: string }) => { + const { id, method } = JSON.parse(init.body) as { id?: number; method: string }; + const result = + method === "tools/list" + ? { tools: [{ name: "preview_snapshot", inputSchema: { type: "object" } }] } + : method === "tools/call" + ? callResult + : {}; + return new Response(id === undefined ? "" : JSON.stringify({ jsonrpc: "2.0", id, result }), { + headers: { "content-type": "application/json" }, + }); + }; + await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, { + process: { env: { T3_MCP_URL: "http://t3.test/mcp", T3_MCP_BEARER_TOKEN: "token" } }, + fetch, + AbortSignal, + Type: { Unsafe: (schema: unknown) => schema }, + pi: { + on: () => undefined, + registerTool: (tool: RegisteredTool) => tools.push(tool), + }, + }); + assert.deepEqual( + tools.map((tool) => tool.name), + ["mcp__t3-code__preview_snapshot"], + ); + return tools[0]!; +} + +describe("Pi MCP tool results", () => { + const image = { type: "image", data: "iVBORw0KGgo=", mimeType: "image/png" }; + + it("passes screenshot image blocks to Pi after the text", async () => { + const tool = await loadMcpTool({ + isError: false, + structuredContent: { url: "https://t3.codes" }, + content: [{ type: "text", text: '{"url":"https://t3.codes"}' }, image], + }); + const result = await tool.execute("call-1", {}); + assert.equal(result.content.length, 2); + assert.include(result.content[0], { type: "text" }); + assert.include((result.content[0] as { text: string }).text, '{"url":"https://t3.codes"}'); + assert.deepEqual(result.content[1], image); + assert.isUndefined(result.isError); + }); + + it("keeps text-only results text-only", async () => { + const tool = await loadMcpTool({ content: [{ type: "text", text: "done" }] }); + assert.deepEqual((await tool.execute("call-1", {})).content, [{ type: "text", text: "done" }]); + }); + + it("sends an image-only result as the image instead of its base64 JSON", async () => { + const tool = await loadMcpTool({ content: [image] }); + assert.deepEqual((await tool.execute("call-1", {})).content, [image]); + }); + + it("drops malformed image blocks", async () => { + const tool = await loadMcpTool({ + content: [ + { type: "text", text: "shot" }, + { type: "image", data: "iVBORw0KGgo=" }, + ], + }); + assert.deepEqual((await tool.execute("call-1", {})).content, [{ type: "text", text: "shot" }]); + }); +}); diff --git a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts index 185bd16f9745..11667ee03202 100644 --- a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts +++ b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts @@ -94,25 +94,47 @@ function jsonSchemaToTypebox(schema: Record | undefined) { return Type.Object({}, { additionalProperties: true }); } -function formatMcpContent(result: unknown): string { - if (result === null || result === undefined) return ""; - if (typeof result !== "object") return String(result); +type ToolContent = + | { readonly type: "text"; readonly text: string } + | { readonly type: "image"; readonly data: string; readonly mimeType: string }; + +/** + * Maps an MCP tool result to Pi tool-result content: the text as one block, + * then any images (such as preview screenshots). Pi swaps images for a + * placeholder when the active model has no image input. + */ +function mcpToolContent(result: unknown): ToolContent[] { + if (result === null || result === undefined) return [{ type: "text", text: "" }]; + if (typeof result !== "object") return [{ type: "text", text: String(result) }]; const record = result as { - readonly content?: ReadonlyArray<{ readonly type?: string; readonly text?: string }>; + readonly content?: ReadonlyArray<{ + readonly type?: string; + readonly text?: string; + readonly data?: string; + readonly mimeType?: string; + }>; readonly structuredContent?: unknown; - readonly isError?: boolean; }; const texts: string[] = []; + const images: ToolContent[] = []; if (Array.isArray(record.content)) { for (const part of record.content) { if (part?.type === "text" && typeof part.text === "string") texts.push(part.text); + if ( + part?.type === "image" && + typeof part.data === "string" && + typeof part.mimeType === "string" + ) { + images.push({ type: "image", data: part.data, mimeType: part.mimeType }); + } } } if (record.structuredContent !== undefined) { texts.push(JSON.stringify(record.structuredContent)); } - if (texts.length > 0) return texts.join("\\n"); - return JSON.stringify(result); + if (texts.length > 0) return [{ type: "text", text: texts.join("\\n") }, ...images]; + if (images.length > 0) return images; + return [{ type: "text", text: JSON.stringify(result) }]; } function isMcpToolError(result: unknown): boolean { @@ -286,9 +308,8 @@ export default async function t3McpExtension(pi: ExtensionAPI) { (params ?? {}) as Record, signal, ); - const text = formatMcpContent(result); return { - content: [{ type: "text", text }], + content: mcpToolContent(result), details: { server: "t3-code", tool: name }, ...(isMcpToolError(result) ? { isError: true } : {}), }; From af9ad9db95159649a78f2e5f3b3cf9a993436367 Mon Sep 17 00:00:00 2001 From: anntnzrb Date: Sun, 4 Oct 2026 22:49:44 -0500 Subject: [PATCH 2/2] test(server): deduplicate runnable extension source in Pi tests - Extract stripped TypeScript source preparation to a top-level constant across test helper functions. --- .../Adapters/piT3McpExtensionSource.test.ts | 24 ++++++++----------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts index cc0cdb89aa63..2bb0fc2da733 100644 --- a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts @@ -4,6 +4,14 @@ import { assert, describe, it } from "@effect/vitest"; import { PI_T3_MCP_EXTENSION_SOURCE } from "./piT3McpExtensionSource.ts"; +// The shipped extension as a plain script. Tests supply Typebox through the VM context. +const runnableSource = NodeModule.stripTypeScriptTypes( + PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace( + "export default async function", + "async function", + ), +); + type RequestHook = ( event: { payload: unknown }, ctx: { model: { provider: string } }, @@ -12,13 +20,7 @@ type RequestHook = ( async function loadRequestHook(): Promise { const handlers = new Map(); // Execute the shipped extension with MCP disabled; this path needs no Typebox. - const source = NodeModule.stripTypeScriptTypes( - PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace( - "export default async function", - "async function", - ), - ); - await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, { + await NodeVM.runInNewContext(`${runnableSource}\nt3McpExtension(pi)`, { process: { env: {} }, pi: { on: (name: string, handler: RequestHook) => handlers.set(name, handler) }, }); @@ -70,12 +72,6 @@ type RegisteredTool = { // Loads the extension against an MCP endpoint that answers tools/call with `callResult`. async function loadMcpTool(callResult: unknown): Promise { const tools: RegisteredTool[] = []; - const source = NodeModule.stripTypeScriptTypes( - PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace( - "export default async function", - "async function", - ), - ); const fetch = async (_url: string, init: { body: string }) => { const { id, method } = JSON.parse(init.body) as { id?: number; method: string }; const result = @@ -88,7 +84,7 @@ async function loadMcpTool(callResult: unknown): Promise { headers: { "content-type": "application/json" }, }); }; - await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, { + await NodeVM.runInNewContext(`${runnableSource}\nt3McpExtension(pi)`, { process: { env: { T3_MCP_URL: "http://t3.test/mcp", T3_MCP_BEARER_TOKEN: "token" } }, fetch, AbortSignal,