fix(server): preserve prompt identity across hydration

packages/server/src/server/agent/agent-manager.test.ts

@@ -8205,6 +8205,23 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
     override async revertFiles({ messageId }: { messageId: string }): Promise<void> {
       this.rewindMessageIds.push(messageId);
     }
+
+    override async *streamHistory(): AsyncGenerator<AgentStreamEvent> {
+      yield {
+        type: "timeline",
+        provider: this.provider,
+        item: {
+          type: "user_message",
+          text: "hello from composer",
+          messageId: "provider-message-1",
+        },
+      };
+      yield {
+        type: "timeline",
+        provider: this.provider,
+        item: { type: "assistant_message", text: "output before provider echo" },
+      };
+    }
   }
 
   class SubmittedUserMessageClient extends TestAgentClient {
@@ -8214,6 +8231,17 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
       this.session = new SubmittedUserMessageSession(config);
       return this.session;
     }
+
+    override async resumeSession(
+      _handle: AgentPersistenceHandle,
+      config?: Partial<AgentSessionConfig>,
+    ): Promise<AgentSession> {
+      this.session = new SubmittedUserMessageSession({
+        provider: "codex",
+        cwd: config?.cwd ?? workdir,
+      });
+      return this.session;
+    }
   }
 
   const client = new SubmittedUserMessageClient();
@@ -8223,6 +8251,8 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
     logger,
     idFactory: () => "00000000-0000-4000-8000-000000000402",
   });
+  let restartedManager: AgentManager | undefined;
+  let restartedStorage: AgentStorage | undefined;
   const events: AgentManagerEvent[] = [];
   manager.subscribe((event) => events.push(event), { replayState: false });
   const streamEvents = () =>
@@ -8302,7 +8332,35 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
       "provider-message-1",
       "provider-native-message",
     ]);
+
+    await manager.closeAgent(snapshot.id);
+    expect((await storage.get(snapshot.id))?.submittedPromptIdentities).toEqual({
+      "msg-client-1": "provider-message-1",
+    });
+    restartedStorage = new AgentStorage(storagePath, logger);
+    const restartedClient = new SubmittedUserMessageClient();
+    restartedManager = new AgentManager({
+      clients: { codex: restartedClient },
+      registry: restartedStorage,
+      logger,
+    });
+    await ensureAgentLoaded(snapshot.id, {
+      agentManager: restartedManager,
+      agentStorage: restartedStorage,
+      logger,
+    });
+
+    expect(
+      restartedManager.fetchTimeline(snapshot.id, { direction: "tail", limit: 20 }).rows[0],
+    ).toMatchObject({
+      providerMessageId: "provider-message-1",
+      item: { clientMessageId: "msg-client-1", messageId: "msg-client-1" },
+    });
+    await restartedManager.rewind(snapshot.id, "msg-client-1", "files");
+    expect(restartedClient.session?.rewindMessageIds).toEqual(["provider-message-1"]);
   } finally {
+    await restartedManager?.flush().catch(() => undefined);
+    await restartedStorage?.flush().catch(() => undefined);
     await manager.flush().catch(() => undefined);
     await storage.flush().catch(() => undefined);
     rmSync(workdir, { recursive: true, force: true });

packages/server/src/server/agent/agent-manager.ts

@@ -601,6 +601,17 @@ function getFirstUserMessageTextFromRows(rows: readonly AgentTimelineRow[]): str
   return null;
 }
 
