diff --git a/src/agents/pi-embedded-subscribe.handlers.tools.media.test.ts b/src/agents/pi-embedded-subscribe.handlers.tools.media.test.ts index a60986ec48f..e3641a4317b 100644 --- a/src/agents/pi-embedded-subscribe.handlers.tools.media.test.ts +++ b/src/agents/pi-embedded-subscribe.handlers.tools.media.test.ts @@ -97,6 +97,22 @@ async function emitUntrustedToolMediaResult( }); } +async function emitMcpMediaToolResult(ctx: EmbeddedPiSubscribeContext, mediaPathOrUrl: string) { + await handleToolExecutionEnd(ctx, { + type: "tool_execution_end", + toolName: "browser", + toolCallId: "tc-1", + isError: false, + result: { + content: [{ type: "text", text: `MEDIA:${mediaPathOrUrl}` }], + details: { + mcpServer: "probe", + mcpTool: "browser", + }, + }, + }); +} + describe("handleToolExecutionEnd media emission", () => { it("does not warn for read tool when path is provided via file_path alias", async () => { const ctx = createMockContext(); @@ -141,6 +157,26 @@ describe("handleToolExecutionEnd media emission", () => { expect(ctx.state.pendingToolMediaUrls).toEqual(["https://example.com/file.png"]); }); + it("does NOT emit local media for MCP-provenance results", async () => { + const onToolResult = vi.fn(); + const ctx = createMockContext({ shouldEmitToolOutput: false, onToolResult }); + + await emitMcpMediaToolResult(ctx, "/tmp/secret.png"); + + expect(onToolResult).not.toHaveBeenCalled(); + expect(ctx.state.pendingToolMediaUrls).toEqual([]); + }); + + it("emits remote media for MCP-provenance results", async () => { + const onToolResult = vi.fn(); + const ctx = createMockContext({ shouldEmitToolOutput: false, onToolResult }); + + await emitMcpMediaToolResult(ctx, "https://example.com/file.png"); + + expect(onToolResult).not.toHaveBeenCalled(); + expect(ctx.state.pendingToolMediaUrls).toEqual(["https://example.com/file.png"]); + }); + it("does NOT queue legacy MEDIA paths when verbose is full", async () => { const onToolResult = vi.fn(); const ctx = createMockContext({ shouldEmitToolOutput: true, onToolResult }); diff --git a/src/agents/pi-embedded-subscribe.handlers.tools.ts b/src/agents/pi-embedded-subscribe.handlers.tools.ts index 42f2ea33953..1662d2b2de6 100644 --- a/src/agents/pi-embedded-subscribe.handlers.tools.ts +++ b/src/agents/pi-embedded-subscribe.handlers.tools.ts @@ -301,7 +301,7 @@ async function emitToolResultOutput(params: { if (ctx.shouldEmitToolOutput()) { const outputText = extractToolResultText(sanitizedResult); if (outputText) { - ctx.emitToolOutput(toolName, meta, outputText); + ctx.emitToolOutput(toolName, meta, outputText, result); } if (!hasStructuredMedia) { return; @@ -316,7 +316,7 @@ async function emitToolResultOutput(params: { if (!mediaReply) { return; } - const mediaUrls = filterToolResultMediaUrls(toolName, mediaReply.mediaUrls); + const mediaUrls = filterToolResultMediaUrls(toolName, mediaReply.mediaUrls, result); if (mediaUrls.length === 0) { return; } diff --git a/src/agents/pi-embedded-subscribe.handlers.types.ts b/src/agents/pi-embedded-subscribe.handlers.types.ts index d7c07cce79e..211ac887642 100644 --- a/src/agents/pi-embedded-subscribe.handlers.types.ts +++ b/src/agents/pi-embedded-subscribe.handlers.types.ts @@ -95,7 +95,7 @@ export type EmbeddedPiSubscribeContext = { shouldEmitToolResult: () => boolean; shouldEmitToolOutput: () => boolean; emitToolSummary: (toolName?: string, meta?: string) => void; - emitToolOutput: (toolName?: string, meta?: string, output?: string) => void; + emitToolOutput: (toolName?: string, meta?: string, output?: string, result?: unknown) => void; stripBlockTags: ( text: string, state: { thinking: boolean; final: boolean; inlineCode?: InlineCodeState }, @@ -174,7 +174,7 @@ export type ToolHandlerContext = { shouldEmitToolResult: () => boolean; shouldEmitToolOutput: () => boolean; emitToolSummary: (toolName?: string, meta?: string) => void; - emitToolOutput: (toolName?: string, meta?: string, output?: string) => void; + emitToolOutput: (toolName?: string, meta?: string, output?: string, result?: unknown) => void; trimMessagingToolSent: () => void; }; diff --git a/src/agents/pi-embedded-subscribe.tools.media.test.ts b/src/agents/pi-embedded-subscribe.tools.media.test.ts index cb50f6e8a40..9d61d7f30e5 100644 --- a/src/agents/pi-embedded-subscribe.tools.media.test.ts +++ b/src/agents/pi-embedded-subscribe.tools.media.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"; import { extractToolResultMediaArtifact, extractToolResultMediaPaths, + filterToolResultMediaUrls, isToolResultMediaTrusted, } from "./pi-embedded-subscribe.tools.js"; @@ -263,4 +264,26 @@ describe("extractToolResultMediaPaths", () => { it("trusts image_generate local MEDIA paths", () => { expect(isToolResultMediaTrusted("image_generate")).toBe(true); }); + + it("does not trust local MEDIA paths for MCP-provenance results", () => { + expect( + filterToolResultMediaUrls("browser", ["/tmp/screenshot.png"], { + details: { + mcpServer: "probe", + mcpTool: "browser", + }, + }), + ).toEqual([]); + }); + + it("still allows remote MEDIA urls for MCP-provenance results", () => { + expect( + filterToolResultMediaUrls("browser", ["https://example.com/screenshot.png"], { + details: { + mcpServer: "probe", + mcpTool: "browser", + }, + }), + ).toEqual(["https://example.com/screenshot.png"]); + }); }); diff --git a/src/agents/pi-embedded-subscribe.tools.ts b/src/agents/pi-embedded-subscribe.tools.ts index 20f0f417531..ab969706d77 100644 --- a/src/agents/pi-embedded-subscribe.tools.ts +++ b/src/agents/pi-embedded-subscribe.tools.ts @@ -162,8 +162,26 @@ const TRUSTED_TOOL_RESULT_MEDIA = new Set([ ]); const HTTP_URL_RE = /^https?:\/\//i; -export function isToolResultMediaTrusted(toolName?: string): boolean { - if (!toolName) { +function readToolResultDetails(result: unknown): Record | undefined { + if (!result || typeof result !== "object") { + return undefined; + } + const record = result as Record; + return record.details && typeof record.details === "object" && !Array.isArray(record.details) + ? (record.details as Record) + : undefined; +} + +function isExternalToolResult(result: unknown): boolean { + const details = readToolResultDetails(result); + if (!details) { + return false; + } + return typeof details.mcpServer === "string" || typeof details.mcpTool === "string"; +} + +export function isToolResultMediaTrusted(toolName?: string, result?: unknown): boolean { + if (!toolName || isExternalToolResult(result)) { return false; } const normalized = normalizeToolName(toolName); @@ -173,11 +191,12 @@ export function isToolResultMediaTrusted(toolName?: string): boolean { export function filterToolResultMediaUrls( toolName: string | undefined, mediaUrls: string[], + result?: unknown, ): string[] { if (mediaUrls.length === 0) { return mediaUrls; } - if (isToolResultMediaTrusted(toolName)) { + if (isToolResultMediaTrusted(toolName, result)) { return mediaUrls; } return mediaUrls.filter((url) => HTTP_URL_RE.test(url.trim())); @@ -203,10 +222,7 @@ export type ToolResultMediaArtifact = { function readToolResultDetailsMedia( result: Record, ): Record | undefined { - const details = - result.details && typeof result.details === "object" && !Array.isArray(result.details) - ? (result.details as Record) - : undefined; + const details = readToolResultDetails(result); const media = details?.media && typeof details.media === "object" && !Array.isArray(details.media) ? (details.media as Record) diff --git a/src/agents/pi-embedded-subscribe.ts b/src/agents/pi-embedded-subscribe.ts index e2fd18cab70..62aea0e5450 100644 --- a/src/agents/pi-embedded-subscribe.ts +++ b/src/agents/pi-embedded-subscribe.ts @@ -337,12 +337,16 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar } return `\`\`\`txt\n${trimmed}\n\`\`\``; }; - const emitToolResultMessage = (toolName: string | undefined, message: string) => { + const emitToolResultMessage = ( + toolName: string | undefined, + message: string, + result?: unknown, + ) => { if (!params.onToolResult) { return; } const { text: cleanedText, mediaUrls } = parseReplyDirectives(message); - const filteredMediaUrls = filterToolResultMediaUrls(toolName, mediaUrls ?? []); + const filteredMediaUrls = filterToolResultMediaUrls(toolName, mediaUrls ?? [], result); if (!cleanedText && filteredMediaUrls.length === 0) { return; } @@ -361,7 +365,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar }); emitToolResultMessage(toolName, agg); }; - const emitToolOutput = (toolName?: string, meta?: string, output?: string) => { + const emitToolOutput = (toolName?: string, meta?: string, output?: string, result?: unknown) => { if (!output) { return; } @@ -369,7 +373,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar markdown: useMarkdown, }); const message = `${agg}\n${formatToolOutputBlock(output)}`; - emitToolResultMessage(toolName, message); + emitToolResultMessage(toolName, message, result); }; const stripBlockTags = (