@axiom-lattice/pg-stores 3.2.1 → 3.3.0

package/dist/index.mjs

@@ -1,8 +1,8 @@
 // src/index.ts
-import { Pool as Pool31 } from "pg";
+import { Pool as Pool33 } from "pg";
 
 // src/createPgStoreConfig.ts
-import { Pool as Pool28 } from "pg";
+import { Pool as Pool30 } from "pg";
 
 // src/migrations/migration.ts
 var MigrationManager = class {
@@ -5210,6 +5210,44 @@ var addTrustedRunContextColumn = {
   }
 };
 
+// src/migrations/add_queue_claim_lease.ts
+var addQueueClaimLease = {
+  version: 186,
+  name: "add_queue_claim_lease_columns",
+  up: async (client) => {
+    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) => {
+    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
+    `);
+  }
+};
+
+// src/migrations/queue_migrations.ts
+function registerThreadMessageQueueMigrations(manager) {
+  manager.register(createThreadMessageQueueTable);
+  manager.register(addPriorityAndCommandColumns);
+  manager.register(alterMessageQueueIdColumn);
+  manager.register(addCustomRunConfigColumn);
+  manager.register(addWorkspaceProjectToQueue);
+  manager.register(addTrustedRunContextColumn);
+  manager.register(addQueueClaimLease);
+}
+
 // src/stores/ThreadMessageQueueStore.ts
 var ThreadMessageQueueStore = class {
   constructor(options) {
@@ -5230,12 +5268,7 @@ var ThreadMessageQueueStore = class {
       throw new Error("Either pool or poolConfig must be provided");
     }
     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) => {
         console.error("Failed to initialize ThreadMessageQueueStore:", error);
@@ -5360,6 +5393,44 @@ var ThreadMessageQueueStore = class {
     );
     return result.rows.map((row) => this.rowToMessage(row));
   }
+  /**
+   * Atomically claims pending messages for a replica, stamping a lease.
+   */
+  async claimPendingMessages(threadId, scope, options) {
+    const filter = scopeClause(scope, 4);
+    const limitClause = options.limit === void 0 ? "" : ` 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, scope) {
+    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
    */
@@ -5450,7 +5521,7 @@ var ThreadMessageQueueStore = class {
     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]
@@ -5460,7 +5531,7 @@ var ThreadMessageQueueStore = class {
   rowToMessage(row) {
     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),
@@ -7191,6 +7262,54 @@ var PostgreSQLScheduleStorage = class {
       values
     );
   }
+  /**
+   * 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) {
+    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,
+        /* @__PURE__ */ 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) {
+    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
    */
@@ -9961,6 +10080,52 @@ var createUserPushSubscriptions = {
   }
 };
 
+// src/migrations/distributed_lease_migration.ts
+var createDistributedLeases = {
+  version: 185,
+  name: "create_lattice_distributed_leases",
+  up: async (client) => {
+    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) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_leases_expires");
+    await client.query("DROP TABLE IF EXISTS lattice_leases");
+  }
+};
+
+// src/migrations/outbox_migration.ts
+var createOutboxTable = {
+  version: 187,
+  name: "create_lattice_outbox",
+  up: async (client) => {
+    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) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_outbox_id");
+    await client.query("DROP TABLE IF EXISTS lattice_outbox");
+  }
+};
+
 // src/migrations/model_provider_migrations.ts
 var createModelProvidersTable = {
   version: 181,
@@ -11536,9 +11701,217 @@ var PostgreSQLModelProviderStore = class {
   }
 };
 
