From 43a2a7969cbf049455b998d38ccd18ce179a88d1 Mon Sep 17 00:00:00 2001 From: "paseo-bot[bot]" <266920839+paseo-bot[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 13:16:40 +0200 Subject: [PATCH] Rewind a Pi chat to the message you pick after an earlier rewind (#5383) * Rewind a Pi chat to the message you pick after an earlier rewind After a rewind, Pi's session file keeps the abandoned branch. Paseo captured user entry ids from the whole file but replayed history from the current branch, pairing them by position. Each abandoned user message shifted later rows onto an earlier entry id, so rewinding a replayed row went back further than the selected message. Capture user entries from the same compaction-aware branch Pi replays, so the nth capture is the nth replayed user message. * Assert Pi entry captures as data with named test helpers * Name the fields of Pi user entry test fixtures * Rewind a Pi row whose entry compaction summarized away The capture fed both replay pairing and the rewind target lookup. Narrowing it to the context branch broke rewinding a live row from before /compact, because that row's entry is off the branch. The extension now reports the session's user entries for the lookup apart from the branch entries replay pairs with. --------- Co-authored-by: paseo-bot[bot] <266920839+paseo-bot[bot]@users.noreply.github.com> --- .../server/agent/providers/pi/agent.test.ts | 100 ++++++++++++++++- .../src/server/agent/providers/pi/agent.ts | 43 +++++--- .../agent/providers/pi/test-utils/fake-pi.ts | 8 +- .../daemon-e2e/pi-rewind.real.e2e.test.ts | 104 +++++++++++++++++- 4 files changed, 233 insertions(+), 22 deletions(-) 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 0712160ecd..917829908e 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -120,6 +120,42 @@ function readUtf8File(pathname: string): string { type PaseoExtensionListener = (event: unknown, context?: unknown) => unknown; +interface PiSessionEntry { + type: "message"; + id: string; + parentId: string | null; + message: { role: string; content: unknown }; +} + +function piUserEntry(input: { id: string; parentId: string | null; text: string }): PiSessionEntry { + return { + type: "message", + id: input.id, + parentId: input.parentId, + message: { role: "user", content: input.text }, + }; +} + +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); + return JSON.parse(notification.slice(prefix.length)); +} + async function loadPaseoExtensionListeners( extensionPath: string, ): Promise> { @@ -921,7 +957,8 @@ describe("PiRpcAgentSession", () => { })) as PiRpcAgentSession; const events = new SessionEvents(session); const fakeSession = pi.latestSession(); - fakeSession.capturedUserEntries = [{ id: "entry-old", parentId: null, text: "old prompt" }]; + fakeSession.treeUserEntries = [{ id: "entry-old", parentId: null, text: "old prompt" }]; + fakeSession.contextUserEntries = fakeSession.treeUserEntries; await session.startTurn("new prompt", { clientMessageId: "client-new" }); fakeSession.finishSubmittedUserMessage({ @@ -1511,6 +1548,47 @@ describe("PiRpcAgentSession", () => { await session.close(); }); + test("captures the session's user entries apart from the ones on the current branch", async () => { + const pi = new FakePi(); + const session = await createClient(pi).createSession(createConfig()); + onTestFinished(() => session.close()); + const listeners = await loadPaseoExtensionListeners(pi.recordedLaunches[0]!.extensionPaths[0]!); + // "abandoned" was rewound; "two" was sent from the same parent afterwards. + const entries = [ + piUserEntry({ id: "one", parentId: null, text: "first" }), + piAssistantEntry("one-reply", "one"), + piUserEntry({ id: "abandoned", parentId: "one-reply", text: "rewound away" }), + piAssistantEntry("abandoned-reply", "abandoned"), + piUserEntry({ id: "two", parentId: "one-reply", text: "second" }), + piAssistantEntry("two-reply", "two"), + ]; + const notifications: string[] = []; + const context = { + sessionManager: { + getEntries: () => entries, + buildContextEntries: () => piBranchTo(entries, "two-reply"), + }, + ui: { notify: (message: string) => notifications.push(message) }, + }; + + await listeners.get("session_start")?.({}, context); + + expect(notifications.map(parseEntryCapture)).toEqual([ + { + reason: "session_start", + treeEntries: [ + { id: "one", parentId: null, text: "first" }, + { 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" }, + ], + }, + ]); + }); + test("appends agent and daemon prompts after Pi's discovered system prompt", async () => { const pi = new FakePi(); const client = createClient(pi); @@ -2856,10 +2934,11 @@ describe("PiRpcAgentClient", () => { test("rewinds conversation through the Pi tree navigation bridge", async () => { const { pi, session, events } = await createSession(); - pi.latestSession().capturedUserEntries = [ + pi.latestSession().treeUserEntries = [ { id: "entry-1", parentId: null, text: "first prompt" }, { id: "entry-3", parentId: "entry-2", text: "second prompt" }, ]; + pi.latestSession().contextUserEntries = pi.latestSession().treeUserEntries; await session.startTurn("first prompt"); pi.latestSession().finishTurn({ role: "assistant", content: [] }); @@ -2875,6 +2954,23 @@ describe("PiRpcAgentClient", () => { expect(pi.latestSession().treeNavigationRequests).toEqual(["entry-1"]); }); + test("rewinds a row whose Pi entry compaction summarized away", async () => { + const { pi, session } = await createSession(); + const fakeSession = pi.latestSession(); + // Compaction summarized entry-1, so the replayed branch starts at entry-3. + fakeSession.treeUserEntries = [ + { id: "entry-1", parentId: null, text: "first prompt" }, + { id: "entry-3", parentId: "entry-2", text: "second prompt" }, + ]; + fakeSession.contextUserEntries = [ + { id: "entry-3", parentId: "entry-2", text: "second prompt" }, + ]; + + await session.revertConversation?.({ messageId: "entry-1" }); + + expect(fakeSession.treeNavigationRequests).toEqual(["entry-1"]); + }); + test("injects MCP servers without replacing the Pi global MCP config", async () => { const agentDir = mkdtempSync(path.join(tmpdir(), "paseo-pi-agent-")); onTestFinished(() => rmSync(agentDir, { recursive: true, force: true })); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 4c2760cdba..c4fe1dadaf 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -622,9 +622,8 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile { .join("\\n\\n"); } - function getCapturedUserEntries(ctx) { - return ctx.sessionManager - .getEntries() + function toCapturedUserEntries(entries) { + return entries .filter((entry) => entry.type === "message" && entry.message?.role === "user") .map(toCapturedUserEntry); } @@ -640,7 +639,14 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile { function emitEntryCapture(ctx, reason, requestId) { ctx.ui.notify( "${PASEO_PI_ENTRY_CAPTURE_MARKER} " + - JSON.stringify({ reason, requestId, entries: getCapturedUserEntries(ctx) }), + JSON.stringify({ + reason, + 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", ); } @@ -1265,8 +1271,8 @@ export class PiRpcAgentSession implements AgentSession { private readonly pendingPromptResults = new Map(); private readonly pendingSteerSubmissions: PiPendingSteerSubmission[] = []; currentLeafOverrideId: string | null | undefined; - private readonly capturedUserEntries: PiCapturedEntry[] = []; - private readonly capturedUserEntriesById = new Map(); + private readonly contextUserEntries: PiCapturedEntry[] = []; + private readonly treeUserEntriesById = new Map(); private readonly pendingExtensionResults = new Map(); private outOfBandCompactionEmit: ((event: AgentStreamEvent) => void) | null = null; private outOfBandCompactionStarted = false; @@ -1475,7 +1481,7 @@ export class PiRpcAgentSession implements AgentSession { yield* streamPiHistory( this.provider, await this.runtimeSession.getMessages(), - this.capturedUserEntries, + this.contextUserEntries, ); } @@ -1619,7 +1625,7 @@ export class PiRpcAgentSession implements AgentSession { } await this.refreshState().catch(() => undefined); await this.requestEntryCapture("rewind"); - const targetEntry = this.capturedUserEntriesById.get(input.messageId); + const targetEntry = this.treeUserEntriesById.get(input.messageId); if (!targetEntry) { throw new Error(`Pi rewind target ${input.messageId} was not found in captured tree entries`); } @@ -1965,11 +1971,14 @@ export class PiRpcAgentSession implements AgentSession { } } - private recordCapturedUserEntries(entries: PiCapturedEntry[]): void { - this.capturedUserEntries.splice(0, this.capturedUserEntries.length, ...entries); - this.capturedUserEntriesById.clear(); - for (const entry of entries) { - this.capturedUserEntriesById.set(entry.id, entry); + private recordCapturedUserEntries(input: { + treeEntries: PiCapturedEntry[]; + contextEntries: PiCapturedEntry[]; + }): void { + this.contextUserEntries.splice(0, this.contextUserEntries.length, ...input.contextEntries); + this.treeUserEntriesById.clear(); + for (const entry of input.treeEntries) { + this.treeUserEntriesById.set(entry.id, entry); } } @@ -2005,10 +2014,12 @@ export class PiRpcAgentSession implements AgentSession { if (!payload) { return false; } - const entries = parseCapturedEntries(payload.entries); - this.recordCapturedUserEntries(entries); + this.recordCapturedUserEntries({ + treeEntries: parseCapturedEntries(payload.treeEntries), + contextEntries: parseCapturedEntries(payload.contextEntries), + }); if (typeof payload.requestId === "string") { - this.resolveExtensionResult(payload.requestId, entries); + this.resolveExtensionResult(payload.requestId, undefined); } return true; } 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 ce20b4d8ee..50a0bf9532 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 @@ -108,7 +108,10 @@ export class FakePiSession implements PiRuntimeSession { readonly handoffRequests: Array<{ customInstructions?: string }> = []; readonly sessionNameRequests: string[] = []; readonly rawFrames: Array = []; - capturedUserEntries: Array<{ id: string; parentId: string | null; text: 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[] = []; abortRequested = false; readonly canceledExtensionUiRequests: string[] = []; readonly extensionUiResponses: Array<{ @@ -454,7 +457,8 @@ export class FakePiSession implements PiRuntimeSession { message: `PASEO_ENTRY_CAPTURE ${JSON.stringify({ reason, requestId, - entries: this.capturedUserEntries, + treeEntries: this.treeUserEntries, + contextEntries: this.contextUserEntries, })}`, }); } 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 51efc705cf..fb25eaef01 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 @@ -63,14 +63,31 @@ function expectPiSessionId(value: string): void { expect(value.length).toBeGreaterThan(0); } -function piPrompt(input: { promptToken: string; doneToken: string }): string { +interface PiTurn { + promptToken: string; + doneToken: string; + filler?: string; +} + +function piPrompt(input: PiTurn): string { return [ `PASEO_PI_REWIND_PROMPT_${input.promptToken}.`, "Remember this marker for the conversation.", + ...(input.filler ? [`Ignore this filler: ${input.filler}`] : []), `Reply exactly: ${input.doneToken}`, ].join(" "); } +function turn(token: string): PiTurn { + return { promptToken: token, doneToken: `PI_${token}_DONE` }; +} + +// Pi keeps the most recent 20k estimated tokens (4 characters each) when it compacts. +// A turn larger than that makes /compact summarize every turn before it. +function turnOutlastingCompaction(token: string): PiTurn { + return { ...turn(token), filler: "lorem ipsum ".repeat(8_000) }; +} + function roleItems(items: AgentTimelineItem[], role: "user_message" | "assistant_message") { return items.filter((item) => item.type === role); } @@ -91,10 +108,29 @@ async function expectNoCreatedFiles(session: PiRewindSession): Promise { await expect(readdir(session.cwd)).resolves.toEqual([]); } +async function compactPi(harness: PiRewindHarness, session: PiRewindSession): Promise { + await harness.client.sendMessage(session.agentId, "/compact"); + const deadline = Date.now() + TURN_TIMEOUT_MS; + while (Date.now() < deadline) { + const items = await fetchTimelineItems(harness.client, session.agentId); + const failure = items.find( + (item) => item.type === "assistant_message" && item.text.includes("Failed to compact"), + ); + if (failure) { + throw new Error(`Pi compaction failed: ${JSON.stringify(failure)}`); + } + if (items.some((item) => item.type === "compaction" && item.status === "completed")) { + return; + } + await new Promise((resolve) => setTimeout(resolve, 500)); + } + throw new Error("Pi compaction did not complete"); +} + async function askPi( harness: PiRewindHarness, session: PiRewindSession, - input: { promptToken: string; doneToken: string }, + input: PiTurn, ): Promise { await harness.client.sendMessage(session.agentId, piPrompt(input)); const finish = await harness.client.waitForFinish(session.agentId, TURN_TIMEOUT_MS); @@ -204,4 +240,68 @@ describe("daemon E2E (real pi) - rewind", () => { await closePiRewindSession(session); } }, 420_000); + + test("rewinds a replayed row to the selected message after an earlier rewind branched the session", async () => { + const session = await launchPiRewindSession(harness, "pi-rewind-branched-replay-real"); + + try { + await askPi(harness, session, turn("ONE")); + await askPi(harness, session, turn("ABANDONED")); + const abandonedId = userMessageIdForToken( + await fetchTimelineItems(harness.client, session.agentId), + "PI_REWIND_PROMPT_ABANDONED", + ); + await harness.client.rewindAgent(session.agentId, abandonedId, "conversation"); + + await askPi(harness, session, turn("TWO")); + await askPi(harness, session, turn("THREE")); + await askPi(harness, session, turn("FOUR")); + const fourId = userMessageIdForToken( + await fetchTimelineItems(harness.client, session.agentId), + "PI_REWIND_PROMPT_FOUR", + ); + await harness.client.rewindAgent(session.agentId, fourId, "conversation"); + + const replayedTimeline = await fetchTimelineItems(harness.client, session.agentId); + expectTimeline(replayedTimeline, { + userTexts: [piPrompt(turn("ONE")), piPrompt(turn("TWO")), piPrompt(turn("THREE"))], + assistantCount: 3, + }); + const replayedThreeId = userMessageIdForToken(replayedTimeline, "PI_REWIND_PROMPT_THREE"); + await harness.client.rewindAgent(session.agentId, replayedThreeId, "conversation"); + + expectTimeline(await fetchTimelineItems(harness.client, session.agentId), { + userTexts: [piPrompt(turn("ONE")), piPrompt(turn("TWO"))], + assistantCount: 2, + }); + await expectNoCreatedFiles(session); + } finally { + await closePiRewindSession(session); + } + }, 900_000); + + test("rewinds a row whose entry an earlier compaction summarized away", async () => { + const session = await launchPiRewindSession(harness, "pi-rewind-compacted-real"); + + try { + await askPi(harness, session, turn("ONE")); + await askPi(harness, session, turn("TWO")); + await askPi(harness, session, turnOutlastingCompaction("THREE")); + const twoId = userMessageIdForToken( + await fetchTimelineItems(harness.client, session.agentId), + "PI_REWIND_PROMPT_TWO", + ); + await compactPi(harness, session); + + await harness.client.rewindAgent(session.agentId, twoId, "conversation"); + + expectTimeline(await fetchTimelineItems(harness.client, session.agentId), { + userTexts: [piPrompt(turn("ONE"))], + assistantCount: 1, + }); + await expectNoCreatedFiles(session); + } finally { + await closePiRewindSession(session); + } + }, 900_000); });