@axiom-lattice/pg-stores 3.2.1 → 3.3.0

package/src/migrations/add_queue_claim_lease.ts

@@ -0,0 +1,35 @@
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Migration: add atomic-claim lease columns to the thread message queue.
+ *
+ * @remarks Version 186 follows the current ceiling (185, distributed leases).
+ * `locked_by` identifies the claiming replica; `lease_expires_at` lets another
+ * replica reclaim an abandoned claim.
+ */
+export const addQueueClaimLease: Migration = {
+  version: 186,
+  name: "add_queue_claim_lease_columns",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_thread_message_queue
+        ADD COLUMN IF NOT EXISTS locked_by VARCHAR(255),
+        ADD COLUMN IF NOT EXISTS locked_at TIMESTAMPTZ,
+        ADD COLUMN IF NOT EXISTS lease_expires_at TIMESTAMPTZ
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_thread_message_queue_lease
+      ON lattice_thread_message_queue(status, lease_expires_at)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_thread_message_queue_lease");
+    await client.query(`
+      ALTER TABLE lattice_thread_message_queue
+        DROP COLUMN IF EXISTS lease_expires_at,
+        DROP COLUMN IF EXISTS locked_at,
+        DROP COLUMN IF EXISTS locked_by
+    `);
+  },
+};

package/src/migrations/distributed_lease_migration.ts

@@ -0,0 +1,32 @@
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Migration: create the cross-replica distributed lease table.
+ *
+ * @remarks Version 185 follows the current ceiling (184). The table is generic:
+ * later plans key task-run ownership, project mutation locks, project-task
+ * machine locks, and leader election on it.
+ */
+export const createDistributedLeases: Migration = {
+  version: 185,
+  name: "create_lattice_distributed_leases",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_leases (
+        key TEXT PRIMARY KEY,
+        holder TEXT NOT NULL,
+        acquired_at TIMESTAMPTZ NOT NULL DEFAULT now(),
+        expires_at TIMESTAMPTZ NOT NULL
+      )
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_leases_expires
+      ON lattice_leases(expires_at)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_leases_expires");
+    await client.query("DROP TABLE IF EXISTS lattice_leases");
+  },
+};

package/src/migrations/outbox_migration.ts

