Compare commits

...
Author SHA1 Message Date
Mohamed Boudra 8764ed25c1 fix(pi): accept native system messages with text blocks 2026-11-02 05:26:03 +01:00
Mohamed Boudra 017cfe861c refactor(pi): name runtime options and assert restored tool output 2026-11-02 05:16:24 +01:00
Mohamed Boudra a56442976a fix(pi): restore complete active-branch chat history
Replay the saved transcript ancestry instead of the compacted model context so older messages retain their native rewind identities and Fork chat boundaries.
2026-11-02 05:08:30 +01:00
9 changed files with 741 additions and 61 deletions
@@ -145,16 +145,6 @@ function piAssistantEntry(id: string, parentId: string): PiSessionEntry {
return { type: "message", id, parentId, message: { role: "assistant", content: [] } };
}
// Pi's context path: the entries from the root to the leaf.
function piBranchTo(entries: PiSessionEntry[], leafId: string): PiSessionEntry[] {
const byId = new Map(entries.map((entry) => [entry.id, entry]));
const branch: PiSessionEntry[] = [];
for (let entry = byId.get(leafId); entry; entry = byId.get(entry.parentId ?? "")) {
branch.unshift(entry);
}
return branch;
}
function parseEntryCapture(notification: string): unknown {
const prefix = "PASEO_ENTRY_CAPTURE ";
expect(notification.startsWith(prefix)).toBe(true);
@@ -1018,7 +1008,7 @@ describe("PiRpcAgentSession", () => {
const events = new SessionEvents(session);
const fakeSession = pi.latestSession();
fakeSession.treeUserEntries = [{ id: "entry-old", parentId: null, text: "old prompt" }];
fakeSession.contextUserEntries = fakeSession.treeUserEntries;
fakeSession.historyUserEntries = fakeSession.treeUserEntries;
await session.startTurn("new prompt", { clientMessageId: "client-new" });
fakeSession.finishSubmittedUserMessage({
@@ -1753,7 +1743,7 @@ describe("PiRpcAgentSession", () => {
await session.close();
});
test("captures the session's user entries apart from the ones on the current branch", async () => {
test("captures all session user entries as rewind targets, including abandoned branches", async () => {
const pi = new FakePi();
const session = await createClient(pi).createSession(createConfig());
onTestFinished(() => session.close());
@@ -1771,7 +1761,6 @@ describe("PiRpcAgentSession", () => {
const context = {
sessionManager: {
getEntries: () => entries,
buildContextEntries: () => piBranchTo(entries, "two-reply"),
},
ui: { notify: (message: string) => notifications.push(message) },
};
@@ -1786,10 +1775,6 @@ describe("PiRpcAgentSession", () => {
{ id: "abandoned", parentId: "one-reply", text: "rewound away" },
{ id: "two", parentId: "one-reply", text: "second" },
],
contextEntries: [
{ id: "one", parentId: null, text: "first" },
{ id: "two", parentId: "one-reply", text: "second" },
],
},
]);
});
@@ -3264,7 +3249,7 @@ describe("PiRpcAgentClient", () => {
{ id: "entry-1", parentId: null, text: "first prompt" },
{ id: "entry-3", parentId: "entry-2", text: "second prompt" },
];
pi.latestSession().contextUserEntries = pi.latestSession().treeUserEntries;
pi.latestSession().historyUserEntries = pi.latestSession().treeUserEntries;
await session.startTurn("first prompt");
pi.latestSession().finishTurn({ role: "assistant", content: [] });
@@ -3288,7 +3273,7 @@ describe("PiRpcAgentClient", () => {
{ id: "entry-1", parentId: null, text: "first prompt" },
{ id: "entry-3", parentId: "entry-2", text: "second prompt" },
];
fakeSession.contextUserEntries = [
fakeSession.historyUserEntries = [
{ id: "entry-3", parentId: "entry-2", text: "second prompt" },
];
@@ -687,8 +687,6 @@ function createPiPaseoExtensionFile({
requestId,
// Rewind targets: rows can still show entries that a rewind or compaction left off the branch.
treeEntries: toCapturedUserEntries(ctx.sessionManager.getEntries()),
// The entries getMessages() replays, so the nth one is the nth replayed user message.
contextEntries: toCapturedUserEntries(ctx.sessionManager.buildContextEntries()),
}),
"info",
);
@@ -1212,7 +1210,6 @@ export class PiRpcAgentSession implements AgentSession {
private readonly pendingPromptResults = new Map<string, boolean>();
private readonly pendingSteerSubmissions: PiPendingSteerSubmission[] = [];
currentLeafOverrideId: string | null | undefined;
private readonly contextUserEntries: PiCapturedEntry[] = [];
private readonly treeUserEntriesById = new Map<string, PiCapturedEntry>();
private readonly pendingExtensionResults = new Map<string, PendingExtensionResult>();
private compactionPromise: Promise<void> | null = null;
@@ -1455,10 +1452,11 @@ export class PiRpcAgentSession implements AgentSession {
async *streamHistory(): AsyncGenerator<AgentStreamEvent> {
await this.requestEntryCapture("history");
const history = await this.runtimeSession.getHistory();
yield* streamPiHistory(
this.provider,
await this.runtimeSession.getMessages(),
this.contextUserEntries,
history.messages,
history.userEntries,
{},
// At most eight 2 MiB child files per replay; later cards retain their summaries.
createPiExtensionHost(this.logger, undefined, 16 * 1024 * 1024),
@@ -2010,13 +2008,9 @@ export class PiRpcAgentSession implements AgentSession {
}
}
private recordCapturedUserEntries(input: {
treeEntries: PiCapturedEntry[];
contextEntries: PiCapturedEntry[];
}): void {
this.contextUserEntries.splice(0, this.contextUserEntries.length, ...input.contextEntries);
private recordCapturedUserEntries(entries: PiCapturedEntry[]): void {
this.treeUserEntriesById.clear();
for (const entry of input.treeEntries) {
for (const entry of entries) {
this.treeUserEntriesById.set(entry.id, entry);
}
}
@@ -2053,10 +2047,7 @@ export class PiRpcAgentSession implements AgentSession {
if (!payload) {
return false;
}
this.recordCapturedUserEntries({
treeEntries: parseCapturedEntries(payload.treeEntries),
contextEntries: parseCapturedEntries(payload.contextEntries),
});
this.recordCapturedUserEntries(parseCapturedEntries(payload.treeEntries));
if (typeof payload.requestId === "string") {
this.resolveExtensionResult(payload.requestId, undefined);
}
@@ -1,3 +1,5 @@
import { mkdir, mkdtemp, writeFile, readFile, rm } from "node:fs/promises";
import path from "node:path";
import type { ChildProcessWithoutNullStreams } from "node:child_process";
import { EventEmitter } from "node:events";
import { PassThrough } from "node:stream";
@@ -545,3 +547,102 @@ describe("PiCliRuntime", () => {
await expect(session.abort()).resolves.toBeUndefined();
});
});
async function createHistoryTestDirectory(prefix: string): Promise<string> {
const artifacts = path.resolve(".artifacts/pi-timeline");
await mkdir(artifacts, { recursive: true });
return mkdtemp(path.join(artifacts, prefix));
}
test("reads complete persisted history instead of compacted RPC context", async () => {
const root = await createHistoryTestDirectory("runtime-");
const file = path.join(root, "session.jsonl");
const records = [
{ type: "session", version: 3, id: "session" },
{
type: "message",
id: "old-user",
parentId: "system",
message: { role: "user", content: "before compaction" },
},
{
type: "message",
id: "old-answer",
parentId: "old-user",
message: { role: "assistant", content: [{ type: "text", text: "old answer" }] },
},
{ type: "compaction", id: "compact", parentId: "old-answer", summary: "summary" },
{
type: "message",
id: "new-user",
parentId: "compact",
message: { role: "user", content: "after compaction" },
},
];
records.splice(1, 0, {
type: "message",
id: "system",
parentId: null,
message: { role: "system", content: "private instructions" },
});
const contents = records.map((row) => JSON.stringify(row)).join("\n");
await writeFile(file, contents);
const child = createPiChild();
onPiCommand(child, (command) =>
writePiResponse(
child,
command,
command.type === "get_state" ? { sessionFile: file } : { messages: [records[5].message] },
),
);
const runtime = await createRuntime(child).startSession({ cwd: root });
try {
const history = await runtime.getHistory();
expect(history.userEntries.map((entry) => entry.id)).toEqual(["old-user", "new-user"]);
expect(history.messages).toEqual([records[2].message, records[3].message, records[5].message]);
expect(await readFile(file, "utf8")).toBe(contents);
await writeFile(
file,
contents +
'\n{"type":"custom","id":"rewind","parentId":"old-answer","customType":"paseo_rewind"}',
);
expect((await runtime.getHistory()).userEntries.map((entry) => entry.id)).toEqual(["old-user"]);
await writeFile(
file,
contents + '\n{"type":"custom","id":"rewind","parentId":null,"customType":"paseo_rewind"}',
);
expect((await runtime.getHistory()).messages).toEqual([]);
await writeFile(file, contents + '\n{"type":"custom","id":"bad","parentId":"missing"}');
await expect(runtime.getHistory()).rejects.toThrow("Missing Pi transcript ancestor");
} finally {
await runtime.close();
await rm(root, { recursive: true, force: true });
}
});
test("permits the unpersisted empty new session but refuses a missing resumed transcript", async () => {
const root = await createHistoryTestDirectory("empty-");
const file = path.join(root, "not-created.jsonl");
const child = createPiChild();
onPiCommand(child, (command) =>
writePiResponse(child, command, { sessionFile: file, messageCount: 0 }),
);
const runtime = createRuntime(child);
const fresh = await runtime.startSession({ cwd: root });
try {
expect(await fresh.getHistory()).toEqual({ messages: [], userEntries: [] });
} finally {
await fresh.close();
}
const resumedChild = createPiChild();
onPiCommand(resumedChild, (command) =>
writePiResponse(resumedChild, command, { sessionFile: file, messageCount: 0 }),
);
const resumed = await createRuntime(resumedChild).startSession({ cwd: root, session: file });
try {
await expect(resumed.getHistory()).rejects.toThrow("ENOENT");
} finally {
await resumed.close();
await rm(root, { recursive: true, force: true });
}
});
@@ -1,3 +1,5 @@
import { readFile } from "node:fs/promises";
import { readPiTranscript, type PiTranscript } from "./transcript.js";
import { createExternalProcessEnv } from "../../../paseo-env.js";
import type { ChildProcessWithoutNullStreams } from "node:child_process";
import type { Logger } from "pino";
@@ -16,7 +18,6 @@ import {
type PiStartSessionInput,
} from "./runtime.js";
import type {
PiAgentMessage,
PiModel,
PiPromptAck,
PiRpcCommand,
@@ -78,18 +79,27 @@ export class PiCliRuntime implements PiRuntime {
await process.close(input.signal.reason);
input.signal.throwIfAborted();
}
return new PiCliRuntimeSession(process, this.commandsRpcName, launch.env);
return new PiCliRuntimeSession({
process,
commandsRpcName: this.commandsRpcName,
launchEnvironment: launch.env,
newSession: input.session === undefined,
});
}
}
interface PiCliRuntimeSessionOptions {
process: JsonlRpcProcess;
commandsRpcName: string;
launchEnvironment: Record<string, string>;
newSession: boolean;
}
class PiCliRuntimeSession implements PiRuntimeSession {
private readonly subscribers = new Set<(event: PiRuntimeEvent) => void>();
constructor(
private readonly process: JsonlRpcProcess,
private readonly commandsRpcName: string,
private readonly launchEnvironment: Record<string, string>,
) {
constructor(private readonly options: PiCliRuntimeSessionOptions) {
const { process } = options;
process.onMessage((message) => {
this.emit(message as PiRuntimeEvent);
});
@@ -99,7 +109,7 @@ class PiCliRuntimeSession implements PiRuntimeSession {
}
get environment(): Record<string, string> {
return this.launchEnvironment;
return this.options.launchEnvironment;
}
onEvent(callback: (event: PiRuntimeEvent) => void): () => void {
@@ -113,7 +123,7 @@ class PiCliRuntimeSession implements PiRuntimeSession {
message: string,
images?: Array<{ type: "image"; data: string; mimeType: string }>,
): Promise<PiPromptAck> {
const { id: requestId, promise } = this.process.startRequest({
const { id: requestId, promise } = this.options.process.startRequest({
type: "prompt",
message,
...(images?.length ? { images } : {}),
@@ -162,9 +172,23 @@ class PiCliRuntimeSession implements PiRuntimeSession {
return (await this.request({ type: "get_state" })) as PiSessionState;
}
async getMessages(): Promise<PiAgentMessage[]> {
const data = (await this.request({ type: "get_messages" })) as { messages?: PiAgentMessage[] };
return data.messages ?? [];
async getHistory(): Promise<PiTranscript> {
const state = await this.getState();
if (!state.sessionFile) throw new Error("Pi history requires a persisted session file");
try {
return readPiTranscript(await readFile(state.sessionFile, "utf8"));
} catch (error) {
// Pi does not create the file until the first conversation message is saved.
if (
this.options.newSession &&
state.messageCount === 0 &&
error instanceof Error &&
"code" in error &&
error.code === "ENOENT"
)
return { messages: [], userEntries: [] };
throw error;
}
}
async getAvailableModels(timeoutMs?: number | null): Promise<PiModel[]> {
@@ -215,21 +239,21 @@ class PiCliRuntimeSession implements PiRuntimeSession {
}
async getCommands(): Promise<PiRpcSlashCommand[]> {
const data = (await this.request({ type: this.commandsRpcName })) as {
const data = (await this.request({ type: this.options.commandsRpcName })) as {
commands?: PiRpcSlashCommand[];
};
return data.commands ?? [];
}
sendRawFrame(frame: object & { type: string }): void {
this.process.send(frame);
this.options.process.send(frame);
}
respondToExtensionUiRequest(
id: string,
response: { value?: string; confirmed?: boolean; cancelled?: boolean },
): void {
this.process.send({ type: "extension_ui_response", id, ...response });
this.options.process.send({ type: "extension_ui_response", id, ...response });
}
cancelExtensionUiRequest(id: string): void {
@@ -237,22 +261,22 @@ class PiCliRuntimeSession implements PiRuntimeSession {
}
async close(): Promise<void> {
await this.process.close(new Error("Pi RPC session is closed"));
await this.options.process.close(new Error("Pi RPC session is closed"));
}
request(command: PiRpcCommand, timeoutMs?: number | null): Promise<unknown> {
return this.process.request(command, timeoutMs);
return this.options.process.request(command, timeoutMs);
}
private requestStopWork(command: PiRpcCommand): Promise<void> {
return this.process.requestStopWork(command);
return this.options.process.requestStopWork(command);
}
private async waitForCompletion(command: PiRpcCommand): Promise<void> {
// Pi only replies after its compaction work is durable. Its child process and
// session close paths already reject pending RPCs, so no elapsed-time failure
// is useful here.
await this.process.request(command, JSONL_RPC_NO_TIMEOUT);
await this.options.process.request(command, JSONL_RPC_NO_TIMEOUT);
}
private emit(event: PiRuntimeEvent): void {
@@ -1,3 +1,5 @@
import { buildAgentForkContextAttachment } from "../../activity-curator.js";
import { readPiTranscript } from "./transcript.js";
import { describe, expect, test } from "vitest";
import type { AgentStreamEvent } from "../../agent-sdk-types.js";
@@ -293,3 +295,235 @@ test("replayed Pi images keep their position before the following reply", async
),
).toEqual(["tool_call", expect.stringContaining("![Image]("), "Image read complete."]);
});
test("replays both sides of compaction with native IDs and isolates rewound branches", async () => {
const records = [
{ type: "session", version: 3, id: "session" },
{
type: "message",
id: "u1",
parentId: null,
message: { role: "user", content: "old question" },
},
{
type: "message",
id: "a1",
parentId: "u1",
message: { role: "assistant", content: [{ type: "text", text: "old answer" }] },
},
{ type: "compaction", id: "c1", parentId: "a1", summary: "summary", firstKeptEntryId: "a1" },
{ type: "message", id: "u2", parentId: "c1", message: { role: "user", content: "follow up" } },
];
const decode = (rows: unknown[]) =>
readPiTranscript(rows.map((row) => JSON.stringify(row)).join("\n"));
const history = decode(records);
expect(history.userEntries.map((entry) => entry.id)).toEqual(["u1", "u2"]);
expect(history.messages.map((message) => message.role)).toEqual(["user", "assistant", "user"]);
expect(await collectHistory(history.messages, history.userEntries)).toHaveLength(3);
const rewound = [
...records,
{ type: "custom", id: "r1", parentId: null, customType: "paseo_rewind" },
];
expect(decode(rewound).messages).toEqual([]);
const replacement = [
...rewound,
{
type: "message",
id: "u3",
parentId: "r1",
message: { role: "user", content: "replacement" },
},
];
expect(decode(replacement).userEntries.map((entry) => entry.id)).toEqual(["u3"]);
expect(
decode([
...replacement,
{ type: "custom", id: "r2", parentId: "a1", customType: "paseo_rewind" },
]).userEntries.map((entry) => entry.id),
).toEqual(["u1"]);
});
test("preserves native custom entries and legacy linear messages without inventing native IDs", async () => {
const custom = {
type: "custom_message",
id: "custom",
parentId: null,
customType: "notice",
content: "visible notice",
display: true,
details: { key: 1 },
};
const history = readPiTranscript(
[JSON.stringify({ type: "session", version: 3 }), JSON.stringify(custom)].join("\n"),
);
expect(history.messages).toEqual([{ ...custom, role: "custom" }]);
expect(await collectHistory(history.messages, history.userEntries)).toHaveLength(1);
const legacy = readPiTranscript(
'{"type":"session","version":1}\n{"type":"message","message":{"role":"user","content":"legacy"}}',
);
expect(legacy.messages).toEqual([{ role: "user", content: "legacy" }]);
expect(legacy.userEntries).toEqual([]);
});
test.each([
[{ type: "message", id: "x", parentId: "x", message: { role: "user", content: "cycle" } }],
[{ type: "message", id: "x", parentId: "missing", message: { role: "user", content: "orphan" } }],
[
{ type: "custom", id: "x", parentId: null },
{ type: "custom", id: "x", parentId: null },
],
[{ type: "custom", id: "x" }],
])("rejects corrupt ancestry rather than replaying a plausible branch: %j", (...entries) => {
expect(() =>
readPiTranscript(
[
JSON.stringify({ type: "session", version: 3 }),
...entries.map((entry) => JSON.stringify(entry)),
].join("\n"),
),
).toThrow();
});
test("Fork chat stops at a restored pre-compaction reply on the active branch", async () => {
const entries = [
{ type: "session", version: 3 },
{
type: "message",
id: "u1",
parentId: null,
message: { role: "user", content: "retained question" },
},
{
type: "message",
id: "a1",
parentId: "u1",
message: { role: "assistant", content: [{ type: "text", text: "retained answer" }] },
},
{
type: "message",
id: "abandoned",
parentId: "a1",
message: { role: "user", content: "abandoned branch" },
},
{ type: "compaction", id: "c1", parentId: "a1", summary: "summary" },
{
type: "message",
id: "later",
parentId: "c1",
message: { role: "user", content: "later question" },
},
];
const transcript = readPiTranscript(entries.map((entry) => JSON.stringify(entry)).join("\n"));
const events = await collectHistory(transcript.messages, transcript.userEntries);
const rows = events
.filter((event) => event.type === "timeline")
.map((event, index) => ({
seq: index + 1,
timestamp: "2026-10-31T00:00:00.000Z",
item: event.item,
}));
const fork = buildAgentForkContextAttachment({
rows,
cursorBoundary: { cursor: { epoch: "test", seq: 2 }, timelineEpoch: "test" },
});
expect(fork.attachment.text).toContain("retained question");
expect(fork.attachment.text).toContain("retained answer");
expect(fork.attachment.text).not.toContain("abandoned branch");
expect(fork.attachment.text).not.toContain("later question");
expect(transcript.userEntries.map((entry) => entry.id)).toEqual(["u1", "later"]);
});
test("keeps reasoning, image content and tool results across structural entries", async () => {
const messages: PiAgentMessage[] = [
{
role: "user",
content: [
{ type: "image", data: "aGVsbG8=", mimeType: "image/png" },
{ type: "text", text: "inspect" },
],
},
{
role: "assistant",
content: [
{ type: "thinking", thinking: "consider" },
{ type: "toolCall", id: "call", name: "read", arguments: { path: "a.txt" } },
],
},
{
role: "toolResult",
toolCallId: "call",
toolName: "read",
content: [{ type: "text", text: "file contents" }],
},
{ role: "assistant", content: [{ type: "text", text: "answer" }] },
];
const records = [
{ type: "session", version: 3 },
...messages.map((message, index) => ({
type: "message",
id: String(index),
parentId: index === 0 ? null : String(index - 1),
message,
})),
{ type: "compaction", id: "compact", parentId: "3", summary: "summary" },
];
const history = readPiTranscript(records.map((record) => JSON.stringify(record)).join("\n"));
expect(history.messages).toEqual(messages);
expect(history.userEntries).toEqual([{ id: "0", text: "inspect" }]);
expect(await collectHistory(history.messages, history.userEntries)).toEqual(
await collectHistory(messages, [{ id: "0", text: "inspect" }]),
);
});
test("traverses native system and summary messages without displaying private context", () => {
const roles = [
{ role: "system", content: "private instructions" },
{ role: "user", content: "question" },
{ role: "branchSummary", summary: "branch context" },
{ role: "compactionSummary", summary: "compacted context" },
{ role: "system", content: null },
];
const rows = [
{ type: "session", version: 3 },
...roles.map((message, index) => ({
type: "message",
id: String(index),
parentId: index === 0 ? null : String(index - 1),
message,
})),
];
expect(readPiTranscript(rows.map((row) => JSON.stringify(row)).join("\n"))).toEqual({
messages: [{ role: "user", content: "question" }],
userEntries: [{ id: "1", text: "question" }],
});
});
test("restores surrounding chat when native system content uses text blocks", async () => {
const records = [
{ type: "session", version: 3 },
{ type: "message", id: "u", parentId: null, message: { role: "user", content: "question" } },
{
type: "message",
id: "s",
parentId: "u",
message: { role: "system", content: [{ type: "text", text: "private instructions" }] },
},
{
type: "message",
id: "a",
parentId: "s",
message: { role: "assistant", content: [{ type: "text", text: "answer" }] },
},
];
const history = readPiTranscript(records.map((record) => JSON.stringify(record)).join("\n"));
expect(history.messages).toEqual([
{ role: "user", content: "question" },
{ role: "assistant", content: [{ type: "text", text: "answer" }] },
]);
expect(history.userEntries).toEqual([{ id: "u", text: "question" }]);
const events = await collectHistory(history.messages, history.userEntries);
expect(events.map((event) => event.item)).toEqual([
{ type: "user_message", text: "question", messageId: "u" },
{ type: "assistant_message", text: "answer", messageId: "pi-history-assistant-1" },
]);
});
@@ -1,5 +1,5 @@
import type { PiTranscript } from "./transcript.js";
import type {
PiAgentMessage,
PiModel,
PiPromptAck,
PiRpcSlashCommand,
@@ -55,7 +55,7 @@ export interface PiRuntimeSession {
setAutoCompaction(enabled: boolean): Promise<void>;
abort(): Promise<void>;
getState(): Promise<PiSessionState>;
getMessages(): Promise<PiAgentMessage[]>;
getHistory(): Promise<PiTranscript>;
getAvailableModels(timeoutMs?: number | null): Promise<PiModel[]>;
setModel(provider: string, modelId: string): Promise<PiModel>;
setThinkingLevel(level: string): Promise<void>;
@@ -150,8 +150,8 @@ export class FakePiSession implements PiRuntimeSession {
readonly rawFrames: Array<object & { type: string }> = [];
// Every user entry in the session file, including rewound and compacted ones.
treeUserEntries: FakePiUserEntry[] = [];
// The user entries on the current branch that getMessages() replays.
contextUserEntries: FakePiUserEntry[] = [];
// Native user entries paired with the complete active-branch history.
historyUserEntries: FakePiUserEntry[] = [];
abortRequested = false;
readonly canceledExtensionUiRequests: string[] = [];
readonly extensionUiResponses: Array<{
@@ -302,8 +302,8 @@ export class FakePiSession implements PiRuntimeSession {
return this.state;
}
async getMessages(): Promise<PiAgentMessage[]> {
return this.messages;
async getHistory() {
return { messages: this.messages, userEntries: this.historyUserEntries };
}
async getAvailableModels(_timeoutMs?: number | null): Promise<PiModel[]> {
@@ -508,7 +508,6 @@ export class FakePiSession implements PiRuntimeSession {
reason,
requestId,
treeEntries: this.treeUserEntries,
contextEntries: this.contextUserEntries,
})}`,
});
}
@@ -0,0 +1,134 @@
import { z } from "zod";
import { getUserMessageText, type PiCapturedUserMessageEntry } from "./history-mapper.js";
import type { PiAgentMessage } from "./rpc-types.js";
export interface PiTranscript {
messages: PiAgentMessage[];
userEntries: PiCapturedUserMessageEntry[];
}
const text = z.object({ type: z.literal("text"), text: z.string() }).passthrough();
const image = z
.object({ type: z.literal("image"), data: z.string(), mimeType: z.string() })
.passthrough();
const content = z.union([z.string(), z.array(z.union([text, image]))]);
const message = z.discriminatedUnion("role", [
z
.object({ role: z.literal("system"), content: z.union([z.string(), z.array(text)]).nullish() })
.passthrough(),
z.object({ role: z.literal("branchSummary"), summary: z.string() }).passthrough(),
z.object({ role: z.literal("compactionSummary"), summary: z.string() }).passthrough(),
z.object({ role: z.literal("user"), content }).passthrough(),
z
.object({
role: z.literal("custom"),
content,
customType: z.string().optional(),
display: z.boolean().optional(),
details: z.unknown().optional(),
})
.passthrough(),
z
.object({
role: z.literal("assistant"),
content: z.array(
z.union([
text,
z.object({ type: z.literal("thinking"), thinking: z.string() }).passthrough(),
z
.object({
type: z.literal("toolCall"),
id: z.string(),
name: z.string(),
arguments: z.unknown(),
})
.passthrough(),
]),
),
})
.passthrough(),
z
.object({
role: z.literal("toolResult"),
toolCallId: z.string(),
toolName: z.string(),
content: z.unknown(),
isError: z.boolean().optional(),
details: z.unknown().optional(),
})
.passthrough(),
z
.object({
role: z.literal("bashExecution"),
command: z.string(),
output: z.string().optional(),
exitCode: z.number().nullable().optional(),
cancelled: z.boolean().optional(),
timestamp: z.number(),
})
.passthrough(),
]);
const record = z
.object({
type: z.string(),
id: z.string().optional(),
parentId: z.string().nullable().optional(),
message: z.unknown().optional(),
})
.passthrough();
/** Select the persisted active ancestry, retaining structural links across compaction. */
export function readPiTranscript(contents: string): PiTranscript {
const records = contents
.split("\n")
.filter((line) => line.trim())
.map((line) => record.parse(JSON.parse(line)));
const header = records.shift();
if (header?.type !== "session") throw new Error("Pi transcript is missing its session header");
const legacy = header.version === 1;
const branch = legacy ? records : activeBranch(records);
const result: PiTranscript = { messages: [], userEntries: [] };
for (const entry of branch) {
if (entry.type !== "message" && entry.type !== "custom_message") continue;
const parsed = message.parse(
entry.type === "custom_message" ? { ...entry, role: "custom" } : entry.message,
);
// These native context messages had no display mapping in get_messages replay.
if (
parsed.role === "system" ||
parsed.role === "branchSummary" ||
parsed.role === "compactionSummary"
)
continue;
result.messages.push(parsed);
if (parsed.role === "user" && entry.id)
result.userEntries.push({
id: entry.id,
text: getUserMessageText(parsed.content),
});
}
return result;
}
function activeBranch(records: z.infer<typeof record>[]): z.infer<typeof record>[] {
const byId = new Map<string, (typeof records)[number]>();
for (const entry of records) {
if (!entry.id || entry.parentId === undefined || byId.has(entry.id))
throw new Error("Invalid Pi transcript tree entry");
byId.set(entry.id, entry);
}
const branch: z.infer<typeof record>[] = [];
const seen = new Set<string>();
let entry = records.at(-1);
while (entry) {
const id = entry.id!;
if (seen.has(id)) throw new Error("Cycle in Pi transcript ancestry");
seen.add(id);
branch.push(entry);
if (entry.parentId === null) break;
entry = byId.get(entry.parentId!);
if (!entry) throw new Error("Missing Pi transcript ancestor");
}
branch.reverse();
return branch;
}
@@ -1,5 +1,7 @@
import { mkdtempSync, rmSync } from "node:fs";
import { readdir } from "node:fs/promises";
import { readdir, mkdir, mkdtemp, writeFile, readFile, rm } from "node:fs/promises";
import { randomUUID } from "node:crypto";
import { PiRpcAgentClient } from "../agent/providers/pi/agent.js";
import { tmpdir } from "node:os";
import path from "node:path";
import pino from "pino";
@@ -366,3 +368,213 @@ describe("daemon E2E (real pi) - rewind", () => {
}
}, 900_000);
});
test("seeded compacted Pi history survives daemon reload, native rewind and Fork chat", async (context) => {
if (!(await canRunRealProvider("pi"))) {
context.skip();
return;
}
const artifacts = path.resolve(".artifacts/pi-timeline");
await mkdir(artifacts, { recursive: true });
const root = await mkdtemp(path.join(artifacts, "daemon-history-"));
const cwd = path.join(root, "cwd");
await mkdir(cwd);
const file = path.join(root, "session.jsonl");
const nativeId = randomUUID();
const timestamp = new Date().toISOString();
const assistant = (text: string) => ({
role: "assistant",
content: [{ type: "text", text }],
api: "openai-completions",
provider: "openrouter",
model: "google/gemini-2.5-flash-lite",
stopReason: "stop",
timestamp: Date.now(),
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
});
const records = [
{ type: "session", version: 3, id: nativeId, cwd, timestamp },
{
type: "message",
id: "system",
parentId: null,
timestamp,
message: {
role: "system",
content: [{ type: "text", text: "private system context" }],
timestamp: Date.now(),
},
},
{
type: "message",
id: "u1",
parentId: "system",
timestamp,
message: { role: "user", content: "old question", timestamp: Date.now() },
},
{
type: "message",
id: "tool-request",
parentId: "u1",
timestamp,
message: {
...assistant(""),
content: [
{ type: "toolCall", id: "read-call", name: "read", arguments: { path: "readme.txt" } },
],
},
},
{
type: "message",
id: "tool-result",
parentId: "tool-request",
timestamp,
message: {
role: "toolResult",
toolCallId: "read-call",
toolName: "read",
content: [{ type: "text", text: "saved tool output" }],
isError: false,
timestamp: Date.now(),
},
},
{
type: "message",
id: "a1",
parentId: "tool-result",
timestamp,
message: assistant("old answer"),
},
{
type: "message",
id: "abandoned",
parentId: "a1",
timestamp,
message: { role: "user", content: "abandoned question", timestamp: Date.now() },
},
{
type: "compaction",
id: "compact",
parentId: "a1",
timestamp,
summary: "old summary",
firstKeptEntryId: null,
tokensBefore: 50000,
},
{
type: "message",
id: "u2",
parentId: "compact",
timestamp,
message: { role: "user", content: "follow up", timestamp: Date.now() },
},
{ type: "message", id: "a2", parentId: "u2", timestamp, message: assistant("new answer") },
];
await writeFile(file, records.map((record) => JSON.stringify(record)).join("\n") + "\n");
const logger = pino({ level: "silent" });
const daemon = await createTestPaseoDaemon({
agentClients: { pi: new PiRpcAgentClient({ logger }) },
logger,
paseoHomeRoot: path.join(root, "home"),
});
const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` });
try {
await client.connect();
await client.fetchAgents({ subscribe: {} });
const imported = await client.importAgent({ provider: "pi", sessionId: file, cwd });
const session = { agentId: imported.id, cwd };
const expected = { userTexts: ["old question", "follow up"], assistantCount: 2 };
expectTimeline(await fetchTimelineItems(client, imported.id), expected);
expect(
userMessageIdForToken(await fetchTimelineItems(client, imported.id), "old question"),
).toBe("u1");
expect(userMessageIdForToken(await fetchTimelineItems(client, imported.id), "follow up")).toBe(
"u2",
);
expect(await runtimeSessionId({ client, daemon }, session)).toBe(nativeId);
await client.refreshAgent(imported.id);
const restored = await fetchTimelineItems(client, imported.id);
expectTimeline(restored, expected);
const tools = restored.filter((item) => item.type === "tool_call");
expect(tools).toHaveLength(1);
expect(tools[0].detail).toEqual({
type: "read",
filePath: "readme.txt",
content: "saved tool output",
});
expect(
restored.filter((item) => item.type === "user_message").map((item) => item.messageId),
).toEqual(["u1", "u2"]);
const oldAnswer = restored.find(
(item) => item.type === "assistant_message" && item.text === "old answer",
);
if (oldAnswer?.type !== "assistant_message" || !oldAnswer.messageId)
throw new Error("Restored reply missing identity");
const beforeFork = await readFile(file, "utf8");
const fork = await client.buildAgentForkContext(imported.id, {
boundaryMessageId: oldAnswer.messageId,
});
expect(fork.attachment?.text).toContain("old question");
expect(fork.attachment?.text).toContain("old answer");
expect(fork.attachment?.text).not.toContain("follow up");
expect(fork.attachment?.text).not.toContain("abandoned question");
expect(await readFile(file, "utf8")).toBe(beforeFork);
const forkFile = path.join(root, "native-fork.jsonl");
const forkId = randomUUID();
await writeFile(
forkFile,
[
{ ...records[0], id: forkId },
...records.slice(1),
{
type: "custom",
id: "fork-leaf",
parentId: "a1",
timestamp,
customType: "paseo_rewind",
data: { targetId: "a1" },
},
]
.map((record) => JSON.stringify(record))
.join("\n") + "\n",
);
const importedFork = await client.importAgent({ provider: "pi", sessionId: forkFile, cwd });
expectTimeline(await fetchTimelineItems(client, importedFork.id), {
userTexts: ["old question"],
assistantCount: 1,
});
expect(await runtimeSessionId({ client, daemon }, { agentId: importedFork.id, cwd })).toBe(
forkId,
);
expect(await readFile(file, "utf8")).toBe(beforeFork);
await client.rewindAgent(imported.id, "u2", "conversation");
expectTimeline(await fetchTimelineItems(client, imported.id), {
userTexts: ["old question"],
assistantCount: 1,
});
await client.refreshAgent(imported.id);
expectTimeline(await fetchTimelineItems(client, imported.id), {
userTexts: ["old question"],
assistantCount: 1,
});
expect(await runtimeSessionId({ client, daemon }, session)).toBe(nativeId);
await client.rewindAgent(imported.id, "u1", "conversation");
await client.refreshAgent(imported.id);
expectTimeline(await fetchTimelineItems(client, imported.id), {
userTexts: [],
assistantCount: 0,
});
expect(await runtimeSessionId({ client, daemon }, session)).toBe(nativeId);
} finally {
await client.close();
await daemon.close();
await rm(root, { recursive: true, force: true });
}
}, 120_000);