@axiom-lattice/pg-stores 1.0.89 → 1.0.90

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "1.0.89",
+  "version": "1.0.90",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -25,14 +25,17 @@
     "@langchain/core": "1.1.30",
     "pg": "^8.16.3",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "2.1.98",
-    "@axiom-lattice/protocols": "2.1.50"
+    "@axiom-lattice/core": "2.1.99",
+    "@axiom-lattice/protocols": "2.1.51"
   },
   "devDependencies": {
+    "@types/jest": "^29.5.14",
     "@types/node": "^20.11.24",
     "@types/pg": "^8.11.10",
     "@types/uuid": "^9.0.8",
+    "jest": "^29.7.0",
     "rimraf": "^5.0.5",
+    "ts-jest": "^29.4.0",
     "tsup": "^8.0.1",
     "typescript": "^5.4.2"
   },

package/src/__tests__/PostgreSQLWorkflowTrackingStore.test.ts

@@ -4,7 +4,7 @@
  * These are pure-function tests — no database required.
  * They verify that corrupted JSON in stored columns won't crash the API.
  */
-import { safeParse, mapRowToWorkflowRun, mapRowToRunStep } from "../stores/PostgreSQLWorkflowTrackingStore";
+import { safeParse, mapRowToWorkflowRun, mapRowToRunStep, PostgreSQLWorkflowTrackingStore } from "../stores/PostgreSQLWorkflowTrackingStore";
 
 // ─── safeParse ──────────────────────────────────────────────────────────────
 
@@ -231,3 +231,73 @@ describe("mapRowToRunStep", () => {
     expect(result.output).toEqual({ deep: { nested: { arr: [1, { x: "y" }] } } });
   });
 });
+
+// ─── queryWorkflowRuns ──────────────────────────────────────────────────────
+
+describe("queryWorkflowRuns", () => {
+  const runRow = {
+    id: "run-1",
+    tenant_id: "t1",
+    assistant_id: "a1",
+    thread_id: "th1",
+    status: "interrupted",
+    topology_edges: [],
+    total_edges: 0,
+    completed_edges: 0,
+    error_message: null,
+    metadata: {},
+    started_at: new Date("2026-01-02"),
+    completed_at: null,
+    created_at: new Date("2026-01-02"),
+    updated_at: new Date("2026-01-02"),
+  };
+
+  function makeStore(queryImpl: (sql: string, params: unknown[]) => { rows: unknown[] }) {
+    const query = jest.fn(async (sql: unknown, params: unknown) =>
+      queryImpl(sql as string, params as unknown[]));
+    const store = new PostgreSQLWorkflowTrackingStore({ pool: { query } as never });
+    return { store, query };
+  }
+
+  it("filters by status, paginates, and parses COUNT string to number", async () => {
+    const { store, query } = makeStore((sql) =>
+      sql.includes("COUNT(*)") ? { rows: [{ count: "7" }] } : { rows: [runRow] });
+
+    const result = await store.queryWorkflowRuns("t1", { status: "interrupted", limit: 20, offset: 40 });
+
+    const [countSql, countParams] = query.mock.calls[0] as [string, unknown[]];
+    expect(countSql).toBe("SELECT COUNT(*) AS count FROM lattice_workflow_runs WHERE tenant_id = $1 AND status = $2");
+    expect(countParams).toEqual(["t1", "interrupted"]);
+
+    const [selectSql, selectParams] = query.mock.calls[1] as [string, unknown[]];
+    expect(selectSql).toBe("SELECT * FROM lattice_workflow_runs WHERE tenant_id = $1 AND status = $2 ORDER BY updated_at DESC, id DESC LIMIT $3 OFFSET $4");
+    expect(selectParams).toEqual(["t1", "interrupted", 20, 40]);
+
+    expect(result.total).toBe(7);
+    expect(result.records).toHaveLength(1);
+    expect(result.records[0].id).toBe("run-1");
+  });
+
+  it("omits LIMIT/OFFSET when not paginating", async () => {
+    const { store, query } = makeStore((sql) =>
+      sql.includes("COUNT(*)") ? { rows: [{ count: "1" }] } : { rows: [runRow] });
+
+    await store.queryWorkflowRuns("t1");
+
+    const [selectSql, selectParams] = query.mock.calls[1] as [string, unknown[]];
+    expect(selectSql).toBe("SELECT * FROM lattice_workflow_runs WHERE tenant_id = $1 ORDER BY updated_at DESC, id DESC");
+    expect(selectParams).toEqual(["t1"]);
+    expect(selectSql).not.toContain("LIMIT");
+  });
+
+  it("adds assistant_id condition when assistantId is given", async () => {
+    const { store, query } = makeStore((sql) =>
+      sql.includes("COUNT(*)") ? { rows: [{ count: "0" }] } : { rows: [] });
+
+    await store.queryWorkflowRuns("t1", { assistantId: "a9", limit: 5 });
+
+    const [selectSql, selectParams] = query.mock.calls[1] as [string, unknown[]];
+    expect(selectSql).toBe("SELECT * FROM lattice_workflow_runs WHERE tenant_id = $1 AND assistant_id = $2 ORDER BY updated_at DESC, id DESC LIMIT $3");
+    expect(selectParams).toEqual(["t1", "a9", 5]);
+  });
+});

