@axiom-lattice/pg-stores 3.3.13 → 4.0.0

package/src/__tests__/PostgreSQLStore.test.ts

@@ -0,0 +1,173 @@
+import { afterAll, beforeAll, beforeEach, describe, expect, it } from "@jest/globals";
+import { Pool } from "pg";
+import { PostgreSQLStore } from "../stores/PostgreSQLStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+
+maybeDescribe("PostgreSQLStore (PostgreSQL)", () => {
+  const tenant = "store-test-tenant";
+  let store: PostgreSQLStore;
+  let pool: Pool;
+
+  beforeAll(async () => {
+    pool = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await pool.query(`
+      CREATE TABLE IF NOT EXISTS lattice_store_items (
+        tenant_id TEXT NOT NULL,
+        namespace TEXT[] NOT NULL,
+        key TEXT NOT NULL,
+        value JSONB NOT NULL,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        expires_at TIMESTAMPTZ,
+        PRIMARY KEY (tenant_id, namespace, key)
+      )
+    `);
+    store = new PostgreSQLStore(pool);
+  });
+
+  beforeEach(async () => {
+    await pool.query("DELETE FROM lattice_store_items WHERE tenant_id LIKE 'store-test-%'");
+  });
+
+  afterAll(async () => {
+    await pool.query("DELETE FROM lattice_store_items WHERE tenant_id LIKE 'store-test-%'");
+    await pool.end();
+  });
+
+  it("round-trips items and isolates tenants", async () => {
+    await store.put(tenant, { namespace: ["a", "b"], key: "one", value: { n: 1 } });
+
+    expect(await store.get(tenant, ["a", "b"], "one")).toMatchObject({
+      namespace: ["a", "b"], key: "one", value: { n: 1 },
+    });
+    expect(await store.get("store-test-other", ["a", "b"], "one")).toBeNull();
+    expect(await store.delete(tenant, ["a", "b"], "one")).toBe(true);
+    expect(await store.get(tenant, ["a", "b"], "one")).toBeNull();
+  });
+
+  it("searches by namespace prefix, JSONB filter, and substring query", async () => {
+    await store.put(tenant, {
+      namespace: ["app", "a"], key: "one",
+      value: { message: "Hello World", nested: { tag: "x" }, list: [1, 2] },
+    });
+    await store.put(tenant, {
+      namespace: ["app", "b"], key: "two",
+      value: { message: "other", nested: { tag: "y" }, list: [1, 2] },
+    });
+    await store.put(tenant, {
+      namespace: ["other"], key: "three", value: { message: "Hello other" },
+    });
+
+    expect((await store.search(tenant, { namespacePrefix: ["app"] })).map((item) => item.key).sort())
+      .toEqual(["one", "two"]);
+    expect((await store.search(tenant, { filter: { nested: { tag: "x" } } })).map((item) => item.key))
+      .toEqual(["one"]);
+    expect((await store.search(tenant, { query: "HELLO" })).map((item) => item.key).sort())
+      .toEqual(["one", "three"]);
+    expect((await store.search(tenant, { filter: { list: [2] } })).map((item) => item.key).sort())
+      .toEqual(["one", "two"]);
+  });
+
+  it("treats percent and underscore literally in query", async () => {
+    await store.put(tenant, { namespace: ["q"], key: "percent", value: { text: "100% done" } });
+    await store.put(tenant, { namespace: ["q"], key: "underscore", value: { text: "under_score" } });
+
+    expect((await store.search(tenant, { query: "%" })).map((item) => item.key))
+      .toEqual(["percent"]);
+    expect((await store.search(tenant, { query: "_" })).map((item) => item.key))
+      .toEqual(["underscore"]);
+  });
+
+  it("upserts per tenant and does not leak across tenants", async () => {
+    await store.put(tenant, { namespace: ["shared"], key: "one", value: { owner: "a" } });
+    await store.put("store-test-other", { namespace: ["shared"], key: "one", value: { owner: "b" } });
+
+    expect((await store.get(tenant, ["shared"], "one"))?.value.owner).toBe("a");
+    expect((await store.get("store-test-other", ["shared"], "one"))?.value.owner).toBe("b");
+    expect((await store.search(tenant, { namespacePrefix: ["shared"] })).map((item) => item.value.owner))
+      .toEqual(["a"]);
+
+    expect(await store.delete(tenant, ["shared"], "one")).toBe(true);
+    expect((await store.get("store-test-other", ["shared"], "one"))?.value.owner).toBe("b");
+  });
+
+  it("clears a previous TTL when upserting without one", async () => {
+    await store.put(tenant, {
+      namespace: ["ttl"], key: "reset", value: { v: 1 }, ttlMinutes: 1,
+    });
+    await store.put(tenant, { namespace: ["ttl"], key: "reset", value: { v: 2 } });
+
+    expect((await store.get(tenant, ["ttl"], "reset"))?.value.v).toBe(2);
+    expect((await store.get(tenant, ["ttl"], "reset"))?.expiresAt).toBeUndefined();
+  });
+
+  it("rejects identity collisions and invalid TTL/pagination at the store", async () => {
+    await expect(store.put(tenant, { namespace: ["a\u001fb"], key: "k", value: {} }))
+      .rejects.toThrow("reserved separator");
+    await expect(store.put(tenant, { namespace: ["a"], key: "b\u001fc", value: {} }))
+      .rejects.toThrow("reserved separator");
+    await expect(store.put(tenant, {
+      namespace: ["a"], key: "k", value: {}, ttlMinutes: 0,
+    })).rejects.toThrow("ttlMinutes");
+    await expect(store.search(tenant, { limit: -1 })).rejects.toThrow("limit/offset");
+  });
+
+  it("uses namespace and key as stable tie-breakers", async () => {
+    await store.put(tenant, { namespace: ["b"], key: "one", value: {} });
+    await store.put(tenant, { namespace: ["a"], key: "two", value: {} });
+    await store.put(tenant, { namespace: ["a"], key: "one", value: {} });
+    await pool.query(
+      "UPDATE lattice_store_items SET updated_at = now() WHERE tenant_id = $1",
+      [tenant],
+    );
+
+    expect((await store.search(tenant, {})).map((item) => `${item.namespace[0]}:${item.key}`))
+      .toEqual(["a:one", "a:two", "b:one"]);
+  });
+
+  it("paginates in updated-at descending order", async () => {
+    for (const [index, key] of ["one", "two", "three"].entries()) {
+      await store.put(tenant, { namespace: ["page"], key, value: { index } });
+      await pool.query(
+        "UPDATE lattice_store_items SET updated_at = $1 WHERE tenant_id = $2 AND key = $3",
+        [new Date(Date.UTC(2026, 0, 1, 0, 0, index)), tenant, key],
+      );
+    }
+
+    expect((await store.search(tenant, { limit: 2 })).map((item) => item.key))
+      .toEqual(["three", "two"]);
+    expect((await store.search(tenant, { limit: 2, offset: 1 })).map((item) => item.key))
+      .toEqual(["two", "one"]);
+  });
+
+  it("hides expired items and deleteExpired removes them across tenants", async () => {
+    await store.put(tenant, {
+      namespace: ["ttl"], key: "short", value: { ok: true }, ttlMinutes: 1,
+    });
+    await store.put("store-test-other", {
+      namespace: ["ttl"], key: "short", value: { ok: true }, ttlMinutes: 1,
+    });
+    await pool.query(
+      `UPDATE lattice_store_items SET expires_at = now() - interval '1 minute'
+       WHERE tenant_id LIKE 'store-test-%'`,
+    );
+
+    expect(await store.get(tenant, ["ttl"], "short")).toBeNull();
+    expect(await store.search(tenant, { namespacePrefix: ["ttl"] })).toEqual([]);
+    expect(await store.listNamespaces(tenant)).toEqual([]);
+    expect(await store.deleteExpired(new Date(), 10)).toBe(2);
+    expect(await store.deleteExpired(new Date(), 10)).toBe(0);
+  });
+
+  it("lists distinct sorted namespaces with limit and prefix", async () => {
+    await store.put(tenant, { namespace: ["b"], key: "1", value: {} });
+    await store.put(tenant, { namespace: ["a", "x"], key: "2", value: {} });
+    await store.put(tenant, { namespace: ["a", "x"], key: "3", value: {} });
+
+    expect(await store.listNamespaces(tenant)).toEqual([["a", "x"], ["b"]]);
+    expect(await store.listNamespaces(tenant, [], 1)).toEqual([["a", "x"]]);
+    expect(await store.listNamespaces(tenant, ["a"])).toEqual([["a", "x"]]);
+  });
+});

