@axiom-lattice/pg-stores 3.4.0 → 4.0.1

package/package.json

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

package/src/PGVectorStoreProvider.ts

@@ -9,7 +9,11 @@ import { Pool } from "pg";
 import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
 import { Document } from "@langchain/core/documents";
 import { embeddingsLatticeManager } from "@axiom-lattice/core";
-import type { VectorStoreProvider, VectorStoreCreateParams } from "@axiom-lattice/protocols";
+import {
+  EMBEDDING_DEFAULT_KEY,
+  type VectorStoreProvider,
+  type VectorStoreCreateParams,
+} from "@axiom-lattice/protocols";
 import type { DocumentInterface } from "@langchain/core/documents";
 
 function safeTableName(name: string): string {
@@ -31,15 +35,22 @@ export class PGVectorStoreProvider implements VectorStoreProvider {
   async create(params: VectorStoreCreateParams): Promise<any> {
     const { tableName, embeddingKey } = params;
     const tenantId = params.tenantId ?? "default";
+    const requestedKey = embeddingKey || EMBEDDING_DEFAULT_KEY;
 
     // 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.
+    // uses. A concrete key that misses falls back to the tenant's embedding
+    // default, then to the deployment registration.
     const embeddings =
       (await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
         tenantId,
-        embeddingKey
+        requestedKey
       )) ??
+      (requestedKey === EMBEDDING_DEFAULT_KEY
+        ? undefined
+        : await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
+            tenantId,
+            EMBEDDING_DEFAULT_KEY
+          )) ??
       (await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
         tenantId,
         "default"
@@ -48,7 +59,7 @@ export class PGVectorStoreProvider implements VectorStoreProvider {
     if (!embeddings) {
       throw new Error(
         `No embedding model "${embeddingKey}" registered for tenant "${tenantId}", ` +
-          `and no deployment default. Configure one in the model provider panel.`
+          `and no tenant or deployment default. Configure one in the model service panel.`
       );
     }
 

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

@@ -34,7 +34,9 @@ describe("agent run retention migration", () => {
         (name) => name === createAgentRunRetentionMigration.name,
       ),
     ).toHaveLength(1);
-    expect(registeredMigrationNames.at(-1)).toBe(
+    // Registration order is not the apply order — the manager sorts by version —
+    // so assert presence rather than that this is the last line registered.
+    expect(registeredMigrationNames).toContain(
       agentRunProjectTaskRunningIndexMigration.name,
     );
   });

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/createPgStoreConfig.ts

@@ -109,6 +109,7 @@ import {
   addModelProviderModelExtra,
   addModelProviderModelDefault,
   addModelProviderModelKind,
+  addModelProviderKindDefault,
 } from "./migrations/model_provider_migrations";
 import { PostgreSQLCapabilityBundleStore } from "./stores/PostgreSQLCapabilityBundleStore";
 import { PostgreSQLProjectRoomStore } from "./stores/PostgreSQLProjectRoomStore";
@@ -222,6 +223,7 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createAgentRunRetentionMigration); // v201
   mm.register(createAgentRunProjectTaskMigration); // v203
   mm.register(agentRunProjectTaskRunningIndexMigration); // v204
+  mm.register(addModelProviderKindDefault);     // v207
 
   await mm.migrate();
 

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/model_provider_migrations.ts

@@ -256,3 +256,41 @@ export const addModelProviderModelKind: Migration = {
     `);
   },
 };
+
+/**
+ * One default per tenant **per kind**.
+ *
+ * v192's index was tenant-wide, so a tenant could have a default chat model
+ * *or* a default embedding model, never both. Each kind publishes its own
+ * alias key (`default`, `default:embedding`, `default:asr`), so the uniqueness
+ * constraint has to be per kind too.
+ *
+ * Existing rows satisfy the wider index unchanged: before this, at most one
+ * row per tenant could be the default.
+ *
+ * Note the `down` migration can legitimately fail once a tenant has one
+ * default per kind — the narrower index cannot represent that state. That is a
+ * deliberate one-way step.
+ */
+export const addModelProviderKindDefault: Migration = {
+  version: 207,
+  name: 'add_model_provider_kind_default',
+  up: async (client: PoolClient) => {
+    await client.query(`
+      DROP INDEX IF EXISTS uk_lattice_model_provider_models_tenant_default
+    `);
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS uk_lattice_model_provider_models_tenant_kind_default
+      ON lattice_model_provider_models (tenant_id, kind) WHERE is_default
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(`
+      DROP INDEX IF EXISTS uk_lattice_model_provider_models_tenant_kind_default
+    `);
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS uk_lattice_model_provider_models_tenant_default
+      ON lattice_model_provider_models (tenant_id) WHERE is_default
+    `);
+  },
+};

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/PostgreSQLModelProviderStore.ts

