diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index c107b3da0b..af8c6f9589 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -692,6 +692,137 @@ test("uses an injected timeline store without making it a production requirement } }); +test("refreshing an agent replaces the injected timeline store instead of appending a second copy", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-timeline-rehydrate-")); + const store = new RecordingTimelineStore(); + class HistorySession extends TestAgentSession { + override async *streamHistory(): AsyncGenerator { + yield { + type: "timeline", + provider: "codex", + item: { type: "user_message", text: "continue" }, + }; + yield { + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "done" }, + }; + } + } + class HistoryClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + return new HistorySession(config); + } + + override async resumeSession( + _handle: AgentPersistenceHandle, + config?: Partial, + ): Promise { + return new HistorySession({ provider: "codex", cwd: config?.cwd ?? workdir }); + } + } + const manager = new AgentManager({ + clients: { codex: new HistoryClient() }, + durableTimelineStore: store, + logger, + }); + let agentId: string | null = null; + try { + const agent = await manager.createAgent({ provider: "codex", cwd: workdir }, undefined, { + workspaceId: undefined, + }); + agentId = agent.id; + + // What session.handleRefreshAgentRequest does for a loaded agent, twice over. + for (let refresh = 0; refresh < 2; refresh += 1) { + await manager.reloadAgentSession(agent.id, undefined, { rehydrateFromDisk: true }); + await manager.hydrateTimelineFromProvider(agent.id, { broadcast: true }); + await manager.flush(); + } + + const rows = await manager.getTimelineRows(agent.id); + expect(rows.map((row) => row.item)).toEqual([ + { type: "user_message", text: "continue" }, + { type: "assistant_message", text: "done" }, + ]); + } finally { + if (agentId) await manager.closeAgent(agentId).catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + +test("a failed history replay leaves the committed timeline intact", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-timeline-replay-fail-")); + const store = new RecordingTimelineStore(); + let failReplay = false; + class FlakyHistorySession extends TestAgentSession { + override async *streamHistory(): AsyncGenerator { + yield { + type: "timeline", + provider: "codex", + item: { type: "user_message", text: "continue" }, + }; + if (failReplay) { + throw new Error("history stream closed"); + } + yield { + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "done" }, + }; + } + } + class FlakyHistoryClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + return new FlakyHistorySession(config); + } + + override async resumeSession( + _handle: AgentPersistenceHandle, + config?: Partial, + ): Promise { + return new FlakyHistorySession({ provider: "codex", cwd: config?.cwd ?? workdir }); + } + } + const manager = new AgentManager({ + clients: { codex: new FlakyHistoryClient() }, + durableTimelineStore: store, + logger, + }); + let agentId: string | null = null; + try { + const agent = await manager.createAgent({ provider: "codex", cwd: workdir }, undefined, { + workspaceId: undefined, + }); + agentId = agent.id; + await manager.reloadAgentSession(agent.id, undefined, { rehydrateFromDisk: true }); + await manager.hydrateTimelineFromProvider(agent.id, { broadcast: true }); + await manager.flush(); + const committed = await manager.getTimelineRows(agent.id); + + failReplay = true; + await manager.reloadAgentSession(agent.id, undefined, { rehydrateFromDisk: true }); + await expect( + manager.hydrateTimelineFromProvider(agent.id, { broadcast: true }), + ).rejects.toThrow("history stream closed"); + await manager.flush(); + + expect(await manager.getTimelineRows(agent.id)).toEqual(committed); + + // A retry once the provider recovers replaces that history rather than stacking on it. + failReplay = false; + await manager.hydrateTimelineFromProvider(agent.id, { broadcast: true }); + await manager.flush(); + expect((await manager.getTimelineRows(agent.id)).map((row) => row.item)).toEqual([ + { type: "user_message", text: "continue" }, + { type: "assistant_message", text: "done" }, + ]); + } finally { + if (agentId) await manager.closeAgent(agentId).catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + test("retries provider history hydration after a stream failure", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-history-retry-")); let attempts = 0; diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index c80d257a3e..2fc062d484 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -4025,23 +4025,17 @@ export class AgentManager { broadcast: boolean | (() => boolean), ): Promise { const deferredBroadcast = typeof broadcast === "function"; - const timelineEvents: Array<{ - event: Extract; - row: AgentTimelineRow; - }> = []; - const providerSubagentEvents: AgentManagerEvent[] = []; + const historyEvents: Extract[] = []; + const historySubagentEvents: Extract[] = []; agent.historyPrimed = false; try { + // Collect the whole replay before touching either store. A stream that fails + // halfway then leaves the committed timeline as it was, instead of a partial + // copy the next attempt would append to. for await (const rawEvent of agent.session.streamHistory()) { const event = limitAgentStreamEventContent(rawEvent); if (event.type === "provider_subagent") { - const update = this.providerSubagents.apply(agent.id, event.provider, event.event); - const managerEvent: AgentManagerEvent = { type: "provider_subagent", event: update }; - if (deferredBroadcast) { - providerSubagentEvents.push(managerEvent); - } else if (broadcast) { - this.dispatch(managerEvent); - } + historySubagentEvents.push(event); continue; } if (event.type !== "timeline") { @@ -4050,25 +4044,47 @@ 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, - ); - if (deferredBroadcast) { - timelineEvents.push({ event, row }); - } else if (broadcast) { - this.dispatchStream(agent.id, event, { - seq: row.seq, - epoch: this.timelineStore.getEpoch(agent.id), - timestamp: row.timestamp, - }); - } + historyEvents.push(event); } } catch (error) { this.logger.warn({ err: error, agentId: agent.id }, "Failed to hydrate provider history"); throw error; } + + // The replay is the timeline, so drop the rows a previous hydration committed. + // Keeping them would leave getTimelineRows reading one copy per hydration. + await this.deleteCommittedTimeline(agent.id); + + const timelineEvents: Array<{ + event: Extract; + row: AgentTimelineRow; + }> = []; + const providerSubagentEvents: AgentManagerEvent[] = []; + for (const event of historySubagentEvents) { + const update = this.providerSubagents.apply(agent.id, event.provider, event.event); + const managerEvent: AgentManagerEvent = { type: "provider_subagent", event: update }; + if (deferredBroadcast) { + providerSubagentEvents.push(managerEvent); + } else if (broadcast) { + this.dispatch(managerEvent); + } + } + for (const event of historyEvents) { + const row = this.recordTimeline( + agent.id, + event.item, + event.timestamp ? { timestamp: event.timestamp } : undefined, + ); + if (deferredBroadcast) { + timelineEvents.push({ event, row }); + } else if (broadcast) { + this.dispatchStream(agent.id, event, { + seq: row.seq, + epoch: this.timelineStore.getEpoch(agent.id), + timestamp: row.timestamp, + }); + } + } agent.historyPrimed = true; if (typeof broadcast !== "function" || !broadcast()) {