@axiom-lattice/pg-stores 3.3.12 → 3.4.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "3.3.12",
+  "version": "3.4.0",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -25,8 +25,8 @@
     "@langchain/core": "1.1.30",
     "pg": "^8.16.3",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "8.0.2",
-    "@axiom-lattice/protocols": "4.10.0"
+    "@axiom-lattice/core": "9.0.0",
+    "@axiom-lattice/protocols": "5.0.0"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/PGVectorStoreProvider.ts

@@ -30,20 +30,26 @@ export class PGVectorStoreProvider implements VectorStoreProvider {
 
   async create(params: VectorStoreCreateParams): Promise<any> {
     const { tableName, embeddingKey } = params;
+    const tenantId = params.tenantId ?? "default";
 
-    let embeddings;
-    try {
-      embeddings = embeddingsLatticeManager.getEmbeddingsClient(embeddingKey);
-    } catch {
-      try {
-        embeddings = embeddingsLatticeManager.getEmbeddingsClient("default");
-      } catch {
-        throw new Error(
-          `No embedding model registered. Register one first:\n` +
-          `  import { registerEmbeddingsLattice } from "@axiom-lattice/core";\n` +
-          `  registerEmbeddingsLattice("default", new SomeEmbeddings(...));`
-        );
-      }
+    // Tenant first, deployment second — the same order every model resolver
+    // uses. A tenant model key is `${providerId}/${modelId}` and only resolves
+    // in that tenant's namespace.
+    const embeddings =
+      (await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
+        tenantId,
+        embeddingKey
+      )) ??
+      (await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
+        tenantId,
+        "default"
+      ));
+
+    if (!embeddings) {
+      throw new Error(
+        `No embedding model "${embeddingKey}" registered for tenant "${tenantId}", ` +
+          `and no deployment default. Configure one in the model provider panel.`
+      );
     }
 
     return PGVectorStore.initialize(embeddings, {

package/src/__tests__/AgentRunRetention.migrations.test.ts

@@ -0,0 +1,41 @@
+import { describe, expect, it, jest } from "@jest/globals";
+
+const registeredMigrationNames: string[] = [];
+
+jest.mock("@axiom-lattice/core", () => ({}));
+jest.mock("pg", () => ({ Pool: class Pool {} }));
+jest.mock("../migrations/migration", () => ({
+  MigrationManager: class MigrationManager {
+    register(migration: { name: string }): void {
+      registeredMigrationNames.push(migration.name);
+    }
+    async migrate(): Promise<void> {
+      throw new Error("stop after migrations");
+    }
+  },
+}));
+
+import { createPgStoreConfig } from "../createPgStoreConfig";
+import { createAgentRunRetentionMigration } from "../migrations/agent_run_retention_migration";
+import { agentRunProjectTaskRunningIndexMigration } from "../migrations/agent_run_project_task_running_index_migration";
+
+describe("agent run retention migration", () => {
+  it("is centrally registered exactly once", async () => {
+    expect(createAgentRunRetentionMigration.version).toBe(201);
+    expect(createAgentRunRetentionMigration.name).toBe("add_agent_run_retention_index");
+
+    registeredMigrationNames.length = 0;
+    await expect(createPgStoreConfig("postgres://unused")).rejects.toThrow(
+      "stop after migrations",
+    );
+
+    expect(
+      registeredMigrationNames.filter(
+        (name) => name === createAgentRunRetentionMigration.name,
+      ),
+    ).toHaveLength(1);
+    expect(registeredMigrationNames.at(-1)).toBe(
+      agentRunProjectTaskRunningIndexMigration.name,
+    );
+  });
+});

package/src/__tests__/PostgreSQLAgentRunLedger.test.ts

@@ -0,0 +1,305 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+// The @axiom-lattice/core barrel cannot load under jest in this package (ESM-only
+// transitive deps), so mock it with the real ledger module.
+jest.mock("@axiom-lattice/core", () => ({
+  ...jest.requireActual("../../../core/src/services/AgentRunLedger"),
+}));
+
+import { PostgreSQLAgentRunLedger } from "../stores/PostgreSQLAgentRunLedger";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+const TENANT_PREFIX = "pg-agent-run-ledger:";
+// Each Jest worker gets its own schema so global reads (`counts`) and cleanup
+// cannot observe another suite's rows when suites run in parallel.
+const SCHEMA = `agent_run_ledger_test_${randomUUID().replace(/-/g, "").slice(0, 16)}`;
+
+function testTenantId(): string {
+  return `${TENANT_PREFIX}${randomUUID()}`;
+}
+
+maybeDescribe("PostgreSQLAgentRunLedger (PostgreSQL)", () => {
+  let store: PostgreSQLAgentRunLedger;
+  let cleanupPool: Pool;
+
+  beforeAll(async () => {
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`CREATE SCHEMA ${SCHEMA}`);
+    await admin.end();
+
+    const poolConfig = {
+      connectionString: TEST_DATABASE_URL!,
+      options: `-c search_path=${SCHEMA}`,
+    };
+    store = new PostgreSQLAgentRunLedger({
+      poolConfig,
+      autoMigrate: false,
+    });
+    await store.initialize();
+    cleanupPool = new Pool(poolConfig);
+  });
+
+  beforeEach(async () => {
+    await cleanupPool.query(
+      "DELETE FROM lattice_agent_runs WHERE tenant_id LIKE $1",
+      [`${TENANT_PREFIX}%`],
+    );
+  });
+
+  afterAll(async () => {
+    await cleanupPool.end();
+    await store.dispose();
+
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`DROP SCHEMA IF EXISTS ${SCHEMA} CASCADE`);
+    await admin.end();
+  });
+
+  function input(overrides: Partial<Parameters<typeof store.beginRun>[0]> = {}) {
+    const tenantId = testTenantId();
+    const threadId = randomUUID();
+    return {
+      runId: randomUUID(),
+      tenantId,
+      workspaceId: randomUUID(),
+      projectId: randomUUID(),
+      assistantId: randomUUID(),
+      threadId,
+      owner: "replica-a",
+      leaseMs: 60_000,
+      inputMessageIds: [randomUUID(), randomUUID()],
+      ...overrides,
+    };
+  }
+
+  async function expire(runId: string): Promise<void> {
+    await cleanupPool.query(
+      `UPDATE lattice_agent_runs
+          SET lease_expires_at = now() - interval '1 second'
+        WHERE run_id = $1`,
+      [runId],
+    );
+  }
+
+  it("inserts a running run with owner, lease and input message ids", async () => {
+    const begin = input();
+    const record = await store.beginRun(begin);
+
+    expect(record.runId).toBe(begin.runId);
+    expect(record.tenantId).toBe(begin.tenantId);
+    expect(record.workspaceId).toBe(begin.workspaceId);
+    expect(record.projectId).toBe(begin.projectId);
+    expect(record.assistantId).toBe(begin.assistantId);
+    expect(record.threadId).toBe(begin.threadId);
+    expect(record.status).toBe("running");
+    expect(record.owner).toBe(begin.owner);
+    expect(record.inputMessageIds).toEqual(begin.inputMessageIds);
+    expect(record.attempts).toBe(1);
+    expect(record.error).toBeNull();
+    expect(record.cancelRequested).toBe(false);
+    expect(record.finishedAt).toBeNull();
+    expect(record.leaseExpiresAt).toBeGreaterThan(Date.now());
+    expect(typeof record.createdAt).toBe("number");
+    expect(typeof record.updatedAt).toBe("number");
+    expect(typeof record.startedAt).toBe("number");
+  });
+
+  it("reconciles a stale live running row when a new lease holder begins", async () => {
+    const first = input();
+    await store.beginRun(first);
+
+    // The new caller holds the thread-run lease, so the previous row is stale
+    // even though its ledger lease has not expired.
+    const second = await store.beginRun(
+      input({ tenantId: first.tenantId, threadId: first.threadId, owner: "replica-b" }),
+    );
+
+    const reconciled = await store.getRun(first.tenantId, first.runId);
+    expect(reconciled?.status).toBe("interrupted");
+    expect(reconciled?.error).toBe("superseded by a new run on the same thread");
+    expect(reconciled?.finishedAt).not.toBeNull();
+    expect(second.status).toBe("running");
+    expect((await store.findRunning(first.tenantId, first.threadId))?.runId).toBe(second.runId);
+  });
+
+  it("reaps an expired running row before inserting the next run", async () => {
+    const first = input();
+    await store.beginRun(first);
+    await expire(first.runId);
+
+    const second = await store.beginRun(
+      input({ tenantId: first.tenantId, threadId: first.threadId }),
+    );
+
+    const reaped = await store.getRun(first.tenantId, first.runId);
+    expect(reaped?.status).toBe("interrupted");
+    expect(reaped?.finishedAt).not.toBeNull();
+    expect(second.status).toBe("running");
+    expect((await store.findRunning(first.tenantId, first.threadId))?.runId).toBe(second.runId);
+  });
+
+  it("heartbeats only for the current owner", async () => {
+    const begin = input();
+    const record = await store.beginRun(begin);
+
+    await expect(store.heartbeatRun(record.runId, "someone-else", 120_000))
+      .resolves.toBe(false);
+    await expect(store.heartbeatRun(record.runId, begin.owner, 120_000))
+      .resolves.toBe(true);
+
+    const updated = await store.getRun(begin.tenantId, begin.runId);
+    expect(updated?.leaseExpiresAt).toBeGreaterThan(record.leaseExpiresAt);
+  });
+
+  it("settles only for the owner and only once", async () => {
+    const begin = input();
+    await store.beginRun(begin);
+
+    await expect(store.settleRun(begin.runId, "someone-else", "error", "nope"))
+      .resolves.toBe(false);
+    await expect(store.settleRun(begin.runId, begin.owner, "error", "boom"))
+      .resolves.toBe(true);
+    await expect(store.settleRun(begin.runId, begin.owner, "success"))
+      .resolves.toBe(false);
+
+    const settled = await store.getRun(begin.tenantId, begin.runId);
+    expect(settled?.status).toBe("error");
+    expect(settled?.error).toBe("boom");
+    expect(settled?.finishedAt).not.toBeNull();
+  });
+
+  it("reports not_found and terminal cancellations", async () => {
+    const tenantId = testTenantId();
+    await expect(store.requestCancel(tenantId, randomUUID())).resolves.toBe("not_found");
+
+    const begin = input({ tenantId });
+    await store.beginRun(begin);
+    await store.settleRun(begin.runId, begin.owner, "success");
+
+    await expect(store.requestCancel(tenantId, begin.runId)).resolves.toBe("terminal");
+  });
+
+  it("flags a live running run for cancellation", async () => {
+    const begin = input();
+    await store.beginRun(begin);
+
+    await expect(store.requestCancel(begin.tenantId, begin.runId)).resolves.toBe("requested");
+
+    const record = await store.getRun(begin.tenantId, begin.runId);
+    expect(record?.status).toBe("running");
+    expect(record?.cancelRequested).toBe(true);
+  });
+
+  it("interrupts an expired running run immediately", async () => {
+    const begin = input();
+    await store.beginRun(begin);
+    await expire(begin.runId);
+
+    await expect(store.requestCancel(begin.tenantId, begin.runId)).resolves.toBe("cancelled");
+
+    const record = await store.getRun(begin.tenantId, begin.runId);
+    expect(record?.status).toBe("interrupted");
+    expect(record?.cancelRequested).toBe(true);
+    expect(record?.finishedAt).not.toBeNull();
+  });
+
+  it("reaps expired running rows with their scope", async () => {
+    const first = input();
+    const second = input();
+    await store.beginRun(first);
+    await store.beginRun(second);
+    await expire(first.runId);
+    await expire(second.runId);
+
+    const reaped = await store.reapLegacyExpired!();
+
+    expect(reaped).toEqual(
+      expect.arrayContaining([
+        {
+          runId: first.runId,
+          tenantId: first.tenantId,
+          assistantId: first.assistantId,
+          threadId: first.threadId,
+          workspaceId: first.workspaceId,
+          projectId: first.projectId,
+        },
+        {
+          runId: second.runId,
+          tenantId: second.tenantId,
+          assistantId: second.assistantId,
+          threadId: second.threadId,
+          workspaceId: second.workspaceId,
+          projectId: second.projectId,
+        },
+      ]),
+    );
+
+    const record = await store.getRun(first.tenantId, first.runId);
+    expect(record?.status).toBe("interrupted");
+    expect(record?.error).toBe("run owner lost");
+    expect(record?.finishedAt).not.toBeNull();
+  });
+
+  it("lists a thread's runs newest first with clamped pagination", async () => {
+    const first = input();
+    await store.beginRun(first);
+    await store.settleRun(first.runId, first.owner, "success");
+
+    const second = await store.beginRun(
+      input({ tenantId: first.tenantId, threadId: first.threadId }),
+    );
+
+    const runs = await store.listRuns(first.tenantId, first.threadId, { limit: 500 });
+    expect(runs.map((run) => run.runId)).toEqual([second.runId, first.runId]);
+
+    const page = await store.listRuns(first.tenantId, first.threadId, {
+      limit: 0,
+      offset: 1,
+    });
+    expect(page.map((run) => run.runId)).toEqual([first.runId]);
+  });
+
+  it("maps SQL NULL workspace/project ids to undefined", async () => {
+    const begin = input({ workspaceId: undefined, projectId: undefined });
+    await store.beginRun(begin);
+
+    const record = await store.getRun(begin.tenantId, begin.runId);
+    expect(record?.workspaceId).toBeUndefined();
+    expect(record?.projectId).toBeUndefined();
+  });
+
+  it("counts runs by status and reports zero for missing statuses", async () => {
+    await expect(store.counts()).resolves.toEqual({
+      pending: 0,
+      running: 0,
+      success: 0,
+      error: 0,
+      interrupted: 0,
+      cancelled: 0,
+      timeout: 0,
+    });
+
+    const running = input();
+    await store.beginRun(running);
+
+    const success = input();
+    await store.beginRun(success);
+    await store.settleRun(success.runId, success.owner, "success");
+
+    const error = input();
+    await store.beginRun(error);
+    await store.settleRun(error.runId, error.owner, "error");
+
+    await expect(store.counts()).resolves.toEqual({
+      pending: 0,
+      running: 1,
+      success: 1,
+      error: 1,
+      interrupted: 0,
+      cancelled: 0,
+      timeout: 0,
+    });
+  });
+});

package/src/__tests__/PostgreSQLAgentRunRetention.test.ts

@@ -0,0 +1,126 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+jest.mock("@axiom-lattice/core", () => ({
+  ...jest.requireActual("../../../core/src/services/AgentRunLedger"),
+}));
+
+import { PostgreSQLAgentRunLedger } from "../stores/PostgreSQLAgentRunLedger";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+const TENANT = `pg-agent-run-retention:${randomUUID()}`;
+const SCHEMA = `agent_run_retention_test_${randomUUID().replace(/-/g, "").slice(0, 16)}`;
+const CUTOFF = Date.now() - 30 * 24 * 60 * 60 * 1000;
+
+maybeDescribe("PostgreSQLAgentRunLedger.deleteTerminalBefore (PostgreSQL)", () => {
+  let store: PostgreSQLAgentRunLedger;
+  let cleanupPool: Pool;
+
+  beforeAll(async () => {
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`CREATE SCHEMA ${SCHEMA}`);
+    await admin.end();
+
+    const poolConfig = {
+      connectionString: TEST_DATABASE_URL!,
+      options: `-c search_path=${SCHEMA}`,
+    };
+    store = new PostgreSQLAgentRunLedger({ poolConfig, autoMigrate: false });
+    await store.initialize();
+    cleanupPool = new Pool(poolConfig);
+  });
+
+  afterAll(async () => {
+    await cleanupPool.end();
+    await store.dispose();
+
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`DROP SCHEMA IF EXISTS ${SCHEMA} CASCADE`);
+    await admin.end();
+  });
+
+  beforeEach(async () => {
+    await cleanupPool.query(
+      "DELETE FROM lattice_agent_runs WHERE tenant_id = $1",
+      [TENANT],
+    );
+  });
+
+  async function terminal(
+    status: "success" | "error" | "interrupted" | "cancelled" | "timeout",
+    finishedAt: number,
+  ): Promise<string> {
+    const runId = randomUUID();
+    const input = {
+      runId,
+      tenantId: TENANT,
+      assistantId: randomUUID(),
+      threadId: randomUUID(),
+      owner: "replica-a",
+      leaseMs: 60_000,
+      inputMessageIds: [randomUUID()],
+    };
+    await store.beginRun(input);
+    await store.settleRun(runId, input.owner, status);
+    await cleanupPool.query(
+      "UPDATE lattice_agent_runs SET finished_at = to_timestamp($2::double precision / 1000) WHERE run_id = $1",
+      [runId, finishedAt],
+    );
+    return runId;
+  }
+
+  it("deletes only terminal rows before the cutoff and reports the count", async () => {
+    const oldSuccess = await terminal("success", CUTOFF - 2_000);
+    const oldTimeout = await terminal("timeout", CUTOFF - 1_000);
+    const atCutoff = await terminal("error", CUTOFF);
+    const newerCancelled = await terminal("cancelled", CUTOFF + 1_000);
+
+    const pendingId = randomUUID();
+    await store.createPendingRun({
+      runId: pendingId,
+      tenantId: TENANT,
+      assistantId: randomUUID(),
+      threadId: randomUUID(),
+      input: { id: randomUUID() },
+    });
+    const runningId = randomUUID();
+    await store.beginRun({
+      runId: runningId,
+      tenantId: TENANT,
+      assistantId: randomUUID(),
+      threadId: randomUUID(),
+      owner: "replica-a",
+      leaseMs: 60_000,
+      inputMessageIds: [randomUUID()],
+    });
+    await cleanupPool.query(
+      "UPDATE lattice_agent_runs SET created_at = now() - interval '60 days', updated_at = now() - interval '60 days' WHERE run_id = ANY($1::uuid[])",
+      [[pendingId, runningId]],
+    );
+
+    await expect(store.deleteTerminalBefore(CUTOFF, 100)).resolves.toBe(2);
+
+    await expect(store.getRun(TENANT, oldSuccess)).resolves.toBeUndefined();
+    await expect(store.getRun(TENANT, oldTimeout)).resolves.toBeUndefined();
+    await expect(store.getRun(TENANT, atCutoff)).resolves.toMatchObject({ status: "error" });
+    await expect(store.getRun(TENANT, newerCancelled)).resolves.toMatchObject({ status: "cancelled" });
+    await expect(store.getRun(TENANT, pendingId)).resolves.toMatchObject({ status: "pending" });
+    await expect(store.getRun(TENANT, runningId)).resolves.toMatchObject({ status: "running" });
+  });
+
+  it("deletes oldest terminal rows first and respects the limit", async () => {
+    const oldest = await terminal("error", CUTOFF - 3_000);
+    const middle = await terminal("interrupted", CUTOFF - 2_000);
+    const newest = await terminal("success", CUTOFF - 1_000);
+
+    await expect(store.deleteTerminalBefore(CUTOFF, 2)).resolves.toBe(2);
+
+    await expect(store.getRun(TENANT, oldest)).resolves.toBeUndefined();
+    await expect(store.getRun(TENANT, middle)).resolves.toBeUndefined();
+    await expect(store.getRun(TENANT, newest)).resolves.toMatchObject({ status: "success" });
+
+    await expect(store.deleteTerminalBefore(CUTOFF, 2)).resolves.toBe(1);
+    await expect(store.getRun(TENANT, newest)).resolves.toBeUndefined();
+  });
+});