From 6adcf7e5b0bc8601761af59b461a746cb5e395fa Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 1 Aug 2026 12:45:14 +0200 Subject: [PATCH] fix(server): preserve prompt identity across hydration --- .../src/server/agent/agent-manager.test.ts | 58 +++++++++++++++++++ .../server/src/server/agent/agent-manager.ts | 52 ++++++++++++----- .../server/src/server/agent/agent-storage.ts | 4 ++ 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 2a3ec15d45..dda92aab7a 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/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 { this.rewindMessageIds.push(messageId); } + + override async *streamHistory(): AsyncGenerator { + 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, + ): Promise { + 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 }); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index a0b8f78003..d5c55d077b 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/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, +): { 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(); private readonly providerEnabled = new Map(); @@ -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 = {}; + 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 { + const identities = (await this.registry?.get(agent.id))?.submittedPromptIdentities ?? {}; const historyEvents: Extract[] = []; const providerSubagentEvents: Extract[] = []; 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 { + const identities = (await this.registry?.get(agent.id))?.submittedPromptIdentities ?? {}; const deferredBroadcast = typeof broadcast === "function"; const timelineEvents: Array<{ event: Extract; @@ -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; } diff --git a/packages/server/src/server/agent/agent-storage.ts b/packages/server/src/server/agent/agent-storage.ts index 5e504cee64..161d4b2713 100644 --- a/packages/server/src/server/agent/agent-storage.ts +++ b/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, ): Promise { 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