@@ -26,6 +26,7 @@ import {
   addModelProviderModelExtra,
   addModelProviderModelDefault,
   addModelProviderModelKind,
+  addModelProviderKindDefault,
 } from '../migrations/model_provider_migrations';
 import { encrypt, safeDecrypt } from '@axiom-lattice/core';
 
@@ -148,6 +149,7 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
     this.migrationManager.register(addModelProviderModelExtra);     // v191
     this.migrationManager.register(addModelProviderModelDefault);   // v192
     this.migrationManager.register(addModelProviderModelKind);      // v193
+    this.migrationManager.register(addModelProviderKindDefault);    // v207
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -610,30 +612,42 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
         (model) => model.kind ?? currentKind.get(model.id) ?? "chat"
       );
 
-      // Whether this write is naming a default at all.
+      // Which kinds this write names a default for.
       //
-      // The unique index spans the *tenant*, while this method writes one
+      // The unique index is per *(tenant, kind)*, while this method writes one
       // provider. Two consequences:
       //
-      // 1. A swap (A was the default, B now is) would see two true rows
-      //    mid-statement — PostgreSQL checks a unique index per row, so whether
-      //    it fails would depend on which row the batch happened to process
-      //    first. Clearing tenant-wide *before* the upsert makes order
+      // 1. A swap (A was kind K's default, B now is) would see two true rows of
+      //    kind K mid-statement — PostgreSQL checks a unique index per row, so
+      //    whether it fails would depend on which row the batch happened to
+      //    process first. Clearing that kind before the upsert makes order
       //    irrelevant.
-      // 2. Once cleared, an omitted flag must resolve to `false`, not "keep":
-      //    keeping it would restore the old default alongside the new one.
-      const namesDefault = rows.some((model) => model.isDefault === true);
-      if (namesDefault) {
+      // 2. Once cleared, an omitted flag must resolve to `false` for the kinds
+      //    this write names — keeping it would restore the old default
+      //    alongside the new one. Other kinds keep their stored flag, so a chat
+      //    write cannot drop the tenant's embedding default.
+      const namedDefaultKinds = new Set<ModelKind>();
+      rows.forEach((model, index) => {
+        if (model.isDefault === true) {
+          namedDefaultKinds.add(effectiveKind[index] as ModelKind);
+        }
+      });
+
+      for (const kind of namedDefaultKinds) {
         await client.query(
-          `UPDATE lattice_model_provider_models SET is_default = FALSE WHERE tenant_id = $1`,
-          [tenantId]
+          `UPDATE lattice_model_provider_models
+              SET is_default = FALSE
+            WHERE tenant_id = $1 AND kind = $2 AND is_default`,
+          [tenantId, kind]
         );
       }
 
       const effectiveIsDefault = rows.map(
-        (model) =>
+        (model, index) =>
           model.isDefault ??
-          (namesDefault ? false : currentIsDefault.get(model.id)) ??
+          (namedDefaultKinds.has(effectiveKind[index] as ModelKind)
+            ? false
+            : currentIsDefault.get(model.id)) ??
           false
       );
 

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);