From 38c187bfb7a0020db36ccc1479e7c3501b205fb8 Mon Sep 17 00:00:00 2001 From: Wout Stiens <71498452+StiensWout@users.noreply.github.com> Date: Sat, 10 Oct 2026 23:22:40 +0200 Subject: [PATCH] fix(pi): preserve tool images and structured results (#17851) --- .../provider-pi/src/server/adapter.test.ts | 301 +++++++++++++++++- packages/provider-pi/src/server/adapter.ts | 103 +++++- .../src/server/mcpBridge.testkit.ts | 140 ++++++++ .../src/server/mcpExtensionSource.test.ts | 175 ++++------ .../src/server/mcpExtensionSource.ts | 71 ++++- packages/provider-pi/src/server/rpc.ts | 10 +- 6 files changed, 659 insertions(+), 141 deletions(-) create mode 100644 packages/provider-pi/src/server/mcpBridge.testkit.ts diff --git a/packages/provider-pi/src/server/adapter.test.ts b/packages/provider-pi/src/server/adapter.test.ts index e13f8213077f58fd6107a7dcbee6303f0e093ee3..a5e2d53884eeaf4217499129e33d20a8a23d8823 100644 --- a/packages/provider-pi/src/server/adapter.test.ts +++ b/packages/provider-pi/src/server/adapter.test.ts @@ -39,6 +39,7 @@ import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator"; import * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; import { handoffBudget } from "@t3tools/provider-core/server/handoffBudget"; import * as HostProcess from "@t3tools/shared/HostProcess"; +import { toolOutputImages, compactDynamicToolOutput } from "@t3tools/shared/toolOutput"; import { makePiAdapterV2, PiAdapterV2Driver, @@ -46,6 +47,8 @@ import { type PiAdapterV2Options, } from "./adapter.ts"; import { makePiRpcConnection, type PiRpcRecord } from "./rpc.ts"; +import { loadMcpBridge } from "./mcpBridge.testkit.ts"; +import { turnItemOutputText } from "../../../client-runtime/src/work-log/itemDetail.ts"; const layerTest = Layer.mergeAll( NodeServices.layer, @@ -77,7 +80,7 @@ const modelSelection = (model: string): ModelSelection => ({ interface FakePi { readonly spawner: ChildProcessSpawner.ChildProcessSpawner["Service"]; - readonly emit: (record: PiRpcRecord) => Effect.Effect; + readonly emit: (record: PiRpcRecord, chunkBytes?: number) => Effect.Effect; readonly takeRequest: (type: string) => Effect.Effect; /** Data returned by the next `get_entries` acks, consumed in order. */ readonly queueEntries: (data: unknown) => void; @@ -153,10 +156,14 @@ const makeFakePi: Effect.Effect = Effect.gen(function* () { let models: ReadonlyArray = []; let stdinBuffer = ""; - const emit = (record: PiRpcRecord) => - Queue.offer(stdout, new TextEncoder().encode(`${encodeJsonLine(record)}\n`)).pipe( - Effect.asVoid, - ); + const emit = (record: PiRpcRecord, chunkBytes?: number) => + Effect.gen(function* () { + const bytes = new TextEncoder().encode(`${encodeJsonLine(record)}\n`); + const size = chunkBytes ?? bytes.byteLength; + for (let offset = 0; offset < bytes.byteLength; offset += size) { + yield* Queue.offer(stdout, bytes.subarray(offset, offset + size)); + } + }); const respondTo = (record: PiRpcRecord): PiRpcRecord | null => { if (typeof record["id"] !== "string") return null; @@ -1852,6 +1859,284 @@ describe("PiAdapterV2", () => { }).pipe(Effect.scoped, Effect.provide(layerTest)), ); + it.effect( + "keeps native and MCP images available to the shared asset reader with typed results", + () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const { runtime, takeEvent } = yield* openRuntime(fake); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* startTurn(runtime, providerThread); + yield* fake.takeRequest("prompt"); + yield* fake.emit({ type: "agent_start" }); + const image = { + type: "image", + mimeType: "image/png", + data: Buffer.alloc(64 * 1024, 0x61).toString("base64"), + }; + const note = "Read image file [image/png]"; + const content = [{ type: "text", text: note }, image]; + const structured = { count: 2, ready: true, paths: ["one.png", "two.png"] }; + const scriptResult = { content, structuredContent: { threadId: "child-thread" } }; + const cases = [ + { + toolName: "read", + result: { content, structuredContent: { ...image, note } }, + expected: { content }, + }, + { + toolName: "image_generate", + result: { content, structuredContent: structured }, + expected: { content, structuredContent: structured }, + }, + { + toolName: "mcp__t3-code__preview_snapshot", + result: { + content, + structuredContent: scriptResult, + details: { server: "t3-code", tool: "preview_snapshot" }, + }, + expected: { content, structuredContent: { threadId: "child-thread" } }, + }, + { + toolName: "structured_tool", + result: { content: [], structuredContent: [false, 2, null, { ready: true }] }, + expected: { content: [], structuredContent: [false, 2, null, { ready: true }] }, + }, + { + toolName: "mcp__custom__extension", + result: { content, structuredContent: { content: ["domain content"], ready: true } }, + expected: { content, structuredContent: { content: ["domain content"], ready: true } }, + }, + { + toolName: "mcp__t3-code__preview_snapshot", + result: { + content: [ + { type: "text", text: "first" }, + { type: "text", text: "second" }, + ], + structuredContent: { + content: [ + { type: "text", text: "first" }, + { type: "text", text: "second" }, + ], + }, + details: { server: "t3-code", tool: "preview_snapshot" }, + }, + expected: { + content: [ + { type: "text", text: "first" }, + { type: "text", text: "second" }, + ], + }, + }, + ]; + for (const [index, test] of cases.entries()) { + for (const phase of ["update", "end"] as const) { + yield* fake.emit({ + type: `tool_execution_${phase}`, + toolCallId: `image-${index}`, + toolName: test.toolName, + ...(phase === "end" ? { result: test.result } : { partialResult: test.result }), + isError: false, + }); + const event = yield* takeEvent( + (event) => + event.type === "turn_item.updated" && event.turnItem.type === "dynamic_tool", + ); + if (event.type !== "turn_item.updated" || event.turnItem.type !== "dynamic_tool") + return yield* Effect.die("Expected an image tool item"); + assert.equal(event.turnItem.status, phase === "end" ? "completed" : "running"); + assert.deepEqual(event.turnItem.output, test.expected); + assert.deepEqual( + toolOutputImages(event.turnItem.output), + test.expected.content.some((block) => block.type === "image") + ? [{ mimeType: "image/png", data: image.data }] + : [], + ); + if (index === 0) assert.equal(turnItemOutputText(event.turnItem), note); + if (index === 5) assert.equal(turnItemOutputText(event.turnItem), "first\nsecond"); + if (index === 2) { + assert.deepEqual(compactDynamicToolOutput(event.turnItem.output), { + threadId: "child-thread", + }); + assert.equal(JSON.stringify(event.turnItem.output).split(image.data).length - 1, 1); + } + } + } + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect( + "keeps large tool text readable, drops identical read mirrors, and bounds structured values", + () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const { runtime, takeEvent } = yield* openRuntime(fake); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* startTurn(runtime, providerThread); + yield* fake.takeRequest("prompt"); + yield* fake.emit({ type: "agent_start" }); + const text = "hello\nworld\n" + "🙂".repeat(40_000); + const content = [{ type: "text", text }]; + const cases = [ + { content, structuredContent: text }, + { content, structuredContent: { oversized: "x".repeat(128 * 1024) } }, + { content, structuredContent: Array.from({ length: 10_000 }, () => 1) }, + { + content, + structuredContent: Array.from({ length: 40 }).reduce( + (nested) => ({ nested }), + null, + ), + }, + ]; + for (const [index, result] of cases.entries()) { + yield* fake.emit({ + type: "tool_execution_end", + toolCallId: `structured-${index}`, + toolName: "read", + result, + isError: true, + }); + const event = yield* takeEvent( + (event) => event.type === "turn_item.updated" && event.turnItem.type === "dynamic_tool", + ); + if (event.type !== "turn_item.updated" || event.turnItem.type !== "dynamic_tool") + return yield* Effect.die("Expected a tool item"); + assert.equal(event.turnItem.status, "failed"); + assert.notProperty(event.turnItem.output, "structuredContent"); + const displayed = turnItemOutputText(event.turnItem); + if (index === 0) { + assert.equal(event.turnItem.output, text); + assert.equal(displayed, text); + } else { + assert.equal( + displayed, + text + "\nStructured output omitted because it exceeds the stored result limit.", + ); + assert.deepEqual(Object.keys(event.turnItem.output as object), ["content"]); + } + } + // Distinct typed metadata remains stored even when readable text exists. + yield* fake.emit({ + type: "tool_execution_end", + toolCallId: "distinct-metadata", + toolName: "read", + result: { content, structuredContent: { path: "read.ts", ready: true } }, + }); + const event = yield* takeEvent( + (event) => event.type === "turn_item.updated" && event.turnItem.type === "dynamic_tool", + ); + if (event.type !== "turn_item.updated" || event.turnItem.type !== "dynamic_tool") + return yield* Effect.die("Expected a tool item"); + assert.deepEqual(event.turnItem.output, { + content, + structuredContent: { path: "read.ts", ready: true }, + }); + assert.equal(turnItemOutputText(event.turnItem), text); + const note = "Read image file [image/png]"; + const image = { type: "image", data: "AAAA", mimeType: "image/png" }; + const imageContent = [{ type: "text", text: note }, image]; + // An extension can attach distinct typed metadata to a read result. + for (const structuredContent of [ + { ...image, note, width: 1200 }, + { ...image, note: "A distinct note" }, + { ...image, note, data: "AQID" }, + ]) { + yield* fake.emit({ + type: "tool_execution_end", + toolCallId: "distinct-image-metadata", + toolName: "read", + result: { content: imageContent, structuredContent }, + }); + const event = yield* takeEvent( + (event) => event.type === "turn_item.updated" && event.turnItem.type === "dynamic_tool", + ); + if (event.type !== "turn_item.updated" || event.turnItem.type !== "dynamic_tool") + return yield* Effect.die("Expected a read image tool item"); + assert.deepEqual(event.turnItem.output, { content: imageContent, structuredContent }); + assert.equal(turnItemOutputText(event.turnItem), note); + } + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + + it.effect( + "delivers realistic MCP screenshots through the bridge, JSONL transport, and adapter", + () => + Effect.gen(function* () { + const fake = yield* makeFakePi; + const { runtime, takeEvent } = yield* openRuntime(fake); + const providerThread = yield* runtime.ensureThread({ + threadId: THREAD_ID, + modelSelection: modelSelection("default"), + runtimePolicy, + }); + yield* startTurn(runtime, providerThread); + yield* fake.takeRequest("prompt"); + yield* fake.emit({ type: "agent_start" }); + for (const size of [3.25 * 1024 * 1024, 10 * 1024 * 1024]) { + const pixels = Buffer.alloc(size, 0x61); + const data = pixels.toString("base64"); + const content = [ + { type: "text", text: "Screenshot captured" }, + { type: "image", data, mimeType: "image/png" }, + ]; + const metadata = { screenshot: { width: 1200, height: 800 }, ready: true }; + const bridge = yield* Effect.promise(() => + loadMcpBridge({ modern: true, result: { content, structuredContent: metadata } }), + ); + const tool = bridge.tools.find((tool) => tool.name === "mcp__t3-code__preview_snapshot")!; + const result = yield* Effect.promise(() => tool.execute("screenshot", {})); + assert.deepEqual(result.structuredContent, { content, structuredContent: metadata }); + yield* fake.emit( + { + type: "tool_execution_end", + toolCallId: "screenshot", + toolName: tool.name, + result, + isError: false, + }, + 64 * 1024, + ); + // A following small completion proves a dropped screenshot immediately, + // without waiting for a timeout when framing loses the large event. + yield* fake.emit({ + type: "tool_execution_end", + toolCallId: "after-screenshot", + toolName: "marker", + result: { content: [{ type: "text", text: "after" }] }, + }); + const event = yield* takeEvent( + (event) => event.type === "turn_item.updated" && event.turnItem.type === "dynamic_tool", + ); + if (event.type !== "turn_item.updated" || event.turnItem.type !== "dynamic_tool") + return yield* Effect.die("Expected a screenshot tool item"); + assert.equal(event.turnItem.toolName, tool.name); + assert.equal(event.turnItem.status, "completed"); + assert.deepEqual(event.turnItem.output, { content, structuredContent: metadata }); + assert.equal(turnItemOutputText(event.turnItem), "Screenshot captured"); + const image = toolOutputImages(event.turnItem.output)[0]; + assert.equal(image?.mimeType, "image/png"); + assert.deepEqual(Buffer.from(image!.data!, "base64"), pixels); + yield* takeEvent( + (event) => + event.type === "turn_item.updated" && + event.turnItem.type === "dynamic_tool" && + event.turnItem.toolName === "marker", + ); + } + }).pipe(Effect.scoped, Effect.provide(layerTest)), + ); + it.effect("observes official subagent results without inventing child threads", () => Effect.gen(function* () { const fake = yield* makeFakePi; @@ -3512,8 +3797,10 @@ describe("PiRpc framing", () => { yield* push('{"type":"agent_'); yield* push('start"}\r\n{"type":"agent_settled"}\nnot json\n{"type":"queue_update"}\n'); - yield* push("x".repeat(8 * 1024 * 1024)); - yield* push('x{"type":"must_not_emit"}\n{"type":"after_oversized"}\n'); + yield* push('{"type":"must_not_emit","text":"'); + const chunk = new TextEncoder().encode("x".repeat(8 * 1024 * 1024)); + for (let index = 0; index < 32; index++) yield* Queue.offer(stdout, chunk); + yield* push('"}\n{"type":"after_oversized"}\n'); const first = yield* Queue.take(connection.events); assert.equal(first["type"], "agent_start"); diff --git a/packages/provider-pi/src/server/adapter.ts b/packages/provider-pi/src/server/adapter.ts index 7fd91ebc760933cc31f554d53fc7d14038b08abd..973ef4ebb6aa1321296ea76aaee94f6c19779683 100644 --- a/packages/provider-pi/src/server/adapter.ts +++ b/packages/provider-pi/src/server/adapter.ts @@ -26,6 +26,8 @@ import * as HostProcess from "@t3tools/shared/HostProcess"; import { AgentScope } from "@t3tools/shared/AgentScope"; import { getModelSelectionStringOptionValue } from "@t3tools/shared/model"; +import { readToolOutputImage } from "@t3tools/shared/toolOutput"; +import * as Predicate from "effect/Predicate"; import { defaultInstanceIdForDriver, ProviderDriverKind, @@ -247,6 +249,104 @@ function contentText(content: unknown): string { .join(""); } +const MAX_STRUCTURED_TOOL_OUTPUT_BYTES = 64 * 1024; + +/** Preserve complete structured JSON only when its size and traversal remain bounded. */ +function isBoundedStructuredToolValue(value: unknown): boolean { + let bytes = MAX_STRUCTURED_TOOL_OUTPUT_BYTES; + let nodes = 4_096; + const visit = (entry: unknown, depth: number): boolean => { + if (depth > 32 || --nodes < 0 || bytes < 0) return false; + if (Array.isArray(entry)) { + bytes -= 2 + entry.length; + return entry.every((item) => visit(item, depth + 1)); + } + if (Predicate.isObject(entry)) { + bytes -= 2; + for (const key in entry) { + if (!Object.hasOwn(entry, key)) continue; + if (key.length > bytes) return false; + bytes -= Buffer.byteLength(JSON.stringify(key), "utf8") + 2; + if (!visit(entry[key], depth + 1)) return false; + } + return bytes >= 0; + } + if (typeof entry === "string" && entry.length > bytes) return false; + bytes -= Buffer.byteLength(JSON.stringify(entry) ?? "null", "utf8"); + return bytes >= 0; + }; + return visit(value, 0); +} + +/** Keep model content and distinct typed values; server ingestion owns image byte policy. */ +function piToolOutput(result: unknown, toolName: string): unknown { + const rawContent = recordField(result, "content"); + const content = Array.isArray(rawContent) + ? rawContent + : typeof rawContent === "string" + ? [{ type: "text", text: rawContent }] + : []; + let structured = recordField(result, "structuredContent"); + // Pi MCP tools return the full CallToolResult to scripts. The model content + // is already captured, so store only its distinct structured payload. + const details = recordField(result, "details"); + const isMcpResult = + toolName.startsWith("mcp__") && + recordString(details, "server") !== undefined && + recordString(details, "tool") !== undefined && + Array.isArray(recordField(structured, "content")); + if (isMcpResult) { + structured = recordField(structured, "structuredContent"); + } + if (typeof structured === "string" && structured === contentText(content)) { + structured = undefined; + } + // Native read repeats its image and first text note as structured output. + // Drop only that exact mirror; an extension's extra metadata stays distinct. + if (toolName === "read" && Predicate.isObject(structured)) { + const image = readToolOutputImage(structured); + const note = + recordString( + content.find((block) => recordField(block, "type") === "text"), + "text", + ) ?? ""; + if ( + image !== null && + Object.keys(structured).every((key) => ["type", "data", "mimeType", "note"].includes(key)) && + recordString(structured, "note") === note && + content.some((block) => { + const match = readToolOutputImage(block); + return match !== null && match.mimeType === image.mimeType && match.data === image.data; + }) + ) { + structured = undefined; + } + } + if (structured !== undefined && !isBoundedStructuredToolValue(structured)) { + return { + content: [ + ...content, + { + type: "text", + text: "Structured output omitted because it exceeds the stored result limit.", + }, + ], + }; + } + if ( + !isMcpResult && + structured === undefined && + content.every((block) => recordField(block, "type") === "text") + ) { + const text = contentText(content); + return text.length > 0 ? text : undefined; + } + return { + content, + ...(structured === undefined ? {} : { structuredContent: structured }), + }; +} + function providerRef( nativeId: string, strength: "strong" | "weak" = "strong", @@ -1047,6 +1147,7 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* ( return; } } + const output = piToolOutput(resultRecord, toolName); yield* emit({ type: "turn_item.updated", driver: PI_PROVIDER, @@ -1057,7 +1158,7 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* ( ...mcpToolPresentation({ toolName }), toolName, input: args ?? {}, - ...(outputText.length > 0 ? { output: outputText } : {}), + ...(output === undefined ? {} : { output }), }, }); if (toolName === "subagent") { diff --git a/packages/provider-pi/src/server/mcpBridge.testkit.ts b/packages/provider-pi/src/server/mcpBridge.testkit.ts new file mode 100644 index 0000000000000000000000000000000000000000..6061c2a3b556b97731a4d331ca83ec8841801ca5 --- /dev/null +++ b/packages/provider-pi/src/server/mcpBridge.testkit.ts @@ -0,0 +1,140 @@ +// @effect-diagnostics nodeBuiltinImport:off -- Executes the shipped Pi extension at its native JavaScript boundary. +import * as NodeModule from "node:module"; +import * as NodeVM from "node:vm"; +import { PI_T3_MCP_EXTENSION_SOURCE } from "./mcpExtensionSource.ts"; + +interface RegisteredTool { + readonly name: string; + readonly description: string; + readonly parameters: unknown; + readonly outputSchema?: unknown; + readonly exposure?: string; + readonly promptSnippet?: string; + readonly promptGuidelines?: ReadonlyArray; + readonly execute: ( + id: string, + args: Record, + signal?: AbortSignal, + ) => Promise<{ + readonly content: ReadonlyArray<{ + readonly type: string; + readonly text?: string; + readonly data?: string; + readonly mimeType?: string; + }>; + readonly structuredContent?: unknown; + readonly isError?: boolean; + }>; +} + +type AgentStartHook = ( + event: { systemPrompt: string }, + ctx: { ui: { notify: (message: string, severity: string) => void } }, +) => Promise<{ systemPrompt: string }>; + +export async function loadMcpBridge( + options: { + readonly modern?: boolean; + readonly result?: unknown; + readonly toolSearchAvailable?: boolean; + readonly toolSearchDisabled?: boolean; + readonly allowsTool?: (name: string) => boolean; + } = {}, +) { + const handlers = new Map(); + const tools: RegisteredTool[] = []; + const requests: Array<{ readonly method: string; readonly params?: unknown }> = []; + let activeTools = ["read"]; + const transports: Array<{ + readonly url: string; + readonly authorization: string; + readonly signal: AbortSignal | undefined; + }> = []; + const servers: Array<{ readonly name: string; readonly config: Record }> = []; + const catalog = [ + { name: "orchestrator_capabilities", description: "Discover available providers and models." }, + { name: "delegate_task", description: "Delegate work to another agent." }, + { name: "task_status", description: "Check delegated work." }, + { name: "preview_snapshot", description: "Inspect the collaborative browser." }, + ].map((tool) => ({ + ...tool, + inputSchema: { type: "object", properties: {} }, + outputSchema: { type: "object", properties: { count: { type: "number" } } }, + })); + const source = NodeModule.stripTypeScriptTypes( + PI_T3_MCP_EXTENSION_SOURCE.replace(/^import .*;$/gm, "").replace( + "export default async function", + "async function", + ), + ); + await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, { + process: { + env: { T3_MCP_URL: "http://fixture.invalid/mcp", T3_MCP_BEARER_TOKEN: "fixture-token" }, + }, + AbortSignal, + Type: { Unsafe: (schema: unknown) => schema }, + fetch: async ( + url: string, + transport: { body: string; headers: Record; signal?: AbortSignal }, + ) => { + transport.signal?.throwIfAborted(); + transports.push({ + url, + authorization: transport.headers.authorization!, + signal: transport.signal, + }); + const request = JSON.parse(transport.body) as { + id: number; + method: string; + params?: unknown; + }; + requests.push(request); + const result = + request.method === "tools/list" + ? { tools: catalog } + : request.method === "tools/call" + ? (options.result ?? { content: [{ type: "text", text: "browser snapshot" }] }) + : {}; + return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result }), { + headers: { "content-type": "application/json" }, + }); + }, + pi: { + on: (name: string, handler: AgentStartHook) => handlers.set(name, handler), + registerTool: (tool: RegisteredTool) => { + if (options.allowsTool && !options.allowsTool(tool.name)) return; + const index = tools.findIndex((current) => current.name === tool.name); + if (index === -1) tools.push(tool); + else tools[index] = tool; + }, + getActiveTools: () => activeTools, + setActiveTools: (names: string[]) => { + activeTools = names.filter((name) => options.allowsTool?.(name) ?? true); + }, + getAllTools: () => + options.toolSearchAvailable && + !options.toolSearchDisabled && + (options.allowsTool?.("tool_search") ?? true) + ? [{ name: "tool_search", sourceInfo: { path: "builtin:tool-search" } }] + : [], + ...(options.modern + ? { + registerMcpServer: (name: string, config: Record) => + servers.push({ name, config }), + unregisterMcpServer: () => servers.splice(0), + } + : {}), + }, + }); + return { + handlers, + tools, + requests, + servers, + transports, + getActiveTools: () => activeTools, + restoreActiveTools: (names: string[]) => { + activeTools = names; + }, + }; +} diff --git a/packages/provider-pi/src/server/mcpExtensionSource.test.ts b/packages/provider-pi/src/server/mcpExtensionSource.test.ts index 25b040c2c002ed99eb2fbe238c43b711e1dc4a4c..8fc7889a7937cc40e5aba3cc99c887df5f37860c 100644 --- a/packages/provider-pi/src/server/mcpExtensionSource.test.ts +++ b/packages/provider-pi/src/server/mcpExtensionSource.test.ts @@ -4,9 +4,10 @@ import * as NodePath from "node:path"; import * as NodeOS from "node:os"; import * as NodeModule from "node:module"; import * as NodeVM from "node:vm"; -import { assert, describe, it } from "@effect/vitest"; +import { assert, describe, expect, it } from "@effect/vitest"; import { PI_T3_MCP_EXTENSION_SOURCE } from "./mcpExtensionSource.ts"; +import { loadMcpBridge } from "./mcpBridge.testkit.ts"; type RequestHook = ( event: { payload: unknown }, @@ -47,123 +48,67 @@ async function loadHooks(commands: ReadonlyArray = []) { return handlers; } -interface RegisteredTool { - readonly name: string; - readonly description: string; - readonly parameters: unknown; - readonly exposure?: string; - readonly promptSnippet?: string; - readonly promptGuidelines?: ReadonlyArray; - readonly execute: ( - id: string, - args: Record, - signal?: AbortSignal, - ) => Promise<{ - readonly content: ReadonlyArray<{ readonly type: string; readonly text: string }>; - }>; -} +describe("Pi MCP tool results", () => { + it("returns mixed screenshot blocks and typed script output without repeating structured text", async () => { + const content = [ + { type: "text", text: "Browser screenshot" }, + { type: "image", data: "AAAA", mimeType: "image/png", _meta: { private: true } }, + { type: "text", text: "Screenshot captured" }, + ] as const; + const result = { content, structuredContent: { count: 2 }, _meta: { private: true } }; + const bridge = await loadMcpBridge({ modern: true, result }); + const tool = bridge.tools.find((tool) => tool.name === "mcp__t3-code__preview_snapshot")!; + const output = await tool.execute("snapshot", {}); + assert.deepEqual(output.content, [ + content[0], + { type: "image", data: "AAAA", mimeType: "image/png" }, + content[2], + ]); + assert.deepEqual(output.structuredContent, { content, structuredContent: { count: 2 } }); + assert.deepEqual(tool.outputSchema, { + type: "object", + properties: { + content: { type: "array", items: { type: "object" } }, + structuredContent: { type: "object", properties: { count: { type: "number" } } }, + isError: { type: "boolean" }, + _meta: { type: "object" }, + }, + required: ["content"], + }); + }); -type AgentStartHook = ( - event: { systemPrompt: string }, - ctx: { ui: { notify: (message: string, severity: string) => void } }, -) => Promise<{ systemPrompt: string }>; - -async function loadMcpBridge( - options: { - readonly modern?: boolean; - readonly toolSearchAvailable?: boolean; - readonly toolSearchDisabled?: boolean; - readonly allowsTool?: (name: string) => boolean; - } = {}, -) { - const handlers = new Map(); - const tools: RegisteredTool[] = []; - const requests: Array<{ readonly method: string; readonly params?: unknown }> = []; - let activeTools = ["read"]; - const transports: Array<{ - readonly url: string; - readonly authorization: string; - readonly signal: AbortSignal | undefined; - }> = []; - const servers: Array<{ readonly name: string; readonly config: Record }> = []; - const catalog = [ - { name: "orchestrator_capabilities", description: "Discover available providers and models." }, - { name: "delegate_task", description: "Delegate work to another agent." }, - { name: "task_status", description: "Check delegated work." }, - { name: "preview_snapshot", description: "Inspect the collaborative browser." }, - ].map((tool) => ({ ...tool, inputSchema: { type: "object", properties: {} } })); - const source = NodeModule.stripTypeScriptTypes( - PI_T3_MCP_EXTENSION_SOURCE.replace(/^import .*;$/gm, "").replace( - "export default async function", - "async function", - ), + it.each([false, true])( + "preserves image-only results and error status, isError=%s", + async (isError) => { + const image = { type: "image", data: "AAAA", mimeType: "image/png" }; + const result = { content: [image], isError }; + const bridge = await loadMcpBridge({ result }); + const output = await bridge.tools[0]!.execute("image", {}); + assert.deepEqual(output.content[0], image); + assert.equal(output.content.length, isError ? 2 : 1); + if (isError) assert.include(output.content[1]?.text ?? "", "returned an error"); + assert.equal(output.isError, isError ? true : undefined); + assert.deepEqual(output.structuredContent, result); + }, ); - await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, { - process: { - env: { T3_MCP_URL: "http://fixture.invalid/mcp", T3_MCP_BEARER_TOKEN: "fixture-token" }, - }, - AbortSignal, - Type: { Unsafe: (schema: unknown) => schema }, - fetch: async ( - url: string, - options: { body: string; headers: Record; signal?: AbortSignal }, - ) => { - transports.push({ - url, - authorization: options.headers.authorization!, - signal: options.signal, - }); - const request = JSON.parse(options.body) as { id: number; method: string; params?: unknown }; - requests.push(request); - const result = - request.method === "tools/list" - ? { tools: catalog } - : request.method === "tools/call" - ? { content: [{ type: "text", text: "browser snapshot" }] } - : {}; - return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result }), { - headers: { "content-type": "application/json" }, - }); - }, - pi: { - on: (name: string, handler: AgentStartHook) => handlers.set(name, handler), - registerTool: (tool: RegisteredTool) => { - if (options.allowsTool && !options.allowsTool(tool.name)) return; - const index = tools.findIndex((current) => current.name === tool.name); - if (index === -1) tools.push(tool); - else tools[index] = tool; - }, - getActiveTools: () => activeTools, - setActiveTools: (names: string[]) => { - activeTools = names.filter((name) => options.allowsTool?.(name) ?? true); - }, - getAllTools: () => - options.toolSearchAvailable && - !options.toolSearchDisabled && - (options.allowsTool?.("tool_search") ?? true) - ? [{ name: "tool_search", sourceInfo: { path: "builtin:tool-search" } }] - : [], - ...(options.modern - ? { - registerMcpServer: (name: string, config: Record) => - servers.push({ name, config }), - unregisterMcpServer: () => servers.splice(0), - } - : {}), - }, + + it("uses structured output as model text when content is empty", async () => { + const result = { content: [], structuredContent: { count: 2 } }; + const bridge = await loadMcpBridge({ result }); + const output = await bridge.tools[0]!.execute("structured", {}); + assert.deepEqual(output.content, [{ type: "text", text: '{"count":2}' }]); + assert.deepEqual(output.structuredContent, result); }); - return { - handlers, - tools, - requests, - servers, - transports, - getActiveTools: () => activeTools, - restoreActiveTools: (names: string[]) => { - activeTools = names; - }, - }; -} + + it("propagates cancellation without returning a successful result", async () => { + const bridge = await loadMcpBridge(); + const controller = new AbortController(); + controller.abort(); + await expect(bridge.tools[0]!.execute("cancelled", {}, controller.signal)).rejects.toThrow( + "aborted", + ); + }); +}); describe("Pi MCP tool exposure", () => { it("keeps orchestration direct and optional bridge tools discoverable on modern Pi", async () => { diff --git a/packages/provider-pi/src/server/mcpExtensionSource.ts b/packages/provider-pi/src/server/mcpExtensionSource.ts index 8279a515ed65fc9697c39f3184e2b3b78102cc20..e33b07cfdcd1e8cd933328e0f1181a41e6fc8310 100644 --- a/packages/provider-pi/src/server/mcpExtensionSource.ts +++ b/packages/provider-pi/src/server/mcpExtensionSource.ts @@ -38,9 +38,11 @@ const FILE_CHANGE_TOOLS = new Set(${JSON.stringify(PI_FILE_CHANGE_TOOLS)}); type RuntimeMode = "approval-required" | "auto-accept-edits" | "auto" | "full-access"; +type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue }; + type JsonRpcResponse = { readonly id?: number | string; - readonly result?: unknown; + readonly result?: JsonValue; readonly error?: { readonly message?: string }; }; @@ -48,6 +50,7 @@ type McpTool = { readonly name: string; readonly description?: string; readonly inputSchema?: Record; + readonly outputSchema?: Record; }; function env(name: string): string | undefined { @@ -96,28 +99,57 @@ 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 ModelContent = + | { type: "text"; text: string } + | { type: "image"; data: string; mimeType: string }; + +function mcpModelContent(result: unknown): ModelContent[] { + if (result === null || result === undefined) return []; + 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 content: ModelContent[] = []; 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 === "text" && typeof part.text === "string") { + content.push({ type: "text", text: part.text }); + } else if (part?.type === "image" && typeof part.data === "string" && typeof part.mimeType === "string") { + content.push({ type: "image", data: part.data, mimeType: part.mimeType }); + } } } - // Most T3 tools mirror structuredContent in a text block. Repeating it would - // leave T3's own output parsing two JSON documents instead of one. - if (record.structuredContent !== undefined) { - const structured = JSON.stringify(record.structuredContent); - if (!texts.includes(structured)) texts.push(structured); + // Structured output belongs to scripts. Only use it as model text when + // the server supplied no model-facing blocks, matching Pi's native MCP tools. + if (content.length === 0) { + content.push({ type: "text", text: JSON.stringify(record.structuredContent ?? result) }); } - if (texts.length > 0) return texts.join("\\n"); - return JSON.stringify(result); + return content; +} + +function mcpScriptResult(result: JsonValue | undefined) { + if (typeof result !== "object" || result === null || Array.isArray(result)) return result; + const { _meta: _ignored, ...scriptResult } = result; + return scriptResult; +} + +function mcpOutputSchema(tool: McpTool) { + return jsonSchemaToTypebox({ + type: "object", + properties: { + content: { type: "array", items: { type: "object" } }, + ...(tool.outputSchema === undefined ? {} : { structuredContent: tool.outputSchema }), + isError: { type: "boolean" }, + _meta: { type: "object" }, + }, + required: ["content"], + }); } function isMcpToolError(result: unknown): boolean { @@ -328,6 +360,7 @@ export default async function t3McpExtension(pi: ExtensionAPI) { label: name, description: tool.description ?? name, parameters: jsonSchemaToTypebox(tool.inputSchema), + ...(supportsExposure ? { outputSchema: mcpOutputSchema(tool) } : {}), ...(supportsExposure ? { exposure } : {}), async execute(_toolCallId, params, signal) { const result = await client.callTool( @@ -335,9 +368,13 @@ export default async function t3McpExtension(pi: ExtensionAPI) { (params ?? {}) as Record, signal, ); - const text = formatMcpContent(result); + const content = mcpModelContent(result); + if (isMcpToolError(result) && !content.some((part) => part.type === "text" && part.text.length > 0)) { + content.push({ type: "text", text: "MCP tool t3-code/" + name + " returned an error" }); + } return { - content: [{ type: "text", text }], + content, + structuredContent: mcpScriptResult(result), details: { server: "t3-code", tool: name }, ...(isMcpToolError(result) ? { isError: true } : {}), }; diff --git a/packages/provider-pi/src/server/rpc.ts b/packages/provider-pi/src/server/rpc.ts index 7c58ccc9997671ad615e08ba94f82b71de9e82d9..3f67291330ae13b1fb3989139bf14b7d19b83fd1 100644 --- a/packages/provider-pi/src/server/rpc.ts +++ b/packages/provider-pi/src/server/rpc.ts @@ -30,6 +30,10 @@ import * as ChildProcess from "effect/process/ChildProcess"; import * as ChildProcessSpawner from "effect/process/ChildProcessSpawner"; import * as HostProcess from "@t3tools/shared/HostProcess"; import { resolveSpawnCommand } from "@t3tools/shared/shell"; +import { + MAX_TOOL_OUTPUT_IMAGES, + MAX_TOOL_OUTPUT_IMAGE_BASE64_LENGTH, +} from "@t3tools/shared/toolOutput"; import { signalProcessGroup } from "@t3tools/provider-core/server/processGroup"; @@ -125,7 +129,11 @@ interface PendingPiRequest { readonly deferred: Deferred.Deferred; } -const MAX_PI_RECORD_CHARS = 8 * 1024 * 1024; +// MCP completion events carry images in both model content and the script +// CallToolResult. Allow the shared image budget twice, plus bounded JSON/text +// overhead, before dropping an extension's oversized record. +const MAX_PI_RECORD_CHARS = + 2 * MAX_TOOL_OUTPUT_IMAGES * MAX_TOOL_OUTPUT_IMAGE_BASE64_LENGTH + 8 * 1024 * 1024; function makeJsonlFramer() { let buffer = ""; -- 2.54.0