+function restoreSubmittedPromptIdentity(
+  item: AgentTimelineItem,
+  identities: Record<string, string>,
+): { item: AgentTimelineItem; providerMessageId?: string } {
+  if (item.type !== "user_message" || !item.messageId) return { item };
+  const match = Object.entries(identities).find(([, providerId]) => providerId === item.messageId);
+  if (!match) return { item };
+  const [clientMessageId, providerMessageId] = match;
+  return { item: { ...item, clientMessageId, messageId: clientMessageId }, providerMessageId };
+}
+
 export class AgentManager {
   private readonly clients = new Map<AgentProvider, AgentClient>();
   private readonly providerEnabled = new Map<AgentProvider, boolean>();
@@ -3200,7 +3211,15 @@ export class AgentManager {
     if (agent.internal) {
       return;
     }
-    await this.registry.applySnapshot(agent, options);
+    const timelineRows = this.timelineStore.getRows(agent.id);
+    const submittedPromptIdentities: Record<string, string> = {};
+    for (const row of timelineRows) {
+      if (row.item.type === "user_message" && row.item.clientMessageId && row.providerMessageId) {
+        submittedPromptIdentities[row.item.clientMessageId] = row.providerMessageId;
+      }
+    }
+    const identities = timelineRows.length > 0 ? submittedPromptIdentities : undefined;
+    await this.registry.applySnapshot(agent, options, identities);
   }
 
   private requireRegistry(): AgentStorage {
@@ -3290,6 +3309,7 @@ export class AgentManager {
     agent: ActiveManagedAgent,
     broadcast: boolean,
   ): Promise<void> {
+    const identities = (await this.registry?.get(agent.id))?.submittedPromptIdentities ?? {};
     const historyEvents: Extract<AgentStreamEvent, { type: "timeline" }>[] = [];
     const providerSubagentEvents: Extract<AgentStreamEvent, { type: "provider_subagent" }>[] = [];
     for await (const event of agent.session.streamHistory()) {
@@ -3320,12 +3340,13 @@ export class AgentManager {
         this.dispatch({ type: "provider_subagent", event: update });
       }
     }
-    for (const event of historyEvents) {
-      const row = this.recordTimeline(
-        agent.id,
-        event.item,
-        event.timestamp ? { timestamp: event.timestamp } : undefined,
-      );
+    for (const historyEvent of historyEvents) {
+      const restored = restoreSubmittedPromptIdentity(historyEvent.item, identities);
+      const event = { ...historyEvent, item: restored.item };
+      const row = this.recordTimeline(agent.id, event.item, {
+        timestamp: event.timestamp,
+        providerMessageId: restored.providerMessageId,
+      });
       if (broadcast) {
         this.dispatchStream(agent.id, event, {
           seq: row.seq,
@@ -3342,6 +3363,7 @@ export class AgentManager {
     agent: ActiveManagedAgent,
     broadcast: boolean | (() => boolean),
   ): Promise<void> {
+    const identities = (await this.registry?.get(agent.id))?.submittedPromptIdentities ?? {};
     const deferredBroadcast = typeof broadcast === "function";
     const timelineEvents: Array<{
       event: Extract<AgentStreamEvent, { type: "timeline" }>;
@@ -3367,15 +3389,16 @@ export class AgentManager {
         if (event.item.type === "user_message" && isSystemInjectedEnvelope(event.item.text)) {
           continue;
         }
-        const row = this.recordTimeline(
-          agent.id,
-          event.item,
-          event.timestamp ? { timestamp: event.timestamp } : undefined,
-        );
+        const restored = restoreSubmittedPromptIdentity(event.item, identities);
+        const timelineEvent = { ...event, item: restored.item };
+        const row = this.recordTimeline(agent.id, timelineEvent.item, {
+          timestamp: event.timestamp,
+          providerMessageId: restored.providerMessageId,
+        });
         if (deferredBroadcast) {
-          timelineEvents.push({ event, row });
+          timelineEvents.push({ event: timelineEvent, row });
         } else if (broadcast) {
-          this.dispatchStream(agent.id, event, {
+          this.dispatchStream(agent.id, timelineEvent, {
             seq: row.seq,
             epoch: this.timelineStore.getEpoch(agent.id),
             timestamp: row.timestamp,
@@ -3977,6 +4000,7 @@ export class AgentManager {
         messageId,
       );
       if (enriched) this.enqueueDurableTimelineUpdate(agent.id, enriched);
+      if (enriched) this.enqueueBackgroundPersist(agent);
     }
     return existing;
   }

packages/server/src/server/agent/agent-storage.ts

@@ -66,6 +66,7 @@ const STORED_AGENT_SCHEMA = z.object({
   internal: z.boolean().optional(),
   archivedAt: z.string().nullable().optional(),
   owner: AgentOwnerSchema.optional(),
+  submittedPromptIdentities: z.record(z.string(), z.string()).optional(),
 });
 
 export type SerializableAgentConfig = Pick<
@@ -205,6 +206,7 @@ export class AgentStorage {
   async applySnapshot(
     agent: ManagedAgent,
     options?: { title?: string | null; internal?: boolean },
+    submittedPromptIdentities?: Record<string, string>,
   ): Promise<void> {
     await this.load();
     await this.waitForPendingWrite(agent.id);
@@ -218,6 +220,8 @@ export class AgentStorage {
       createdAt: existing?.createdAt,
       internal: hasInternalOverride ? options?.internal : (agent.internal ?? existing?.internal),
     });
+    record.submittedPromptIdentities =
+      submittedPromptIdentities ?? existing?.submittedPromptIdentities;
 
     // Preserve soft-delete/archive status across snapshot flushes.
     // `archivedAt` is not part of the ManagedAgent snapshot, so a naive projection