package/src/__tests__/PostgreSQLThreadStore.test.ts

@@ -0,0 +1,80 @@
+import { afterAll, beforeAll, beforeEach, describe, expect, it } from "@jest/globals";
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+import { PostgreSQLThreadStore } from "../stores/PostgreSQLThreadStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+
+maybeDescribe("PostgreSQLThreadStore TTL (PostgreSQL)", () => {
+  let store: PostgreSQLThreadStore;
+  let cleanupPool: Pool;
+
+  beforeAll(async () => {
+    store = new PostgreSQLThreadStore({
+      poolConfig: TEST_DATABASE_URL!,
+      autoMigrate: false,
+    });
+    await store.initialize();
+    cleanupPool = new Pool({ connectionString: TEST_DATABASE_URL! });
+  });
+
+  beforeEach(async () => {
+    await cleanupPool.query("DELETE FROM lattice_threads WHERE id LIKE 'ttl-test-%'");
+  });
+
+  afterAll(async () => {
+    await cleanupPool.query("DELETE FROM lattice_threads WHERE id LIKE 'ttl-test-%'");
+    await cleanupPool.end();
+    await store.dispose();
+  });
+
+  function threadId(name: string): string {
+    return `ttl-test-${name}-${randomUUID()}`;
+  }
+
+  it("round-trips expiresAt through create, get, and update", async () => {
+    const tenantId = randomUUID();
+    const id = threadId("roundtrip");
+    const expiresAt = new Date("2026-01-01T00:00:00.000Z");
+
+    await store.createThread(tenantId, "assistant-1", id, { expiresAt });
+    expect((await store.getThreadById(tenantId, id))?.expiresAt).toEqual(expiresAt);
+
+    const next = new Date("2026-02-01T00:00:00.000Z");
+    await store.updateThread(tenantId, id, { expiresAt: next });
+    expect((await store.getThreadById(tenantId, id))?.expiresAt).toEqual(next);
+  });
+
+  it("lists expired threads across tenants, oldest first, up to the limit", async () => {
+    const tenantA = randomUUID();
+    const tenantB = randomUUID();
+    const now = new Date("2026-01-01T00:00:00.000Z");
+    const earlier = threadId("earlier");
+    const later = threadId("later");
+    const future = threadId("future");
+    const never = threadId("never");
+
+    await store.createThread(tenantA, "assistant-1", later, {
+      expiresAt: new Date("2025-12-31T23:00:00.000Z"),
+    });
+    await store.createThread(tenantB, "assistant-2", earlier, {
+      expiresAt: new Date("2025-12-31T22:00:00.000Z"),
+    });
+    await store.createThread(tenantA, "assistant-1", future, {
+      expiresAt: new Date("2026-01-01T01:00:00.000Z"),
+    });
+    await store.createThread(tenantA, "assistant-1", never, {});
+
+    const all = await store.listExpiredThreads(now, 10);
+    const ids = all.map((thread) => thread.id);
+    expect(ids.indexOf(earlier)).toBeGreaterThanOrEqual(0);
+    expect(ids.indexOf(later)).toBeGreaterThan(ids.indexOf(earlier));
+    expect(ids).not.toContain(future);
+    expect(ids).not.toContain(never);
+
+    const limited = await store.listExpiredThreads(now, 1);
+    expect(limited.map((thread) => thread.id)).toEqual([earlier]);
+  });
+});

