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); });