Phase 4c: skip provider-history rehydrate when DB timeline is available

This commit is contained in:
Mohamed Boudra
2026-03-25 20:39:07 +07:00
parent eae83a53d8
commit 6027c38f8c
6 changed files with 208 additions and 76 deletions
@@ -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<AgentStreamEvent> {
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<boolean> {
return true;
}
async createSession(config: AgentSessionConfig): Promise<AgentSession> {
return new HistoryReplayProbeSession(config);
}
async resumeSession(
handle: AgentPersistenceHandle,
overrides?: Partial<AgentSessionConfig>,
): Promise<AgentSession> {
const metadata = (handle.metadata ?? {}) as Partial<AgentSessionConfig>;
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");
@@ -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<AgentSessionConfig>,
@@ -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<void> {
const agent = this.requireAgent(agentId);
if (this.durableTimelineStore) {
return;
}
await this.hydrateTimeline(agent);
}
async deleteCommittedTimeline(agentId: string): Promise<void> {
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<SeedAgentTimelineOptions> {
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>): void {
this.backgroundTasks.add(task);
void task.finally(() => {
@@ -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*\(/);
});
});
@@ -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<Record<string, unknown>
}
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 () =>
@@ -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<AgentSessionConfig>;
}): Promise<ManagedAgent> {
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<ManagedAgent> {
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<ManagedAgent> {
@@ -137,17 +127,6 @@ export class ProviderHistoryCompatibilityService {
);
}
return this.hydrateCompatibilityTimeline({
agentId: options.agentId,
fallbackSnapshot: snapshot,
});
}
private async hydrateCompatibilityTimeline(options: {
agentId: string;
fallbackSnapshot: ManagedAgent;
}): Promise<ManagedAgent> {
await this.agentManager.hydrateTimelineFromProvider(options.agentId);
return this.agentManager.getAgent(options.agentId) ?? options.fallbackSnapshot;
return this.agentManager.getAgent(options.agentId) ?? snapshot;
}
}
+2 -1
View File
@@ -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}`,
);
}