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>

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<Map<string, PaseoExtensionListener>> {
@@ -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 }));

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<string, boolean>();
   private readonly pendingSteerSubmissions: PiPendingSteerSubmission[] = [];
   currentLeafOverrideId: string | null | undefined;
-  private readonly capturedUserEntries: PiCapturedEntry[] = [];
-  private readonly capturedUserEntriesById = new Map<string, PiCapturedEntry>();
+  private readonly contextUserEntries: PiCapturedEntry[] = [];
+  private readonly treeUserEntriesById = new Map<string, PiCapturedEntry>();
   private readonly pendingExtensionResults = new Map<string, PendingExtensionResult>();
   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;
   }

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<object & { type: string }> = [];
-  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,
       })}`,
     });
   }

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<void> {
   await expect(readdir(session.cwd)).resolves.toEqual([]);
 }
 
+async function compactPi(harness: PiRewindHarness, session: PiRewindSession): Promise<void> {
+  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<void> {
   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);
 });