@axiom-lattice/pg-stores 3.3.13 → 3.4.0

package/src/stores/PostgreSQLCheckpointStorage.ts

@@ -0,0 +1,548 @@
+import type { Pool, PoolClient } from "pg";
+import type {
+  CheckpointMaintenance,
+  CheckpointStorage,
+  ThreadCheckpointStorage,
+} from "@axiom-lattice/protocols";
+
+/**
+ * Raised when a checkpoint copy targets a thread that already owns rows.
+ */
+export class CheckpointTargetNotEmptyError extends Error {
+  /** Machine-readable conflict code. */
+  readonly code = "CHECKPOINT_TARGET_NOT_EMPTY" as const;
+
+  /**
+   * @param threadId - Non-empty target thread id.
+   */
+  constructor(threadId: string) {
+    super(`Checkpoint target thread ${threadId} is not empty`);
+    this.name = "CheckpointTargetNotEmptyError";
+  }
+}
+
+/**
+ * Raised when a checkpoint rewind targets a base checkpoint that does not exist.
+ */
+export class CheckpointBaseNotFoundError extends Error {
+  /** Machine-readable failure code. */
+  readonly code = "CHECKPOINT_BASE_NOT_FOUND" as const;
+  /** Thread owning the missing base checkpoint. */
+  readonly threadId: string;
+  /** Base checkpoint id that was not found. */
+  readonly baseCheckpointId: string;
+
+  /**
+   * @param threadId - Thread owning the missing base checkpoint.
+   * @param baseCheckpointId - Base checkpoint id that was not found.
+   */
+  constructor(threadId: string, baseCheckpointId: string) {
+    super(`Checkpoint ${baseCheckpointId} does not exist on thread ${threadId}`);
+    this.name = "CheckpointBaseNotFoundError";
+    this.threadId = threadId;
+    this.baseCheckpointId = baseCheckpointId;
+  }
+}
+
+async function rollbackQuietly(client: PoolClient): Promise<void> {
+  try {
+    await client.query("ROLLBACK");
+  } catch {
+    // Preserve the original failure.
+  }
+}
+
+function isUniqueViolation(error: unknown): boolean {
+  return typeof error === "object"
+    && error !== null
+    && "code" in error
+    && (error as { code?: unknown }).code === "23505";
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function checkpointKey(checkpointNs: string, checkpointId: string): string {
+  return `${checkpointNs}\u0000${checkpointId}`;
+}
+
+function blobKey(checkpointNs: string, channel: string, version: string): string {
+  return `${checkpointNs}\u0000${channel}\u0000${version}`;
+}
+
+function checkpointRecord(value: unknown): Record<string, unknown> | undefined {
+  if (typeof value === "string") {
+    try {
+      const parsed: unknown = JSON.parse(value);
+      return isRecord(parsed) ? parsed : undefined;
+    } catch {
+      return undefined;
+    }
+  }
+  return isRecord(value) ? value : undefined;
+}
+
+function channelVersions(value: unknown): Record<string, unknown> | undefined {
+  const checkpoint = checkpointRecord(value);
+  const versions = checkpoint?.channel_versions;
+  return isRecord(versions) ? versions : undefined;
+}
+
+interface RetainedCheckpointRow {
+  checkpoint_ns: string;
+  checkpoint_id: string;
+  checkpoint: unknown;
+}
+
+interface CheckpointKeyRow {
+  checkpoint_ns: string;
+  checkpoint_id: string;
+}
+
+interface BlobKeyRow {
+  checkpoint_ns: string;
+  channel: string;
+  version: string;
+}
+
+interface WriteKeyRow {
+  checkpoint_ns: string;
+  checkpoint_id: string;
+  task_id: string;
+  idx: number;
+}
+
+const DELETE_CHUNK_SIZE = 100;
+
+function chunk<T>(items: T[], size: number): T[][] {
+  const chunks: T[][] = [];
+  for (let index = 0; index < items.length; index += size) {
+    chunks.push(items.slice(index, index + size));
+  }
+  return chunks;
+}
+
+interface CheckpointStorageRow {
+  checkpoint_count: number | string | null;
+  checkpoint_bytes: number | string | null;
+  blob_count: number | string | null;
+  blob_bytes: number | string | null;
+  write_count: number | string | null;
+  write_bytes: number | string | null;
+}
+
+function toCount(value: number | string | null | undefined): number {
+  const parsed = Number(value ?? 0);
+  return Number.isFinite(parsed) ? parsed : 0;
+}
+
+/**
+ * Checkpoint accounting and maintenance over LangGraph's built-in
+ * `PostgresSaver` tables.
+ *
+ * `get` is a read-only aggregate. `copyThread`, `deleteThread`, `deleteAfter`,
+ * and `pruneThread` are the explicit maintenance operations; each runs in one
+ * transaction and touches only the `checkpoints`, `checkpoint_blobs`, and
+ * `checkpoint_writes` tables created by `PostgresSaver.setup()`. No migration
+ * is added.
+ */
+export class PostgreSQLCheckpointStorage implements CheckpointStorage, CheckpointMaintenance {
+  /**
+   * @param pool - PostgreSQL pool shared with the rest of the store config.
+   */
+  constructor(private readonly pool: Pool) {}
+
+  /**
+   * Measure the checkpoint storage used by one thread.
+   *
+   * @param threadId - Thread id whose checkpoint rows should be measured.
+   * @returns Counts and byte sizes grouped by checkpoint storage category.
+   */
+  async get(threadId: string): Promise<ThreadCheckpointStorage> {
+    const result = await this.pool.query<CheckpointStorageRow>(
+      `SELECT
+        (SELECT count(*)::bigint FROM checkpoints WHERE thread_id=$1) AS checkpoint_count,
+        (SELECT coalesce(sum(pg_column_size(checkpoint) + pg_column_size(metadata)),0)::bigint FROM checkpoints WHERE thread_id=$1) AS checkpoint_bytes,
+        (SELECT count(*)::bigint FROM checkpoint_blobs WHERE thread_id=$1) AS blob_count,
+        (SELECT coalesce(sum(coalesce(octet_length(blob),0)),0)::bigint FROM checkpoint_blobs WHERE thread_id=$1) AS blob_bytes,
+        (SELECT count(*)::bigint FROM checkpoint_writes WHERE thread_id=$1) AS write_count,
+        (SELECT coalesce(sum(octet_length(blob)),0)::bigint FROM checkpoint_writes WHERE thread_id=$1) AS write_bytes`,
+      [threadId],
+    );
+
+    const row = result.rows[0];
+    const checkpointCount = toCount(row?.checkpoint_count);
+    const checkpointBytes = toCount(row?.checkpoint_bytes);
+    const blobCount = toCount(row?.blob_count);
+    const blobBytes = toCount(row?.blob_bytes);
+    const writeCount = toCount(row?.write_count);
+    const writeBytes = toCount(row?.write_bytes);
+
+    return {
+      threadId,
+      totalBytes: checkpointBytes + blobBytes + writeBytes,
+      checkpoints: { count: checkpointCount, bytes: checkpointBytes },
+      channelValues: { count: blobCount, bytes: blobBytes },
+      pendingWrites: { count: writeCount, bytes: writeBytes },
+    };
+  }
+
+  /**
+   * Copy every checkpoint row from a source thread to an empty target thread.
+   *
+   * @param sourceThreadId - Thread whose rows are copied.
+   * @param targetThreadId - Empty thread that receives the rows.
+   * @returns Resolves when the copy commits; an empty source is a no-op.
+   * @throws {CheckpointTargetNotEmptyError} When the target already owns rows.
+   */
+  async copyThread(sourceThreadId: string, targetThreadId: string): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+      const target = await client.query<{ not_empty: boolean }>(
+        `SELECT (
+          EXISTS (SELECT 1 FROM checkpoints WHERE thread_id=$1)
+          OR EXISTS (SELECT 1 FROM checkpoint_blobs WHERE thread_id=$1)
+          OR EXISTS (SELECT 1 FROM checkpoint_writes WHERE thread_id=$1)
+        ) AS not_empty`,
+        [targetThreadId],
+      );
+      if (target.rows[0]?.not_empty) {
+        throw new CheckpointTargetNotEmptyError(targetThreadId);
+      }
+
+      await client.query(
+        `INSERT INTO checkpoints (
+          thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
+        )
+        SELECT $2, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint,
+          jsonb_set(coalesce(metadata, '{}'::jsonb), '{thread_id}', to_jsonb($2::text), true)
+        FROM checkpoints WHERE thread_id=$1`,
+        [sourceThreadId, targetThreadId],
+      );
+      await client.query(
+        `INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
+        SELECT $2, checkpoint_ns, channel, version, type, blob
+        FROM checkpoint_blobs WHERE thread_id=$1`,
+        [sourceThreadId, targetThreadId],
+      );
+      await client.query(
+        `INSERT INTO checkpoint_writes (
+          thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob
+        )
+        SELECT $2, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob
+        FROM checkpoint_writes WHERE thread_id=$1`,
+        [sourceThreadId, targetThreadId],
+      );
+      await client.query("COMMIT");
+    } catch (error) {
+      await rollbackQuietly(client);
+      if (isUniqueViolation(error)) {
+        throw new CheckpointTargetNotEmptyError(targetThreadId);
+      }
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  /**
+   * Retain only the latest N checkpoints per checkpoint namespace.
+   *
+   * @param threadId - Thread whose checkpoint history is pruned.
+   * @param options - Retention options.
+   * @param options.keepLatest - Number of latest checkpoints to retain per namespace.
+   * @returns Counts of deleted checkpoint, blob, and write rows.
+   */
+  async pruneThread(
+    threadId: string,
+    options: { keepLatest: number },
+  ): Promise<{ checkpoints: number; blobs: number; writes: number }> {
+    if (!Number.isSafeInteger(options.keepLatest) || options.keepLatest < 1) {
+      throw new RangeError("keepLatest must be an integer >= 1");
+    }
+    const keepLatest = options.keepLatest;
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+
+      const retainedResult = await client.query<RetainedCheckpointRow>(
+        `SELECT checkpoint_ns, checkpoint_id, checkpoint FROM (
+          SELECT checkpoint_ns, checkpoint_id, checkpoint,
+            row_number() OVER (
+              PARTITION BY checkpoint_ns ORDER BY checkpoint_id DESC
+            ) AS rn
+          FROM checkpoints WHERE thread_id=$1
+        ) t WHERE rn <= $2`,
+        [threadId, keepLatest],
+      );
+
+      const retained = retainedResult.rows;
+      const retainedCheckpoints = new Set(
+        retained.map((row) => checkpointKey(row.checkpoint_ns, row.checkpoint_id)),
+      );
+      const fallbackNamespaces = new Set<string>();
+      const retainedBlobs = new Set<string>();
+      for (const row of retained) {
+        const versions = channelVersions(row.checkpoint);
+        if (!versions) {
+          fallbackNamespaces.add(row.checkpoint_ns);
+          continue;
+        }
+        for (const [channel, version] of Object.entries(versions)) {
+          retainedBlobs.add(blobKey(row.checkpoint_ns, channel, String(version)));
+        }
+      }
+
+      const writeRows = await client.query<WriteKeyRow>(
+        `SELECT checkpoint_ns, checkpoint_id, task_id, idx
+         FROM checkpoint_writes WHERE thread_id=$1`,
+        [threadId],
+      );
+      const staleWrites = writeRows.rows.filter(
+        (row) => !retainedCheckpoints.has(checkpointKey(row.checkpoint_ns, row.checkpoint_id)),
+      );
+      let writes = 0;
+      for (const batch of chunk(staleWrites, DELETE_CHUNK_SIZE)) {
+        const params: unknown[] = [threadId];
+        const tuples = batch.map((row) => {
+          params.push(row.checkpoint_ns, row.checkpoint_id, row.task_id, row.idx);
+          return `($${params.length - 3}, $${params.length - 2}, $${params.length - 1}, $${params.length})`;
+        });
+        const deleted = await client.query(
+          `DELETE FROM checkpoint_writes
+           WHERE thread_id=$1
+             AND (checkpoint_ns, checkpoint_id, task_id, idx) IN (${tuples.join(", ")})`,
+          params,
+        );
+        writes += deleted.rowCount ?? 0;
+      }
+
+      const checkpointRows = await client.query<CheckpointKeyRow>(
+        `SELECT checkpoint_ns, checkpoint_id FROM checkpoints WHERE thread_id=$1`,
+        [threadId],
+      );
+      const staleCheckpoints = checkpointRows.rows.filter(
+        (row) => !retainedCheckpoints.has(checkpointKey(row.checkpoint_ns, row.checkpoint_id)),
+      );
+      let checkpoints = 0;
+      for (const batch of chunk(staleCheckpoints, DELETE_CHUNK_SIZE)) {
+        const params: unknown[] = [threadId];
+        const tuples = batch.map((row) => {
+          params.push(row.checkpoint_ns, row.checkpoint_id);
+          return `($${params.length - 1}, $${params.length})`;
+        });
+        const deleted = await client.query(
+          `DELETE FROM checkpoints
+           WHERE thread_id=$1
+             AND (checkpoint_ns, checkpoint_id) IN (${tuples.join(", ")})`,
+          params,
+        );
+        checkpoints += deleted.rowCount ?? 0;
+      }
+
+      const blobRows = await client.query<BlobKeyRow>(
+        `SELECT checkpoint_ns, channel, version FROM checkpoint_blobs WHERE thread_id=$1`,
+        [threadId],
+      );
+      const staleBlobs = blobRows.rows.filter(
+        (row) => !fallbackNamespaces.has(row.checkpoint_ns)
+          && !retainedBlobs.has(blobKey(row.checkpoint_ns, row.channel, row.version)),
+      );
+      let blobs = 0;
+      for (const batch of chunk(staleBlobs, DELETE_CHUNK_SIZE)) {
+        const params: unknown[] = [threadId];
+        const tuples = batch.map((row) => {
+          params.push(row.checkpoint_ns, row.channel, row.version);
+          return `($${params.length - 2}, $${params.length - 1}, $${params.length})`;
+        });
+        const deleted = await client.query(
+          `DELETE FROM checkpoint_blobs
+           WHERE thread_id=$1
+             AND (checkpoint_ns, channel, version) IN (${tuples.join(", ")})`,
+          params,
+        );
+        blobs += deleted.rowCount ?? 0;
+      }
+
+      await client.query("COMMIT");
+      return { checkpoints, blobs, writes };
+    } catch (error) {
+      await rollbackQuietly(client);
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  /**
+   * Delete the checkpoints that descend from a base checkpoint.
+   *
+   * Descendants are discovered by following both same-namespace
+   * `parent_checkpoint_id` edges and cross-namespace `metadata.parents` edges
+   * (LangGraph's `{ "<parent checkpoint_ns>": "<parent checkpoint id>" }` map on
+   * subgraph checkpoints), recursively across every namespace. Out-of-order or
+   * clock-skewed ids do not affect the result. Writes keyed to the base
+   * checkpoint belong to the rolled-back run's first superstep and are deleted
+   * too. Blobs referenced only by deleted checkpoints are removed.
+   *
+   * @param threadId - Thread whose later checkpoint history is deleted.
+   * @param baseCheckpointId - Checkpoint to retain as the new tip, or `null` to
+   * delete every checkpoint for the thread.
+   * @returns Resolves when the deletion commits; no descendants is a no-op.
+   * @throws {RangeError} When `baseCheckpointId` is an empty string.
+   * @throws {CheckpointBaseNotFoundError} When the base checkpoint does not
+   * exist on the thread.
+   */
+  async deleteAfter(threadId: string, baseCheckpointId: string | null): Promise<void> {
+    if (baseCheckpointId !== null && baseCheckpointId.length === 0) {
+      throw new RangeError("baseCheckpointId must be null or a non-empty checkpoint id");
+    }
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+
+      if (baseCheckpointId === null) {
+        await client.query("DELETE FROM checkpoint_writes WHERE thread_id=$1", [threadId]);
+        await client.query("DELETE FROM checkpoint_blobs WHERE thread_id=$1", [threadId]);
+        await client.query("DELETE FROM checkpoints WHERE thread_id=$1", [threadId]);
+        await client.query("COMMIT");
+        return;
+      }
+
+      const base = await client.query<{ checkpoint_ns: string }>(
+        `SELECT checkpoint_ns
+         FROM checkpoints
+         WHERE thread_id=$1 AND checkpoint_id=$2`,
+        [threadId, baseCheckpointId],
+      );
+      if (base.rows.length === 0) {
+        throw new CheckpointBaseNotFoundError(threadId, baseCheckpointId);
+      }
+      const baseCheckpointNs = base.rows[0]!.checkpoint_ns;
+
+      const descendantResult = await client.query<CheckpointKeyRow>(
+        `WITH RECURSIVE lineage(checkpoint_ns, checkpoint_id) AS (
+           SELECT $2::text, $3::text
+           UNION
+           SELECT c.checkpoint_ns, c.checkpoint_id
+           FROM checkpoints c
+           JOIN lineage l ON (
+             (c.checkpoint_ns = l.checkpoint_ns
+               AND c.parent_checkpoint_id = l.checkpoint_id)
+             OR (c.metadata->'parents'->>l.checkpoint_ns = l.checkpoint_id)
+           )
+           WHERE c.thread_id=$1
+         ),
+         roots AS (
+           SELECT $2::text AS checkpoint_ns, $3::text AS checkpoint_id
+         )
+         SELECT l.checkpoint_ns, l.checkpoint_id
+         FROM lineage l
+         WHERE NOT EXISTS (
+           SELECT 1 FROM roots r
+           WHERE r.checkpoint_ns = l.checkpoint_ns
+             AND r.checkpoint_id = l.checkpoint_id
+         )`,
+        [threadId, baseCheckpointNs, baseCheckpointId],
+      );
+      const descendants = descendantResult.rows;
+
+      await client.query(
+        `DELETE FROM checkpoint_writes
+         WHERE thread_id=$1 AND checkpoint_ns=$2 AND checkpoint_id=$3`,
+        [threadId, baseCheckpointNs, baseCheckpointId],
+      );
+
+      for (const batch of chunk(descendants, DELETE_CHUNK_SIZE)) {
+        const params: unknown[] = [threadId];
+        const tuples = batch.map((row) => {
+          params.push(row.checkpoint_ns, row.checkpoint_id);
+          return `($${params.length - 1}, $${params.length})`;
+        });
+        await client.query(
+          `DELETE FROM checkpoint_writes
+           WHERE thread_id=$1
+             AND (checkpoint_ns, checkpoint_id) IN (${tuples.join(", ")})`,
+          params,
+        );
+        await client.query(
+          `DELETE FROM checkpoints
+           WHERE thread_id=$1
+             AND (checkpoint_ns, checkpoint_id) IN (${tuples.join(", ")})`,
+          params,
+        );
+      }
+
+      const retainedResult = await client.query<RetainedCheckpointRow>(
+        `SELECT checkpoint_ns, checkpoint_id, checkpoint
+         FROM checkpoints WHERE thread_id=$1`,
+        [threadId],
+      );
+      const fallbackNamespaces = new Set<string>();
+      const retainedBlobs = new Set<string>();
+      for (const row of retainedResult.rows) {
+        const versions = channelVersions(row.checkpoint);
+        if (!versions) {
+          fallbackNamespaces.add(row.checkpoint_ns);
+          continue;
+        }
+        for (const [channel, version] of Object.entries(versions)) {
+          retainedBlobs.add(blobKey(row.checkpoint_ns, channel, String(version)));
+        }
+      }
+
+      const blobRows = await client.query<BlobKeyRow>(
+        `SELECT checkpoint_ns, channel, version FROM checkpoint_blobs WHERE thread_id=$1`,
+        [threadId],
+      );
+      const staleBlobs = blobRows.rows.filter(
+        (row) => !fallbackNamespaces.has(row.checkpoint_ns)
+          && !retainedBlobs.has(blobKey(row.checkpoint_ns, row.channel, row.version)),
+      );
+      for (const batch of chunk(staleBlobs, DELETE_CHUNK_SIZE)) {
+        const params: unknown[] = [threadId];
+        const tuples = batch.map((row) => {
+          params.push(row.checkpoint_ns, row.channel, row.version);
+          return `($${params.length - 2}, $${params.length - 1}, $${params.length})`;
+        });
+        await client.query(
+          `DELETE FROM checkpoint_blobs
+           WHERE thread_id=$1
+             AND (checkpoint_ns, channel, version) IN (${tuples.join(", ")})`,
+          params,
+        );
+      }
+
+      await client.query("COMMIT");
+    } catch (error) {
+      await rollbackQuietly(client);
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  /**
+   * Delete every checkpoint row for a thread.
+   *
+   * @param threadId - Thread whose checkpoint rows are deleted.
+   * @returns Resolves when the deletion commits.
+   */
+  async deleteThread(threadId: string): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+      await client.query("DELETE FROM checkpoint_writes WHERE thread_id=$1", [threadId]);
+      await client.query("DELETE FROM checkpoint_blobs WHERE thread_id=$1", [threadId]);
+      await client.query("DELETE FROM checkpoints WHERE thread_id=$1", [threadId]);
+      await client.query("COMMIT");
+    } catch (error) {
+      await rollbackQuietly(client);
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+}