+// src/stores/PostgreSQLLeaseStore.ts
+import { Pool as Pool28 } from "pg";
+var PostgreSQLLeaseStore = class {
+  /** @param options An externally managed pool, or a pool config/URL. */
+  constructor(options) {
+    this.ownsPool = true;
+    this.initialized = false;
+    this.initPromise = null;
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool28({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool28(options.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() {
+    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() {
+    if (this.ownsPool) await this.pool.end();
+  }
+  /** Atomically acquires a free/expired key (or re-acquires the caller's own). */
+  async tryAcquire(key, holder, ttlMs) {
+    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, holder, ttlMs) {
+    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, holder) {
+    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) {
+    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()
+    };
+  }
+};
+
+// src/stores/PostgreSQLOutboxStore.ts
+import { Pool as Pool29 } from "pg";
+var PostgreSQLOutboxStore = class {
+  /** @param options An externally managed pool, or a pool config/URL. */
+  constructor(options) {
+    this.ownsPool = true;
+    this.initialized = false;
+    this.initPromise = null;
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool29({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool29(options.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() {
+    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() {
+    if (this.ownsPool) await this.pool.end();
+  }
+  /** Appends one event and returns it with its assigned id. */
+  async append(input) {
+    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, limit) {
+    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() {
+    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) {
+    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}. */
+  rowToEvent(row) {
+    return {
+      id: String(row.id),
+      topic: row.topic,
+      payload: row.payload,
+      originReplicaId: row.origin_replica_id,
+      createdAt: new Date(row.created_at).toISOString()
+    };
+  }
+};
+
 // src/createPgStoreConfig.ts
 async function createPgStoreConfig(connectionString) {
-  const pool = new Pool28({ connectionString });
+  const pool = new Pool30({ connectionString });
   const mm = new MigrationManager(pool);
   mm.register(createThreadsTable);
   mm.register(createScheduledTasksTable);
@@ -11562,9 +11935,7 @@ async function createPgStoreConfig(connectionString) {
   mm.register(createChannelInstallationsTable);
   mm.register(addThreadTenantId);
   mm.register(addProjectConfigColumn);
-  mm.register(createThreadMessageQueueTable);
-  mm.register(addPriorityAndCommandColumns);
-  mm.register(alterMessageQueueIdColumn);
+  registerThreadMessageQueueMigrations(mm);
   mm.register(createWorkflowTrackingTables);
   for (const m of evalMigrations) {
     mm.register(m);
@@ -11574,9 +11945,7 @@ async function createPgStoreConfig(connectionString) {
   mm.register(createDatabaseConfigsTable);
   mm.register(changeAssistantPrimaryKey);
   mm.register(addStepThreadId);
-  mm.register(addCustomRunConfigColumn);
   mm.register(createA2AApiKeysTable);
-  mm.register(addWorkspaceProjectToQueue);
   mm.register(addAssistantOwnerUserId);
   for (const m of taskMigrations) {
     mm.register(m);
@@ -11600,7 +11969,6 @@ async function createPgStoreConfig(connectionString) {
   mm.register(createCapabilityBundlesTable);
   mm.register(addTaskWorkItemPendingIndexesMigration);
   mm.register(createProjectRoomTables);
-  mm.register(addTrustedRunContextColumn);
   mm.register(addProjectLifecycleEventIndex);
   mm.register(addAgentWebAppsIdentity);
   mm.register(createProjectRoomReadStates);
@@ -11610,6 +11978,8 @@ async function createPgStoreConfig(connectionString) {
   mm.register(createModelProviderModelsTable);
   mm.register(addModelProviderApiStyle);
   mm.register(addModelProviderModelStale);
+  mm.register(createDistributedLeases);
+  mm.register(createOutboxTable);
   await mm.migrate();
   const checkpoint = PostgresSaver.fromConnString(connectionString);
   checkpoint.setup().catch((err) => {
@@ -11651,13 +12021,15 @@ async function createPgStoreConfig(connectionString) {
     projectRoomMessage: new PostgreSQLProjectRoomMessageStore(opts),
     projectRoomReadState: new PostgreSQLProjectRoomReadStateStore(opts),
     userPushSubscription: new PostgreSQLUserPushSubscriptionStore(opts),
+    lease: new PostgreSQLLeaseStore(opts),
+    outbox: new PostgreSQLOutboxStore(opts),
     vectorStoreProvider: new PGVectorStoreProvider(pool, connectionString),
     checkpoint
   };
 }
 
 // src/stores/PostgreSQLSkillStore.ts
-import { Pool as Pool29 } from "pg";
+import { Pool as Pool31 } from "pg";
 var PostgreSQLSkillStore = class {
   constructor(options) {
     this.initialized = false;
@@ -11670,9 +12042,9 @@ var PostgreSQLSkillStore = class {
       return;
     }
     if (typeof options.poolConfig === "string") {
-      this.pool = new Pool29({ connectionString: options.poolConfig });
+      this.pool = new Pool31({ connectionString: options.poolConfig });
     } else if (options.poolConfig) {
-      this.pool = new Pool29(options.poolConfig);
+      this.pool = new Pool31(options.poolConfig);
     } else {
       throw new Error("Either pool or poolConfig must be provided");
     }
@@ -11973,7 +12345,7 @@ var PostgreSQLSkillStore = class {
 };
 
 // src/stores/ChannelIdentityMappingStore.ts
-import { Pool as Pool30 } from "pg";
+import { Pool as Pool32 } from "pg";
 var ChannelIdentityMappingStore = class {
   constructor(options) {
     this.initialized = false;
@@ -11985,7 +12357,7 @@ var ChannelIdentityMappingStore = class {
       this.initialized = true;
       return;
     }
-    this.pool = typeof options.poolConfig === "string" ? new Pool30({ connectionString: options.poolConfig }) : options.poolConfig ? new Pool30(options.poolConfig) : (() => {
+    this.pool = typeof options.poolConfig === "string" ? new Pool32({ connectionString: options.poolConfig }) : options.poolConfig ? new Pool32(options.poolConfig) : (() => {
       throw new Error("Either pool or poolConfig must be provided");
     })();
     this.migrationManager = new MigrationManager(this.pool);
@@ -12211,7 +12583,7 @@ export {
   MenuStore,
   MigrationManager,
   PGVectorStoreProvider,
-  Pool31 as Pool,
+  Pool33 as Pool,
   PostgreSQLA2AApiKeyStore,
   PostgreSQLAgentWebAppStore,
   PostgreSQLAssistantStore,