package/src/__tests__/ThreadMessageQueueStore.integration.test.ts

@@ -79,4 +79,25 @@ maybeDescribe("ThreadMessageQueueStore claim/lease (PostgreSQL)", () => {
     const [pending] = await store.getPendingMessages(threadId);
     expect(pending.content).toBe("hello");
   });
+
+  it("reports hasPending=true for pending rows and false for processing-only rows", async () => {
+    const pendingThread = newThreadId();
+    const processingThread = newThreadId();
+    for (const threadId of [pendingThread, processingThread]) {
+      await store.addMessage({
+        threadId,
+        tenantId: "tenant",
+        assistantId: "assistant",
+        content: { id: randomUUID(), message: "message", queueMessage: { input: {} } },
+        id: randomUUID(),
+      });
+    }
+    await store.claimPendingMessages(processingThread, undefined, { owner: "r1", ttlMs: 60_000 });
+
+    const threads = await store.getThreadsWithPendingMessages();
+    const byThread = new Map(threads.map((thread) => [thread.threadId, thread]));
+
+    expect(byThread.get(pendingThread)?.hasPending).toBe(true);
+    expect(byThread.get(processingThread)?.hasPending).toBe(false);
+  });
 });

package/src/__tests__/ThreadMessageQueueStore.test.ts