package/src/createPgStoreConfig.ts

@@ -59,7 +59,7 @@ import { addPriorityAndCommandColumns } from "./migrations/add_priority_command_
 import { addCustomRunConfigColumn } from "./migrations/add_custom_run_config_column";
 import { alterMessageQueueIdColumn } from "./migrations/alter_message_queue_id_column";
 import { addWorkspaceProjectToQueue } from "./migrations/add_workspace_project_to_queue";
-import { createWorkflowTrackingTables, addStepThreadId } from "./migrations/workflow_tracking_migrations";
+import { createWorkflowTrackingTables, addStepThreadId, addWorkflowRunsTenantStatusUpdatedIndex } from "./migrations/workflow_tracking_migrations";
 import { evalMigrations } from "./migrations/eval_migrations";
 import { createA2AApiKeysTable } from "./migrations/a2a_api_key_migration";
 import { createTasksTable } from "./migrations/task_migration";
@@ -119,6 +119,7 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(addFileContentType);               // v136
   mm.register(createCollectionsTable);           // v137
   mm.register(createConnectionConfigsTable);      // v160
+  mm.register(addWorkflowRunsTenantStatusUpdatedIndex); // v161
 
   await mm.migrate();
 

package/src/migrations/workflow_tracking_migrations.ts

@@ -96,3 +96,21 @@ export const addStepThreadId: Migration = {
     `);
   },
 };
+
+export const addWorkflowRunsTenantStatusUpdatedIndex: Migration = {
+  name: "add_workflow_runs_tenant_status_updated_index",
+  version: 161,
+
+  up: async (client) => {
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_workflow_runs_tenant_status
+        ON lattice_workflow_runs (tenant_id, status, updated_at DESC);
+    `);
+  },
+
+  down: async (client) => {
+    await client.query(`
+      DROP INDEX IF EXISTS idx_workflow_runs_tenant_status;
+    `);
+  },
+};

package/src/stores/PostgreSQLWorkflowTrackingStore.ts

@@ -14,9 +14,11 @@ import {
   UpdateRunStepRequest,
   StepType,
   WorkflowRunStatus,
+  QueryWorkflowRunsOptions,
+  QueryWorkflowRunsResult,
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
-import { createWorkflowTrackingTables, addStepThreadId } from "../migrations/workflow_tracking_migrations";
+import { createWorkflowTrackingTables, addStepThreadId, addWorkflowRunsTenantStatusUpdatedIndex } from "../migrations/workflow_tracking_migrations";
 
 export interface PostgreSQLWorkflowTrackingStoreOptions {
   pool?: Pool;
@@ -51,6 +53,7 @@ export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createWorkflowTrackingTables);
     this.migrationManager.register(addStepThreadId);
+    this.migrationManager.register(addWorkflowRunsTenantStatusUpdatedIndex);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -180,6 +183,41 @@ export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
     return result.rows.map(mapRowToWorkflowRun);
   }
 
+  async queryWorkflowRuns(tenantId: string, options: QueryWorkflowRunsOptions = {}): Promise<QueryWorkflowRunsResult> {
+    await this.ensureInitialized();
+    const conditions: string[] = ["tenant_id = $1"];
+    const whereParams: unknown[] = [tenantId];
+    if (options.status) {
+      whereParams.push(options.status);
+      conditions.push(`status = $${whereParams.length}`);
+    }
+    if (options.assistantId) {
+      whereParams.push(options.assistantId);
+      conditions.push(`assistant_id = $${whereParams.length}`);
+    }
+    const where = conditions.join(" AND ");
+
+    const countResult = await this.pool.query(
+      `SELECT COUNT(*) AS count FROM lattice_workflow_runs WHERE ${where}`,
+      whereParams
+    );
+    const total = parseInt(countResult.rows[0].count, 10);
+
+    const selectParams: unknown[] = [...whereParams];
+    let sql = `SELECT * FROM lattice_workflow_runs WHERE ${where} ORDER BY updated_at DESC, id DESC`;
+    if (options.limit !== undefined) {
+      selectParams.push(options.limit);
+      sql += ` LIMIT $${selectParams.length}`;
+    }
+    if (options.offset !== undefined) {
+      selectParams.push(options.offset);
+      sql += ` OFFSET $${selectParams.length}`;
+    }
+
+    const result = await this.pool.query(sql, selectParams);
+    return { records: result.rows.map(mapRowToWorkflowRun), total };
+  }
+
   async createRunStep(request: CreateRunStepRequest): Promise<RunStep> {
     await this.ensureInitialized();
     const now = new Date();