diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 69c8f4c83e..961417585a 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -36,9 +36,23 @@ import { InMemoryAgentTimelineStore, type SeedAgentTimelineOptions, } from "./agent-timeline-store.js"; +import type { + AgentTimelineFetchOptions, + AgentTimelineFetchResult, + AgentTimelineRow, + AgentTimelineStore, +} from "./agent-timeline-store-types.js"; import { AGENT_PROVIDER_IDS } from "./provider-manifest.js"; export { AGENT_LIFECYCLE_STATUSES, type AgentLifecycleStatus }; +export type { + AgentTimelineCursor, + AgentTimelineFetchDirection, + AgentTimelineFetchOptions, + AgentTimelineFetchResult, + AgentTimelineRow, + AgentTimelineWindow, +} from "./agent-timeline-store-types.js"; export type AgentManagerEvent = | { type: "agent_state"; agent: ManagedAgent } @@ -77,6 +91,7 @@ export type AgentManagerOptions = { idFactory?: () => string; registry?: AgentSnapshotStore; onAgentAttention?: AgentAttentionCallback; + durableTimelineStore?: AgentTimelineStore; logger: Logger; }; @@ -95,43 +110,6 @@ export type WaitForAgentStartOptions = { signal?: AbortSignal; }; -export type AgentTimelineRow = { - seq: number; - timestamp: string; - item: AgentTimelineItem; -}; - -export type AgentTimelineCursor = { - seq: number; -}; - -export type AgentTimelineFetchDirection = "tail" | "before" | "after"; - -export type AgentTimelineFetchOptions = { - direction?: AgentTimelineFetchDirection; - cursor?: AgentTimelineCursor; - /** - * Number of canonical rows to return. - * - undefined: manager default - * - 0: all rows in the selected window - */ - limit?: number; -}; - -export type AgentTimelineWindow = { - minSeq: number; - maxSeq: number; - nextSeq: number; -}; - -export type AgentTimelineFetchResult = { - direction: AgentTimelineFetchDirection; - window: AgentTimelineWindow; - hasOlder: boolean; - hasNewer: boolean; - rows: AgentTimelineRow[]; -}; - type AttentionState = | { requiresAttention: false } | { @@ -305,6 +283,7 @@ export class AgentManager { private readonly subscribers = new Set(); private readonly idFactory: () => string; private readonly registry?: AgentSnapshotStore; + private readonly durableTimelineStore?: AgentTimelineStore; private readonly previousStatuses = new Map(); private readonly backgroundTasks = new Set>(); private onAgentAttention?: AgentAttentionCallback; @@ -313,6 +292,7 @@ export class AgentManager { constructor(options: AgentManagerOptions) { this.idFactory = options?.idFactory ?? (() => randomUUID()); this.registry = options?.registry; + this.durableTimelineStore = options?.durableTimelineStore; this.onAgentAttention = options?.onAgentAttention; this.logger = options.logger.child({ module: "agent", component: "agent-manager" }); if (options?.clients) { @@ -671,6 +651,7 @@ export class AgentManager { }; await session.close(); this.timelineStore.delete(agentId); + this.enqueueDurableTimelineDelete(agentId); this.emitState(closedAgent); this.logger.trace({ agentId }, "closeAgent: completed"); } @@ -1630,6 +1611,9 @@ export class AgentManager { if (timelineSeed || !this.timelineStore.has(resolvedAgentId)) { this.timelineStore.initialize(resolvedAgentId, timelineSeed ?? { timestamp: now.toISOString() }); } + if (options?.timelineRows?.length) { + this.enqueueDurableTimelineBulkInsert(resolvedAgentId, options.timelineRows); + } const managed = { id: resolvedAgentId, @@ -2221,7 +2205,9 @@ export class AgentManager { } private recordTimeline(agentId: string, item: AgentTimelineItem): AgentTimelineRow { - return this.timelineStore.append(agentId, item); + const row = this.timelineStore.append(agentId, item); + this.enqueueDurableTimelineAppend(agentId, row); + return row; } private emitState(agent: ManagedAgent, options?: { persist?: boolean }): void { @@ -2284,6 +2270,48 @@ export class AgentManager { this.trackBackgroundTask(task); } + private enqueueDurableTimelineAppend(agentId: string, row: AgentTimelineRow): void { + if (!this.durableTimelineStore) { + return; + } + const task = this.durableTimelineStore + .appendCommitted(agentId, row.item, { timestamp: row.timestamp }) + .then(() => undefined) + .catch((err) => { + this.logger.error( + { err, agentId, seq: row.seq, itemType: row.item.type }, + "Failed to append timeline row to durable store", + ); + }); + this.trackBackgroundTask(task); + } + + private enqueueDurableTimelineBulkInsert( + agentId: string, + rows: readonly AgentTimelineRow[], + ): void { + if (!this.durableTimelineStore || rows.length === 0) { + return; + } + const task = this.durableTimelineStore.bulkInsert(agentId, rows).catch((err) => { + this.logger.error( + { err, agentId, rowCount: rows.length }, + "Failed to seed durable timeline store", + ); + }); + 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/agent/agent-timeline-store-types.ts b/packages/server/src/server/agent/agent-timeline-store-types.ts new file mode 100644 index 0000000000..8882c1c651 --- /dev/null +++ b/packages/server/src/server/agent/agent-timeline-store-types.ts @@ -0,0 +1,54 @@ +import type { AgentTimelineItem } from "./agent-sdk-types.js"; + +export type AgentTimelineRow = { + seq: number; + timestamp: string; + item: AgentTimelineItem; +}; + +export type AgentTimelineCursor = { + seq: number; +}; + +export type AgentTimelineFetchDirection = "tail" | "before" | "after"; + +export type AgentTimelineFetchOptions = { + direction?: AgentTimelineFetchDirection; + cursor?: AgentTimelineCursor; + /** + * Number of canonical rows to return. + * - undefined: store default + * - 0: all rows in the selected window + */ + limit?: number; +}; + +export type AgentTimelineWindow = { + minSeq: number; + maxSeq: number; + nextSeq: number; +}; + +export type AgentTimelineFetchResult = { + direction: AgentTimelineFetchDirection; + window: AgentTimelineWindow; + hasOlder: boolean; + hasNewer: boolean; + rows: AgentTimelineRow[]; +}; + +export interface AgentTimelineStore { + appendCommitted( + agentId: string, + item: AgentTimelineItem, + options?: { timestamp?: string }, + ): Promise; + fetchCommitted( + agentId: string, + options?: AgentTimelineFetchOptions, + ): Promise; + getLatestCommittedSeq(agentId: string): Promise; + getCommittedRows(agentId: string): Promise; + deleteAgent(agentId: string): Promise; + bulkInsert(agentId: string, rows: readonly AgentTimelineRow[]): Promise; +} diff --git a/packages/server/src/server/agent/agent-timeline-store.ts b/packages/server/src/server/agent/agent-timeline-store.ts index 686f30827a..1a966de8c7 100644 --- a/packages/server/src/server/agent/agent-timeline-store.ts +++ b/packages/server/src/server/agent/agent-timeline-store.ts @@ -3,7 +3,7 @@ import type { AgentTimelineFetchOptions, AgentTimelineFetchResult, AgentTimelineRow, -} from "./agent-manager.js"; +} from "./agent-timeline-store-types.js"; export type SeedAgentTimelineOptions = { items?: readonly AgentTimelineItem[]; diff --git a/packages/server/src/server/bootstrap.ts b/packages/server/src/server/bootstrap.ts index a96e69fd8e..0451ca8e35 100644 --- a/packages/server/src/server/bootstrap.ts +++ b/packages/server/src/server/bootstrap.ts @@ -89,6 +89,7 @@ import { createAgentMcpServer } from "./agent/mcp-server.js"; import { createAllClients, shutdownProviders } from "./agent/provider-registry.js"; import { bootstrapWorkspaceRegistries } from "./workspace-registry-bootstrap.js"; import { DbAgentSnapshotStore } from "./db/db-agent-snapshot-store.js"; +import { DbAgentTimelineStore } from "./db/db-agent-timeline-store.js"; import { DbProjectRegistry } from "./db/db-project-registry.js"; import { DbWorkspaceRegistry } from "./db/db-workspace-registry.js"; import { importLegacyAgentSnapshots } from "./db/legacy-agent-snapshot-import.js"; @@ -364,6 +365,7 @@ export async function createPaseoDaemon( database = await openPaseoDatabase(path.join(config.paseoHome, "db")); logger.info({ elapsed: elapsed() }, "Paseo database opened"); const agentStorage = new DbAgentSnapshotStore(database.db); + const durableTimelineStore = new DbAgentTimelineStore(database.db); const agentManager = new AgentManager({ clients: { ...createAllClients(logger, { @@ -372,6 +374,7 @@ export async function createPaseoDaemon( ...config.agentClients, }, registry: agentStorage, + durableTimelineStore, logger, }); diff --git a/packages/server/src/server/db/db-agent-timeline-store.test.ts b/packages/server/src/server/db/db-agent-timeline-store.test.ts new file mode 100644 index 0000000000..3d586701d3 --- /dev/null +++ b/packages/server/src/server/db/db-agent-timeline-store.test.ts @@ -0,0 +1,214 @@ +import os from "node:os"; +import path from "node:path"; +import { mkdtempSync, rmSync } from "node:fs"; + +import { afterEach, beforeEach, describe, expect, test } from "vitest"; + +import type { AgentTimelineItem } from "../agent/agent-sdk-types.js"; +import type { AgentTimelineRow } from "../agent/agent-timeline-store-types.js"; +import { openPaseoDatabase, type PaseoDatabaseHandle } from "./pglite-database.js"; +import { DbAgentTimelineStore } from "./db-agent-timeline-store.js"; +import { agentTimelineRows } from "./schema.js"; + +function createTimestamp(seq: number): string { + return new Date(Date.UTC(2026, 2, 1, 0, 0, seq)).toISOString(); +} + +function createTimelineItem( + type: Extract, + value: string, +): AgentTimelineItem { + if (type === "user_message") { + return { + type, + text: `user-${value}`, + messageId: `message-${value}`, + }; + } + + return { + type, + text: `assistant-${value}`, + }; +} + +function createRow(seq: number, item?: AgentTimelineItem): AgentTimelineRow { + return { + seq, + timestamp: createTimestamp(seq), + item: item ?? createTimelineItem("assistant_message", String(seq)), + }; +} + +describe("DbAgentTimelineStore", () => { + let tmpDir: string; + let dataDir: string; + let database: PaseoDatabaseHandle; + let store: DbAgentTimelineStore; + + beforeEach(async () => { + tmpDir = mkdtempSync(path.join(os.tmpdir(), "db-agent-timeline-store-")); + dataDir = path.join(tmpDir, "db"); + database = await openPaseoDatabase(dataDir); + store = new DbAgentTimelineStore(database.db); + }); + + afterEach(async () => { + await database.close(); + rmSync(tmpDir, { recursive: true, force: true }); + }); + + test("appendCommitted assigns sequential seq numbers per agent", async () => { + expect( + await store.appendCommitted("agent-1", createTimelineItem("assistant_message", "1")), + ).toEqual({ + seq: 1, + timestamp: expect.any(String), + item: createTimelineItem("assistant_message", "1"), + }); + + expect( + await store.appendCommitted("agent-1", createTimelineItem("assistant_message", "2")), + ).toEqual({ + seq: 2, + timestamp: expect.any(String), + item: createTimelineItem("assistant_message", "2"), + }); + + expect( + await store.appendCommitted("agent-1", createTimelineItem("assistant_message", "3")), + ).toEqual({ + seq: 3, + timestamp: expect.any(String), + item: createTimelineItem("assistant_message", "3"), + }); + }); + + test("appendCommitted for different agents has independent seq sequences", async () => { + const firstAgentFirstRow = await store.appendCommitted( + "agent-1", + createTimelineItem("assistant_message", "a1"), + ); + const secondAgentFirstRow = await store.appendCommitted( + "agent-2", + createTimelineItem("assistant_message", "b1"), + ); + const firstAgentSecondRow = await store.appendCommitted( + "agent-1", + createTimelineItem("assistant_message", "a2"), + ); + + expect(firstAgentFirstRow.seq).toBe(1); + expect(secondAgentFirstRow.seq).toBe(1); + expect(firstAgentSecondRow.seq).toBe(2); + }); + + test("fetchCommitted tail returns the last N rows", async () => { + await store.bulkInsert("agent-1", [1, 2, 3, 4, 5].map((seq) => createRow(seq))); + + await expect( + store.fetchCommitted("agent-1", { + direction: "tail", + limit: 2, + }), + ).resolves.toEqual({ + direction: "tail", + window: { + minSeq: 1, + maxSeq: 5, + nextSeq: 6, + }, + hasOlder: true, + hasNewer: false, + rows: [createRow(4), createRow(5)], + }); + }); + + test("fetchCommitted after-cursor returns rows after a given seq", async () => { + await store.bulkInsert("agent-1", [1, 2, 3, 4, 5].map((seq) => createRow(seq))); + + await expect( + store.fetchCommitted("agent-1", { + direction: "after", + cursor: { seq: 2 }, + limit: 2, + }), + ).resolves.toEqual({ + direction: "after", + window: { + minSeq: 1, + maxSeq: 5, + nextSeq: 6, + }, + hasOlder: true, + hasNewer: true, + rows: [createRow(3), createRow(4)], + }); + }); + + test("fetchCommitted before-cursor returns rows before a given seq", async () => { + await store.bulkInsert("agent-1", [1, 2, 3, 4, 5].map((seq) => createRow(seq))); + + await expect( + store.fetchCommitted("agent-1", { + direction: "before", + cursor: { seq: 4 }, + limit: 2, + }), + ).resolves.toEqual({ + direction: "before", + window: { + minSeq: 1, + maxSeq: 5, + nextSeq: 6, + }, + hasOlder: true, + hasNewer: true, + rows: [createRow(2), createRow(3)], + }); + }); + + test("getLatestCommittedSeq returns 0 for an unknown agent", async () => { + await expect(store.getLatestCommittedSeq("missing-agent")).resolves.toBe(0); + }); + + test("getLatestCommittedSeq returns the latest seq after appends", async () => { + await store.appendCommitted("agent-1", createTimelineItem("assistant_message", "1")); + await store.appendCommitted("agent-1", createTimelineItem("assistant_message", "2")); + + await expect(store.getLatestCommittedSeq("agent-1")).resolves.toBe(2); + }); + + test("deleteAgent removes all rows for the target agent", async () => { + await store.bulkInsert("agent-1", [createRow(1), createRow(2)]); + await store.bulkInsert("agent-2", [createRow(1)]); + + await store.deleteAgent("agent-1"); + + await expect(store.getCommittedRows("agent-1")).resolves.toEqual([]); + await expect(store.getCommittedRows("agent-2")).resolves.toEqual([createRow(1)]); + }); + + test("bulkInsert preserves provided seq numbers", async () => { + const rows = [createRow(3), createRow(7)]; + + await store.bulkInsert("agent-1", rows); + + await expect(store.getCommittedRows("agent-1")).resolves.toEqual(rows); + }); + + test("item_kind is populated from item.type", async () => { + await store.appendCommitted("agent-1", createTimelineItem("user_message", "kind-check"), { + timestamp: createTimestamp(1), + }); + + await expect(database.db.select().from(agentTimelineRows)).resolves.toEqual([ + expect.objectContaining({ + agentId: "agent-1", + seq: 1, + committedAt: createTimestamp(1), + itemKind: "user_message", + }), + ]); + }); +}); diff --git a/packages/server/src/server/db/db-agent-timeline-store.ts b/packages/server/src/server/db/db-agent-timeline-store.ts new file mode 100644 index 0000000000..6a6e78da78 --- /dev/null +++ b/packages/server/src/server/db/db-agent-timeline-store.ts @@ -0,0 +1,225 @@ +import { and, asc, desc, eq, gt, lt, sql } from "drizzle-orm"; + +import type { + AgentTimelineFetchOptions, + AgentTimelineFetchResult, + AgentTimelineRow, + AgentTimelineStore, + AgentTimelineWindow, +} from "../agent/agent-timeline-store-types.js"; +import type { AgentTimelineItem } from "../agent/agent-sdk-types.js"; +import type { PaseoDatabaseHandle } from "./pglite-database.js"; +import { agentTimelineRows } from "./schema.js"; + +type AgentTimelineRowRecord = typeof agentTimelineRows.$inferSelect; +type AgentTimelineRowInsert = typeof agentTimelineRows.$inferInsert; + +const DEFAULT_TIMELINE_FETCH_LIMIT = 200; + +function toTimelineRow(row: AgentTimelineRowRecord): AgentTimelineRow { + return { + seq: row.seq, + timestamp: row.committedAt, + item: row.item, + }; +} + +function toInsertValues(agentId: string, row: AgentTimelineRow): AgentTimelineRowInsert { + return { + agentId, + seq: row.seq, + committedAt: row.timestamp, + item: row.item, + itemKind: row.item.type, + }; +} + +function normalizeFetchLimit(limit: number | undefined): number { + if (limit === undefined) { + return DEFAULT_TIMELINE_FETCH_LIMIT; + } + return Math.max(0, Math.floor(limit)); +} + +export class DbAgentTimelineStore implements AgentTimelineStore { + private readonly db: PaseoDatabaseHandle["db"]; + + constructor(db: PaseoDatabaseHandle["db"]) { + this.db = db; + } + + async appendCommitted( + agentId: string, + item: AgentTimelineItem, + options?: { timestamp?: string }, + ): Promise { + const nextSeq = (await this.getMaxSeq(agentId)) + 1; + const row: AgentTimelineRow = { + seq: nextSeq, + timestamp: options?.timestamp ?? new Date().toISOString(), + item, + }; + + await this.db.insert(agentTimelineRows).values(toInsertValues(agentId, row)); + return row; + } + + async fetchCommitted( + agentId: string, + options?: AgentTimelineFetchOptions, + ): Promise { + const direction = options?.direction ?? "tail"; + const limit = normalizeFetchLimit(options?.limit); + const selectAll = limit === 0; + const window = await this.getWindow(agentId); + + if (window.maxSeq === 0) { + return { + direction, + window, + hasOlder: false, + hasNewer: false, + rows: [], + }; + } + + if (direction === "tail") { + const rows = selectAll + ? await this.db + .select() + .from(agentTimelineRows) + .where(eq(agentTimelineRows.agentId, agentId)) + .orderBy(asc(agentTimelineRows.seq)) + : ( + await this.db + .select() + .from(agentTimelineRows) + .where(eq(agentTimelineRows.agentId, agentId)) + .orderBy(desc(agentTimelineRows.seq)) + .limit(limit) + ).reverse(); + const selected = rows.map(toTimelineRow); + return { + direction, + window, + hasOlder: selected.length > 0 && selected[0]!.seq > window.minSeq, + hasNewer: false, + rows: selected, + }; + } + + if (direction === "after") { + const baseSeq = options?.cursor?.seq ?? 0; + const rows = ( + selectAll + ? await this.db + .select() + .from(agentTimelineRows) + .where(and(eq(agentTimelineRows.agentId, agentId), gt(agentTimelineRows.seq, baseSeq))) + .orderBy(asc(agentTimelineRows.seq)) + : await this.db + .select() + .from(agentTimelineRows) + .where(and(eq(agentTimelineRows.agentId, agentId), gt(agentTimelineRows.seq, baseSeq))) + .orderBy(asc(agentTimelineRows.seq)) + .limit(limit) + ).map(toTimelineRow); + + if (rows.length === 0) { + return { + direction, + window, + hasOlder: baseSeq >= window.minSeq, + hasNewer: false, + rows, + }; + } + + const lastSelected = rows[rows.length - 1]!; + return { + direction, + window, + hasOlder: rows[0]!.seq > window.minSeq, + hasNewer: lastSelected.seq < window.maxSeq, + rows, + }; + } + + const beforeSeq = options?.cursor?.seq ?? window.nextSeq; + const rows = ( + selectAll + ? await this.db + .select() + .from(agentTimelineRows) + .where(and(eq(agentTimelineRows.agentId, agentId), lt(agentTimelineRows.seq, beforeSeq))) + .orderBy(asc(agentTimelineRows.seq)) + : ( + await this.db + .select() + .from(agentTimelineRows) + .where(and(eq(agentTimelineRows.agentId, agentId), lt(agentTimelineRows.seq, beforeSeq))) + .orderBy(desc(agentTimelineRows.seq)) + .limit(limit) + ).reverse() + ).map(toTimelineRow); + + return { + direction, + window, + hasOlder: rows.length > 0 && rows[0]!.seq > window.minSeq, + hasNewer: beforeSeq <= window.maxSeq, + rows, + }; + } + + async getLatestCommittedSeq(agentId: string): Promise { + return this.getMaxSeq(agentId); + } + + async getCommittedRows(agentId: string): Promise { + const rows = await this.db + .select() + .from(agentTimelineRows) + .where(eq(agentTimelineRows.agentId, agentId)) + .orderBy(asc(agentTimelineRows.seq)); + return rows.map(toTimelineRow); + } + + async deleteAgent(agentId: string): Promise { + await this.db.delete(agentTimelineRows).where(eq(agentTimelineRows.agentId, agentId)); + } + + async bulkInsert(agentId: string, rows: readonly AgentTimelineRow[]): Promise { + if (rows.length === 0) { + return; + } + await this.db.insert(agentTimelineRows).values(rows.map((row) => toInsertValues(agentId, row))); + } + + private async getMaxSeq(agentId: string): Promise { + const [row] = await this.db + .select({ + maxSeq: sql`coalesce(max(${agentTimelineRows.seq}), 0)`, + }) + .from(agentTimelineRows) + .where(eq(agentTimelineRows.agentId, agentId)); + return Number(row?.maxSeq ?? 0); + } + + private async getWindow(agentId: string): Promise { + const [row] = await this.db + .select({ + minSeq: sql`coalesce(min(${agentTimelineRows.seq}), 0)`, + maxSeq: sql`coalesce(max(${agentTimelineRows.seq}), 0)`, + }) + .from(agentTimelineRows) + .where(eq(agentTimelineRows.agentId, agentId)); + const minSeq = Number(row?.minSeq ?? 0); + const maxSeq = Number(row?.maxSeq ?? 0); + return { + minSeq, + maxSeq, + nextSeq: maxSeq + 1, + }; + } +}