@@ -0,0 +1,30 @@
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Migration: create the durable outbox fan-out log.
+ *
+ * @remarks Version 187 follows the current ceiling (186). The table is an
+ * append-only, bare-serial-ordered event log; replicas poll it with a cursor to
+ * deliver notifications to their process-local listeners.
+ */
+export const createOutboxTable: Migration = {
+  version: 187,
+  name: "create_lattice_outbox",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_outbox (
+        id BIGSERIAL PRIMARY KEY,
+        topic TEXT NOT NULL,
+        payload JSONB NOT NULL,
+        origin_replica_id TEXT NOT NULL,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT now()
+      )
+    `);
+    await client.query(`CREATE INDEX IF NOT EXISTS idx_lattice_outbox_id ON lattice_outbox(id)`);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_outbox_id");
+    await client.query("DROP TABLE IF EXISTS lattice_outbox");
+  },
+};

package/src/migrations/queue_migrations.ts

@@ -0,0 +1,28 @@
+import type { MigrationManager } from "./migration";
+import { createThreadMessageQueueTable } from "./thread_message_queue_migrations";
+import { addPriorityAndCommandColumns } from "./add_priority_command_columns";
+import { addCustomRunConfigColumn } from "./add_custom_run_config_column";
+import { alterMessageQueueIdColumn } from "./alter_message_queue_id_column";
+import { addWorkspaceProjectToQueue } from "./add_workspace_project_to_queue";
+import { addTrustedRunContextColumn } from "./add_trusted_run_context_column";
+import { addQueueClaimLease } from "./add_queue_claim_lease";
+
+/**
+ * Single source of truth for the thread-message-queue schema migrations.
+ *
+ * @param manager Migration manager to register the queue migrations on.
+ * @remarks
+ * Both the own-pool `ThreadMessageQueueStore` constructor and the shared-pool
+ * `createPgStoreConfig` path must register the SAME set. Keeping the list here
+ * prevents the two sites from drifting (a missing column on one path once
+ * slipped through and only surfaced in a real two-process run).
+ */
+export function registerThreadMessageQueueMigrations(manager: MigrationManager): void {
+  manager.register(createThreadMessageQueueTable); // v100
+  manager.register(addPriorityAndCommandColumns); // v101
+  manager.register(alterMessageQueueIdColumn); // v102
+  manager.register(addCustomRunConfigColumn); // v114
+  manager.register(addWorkspaceProjectToQueue); // v131
+  manager.register(addTrustedRunContextColumn); // v173
+  manager.register(addQueueClaimLease); // v186
+}

package/src/stores/PostgreSQLLeaseStore.ts

@@ -0,0 +1,126 @@
+import { Pool, type PoolConfig } from "pg";
+import type { LeaseRecord, LeaseStore } from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createDistributedLeases } from "../migrations/distributed_lease_migration";
+
+export interface PostgreSQLLeaseStoreOptions {
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+/**
+ * PostgreSQL {@link LeaseStore}: cross-replica exclusion with self-expiring rows.
+ *
+ * @remarks `tryAcquire` is a single atomic upsert guarded by `expires_at`; no
+ * external lock manager is required. Uses the shared `lattice_leases` table.
+ */
+export class PostgreSQLLeaseStore implements LeaseStore {
+  private pool: Pool;
+  private ownsPool = true;
+  private initialized = false;
+  private initPromise: Promise<void> | null = null;
+  private migrationManager!: MigrationManager;
+
+  /** @param options An externally managed pool, or a pool config/URL. */
+  constructor(options: PostgreSQLLeaseStoreOptions) {
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
+    } else {
+      throw new Error("Either pool or poolConfig must be provided");
+    }
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createDistributedLeases);
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLLeaseStore:", error);
+        throw error;
+      });
+    }
+  }
+
+  /** Runs pending migrations once. */
+  async initialize(): Promise<void> {
+    if (this.initialized) return;
+    if (this.initPromise) return this.initPromise;
+    this.initPromise = (async () => {
+      try {
+        await this.migrationManager.migrate();
+        this.initialized = true;
+      } finally {
+        this.initPromise = null;
+      }
+    })();
+    return this.initPromise;
+  }
+
+  /** Ends the pool when this store owns it. */
+  async dispose(): Promise<void> {
+    if (this.ownsPool) await this.pool.end();
+  }
+
+  /** Atomically acquires a free/expired key (or re-acquires the caller's own). */
+  async tryAcquire(key: string, holder: string, ttlMs: number): Promise<boolean> {
+    const result = await this.pool.query(
+      `INSERT INTO lattice_leases (key, holder, acquired_at, expires_at)
+       VALUES ($1, $2, now(), now() + ($3::bigint || ' milliseconds')::interval)
+       ON CONFLICT (key) DO UPDATE
+         SET holder = EXCLUDED.holder,
+             acquired_at = CASE WHEN lattice_leases.holder = EXCLUDED.holder
+                                THEN lattice_leases.acquired_at ELSE now() END,
+             expires_at = EXCLUDED.expires_at
+         WHERE lattice_leases.expires_at <= now()
+            OR lattice_leases.holder = EXCLUDED.holder
+       RETURNING holder`,
+      [key, holder, ttlMs],
+    );
+    return result.rowCount === 1 && result.rows[0].holder === holder;
+  }
+
+  /** Extends a lease only while `holder` still owns it. */
+  async renew(key: string, holder: string, ttlMs: number): Promise<boolean> {
+    const result = await this.pool.query(
+      `UPDATE lattice_leases
+          SET expires_at = now() + ($3::bigint || ' milliseconds')::interval
+        WHERE key = $1 AND holder = $2
+        RETURNING key`,
+      [key, holder, ttlMs],
+    );
+    return (result.rowCount ?? 0) > 0;
+  }
+
+  /** Deletes a lease only when `holder` owns it. */
+  async release(key: string, holder: string): Promise<boolean> {
+    const result = await this.pool.query(
+      `DELETE FROM lattice_leases WHERE key = $1 AND holder = $2 RETURNING key`,
+      [key, holder],
+    );
+    return (result.rowCount ?? 0) > 0;
+  }
+
+  /** Returns the live lease, or `null` when free/expired/absent. */
+  async get(key: string): Promise<LeaseRecord | null> {
+    const result = await this.pool.query(
+      `SELECT key, holder, acquired_at, expires_at
+         FROM lattice_leases
+        WHERE key = $1 AND expires_at > now()`,
+      [key],
+    );
+    if (result.rowCount === 0) return null;
+    const row = result.rows[0];
+    return {
+      key: row.key,
+      holder: row.holder,
+      acquiredAt: new Date(row.acquired_at).toISOString(),
+      expiresAt: new Date(row.expires_at).toISOString(),
+    };
+  }
+}

package/src/stores/PostgreSQLOutboxStore.ts

@@ -0,0 +1,136 @@
+import { Pool, type PoolConfig } from "pg";
+import type { OutboxEvent, OutboxStore } from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createOutboxTable } from "../migrations/outbox_migration";
+
+export interface PostgreSQLOutboxStoreOptions {
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+/**
+ * PostgreSQL {@link OutboxStore}: an append-only, cross-replica notification log.
+ *
+ * @remarks Ids come from a `BIGSERIAL`, giving a total order across replicas.
+ * Callers poll with a cursor (`listAfter`) and skip their own `originReplicaId`;
+ * delivery is therefore at-least-once. Uses the shared `lattice_outbox` table.
+ */
+export class PostgreSQLOutboxStore implements OutboxStore {
+  private pool: Pool;
+  private ownsPool = true;
+  private initialized = false;
+  private initPromise: Promise<void> | null = null;
+  private migrationManager!: MigrationManager;
+
+  /** @param options An externally managed pool, or a pool config/URL. */
+  constructor(options: PostgreSQLOutboxStoreOptions) {
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
+    } else {
+      throw new Error("Either pool or poolConfig must be provided");
+    }
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createOutboxTable);
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLOutboxStore:", error);
+        throw error;
+      });
+    }
+  }
+
+  /** Runs pending migrations once. */
+  async initialize(): Promise<void> {
+    if (this.initialized) return;
+    if (this.initPromise) return this.initPromise;
+    this.initPromise = (async () => {
+      try {
+        await this.migrationManager.migrate();
+        this.initialized = true;
+      } finally {
+        this.initPromise = null;
+      }
+    })();
+    return this.initPromise;
+  }
+
+  /** Ends the pool when this store owns it. */
+  async dispose(): Promise<void> {
+    if (this.ownsPool) await this.pool.end();
+  }
+
+  /** Appends one event and returns it with its assigned id. */
+  async append(input: { topic: string; payload: unknown; originReplicaId: string }): Promise<OutboxEvent> {
+    const result = await this.pool.query(
+      `INSERT INTO lattice_outbox (topic, payload, origin_replica_id)
+       VALUES ($1, $2::jsonb, $3)
+       RETURNING id, topic, payload, origin_replica_id, created_at`,
+      [input.topic, JSON.stringify(input.payload), input.originReplicaId],
+    );
+    return this.rowToEvent(result.rows[0]);
+  }
+
+  /** Lists up to `limit` events with `id > cursor`, ascending. */
+  async listAfter(cursor: string | null, limit: number): Promise<OutboxEvent[]> {
+    const result = await this.pool.query(
+      `SELECT id, topic, payload, origin_replica_id, created_at
+         FROM lattice_outbox
+        WHERE id > $1
+        ORDER BY id ASC
+        LIMIT $2`,
+      [cursor === null ? 0 : Number(cursor), limit],
+    );
+    return result.rows.map((row) => this.rowToEvent(row));
+  }
+
+  /** Returns the newest id, or `null` when the log is empty. */
+  async latestId(): Promise<string | null> {
+    const result = await this.pool.query(
+      `SELECT id FROM lattice_outbox ORDER BY id DESC LIMIT 1`,
+    );
+    if (result.rowCount === 0) return null;
+    return String(result.rows[0].id);
+  }
+
+  /**
+   * Deletes events older than `olderThanMs` (evaluated against the database clock).
+   *
+   * @param olderThanMs Age threshold in milliseconds.
+   * @returns The number of rows deleted.
+   */
+  async pruneOlderThan(olderThanMs: number): Promise<number> {
+    const result = await this.pool.query(
+      `DELETE FROM lattice_outbox
+        WHERE created_at < now() - ($1::bigint || ' milliseconds')::interval
+        RETURNING id`,
+      [olderThanMs],
+    );
+    return result.rowCount ?? 0;
+  }
+
+  /** Maps a `lattice_outbox` row to an {@link OutboxEvent}. */
+  private rowToEvent(row: {
+    id: string | number;
+    topic: string;
+    payload: unknown;
+    origin_replica_id: string;
+    created_at: Date | string;
+  }): OutboxEvent {
+    return {
+      id: String(row.id),
+      topic: row.topic,
+      payload: row.payload,
+      originReplicaId: row.origin_replica_id,
+      createdAt: new Date(row.created_at).toISOString(),
+    };
+  }
+}

package/src/stores/PostgreSQLScheduleStorage.ts

@@ -339,6 +339,60 @@ export class PostgreSQLScheduleStorage implements ScheduleStorage {
     );
   }
 
+  /**
+   * Atomically claims one due fire by transitioning the task `pending -> running`.
+   *
+   * Only one caller wins the `UPDATE ... WHERE status = 'pending'`; all others
+   * observe zero affected rows and must skip execution.
+   *
+   * @param taskId Task to claim.
+   * @returns True when this caller won the claim (task was `pending`); false if
+   *   another replica already claimed it or it is not pending.
+   */
+  async claimScheduledRun(taskId: string): Promise<boolean> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `UPDATE lattice_scheduled_tasks
+       SET status = $1, updated_at = $2
+       WHERE task_id = $3 AND status = $4
+       RETURNING task_id`,
+      [
+        ScheduledTaskStatus.RUNNING,
+        new Date(),
+        taskId,
+        ScheduledTaskStatus.PENDING,
+      ]
+    );
+
+    return (result.rowCount ?? 0) > 0;
+  }
+
+  /**
+   * Requeues `running` tasks whose `updated_at` is older than `olderThanMs` back
+   * to `pending` (recovery after a crash that left a claim unsettled).
+   *
+   * @param olderThanMs Age threshold in milliseconds.
+   * @returns The number of tasks requeued.
+   */
+  async requeueStaleRunning(olderThanMs: number): Promise<number> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `UPDATE lattice_scheduled_tasks
+       SET status = $1, updated_at = now()
+       WHERE status = $2 AND updated_at < now() - ($3::bigint || ' milliseconds')::interval
+       RETURNING task_id`,
+      [
+        ScheduledTaskStatus.PENDING,
+        ScheduledTaskStatus.RUNNING,
+        olderThanMs,
+      ]
+    );
+
+    return result.rowCount ?? 0;
+  }
+
   /**
    * Delete task
    */

package/src/stores/ThreadMessageQueueStore.ts

@@ -12,6 +12,7 @@ import type {
   AddMessageParams,
   ThreadInfo,
   QueueScope,
+  ClaimMessagesOptions,
 } from "@axiom-lattice/core";
 import {
   parseQueuedExecutionMode,
@@ -20,12 +21,7 @@ import {
   type TrustedRunContext,
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
-import { createThreadMessageQueueTable } from "../migrations/thread_message_queue_migrations";
-import { addPriorityAndCommandColumns } from "../migrations/add_priority_command_columns";
-import { addCustomRunConfigColumn } from "../migrations/add_custom_run_config_column";
-import { alterMessageQueueIdColumn } from "../migrations/alter_message_queue_id_column";
-import { addWorkspaceProjectToQueue } from "../migrations/add_workspace_project_to_queue";
-import { addTrustedRunContextColumn } from "../migrations/add_trusted_run_context_column";
+import { registerThreadMessageQueueMigrations } from "../migrations/queue_migrations";
 
 export type { PendingMessage, AddMessageParams, ThreadInfo };
 
@@ -61,12 +57,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
     }
 
     this.migrationManager = new MigrationManager(this.pool);
-    this.migrationManager.register(createThreadMessageQueueTable);
-    this.migrationManager.register(addPriorityAndCommandColumns);
-    this.migrationManager.register(addCustomRunConfigColumn);
-    this.migrationManager.register(alterMessageQueueIdColumn);
-    this.migrationManager.register(addWorkspaceProjectToQueue);
-    this.migrationManager.register(addTrustedRunContextColumn);
+    registerThreadMessageQueueMigrations(this.migrationManager);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -211,6 +202,50 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
     return result.rows.map(row => this.rowToMessage(row));
   }
 
+  /**
+   * Atomically claims pending messages for a replica, stamping a lease.
+   */
+  async claimPendingMessages(
+    threadId: string,
+    scope: QueueScope | undefined,
+    options: ClaimMessagesOptions,
+  ): Promise<PendingMessage[]> {
+    const filter = scopeClause(scope, 4);
+    const limitClause = options.limit === undefined ? "" : ` LIMIT ${Math.max(0, Math.floor(options.limit))}`;
+    const result = await this.pool.query(
+      `UPDATE lattice_thread_message_queue q
+          SET status = 'processing',
+              locked_by = $2,
+              locked_at = now(),
+              lease_expires_at = now() + ($3::bigint || ' milliseconds')::interval
+        WHERE q.id IN (
+          SELECT id FROM lattice_thread_message_queue
+           WHERE thread_id = $1 AND status = 'pending'${filter.sql}
+           ORDER BY priority DESC, sequence_order ASC${limitClause}
+           FOR UPDATE SKIP LOCKED
+        )
+        RETURNING q.*`,
+      [threadId, options.owner, options.ttlMs, ...filter.params],
+    );
+    return result.rows.map((row) => this.rowToMessage(row));
+  }
+
+  /**
+   * Returns expired processing rows to pending and clears their lease.
+   */
+  async reclaimExpiredProcessing(threadId: string, scope?: QueueScope): Promise<number> {
+    const filter = scopeClause(scope, 2);
+    const result = await this.pool.query(
+      `UPDATE lattice_thread_message_queue
+          SET status = 'pending', locked_by = NULL, locked_at = NULL, lease_expires_at = NULL
+        WHERE thread_id = $1 AND status = 'processing'
+          AND lease_expires_at IS NOT NULL AND lease_expires_at <= now()${filter.sql}
+        RETURNING id`,
+      [threadId, ...filter.params],
+    );
+    return result.rowCount ?? 0;
+  }
+
   /**
    * Get processing messages for a thread
    */
@@ -311,7 +346,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
     const filter = scopeClause(scope, 2);
     const result = await this.pool.query(
       `UPDATE lattice_thread_message_queue 
-       SET status = 'pending'
+       SET status = 'pending', locked_by = NULL, locked_at = NULL, lease_expires_at = NULL
         WHERE thread_id = $1 AND status = 'processing'${filter.sql}
        RETURNING id`,
       [threadId, ...filter.params]
@@ -322,9 +357,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
   private rowToMessage(row: any): PendingMessage {
     return {
       id: row.id,
-      content: typeof row.message_content === "string"
-        ? JSON.parse(row.message_content)
-        : row.message_content,
+      content: row.message_content,
       type: row.message_type,
       sequence: row.sequence_order,
       createdAt: new Date(row.created_at),