@axiom-lattice/pg-stores 3.4.0 → 4.0.0

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);
@@ -12741,7 +12788,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 +12838,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 +13360,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
@@ -15157,6 +15210,7 @@ export {
   ProjectRoomMessageIdConflictError,
   ThreadMessageQueueStore,
   addAssistantTenantId,
+  addCancelRequestedToEvalRuns,
   addEnvToEvalRuns,
   addFileContentType,
   addHoldoutToEvalRuns,