@axiom-lattice/pg-stores 3.3.13 → 4.0.0

package/src/migrations/agent_run_retention_migration.ts

@@ -0,0 +1,25 @@
+import type { Migration } from "./migration";
+
+/**
+ * Indexes terminal run completion times for retention cleanup.
+ *
+ * @remarks Only terminal rows are ever eligible for deletion, so the partial
+ * index keeps the bounded oldest-first cleanup scan cheap without indexing
+ * pending or running queue work.
+ */
+export const createAgentRunRetentionMigration: Migration = {
+  version: 201,
+  name: "add_agent_run_retention_index",
+  up: async (client) => {
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_agent_runs_terminal_finished_at
+        ON lattice_agent_runs (finished_at)
+        WHERE status IN ('success', 'error', 'interrupted', 'cancelled', 'timeout')
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      DROP INDEX IF EXISTS idx_lattice_agent_runs_terminal_finished_at
+    `);
+  },
+};

package/src/migrations/agent_run_rollback_failure_migration.ts

@@ -0,0 +1,28 @@
+import type { Migration } from "./migration";
+
+/**
+ * Adds bounded rollback failure tracking to run records.
+ *
+ * @remarks Recovery increments `rollback_attempts` and stores
+ * `rollback_error`. When the attempt cap is reached the row settles `error`
+ * with `rollback_requested=false`, so join/list converge instead of retrying a
+ * permanently failing checkpoint rewind forever.
+ */
+export const createAgentRunRollbackFailureMigration: Migration = {
+  version: 199,
+  name: "add_agent_run_rollback_failure",
+  up: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ADD COLUMN IF NOT EXISTS rollback_attempts INTEGER NOT NULL DEFAULT 0,
+        ADD COLUMN IF NOT EXISTS rollback_error TEXT
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        DROP COLUMN IF EXISTS rollback_error,
+        DROP COLUMN IF EXISTS rollback_attempts
+    `);
+  },
+};

package/src/migrations/agent_run_running_index_source_migration.ts

