From 1f8a4559a86e98c51a64b8acffa47e404870ec3b Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Mon, 5 Oct 2026 00:03:20 +0100 Subject: [PATCH] fix(providers): show custom context in expandable tool rows (#6090) * fix(providers): show custom context in expandable tool rows Preserve custom message type and data instead of attributing extension context to the assistant. Apply the same mapping in live output and history for OMP and Pi, and honor hidden Pi messages without suppressing structured extension updates. * fix(server): mark truncated plain-text tool content Keep the existing timeline content budget while making omitted text visible to readers of extension context and other plain-text tools. --- .../agent/agent-timeline-content.test.ts | 5 +- .../server/agent/agent-timeline-content.ts | 3 +- .../server/agent/providers/custom-message.ts | 26 +++ .../server/agent/providers/omp/agent.test.ts | 56 ++++++- .../src/server/agent/providers/omp/agent.ts | 18 +- .../providers/omp/history-mapper.test.ts | 155 +++++++++++++++--- .../agent/providers/omp/message-history.ts | 3 +- .../server/agent/providers/pi/agent.test.ts | 64 +++++++- .../src/server/agent/providers/pi/agent.ts | 9 +- .../gotgenes-pi-subagents/index.test.ts | 19 ++- .../pi/extensions/subagent-fixture-test.ts | 50 ++++-- .../agent/providers/pi/history-mapper.test.ts | 41 ++++- .../agent/providers/pi/history-mapper.ts | 9 +- 13 files changed, 393 insertions(+), 65 deletions(-) create mode 100644 packages/server/src/server/agent/providers/custom-message.ts diff --git a/packages/server/src/server/agent/agent-timeline-content.test.ts b/packages/server/src/server/agent/agent-timeline-content.test.ts index f3b94e248e..c40b712fad 100644 --- a/packages/server/src/server/agent/agent-timeline-content.test.ts +++ b/packages/server/src/server/agent/agent-timeline-content.test.ts @@ -3,7 +3,7 @@ import { describe, expect, test } from "vitest"; import { limitAgentTimelineItemContent } from "./agent-timeline-content.js"; describe("agent timeline content", () => { - test("limits terminal input to the tool-call content budget", () => { + test("marks truncated plain-text tool content within the content budget", () => { const oversizedInput = "x".repeat(64 * 1024 + 1); const item = limitAgentTimelineItemContent({ @@ -19,6 +19,7 @@ describe("agent timeline content", () => { }, }); + expect(limitAgentTimelineItemContent(item)).toEqual(item); expect(item).toEqual({ type: "tool_call", callId: "terminal-session-4242", @@ -27,7 +28,7 @@ describe("agent timeline content", () => { error: null, detail: { type: "plain_text", - text: "x".repeat(64 * 1024), + text: "x".repeat(64 * 1024 - "\n\n[Content truncated]".length) + "\n\n[Content truncated]", icon: "square_terminal", }, }); diff --git a/packages/server/src/server/agent/agent-timeline-content.ts b/packages/server/src/server/agent/agent-timeline-content.ts index 8cf92381fa..d9bb08c0bc 100644 --- a/packages/server/src/server/agent/agent-timeline-content.ts +++ b/packages/server/src/server/agent/agent-timeline-content.ts @@ -42,11 +42,12 @@ function limitPlainText(item: AgentTimelineItem): AgentTimelineItem { ) { return item; } + const notice = "\n\n[Content truncated]"; return { ...item, detail: { ...item.detail, - text: item.detail.text.slice(0, TOOL_CALL_CONTENT_MAX_LENGTH), + text: item.detail.text.slice(0, TOOL_CALL_CONTENT_MAX_LENGTH - notice.length) + notice, }, }; } diff --git a/packages/server/src/server/agent/providers/custom-message.ts b/packages/server/src/server/agent/providers/custom-message.ts new file mode 100644 index 0000000000..3dd15b2d9c --- /dev/null +++ b/packages/server/src/server/agent/providers/custom-message.ts @@ -0,0 +1,26 @@ +import type { AgentTimelineItem } from "../agent-sdk-types.js"; + +/** Preserve extension context as an expandable tool row instead of assistant speech. */ +export function mapCustomMessageToToolCall( + message: { role: "custom"; customType?: unknown; details?: unknown }, + text: string, + callId: string, +): Extract { + const customType = + typeof message.customType === "string" && message.customType.trim() + ? message.customType + : "custom-message"; + return { + type: "tool_call", + callId, + name: customType, + status: "completed", + detail: { type: "plain_text", text }, + metadata: { + synthetic: true, + customType, + ...(message.details === undefined ? {} : { details: message.details }), + }, + error: null, + }; +} diff --git a/packages/server/src/server/agent/providers/omp/agent.test.ts b/packages/server/src/server/agent/providers/omp/agent.test.ts index 640ead5d30..f06bb62f57 100644 --- a/packages/server/src/server/agent/providers/omp/agent.test.ts +++ b/packages/server/src/server/agent/providers/omp/agent.test.ts @@ -662,6 +662,46 @@ describe("OMP agent client and session", () => { expect(omp.completedTurnCount()).toBe(1); }); + test("keeps custom context in separate tools while a turn continues", async () => { + const omp = new OmpHarness(); + await omp.start(); + await omp.requireStartTurn("Explain the project"); + omp.runtime().beginTurn(); + for (const display of [true, false, true]) { + omp.emit({ + type: "message_end", + message: { + role: "custom", + customType: "project-context", + content: [{ type: "text", text: "Project instructions" }], + details: { project: "example" }, + display, + }, + }); + } + const items = omp.timeline(); + expect(items).toEqual( + [1, 2].map(() => ({ + type: "tool_call", + callId: expect.stringMatching(/^omp-custom-/), + name: "project-context", + status: "completed", + detail: { type: "plain_text", text: "Project instructions" }, + metadata: { + synthetic: true, + customType: "project-context", + details: { project: "example" }, + }, + error: null, + })), + ); + expect(new Set(items.map((item) => item.type === "tool_call" && item.callId)).size).toBe(2); + expect(omp.completedTurnCount()).toBe(0); + omp.runtime().finishTurn(); + await waitForImmediate(); + expect(omp.completedTurnCount()).toBe(1); + }); + test("omits live custom messages when display is false", async () => { const omp = new OmpHarness(); await omp.start(); @@ -705,11 +745,19 @@ describe("OMP agent client and session", () => { message: "Background job DocsSmokeTwo completed", }, ]); - // Non-notice custom messages still fall through as assistant messages with - // their own id so the stream coalescer never glues them onto the open reply. - expect(omp.timeline().filter((item) => item.type === "assistant_message")).toEqual([ + expect( + omp.timeline().filter((item) => item.type !== "notification" && item.type !== "user_message"), + ).toEqual([ { type: "assistant_message", text: "done", messageId: "omp-assistant-1" }, - { type: "assistant_message", text: "plain custom status text", messageId: "omp-custom-1" }, + { + type: "tool_call", + callId: expect.stringMatching(/^omp-custom-/), + name: "custom-message", + status: "completed", + detail: { type: "plain_text", text: "plain custom status text" }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, + }, ]); }); diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index a04b4c8fb0..4c97014fe8 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -1,3 +1,4 @@ +import { mapCustomMessageToToolCall } from "../custom-message.js"; import { randomUUID } from "node:crypto"; import { existsSync } from "node:fs"; import { homedir } from "node:os"; @@ -69,11 +70,7 @@ import { } from "./provider-config.js"; export { formatOmpVersionSupport, resolveOmpDiagnosticPaths } from "./provider-config.js"; import { OmpSubagentCardTracker, type OmpSubagentCardScheduler } from "./subagent-card-tracker.js"; -import { - ompCustomMessageId, - ompSkillPromptUserText, - shouldDisplayOmpCustomMessage, -} from "./custom-message.js"; +import { ompSkillPromptUserText, shouldDisplayOmpCustomMessage } from "./custom-message.js"; import { getUserMessageText } from "./message-history.js"; import { mapOmpSystemNoticeToNotification } from "./system-notice.js"; import { materializeProviderImage } from "../provider-image-output.js"; @@ -720,7 +717,6 @@ export class OmpAgentSession implements AgentSession { private closed = false; private live: boolean; private readonly emittedUserMessageIds = new Set(); - private customMessageIndex = 0; private readonly usageSessionKey = randomUUID(); @@ -2000,14 +1996,8 @@ export class OmpAgentSession implements AgentSession { type: "timeline", provider: this.provider, turnId, - item: item ?? { - type: "assistant_message", - text, - messageId: ompCustomMessageId(event.message, () => { - this.customMessageIndex += 1; - return this.customMessageIndex; - }), - }, + item: + item ?? mapCustomMessageToToolCall(event.message, text, `omp-custom-${randomUUID()}`), }); } } diff --git a/packages/server/src/server/agent/providers/omp/history-mapper.test.ts b/packages/server/src/server/agent/providers/omp/history-mapper.test.ts index 02edb43af6..a747d3bee0 100644 --- a/packages/server/src/server/agent/providers/omp/history-mapper.test.ts +++ b/packages/server/src/server/agent/providers/omp/history-mapper.test.ts @@ -27,6 +27,40 @@ async function collectHistory( } describe("OMP history mapper", () => { + test("renders visible custom messages as completed tools with their type and content", async () => { + const events = await collectHistory([ + { + role: "custom", + customType: "project-context", + content: "Project instructions", + display: true, + }, + { + role: "custom", + customType: "private-context", + content: "Hidden instructions", + display: false, + }, + { + role: "custom", + customType: "project-context", + content: "Project instructions", + display: true, + }, + ]); + expect(events.map((event) => event.item)).toEqual( + [1, 2].map((index) => ({ + type: "tool_call", + callId: `omp-custom-${index}`, + name: "project-context", + status: "completed", + detail: { type: "plain_text", text: "Project instructions" }, + metadata: { synthetic: true, customType: "project-context" }, + error: null, + })), + ); + }); + test("replays a web search details error as failed when OMP sets isError false", async () => { const events = await collectHistory([ { @@ -204,9 +238,16 @@ describe("OMP history mapper", () => { events.push(event); expect(events.map((event) => event.item)).toEqual([ { - type: "assistant_message", - text: "[developer] External instruction", - messageId: "omp-custom-1", + type: "tool_call", + callId: "omp-custom-1", + name: "custom-message", + status: "completed", + detail: { + type: "plain_text", + text: "[developer] External instruction", + }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, }, ]); }); @@ -408,18 +449,26 @@ describe("OMP history mapper", () => { type: "timeline", provider: "omp", item: { - type: "assistant_message", - text: "visible explicit custom", - messageId: "omp-custom-1", + type: "tool_call", + callId: "omp-custom-1", + name: "custom-message", + status: "completed", + detail: { type: "plain_text", text: "visible explicit custom" }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, }, }, { type: "timeline", provider: "omp", item: { - type: "assistant_message", - text: "visible legacy custom", - messageId: "omp-custom-2", + type: "tool_call", + callId: "omp-custom-2", + name: "custom-message", + status: "completed", + detail: { type: "plain_text", text: "visible legacy custom" }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, }, }, { @@ -651,11 +700,26 @@ describe("OMP history mapper", () => { expect(events.map((event) => event.item)).toEqual([ { type: "user_message", text: "active branch", messageId: "user-active" }, { - type: "assistant_message", - text: "[future_control] Unsupported history record", - messageId: "omp-custom-1", + type: "tool_call", + callId: "omp-custom-1", + name: "custom-message", + status: "completed", + detail: { + type: "plain_text", + text: "[future_control] Unsupported history record", + }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, + }, + { + type: "tool_call", + callId: "omp-custom-2", + name: "custom-message", + status: "completed", + detail: { type: "plain_text", text: "[developer] developer note" }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, }, - { type: "assistant_message", text: "[developer] developer note", messageId: "omp-custom-2" }, ]); const omp = new FakeOmp(); @@ -758,11 +822,30 @@ describe("OMP history mapper", () => { expect(events.map((event) => event.item)).toEqual([ { type: "assistant_message", text: "Done.", messageId: "resp-1" }, { type: "user_message", text: "/skill:commit", messageId: "omp-custom-skill-1-user" }, - { type: "assistant_message", text: ircMessage, messageId: "omp-custom-irc-1" }, { - type: "assistant_message", - text: "visible without display flag", - messageId: "omp-custom-legacy-1", + type: "tool_call", + callId: "omp-custom-irc-1", + name: "irc:incoming", + status: "completed", + detail: { type: "plain_text", text: ircMessage }, + metadata: { + synthetic: true, + customType: "irc:incoming", + details: { from: "worker-1", message: "ready for review" }, + }, + error: null, + }, + { + type: "tool_call", + callId: "omp-custom-legacy-1", + name: "legacy-no-display", + status: "completed", + detail: { + type: "plain_text", + text: "visible without display flag", + }, + metadata: { synthetic: true, customType: "legacy-no-display" }, + error: null, }, ]); }); @@ -821,18 +904,44 @@ describe("OMP history mapper", () => { type: "timeline", provider: "omp", item: { - type: "assistant_message", - text: '[IMPORTANT: Agent invoked the "improve" skill.]', - messageId: "omp-custom-skill-agent", + type: "tool_call", + callId: "omp-custom-skill-agent", + name: "skill-prompt", + status: "completed", + detail: { + type: "plain_text", + text: '[IMPORTANT: Agent invoked the "improve" skill.]', + }, + metadata: { + synthetic: true, + customType: "skill-prompt", + details: { + name: "improve", + path: "/home/me/.agents/skills/improve/SKILL.md", + lineCount: 9, + }, + }, + error: null, }, }, { type: "timeline", provider: "omp", item: { - type: "assistant_message", - text: "\nworker-1\nhi\n", - messageId: "omp-custom-irc-user", + type: "tool_call", + callId: "omp-custom-irc-user", + name: "irc:incoming", + status: "completed", + detail: { + type: "plain_text", + text: "\nworker-1\nhi\n", + }, + metadata: { + synthetic: true, + customType: "irc:incoming", + details: { from: "worker-1", message: "hi" }, + }, + error: null, }, }, ]); diff --git a/packages/server/src/server/agent/providers/omp/message-history.ts b/packages/server/src/server/agent/providers/omp/message-history.ts index 112a942b64..1950502c89 100644 --- a/packages/server/src/server/agent/providers/omp/message-history.ts +++ b/packages/server/src/server/agent/providers/omp/message-history.ts @@ -1,3 +1,4 @@ +import { mapCustomMessageToToolCall } from "../custom-message.js"; import type { AgentStreamEvent, AgentTimelineItem, ToolCallDetail } from "../../agent-sdk-types.js"; import type { OmpAgentMessage, OmpImageContent, OmpTextContent } from "./rpc-types.js"; import type { OmpBridgedToolIdentity } from "./mcp-bridge.js"; @@ -155,7 +156,7 @@ export class OmpHistoryMapper { { type: "timeline", provider: this.provider, - item: { type: "assistant_message", text, messageId }, + item: mapCustomMessageToToolCall(message, text, messageId), }, ]; } 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 8035d59830..c6384e1fa2 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -1036,6 +1036,47 @@ describe("PiRpcAgentSession", () => { ]); }); + test("keeps custom context in separate tools while a turn continues", async () => { + const { pi, session, events } = await createSession(); + const runtime = pi.latestSession(); + await session.startTurn("Explain the project"); + runtime.emit({ type: "agent_start" }); + runtime.emit({ type: "turn_start" }); + for (const display of [true, false, true]) { + runtime.emit({ + type: "message_end", + message: { + role: "custom", + customType: "project-context", + content: [{ type: "text", text: "Project instructions" }], + details: { project: "example" }, + display, + }, + }); + } + const items = events.timelineItems(); + expect(items).toEqual( + [1, 2].map(() => ({ + type: "tool_call", + callId: expect.stringMatching(/^pi-custom-/), + name: "project-context", + status: "completed", + detail: { type: "plain_text", text: "Project instructions" }, + metadata: { + synthetic: true, + customType: "project-context", + details: { project: "example" }, + }, + error: null, + })), + ); + expect(new Set(items.map((item) => item.type === "tool_call" && item.callId)).size).toBe(2); + expect(events.turnCompletedEvents()).toHaveLength(0); + runtime.finishTurn(); + await events.nextTurnCompletion(); + await session.close(); + }); + test("surfaces Pi extension command messages and completes when no agent turn starts", async () => { const { pi, session, events } = await createSession(); const fakeSession = pi.latestSession(); @@ -1052,7 +1093,15 @@ describe("PiRpcAgentSession", () => { expect(events.timelineAndCompletionEvents()).toEqual([ { type: "timeline", - item: { type: "assistant_message", text: "Extension command output" }, + item: { + type: "tool_call", + callId: expect.stringMatching(/^pi-custom-/), + name: "custom-message", + status: "completed", + detail: { type: "plain_text", text: "Extension command output" }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, + }, }, { type: "turn_completed" }, ]); @@ -1081,7 +1130,18 @@ describe("PiRpcAgentSession", () => { }); expect(events.timelineItems()).toEqual([ - { type: "assistant_message", text: "Background process completed" }, + { + type: "tool_call", + callId: expect.stringMatching(/^pi-custom-/), + name: "custom-message", + status: "completed", + detail: { + type: "plain_text", + text: "Background process completed", + }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, + }, ]); expect(events.turnLifecycleEvents()).toEqual([{ type: "turn_started", turnId: undefined }]); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 5c9c7bb949..38a1f9815f 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -1,3 +1,4 @@ +import { mapCustomMessageToToolCall } from "../custom-message.js"; import { randomUUID } from "node:crypto"; import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { homedir, tmpdir } from "node:os"; @@ -2368,12 +2369,16 @@ export class PiRpcAgentSession implements AgentSession { const customMapping = this.extensionHost.mapCustomMessage(event.message); this.emitExtensionOutput(customMapping, turnId); const text = getUserMessageText(event.message.content); - if (text) { + if (event.message.display !== false && text) { this.emit({ type: "timeline", provider: this.provider, turnId, - item: { type: "assistant_message", text }, + item: mapCustomMessageToToolCall( + event.message, + text, + `${this.provider}-custom-${randomUUID()}`, + ), }); } if (!this.activeTurnStarted) { diff --git a/packages/server/src/server/agent/providers/pi/extensions/gotgenes-pi-subagents/index.test.ts b/packages/server/src/server/agent/providers/pi/extensions/gotgenes-pi-subagents/index.test.ts index 8e3d6e64e9..6bbe1b3b06 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/gotgenes-pi-subagents/index.test.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/gotgenes-pi-subagents/index.test.ts @@ -129,15 +129,28 @@ describe("@gotgenes/pi-subagents adapter", () => { expect(new Set(upserts.map((event) => event?.id)).size).toBe(1); expect(events.filter((event) => event.event.type === "timeline").length).toBeGreaterThan(0); }); - test("keeps claimed notification text visible in history", async () => { + test("keeps claimed notification text in a tool row in history", async () => { const fixture = readSubagentFixture(new URL("./fixtures/background.json", import.meta.url)); - const text = fixture.messages.find((message) => message.role === "custom")?.content; + const notification = fixture.messages.find((message) => message.role === "custom"); + const text = notification?.content; const events = []; for await (const event of streamPiHistory("pi", fixture.messages)) events.push(event); expect(events).toContainEqual({ type: "timeline", provider: "pi", - item: { type: "assistant_message", text }, + item: { + type: "tool_call", + callId: "pi-custom-1", + name: notification?.customType, + status: "completed", + detail: { type: "plain_text", text }, + metadata: { + synthetic: true, + customType: notification?.customType, + details: notification?.details, + }, + error: null, + }, }); expect(events.some((event) => event.type === "provider_subagent")).toBe(true); }); diff --git a/packages/server/src/server/agent/providers/pi/extensions/subagent-fixture-test.ts b/packages/server/src/server/agent/providers/pi/extensions/subagent-fixture-test.ts index 56d541a680..8a9ca61d08 100644 --- a/packages/server/src/server/agent/providers/pi/extensions/subagent-fixture-test.ts +++ b/packages/server/src/server/agent/providers/pi/extensions/subagent-fixture-test.ts @@ -65,16 +65,46 @@ export async function verifySubagentFixture( }); for (const message of fixture.messages) { if (message.role !== "custom" || typeof message.content !== "string") continue; - const assistantText = { - type: "assistant_message" as const, - text: message.content, - }; - expect(live).toContainEqual( - expect.objectContaining({ type: "timeline", item: assistantText }), - ); - expect(replay).toContainEqual( - expect.objectContaining({ type: "timeline", item: assistantText }), - ); + for (const events of [live, replay]) { + const tools = events.flatMap((event) => + event.type === "timeline" && + event.item.type === "tool_call" && + event.item.detail.type === "plain_text" && + event.item.detail.text === message.content + ? [event.item] + : [], + ); + expect(tools).toEqual( + message.display === false + ? [] + : [ + { + type: "tool_call", + callId: expect.stringMatching(/^pi-custom-/), + name: message.customType ?? "custom-message", + status: "completed", + detail: { + type: "plain_text", + text: message.content, + }, + metadata: { + synthetic: true, + customType: message.customType ?? "custom-message", + ...(message.details === undefined ? {} : { details: message.details }), + }, + error: null, + }, + ], + ); + expect( + events.filter( + (event) => + event.type === "timeline" && + event.item.type === "assistant_message" && + event.item.text === message.content, + ), + ).toEqual([]); + } } const liveSubagents = live.filter((event) => event.type === "provider_subagent"); const firstTimeline = liveSubagents.findIndex((event) => event.event.type === "timeline"); 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 380cf4ee75..24d18306b2 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 @@ -16,6 +16,35 @@ async function collectHistory( } describe("Pi history mapper", () => { + test("renders visible custom messages as completed tools and hides private context", async () => { + const events = await collectHistory([ + { + role: "custom", + customType: "project-context", + content: "Project instructions", + display: true, + }, + { + role: "custom", + customType: "private-context", + content: "Hidden instructions", + display: false, + }, + { role: "custom", customType: "project-context", content: "Project instructions" }, + ]); + expect(events.map((event) => event.item)).toEqual( + [1, 2].map((index) => ({ + type: "tool_call", + callId: `pi-custom-${index}`, + name: "project-context", + status: "completed", + detail: { type: "plain_text", text: "Project instructions" }, + metadata: { synthetic: true, customType: "project-context" }, + error: null, + })), + ); + }); + test("replays user, assistant, reasoning, and completed tool calls", async () => { await expect( collectHistory([ @@ -128,14 +157,22 @@ describe("Pi history mapper", () => { ]); }); - test("replays non-notice custom messages as assistant text, matching the live path", async () => { + test("replays custom messages as completed tools, matching the live path", async () => { await expect( collectHistory([{ role: "custom", content: "Extension command output" }]), ).resolves.toEqual([ { type: "timeline", provider: "pi", - item: { type: "assistant_message", text: "Extension command output" }, + item: { + type: "tool_call", + callId: "pi-custom-1", + name: "custom-message", + status: "completed", + detail: { type: "plain_text", text: "Extension command output" }, + metadata: { synthetic: true, customType: "custom-message" }, + error: null, + }, }, ]); }); 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 4ec02e0133..27a0250b00 100644 --- a/packages/server/src/server/agent/providers/pi/history-mapper.ts +++ b/packages/server/src/server/agent/providers/pi/history-mapper.ts @@ -1,3 +1,4 @@ +import { mapCustomMessageToToolCall } from "../custom-message.js"; import type { AgentStreamEvent, AgentTimelineItem, ToolCallDetail } from "../../agent-sdk-types.js"; import { createPiExtensionHost, @@ -61,6 +62,7 @@ export class PiHistoryMapper { private readonly hydrations: Promise[] = []; private userIndex = 0; private assistantIndex = 0; + private customIndex = 0; constructor( private readonly provider: string, @@ -125,6 +127,7 @@ export class PiHistoryMapper { ): AgentStreamEvent[] { const extensionMapping = this.extensionHost.mapCustomMessage(message); const extensionEvents = this.extensionEvents(extensionMapping); + if (message.display === false) return extensionEvents; const text = getUserMessageText(message.content); const mappedEvent = text ? this.hooks.mapCustomMessage?.(text, this.provider) : null; if (mappedEvent) { @@ -137,7 +140,11 @@ export class PiHistoryMapper { { type: "timeline", provider: this.provider, - item: { type: "assistant_message", text }, + item: mapCustomMessageToToolCall( + message, + text, + `${this.provider}-custom-${++this.customIndex}`, + ), } as AgentStreamEvent, ] : []),