fix(server): preserve Pi message IDs after resume (#2313)

Author: boudra

packages/server/src/server/agent/providers/pi/agent.test.ts

@@ -57,21 +57,30 @@ function readUtf8File(pathname: string): string {
   }
 }
 
-async function applyPaseoExtensionSystemPrompt(
+type PaseoExtensionListener = (event: unknown, context?: unknown) => unknown;
+
+async function loadPaseoExtensionListeners(
   extensionPath: string,
-  systemPrompt: string,
-): Promise<string | undefined> {
-  const listeners = new Map<string, (event: { systemPrompt: string }) => unknown>();
+): Promise<Map<string, PaseoExtensionListener>> {
+  const listeners = new Map<string, PaseoExtensionListener>();
   const extension = (await import(pathToFileURL(extensionPath).href)) as {
     default: (piApi: {
-      on: (event: string, listener: (event: { systemPrompt: string }) => unknown) => void;
+      on: (event: string, listener: PaseoExtensionListener) => void;
       registerCommand: () => void;
     }) => void;
   };
   extension.default({
     on: (event, listener) => listeners.set(event, listener),
     registerCommand: () => undefined,
   });
+  return listeners;
+}
+
+async function applyPaseoExtensionSystemPrompt(
+  extensionPath: string,
+  systemPrompt: string,
+): Promise<string | undefined> {
+  const listeners = await loadPaseoExtensionListeners(extensionPath);
   const result = await listeners.get("before_agent_start")?.({ systemPrompt });
   return (result as { systemPrompt?: string } | undefined)?.systemPrompt;
 }
@@ -594,15 +603,15 @@ describe("PiRpcAgentSession", () => {
     ]);
   });
 
