@axiom-lattice/pg-stores 5.0.0 → 5.0.2

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

@@ -0,0 +1,447 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+import { PostgreSQLWorkflowTrackingStore } from "../stores/PostgreSQLWorkflowTrackingStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const describeIfDatabase = TEST_DATABASE_URL ? describe : describe.skip;
+
+describeIfDatabase("PostgreSQLWorkflowTrackingStore integration", () => {
+  let store: PostgreSQLWorkflowTrackingStore;
+  let pool: Pool;
+  const tenants = new Set<string>();
+
+  beforeAll(async () => {
+    store = new PostgreSQLWorkflowTrackingStore({
+      poolConfig: TEST_DATABASE_URL!,
+      autoMigrate: true,
+    });
+    await store.initialize();
+    pool = new Pool({ connectionString: TEST_DATABASE_URL });
+  });
+
+  afterEach(async () => {
+    if (tenants.size === 0) {
+      return;
+    }
+    await pool.query("DELETE FROM lattice_workflow_runs WHERE tenant_id = ANY($1::text[])", [
+      [...tenants],
+    ]);
+    tenants.clear();
+  });
+
+  afterAll(async () => {
+    await pool.end();
+    await store.dispose();
+  });
+
+  function tenantId(): string {
+    const id = `workflow-${randomUUID()}`;
+    tenants.add(id);
+    return id;
+  }
+
+  function topologyEdges() {
+    return [
+      { from: "input", to: "planner", purpose: "kickoff" },
+      { from: "planner", to: "worker", purpose: "delegate" },
+    ];
+  }
+
+  async function createRun(
+    tenant: string,
+    overrides: {
+      threadId?: string;
+      assistantId?: string;
+      metadata?: Record<string, unknown>;
+    } = {},
+  ) {
+    return store.createWorkflowRun({
+      tenantId: tenant,
+      assistantId: overrides.assistantId ?? `assistant-${randomUUID()}`,
+      threadId: overrides.threadId ?? `thread-${randomUUID()}`,
+      topologyEdges: topologyEdges(),
+      metadata: overrides.metadata ?? {
+        source: "test",
+        nested: { workflow: "demo", tags: ["a", "b"] },
+      },
+    });
+  }
+
+  describe("workflow runs", () => {
+    it("creates a run and reads back topology edges and metadata", async () => {
+      const tenant = tenantId();
+      const threadId = `thread-${randomUUID()}`;
+      const assistantId = `assistant-${randomUUID()}`;
+
+      const run = await store.createWorkflowRun({
+        tenantId: tenant,
+        assistantId,
+        threadId,
+        topologyEdges: topologyEdges(),
+        metadata: { source: "test", nested: { workflow: "demo", tags: ["a", "b"] } },
+      });
+
+      expect(run).toMatchObject({
+        tenantId: tenant,
+        assistantId,
+        threadId,
+        status: "running",
+        totalEdges: 2,
+        completedEdges: 0,
+      });
+      expect(run.topologyEdges).toEqual(topologyEdges());
+      expect(run.metadata).toEqual({ source: "test", nested: { workflow: "demo", tags: ["a", "b"] } });
+      expect(run.startedAt).toBeInstanceOf(Date);
+
+      const read = await store.getWorkflowRun(run.id);
+      expect(read).toMatchObject({ id: run.id, tenantId: tenant, status: "running" });
+      expect(read!.topologyEdges).toEqual(topologyEdges());
+      expect(read!.metadata).toEqual(run.metadata);
+    });
+
+    it("returns null for an absent run", async () => {
+      expect(await store.getWorkflowRun(randomUUID())).toBeNull();
+    });
+
+    it("creates distinct runs for the same thread and timestamp", async () => {
+      const tenant = tenantId();
+      const threadId = `thread-${randomUUID()}`;
+      const assistantId = `assistant-${randomUUID()}`;
+      const nowSpy = jest.spyOn(Date, "now").mockReturnValue(1700000000000);
+
+      try {
+        const first = await store.createWorkflowRun({
+          tenantId: tenant,
+          assistantId,
+          threadId,
+          topologyEdges: topologyEdges(),
+        });
+        const second = await store.createWorkflowRun({
+          tenantId: tenant,
+          assistantId,
+          threadId,
+          topologyEdges: topologyEdges(),
+        });
+
+        expect(second.id).not.toBe(first.id);
+        expect(first.id.startsWith(`${threadId}_`)).toBe(true);
+        expect(second.id.startsWith(`${threadId}_`)).toBe(true);
+
+        const listed = await store.getWorkflowRunsByThreadId(tenant, threadId);
+        expect(listed.map((r) => r.id).sort()).toEqual([first.id, second.id].sort());
+      } finally {
+        nowSpy.mockRestore();
+      }
+    });
+
+    it("updates run status, counts, error, completion time, and metadata", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+      const completedAt = new Date("2026-02-03T04:05:06.000Z");
+
+      const updated = await store.updateWorkflowRun(run.id, {
+        status: "failed",
+        completedEdges: 1,
+        errorMessage: "worker exploded",
+        completedAt,
+        metadata: { failure: { code: "E_WORKER", retryable: false } },
+      });
+
+      expect(updated).toMatchObject({
+        id: run.id,
+        status: "failed",
+        completedEdges: 1,
+        totalEdges: 2,
+        errorMessage: "worker exploded",
+      });
+      expect(updated!.completedAt?.toISOString()).toBe(completedAt.toISOString());
+      expect(updated!.metadata).toEqual({ failure: { code: "E_WORKER", retryable: false } });
+      expect(updated!.topologyEdges).toEqual(topologyEdges());
+    });
+
+    it("returns null when updating an absent run", async () => {
+      expect(await store.updateWorkflowRun(randomUUID(), { status: "completed" })).toBeNull();
+    });
+
+    it("deletes a run and treats a repeat delete as a no-op", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+
+      await store.deleteWorkflowRun(run.id);
+      expect(await store.getWorkflowRun(run.id)).toBeNull();
+      await expect(store.deleteWorkflowRun(run.id)).resolves.toBeUndefined();
+    });
+
+    it("lists runs by thread, assistant, and tenant with tenant isolation", async () => {
+      const tenantA = tenantId();
+      const tenantB = tenantId();
+      const assistantA = `assistant-${randomUUID()}`;
+      const assistantB = `assistant-${randomUUID()}`;
+      const threadA1 = `thread-${randomUUID()}`;
+      const threadA2 = `thread-${randomUUID()}`;
+      const threadA3 = `thread-${randomUUID()}`;
+
+      const r1 = await createRun(tenantA, { threadId: threadA1, assistantId: assistantA });
+      const r2 = await createRun(tenantA, { threadId: threadA2, assistantId: assistantA });
+      const r3 = await createRun(tenantA, { threadId: threadA3, assistantId: assistantB });
+      const r4 = await createRun(tenantB, { threadId: `thread-${randomUUID()}`, assistantId: assistantA });
+
+      expect((await store.getWorkflowRunsByThreadId(tenantA, threadA1)).map((r) => r.id)).toEqual([
+        r1.id,
+      ]);
+      expect((await store.getWorkflowRunsByThreadId(tenantA, threadA2)).map((r) => r.id)).toEqual([
+        r2.id,
+      ]);
+      expect((await store.getWorkflowRunsByThreadId(tenantA, threadA3)).map((r) => r.id)).toEqual([
+        r3.id,
+      ]);
+      expect(
+        (await store.getWorkflowRunsByAssistantId(tenantA, assistantA)).map((r) => r.id).sort(),
+      ).toEqual([r1.id, r2.id].sort());
+      expect((await store.getWorkflowRunsByTenantId(tenantA)).map((r) => r.id).sort()).toEqual(
+        [r1.id, r2.id, r3.id].sort(),
+      );
+      expect((await store.getWorkflowRunsByTenantId(tenantB)).map((r) => r.id)).toEqual([r4.id]);
+
+      expect(await store.getWorkflowRunsByThreadId(tenantB, threadA1)).toEqual([]);
+      expect(await store.getWorkflowRunsByAssistantId(tenantB, assistantB)).toEqual([]);
+    });
+
+    it("filters queryWorkflowRuns by status and assistant", async () => {
+      const tenant = tenantId();
+      const otherTenant = tenantId();
+      const assistantA = `assistant-${randomUUID()}`;
+      const assistantB = `assistant-${randomUUID()}`;
+      const r1 = await createRun(tenant, { assistantId: assistantA });
+      const r2 = await createRun(tenant, { assistantId: assistantA });
+      const r3 = await createRun(tenant, { assistantId: assistantB });
+      await store.updateWorkflowRun(r2.id, { status: "completed" });
+
+      const running = await store.queryWorkflowRuns(tenant, { status: "running" });
+      expect(running.total).toBe(2);
+      expect(running.records.map((r) => r.id).sort()).toEqual([r1.id, r3.id].sort());
+
+      const completed = await store.queryWorkflowRuns(tenant, { status: "completed" });
+      expect(completed.total).toBe(1);
+      expect(completed.records.map((r) => r.id)).toEqual([r2.id]);
+
+      const byAssistant = await store.queryWorkflowRuns(tenant, { assistantId: assistantB });
+      expect(byAssistant.total).toBe(1);
+      expect(byAssistant.records.map((r) => r.id)).toEqual([r3.id]);
+
+      const combined = await store.queryWorkflowRuns(tenant, {
+        status: "running",
+        assistantId: assistantA,
+      });
+      expect(combined.total).toBe(1);
+      expect(combined.records.map((r) => r.id)).toEqual([r1.id]);
+
+      const other = await store.queryWorkflowRuns(otherTenant, { status: "running" });
+      expect(other).toEqual({ records: [], total: 0 });
+    });
+
+    it("paginates queryWorkflowRuns while counting all matches", async () => {
+      const tenant = tenantId();
+      const runs = [await createRun(tenant), await createRun(tenant), await createRun(tenant)];
+
+      const firstPage = await store.queryWorkflowRuns(tenant, { limit: 2, offset: 0 });
+      const secondPage = await store.queryWorkflowRuns(tenant, { limit: 2, offset: 2 });
+      const all = await store.queryWorkflowRuns(tenant);
+
+      expect(firstPage.total).toBe(3);
+      expect(secondPage.total).toBe(3);
+      expect(firstPage.records).toHaveLength(2);
+      expect(secondPage.records).toHaveLength(1);
+
+      const firstIds = firstPage.records.map((r) => r.id);
+      const secondIds = secondPage.records.map((r) => r.id);
+      expect(firstIds.filter((id) => secondIds.includes(id))).toEqual([]);
+      expect([...firstIds, ...secondIds].sort()).toEqual(runs.map((r) => r.id).sort());
+
+      expect(all.total).toBe(3);
+      expect(all.records).toHaveLength(3);
+    });
+
+    it("falls back to empty metadata when the stored column is null", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+
+      await pool.query("UPDATE lattice_workflow_runs SET metadata = NULL WHERE id = $1", [run.id]);
+
+      const read = await store.getWorkflowRun(run.id);
+      expect(read!.metadata).toEqual({});
+      expect(read!.topologyEdges).toEqual(topologyEdges());
+    });
+  });
+
+  describe("run steps", () => {
+    it("creates a step and reads back input, edge data, and thread id", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+      const threadId = `thread-${randomUUID()}`;
+
+      const step = await store.createRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "tool_call",
+        stepName: "search",
+        threadId,
+        edgeFrom: "planner",
+        edgeTo: "worker",
+        edgePurpose: "delegate",
+        input: { query: "weather", nested: { limit: 3 } },
+      });
+
+      expect(step).toMatchObject({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "tool_call",
+        stepName: "search",
+        threadId,
+        edgeFrom: "planner",
+        edgeTo: "worker",
+        edgePurpose: "delegate",
+        status: "running",
+      });
+      expect(step.input).toEqual({ query: "weather", nested: { limit: 3 } });
+      expect(step.output).toBeUndefined();
+      expect(step.startedAt).toBeInstanceOf(Date);
+
+      const steps = await store.getRunSteps(run.id);
+      expect(steps.map((s) => s.id)).toEqual([step.id]);
+      expect(steps[0].input).toEqual({ query: "weather", nested: { limit: 3 } });
+    });
+
+    it("upserts by run, step type, and step name", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+      const first = await store.createRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "agent",
+        stepName: "Writer",
+        input: { brief: "original" },
+      });
+
+      const upserted = await store.upsertRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "agent",
+        stepName: "Writer",
+        threadId: "sub-thread-1",
+        input: { brief: "replacement" },
+      });
+
+      expect(upserted.id).toBe(first.id);
+      expect(upserted.threadId).toBe("sub-thread-1");
+      expect(upserted.input).toEqual({ brief: "original" });
+
+      const again = await store.upsertRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "agent",
+        stepName: "Writer",
+        threadId: "sub-thread-2",
+      });
+      expect(again.id).toBe(first.id);
+      expect(again.threadId).toBe("sub-thread-1");
+
+      const distinct = await store.upsertRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "agent",
+        stepName: "Reviewer",
+      });
+      expect(distinct.id).not.toBe(first.id);
+      expect(await store.getRunSteps(run.id)).toHaveLength(2);
+    });
+
+    it("updates a step result and returns null for an absent step", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+      const step = await store.createRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "agent",
+        stepName: "Writer",
+      });
+      const completedAt = new Date("2026-03-04T05:06:07.000Z");
+
+      const updated = await store.updateRunStep(run.id, step.id, {
+        status: "completed",
+        output: { draft: "hello", tokens: { input: 10, output: 20 } },
+        completedAt,
+        durationMs: 1234,
+      });
+
+      expect(updated).toMatchObject({
+        id: step.id,
+        status: "completed",
+        durationMs: 1234,
+      });
+      expect(updated!.output).toEqual({ draft: "hello", tokens: { input: 10, output: 20 } });
+      expect(updated!.completedAt?.toISOString()).toBe(completedAt.toISOString());
+
+      expect(
+        await store.updateRunStep(run.id, randomUUID(), { status: "completed" }),
+      ).toBeNull();
+    });
+
+    it("finds steps by type and interrupted steps", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+      const toolStep = await store.createRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "tool_call",
+        stepName: "search",
+      });
+      const agentStep = await store.createRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "agent",
+        stepName: "Writer",
+      });
+      const humanStep = await store.createRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "human_in_loop",
+        stepName: "approve",
+      });
+      await store.updateRunStep(run.id, humanStep.id, { status: "interrupted" });
+      await store.updateRunStep(run.id, agentStep.id, { status: "completed" });
+
+      const toolSteps = await store.getRunStepsByType(run.id, "tool_call");
+      expect(toolSteps.map((s) => s.id)).toEqual([toolStep.id]);
+
+      const interrupted = await store.getInterruptedSteps(run.id);
+      expect(interrupted.map((s) => s.id)).toEqual([humanStep.id]);
+      expect(interrupted[0].status).toBe("interrupted");
+    });
+
+    it("returns empty step lists for an absent run", async () => {
+      const missingRun = randomUUID();
+
+      expect(await store.getRunSteps(missingRun)).toEqual([]);
+      expect(await store.getRunStepsByType(missingRun, "agent")).toEqual([]);
+      expect(await store.getInterruptedSteps(missingRun)).toEqual([]);
+    });
+
+    it("falls back to undefined for null step input and output", async () => {
+      const tenant = tenantId();
+      const run = await createRun(tenant);
+      const step = await store.createRunStep({
+        runId: run.id,
+        tenantId: tenant,
+        stepType: "input",
+        stepName: "start",
+      });
+
+      const [read] = await store.getRunSteps(run.id);
+      expect(read.id).toBe(step.id);
+      expect(read.input).toBeUndefined();
+      expect(read.output).toBeUndefined();
+    });
+  });
+});