@@ -0,0 +1,31 @@
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Scope the one-running-row-per-thread index to the durable run queue.
+ *
+ * @remarks Legacy message-queue ledger rows also use `status='running'`. A
+ * crashed legacy owner can leave such a row behind; it must not make a
+ * `run_api` pending row unclaimable after the `thread-run` lease expires.
+ * Real execution exclusion remains the shared `thread-run` lease.
+ */
+export const agentRunRunningIndexSourceMigration: Migration = {
+  version: 202,
+  name: "scope_agent_run_running_index_to_run_api",
+  up: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS uq_lattice_agent_runs_one_running");
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS uq_lattice_agent_runs_one_running
+        ON lattice_agent_runs (tenant_id, thread_id)
+        WHERE status = 'running' AND source = 'run_api'
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS uq_lattice_agent_runs_one_running");
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS uq_lattice_agent_runs_one_running
+        ON lattice_agent_runs (tenant_id, thread_id)
+        WHERE status = 'running'
+    `);
+  },
+};

package/src/migrations/agent_run_task_queue_migration.ts

@@ -0,0 +1,65 @@
+import type { Migration } from "./migration";
+
+/**
+ * Expands `lattice_agent_runs` into a durable run task queue.
+ *
+ * @remarks
+ * - `source` separates new run-queue work (`run_api`) from the existing
+ *   observational ledger path (`legacy`, the default for old/inserted rows).
+ * - Pending rows have no owner, lease or start time, so those columns become
+ *   nullable.
+ * - Payload columns are JSONB; `available_at` is the earliest claim time.
+ * - The partial unique running index is intentionally left untouched.
+ */
+export const createAgentRunTaskQueueMigration: Migration = {
+  version: 196,
+  name: "expand_lattice_agent_runs_task_queue",
+  up: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ALTER COLUMN owner DROP NOT NULL,
+        ALTER COLUMN lease_expires_at DROP NOT NULL,
+        ALTER COLUMN started_at DROP NOT NULL,
+        ALTER COLUMN started_at DROP DEFAULT
+    `);
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ADD COLUMN IF NOT EXISTS source VARCHAR(32) NOT NULL DEFAULT 'legacy',
+        ADD COLUMN IF NOT EXISTS input JSONB,
+        ADD COLUMN IF NOT EXISTS command JSONB,
+        ADD COLUMN IF NOT EXISTS config JSONB,
+        ADD COLUMN IF NOT EXISTS metadata JSONB,
+        ADD COLUMN IF NOT EXISTS base_checkpoint_id TEXT,
+        ADD COLUMN IF NOT EXISTS available_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+        ADD COLUMN IF NOT EXISTS rollback_requested BOOLEAN NOT NULL DEFAULT FALSE
+    `);
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ALTER COLUMN source SET DEFAULT 'legacy'
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_agent_runs_run_queue
+        ON lattice_agent_runs (source, status, available_at, created_at)
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      DROP INDEX IF EXISTS idx_lattice_agent_runs_run_queue
+    `);
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        DROP COLUMN IF EXISTS rollback_requested,
+        DROP COLUMN IF EXISTS available_at,
+        DROP COLUMN IF EXISTS base_checkpoint_id,
+        DROP COLUMN IF EXISTS metadata,
+        DROP COLUMN IF EXISTS config,
+        DROP COLUMN IF EXISTS command,
+        DROP COLUMN IF EXISTS input,
+        DROP COLUMN IF EXISTS source
+    `);
+    await client.query(`
+      ALTER TABLE lattice_agent_runs
+        ALTER COLUMN started_at SET DEFAULT now()
+    `);
+  },
+};

package/src/migrations/channel_reply_job_migration.ts

@@ -0,0 +1,55 @@
+import type { PoolClient } from "pg";
+import type { Migration } from "./migration";
+
+const DB_CLOCK_MS =
+  "FLOOR(EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::bigint";
+
+/**
+ * Durable channel reply jobs used by the cross-replica reply dispatcher.
+ *
+ * @remarks Lease-based claims make the row the single source of truth for
+ * "which replica is delivering this run's channel reply". Timestamps are epoch
+ * milliseconds from the database clock.
+ */
+export const createChannelReplyJobsTable: Migration = {
+  version: 200,
+  name: "create_lattice_channel_reply_jobs",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_channel_reply_jobs (
+        run_id TEXT PRIMARY KEY,
+        tenant_id VARCHAR(255) NOT NULL,
+        assistant_id TEXT NOT NULL,
+        thread_id TEXT NOT NULL,
+        adapter_channel TEXT NOT NULL,
+        channel_installation_id TEXT NOT NULL,
+        reply_target JSONB NOT NULL,
+        input_message_id TEXT NOT NULL,
+        stream_cursor TEXT NULL,
+        status TEXT NOT NULL,
+        attempts INT NOT NULL DEFAULT 0,
+        available_at BIGINT NOT NULL,
+        lease_owner TEXT NULL,
+        lease_expires_at BIGINT NULL,
+        last_error TEXT NULL,
+        output_message_id TEXT NULL,
+        created_at BIGINT NOT NULL DEFAULT (${DB_CLOCK_MS}),
+        updated_at BIGINT NOT NULL DEFAULT (${DB_CLOCK_MS}),
+        delivered_at BIGINT NULL
+      )
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_channel_reply_jobs_due
+      ON lattice_channel_reply_jobs (status, available_at)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_channel_reply_jobs_tenant_thread
+      ON lattice_channel_reply_jobs (tenant_id, thread_id)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_channel_reply_jobs_tenant_thread");
+    await client.query("DROP INDEX IF EXISTS idx_lattice_channel_reply_jobs_due");
+    await client.query("DROP TABLE IF EXISTS lattice_channel_reply_jobs");
+  },
+};

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

@@ -0,0 +1,50 @@
+/**
+ * Namespaced KV store migrations.
+ */
+
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Migration: create the tenant-scoped namespaced KV store table.
+ *
+ * GIN indexes support namespace prefix and JSONB containment filters; the
+ * partial expiry index supports the TTL sweeper.
+ */
+export const createStoreItemsTable: Migration = {
+  version: 195,
+  name: "create_lattice_store_items_table",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_store_items (
+        tenant_id TEXT NOT NULL,
+        namespace TEXT[] NOT NULL,
+        key TEXT NOT NULL,
+        value JSONB NOT NULL,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        expires_at TIMESTAMPTZ,
+        PRIMARY KEY (tenant_id, namespace, key)
+      )
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_store_items_namespace
+      ON lattice_store_items USING GIN (namespace)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_store_items_value
+      ON lattice_store_items USING GIN (value)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_store_items_expires_at
+      ON lattice_store_items(expires_at)
+      WHERE expires_at IS NOT NULL
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_store_items_expires_at");
+    await client.query("DROP INDEX IF EXISTS idx_lattice_store_items_value");
+    await client.query("DROP INDEX IF EXISTS idx_lattice_store_items_namespace");
+    await client.query("DROP TABLE IF EXISTS lattice_store_items");
+  },
+};

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

@@ -40,3 +40,28 @@ export const createThreadsTable: Migration = {
     await client.query("DROP TABLE IF EXISTS lattice_threads");
   },
 };
+
+/**
+ * Migration: Add per-thread expiry (`expires_at`) for the TTL sweeper.
+ *
+ * A partial index keeps expiry scans cheap while leaving non-TTL threads out.
+ */
+export const addThreadExpiresAt: Migration = {
+  version: 194,
+  name: "add_thread_expires_at",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_threads
+      ADD COLUMN IF NOT EXISTS expires_at TIMESTAMPTZ
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_threads_expires_at
+      ON lattice_threads(expires_at)
+      WHERE expires_at IS NOT NULL
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_threads_expires_at");
+    await client.query("ALTER TABLE lattice_threads DROP COLUMN IF EXISTS expires_at");
+  },
+};