@axiom-lattice/pg-stores 3.4.0 → 4.0.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "3.4.0",
+  "version": "4.0.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": "9.0.0",
-    "@axiom-lattice/protocols": "5.0.0"
+    "@axiom-lattice/core": "10.0.0",
+    "@axiom-lattice/protocols": "6.0.0"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/PostgreSQLChannelReplyJobStore.test.ts

@@ -19,7 +19,9 @@ function input(
     channelInstallationId: "install-1",
     replyTarget: { chatId: "chat-1" },
     inputMessageId: "input-1",
-    availableAt: Date.now() - 1_000,
+    // Well in the past: the store compares against the database clock, which
+    // can trail the Node process clock by a few milliseconds.
+    availableAt: Date.now() - 60_000,
     ...overrides,
   };
 }
@@ -194,7 +196,7 @@ maybeDescribe("PostgreSQLChannelReplyJobStore (PostgreSQL)", () => {
       store.reschedule(created.tenantId, created.runId, "owner-b", "nope", Date.now(), 1),
     ).resolves.toBeNull();
 
-    const availableAt = Date.now() - 1;
+    const availableAt = Date.now() - 60_000;
     const rescheduled = await store.reschedule(
       created.tenantId,
       created.runId,

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/migrations/eval_migrations.ts

@@ -381,6 +381,24 @@ export const addTaskIdToEvalRuns: Migration = {
   },
 };
 
+
+/** Cooperative cancellation flag for running eval runs. */
+export const addCancelRequestedToEvalRuns: Migration = {
+  version: 205,
+  name: "add_cancel_requested_to_eval_runs",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_eval_runs
+        ADD COLUMN IF NOT EXISTS cancel_requested BOOLEAN NOT NULL DEFAULT FALSE
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_eval_runs DROP COLUMN IF EXISTS cancel_requested
+    `);
+  },
+};
+
 /** All eval migrations in version order */
 export const evalMigrations: Migration[] = [
   createEvalProjectsTable,
@@ -394,4 +412,5 @@ export const evalMigrations: Migration[] = [
   addInterruptPolicyToEvalCases,
   enforceUniqueEvalProjectName,
   addTaskIdToEvalRuns,
+  addCancelRequestedToEvalRuns,
 ];

package/src/migrations/task_migration.ts

@@ -112,6 +112,22 @@ export const addTaskDependenciesGinIndex: Migration = {
   },
 };
 
+/** Index task execution threads for ownership checks on isolated worker threads. */
+export const addTaskContextThreadIndex: Migration = {
+  version: 206,
+  name: "add_task_context_thread_index",
+  up: async (client) => {
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_tasks_context_thread
+      ON lattice_tasks (tenant_id, (context->>'thread_id'))
+      WHERE context->>'thread_id' IS NOT NULL
+    `);
+  },
+  down: async (client) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_tasks_context_thread");
+  },
+};
+
 /** All task migrations in version order */
 export const taskMigrations: Migration[] = [
   createTasksTable,
@@ -119,4 +135,5 @@ export const taskMigrations: Migration[] = [
   addTaskProjectFieldsMigration,
   addFilesToTasks,
   addTaskDependenciesGinIndex,
+  addTaskContextThreadIndex,
 ];

package/src/stores/PostgreSQLAgentRunLedger.ts

@@ -274,7 +274,10 @@ export class PostgreSQLAgentRunLedger implements IAgentRunLedger {
     const inputMessageIds = input.inputMessageIds
       ? normalizeAgentRunMessageIds(input.inputMessageIds)
       : extractAgentRunInputMessageIds(input.input);
-    const availableAt = input.availableAt ?? Date.now();
+    // Default to the database clock. Mixing Node's `Date.now()` with the
+    // `available_at <= now()` claim predicate makes freshly created rows
+    // briefly unclaimable whenever the two clocks drift by a few milliseconds.
+    const availableAt = input.availableAt ?? null;
     const limit = input.maxActiveRunsForTenant;
     const client = await this.pool.connect();
     try {
@@ -333,7 +336,10 @@ export class PostgreSQLAgentRunLedger implements IAgentRunLedger {
             available_at, created_at, updated_at, task_id)
          VALUES ($1, $2, $3, $4, $5, $6, 'pending', $7, NULL, NULL, $8, 0,
                  $9::jsonb, $10::jsonb, $11::jsonb, $12::jsonb, $13,
-                 to_timestamp($14::double precision / 1000.0),
+                 COALESCE(
+                   to_timestamp($14::double precision / 1000.0),
+                   date_trunc('milliseconds', now())
+                 ),
                  date_trunc('milliseconds', now()),
                  date_trunc('milliseconds', now()), $15)
          RETURNING *`,
@@ -971,14 +977,19 @@ export class PostgreSQLAgentRunLedger implements IAgentRunLedger {
     options?: SettleExpiredRunOptions,
   ): Promise<boolean> {
     if (outcome === "pending") {
-      const availableAt = options?.availableAt ?? Date.now();
+      // Same clock rule as createPendingRun: only explicit backoff timestamps
+      // are caller-supplied; the default must come from the database.
+      const availableAt = options?.availableAt ?? null;
       const result = await this.pool.query(
         `UPDATE lattice_agent_runs
             SET status = 'pending',
                 owner = NULL,
                 lease_expires_at = NULL,
                 started_at = NULL,
-                available_at = to_timestamp($2::double precision / 1000.0),
+                available_at = COALESCE(
+                  to_timestamp($2::double precision / 1000.0),
+                  date_trunc('milliseconds', now())
+                ),
                 finished_at = NULL,
                 updated_at = now()
           WHERE run_id = $1

package/src/stores/PostgreSQLEvalStore.ts

@@ -180,6 +180,7 @@ export class PostgreSQLEvalStore implements EvalStore {
       interruptedCases: (row.interrupted_cases as number) ?? 0,
       avgScore: (row.avg_score as number) ?? 0,
       error: row.error as string | undefined,
+      cancelRequested: (row.cancel_requested as boolean) ?? false,
       holdout: row.holdout as boolean | undefined,
       envProjectId: row.env_project_id as string | undefined,
       envWorkspaceId: row.env_workspace_id as string | undefined,
@@ -610,7 +611,7 @@ export class PostgreSQLEvalStore implements EvalStore {
       `SELECT id, project_id, tenant_id,
               status, concurrency, total_cases,
               passed_cases, failed_cases, interrupted_cases, avg_score,
-              error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at
+              error, holdout, env_project_id, env_workspace_id, task_id, cancel_requested, created_at, started_at, completed_at
        FROM lattice_eval_runs
        WHERE ${conditions.join(" AND ")}
        ORDER BY created_at DESC`,
@@ -632,7 +633,7 @@ export class PostgreSQLEvalStore implements EvalStore {
       `SELECT id, project_id, tenant_id,
               status, concurrency, total_cases,
               passed_cases, failed_cases, interrupted_cases, avg_score,
-              error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at
+              error, holdout, env_project_id, env_workspace_id, task_id, cancel_requested, created_at, started_at, completed_at
        FROM lattice_eval_runs
        WHERE id = $1 AND tenant_id = $2`,
       [id, tenantId]
@@ -662,7 +663,7 @@ export class PostgreSQLEvalStore implements EvalStore {
        RETURNING id, project_id, tenant_id,
                  status, concurrency, total_cases,
                  passed_cases, failed_cases, interrupted_cases, avg_score,
-                 error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at`,
+                 error, holdout, env_project_id, env_workspace_id, task_id, cancel_requested, created_at, started_at, completed_at`,
       [
         actualId,
         projectId,
@@ -722,12 +723,24 @@ export class PostgreSQLEvalStore implements EvalStore {
        RETURNING id, project_id, tenant_id,
                  status, concurrency, total_cases,
                  passed_cases, failed_cases, interrupted_cases, avg_score,
-                 error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at`,
+                 error, holdout, env_project_id, env_workspace_id, task_id, cancel_requested, created_at, started_at, completed_at`,
       vals
     );
     return rows.length ? this.mapRowToRun(rows[0] as unknown as Record<string, unknown>) : null;
   }
 
+  /** Request cooperative cancellation for a running evaluation. */
+  async requestCancel(tenantId: string, id: string): Promise<boolean> {
+    await this.ensureInitialized();
+    const { rowCount } = await this.pool.query(
+      `UPDATE lattice_eval_runs
+       SET cancel_requested = TRUE
+       WHERE tenant_id = $1 AND id = $2 AND status = 'running'`,
+      [tenantId, id],
+    );
+    return (rowCount ?? 0) > 0;
+  }
+
   /** Delete a run and its results */
   async deleteRun(tenantId: string, id: string): Promise<boolean> {
     await this.ensureInitialized();

package/src/stores/PostgreSQLTaskStore.ts

@@ -283,6 +283,10 @@ export class PostgreSQLTaskStore implements TaskStore {
       conditions.push(`priority = $${paramIndex++}`);
       params.push(filter.priority);
     }
+    if (filter.threadId) {
+      conditions.push(`context->>'thread_id' = $${paramIndex++}`);
+      params.push(filter.threadId);
+    }
     if (filter.workspaceId) {
       conditions.push(`workspace_id = $${paramIndex++}`);
       params.push(filter.workspaceId);