From 6027c38f8cddcbc9fd350305eefa54736c94933b Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Wed, 25 Mar 2026 12:00:35 +0700 Subject: [PATCH] Phase 4c: skip provider-history rehydrate when DB timeline is available --- .../src/server/agent/agent-manager.test.ts | 131 ++++++++++++++++++ .../server/src/server/agent/agent-manager.ts | 62 +++++++-- ...der-history-compatibility-boundary.test.ts | 2 +- ...ider-history-compatibility-service.test.ts | 53 +++---- .../provider-history-compatibility-service.ts | 33 +---- packages/server/src/server/session.ts | 3 +- 6 files changed, 208 insertions(+), 76 deletions(-) diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 23646351d4..5e8d981585 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -712,6 +712,137 @@ describe("AgentManager", () => { expect(afterReload?.config?.title).toBeUndefined(); }); + test("resumeAgentFromPersistence seeds live helpers from durable rows and skips provider replay", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-durable-seed-")); + const storagePath = join(workdir, "agents"); + const dataDir = join(workdir, "db"); + const storage = new AgentStorage(storagePath, logger); + const database = await openPaseoDatabase(dataDir); + let historyReplayCount = 0; + let manager: AgentManager | null = null; + + class HistoryReplayProbeSession extends TestAgentSession { + async *streamHistory(): AsyncGenerator { + historyReplayCount += 1; + yield { + type: "timeline", + provider: this.provider, + item: { type: "assistant_message", text: "provider history replay" }, + }; + } + } + + class HistoryReplayProbeClient implements AgentClient { + readonly provider = "codex" as const; + readonly capabilities = TEST_CAPABILITIES; + + async isAvailable(): Promise { + return true; + } + + async createSession(config: AgentSessionConfig): Promise { + return new HistoryReplayProbeSession(config); + } + + async resumeSession( + handle: AgentPersistenceHandle, + overrides?: Partial, + ): Promise { + const metadata = (handle.metadata ?? {}) as Partial; + return new HistoryReplayProbeSession({ + ...metadata, + ...overrides, + provider: "codex", + cwd: overrides?.cwd ?? metadata.cwd ?? process.cwd(), + }); + } + } + + try { + const durableTimelineStore = new DbAgentTimelineStore(database.db); + manager = new AgentManager({ + clients: { + codex: new HistoryReplayProbeClient(), + }, + registry: storage, + durableTimelineStore, + logger, + idFactory: () => "00000000-0000-4000-8000-000000000128", + }); + + const snapshot = await manager.createAgent({ + provider: "codex", + cwd: workdir, + }); + + await manager.appendTimelineItem(snapshot.id, { + type: "assistant_message", + text: "durable only", + }); + await manager.flush(); + + const handle = manager.getAgent(snapshot.id)?.persistence; + expect(handle).not.toBeNull(); + if (!handle) { + throw new Error("Expected persistence handle to be available"); + } + + await manager.closeAgent(snapshot.id); + + await expect(durableTimelineStore.getCommittedRows(snapshot.id)).resolves.toEqual([ + { + seq: 1, + timestamp: expect.any(String), + item: { + type: "assistant_message", + text: "durable only", + }, + }, + ]); + + const resumed = await manager.resumeAgentFromPersistence(handle, undefined, snapshot.id); + + expect(resumed.id).toBe(snapshot.id); + expect(manager.getTimeline(snapshot.id)).toEqual([ + { + type: "assistant_message", + text: "durable only", + }, + ]); + expect(manager.getLastAssistantMessage(snapshot.id)).toBe("durable only"); + await expect(manager.getTimelineRows(snapshot.id)).resolves.toEqual([ + { + seq: 1, + timestamp: expect.any(String), + item: { + type: "assistant_message", + text: "durable only", + }, + }, + ]); + + await manager.hydrateTimelineFromProvider(snapshot.id); + + expect(historyReplayCount).toBe(0); + expect(manager.getTimeline(snapshot.id)).toEqual([ + { + type: "assistant_message", + text: "durable only", + }, + ]); + + await manager.closeAgent(snapshot.id); + await manager.deleteCommittedTimeline(snapshot.id); + + await expect(durableTimelineStore.getCommittedRows(snapshot.id)).resolves.toEqual([]); + } finally { + await manager?.flush().catch(() => undefined); + await storage.flush().catch(() => undefined); + await database.close(); + rmSync(workdir, { recursive: true, force: true }); + } + }); + test("setTitle bumps updatedAt and persists title in the same snapshot write", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-set-title-updated-at-")); const storagePath = join(workdir, "agents"); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 2c2ec83e8d..050bf46920 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -525,8 +525,10 @@ export class AgentManager { }); } - // Reconstruct an agent from provider persistence. Callers should explicitly - // hydrate timeline history after resume. + // Reconstruct an agent from provider persistence. When a durable timeline + // store is configured, committed history is seeded from the durable store. + // Tests without a durable timeline store can still call + // hydrateTimelineFromProvider() for backward compatibility. async resumeAgentFromPersistence( handle: AgentPersistenceHandle, overrides?: Partial, @@ -660,7 +662,6 @@ export class AgentManager { }; await session.close(); this.timelineStore.delete(agentId); - this.enqueueDurableTimelineDelete(agentId); this.emitState(closedAgent); this.logger.trace({ agentId }, "closeAgent: completed"); } @@ -1398,9 +1399,19 @@ export class AgentManager { async hydrateTimelineFromProvider(agentId: string): Promise { const agent = this.requireAgent(agentId); + if (this.durableTimelineStore) { + return; + } await this.hydrateTimeline(agent); } + async deleteCommittedTimeline(agentId: string): Promise { + if (!this.durableTimelineStore) { + return; + } + await this.durableTimelineStore.deleteAgent(agentId); + } + getLastAssistantMessage(agentId: string): string | null { const agent = this.agents.get(agentId); if (!agent) { @@ -1606,7 +1617,7 @@ export class AgentManager { const initialPersistedTitle = await this.resolveInitialPersistedTitle(resolvedAgentId, config); const now = new Date(); - const timelineSeed: SeedAgentTimelineOptions | null = + const explicitTimelineSeed: SeedAgentTimelineOptions | null = options?.timeline?.length || options?.timelineRows?.length || options?.timelineNextSeq !== undefined @@ -1617,6 +1628,14 @@ export class AgentManager { timestamp: (options?.updatedAt ?? options?.createdAt ?? now).toISOString(), } : null; + const shouldSeedFromDurable = + !explicitTimelineSeed && + !this.timelineStore.has(resolvedAgentId) && + this.durableTimelineStore !== undefined; + const durableTimelineSeed = shouldSeedFromDurable + ? await this.loadCommittedTimelineSeed(resolvedAgentId, now) + : null; + const timelineSeed = explicitTimelineSeed ?? durableTimelineSeed; if (timelineSeed || !this.timelineStore.has(resolvedAgentId)) { this.timelineStore.initialize(resolvedAgentId, timelineSeed ?? { timestamp: now.toISOString() }); } @@ -1644,7 +1663,7 @@ export class AgentManager { unsubscribeSession: null, provisionalAssistantText: options?.provisionalAssistantText ?? null, persistence: attachPersistenceCwd(session.describePersistence(), config.cwd), - historyPrimed: options?.historyPrimed ?? false, + historyPrimed: options?.historyPrimed ?? shouldSeedFromDurable, lastUserMessageAt: options?.lastUserMessageAt ?? null, lastUsage: options?.lastUsage, lastError: options?.lastError, @@ -1679,6 +1698,29 @@ export class AgentManager { return { ...managed }; } + private async loadCommittedTimelineSeed( + agentId: string, + now: Date, + ): Promise { + if (!this.durableTimelineStore) { + return { timestamp: now.toISOString() }; + } + + const rows = await this.durableTimelineStore.getCommittedRows(agentId); + if (rows.length === 0) { + return { + nextSeq: 1, + timestamp: now.toISOString(), + }; + } + + return { + rows, + nextSeq: rows[rows.length - 1]!.seq + 1, + timestamp: rows[rows.length - 1]!.timestamp, + }; + } + private subscribeToSession(agent: ActiveManagedAgent): void { if (agent.unsubscribeSession) { return; @@ -2311,16 +2353,6 @@ export class AgentManager { this.trackBackgroundTask(task); } - private enqueueDurableTimelineDelete(agentId: string): void { - if (!this.durableTimelineStore) { - return; - } - const task = this.durableTimelineStore.deleteAgent(agentId).catch((err) => { - this.logger.error({ err, agentId }, "Failed to delete durable timeline rows"); - }); - this.trackBackgroundTask(task); - } - private trackBackgroundTask(task: Promise): void { this.backgroundTasks.add(task); void task.finally(() => { diff --git a/packages/server/src/server/provider-history-compatibility-boundary.test.ts b/packages/server/src/server/provider-history-compatibility-boundary.test.ts index b0d49103d5..ac74d5fb72 100644 --- a/packages/server/src/server/provider-history-compatibility-boundary.test.ts +++ b/packages/server/src/server/provider-history-compatibility-boundary.test.ts @@ -10,6 +10,6 @@ describe("provider history compatibility boundary", () => { ); expect(sessionSource).not.toMatch(/hydrateTimelineFromProvider\s*\(/); - expect(compatibilitySource.match(/hydrateTimelineFromProvider\s*\(/g)).toHaveLength(1); + expect(compatibilitySource).not.toMatch(/hydrateTimelineFromProvider\s*\(/); }); }); diff --git a/packages/server/src/server/provider-history-compatibility-service.test.ts b/packages/server/src/server/provider-history-compatibility-service.test.ts index 95a1a74131..16493e6afc 100644 --- a/packages/server/src/server/provider-history-compatibility-service.test.ts +++ b/packages/server/src/server/provider-history-compatibility-service.test.ts @@ -6,6 +6,8 @@ import { describe, expect, test, vi } from "vitest"; import { AgentManager } from "./agent/agent-manager.js"; import { AgentStorage, type StoredAgentRecord } from "./agent/agent-storage.js"; +import { DbAgentTimelineStore } from "./db/db-agent-timeline-store.js"; +import { openPaseoDatabase } from "./db/pglite-database.js"; import { ProviderHistoryCompatibilityService } from "./provider-history-compatibility-service.js"; import { createTestAgentClients } from "./test-utils/fake-agent-client.js"; @@ -55,15 +57,17 @@ function createCompatibilitySnapshot(overrides?: Partial } describe("ProviderHistoryCompatibilityService", () => { - test("ensureAgentLoaded replays provider history for an unloaded persisted agent", async () => { + test("ensureAgentLoaded seeds the live timeline from durable rows for an unloaded persisted agent", async () => { const workspaceRoot = mkdtempSync(path.join(os.tmpdir(), "provider-history-compat-load-")); const logger = pino({ level: "silent" }); + const database = await openPaseoDatabase(path.join(workspaceRoot, "db")); try { const storage = new AgentStorage(path.join(workspaceRoot, "agents"), logger); const manager = new AgentManager({ clients: createTestAgentClients(), registry: storage, + durableTimelineStore: new DbAgentTimelineStore(database.db), logger, idFactory: () => "00000000-0000-4000-8000-000000000301", }); @@ -81,6 +85,15 @@ describe("ProviderHistoryCompatibilityService", () => { await manager.runAgent(snapshot.id, "say 'timeline test'"); await manager.flush(); await storage.flush(); + rmSync( + path.join( + os.tmpdir(), + "paseo-fake-provider-history", + "codex", + `${snapshot.persistence?.sessionId}.jsonl`, + ), + { force: true }, + ); await manager.closeAgent(snapshot.id); const loaded = await service.ensureAgentLoaded({ agentId: snapshot.id }); @@ -88,6 +101,7 @@ describe("ProviderHistoryCompatibilityService", () => { expect(loaded.id).toBe(snapshot.id); expect(manager.getTimeline(snapshot.id)).toEqual([{ type: "assistant_message", text: "timeline test" }]); } finally { + await database.close(); rmSync(workspaceRoot, { recursive: true, force: true }); } }); @@ -149,9 +163,6 @@ describe("ProviderHistoryCompatibilityService", () => { resumeAgentFromPersistence: vi.fn(async () => deferred.promise), createAgent: vi.fn(), reloadAgentSession: vi.fn(), - hydrateTimelineFromProvider: vi.fn(async () => { - currentAgent = snapshot; - }), }; const logger = { child: () => logger, @@ -174,20 +185,15 @@ describe("ProviderHistoryCompatibilityService", () => { expect(secondResult).toEqual(snapshot); expect(agentStorage.get).toHaveBeenCalledTimes(1); expect(agentManager.resumeAgentFromPersistence).toHaveBeenCalledTimes(1); - expect(agentManager.hydrateTimelineFromProvider).toHaveBeenCalledTimes(1); }); - test("resumeAgent delegates to manager resume then compatibility hydrate", async () => { + test("resumeAgent delegates to manager resume", async () => { const snapshot = createCompatibilitySnapshot({ id: "agent-compat-resume" }); - let currentAgent: any = null; const agentManager = { - getAgent: vi.fn(() => currentAgent), + getAgent: vi.fn(() => null), resumeAgentFromPersistence: vi.fn(async () => snapshot), createAgent: vi.fn(), reloadAgentSession: vi.fn(), - hydrateTimelineFromProvider: vi.fn(async () => { - currentAgent = snapshot; - }), }; const logger = { child: () => logger, @@ -221,11 +227,10 @@ describe("ProviderHistoryCompatibilityService", () => { model: "gpt-5.4", }, ); - expect(agentManager.hydrateTimelineFromProvider).toHaveBeenCalledWith("agent-compat-resume"); expect(result).toEqual(snapshot); }); - test("refreshAgent reloads loaded persisted agents through the compatibility hydrate path", async () => { + test("refreshAgent reloads loaded persisted agents", async () => { const existing = createCompatibilitySnapshot({ id: "agent-compat-refresh-loaded" }); const reloaded = createCompatibilitySnapshot({ id: "agent-compat-refresh-loaded" }); let currentAgent: any = existing; @@ -237,7 +242,6 @@ describe("ProviderHistoryCompatibilityService", () => { currentAgent = reloaded; return reloaded; }), - hydrateTimelineFromProvider: vi.fn(async () => undefined), }; const logger = { child: () => logger, @@ -255,13 +259,10 @@ describe("ProviderHistoryCompatibilityService", () => { const result = await service.refreshAgent({ agentId: "agent-compat-refresh-loaded" }); expect(agentManager.reloadAgentSession).toHaveBeenCalledWith("agent-compat-refresh-loaded"); - expect(agentManager.hydrateTimelineFromProvider).toHaveBeenCalledWith( - "agent-compat-refresh-loaded", - ); expect(result).toEqual(reloaded); }); - test("refreshAgent rehydrates loaded non-persisted agents without reloading", async () => { + test("refreshAgent keeps loaded non-persisted agents without reloading", async () => { const existing = createCompatibilitySnapshot({ id: "agent-compat-refresh-live", persistence: null, @@ -271,7 +272,6 @@ describe("ProviderHistoryCompatibilityService", () => { resumeAgentFromPersistence: vi.fn(), createAgent: vi.fn(), reloadAgentSession: vi.fn(), - hydrateTimelineFromProvider: vi.fn(async () => undefined), }; const logger = { child: () => logger, @@ -289,15 +289,11 @@ describe("ProviderHistoryCompatibilityService", () => { const result = await service.refreshAgent({ agentId: "agent-compat-refresh-live" }); expect(agentManager.reloadAgentSession).not.toHaveBeenCalled(); - expect(agentManager.hydrateTimelineFromProvider).toHaveBeenCalledWith( - "agent-compat-refresh-live", - ); expect(result).toEqual(existing); }); - test("refreshAgent resumes unloaded persisted agents through the compatibility hydrate path", async () => { + test("refreshAgent resumes unloaded persisted agents", async () => { const snapshot = createCompatibilitySnapshot({ id: "agent-compat-refresh-cold" }); - let currentAgent: any = null; const record = createStoredAgentRecord({ id: "agent-compat-refresh-cold", cwd: "/tmp/refresh-cold", @@ -307,13 +303,10 @@ describe("ProviderHistoryCompatibilityService", () => { }, }); const agentManager = { - getAgent: vi.fn(() => currentAgent), + getAgent: vi.fn(() => null), resumeAgentFromPersistence: vi.fn(async () => snapshot), createAgent: vi.fn(), reloadAgentSession: vi.fn(), - hydrateTimelineFromProvider: vi.fn(async () => { - currentAgent = snapshot; - }), }; const logger = { child: () => logger, @@ -355,9 +348,6 @@ describe("ProviderHistoryCompatibilityService", () => { labels: {}, }, ); - expect(agentManager.hydrateTimelineFromProvider).toHaveBeenCalledWith( - "agent-compat-refresh-cold", - ); expect(result).toEqual(snapshot); }); @@ -368,7 +358,6 @@ describe("ProviderHistoryCompatibilityService", () => { resumeAgentFromPersistence: vi.fn(), createAgent: vi.fn(), reloadAgentSession: vi.fn(), - hydrateTimelineFromProvider: vi.fn(), } as any, agentStorage: { get: async () => diff --git a/packages/server/src/server/provider-history-compatibility-service.ts b/packages/server/src/server/provider-history-compatibility-service.ts index 44004560be..a57388ed8a 100644 --- a/packages/server/src/server/provider-history-compatibility-service.ts +++ b/packages/server/src/server/provider-history-compatibility-service.ts @@ -18,7 +18,6 @@ export type ProviderHistoryCompatibilityServiceOptions = { AgentManager, | "createAgent" | "getAgent" - | "hydrateTimelineFromProvider" | "reloadAgentSession" | "resumeAgentFromPersistence" >; @@ -26,8 +25,8 @@ export type ProviderHistoryCompatibilityServiceOptions = { logger: pino.Logger; }; -// Compatibility-only bridge for runtime paths that still need provider history -// replay during cold load, explicit resume, or refresh. +// Compatibility-only bridge for runtime paths that still need coordinated cold +// load, explicit resume, or refresh behavior. export class ProviderHistoryCompatibilityService { private readonly agentManager: ProviderHistoryCompatibilityServiceOptions["agentManager"]; private readonly agentStorage: ProviderHistoryCompatibilityServiceOptions["agentStorage"]; @@ -67,23 +66,18 @@ export class ProviderHistoryCompatibilityService { handle: AgentPersistenceHandle; overrides?: Partial; }): Promise { - const snapshot = await this.agentManager.resumeAgentFromPersistence( + return this.agentManager.resumeAgentFromPersistence( options.handle, options.overrides, ); - return this.hydrateCompatibilityTimeline({ agentId: snapshot.id, fallbackSnapshot: snapshot }); } async refreshAgent(options: { agentId: string }): Promise { const existing = this.agentManager.getAgent(options.agentId); if (existing) { - const snapshot = existing.persistence + return existing.persistence ? await this.agentManager.reloadAgentSession(options.agentId) : existing; - return this.hydrateCompatibilityTimeline({ - agentId: options.agentId, - fallbackSnapshot: snapshot, - }); } const record = await this.agentStorage.get(options.agentId); @@ -96,16 +90,12 @@ export class ProviderHistoryCompatibilityService { throw new Error(`Agent ${options.agentId} cannot be refreshed because it lacks persistence`); } - const snapshot = await this.agentManager.resumeAgentFromPersistence( + return this.agentManager.resumeAgentFromPersistence( handle, buildConfigOverrides(record), options.agentId, extractTimestamps(record), ); - return this.hydrateCompatibilityTimeline({ - agentId: options.agentId, - fallbackSnapshot: snapshot, - }); } private async loadStoredAgent(options: { agentId: string }): Promise { @@ -137,17 +127,6 @@ export class ProviderHistoryCompatibilityService { ); } - return this.hydrateCompatibilityTimeline({ - agentId: options.agentId, - fallbackSnapshot: snapshot, - }); - } - - private async hydrateCompatibilityTimeline(options: { - agentId: string; - fallbackSnapshot: ManagedAgent; - }): Promise { - await this.agentManager.hydrateTimelineFromProvider(options.agentId); - return this.agentManager.getAgent(options.agentId) ?? options.fallbackSnapshot; + return this.agentManager.getAgent(options.agentId) ?? snapshot; } } diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 83a3757143..1412dd5334 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -1689,10 +1689,11 @@ export class Session { try { await this.agentStorage.remove(agentId); + await this.agentManager.deleteCommittedTimeline(agentId); } catch (error: any) { this.sessionLogger.error( { err: error, agentId }, - `Failed to remove agent ${agentId} from registry`, + `Failed to fully delete agent ${agentId}`, ); }