@@ -457,9 +457,9 @@ describe('ThreadMessageQueueStore', () => {
   describe('getThreadsWithPendingMessages', () => {
     it('should return unique threads with pending messages', async () => {
       const mockThreads = [
-        { tenant_id: 'tenant-1', assistant_id: 'assistant-1', thread_id: 'thread-1' },
-        { tenant_id: 'tenant-1', assistant_id: 'assistant-1', thread_id: 'thread-2' },
-        { tenant_id: 'tenant-2', assistant_id: 'assistant-2', thread_id: 'thread-3' },
+        { tenant_id: 'tenant-1', assistant_id: 'assistant-1', thread_id: 'thread-1', has_pending: true },
+        { tenant_id: 'tenant-1', assistant_id: 'assistant-1', thread_id: 'thread-2', has_pending: true },
+        { tenant_id: 'tenant-2', assistant_id: 'assistant-2', thread_id: 'thread-3', has_pending: false },
       ];
 
       mockQuery.mockResolvedValueOnce({ rows: mockThreads });
@@ -470,6 +470,9 @@ describe('ThreadMessageQueueStore', () => {
       expect(result[0].threadId).toBe('thread-1');
       expect(result[1].threadId).toBe('thread-2');
       expect(result[2].threadId).toBe('thread-3');
+      expect(result.map((row) => row.hasPending)).toEqual([true, true, false]);
+      expect(mockQuery.mock.calls[0][0]).toMatch(/bool_or\(status = 'pending'\) AS has_pending/);
+      expect(mockQuery.mock.calls[0][0]).toMatch(/GROUP BY tenant_id, assistant_id, thread_id, workspace_id, project_id/);
     });
 
     it('preserves NULL workspace and project identity for legacy rows', async () => {
@@ -480,6 +483,7 @@ describe('ThreadMessageQueueStore', () => {
           thread_id: mockThreadId,
           workspace_id: null,
           project_id: null,
+          has_pending: true,
         }],
       });
 
@@ -489,6 +493,7 @@ describe('ThreadMessageQueueStore', () => {
         threadId: mockThreadId,
         workspaceId: null,
         projectId: null,
+        hasPending: true,
       }]);
     });
   });

package/src/__tests__/createPgStoreConfig.test.ts

@@ -0,0 +1,38 @@
+import { describe, expect, it, jest } from "@jest/globals";
+
+jest.mock("@axiom-lattice/core", () => ({}));
+
+jest.mock("pg", () => ({
+  Pool: class Pool {
+    query = async (): Promise<{ rows: unknown[] }> => ({ rows: [] });
+    connect = async (): Promise<unknown> => ({
+      query: async (): Promise<{ rows: unknown[] }> => ({ rows: [] }),
+      release: (): void => undefined,
+    });
+    end = async (): Promise<void> => undefined;
+  },
+}));
+
+jest.mock("../migrations/migration", () => ({
+  MigrationManager: class MigrationManager {
+    register(): void {}
+    async migrate(): Promise<void> {}
+  },
+}));
+
+jest.mock("@langchain/langgraph-checkpoint-postgres", () => ({
+  PostgresSaver: {
+    fromConnString: () => ({ setup: async (): Promise<void> => undefined }),
+  },
+}));
+
+import { createPgStoreConfig } from "../createPgStoreConfig";
+import { PostgreSQLAgentRunLedger } from "../stores/PostgreSQLAgentRunLedger";
+
+describe("createPgStoreConfig", () => {
+  it("registers the agentRun ledger store", async () => {
+    const stores = await createPgStoreConfig("postgres://unused");
+
+    expect(stores.agentRun).toBeInstanceOf(PostgreSQLAgentRunLedger);
+  });
+});

package/src/__tests__/eval-task-id.test.ts

@@ -6,16 +6,27 @@
  */
 
 import { describe, expect, it, jest } from "@jest/globals";
-import { addTaskIdToEvalRuns, evalMigrations } from "../migrations/eval_migrations";
+import { addCancelRequestedToEvalRuns, addTaskIdToEvalRuns, evalMigrations } from "../migrations/eval_migrations";
 
 describe("add_task_id_to_eval_runs migration", () => {
   it("is registered in evalMigrations", () => {
     expect(evalMigrations.map((m) => m.name)).toContain("add_task_id_to_eval_runs");
   });
 
-  it("has a version greater than all existing eval migrations", () => {
+  it("keeps task_id before the cancellation migration", () => {
     const versions = evalMigrations.map((m) => m.version);
-    expect(addTaskIdToEvalRuns.version).toBe(Math.max(...versions));
+    expect(addCancelRequestedToEvalRuns.version).toBe(Math.max(...versions));
+    expect(addTaskIdToEvalRuns.version).toBeLessThan(addCancelRequestedToEvalRuns.version);
+  });
+
+  it("emits the cooperative cancel column in up() and down()", async () => {
+    const upQueries: string[] = [];
+    await addCancelRequestedToEvalRuns.up({ query: jest.fn(async (sql: string) => { upQueries.push(sql); }) } as never);
+    expect(upQueries.some((q) => q.includes("cancel_requested BOOLEAN NOT NULL DEFAULT FALSE"))).toBe(true);
+
+    const downQueries: string[] = [];
+    await addCancelRequestedToEvalRuns.down({ query: jest.fn(async (sql: string) => { downQueries.push(sql); }) } as never);
+    expect(downQueries.some((q) => q.includes("DROP COLUMN IF EXISTS cancel_requested"))).toBe(true);
   });
 
   it("emits ADD COLUMN task_id in up()", async () => {

package/src/createPgStoreConfig.ts

@@ -30,7 +30,7 @@ import { PGVectorStoreProvider } from "./PGVectorStoreProvider";
 
 // Migration imports — all migrations registered centrally to guarantee
 // ordering and prevent name collisions across features.
-import { createThreadsTable } from "./migrations/thread_migrations";
+import { addThreadExpiresAt, createThreadsTable } from "./migrations/thread_migrations";
 import { addThreadTenantId } from "./migrations/thread_tenant_migration";
 import { changeThreadPrimaryKey } from "./migrations/thread_pk_migration";
 import { createScheduledTasksTable } from "./migrations/schedule_migrations";
@@ -80,6 +80,16 @@ import {
   addProjectLifecycleEventIndex,
 } from "./migrations/task_work_items_migration";
 import { createAgentWebAppsTable } from "./migrations/agent_web_apps_migration";
+import { createAgentRunLedgerTable } from "./migrations/agent_run_ledger_migration";
+import { createChannelReplyJobsTable } from "./migrations/channel_reply_job_migration";
+import { createAgentRunTaskQueueMigration } from "./migrations/agent_run_task_queue_migration";
+import { createAgentRunCapabilityRevisionMigration } from "./migrations/agent_run_capability_revision_migration";
+import { createAgentRunBaseCheckpointCapturedMigration } from "./migrations/agent_run_base_checkpoint_captured_migration";
+import { createAgentRunRollbackFailureMigration } from "./migrations/agent_run_rollback_failure_migration";
+import { agentRunRunningIndexSourceMigration } from "./migrations/agent_run_running_index_source_migration";
+import { createAgentRunRetentionMigration } from "./migrations/agent_run_retention_migration";
+import { createAgentRunProjectTaskMigration } from "./migrations/agent_run_project_task_migration";
+import { agentRunProjectTaskRunningIndexMigration } from "./migrations/agent_run_project_task_running_index_migration";
 import { addAgentWebAppsIdentity } from "./migrations/add_agent_web_apps_identity_migration";
 import { createCapabilityBundlesTable } from "./migrations/capability_bundle_migration";
 import { createProjectRoomReadStates, createProjectRoomTables } from "./migrations/project_room_migration";
@@ -111,6 +121,11 @@ import { PostgreSQLModelProviderStore } from "./stores/PostgreSQLModelProviderSt
 import { PostgreSQLLeaseStore } from "./stores/PostgreSQLLeaseStore";
 import { PostgreSQLOutboxStore } from "./stores/PostgreSQLOutboxStore";
 import { PostgreSQLPluginDeploymentStore } from "./stores/PostgreSQLPluginDeploymentStore";
+import { PostgreSQLAgentRunLedger } from "./stores/PostgreSQLAgentRunLedger";
+import { PostgreSQLChannelReplyJobStore } from "./stores/PostgreSQLChannelReplyJobStore";
+import { PostgreSQLCheckpointStorage } from "./stores/PostgreSQLCheckpointStorage";
+import { PostgreSQLStore } from "./stores/PostgreSQLStore";
+import { createStoreItemsTable } from "./migrations/store_migrations";
 
 export async function createPgStoreConfig(connectionString: string) {
   const pool = new Pool({ connectionString });
@@ -192,7 +207,21 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(dropPluginDeploymentAuthToken);  // v190
   mm.register(addModelProviderModelExtra);      // v191
   mm.register(addModelProviderModelDefault);    // v192
-  mm.register(addModelProviderModelKind);       // v193
+  mm.register(addModelProviderModelKind);       // v193 (model kind)
+  // v193 is intentionally shared with the agent-run ledger (versions only order
+  // registrations; migration identity is the name). Keep both.
+  mm.register(createAgentRunLedgerTable);       // v193
+  mm.register(createAgentRunTaskQueueMigration); // v196
+  mm.register(createAgentRunCapabilityRevisionMigration); // v197
+  mm.register(createAgentRunBaseCheckpointCapturedMigration); // v198
+  mm.register(createAgentRunRollbackFailureMigration);
+  mm.register(agentRunRunningIndexSourceMigration); // v202
+  mm.register(addThreadExpiresAt);              // v194
+  mm.register(createStoreItemsTable);           // v195
+  mm.register(createChannelReplyJobsTable);    // v200
+  mm.register(createAgentRunRetentionMigration); // v201
+  mm.register(createAgentRunProjectTaskMigration); // v203
+  mm.register(agentRunProjectTaskRunningIndexMigration); // v204
 
   await mm.migrate();
 
@@ -242,7 +271,11 @@ export async function createPgStoreConfig(connectionString: string) {
     lease: new PostgreSQLLeaseStore(opts),
     outbox: new PostgreSQLOutboxStore(opts),
     pluginDeployment: new PostgreSQLPluginDeploymentStore(opts),
+    agentRun: new PostgreSQLAgentRunLedger(opts),
+    channelReplyJob: new PostgreSQLChannelReplyJobStore(opts),
     vectorStoreProvider: new PGVectorStoreProvider(pool, connectionString),
     checkpoint,
+    checkpointStorage: new PostgreSQLCheckpointStorage(pool),
+    store: new PostgreSQLStore(pool),
   };
 }

package/src/index.ts

@@ -82,6 +82,13 @@ export * from "./stores/PostgreSQLProjectRoomReadStateStore";
 export * from "./migrations/project_room_migration";
 export * from "./stores/PostgreSQLPluginDeploymentStore";
 export * from "./migrations/plugin_deployment_migration";
+export * from "./migrations/agent_run_ledger_migration";
+export * from "./migrations/channel_reply_job_migration";
+export * from "./migrations/store_migrations";
+export * from "./stores/PostgreSQLAgentRunLedger";
+export * from "./stores/PostgreSQLChannelReplyJobStore";
+export * from "./stores/PostgreSQLCheckpointStorage";
+export * from "./stores/PostgreSQLStore";
 
 export * from "./PGVectorStoreProvider";
 

package/src/migrations/agent_run_base_checkpoint_captured_migration.ts

@@ -0,0 +1,34 @@
+import type { Migration } from "./migration";
+
+/**
+ * Records whether a run captured its rollback base checkpoint.
+ *
+ * @remarks A crash after claim but before `setBaseCheckpoint` leaves
+ * `base_checkpoint_id` null without meaning "no checkpoint". Recovery must only
+ * delete after `null` when this marker is true; otherwise it terminates the
+ * rollback with a repair signal instead of deleting unrelated history.
+ */
+export const createAgentRunBaseCheckpointCapturedMigration: Migration = {
+  version: 198,
+  name: "add_agent_run_base_checkpoint_captured",
+  up: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ADD COLUMN IF NOT EXISTS base_checkpoint_captured BOOLEAN NOT NULL DEFAULT FALSE
+    `);
+    // Rows that already persisted a base were captured by an older worker.
+    // A captured null base cannot be distinguished from an early crash, so it
+    // stays false (the safe choice: no checkpoint deletion).
+    await client.query(`
+      UPDATE lattice_agent_runs
+         SET base_checkpoint_captured = TRUE
+       WHERE base_checkpoint_id IS NOT NULL
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        DROP COLUMN IF EXISTS base_checkpoint_captured
+    `);
+  },
+};

package/src/migrations/agent_run_capability_revision_migration.ts

@@ -0,0 +1,24 @@
+import type { Migration } from "./migration";
+
+/**
+ * Adds the resolved capability runtime revision to run records.
+ *
+ * @remarks The run worker persists the revision before settling/interruption so
+ * HITL command runs and retries can pin the same runtime policy.
+ */
+export const createAgentRunCapabilityRevisionMigration: Migration = {
+  version: 197,
+  name: "add_agent_run_capability_revision",
+  up: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ADD COLUMN IF NOT EXISTS capability_runtime_revision TEXT
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        DROP COLUMN IF EXISTS capability_runtime_revision
+    `);
+  },
+};

package/src/migrations/agent_run_ledger_migration.ts

@@ -0,0 +1,51 @@
+import type { Migration } from "./migration";
+
+/**
+ * Creates the durable agent run ledger.
+ *
+ * @remarks
+ * The partial unique index enforces the ledger invariant: at most one
+ * `running` row per `(tenant_id, thread_id)`.
+ */
+export const createAgentRunLedgerTable: Migration = {
+  version: 193,
+  name: "create_lattice_agent_runs_table",
+  up: async (client) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_agent_runs (
+        run_id UUID PRIMARY KEY,
+        tenant_id VARCHAR(255) NOT NULL,
+        workspace_id VARCHAR(255),
+        project_id VARCHAR(255),
+        assistant_id VARCHAR(255) NOT NULL,
+        thread_id VARCHAR(255) NOT NULL,
+        status VARCHAR(32) NOT NULL,
+        owner VARCHAR(255) NOT NULL,
+        lease_expires_at TIMESTAMPTZ NOT NULL,
+        input_message_ids TEXT[] NOT NULL DEFAULT '{}',
+        attempts INTEGER NOT NULL DEFAULT 1,
+        error TEXT,
+        cancel_requested BOOLEAN NOT NULL DEFAULT FALSE,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+        started_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+        finished_at TIMESTAMPTZ
+      )
+    `);
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS uq_lattice_agent_runs_one_running
+        ON lattice_agent_runs (tenant_id, thread_id) WHERE status = 'running'
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_agent_runs_history
+        ON lattice_agent_runs (tenant_id, thread_id, created_at DESC)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_agent_runs_reap
+        ON lattice_agent_runs (status, lease_expires_at)
+    `);
+  },
+  down: async (client) => {
+    await client.query("DROP TABLE IF EXISTS lattice_agent_runs");
+  },
+};

package/src/migrations/agent_run_project_task_migration.ts

@@ -0,0 +1,34 @@
+import type { Migration } from "./migration";
+
+/**
+ * Adds ProjectTask identity to durable agent runs.
+ *
+ * @remarks
+ * `task_id` links a `source='project_task'` run to the canonical TaskItem it
+ * executes. The partial index resolves the single active run for a task without
+ * indexing ordinary run-API or legacy rows.
+ */
+export const createAgentRunProjectTaskMigration: Migration = {
+  version: 203,
+  name: "add_agent_run_project_task",
+  up: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ADD COLUMN IF NOT EXISTS task_id TEXT NULL
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_agent_runs_project_task
+        ON lattice_agent_runs (tenant_id, task_id, status, created_at DESC)
+        WHERE task_id IS NOT NULL
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      DROP INDEX IF EXISTS idx_lattice_agent_runs_project_task
+    `);
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        DROP COLUMN IF EXISTS task_id
+    `);
+  },
+};

package/src/migrations/agent_run_project_task_running_index_migration.ts

@@ -0,0 +1,33 @@
+import type { Migration } from "./migration";
+
+/**
+ * Enforces one running row per thread across both durable queue sources.
+ *
+ * @remarks
+ * The previous index only covered `run_api`, so concurrent `project_task` or
+ * cross-source claims could both mark a thread running. The queue sources share
+ * the same `thread-run` execution authority, so the unique predicate now spans
+ * `run_api` and `project_task`. Legacy rows keep their separate supervisor and
+ * are intentionally excluded.
+ */
+export const agentRunProjectTaskRunningIndexMigration: Migration = {
+  version: 204,
+  name: "scope_agent_run_running_index_to_queue_sources",
+  up: async (client) => {
+    await client.query("DROP INDEX IF EXISTS uq_lattice_agent_runs_one_running");
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS uq_lattice_agent_runs_one_running
+        ON lattice_agent_runs (tenant_id, thread_id)
+        WHERE status = 'running'
+          AND source IN ('run_api', 'project_task')
+    `);
+  },
+  down: async (client) => {
+    await client.query("DROP INDEX IF EXISTS uq_lattice_agent_runs_one_running");
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS uq_lattice_agent_runs_one_running
+        ON lattice_agent_runs (tenant_id, thread_id)
+        WHERE status = 'running' AND source = 'run_api'
+    `);
+  },
+};