package/src/stores/PostgreSQLEvalStore.ts

@@ -300,8 +300,8 @@ export class PostgreSQLEvalStore implements EvalStore {
         data.name,
         data.description || null,
         data.version || null,
-        JSON.stringify(data.judgeModelConfig),
-        JSON.stringify(data.targetServerConfig),
+        JSON.stringify(data.judgeModelConfig ?? {}),
+        JSON.stringify(data.targetServerConfig ?? {}),
         data.concurrency ?? 3,
         data.reportConfig === undefined ? null : JSON.stringify(data.reportConfig),
       ]
@@ -857,7 +857,7 @@ export class PostgreSQLEvalStore implements EvalStore {
          AND run_id IN (SELECT id FROM lattice_eval_runs WHERE tenant_id = $${i + 1})
        RETURNING id, run_id, suite_name, case_id,
                  pass, score, summary, dimension_results,
-                 duration_ms, messages, logs, error, created_at`,
+                 duration_ms, messages, logs, error, interrupted, created_at`,
       vals
     );
     return rows.length ? this.mapRowToRunResult(rows[0] as unknown as Record<string, unknown>) : null;

package/src/stores/PostgreSQLSkillStore.ts

@@ -502,7 +502,7 @@ export class PostgreSQLSkillStore implements SkillStore {
       WHERE tenant_id = $1 AND sub_skills @> to_jsonb($2::text)
       ORDER BY created_at DESC
     `,
-      [tenantId, JSON.stringify([parentSkillName])]
+      [tenantId, parentSkillName]
     );
 
     return result.rows.map(this.mapRowToSkill);

package/src/stores/PostgreSQLWorkflowTrackingStore.ts

@@ -2,6 +2,7 @@
  * PostgreSQL implementation of WorkflowTrackingStore
  */
 
+import { randomUUID } from "node:crypto";
 import { Pool } from "pg";
 import type { PoolConfig } from "pg";
 import {
@@ -92,7 +93,7 @@ export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
   async createWorkflowRun(request: CreateWorkflowRunRequest): Promise<WorkflowRun> {
     await this.ensureInitialized();
     const now = new Date();
-    const id = `${request.threadId}_${Date.now()}`;
+    const id = `${request.threadId}_${Date.now()}_${randomUUID()}`;
 
     await this.pool.query(
       `INSERT INTO lattice_workflow_runs (id, tenant_id, assistant_id, thread_id, status, topology_edges, total_edges, completed_edges, metadata, started_at, created_at, updated_at)