-  test("emits live user messages with captured Pi tree entry ids", async () => {
+  test("emits live user messages with submitted Pi tree entry ids", async () => {
     const { pi, session, events } = await createSession();
     const fakeSession = pi.latestSession();
 
-    fakeSession.capturedUserEntries = [{ id: "entry-user-1", parentId: null, text: "hello" }];
     await session.startTurn("hello");
-    fakeSession.emit({
-      type: "message_end",
-      message: { role: "user", content: "hello" },
+    fakeSession.finishSubmittedUserMessage({
+      id: "entry-user-1",
+      parentId: null,
+      text: "hello",
     });
 
     await events.nextTimelineEvent();
@@ -612,6 +621,38 @@ describe("PiRpcAgentSession", () => {
     ]);
   });
 
+  test("uses the Pi entry attached to a submitted prompt after resuming old history", async () => {
+    const pi = new FakePi();
+    const client = createClient(pi);
+    const session = (await client.resumeSession({
+      provider: "pi",
+      sessionId: "pi-session-1",
+      nativeHandle: "/tmp/native-pi-session",
+      metadata: { cwd: "/workspace/project" },
+    })) as PiRpcAgentSession;
+    const events = new SessionEvents(session);
+    const fakeSession = pi.latestSession();
+    fakeSession.capturedUserEntries = [{ id: "entry-old", parentId: null, text: "old prompt" }];
+
+    await session.startTurn("new prompt", { clientMessageId: "client-new" });
+    fakeSession.finishSubmittedUserMessage({
+      id: "entry-new",
+      parentId: "entry-old-assistant",
+      text: "new prompt",
+    });
+
+    await events.nextTimelineEvent();
+
+    expect(events.timelineItems()).toEqual([
+      {
+        type: "user_message",
+        text: "new prompt",
+        messageId: "entry-new",
+        clientMessageId: "client-new",
+      },
+    ]);
+  });
+
   test("surfaces Pi extension command messages and completes when no agent turn starts", async () => {
     const { pi, session, events } = await createSession();
     const fakeSession = pi.latestSession();
@@ -734,6 +775,52 @@ describe("PiRpcAgentSession", () => {
     ]);
   });
 
+  test("reports the persisted Pi entry attached to the submitted message", async () => {
+    const pi = new FakePi();
+    const client = createClient(pi);
+    const session = await client.createSession(createConfig());
+    const extensionPath = pi.recordedLaunches[0]?.extensionPaths[0];
+    expect(extensionPath).toBeDefined();
+    const listeners = await loadPaseoExtensionListeners(extensionPath!);
+    const submittedMessage = { role: "user", content: "new prompt" };
+    const entries: Array<{
+      type: string;
+      id: string;
+      parentId: string | null;
+      message: { role: string; content: string };
+    }> = [
+      {
+        type: "message",
+        id: "entry-old",
+        parentId: null,
+        message: { role: "user", content: "old prompt" },
+      },
+    ];
+    const notifications: string[] = [];
+    const context = {
+      sessionManager: { getEntries: () => entries },
+      ui: { notify: (message: string) => notifications.push(message) },
+    };
+
+    await listeners.get("message_end")?.({ message: submittedMessage }, context);
+    entries.push({
+      type: "message",
+      id: "entry-new",
+      parentId: "entry-old-assistant",
+      message: submittedMessage,
+    });
+    await listeners.get("message_start")?.(
+      { message: { role: "assistant", content: [] } },
+      context,
+    );
+
+    expect(notifications).toEqual([
+      'PASEO_SUBMITTED_USER_ENTRY {"entry":{"id":"entry-new","parentId":"entry-old-assistant","text":"new prompt"}}',
+    ]);
+
+    await session.close();
+  });
+
   test("appends agent and daemon prompts after Pi's discovered system prompt", async () => {
     const pi = new FakePi();
     const client = createClient(pi);

packages/server/src/server/agent/providers/pi/agent.ts

@@ -92,6 +92,7 @@ const PI_CATALOG_REQUEST_TIMEOUT_MS = 120_000;
 const PASEO_PI_TREE_EXTENSION_COMMAND = "paseo_tree";
 const PASEO_PI_CAPTURE_EXTENSION_COMMAND = "paseo_capture_entries";
 const PASEO_PI_ENTRY_CAPTURE_MARKER = "PASEO_ENTRY_CAPTURE";
+const PASEO_PI_SUBMITTED_USER_ENTRY_MARKER = "PASEO_SUBMITTED_USER_ENTRY";
 const PASEO_PI_COMMAND_RESULT_MARKER = "PASEO_COMMAND_RESULT";
 const DEFAULT_PI_EXTENSION_RESULT_TIMEOUT_MS = 30_000;
 const QUESTION_RESPONSE_HEADER = "Response";
@@ -256,12 +257,6 @@ interface PiCapturedEntry extends PiCapturedUserMessageEntry {
   parentId: string | null;
 }
 
-interface PendingPiUserMessage {
-  text: string;
-  turnId: string | undefined;
-  clientMessageId?: string;
-}
-
 interface PendingExtensionResult {
   resolve: (value: unknown) => void;
   reject: (error: Error) => void;
@@ -653,11 +648,15 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile {
 	  return ctx.sessionManager
 	    .getEntries()
 	    .filter((entry) => entry.type === "message" && entry.message?.role === "user")
-	    .map((entry) => ({
+	    .map(toCapturedUserEntry);
+	}
+
+	function toCapturedUserEntry(entry) {
+	  return {
 	      id: entry.id,
 	      parentId: entry.parentId ?? null,
 	      text: readTextContent(entry.message.content),
-	    }));
+	    };
 	}
 
 	function emitEntryCapture(ctx, reason, requestId) {
@@ -676,6 +675,30 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile {
 	}
 
 	export default function paseoIntegration(pi) {
+	  const submittedUserMessages = [];
+
+	  function emitSubmittedUserEntries(ctx) {
+	    const entries = ctx.sessionManager.getEntries();
+	    for (let index = 0; index < submittedUserMessages.length; index += 1) {
+	      const message = submittedUserMessages[index];
+	      // Pi assigns the entry ID after message_end, then persists this same message object.
+	      // Reference equality preserves the exact association even when another extension edits it.
+	      const entry = entries.find(
+	        (candidate) => candidate.type === "message" && candidate.message === message,
+	      );
+	      if (!entry) {
+	        continue;
+	      }
+	      submittedUserMessages.splice(index, 1);
+	      index -= 1;
+	      ctx.ui.notify(
+	        "${PASEO_PI_SUBMITTED_USER_ENTRY_MARKER} " +
+	          JSON.stringify({ entry: toCapturedUserEntry(entry) }),
+	        "info",
+	      );
+	    }
+	  }
+
 	  ${
       systemPrompt
         ? `pi.on("before_agent_start", async (event) => ({
@@ -688,7 +711,20 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile {
 	    emitEntryCapture(ctx, "session_start");
 	  });
 
+	  pi.on("message_end", async (event) => {
+	    if (event.message?.role === "user") {
+	      submittedUserMessages.push(event.message);
+	    }
+	  });
+
+	  pi.on("message_start", async (event, ctx) => {
+	    if (event.message?.role === "assistant") {
+	      emitSubmittedUserEntries(ctx);
+	    }
+	  });
+
 	  pi.on("turn_end", async (_event, ctx) => {
+	    emitSubmittedUserEntries(ctx);
 	    emitEntryCapture(ctx, "turn_end");
 	  });
 
@@ -1197,8 +1233,6 @@ export class PiRpcAgentSession implements AgentSession {
   currentLeafOverrideId: string | null | undefined;
   private readonly capturedUserEntries: PiCapturedEntry[] = [];
   private readonly capturedUserEntriesById = new Map<string, PiCapturedEntry>();
-  private readonly seenUserEntryIds = new Set<string>();
-  private readonly pendingUserMessages: PendingPiUserMessage[] = [];
   private readonly pendingExtensionResults = new Map<string, PendingExtensionResult>();
   private outOfBandCompactionEmit: ((event: AgentStreamEvent) => void) | null = null;
   private outOfBandCompactionStarted = false;
@@ -1776,42 +1810,34 @@ export class PiRpcAgentSession implements AgentSession {
   }
 
   private recordCapturedUserEntries(entries: PiCapturedEntry[]): void {
-    const previouslySeenEntryIds = new Set(this.seenUserEntryIds);
     this.capturedUserEntries.splice(0, this.capturedUserEntries.length, ...entries);
     this.capturedUserEntriesById.clear();
     for (const entry of entries) {
       this.capturedUserEntriesById.set(entry.id, entry);
     }
-    this.flushPendingUserMessages(previouslySeenEntryIds);
-    for (const entry of entries) {
-      this.seenUserEntryIds.add(entry.id);
-    }
   }
 
-  private flushPendingUserMessages(previouslySeenEntryIds: Set<string>): void {
-    for (let index = 0; index < this.pendingUserMessages.length; index += 1) {
-      const pending = this.pendingUserMessages[index]!;
-      const entry = this.capturedUserEntries.find(
-        (candidate) => !previouslySeenEntryIds.has(candidate.id),
-      );
-      if (!entry) {
-        continue;
-      }
-      previouslySeenEntryIds.add(entry.id);
-      this.pendingUserMessages.splice(index, 1);
-      index -= 1;
-      this.emit({
-        type: "timeline",
-        provider: this.provider,
-        turnId: pending.turnId,
-        item: {
-          type: "user_message",
-          text: pending.text,
-          messageId: entry.id,
-          ...(pending.clientMessageId ? { clientMessageId: pending.clientMessageId } : {}),
-        },
-      });
+  private handleSubmittedUserEntryMarker(message: string): boolean {
+    const payload = parseExtensionMarkerPayload(message, PASEO_PI_SUBMITTED_USER_ENTRY_MARKER);
+    if (!payload) {
+      return false;
+    }
+    const [entry] = parseCapturedEntries([payload.entry]);
+    if (!entry) {
+      return true;
     }
+    this.emit({
+      type: "timeline",
+      provider: this.provider,
+      turnId: this.currentTurnIdForEvent(),
+      item: {
+        type: "user_message",
+        text: entry.text,
+        messageId: entry.id,
+        ...(this.activeClientMessageId ? { clientMessageId: this.activeClientMessageId } : {}),
+      },
+    });
+    return true;
   }
 
   private handleEntryCaptureMarker(message: string): boolean {
@@ -1849,7 +1875,11 @@ export class PiRpcAgentSession implements AgentSession {
   ): void {
     const message = optionalString(event.message);
     if (event.method === "notify" && message) {
-      if (this.handleEntryCaptureMarker(message) || this.handleCommandResultMarker(message)) {
+      if (
+        this.handleSubmittedUserEntryMarker(message) ||
+        this.handleEntryCaptureMarker(message) ||
+        this.handleCommandResultMarker(message)
+      ) {
         return;
       }
       this.bufferNoTurnOutput(message);
@@ -2176,28 +2206,6 @@ export class PiRpcAgentSession implements AgentSession {
       this.completeTurn(turnId, []);
       return;
     }
-
-    if (event.message.role !== "user") {
-      return;
-    }
-    const text = getUserMessageText(event.message.content);
-    if (!text) {
-      return;
-    }
-    this.pendingUserMessages.push({
-      text,
-      turnId,
-      ...(this.activeClientMessageId ? { clientMessageId: this.activeClientMessageId } : {}),
-    });
-    void this.requestEntryCapture("message_end").catch((error: unknown) => {
-      const message = error instanceof Error ? error.message : String(error);
-      this.emit({
-        type: "turn_failed",
-        provider: this.provider,
-        turnId,
-        error: message,
-      });
-    });
   }
 
   private emitToolCallEvent(

packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts

@@ -45,6 +45,12 @@ export interface FakePiSubagentMessagesResult {
   messages: PiAgentMessage[];
 }
 
+interface FakePiUserEntry {
+  id: string;
+  parentId: string | null;
+  text: string;
+}
+
 export class FakePi implements PiRuntime {
   readonly recordedLaunches: PiRuntimeLaunch[] = [];
   private readonly sessions: FakePiSession[] = [];
@@ -349,6 +355,19 @@ export class FakePiSession implements PiRuntimeSession {
     this.emit({ type: "agent_end", messages: this.messages });
   }
 
+  finishSubmittedUserMessage(entry: FakePiUserEntry): void {
+    this.emit({
+      type: "message_end",
+      message: { role: "user", content: entry.text },
+    });
+    this.emit({
+      type: "extension_ui_request",
+      id: `submitted-user-${entry.id}`,
+      method: "notify",
+      message: `PASEO_SUBMITTED_USER_ENTRY ${JSON.stringify({ entry })}`,
+    });
+  }
+
   private handleTreeNavigationCommand(message: string): void {
     const prefix = "/paseo_tree ";
     if (!message.startsWith(prefix)) {