From e1b20384a10680ca6eaae50fa01fd55a4dcd697d Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 3 Nov 2026 12:20:34 +0700 Subject: [PATCH] fix(pi): restore complete chat history after compaction (#8523) * 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. * refactor(pi): name runtime options and assert restored tool output * fix(pi): accept native system messages with text blocks --- .../server/agent/providers/pi/agent.test.ts | 25 +- .../src/server/agent/providers/pi/agent.ts | 21 +- .../agent/providers/pi/cli-runtime.test.ts | 101 ++++++++ .../server/agent/providers/pi/cli-runtime.ts | 62 +++-- .../agent/providers/pi/history-mapper.test.ts | 234 ++++++++++++++++++ .../src/server/agent/providers/pi/runtime.ts | 4 +- .../agent/providers/pi/test-utils/fake-pi.ts | 13 +- .../server/agent/providers/pi/transcript.ts | 134 ++++++++++ .../daemon-e2e/pi-rewind.real.e2e.test.ts | 214 +++++++++++++++- 9 files changed, 744 insertions(+), 64 deletions(-) create mode 100644 packages/server/src/server/agent/providers/pi/transcript.ts 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 008a3a4fac..039d3eaf22 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -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({ @@ -1670,7 +1660,7 @@ describe("PiRpcAgentSession", () => { test("closes a history-only Pi runtime when replay fails", async () => { const pi = new FakePi(); pi.queueSessionSetup((runtimeSession) => { - runtimeSession.getMessagesError = new Error("History unavailable"); + runtimeSession.getHistoryError = new Error("History unavailable"); }); const session = await createClient(pi).resumeSession( { @@ -1858,7 +1848,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()); @@ -1876,7 +1866,6 @@ describe("PiRpcAgentSession", () => { const context = { sessionManager: { getEntries: () => entries, - buildContextEntries: () => piBranchTo(entries, "two-reply"), }, ui: { notify: (message: string) => notifications.push(message) }, }; @@ -1891,10 +1880,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" }, - ], }, ]); }); @@ -3369,7 +3354,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: [] }); @@ -3393,7 +3378,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" }, ]; diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 368c7405a7..198d596c9b 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -688,8 +688,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", ); @@ -1213,7 +1211,6 @@ export class PiRpcAgentSession implements AgentSession { private readonly pendingPromptResults = new Map(); private readonly pendingSteerSubmissions: PiPendingSteerSubmission[] = []; currentLeafOverrideId: string | null | undefined; - private readonly contextUserEntries: PiCapturedEntry[] = []; private readonly treeUserEntriesById = new Map(); private readonly pendingExtensionResults = new Map(); private compactionPromise: Promise | null = null; @@ -1459,10 +1456,11 @@ export class PiRpcAgentSession implements AgentSession { async *streamHistory(): AsyncGenerator { try { 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), @@ -2017,13 +2015,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); } } @@ -2060,10 +2054,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); } diff --git a/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts b/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts index 463870760f..8a838d74d9 100644 --- a/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts +++ b/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts @@ -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 { + 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 }); + } +}); diff --git a/packages/server/src/server/agent/providers/pi/cli-runtime.ts b/packages/server/src/server/agent/providers/pi/cli-runtime.ts index 62ee01c503..78c1fc7cf6 100644 --- a/packages/server/src/server/agent/providers/pi/cli-runtime.ts +++ b/packages/server/src/server/agent/providers/pi/cli-runtime.ts @@ -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; + 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, - ) { + 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 { - 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 { - 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 { - const data = (await this.request({ type: "get_messages" })) as { messages?: PiAgentMessage[] }; - return data.messages ?? []; + async getHistory(): Promise { + 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 { @@ -215,21 +239,21 @@ class PiCliRuntimeSession implements PiRuntimeSession { } async getCommands(): Promise { - 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 { - 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 { - return this.process.request(command, timeoutMs); + return this.options.process.request(command, timeoutMs); } private requestStopWork(command: PiRpcCommand): Promise { - return this.process.requestStopWork(command); + return this.options.process.requestStopWork(command); } private async waitForCompletion(command: PiRpcCommand): Promise { // 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 { 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 f4136956eb..b6f4349f47 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 @@ -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" }, + ]); +}); diff --git a/packages/server/src/server/agent/providers/pi/runtime.ts b/packages/server/src/server/agent/providers/pi/runtime.ts index 1064952e6d..42399d3275 100644 --- a/packages/server/src/server/agent/providers/pi/runtime.ts +++ b/packages/server/src/server/agent/providers/pi/runtime.ts @@ -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; abort(): Promise; getState(): Promise; - getMessages(): Promise; + getHistory(): Promise; getAvailableModels(timeoutMs?: number | null): Promise; setModel(provider: string, modelId: string): Promise; setThinkingLevel(level: string): Promise; diff --git a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts index 69aab7b3b8..770f24625c 100644 --- a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts +++ b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts @@ -150,8 +150,8 @@ export class FakePiSession implements PiRuntimeSession { readonly rawFrames: Array = []; // 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; closeCalls = 0; readonly canceledExtensionUiRequests: string[] = []; @@ -162,7 +162,7 @@ export class FakePiSession implements PiRuntimeSession { setModelResult: PiModel | null = null; models: PiModel[] = []; messages: PiAgentMessage[] = []; - getMessagesError: Error | null = null; + getHistoryError: Error | null = null; stats: PiSessionStats = { tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, cost: 0, @@ -304,9 +304,9 @@ export class FakePiSession implements PiRuntimeSession { return this.state; } - async getMessages(): Promise { - if (this.getMessagesError) throw this.getMessagesError; - return this.messages; + async getHistory() { + if (this.getHistoryError) throw this.getHistoryError; + return { messages: this.messages, userEntries: this.historyUserEntries }; } async getAvailableModels(_timeoutMs?: number | null): Promise { @@ -513,7 +513,6 @@ export class FakePiSession implements PiRuntimeSession { reason, requestId, treeEntries: this.treeUserEntries, - contextEntries: this.contextUserEntries, })}`, }); } diff --git a/packages/server/src/server/agent/providers/pi/transcript.ts b/packages/server/src/server/agent/providers/pi/transcript.ts new file mode 100644 index 0000000000..b4ead7bc5c --- /dev/null +++ b/packages/server/src/server/agent/providers/pi/transcript.ts @@ -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[]): z.infer[] { + const byId = new Map(); + 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[] = []; + const seen = new Set(); + 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; +} diff --git a/packages/server/src/server/daemon-e2e/pi-rewind.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/pi-rewind.real.e2e.test.ts index d198be9e09..704243d503 100644 --- a/packages/server/src/server/daemon-e2e/pi-rewind.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/pi-rewind.real.e2e.test.ts @@ -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);