fix(server): hide Pi custom messages marked display: false

pi-goal-x sends its `<pi_goal_continuation …/>` checkpoint trigger as a custom
message with `display: false`: the model reads it, and Pi's TUI never renders
it. Paseo ignored the flag and emitted every custom message as assistant text,
which the timeline projection then glued onto the previous answer. The live
and history paths now read custom text through one function that respects the
flag. Extension adapters still see hidden messages, so structured mappings are
unchanged; pi-subagents' hidden `subagent-notify` text is dropped the same way.

Also removes the unused `mapCustomMessage` history hook left over from the
Oh My Pi removal.
This commit is contained in:
Mohamed Boudra
2026-09-26 19:56:48 +02:00
parent 6639847425
commit d3aee8c2df
6 changed files with 60 additions and 21 deletions
@@ -24,6 +24,7 @@ import {
transformPiModels,
} from "./agent.js";
import { FakePi } from "./test-utils/fake-pi.js";
import { HIDDEN_GOAL_CONTINUATION } from "./test-utils/hidden-custom-message.js";
import { createPiExtensionHost } from "./extensions/index.js";
import { PiExtensionHost } from "./extensions/host.js";
import type { PiModel, PiThinkingLevel } from "./rpc-types.js";
@@ -1080,6 +1081,17 @@ describe("PiRpcAgentSession", () => {
]);
});
test("hides Pi custom messages that the extension marks display: false", async () => {
const { pi, events } = await createSession();
const fakeSession = pi.latestSession();
fakeSession.emit({ type: "agent_start" });
fakeSession.emit({ type: "turn_start" });
fakeSession.emit({ type: "message_end", message: HIDDEN_GOAL_CONTINUATION });
expect(events.timelineItems()).toEqual([]);
});
test("canceling a silent Pi extension command leaves the session usable", async () => {
const { pi, session, events } = await createSession();
const fakeSession = pi.latestSession();
@@ -59,7 +59,7 @@ import {
toDiagnosticErrorMessage,
} from "../diagnostic-utils.js";
import {
getUserMessageText,
getCustomMessageText,
streamPiHistory,
type PiCapturedUserMessageEntry,
} from "./history-mapper.js";
@@ -2300,7 +2300,7 @@ export class PiRpcAgentSession implements AgentSession {
if (event.message.role === "custom") {
const customMapping = this.extensionHost.mapCustomMessage(event.message);
this.emitExtensionOutput(customMapping, turnId);
const text = getUserMessageText(event.message.content);
const text = getCustomMessageText(event.message);
if (text) {
this.emit({
type: "timeline",
@@ -65,16 +65,18 @@ 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 }),
);
const assistantText = expect.objectContaining({
type: "timeline",
item: { type: "assistant_message", text: message.content },
});
// Claimed notifications keep their text unless the extension hides it from the user.
if (message.display === false) {
expect(live).not.toContainEqual(assistantText);
expect(replay).not.toContainEqual(assistantText);
} else {
expect(live).toContainEqual(assistantText);
expect(replay).toContainEqual(assistantText);
}
}
const liveSubagents = live.filter((event) => event.type === "provider_subagent");
const firstTimeline = liveSubagents.findIndex((event) => event.event.type === "timeline");
@@ -3,6 +3,7 @@ import { describe, expect, test } from "vitest";
import type { AgentStreamEvent } from "../../agent-sdk-types.js";
import { streamPiHistory, type PiCapturedUserMessageEntry } from "./history-mapper.js";
import type { PiAgentMessage } from "./rpc-types.js";
import { HIDDEN_GOAL_CONTINUATION } from "./test-utils/hidden-custom-message.js";
async function collectHistory(
messages: PiAgentMessage[],
@@ -140,6 +141,10 @@ describe("Pi history mapper", () => {
]);
});
test("drops custom messages that the extension marks display: false", async () => {
await expect(collectHistory([HIDDEN_GOAL_CONTINUATION])).resolves.toEqual([]);
});
test("uses Pi tree entry ids for replayed user messages", async () => {
await expect(
collectHistory(
@@ -20,10 +20,6 @@ export interface PiCapturedUserMessageEntry {
}
export interface PiHistoryMapperHooks {
mapCustomMessage?: (
text: string,
provider: string,
) => Extract<AgentStreamEvent, { type: "timeline" }> | null;
resolveToolCallId?: (toolCallId: string, toolCall: PiTrackedToolCall) => string;
mapToolDetail?: (
toolCall: PiTrackedToolCall,
@@ -56,6 +52,11 @@ export function getUserMessageText(content: string | (PiTextContent | PiImageCon
return textParts.join("\n\n");
}
/** The text Pi shows for a custom message. Extensions hide model-only messages with `display: false`. */
export function getCustomMessageText(message: Extract<PiAgentMessage, { role: "custom" }>): string {
return message.display === false ? "" : getUserMessageText(message.content);
}
export class PiHistoryMapper {
private readonly pendingToolCalls = new Map<string, PiTrackedToolCall>();
private readonly hydrations: Promise<AgentStreamEvent[]>[] = [];
@@ -125,11 +126,7 @@ export class PiHistoryMapper {
): AgentStreamEvent[] {
const extensionMapping = this.extensionHost.mapCustomMessage(message);
const extensionEvents = this.extensionEvents(extensionMapping);
const text = getUserMessageText(message.content);
const mappedEvent = text ? this.hooks.mapCustomMessage?.(text, this.provider) : null;
if (mappedEvent) {
return [...extensionEvents, mappedEvent];
}
const text = getCustomMessageText(message);
return [
...extensionEvents,
...(text
@@ -0,0 +1,23 @@
import type { PiAgentMessage } from "../rpc-types.js";
// A pi-goal-x 0.31.9 (2e8ad767) checkpoint continuation, captured verbatim from a real
// `pi --mode rpc` run (Pi 0.84.4) on 2026-09-26. pi-goal-x sends it with `display: false`:
// the marker is for the model, and Pi's TUI never renders it.
export const HIDDEN_GOAL_CONTINUATION: PiAgentMessage = {
role: "custom",
customType: "pi-goal-event",
content: '<pi_goal_continuation goal_id="muioeg51-6db1a7" kind="checkpoint" v="2"/>',
display: false,
details: {
version: 3,
generation: "13a8082c-e249-4df5-ac0e-fdc863ad289e",
dispatchId: "a9c0ef02-d39c-42a2-8d5a-993a96b4e827",
kind: "checkpoint",
goalId: "muioeg51-6db1a7",
status: "active",
revision: 3,
checkpointSeq: 1,
timestamp: 1790444456255,
},
timestamp: 1790444456256,
};