From 531370c4ce98c1e32a0c0a5c4baf4b31501999a4 Mon Sep 17 00:00:00 2001 From: "paseo-bot[bot]" <266920839+paseo-bot[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:10:39 +0000 Subject: [PATCH] fix(pi): show tool-result images live and after reload (#8280) * pi: render tool-result images as inline image rows (paseo #1717 parity) * fix(pi): preserve turn ownership while rendering images * fix(pi): preserve metadata when stripping nested image payloads --------- Co-authored-by: be-my-bot Co-authored-by: paseo-bot[bot] <266920839+paseo-bot[bot]@users.noreply.github.com> Co-authored-by: duyiliu --- .../server/agent/providers/pi/agent.test.ts | 111 ++++++++++++++++++ .../src/server/agent/providers/pi/agent.ts | 68 ++++++++++- .../providers/pi/extensions/child-session.ts | 11 +- .../providers/pi/extensions/index.test.ts | 6 +- .../rpiv-ask-user-question/index.test.ts | 4 +- .../pi/extensions/task-fixture-test.ts | 2 +- .../agent/providers/pi/history-mapper.test.ts | 72 ++++++++++++ .../agent/providers/pi/history-mapper.ts | 36 ++++-- .../providers/pi/tool-call-mapper.test.ts | 82 ++++++++++++- .../agent/providers/pi/tool-call-mapper.ts | 63 ++++++++++ 10 files changed, 429 insertions(+), 26 deletions(-) diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index 82887bf199..936bdbc973 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -3643,3 +3643,114 @@ test("closing Pi during image conversion rejects its pending start", async () => await session.close(); await rejected; }); + +describe("Pi tool-result images", () => { + test.each([false, true])( + "tool-result images precede the reply and completion (isError=%s)", + async (isError) => { + const { pi, session, events } = await createSession(); + onTestFinished(() => session.close()); + const fakeSession = pi.latestSession(); + await session.startTurn("run"); + fakeSession.emit({ + type: "tool_execution_start", + toolCallId: "shot-1", + toolName: "screenshot", + args: {}, + }); + fakeSession.emit({ + type: "tool_execution_end", + toolCallId: "shot-1", + toolName: "screenshot", + result: { + content: [ + { type: "text", text: "captured" }, + { type: "image", data: Buffer.from("hello").toString("base64"), mimeType: "image/png" }, + ], + structuredContent: { + type: "image", + data: Buffer.from("hello").toString("base64"), + mimeType: "image/png", + note: "captured", + }, + }, + isError, + }); + fakeSession.emit({ + type: "message_update", + assistantMessageEvent: { type: "text_delta", delta: "Image read complete." }, + }); + fakeSession.finishTurn(); + await events.nextTurnCompletion(); + + const items = events.timelineItems(); + const toolIndex = items.findIndex( + (item) => + item.type === "tool_call" && + item.callId === "shot-1" && + item.status === (isError ? "failed" : "completed"), + ); + expect(toolIndex).toBeGreaterThanOrEqual(0); + // Base64 must never reach the tool output stringifier. + expect(JSON.stringify(items[toolIndex])).not.toContain( + Buffer.from("hello").toString("base64"), + ); + expect(items.slice(toolIndex + 1)).toEqual([ + expect.objectContaining({ + type: "assistant_message", + text: expect.stringContaining("![Image]("), + }), + expect.objectContaining({ type: "assistant_message", text: "Image read complete." }), + ]); + expect(events.timelineAndCompletionEvents().at(-1)).toEqual({ type: "turn_completed" }); + }, + ); +}); + +test("Pi image rendering preserves turn ownership across an interrupt and follow-up", async () => { + const { pi, session, events } = await createSession(); + onTestFinished(() => session.close()); + const emitted: AgentStreamEvent[] = []; + session.subscribe((event) => emitted.push(event)); + const native = pi.latestSession(); + const first = await session.startTurn("read image"); + native.emit({ + type: "tool_execution_end", + toolCallId: "large-image", + toolName: "screenshot", + result: { + content: [ + { + type: "image", + data: Buffer.alloc(8 * 1024 * 1024, 73).toString("base64"), + mimeType: "image/png", + }, + ], + }, + isError: false, + }); + native.finishTurn(); + await session.interrupt(); + const followup = await session.startTurn("follow-up after image"); + await waitForImmediate(); + await waitForImmediate(); + await expect + .poll(() => + events + .timelineItems() + .some((item) => item.type === "assistant_message" && item.text.includes("![Image](")), + ) + .toBe(true); + expect(events.turnCompletedEvents().some((event) => event.turnId === followup.turnId)).toBe( + false, + ); + const imageEvent = emitted.find( + (event) => + event.type === "timeline" && + event.item.type === "assistant_message" && + event.item.text.includes("![Image]("), + ); + expect(imageEvent).toMatchObject({ turnId: first.turnId }); + native.finishTurn(); + await events.nextTurnCompletion(); +}); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 5fbffde84b..9474e98cc8 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -64,7 +64,11 @@ import { streamPiHistory, type PiCapturedUserMessageEntry, } from "./history-mapper.js"; -import { materializeProviderImage } from "../provider-image-output.js"; +import { + materializeProviderImage, + renderProviderImageOutputAsAssistantMarkdown, + type ProviderImageOutput, +} from "../provider-image-output.js"; import { PiCliRuntime } from "./cli-runtime.js"; import { createPiExtensionHost, @@ -89,6 +93,7 @@ import { mapToolDetail, parseToolArgs, parseToolResult, + splitPiToolResultImages, type PiToolResult, type PiTrackedToolCall, } from "./tool-call-mapper.js"; @@ -1201,6 +1206,7 @@ export class PiRpcAgentSession implements AgentSession { private closed = false; private readonly closeController = new AbortController(); private readonly pendingExtensionHydrations = new Set>(); + private pendingOutput: Promise | null = null; // Pi publishes the terminal before acknowledging abort. Autonomous runs have no // turn ID; retain their errors too until the cancellation request settles. private interruptingTurn: { turnId: string | undefined; error: string | null } | null = null; @@ -1602,6 +1608,7 @@ export class PiRpcAgentSession implements AgentSession { this.pendingExtensionUiRequests.size || this.pendingExtensionResults.size || this.pendingExtensionHydrations.size || + this.pendingOutput || this.pendingSteerSubmissions.length || this.pendingNoTurnOutputs.length || this.interruptingTurn @@ -1622,7 +1629,7 @@ export class PiRpcAgentSession implements AgentSession { try { await this.runtimeSession.close(); } finally { - await Promise.all(this.pendingExtensionHydrations); + await Promise.all([this.pendingOutput, ...this.pendingExtensionHydrations]); this.rejectAllExtensionResults(new Error("Pi session closed")); this.cleanup?.(); } @@ -1686,9 +1693,31 @@ export class PiRpcAgentSession implements AgentSession { } private emit(event: AgentStreamEvent): void { - for (const subscriber of this.subscribers) { - subscriber(event); + if (this.pendingOutput) { + this.emitAsync(() => Promise.resolve([event])); + return; } + this.publish(event); + } + + private publish(event: AgentStreamEvent): void { + for (const subscriber of this.subscribers) subscriber(event); + } + + // Keep timeline order while images are materialized, without delaying provider state changes. + private emitAsync(render: () => Promise): void { + const pending = (this.pendingOutput ?? Promise.resolve()) + .then(render) + .then((events) => { + if (!this.closed) for (const event of events) this.publish(event); + return undefined; + }) + .catch((error) => this.logger.error({ err: error }, "Pi output rendering failed")); + this.pendingOutput = pending; + void pending.then(() => { + if (this.pendingOutput === pending) this.pendingOutput = null; + return undefined; + }); } private emitExtensionOutput(output: PiExtensionEventOutput | undefined, turnId?: string): void { @@ -2315,8 +2344,11 @@ export class PiRpcAgentSession implements AgentSession { this.activeToolCalls.get(event.toolCallId) ?? parseToolArgs(event.toolName, null); this.activeToolCalls.delete(event.toolCallId); - const result = parseToolResult(event.result); - const error = event.isError ? event.result : null; + // Pull image blocks out of the result so base64 never reaches the tool output, then render + // each one as an assistant_message markdown image after the tool_call (mirroring the Claude + // provider, upstream #1717). + const { images, result } = splitPiToolResultImages(parseToolResult(event.result)); + const error = event.isError ? result : null; const status = event.isError ? "failed" : "completed"; this.extensionHost.onToolEnd({ callId: event.toolCallId, @@ -2326,6 +2358,30 @@ export class PiRpcAgentSession implements AgentSession { result, }); this.emitToolCallEvent(event.toolCallId, toolCall, status, result, error); + if (images.length === 0) return; + const turnId = this.currentTurnIdForEvent(); + this.emitAsync(() => this.renderToolResultImages(images, turnId)); + } + + private async renderToolResultImages( + images: ProviderImageOutput[], + turnId: string | undefined, + ): Promise { + const events: AgentStreamEvent[] = []; + for (const image of images) { + const imageItem = await renderProviderImageOutputAsAssistantMarkdown(image, { + materialize: materializeProviderImage, + }); + if (imageItem && !this.closed) { + events.push({ + type: "timeline", + provider: this.provider, + turnId, + item: imageItem, + }); + } + } + return events; } private emitCompactionTimeline(input: { diff --git a/packages/server/src/server/agent/providers/pi/extensions/child-session.ts b/packages/server/src/server/agent/providers/pi/extensions/child-session.ts index 2013b831be..0c7c7df9aa 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/child-session.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/child-session.ts @@ -21,11 +21,11 @@ export function mapPiChildSession( return new PiChildSessionFollower(id, file, maxBytes).readNew(); } -function mapChildLine( +async function mapChildLine( mapper: PiHistoryMapper, id: string, line: string, -): ProviderSubagentInputEvent[] { +): Promise { let entry: { message?: PiAgentMessage; timestamp?: string }; try { entry = JSON.parse(line); @@ -33,7 +33,7 @@ function mapChildLine( return []; } if (!entry.message || typeof entry.message !== "object" || !("role" in entry.message)) return []; - return mapper.mapMessages([entry.message]).flatMap((mapped) => + return (await mapper.mapMessages([entry.message])).flatMap((mapped) => mapped.type === "timeline" ? [ { @@ -83,7 +83,10 @@ export class PiChildSessionFollower { const events: ProviderSubagentInputEvent[] = []; for (const line of data.subarray(0, lastNewline).toString("utf8").split("\n")) { if (this.items >= MAX_ITEMS) break; - const mapped = mapChildLine(this.mapper, this.id, line).slice(0, MAX_ITEMS - this.items); + const mapped = (await mapChildLine(this.mapper, this.id, line)).slice( + 0, + MAX_ITEMS - this.items, + ); this.items += mapped.length; events.push(...mapped); } diff --git a/packages/server/src/server/agent/providers/pi/extensions/index.test.ts b/packages/server/src/server/agent/providers/pi/extensions/index.test.ts index 9afacae977..0bf32a9f2f 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/index.test.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/index.test.ts @@ -83,7 +83,7 @@ describe("Pi extension host", () => { test("declines a throwing adapter and preserves generic history", async () => { const host = createPiExtensionHost(undefined, [throwingAdapter]); const mapper = new PiHistoryMapper("pi", [], {}, host); - const events = mapper.mapMessages([ + const events = await mapper.mapMessages([ { role: "assistant", content: [{ type: "toolCall", id: "call-1", name: "other", arguments: {} }], @@ -193,7 +193,7 @@ describe("Pi extension host", () => { await output?.hydration; expect(reads).toHaveLength(8); }); - test("claims a tool, declines an unrelated one, and maps history like live events", () => { + test("claims a tool, declines an unrelated one, and maps history like live events", async () => { const host = createPiExtensionHost(); const args = { agent: "scout", task: "Inspect files" }; const result = parseToolResult({ content: [{ type: "text", text: "Found two files" }] }); @@ -225,7 +225,7 @@ describe("Pi extension host", () => { ).toBeUndefined(); const mapper = new PiHistoryMapper("pi"); - const events = mapper.mapMessages([ + const events = await mapper.mapMessages([ { role: "assistant", content: [{ type: "toolCall", id: "call-1", name: "subagent", arguments: args }], diff --git a/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts b/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts index 3ec9256a29..d3c9fb7a9e 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/rpiv-ask-user-question/index.test.ts @@ -79,7 +79,7 @@ test("cancellation declines each remaining real RPC dialog", () => { } }); -test("captured result renders in live and history tool rows, while foreign shapes decline", () => { +test("captured result renders in live and history tool rows, while foreign shapes decline", async () => { const call = { callId: end.toolCallId, toolName: end.toolName, @@ -93,7 +93,7 @@ test("captured result renders in live and history tool rows, while foreign shape label: "Questions", text: "Which snack?: Cookie\nWhich colors?: Red, Blue\nWhat project name?: Nebula", }); - const history = new PiHistoryMapper("pi").mapMessages([ + const history = await new PiHistoryMapper("pi").mapMessages([ { role: "assistant", content: [ diff --git a/packages/server/src/server/agent/providers/pi/extensions/task-fixture-test.ts b/packages/server/src/server/agent/providers/pi/extensions/task-fixture-test.ts index a3b2114cab..f702e0d7cf 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/task-fixture-test.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/task-fixture-test.ts @@ -89,7 +89,7 @@ export async function verifyTaskFixture(path: URL, expectedStatuses: string[]): for (const event of fixture.events) pi.latestSession().emit(event); pi.latestSession().finishTurn(); expect(todos(live)).toEqual(mapped); - const replay = new PiHistoryMapper("pi").mapMessages(fixture.messages); + const replay = await new PiHistoryMapper("pi").mapMessages(fixture.messages); expect(todos(replay)).toEqual(mapped); expect( replay.filter( diff --git a/packages/server/src/server/agent/providers/pi/history-mapper.test.ts b/packages/server/src/server/agent/providers/pi/history-mapper.test.ts index 24d18306b2..f4136956eb 100644 --- a/packages/server/src/server/agent/providers/pi/history-mapper.test.ts +++ b/packages/server/src/server/agent/providers/pi/history-mapper.test.ts @@ -221,3 +221,75 @@ describe("Pi history mapper", () => { ]); }); }); + +describe("Pi tool-result image replay", () => { + test("recreates an image from native tool-result history without its source file", async () => { + const events = await collectHistory([ + { + role: "assistant", + content: [ + { + type: "toolCall", + id: "image-read", + name: "read", + arguments: { path: "missing-preview.png" }, + }, + ], + }, + { + role: "toolResult", + toolCallId: "image-read", + toolName: "read", + content: [ + { type: "text", text: "Read image file [image/png]" }, + { type: "image", data: "aGVsbG8=", mimeType: "image/png" }, + ], + }, + ]); + expect(events.map((event) => event.item?.type)).toEqual([ + "tool_call", + "tool_call", + "assistant_message", + ]); + const image = events.find( + (event) => event.type === "timeline" && event.item.type === "assistant_message", + ); + expect(image).toMatchObject({ + type: "timeline", + item: { type: "assistant_message", text: expect.stringContaining("![Image](file://") }, + }); + if (image?.type !== "timeline" || image.item.type !== "assistant_message") + throw new Error("Missing image"); + const url = image.item.text.match(/\((file:[^)]+)\)/)![1]; + const { readFile, unlink } = await import("node:fs/promises"); + expect(await readFile(new URL(url))).toEqual(Buffer.from("hello")); + await unlink(new URL(url)); + const replay = await collectHistory([ + { + role: "toolResult", + toolCallId: "image-read", + toolName: "read", + content: [{ type: "image", data: "aGVsbG8=", mimeType: "image/png" }], + }, + ]); + expect(replay).toContainEqual(image); + expect(await readFile(new URL(url))).toEqual(Buffer.from("hello")); + }); +}); + +test("replayed Pi images keep their position before the following reply", async () => { + const events = await collectHistory([ + { + role: "toolResult", + toolCallId: "image-position", + toolName: "read", + content: [{ type: "image", data: "aGVsbG8=", mimeType: "image/png" }], + }, + { role: "assistant", content: [{ type: "text", text: "Image read complete." }] }, + ]); + expect( + events.map((event) => + event.item?.type === "assistant_message" ? event.item.text : event.item?.type, + ), + ).toEqual(["tool_call", expect.stringContaining("![Image]("), "Image read complete."]); +}); diff --git a/packages/server/src/server/agent/providers/pi/history-mapper.ts b/packages/server/src/server/agent/providers/pi/history-mapper.ts index 27a0250b00..89d9fcedce 100644 --- a/packages/server/src/server/agent/providers/pi/history-mapper.ts +++ b/packages/server/src/server/agent/providers/pi/history-mapper.ts @@ -1,5 +1,9 @@ import { mapCustomMessageToToolCall } from "../custom-message.js"; import type { AgentStreamEvent, AgentTimelineItem, ToolCallDetail } from "../../agent-sdk-types.js"; +import { + materializeProviderImage, + renderProviderImageOutputAsAssistantMarkdown, +} from "../provider-image-output.js"; import { createPiExtensionHost, type PiExtensionEventOutput, @@ -11,6 +15,7 @@ import { mapToolDetail, parseToolArgs, parseToolResult, + splitPiToolResultImages, type PiToolResult, type PiTrackedToolCall, } from "./tool-call-mapper.js"; @@ -71,7 +76,7 @@ export class PiHistoryMapper { private readonly extensionHost: PiExtensionHost = createPiExtensionHost(), ) {} - mapMessages(messages: readonly PiAgentMessage[]): AgentStreamEvent[] { + async mapMessages(messages: readonly PiAgentMessage[]): Promise { const events: AgentStreamEvent[] = []; for (const message of messages) { @@ -86,7 +91,7 @@ export class PiHistoryMapper { events.push(...this.mapAssistantMessage(message)); break; case "toolResult": { - events.push(...this.mapToolResultMessage(message)); + events.push(...(await this.mapToolResultMessage(message))); break; } case "bashExecution": @@ -207,24 +212,35 @@ export class PiHistoryMapper { return events; } - private mapToolResultMessage( + private async mapToolResultMessage( message: Extract, - ): AgentStreamEvent[] { + ): Promise { const tracked = this.pendingToolCalls.get(message.toolCallId) ?? parseToolArgs(message.toolName, null); this.pendingToolCalls.delete(message.toolCallId); const result = parseToolResult({ content: message.content, details: message.details }); + // Pull image blocks out of the result so base64 never reaches the tool output, then render + // each one as an assistant_message markdown image after the tool_call (mirroring the live + // path and the Claude provider, upstream #1717). + const { images, result: splitResult } = splitPiToolResultImages(result); const mapping = this.extensionHost.mapToolCall({ callId: message.toolCallId, toolName: tracked.toolName, args: tracked.args, status: message.isError ? "failed" : "completed", - result, + result: splitResult, }); - const detail = this.mapToolDetail(message.toolCallId, tracked, result, mapping?.detail); + const detail = this.mapToolDetail(message.toolCallId, tracked, splitResult, mapping?.detail); if (!detail) { return []; } + const imageEvents: AgentStreamEvent[] = []; + for (const image of images) { + const item = await renderProviderImageOutputAsAssistantMarkdown(image, { + materialize: materializeProviderImage, + }); + if (item) imageEvents.push({ type: "timeline", provider: this.provider, item }); + } return [ { type: "timeline", @@ -234,9 +250,10 @@ export class PiHistoryMapper { name: mapping?.name ?? tracked.toolName, isError: Boolean(message.isError), detail, - errorText: extractTextFromToolResult(result) ?? "Tool call failed", + errorText: extractTextFromToolResult(splitResult) ?? "Tool call failed", }), }, + ...imageEvents, ...this.extensionEvents(mapping), ]; } @@ -297,9 +314,10 @@ export async function* streamPiHistory( signal?: AbortSignal, ): AsyncGenerator { const mapper = new PiHistoryMapper(provider, userEntries, hooks, extensionHost); - for (const event of mapper.mapMessages(messages)) { + for (const message of messages) { if (signal?.aborted) return; - if (event) { + for (const event of await mapper.mapMessages([message])) { + if (signal?.aborted) return; yield event; } } diff --git a/packages/server/src/server/agent/providers/pi/tool-call-mapper.test.ts b/packages/server/src/server/agent/providers/pi/tool-call-mapper.test.ts index 63b557f994..0c065e27cd 100644 --- a/packages/server/src/server/agent/providers/pi/tool-call-mapper.test.ts +++ b/packages/server/src/server/agent/providers/pi/tool-call-mapper.test.ts @@ -1,6 +1,12 @@ import { describe, expect, test } from "vitest"; -import { mapToolDetail, parseToolArgs, parseToolResult } from "./tool-call-mapper.js"; +import { + extractTextFromToolResult, + mapToolDetail, + parseToolArgs, + parseToolResult, + splitPiToolResultImages, +} from "./tool-call-mapper.js"; describe("Pi tool call mapper", () => { test("maps bash args and result to shell detail", () => { @@ -56,3 +62,77 @@ describe("Pi tool call mapper", () => { }); }); }); + +describe("splitPiToolResultImages", () => { + test("extracts image blocks and replaces them with [image] placeholders", () => { + const parsed = parseToolResult({ + content: [ + { type: "text", text: "shot taken" }, + { type: "image", data: "aGVsbG8=", mimeType: "image/png" }, + ], + }); + const { images, result } = splitPiToolResultImages(parsed); + + expect(images).toEqual([{ data: "aGVsbG8=", mimeType: "image/png" }]); + expect(result).toEqual({ + content: [ + { type: "text", text: "shot taken" }, + { type: "text", text: "[image]" }, + ], + }); + }); + + test("replaces image-only results with a non-empty placeholder", () => { + const parsed = parseToolResult({ + content: [{ type: "image", data: "aGVsbG8=", mimeType: "image/jpeg" }], + }); + const { images, result } = splitPiToolResultImages(parsed); + + expect(images).toHaveLength(1); + expect(extractTextFromToolResult(result)).toBe("[image]"); + }); + + test("passes through string results and results without images", () => { + expect(splitPiToolResultImages("plain output")).toEqual({ + images: [], + result: "plain output", + }); + const textOnly = parseToolResult({ content: [{ type: "text", text: "no images" }] }); + expect(splitPiToolResultImages(textOnly)).toEqual({ images: [], result: textOnly }); + expect(splitPiToolResultImages(null)).toEqual({ images: [], result: null }); + }); + + test("keeps other fields on the result object", () => { + const parsed = parseToolResult({ + exitCode: 0, + content: [{ type: "image", data: "aGVsbG8=" }], + }); + const { result } = splitPiToolResultImages(parsed); + expect(result).toMatchObject({ exitCode: 0 }); + }); +}); + +test("tool-result images remove nested structured image data while preserving metadata", () => { + const data = Buffer.from("nested screenshot bytes").toString("base64"); + const result = parseToolResult({ + content: [{ type: "image", data, mimeType: "image/png" }], + structuredContent: { + content: [{ type: "image", data, mimeType: "image/png", width: 64 }], + captureId: "capture-1", + read: { type: "image", data, mimeType: "image/png", note: "Read image file" }, + extra: { data: "keep ordinary data" }, + }, + }); + const split = splitPiToolResultImages(result); + const detail = mapToolDetail(parseToolArgs("screenshot", {}), split.result); + expect(JSON.stringify(detail)).not.toContain(data); + expect(split.images).toEqual([{ data, mimeType: "image/png" }]); + expect(split.result).toMatchObject({ + structuredContent: { + content: [{ type: "image", data: "[image]", mimeType: "image/png", width: 64 }], + captureId: "capture-1", + read: { type: "image", data: "[image]", mimeType: "image/png", note: "Read image file" }, + extra: { data: "keep ordinary data" }, + }, + }); +}); diff --git a/packages/server/src/server/agent/providers/pi/tool-call-mapper.ts b/packages/server/src/server/agent/providers/pi/tool-call-mapper.ts index 8c18a8fdbf..8fbba690b8 100644 --- a/packages/server/src/server/agent/providers/pi/tool-call-mapper.ts +++ b/packages/server/src/server/agent/providers/pi/tool-call-mapper.ts @@ -1,6 +1,7 @@ import { z } from "zod"; import type { ToolCallDetail } from "../../agent-sdk-types.js"; +import type { ProviderImageOutput } from "../provider-image-output.js"; interface BashToolInput { command: string; @@ -52,6 +53,7 @@ interface PiToolResultObject { stdout?: string; text?: string; content?: PiToolResultContent[]; + structuredContent?: unknown; exitCode?: number; code?: number; details?: PiToolResultDetails; @@ -67,6 +69,12 @@ interface PiToolResultTextContent { text: string; } +const PiToolResultImageContentSchema = z.object({ + type: z.literal("image"), + data: z.string().min(1), + mimeType: z.string().optional(), +}); + interface PiToolResultUnknownContent { type: string; } @@ -236,6 +244,61 @@ export function parseToolResult(rawResult: unknown): PiToolResult { return null; } +// Pi returns images inside tool-result content as base64 blocks. Left in place they reach the +// tool-output stringifier, which dumps base64 into the tool output. We pull those blocks out to +// render them as image markdown and leave a "[image]" placeholder so image-only results still +// produce non-empty output (mirroring splitClaudeToolResultImages). +export function splitPiToolResultImages(result: PiToolResult): { + images: ProviderImageOutput[]; + result: PiToolResult; +} { + if (!result || typeof result === "string" || !result.content) { + return { images: [], result }; + } + const images: ProviderImageOutput[] = []; + let replaced = false; + const content = result.content.map((block) => { + const image = toPiImageOutput(block); + if (image) { + images.push(image); + replaced = true; + return { type: "text", text: "[image]" } satisfies PiToolResultTextContent; + } + return block; + }); + if (!replaced) { + return { images: [], result }; + } + return { + images, + result: { + ...result, + content, + ...(result.structuredContent === undefined + ? {} + : { structuredContent: withoutImageData(result.structuredContent) }), + }, + }; +} + +function withoutImageData(value: unknown): unknown { + if (Array.isArray(value)) return value.map(withoutImageData); + if (value === null || typeof value !== "object") return value; + const image = toPiImageOutput(value); + return Object.fromEntries( + Object.entries(value).map(([key, child]) => [ + key, + image && key === "data" ? "[image]" : withoutImageData(child), + ]), + ); +} + +function toPiImageOutput(block: unknown): ProviderImageOutput | null { + const parsed = PiToolResultImageContentSchema.safeParse(block); + if (!parsed.success) return null; + return { data: parsed.data.data, mimeType: parsed.data.mimeType ?? null }; +} + export function extractTextFromToolResult(result: PiToolResult): string | undefined { if (typeof result === "string") { return result;