@axiom-lattice/pg-stores 3.4.0 → 4.0.1

package/dist/index.mjs

@@ -4291,6 +4291,21 @@ var addTaskIdToEvalRuns = {
     `);
   }
 };
+var addCancelRequestedToEvalRuns = {
+  version: 205,
+  name: "add_cancel_requested_to_eval_runs",
+  up: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_eval_runs
+        ADD COLUMN IF NOT EXISTS cancel_requested BOOLEAN NOT NULL DEFAULT FALSE
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_eval_runs DROP COLUMN IF EXISTS cancel_requested
+    `);
+  }
+};
 var evalMigrations = [
   createEvalProjectsTable,
   createEvalSuitesTable,
@@ -4302,7 +4317,8 @@ var evalMigrations = [
   addInterruptColumnsToEval,
   addInterruptPolicyToEvalCases,
   enforceUniqueEvalProjectName,
-  addTaskIdToEvalRuns
+  addTaskIdToEvalRuns,
+  addCancelRequestedToEvalRuns
 ];
 
 // src/stores/PostgreSQLEvalStore.ts
@@ -4427,6 +4443,7 @@ var PostgreSQLEvalStore = class {
       interruptedCases: row.interrupted_cases ?? 0,
       avgScore: row.avg_score ?? 0,
       error: row.error,
+      cancelRequested: row.cancel_requested ?? false,
       holdout: row.holdout,
       envProjectId: row.env_project_id,
       envWorkspaceId: row.env_workspace_id,
@@ -4798,7 +4815,7 @@ var PostgreSQLEvalStore = class {
       `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`,
@@ -4813,7 +4830,7 @@ var PostgreSQLEvalStore = class {
       `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]
@@ -4831,7 +4848,7 @@ var PostgreSQLEvalStore = class {
        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,
@@ -4893,11 +4910,22 @@ var PostgreSQLEvalStore = class {
        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]) : null;
   }
+  /** Request cooperative cancellation for a running evaluation. */
+  async requestCancel(tenantId, id) {
+    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, id) {
     await this.ensureInitialized();
@@ -7650,12 +7678,27 @@ var addTaskDependenciesGinIndex = {
     await client.query("DROP INDEX IF EXISTS idx_lattice_tasks_dependencies_gin");
   }
 };
+var addTaskContextThreadIndex = {
+  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");
+  }
+};
 var taskMigrations = [
   createTasksTable,
   addTaskFieldsMigration,
   addTaskProjectFieldsMigration,
   addFilesToTasks,
-  addTaskDependenciesGinIndex
+  addTaskDependenciesGinIndex,
+  addTaskContextThreadIndex
 ];
 
 // src/stores/PostgreSQLTaskStore.ts
@@ -7866,6 +7909,10 @@ var PostgreSQLTaskStore = class {
       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);
@@ -9512,6 +9559,9 @@ import { Pool as Pool25 } from "pg";
 import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
 import { Document } from "@langchain/core/documents";
 import { embeddingsLatticeManager } from "@axiom-lattice/core";
+import {
+  EMBEDDING_DEFAULT_KEY
+} from "@axiom-lattice/protocols";
 function safeTableName(name) {
   if (!/^[a-z][a-z0-9_-]*$/.test(name)) {
     throw new Error(`Invalid table name: ${name}`);
@@ -9526,16 +9576,20 @@ var PGVectorStoreProvider = class {
   async create(params) {
     const { tableName, embeddingKey } = params;
     const tenantId = params.tenantId ?? "default";
+    const requestedKey = embeddingKey || EMBEDDING_DEFAULT_KEY;
     const embeddings = await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
       tenantId,
-      embeddingKey
-    ) ?? await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
+      requestedKey
+    ) ?? (requestedKey === EMBEDDING_DEFAULT_KEY ? void 0 : await embeddingsLatticeManager.getEmbeddingsClientWithTenant(
+      tenantId,
+      EMBEDDING_DEFAULT_KEY
+    )) ?? 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.`
+        `No embedding model "${embeddingKey}" registered for tenant "${tenantId}", and no tenant or deployment default. Configure one in the model service panel.`
       );
     }
     return PGVectorStore.initialize(embeddings, {
@@ -10704,6 +10758,28 @@ var addModelProviderModelKind = {
     `);
   }
 };
+var addModelProviderKindDefault = {
+  version: 207,
+  name: "add_model_provider_kind_default",
+  up: async (client) => {
+    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) => {
+    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
+    `);
+  }
+};
 
 // src/stores/PostgreSQLCapabilityBundleStore.ts
 import { randomUUID as randomUUID3 } from "crypto";
@@ -11698,6 +11774,7 @@ var PostgreSQLModelProviderStore = class {
     this.migrationManager.register(addModelProviderModelExtra);
     this.migrationManager.register(addModelProviderModelDefault);
     this.migrationManager.register(addModelProviderModelKind);
+    this.migrationManager.register(addModelProviderKindDefault);
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
         console.error("Failed to initialize PostgreSQLModelProviderStore:", error);
@@ -12036,15 +12113,22 @@ var PostgreSQLModelProviderStore = class {
       const effectiveKind = rows.map(
         (model) => model.kind ?? currentKind.get(model.id) ?? "chat"
       );
-      const namesDefault = rows.some((model) => model.isDefault === true);
-      if (namesDefault) {
+      const namedDefaultKinds = /* @__PURE__ */ new Set();
+      rows.forEach((model, index) => {
+        if (model.isDefault === true) {
+          namedDefaultKinds.add(effectiveKind[index]);
+        }
+      });
+      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.isDefault ?? (namesDefault ? false : currentIsDefault.get(model.id)) ?? false
+        (model, index) => model.isDefault ?? (namedDefaultKinds.has(effectiveKind[index]) ? false : currentIsDefault.get(model.id)) ?? false
       );
       await client.query(
         `
@@ -12741,7 +12825,7 @@ var PostgreSQLAgentRunLedger = class {
       throw new AgentRunInputCommandConflictError();
     }
     const inputMessageIds = input.inputMessageIds ? normalizeAgentRunMessageIds(input.inputMessageIds) : extractAgentRunInputMessageIds(input.input);
-    const availableAt = input.availableAt ?? Date.now();
+    const availableAt = input.availableAt ?? null;
     const limit = input.maxActiveRunsForTenant;
     const client = await this.pool.connect();
     try {
@@ -12791,7 +12875,10 @@ var PostgreSQLAgentRunLedger = class {
             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 *`,
@@ -13310,14 +13397,17 @@ var PostgreSQLAgentRunLedger = class {
   /** Atomically requeues or interrupts an expired run-queue row. */
   async settleExpiredRun(runId, outcome, options) {
     if (outcome === "pending") {
-      const availableAt = options?.availableAt ?? Date.now();
+      const availableAt = options?.availableAt ?? null;
       const result2 = 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
@@ -14507,6 +14597,7 @@ async function createPgStoreConfig(connectionString) {
   mm.register(createAgentRunRetentionMigration);
   mm.register(createAgentRunProjectTaskMigration);
   mm.register(agentRunProjectTaskRunningIndexMigration);
+  mm.register(addModelProviderKindDefault);
   await mm.migrate();
   const checkpoint = PostgresSaver.fromConnString(connectionString);
   checkpoint.setup().catch((err) => {
@@ -15157,6 +15248,7 @@ export {
   ProjectRoomMessageIdConflictError,
   ThreadMessageQueueStore,
   addAssistantTenantId,
+  addCancelRequestedToEvalRuns,
   addEnvToEvalRuns,
   addFileContentType,
   addHoldoutToEvalRuns,