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.
This commit is contained in:
Mohamed Boudra
2026-10-05 00:03:20 +01:00
committed by GitHub
parent 2187124642
commit 1f8a4559a8
13 changed files with 393 additions and 65 deletions
@@ -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",
},
});
@@ -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,
},
};
}
@@ -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<AgentTimelineItem, { type: "tool_call" }> {
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,
};
}
@@ -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,
},
]);
});
@@ -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<string>();
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()}`),
});
}
}
@@ -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: "<irc>\n<from>worker-1</from>\n<message>hi</message>\n</irc>",
messageId: "omp-custom-irc-user",
type: "tool_call",
callId: "omp-custom-irc-user",
name: "irc:incoming",
status: "completed",
detail: {
type: "plain_text",
text: "<irc>\n<from>worker-1</from>\n<message>hi</message>\n</irc>",
},
metadata: {
synthetic: true,
customType: "irc:incoming",
details: { from: "worker-1", message: "hi" },
},
error: null,
},
},
]);
@@ -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),
},
];
}
@@ -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 }]);
@@ -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) {
@@ -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);
});
@@ -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");
@@ -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,
},
},
]);
});
@@ -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<AgentStreamEvent[]>[] = [];
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,
]
: []),