@axiom-lattice/pg-stores 3.3.12 → 3.4.0

package/src/stores/PostgreSQLChannelReplyJobStore.ts

@@ -0,0 +1,373 @@
+import { Pool, type PoolConfig } from "pg";
+import {
+  assertValidCreateChannelReplyJobInput,
+  normalizeChannelReplyJobListLimit,
+  type ChannelReplyJob,
+  type ChannelReplyJobStore,
+  type CreateChannelReplyJobInput,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createChannelReplyJobsTable } from "../migrations/channel_reply_job_migration";
+
+/** Options accepted by {@link PostgreSQLChannelReplyJobStore}. */
+export interface PostgreSQLChannelReplyJobStoreOptions {
+  /** Externally managed pool. When supplied, migrations are not run here. */
+  pool?: Pool;
+  /** Connection string or pool config used to create a pool. */
+  poolConfig?: string | PoolConfig;
+  /** Set to `false` to skip running this store's migration automatically. */
+  autoMigrate?: boolean;
+}
+
+interface ChannelReplyJobRow {
+  run_id: string;
+  tenant_id: string;
+  assistant_id: string;
+  thread_id: string;
+  adapter_channel: string;
+  channel_installation_id: string;
+  reply_target: Record<string, unknown>;
+  input_message_id: string;
+  stream_cursor: string | null;
+  status: string;
+  attempts: number | string;
+  available_at: number | string;
+  lease_owner: string | null;
+  lease_expires_at: number | string | null;
+  last_error: string | null;
+  output_message_id: string | null;
+  created_at: number | string;
+  updated_at: number | string;
+  delivered_at: number | string | null;
+}
+
+/** Epoch-millisecond database clock, so replicas never share `Date.now()` skew. */
+const DB_CLOCK_MS =
+  "FLOOR(EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::bigint";
+const COLUMNS = `run_id, tenant_id, assistant_id, thread_id, adapter_channel,
+  channel_installation_id, reply_target, input_message_id, stream_cursor, status,
+  attempts, available_at, lease_owner, lease_expires_at, last_error,
+  output_message_id, created_at, updated_at, delivered_at`;
+
+function toNumber(value: number | string): number {
+  if (typeof value === "number") return value;
+  if (typeof value === "string") return Number(value);
+  throw new Error("Invalid channel reply job numeric column");
+}
+
+function toNullableNumber(value: number | string | null): number | null {
+  return value === null ? null : toNumber(value);
+}
+
+function toStatus(value: string): ChannelReplyJob["status"] {
+  if (value === "pending" || value === "delivered" || value === "failed") {
+    return value;
+  }
+  throw new Error(`Invalid channel reply job status: ${value}`);
+}
+
+function mapRow(row: ChannelReplyJobRow): ChannelReplyJob {
+  return {
+    runId: row.run_id,
+    tenantId: row.tenant_id,
+    assistantId: row.assistant_id,
+    threadId: row.thread_id,
+    adapterChannel: row.adapter_channel,
+    channelInstallationId: row.channel_installation_id,
+    replyTarget: row.reply_target,
+    inputMessageId: row.input_message_id,
+    streamCursor: row.stream_cursor,
+    status: toStatus(row.status),
+    attempts: toNumber(row.attempts),
+    availableAt: toNumber(row.available_at),
+    leaseOwner: row.lease_owner,
+    leaseExpiresAt: toNullableNumber(row.lease_expires_at),
+    lastError: row.last_error,
+    outputMessageId: row.output_message_id,
+    createdAt: toNumber(row.created_at),
+    updatedAt: toNumber(row.updated_at),
+    deliveredAt: toNullableNumber(row.delivered_at),
+  };
+}
+
+/**
+ * PostgreSQL {@link ChannelReplyJobStore}.
+ *
+ * @remarks Claims are a single conditional `UPDATE` (CAS) guarded by status,
+ * `available_at`, and lease expiry, so competing replicas cannot both win.
+ * Lease/created/updated/delivered timestamps use the database clock; only the
+ * caller-supplied `availableAt` epoch is stored as-is.
+ */
+export class PostgreSQLChannelReplyJobStore implements ChannelReplyJobStore {
+  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/connection URL. */
+  constructor(options: PostgreSQLChannelReplyJobStoreOptions) {
+    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(createChannelReplyJobsTable);
+
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLChannelReplyJobStore:", 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();
+  }
+
+  /** Creates a pending job with database-clock created/updated timestamps. */
+  async create(input: CreateChannelReplyJobInput): Promise<ChannelReplyJob> {
+    assertValidCreateChannelReplyJobInput(input);
+    const result = await this.pool.query<ChannelReplyJobRow>(
+      `WITH clock AS (SELECT ${DB_CLOCK_MS} AS now_ms)
+       INSERT INTO lattice_channel_reply_jobs (
+         run_id, tenant_id, assistant_id, thread_id, adapter_channel,
+         channel_installation_id, reply_target, input_message_id, stream_cursor,
+         status, attempts, available_at, lease_owner, lease_expires_at,
+         last_error, output_message_id, created_at, updated_at, delivered_at
+       )
+       SELECT
+         $1, $2, $3, $4, $5,
+         $6, $7::jsonb, $8, $9,
+         'pending', 0, $10, NULL, NULL,
+         NULL, NULL, clock.now_ms, clock.now_ms, NULL
+       FROM clock
+       RETURNING ${COLUMNS}`,
+      [
+        input.runId,
+        input.tenantId,
+        input.assistantId,
+        input.threadId,
+        input.adapterChannel,
+        input.channelInstallationId,
+        JSON.stringify(input.replyTarget),
+        input.inputMessageId,
+        input.streamCursor ?? null,
+        input.availableAt,
+      ],
+    );
+    return mapRow(result.rows[0]);
+  }
+
+  /** Reads one job within a tenant scope. */
+  async get(tenantId: string, runId: string): Promise<ChannelReplyJob | null> {
+    const result = await this.pool.query<ChannelReplyJobRow>(
+      `SELECT ${COLUMNS}
+         FROM lattice_channel_reply_jobs
+        WHERE tenant_id = $1 AND run_id = $2`,
+      [tenantId, runId],
+    );
+    return result.rowCount === 0 ? null : mapRow(result.rows[0]);
+  }
+
+  /** Lists claimable pending jobs that are due at `now`, oldest first. */
+  async listPending(now: number, limit?: number): Promise<ChannelReplyJob[]> {
+    const result = await this.pool.query<ChannelReplyJobRow>(
+      `SELECT ${COLUMNS}
+         FROM lattice_channel_reply_jobs
+        WHERE status = 'pending'
+          AND available_at <= $1
+          AND (lease_owner IS NULL OR lease_expires_at IS NULL OR lease_expires_at <= $1)
+        ORDER BY available_at ASC, created_at ASC, run_id ASC
+        LIMIT $2`,
+      [now, normalizeChannelReplyJobListLimit(limit)],
+    );
+    return result.rows.map(mapRow);
+  }
+
+  /**
+   * Atomically leases a claimable job.
+   *
+   * @param now Accepted for interface symmetry; PostgreSQL uses the database clock.
+   */
+  async claim(
+    tenantId: string,
+    runId: string,
+    owner: string,
+    leaseMs: number,
+    now?: number,
+  ): Promise<ChannelReplyJob | null> {
+    void now;
+    const result = await this.pool.query<ChannelReplyJobRow>(
+      `WITH clock AS (SELECT ${DB_CLOCK_MS} AS now_ms)
+       UPDATE lattice_channel_reply_jobs AS j
+          SET lease_owner = $3,
+              lease_expires_at = clock.now_ms + $4,
+              updated_at = clock.now_ms
+         FROM clock
+        WHERE j.tenant_id = $1
+          AND j.run_id = $2
+          AND j.status = 'pending'
+          AND j.available_at <= clock.now_ms
+          AND (j.lease_owner IS NULL OR j.lease_expires_at IS NULL OR j.lease_expires_at <= clock.now_ms)
+        RETURNING ${COLUMNS}`,
+      [tenantId, runId, owner, leaseMs],
+    );
+    return result.rowCount === 0 ? null : mapRow(result.rows[0]);
+  }
+
+  /**
+   * Marks a job delivered when `owner` still holds the lease.
+   *
+   * @param now Accepted for interface symmetry; PostgreSQL uses the database clock.
+   */
+  async markDelivered(
+    tenantId: string,
+    runId: string,
+    owner: string,
+    outputMessageId: string | null = null,
+    now?: number,
+  ): Promise<ChannelReplyJob | null> {
+    void now;
+    const result = await this.pool.query<ChannelReplyJobRow>(
+      `WITH clock AS (SELECT ${DB_CLOCK_MS} AS now_ms)
+       UPDATE lattice_channel_reply_jobs AS j
+          SET status = 'delivered',
+              output_message_id = $4,
+              lease_owner = NULL,
+              lease_expires_at = NULL,
+              delivered_at = clock.now_ms,
+              updated_at = clock.now_ms
+         FROM clock
+        WHERE j.tenant_id = $1
+          AND j.run_id = $2
+          AND j.lease_owner = $3
+          AND j.status = 'pending'
+        RETURNING ${COLUMNS}`,
+      [tenantId, runId, owner, outputMessageId],
+    );
+    return result.rowCount === 0 ? null : mapRow(result.rows[0]);
+  }
+
+  /**
+   * Returns a job to pending with an incremented attempt count.
+   *
+   * @param now Accepted for interface symmetry; PostgreSQL uses the database clock.
+   */
+  async reschedule(
+    tenantId: string,
+    runId: string,
+    owner: string,
+    error: string,
+    availableAt: number,
+    now?: number,
+  ): Promise<ChannelReplyJob | null> {
+    void now;
+    const result = await this.pool.query<ChannelReplyJobRow>(
+      `WITH clock AS (SELECT ${DB_CLOCK_MS} AS now_ms)
+       UPDATE lattice_channel_reply_jobs AS j
+          SET status = 'pending',
+              attempts = attempts + 1,
+              lease_owner = NULL,
+              lease_expires_at = NULL,
+              last_error = $4,
+              available_at = $5,
+              updated_at = clock.now_ms
+         FROM clock
+        WHERE j.tenant_id = $1
+          AND j.run_id = $2
+          AND j.lease_owner = $3
+          AND j.status = 'pending'
+        RETURNING ${COLUMNS}`,
+      [tenantId, runId, owner, error, availableAt],
+    );
+    return result.rowCount === 0 ? null : mapRow(result.rows[0]);
+  }
+
+  /**
+   * Marks a job terminally failed when `owner` still holds the lease.
+   *
+   * @param now Accepted for interface symmetry; PostgreSQL uses the database clock.
+   */
+  async markFailed(
+    tenantId: string,
+    runId: string,
+    owner: string,
+    error: string,
+    now?: number,
+  ): Promise<ChannelReplyJob | null> {
+    void now;
+    const result = await this.pool.query<ChannelReplyJobRow>(
+      `WITH clock AS (SELECT ${DB_CLOCK_MS} AS now_ms)
+       UPDATE lattice_channel_reply_jobs AS j
+          SET status = 'failed',
+              last_error = $4,
+              lease_owner = NULL,
+              lease_expires_at = NULL,
+              updated_at = clock.now_ms
+         FROM clock
+        WHERE j.tenant_id = $1
+          AND j.run_id = $2
+          AND j.lease_owner = $3
+          AND j.status = 'pending'
+        RETURNING ${COLUMNS}`,
+      [tenantId, runId, owner, error],
+    );
+    return result.rowCount === 0 ? null : mapRow(result.rows[0]);
+  }
+
+  /** Deletes one job. Returns whether a row was removed. */
+  async delete(tenantId: string, runId: string): Promise<boolean> {
+    const result = await this.pool.query(
+      `DELETE FROM lattice_channel_reply_jobs WHERE tenant_id = $1 AND run_id = $2`,
+      [tenantId, runId],
+    );
+    return (result.rowCount ?? 0) > 0;
+  }
+
+  /** Deletes terminal jobs whose last update is older than `before`. */
+  async deleteTerminalBefore(before: number, limit?: number): Promise<number> {
+    const result = await this.pool.query(
+      `DELETE FROM lattice_channel_reply_jobs
+        WHERE run_id IN (
+          SELECT run_id
+            FROM lattice_channel_reply_jobs
+           WHERE status IN ('delivered', 'failed')
+             AND updated_at < $1
+           ORDER BY updated_at ASC, run_id ASC
+           LIMIT $2
+        )`,
+      [before, normalizeChannelReplyJobListLimit(limit)],
+    );
+    return result.rowCount ?? 0;
+  }
+}