fix(server): Pi now sees screenshots returned by T3 MCP tools

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<RequestHook> {
   const handlers = new Map<string, RequestHook>();
   // 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) },
   });
@@ -57,3 +59,82 @@ describe("Pi upstream output-budget workaround", () => {
     );
   });
 });
+
+type RegisteredTool = {
+  readonly name: string;
+  readonly execute: (
+    toolCallId: string,
+    params: unknown,
+    signal?: AbortSignal,
+  ) => Promise<{ readonly content: ReadonlyArray<unknown>; readonly isError?: boolean }>;
+};
+
+// Loads the extension against an MCP endpoint that answers tools/call with `callResult`.
+async function loadMcpTool(callResult: unknown): Promise<RegisteredTool> {
+  const tools: RegisteredTool[] = [];
+  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(`${runnableSource}\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" }]);
+  });
+});

apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts

@@ -94,25 +94,47 @@ function jsonSchemaToTypebox(schema: Record<string, unknown> | 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<string, unknown>,
               signal,
             );
-            const text = formatMcpContent(result);
             return {
-              content: [{ type: "text", text }],
+              content: mcpToolContent(result),
               details: { server: "t3-code", tool: name },
               ...(isMcpToolError(result) ? { isError: true } : {}),
             };