@axiom-lattice/pg-stores 3.3.13 → 3.4.0

package/src/__tests__/PostgreSQLCheckpointStorage.deleteAfter.test.ts

@@ -0,0 +1,473 @@
+import { afterAll, beforeAll, beforeEach, describe, expect, it } from "@jest/globals";
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+import {
+  CheckpointBaseNotFoundError,
+  PostgreSQLCheckpointStorage,
+} from "../stores/PostgreSQLCheckpointStorage";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+
+interface CheckpointSeed {
+  threadId: string;
+  checkpointNs: string;
+  checkpointId: string;
+  channelVersions: Record<string, string>;
+  parentCheckpointId?: string | null;
+  metadata?: Record<string, unknown>;
+}
+
+interface BlobSeed {
+  threadId: string;
+  checkpointNs: string;
+  channel: string;
+  version: string;
+}
+
+interface WriteSeed {
+  threadId: string;
+  checkpointNs: string;
+  checkpointId: string;
+  taskId: string;
+  idx: number;
+}
+
+maybeDescribe("PostgreSQLCheckpointStorage.deleteAfter (PostgreSQL)", () => {
+  const threadId = randomUUID();
+  let pool: Pool;
+  let store: PostgreSQLCheckpointStorage;
+
+  beforeAll(async () => {
+    pool = new Pool({ connectionString: TEST_DATABASE_URL! });
+    // Both checkpoint suites create the same schema and Jest runs test files
+    // in parallel. Serialize the DDL so CREATE TABLE IF NOT EXISTS cannot race
+    // on pg_type.
+    const client = await pool.connect();
+    try {
+      await client.query("SELECT pg_advisory_lock(94711231)");
+      await client.query(`
+        CREATE TABLE IF NOT EXISTS checkpoints (
+          thread_id TEXT NOT NULL,
+          checkpoint_ns TEXT NOT NULL DEFAULT '',
+          checkpoint_id TEXT NOT NULL,
+          parent_checkpoint_id TEXT,
+          type TEXT,
+          checkpoint JSONB NOT NULL,
+          metadata JSONB NOT NULL DEFAULT '{}',
+          PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
+        );
+        CREATE TABLE IF NOT EXISTS checkpoint_blobs (
+          thread_id TEXT NOT NULL,
+          checkpoint_ns TEXT NOT NULL DEFAULT '',
+          channel TEXT NOT NULL,
+          version TEXT NOT NULL,
+          type TEXT NOT NULL,
+          blob BYTEA,
+          PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
+        );
+        CREATE TABLE IF NOT EXISTS checkpoint_writes (
+          thread_id TEXT NOT NULL,
+          checkpoint_ns TEXT NOT NULL DEFAULT '',
+          checkpoint_id TEXT NOT NULL,
+          task_id TEXT NOT NULL,
+          idx INTEGER NOT NULL,
+          channel TEXT NOT NULL,
+          type TEXT,
+          blob BYTEA NOT NULL,
+          PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
+        );
+      `);
+    } finally {
+      await client.query("SELECT pg_advisory_unlock(94711231)").catch(() => undefined);
+      client.release();
+    }
+    store = new PostgreSQLCheckpointStorage(pool);
+  });
+
+  beforeEach(async () => {
+    await removeRows(threadId);
+  });
+
+  afterAll(async () => {
+    await removeRows(threadId);
+    await pool.end();
+  });
+
+  async function removeRows(id: string): Promise<void> {
+    await pool.query("DELETE FROM checkpoint_writes WHERE thread_id = $1", [id]);
+    await pool.query("DELETE FROM checkpoint_blobs WHERE thread_id = $1", [id]);
+    await pool.query("DELETE FROM checkpoints WHERE thread_id = $1", [id]);
+  }
+
+  async function insertCheckpoints(rows: CheckpointSeed[]): Promise<void> {
+    for (const row of rows) {
+      await pool.query(
+        `INSERT INTO checkpoints (
+           thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata
+         )
+         VALUES ($1, $2, $3, $4, $5::jsonb, $6::jsonb)`,
+        [
+          row.threadId,
+          row.checkpointNs,
+          row.checkpointId,
+          row.parentCheckpointId ?? null,
+          JSON.stringify({ v: 4, channel_versions: row.channelVersions }),
+          JSON.stringify(row.metadata ?? {}),
+        ],
+      );
+    }
+  }
+
+  async function insertBlobs(rows: BlobSeed[]): Promise<void> {
+    for (const row of rows) {
+      await pool.query(
+        `INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
+         VALUES ($1, $2, $3, $4, 'msgpack', $5)`,
+        [
+          row.threadId,
+          row.checkpointNs,
+          row.channel,
+          row.version,
+          Buffer.from(`${row.checkpointNs}:${row.channel}:${row.version}`),
+        ],
+      );
+    }
+  }
+
+  async function insertWrites(rows: WriteSeed[]): Promise<void> {
+    for (const row of rows) {
+      await pool.query(
+        `INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
+         VALUES ($1, $2, $3, $4, $5, 'messages', 'msgpack', $6)`,
+        [
+          row.threadId,
+          row.checkpointNs,
+          row.checkpointId,
+          row.taskId,
+          row.idx,
+          Buffer.from(`${row.checkpointNs}:${row.checkpointId}:${row.idx}`),
+        ],
+      );
+    }
+  }
+
+  async function counts(id: string): Promise<{ checkpoints: number; blobs: number; writes: number }> {
+    const result = await pool.query(
+      `SELECT
+        (SELECT count(*)::int FROM checkpoints WHERE thread_id = $1) AS checkpoints,
+        (SELECT count(*)::int FROM checkpoint_blobs WHERE thread_id = $1) AS blobs,
+        (SELECT count(*)::int FROM checkpoint_writes WHERE thread_id = $1) AS writes`,
+      [id],
+    );
+    return result.rows[0] as { checkpoints: number; blobs: number; writes: number };
+  }
+
+  it("deletes every checkpoint row when the base is null", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: { messages: "1" } },
+      { threadId, checkpointNs: "sub", checkpointId: "cp-2", channelVersions: { messages: "2" } },
+    ]);
+    await insertBlobs([
+      { threadId, checkpointNs: "", channel: "messages", version: "1" },
+      { threadId, checkpointNs: "sub", channel: "messages", version: "2" },
+    ]);
+    await insertWrites([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", taskId: "task-1", idx: 0 },
+      { threadId, checkpointNs: "sub", checkpointId: "cp-2", taskId: "task-2", idx: 0 },
+    ]);
+
+    await store.deleteAfter(threadId, null);
+
+    await expect(counts(threadId)).resolves.toEqual({ checkpoints: 0, blobs: 0, writes: 0 });
+  });
+
+  it("follows metadata.parents and same-namespace parent links across subgraph namespaces", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: { messages: "1" } },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", parentCheckpointId: "cp-1", channelVersions: { messages: "2" } },
+      { threadId, checkpointNs: "", checkpointId: "cp-3", parentCheckpointId: "cp-2", channelVersions: { messages: "3" } },
+      {
+        threadId,
+        checkpointNs: "sub",
+        checkpointId: "cp-1",
+        channelVersions: { messages: "sub-1" },
+        metadata: { parents: { "": "cp-1" } },
+      },
+      {
+        threadId,
+        checkpointNs: "sub",
+        checkpointId: "cp-4",
+        channelVersions: { messages: "sub-4" },
+        metadata: { parents: { "": "cp-2" } },
+      },
+      { threadId, checkpointNs: "sub", checkpointId: "cp-5", parentCheckpointId: "cp-4", channelVersions: { messages: "sub-5" } },
+    ]);
+    await insertBlobs([
+      { threadId, checkpointNs: "", channel: "messages", version: "1" },
+      { threadId, checkpointNs: "", channel: "messages", version: "2" },
+      { threadId, checkpointNs: "", channel: "messages", version: "3" },
+      { threadId, checkpointNs: "sub", channel: "messages", version: "sub-1" },
+      { threadId, checkpointNs: "sub", channel: "messages", version: "sub-4" },
+      { threadId, checkpointNs: "sub", channel: "messages", version: "sub-5" },
+      { threadId, checkpointNs: "", channel: "messages", version: "orphan" },
+    ]);
+    await insertWrites([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", taskId: "task-1", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", taskId: "task-2", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "cp-3", taskId: "task-3", idx: 0 },
+      { threadId, checkpointNs: "sub", checkpointId: "cp-1", taskId: "task-sub-1", idx: 0 },
+      { threadId, checkpointNs: "sub", checkpointId: "cp-4", taskId: "task-sub-4", idx: 0 },
+      { threadId, checkpointNs: "sub", checkpointId: "cp-5", taskId: "task-sub-5", idx: 0 },
+    ]);
+
+    await store.deleteAfter(threadId, "cp-2");
+
+    const checkpoints = await pool.query(
+      "SELECT checkpoint_ns, checkpoint_id FROM checkpoints WHERE thread_id = $1 ORDER BY checkpoint_ns, checkpoint_id",
+      [threadId],
+    );
+    expect(checkpoints.rows).toEqual([
+      { checkpoint_ns: "", checkpoint_id: "cp-1" },
+      { checkpoint_ns: "", checkpoint_id: "cp-2" },
+      { checkpoint_ns: "sub", checkpoint_id: "cp-1" },
+    ]);
+
+    const blobs = await pool.query(
+      "SELECT checkpoint_ns, version FROM checkpoint_blobs WHERE thread_id = $1 ORDER BY checkpoint_ns, version",
+      [threadId],
+    );
+    expect(blobs.rows).toEqual([
+      { checkpoint_ns: "", version: "1" },
+      { checkpoint_ns: "", version: "2" },
+      { checkpoint_ns: "sub", version: "sub-1" },
+    ]);
+
+    const writes = await pool.query(
+      "SELECT checkpoint_ns, checkpoint_id FROM checkpoint_writes WHERE thread_id = $1 ORDER BY checkpoint_ns, checkpoint_id",
+      [threadId],
+    );
+    expect(writes.rows).toEqual([
+      { checkpoint_ns: "", checkpoint_id: "cp-1" },
+      { checkpoint_ns: "sub", checkpoint_id: "cp-1" },
+    ]);
+  });
+
+  it("deletes a subgraph checkpoint whose metadata.parents points at the base", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: {} },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", parentCheckpointId: "cp-1", channelVersions: {} },
+      {
+        threadId,
+        checkpointNs: "sub",
+        checkpointId: "sub-1",
+        channelVersions: {},
+        metadata: { parents: { "": "cp-2" } },
+      },
+    ]);
+
+    await store.deleteAfter(threadId, "cp-2");
+
+    const checkpoints = await pool.query(
+      "SELECT checkpoint_ns, checkpoint_id FROM checkpoints WHERE thread_id = $1 ORDER BY checkpoint_ns, checkpoint_id",
+      [threadId],
+    );
+    expect(checkpoints.rows).toEqual([
+      { checkpoint_ns: "", checkpoint_id: "cp-1" },
+      { checkpoint_ns: "", checkpoint_id: "cp-2" },
+    ]);
+  });
+
+  it("retains a subgraph checkpoint whose metadata.parents points at an earlier root", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: {} },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", parentCheckpointId: "cp-1", channelVersions: {} },
+      {
+        threadId,
+        checkpointNs: "sub",
+        checkpointId: "sub-1",
+        channelVersions: {},
+        metadata: { parents: { "": "cp-1" } },
+      },
+    ]);
+
+    await store.deleteAfter(threadId, "cp-2");
+
+    const checkpoints = await pool.query(
+      "SELECT checkpoint_ns, checkpoint_id FROM checkpoints WHERE thread_id = $1 ORDER BY checkpoint_ns, checkpoint_id",
+      [threadId],
+    );
+    expect(checkpoints.rows).toEqual([
+      { checkpoint_ns: "", checkpoint_id: "cp-1" },
+      { checkpoint_ns: "", checkpoint_id: "cp-2" },
+      { checkpoint_ns: "sub", checkpoint_id: "sub-1" },
+    ]);
+  });
+
+  it("deletes the rolled-back run's first-step write keyed to the base checkpoint", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: {} },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", parentCheckpointId: "cp-1", channelVersions: {} },
+    ]);
+    await insertWrites([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", taskId: "earlier-task", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", taskId: "rolled-back-task", idx: 0 },
+    ]);
+
+    await store.deleteAfter(threadId, "cp-2");
+
+    const writes = await pool.query(
+      "SELECT checkpoint_id, task_id FROM checkpoint_writes WHERE thread_id = $1 ORDER BY checkpoint_id",
+      [threadId],
+    );
+    expect(writes.rows).toEqual([{ checkpoint_id: "cp-1", task_id: "earlier-task" }]);
+  });
+
+  it("is idempotent when no descendants remain", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: { messages: "1" } },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", parentCheckpointId: "cp-1", channelVersions: { messages: "2" } },
+    ]);
+    await insertBlobs([
+      { threadId, checkpointNs: "", channel: "messages", version: "1" },
+      { threadId, checkpointNs: "", channel: "messages", version: "2" },
+    ]);
+    await insertWrites([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", taskId: "task-1", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "cp-2", taskId: "task-2", idx: 0 },
+    ]);
+
+    await store.deleteAfter(threadId, "cp-2");
+    await store.deleteAfter(threadId, "cp-2");
+
+    await expect(counts(threadId)).resolves.toEqual({ checkpoints: 2, blobs: 2, writes: 1 });
+  });
+
+  it("rejects an empty base checkpoint id instead of deleting every row", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: { messages: "1" } },
+    ]);
+    await insertBlobs([
+      { threadId, checkpointNs: "", channel: "messages", version: "1" },
+    ]);
+    await insertWrites([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", taskId: "task-1", idx: 0 },
+    ]);
+
+    await expect(store.deleteAfter(threadId, "")).rejects.toThrow(
+      "baseCheckpointId must be null or a non-empty checkpoint id",
+    );
+    await expect(counts(threadId)).resolves.toEqual({ checkpoints: 1, blobs: 1, writes: 1 });
+  });
+
+  it("deletes by parent lineage, not checkpoint-id ordering", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "z-earlier", channelVersions: { messages: "earlier" } },
+      { threadId, checkpointNs: "", checkpointId: "m-base", channelVersions: { messages: "base" } },
+      { threadId, checkpointNs: "", checkpointId: "a-unrelated", channelVersions: { messages: "unrelated" } },
+      { threadId, checkpointNs: "", checkpointId: "z-first", parentCheckpointId: "m-base", channelVersions: { messages: "first" } },
+      {
+        threadId,
+        checkpointNs: "sub",
+        checkpointId: "0-sub",
+        channelVersions: { messages: "sub" },
+        metadata: { parents: { "": "z-first" } },
+      },
+      { threadId, checkpointNs: "", checkpointId: "m-unrelated-child", parentCheckpointId: "a-unrelated", channelVersions: { messages: "unrelated-child" } },
+    ]);
+    await insertBlobs([
+      { threadId, checkpointNs: "", channel: "messages", version: "earlier" },
+      { threadId, checkpointNs: "", channel: "messages", version: "base" },
+      { threadId, checkpointNs: "", channel: "messages", version: "unrelated" },
+      { threadId, checkpointNs: "", channel: "messages", version: "first" },
+      { threadId, checkpointNs: "sub", channel: "messages", version: "sub" },
+      { threadId, checkpointNs: "", channel: "messages", version: "unrelated-child" },
+      { threadId, checkpointNs: "", channel: "messages", version: "orphan" },
+    ]);
+    await insertWrites([
+      { threadId, checkpointNs: "", checkpointId: "z-earlier", taskId: "earlier-task", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "m-base", taskId: "base-task", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "a-unrelated", taskId: "unrelated-task", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "z-first", taskId: "first-task", idx: 0 },
+      { threadId, checkpointNs: "sub", checkpointId: "0-sub", taskId: "sub-task", idx: 0 },
+      { threadId, checkpointNs: "", checkpointId: "m-unrelated-child", taskId: "unrelated-child-task", idx: 0 },
+    ]);
+
+    await store.deleteAfter(threadId, "m-base");
+
+    const checkpoints = await pool.query(
+      "SELECT checkpoint_ns, checkpoint_id FROM checkpoints WHERE thread_id = $1 ORDER BY checkpoint_ns, checkpoint_id",
+      [threadId],
+    );
+    expect(checkpoints.rows).toEqual([
+      { checkpoint_ns: "", checkpoint_id: "a-unrelated" },
+      { checkpoint_ns: "", checkpoint_id: "m-base" },
+      { checkpoint_ns: "", checkpoint_id: "m-unrelated-child" },
+      { checkpoint_ns: "", checkpoint_id: "z-earlier" },
+    ]);
+
+    const writes = await pool.query(
+      "SELECT checkpoint_ns, checkpoint_id FROM checkpoint_writes WHERE thread_id = $1 ORDER BY checkpoint_ns, checkpoint_id",
+      [threadId],
+    );
+    expect(writes.rows).toEqual([
+      { checkpoint_ns: "", checkpoint_id: "a-unrelated" },
+      { checkpoint_ns: "", checkpoint_id: "m-unrelated-child" },
+      { checkpoint_ns: "", checkpoint_id: "z-earlier" },
+    ]);
+
+    const blobs = await pool.query(
+      "SELECT checkpoint_ns, version FROM checkpoint_blobs WHERE thread_id = $1 ORDER BY checkpoint_ns, version",
+      [threadId],
+    );
+    expect(blobs.rows).toEqual([
+      { checkpoint_ns: "", version: "base" },
+      { checkpoint_ns: "", version: "earlier" },
+      { checkpoint_ns: "", version: "unrelated" },
+      { checkpoint_ns: "", version: "unrelated-child" },
+    ]);
+  });
+
+  it("fails closed when the base checkpoint does not exist", async () => {
+    await insertCheckpoints([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", channelVersions: { messages: "1" } },
+    ]);
+    await insertBlobs([
+      { threadId, checkpointNs: "", channel: "messages", version: "1" },
+    ]);
+    await insertWrites([
+      { threadId, checkpointNs: "", checkpointId: "cp-1", taskId: "task-1", idx: 0 },
+    ]);
+
+    await expect(store.deleteAfter(threadId, "missing-base")).rejects.toMatchObject({
+      name: "CheckpointBaseNotFoundError",
+      code: "CHECKPOINT_BASE_NOT_FOUND",
+    });
+    await expect(store.deleteAfter(threadId, "missing-base"))
+      .rejects.toBeInstanceOf(CheckpointBaseNotFoundError);
+    await expect(counts(threadId)).resolves.toEqual({ checkpoints: 1, blobs: 1, writes: 1 });
+  });
+
+  it("keeps namespace blobs when a retained checkpoint has no channel versions", async () => {
+    await pool.query(
+      `INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata)
+       VALUES ($1, '', 'cp-1', NULL, '{"v":4}'::jsonb, '{}'::jsonb),
+              ($1, '', 'cp-2', 'cp-1', '{"v":4}'::jsonb, '{}'::jsonb)`,
+      [threadId],
+    );
+    await insertBlobs([
+      { threadId, checkpointNs: "", channel: "messages", version: "1" },
+      { threadId, checkpointNs: "", channel: "messages", version: "2" },
+      { threadId, checkpointNs: "", channel: "messages", version: "3" },
+    ]);
+
+    await store.deleteAfter(threadId, "cp-1");
+
+    const blobs = await pool.query(
+      "SELECT version FROM checkpoint_blobs WHERE thread_id = $1 ORDER BY version",
+      [threadId],
+    );
+    expect(blobs.rows.map((row) => row.version)).toEqual(["1", "2", "3"]);
+  });
+});