@axiom-lattice/pg-stores 3.3.12 → 3.4.0

package/src/__tests__/PostgreSQLAgentRunTaskQueue.test.ts

@@ -0,0 +1,1375 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+// The @axiom-lattice/core barrel cannot load under jest in this package (ESM-only
+// transitive deps), so mock it with the real ledger module.
+jest.mock("@axiom-lattice/core", () => ({
+  ...jest.requireActual("../../../core/src/services/AgentRunLedger"),
+}));
+
+import {
+  AgentRunDuplicateMessageIdError,
+  AgentRunInputCommandConflictError,
+  AgentRunInvalidTaskIdError,
+  AgentRunTenantQuotaExceededError,
+  type CreatePendingAgentRunInput,
+} from "@axiom-lattice/core";
+import { PostgreSQLAgentRunLedger } from "../stores/PostgreSQLAgentRunLedger";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+const TENANT_PREFIX = "pg-agent-run-task-queue:";
+// Each Jest worker gets its own schema so global reads (`counts`) and cleanup
+// cannot observe another suite's rows when suites run in parallel.
+const SCHEMA = `agent_run_task_queue_test_${randomUUID().replace(/-/g, "").slice(0, 16)}`;
+
+function testTenantId(): string {
+  return `${TENANT_PREFIX}${randomUUID()}`;
+}
+
+maybeDescribe("PostgreSQLAgentRunLedger task queue (PostgreSQL)", () => {
+  let store: PostgreSQLAgentRunLedger;
+  let cleanupPool: Pool;
+
+  beforeAll(async () => {
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`CREATE SCHEMA ${SCHEMA}`);
+    await admin.end();
+
+    const poolConfig = {
+      connectionString: TEST_DATABASE_URL!,
+      options: `-c search_path=${SCHEMA}`,
+    };
+    store = new PostgreSQLAgentRunLedger({
+      poolConfig,
+      autoMigrate: false,
+    });
+    await store.initialize();
+    cleanupPool = new Pool(poolConfig);
+  });
+
+  beforeEach(async () => {
+    await cleanupPool.query(
+      "DELETE FROM lattice_agent_runs WHERE tenant_id LIKE $1",
+      [`${TENANT_PREFIX}%`],
+    );
+  });
+
+  afterAll(async () => {
+    await cleanupPool.end();
+    await store.dispose();
+
+    const admin = new Pool({ connectionString: TEST_DATABASE_URL! });
+    await admin.query(`DROP SCHEMA IF EXISTS ${SCHEMA} CASCADE`);
+    await admin.end();
+  });
+
+  function createInput(
+    overrides: Partial<CreatePendingAgentRunInput> = {},
+  ): CreatePendingAgentRunInput {
+    const tenantId = testTenantId();
+    const threadId = randomUUID();
+    return {
+      runId: randomUUID(),
+      tenantId,
+      workspaceId: randomUUID(),
+      projectId: randomUUID(),
+      assistantId: randomUUID(),
+      threadId,
+      input: { message: "hello", id: randomUUID() },
+      config: { temperature: 0 },
+      metadata: { origin: "test" },
+      ...overrides,
+    };
+  }
+
+  async function expire(runId: string): Promise<void> {
+    await cleanupPool.query(
+      `UPDATE lattice_agent_runs
+          SET lease_expires_at = now() - interval '1 second'
+        WHERE run_id = $1`,
+      [runId],
+    );
+  }
+
+  it("adds the task-queue columns with nullable pending ownership fields", async () => {
+    const columns = await cleanupPool.query<{
+      column_name: string;
+      is_nullable: string;
+    }>(
+      `SELECT column_name, is_nullable
+         FROM information_schema.columns
+        WHERE table_schema = current_schema()
+          AND table_name = 'lattice_agent_runs'`,
+    );
+    const byName = new Map(columns.rows.map((row) => [row.column_name, row]));
+
+    for (const column of [
+      "source",
+      "input",
+      "command",
+      "config",
+      "metadata",
+      "base_checkpoint_id",
+      "available_at",
+      "rollback_requested",
+      "task_id",
+    ]) {
+      expect(byName.has(column)).toBe(true);
+    }
+    for (const nullable of ["owner", "lease_expires_at", "started_at"]) {
+      expect(byName.get(nullable)?.is_nullable).toBe("YES");
+    }
+  });
+
+  it("keeps the partial unique running index and adds the queue scan index", async () => {
+    const indexes = await cleanupPool.query<{
+      indexname: string;
+      indexdef: string;
+    }>(
+      `SELECT indexname, indexdef
+         FROM pg_indexes
+        WHERE schemaname = current_schema()
+          AND tablename = 'lattice_agent_runs'`,
+    );
+    const byName = new Map(indexes.rows.map((row) => [row.indexname, row.indexdef]));
+
+    const runningIndex = byName.get("uq_lattice_agent_runs_one_running");
+    expect(runningIndex).toContain("(tenant_id, thread_id)");
+    expect(runningIndex).toContain("WHERE");
+    expect(runningIndex).toContain("running");
+    expect(runningIndex).toContain("run_api");
+    expect(runningIndex).toContain("project_task");
+    expect(byName.get("idx_lattice_agent_runs_run_queue")).toContain(
+      "(source, status, available_at, created_at)",
+    );
+
+    const projectTaskIndex = byName.get("idx_lattice_agent_runs_project_task");
+    expect(projectTaskIndex).toContain("(tenant_id, task_id, status, created_at DESC)");
+    expect(projectTaskIndex).toContain("WHERE");
+    expect(projectTaskIndex).toContain("task_id IS NOT NULL");
+  });
+
+  it("creates an unowned pending run_api row and projects it", async () => {
+    const input = createInput();
+    const record = await store.createPendingRun(input);
+
+    expect(record).toMatchObject({
+      runId: input.runId,
+      tenantId: input.tenantId,
+      workspaceId: input.workspaceId,
+      projectId: input.projectId,
+      assistantId: input.assistantId,
+      threadId: input.threadId,
+      status: "pending",
+      source: "run_api",
+      owner: null,
+      leaseExpiresAt: null,
+      startedAt: null,
+      finishedAt: null,
+      attempts: 0,
+      cancelRequested: false,
+      rollbackRequested: false,
+      baseCheckpointId: null,
+      input: input.input,
+      config: input.config,
+      metadata: input.metadata,
+      command: null,
+      inputMessageIds: [(input.input as { id: string }).id],
+    });
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "pending",
+    });
+    await expect(store.listRuns(input.tenantId, input.threadId)).resolves.toEqual([
+      expect.objectContaining({ runId: input.runId, status: "pending" }),
+    ]);
+    expect(record.availableAt).toBeGreaterThan(0);
+  });
+
+  it("persists a project_task run with taskId and filters pending runs by source", async () => {
+    const taskId = randomUUID();
+    const taskRun = createInput({ source: "project_task", taskId });
+    const apiRun = createInput();
+    await store.createPendingRun(taskRun);
+    await store.createPendingRun(apiRun);
+
+    const created = await store.getRun(taskRun.tenantId, taskRun.runId);
+    expect(created).toMatchObject({
+      source: "project_task",
+      taskId,
+      status: "pending",
+    });
+
+    const projectOnly = await store.listPendingRuns({
+      sources: ["project_task"],
+    });
+    expect(projectOnly.map((run) => run.runId)).toEqual([taskRun.runId]);
+
+    const defaultSources = await store.listPendingRuns();
+    expect(defaultSources.map((run) => run.runId)).toEqual([apiRun.runId]);
+
+    const both = await store.listPendingRuns({
+      sources: ["run_api", "project_task"],
+    });
+    expect(new Set(both.map((run) => run.runId))).toEqual(
+      new Set([taskRun.runId, apiRun.runId]),
+    );
+  });
+
+  it("requires taskId for project_task and rejects it on other sources", async () => {
+    await expect(
+      store.createPendingRun(createInput({ source: "project_task" })),
+    ).rejects.toBeInstanceOf(AgentRunInvalidTaskIdError);
+    await expect(
+      store.createPendingRun(
+        createInput({ source: "project_task", taskId: "   " }),
+      ),
+    ).rejects.toBeInstanceOf(AgentRunInvalidTaskIdError);
+    await expect(
+      store.createPendingRun(
+        createInput({ source: "run_api", taskId: randomUUID() }),
+      ),
+    ).rejects.toBeInstanceOf(AgentRunInvalidTaskIdError);
+    await expect(
+      store.createPendingRun(
+        createInput({ source: "legacy", taskId: randomUUID() }),
+      ),
+    ).rejects.toBeInstanceOf(AgentRunInvalidTaskIdError);
+
+    const taskId = randomUUID();
+    await expect(
+      store.createPendingRun(
+        createInput({ source: "project_task", taskId: ` ${taskId} ` }),
+      ),
+    ).resolves.toMatchObject({ source: "project_task", taskId: ` ${taskId} ` });
+  });
+
+  it("finds a pending run for a task and excludes terminal runs", async () => {
+    const taskId = randomUUID();
+    const input = createInput({ source: "project_task", taskId });
+    await store.createPendingRun(input);
+
+    await expect(
+      store.getActiveRunForTask(input.tenantId, taskId),
+    ).resolves.toMatchObject({ runId: input.runId, status: "pending" });
+
+    await store.claimRun({
+      runId: input.runId,
+      owner: "worker-1",
+      leaseMs: 100,
+      sources: ["project_task"],
+    });
+    await expire(input.runId);
+    await expect(
+      store.getActiveRunForTask(input.tenantId, taskId),
+    ).resolves.toMatchObject({ runId: input.runId, status: "running" });
+
+    await expect(
+      store.requestCancel(input.tenantId, input.runId),
+    ).resolves.toBe("cancelled");
+    await expect(
+      store.getActiveRunForTask(input.tenantId, taskId),
+    ).resolves.toBeUndefined();
+    await expect(
+      store.getActiveRunForTask(input.tenantId, randomUUID()),
+    ).resolves.toBeUndefined();
+  });
+
+  it("claims a project_task run only when its source is requested", async () => {
+    const taskRun = createInput({
+      source: "project_task",
+      taskId: randomUUID(),
+    });
+    await store.createPendingRun(taskRun);
+
+    await expect(
+      store.claimRun({ runId: taskRun.runId, owner: "worker-1", leaseMs: 60_000 }),
+    ).resolves.toBeUndefined();
+    await expect(
+      store.claimRun({
+        runId: taskRun.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        sources: ["run_api"],
+      }),
+    ).resolves.toBeUndefined();
+    await expect(
+      store.getRun(taskRun.tenantId, taskRun.runId),
+    ).resolves.toMatchObject({ source: "project_task", status: "pending" });
+
+    await expect(
+      store.claimRun({
+        runId: taskRun.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        sources: ["project_task"],
+      }),
+    ).resolves.toMatchObject({
+      runId: taskRun.runId,
+      source: "project_task",
+      taskId: taskRun.taskId,
+      status: "running",
+      owner: "worker-1",
+    });
+  });
+
+  it("keeps the one-running-run-per-thread guard generic across selected sources", async () => {
+    const taskRun = createInput({ source: "project_task", taskId: randomUUID() });
+    const apiRun = createInput({ tenantId: taskRun.tenantId, threadId: taskRun.threadId });
+    await store.createPendingRun(taskRun);
+    await store.createPendingRun(apiRun);
+
+    await expect(
+      store.claimRun({
+        runId: taskRun.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        sources: ["project_task"],
+      }),
+    ).resolves.toMatchObject({ status: "running" });
+
+    await expect(
+      store.claimRun({
+        runId: apiRun.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+        sources: ["run_api", "project_task"],
+      }),
+    ).resolves.toBeUndefined();
+    await expect(
+      store.getRun(apiRun.tenantId, apiRun.runId),
+    ).resolves.toMatchObject({ status: "pending" });
+
+    const reverseApi = createInput({ tenantId: taskRun.tenantId, threadId: randomUUID() });
+    const reverseTask = createInput({
+      tenantId: taskRun.tenantId,
+      threadId: reverseApi.threadId,
+      source: "project_task",
+      taskId: randomUUID(),
+    });
+    await store.createPendingRun(reverseApi);
+    await store.createPendingRun(reverseTask);
+    await expect(
+      store.claimRun({
+        runId: reverseApi.runId,
+        owner: "worker-3",
+        leaseMs: 60_000,
+        sources: ["run_api"],
+      }),
+    ).resolves.toMatchObject({ status: "running" });
+    await expect(
+      store.claimRun({
+        runId: reverseTask.runId,
+        owner: "worker-4",
+        leaseMs: 60_000,
+        sources: ["project_task"],
+      }),
+    ).resolves.toBeUndefined();
+  });
+
+  it("applies the per-tenant running cap across selected sources", async () => {
+    const first = createInput({ source: "project_task", taskId: randomUUID() });
+    const second = createInput({
+      tenantId: first.tenantId,
+      source: "project_task",
+      taskId: randomUUID(),
+    });
+    const apiRunning = createInput({ tenantId: first.tenantId });
+    await store.createPendingRun(apiRunning);
+    await store.claimRun({
+      runId: apiRunning.runId,
+      owner: "api-worker",
+      leaseMs: 60_000,
+      sources: ["run_api"],
+    });
+    await store.createPendingRun(first);
+    await store.createPendingRun(second);
+
+    // The running run_api row is not selected, so it must not count toward the
+    // project_task cap.
+    const capped = await store.listPendingRuns({
+      sources: ["project_task"],
+      maxRunningPerTenant: 1,
+    });
+    expect(new Set(capped.map((run) => run.runId))).toEqual(
+      new Set([first.runId, second.runId]),
+    );
+
+    await expect(
+      store.claimRun({
+        runId: first.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        sources: ["project_task"],
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toMatchObject({ status: "running" });
+    await expect(
+      store.claimRun({
+        runId: second.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+        sources: ["project_task"],
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toBeUndefined();
+    await expect(
+      store.listPendingRuns({
+        sources: ["project_task"],
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toEqual([]);
+  });
+
+  it("filters expired reaps by source and recovers an expired project_task run", async () => {
+    const taskRun = createInput({ source: "project_task", taskId: randomUUID() });
+    const apiRun = createInput();
+    await store.createPendingRun(taskRun);
+    await store.createPendingRun(apiRun);
+
+    await expect(
+      store.claimRun({
+        runId: taskRun.runId,
+        owner: "dead-owner",
+        leaseMs: 100,
+        sources: ["project_task"],
+      }),
+    ).resolves.toMatchObject({ status: "running" });
+    await expect(
+      store.claimRun({ runId: apiRun.runId, owner: "dead-owner", leaseMs: 100 }),
+    ).resolves.toMatchObject({ status: "running" });
+    await expire(taskRun.runId);
+    await expire(apiRun.runId);
+
+    const legacyRunId = randomUUID();
+    await store.beginRun({
+      runId: legacyRunId,
+      tenantId: taskRun.tenantId,
+      assistantId: taskRun.assistantId,
+      threadId: randomUUID(),
+      owner: "legacy-worker",
+      leaseMs: 60_000,
+      inputMessageIds: [],
+    });
+    await expire(legacyRunId);
+
+    const defaultReaped = await store.reapExpired();
+    expect(defaultReaped.map((run) => run.runId)).toContain(apiRun.runId);
+    expect(defaultReaped.map((run) => run.runId)).not.toContain(taskRun.runId);
+
+    const taskReaped = await store.reapExpired(["project_task"]);
+    expect(taskReaped.map((run) => run.runId)).toEqual([taskRun.runId]);
+    expect(taskReaped[0]).toMatchObject({
+      source: "project_task",
+      taskId: taskRun.taskId,
+    });
+
+    await expect(
+      store.recoverExpiredRun(taskRun.runId, "recovery-1", 60_000),
+    ).resolves.toMatchObject({
+      runId: taskRun.runId,
+      source: "project_task",
+      status: "running",
+      owner: "recovery-1",
+      attempts: 2,
+    });
+    await expire(taskRun.runId);
+    await expect(
+      store.settleExpiredRun(taskRun.runId, "interrupted", { error: "recovered" }),
+    ).resolves.toBe(true);
+    await expect(
+      store.getRun(taskRun.tenantId, taskRun.runId),
+    ).resolves.toMatchObject({ source: "project_task", status: "interrupted" });
+
+    const legacyReaped = await store.reapLegacyExpired?.();
+    expect(legacyReaped?.map((run) => run.runId)).toEqual([legacyRunId]);
+  });
+
+  it("rejects legacy as a queue source selection", async () => {
+    await expect(
+      store.listPendingRuns({ sources: ["legacy"] }),
+    ).rejects.toThrow("Legacy runs cannot be selected as queue sources");
+    await expect(
+      store.claimRun({
+        runId: randomUUID(),
+        owner: "worker-1",
+        leaseMs: 60_000,
+        sources: ["legacy"],
+      }),
+    ).rejects.toThrow("Legacy runs cannot be selected as queue sources");
+    await expect(
+      store.reapExpired(["legacy"]),
+    ).rejects.toThrow("Legacy runs cannot be selected as queue sources");
+  });
+
+  it("allows only one concurrent project_task claim on a thread", async () => {
+    const first = createInput({ source: "project_task", taskId: randomUUID() });
+    const second = createInput({
+      tenantId: first.tenantId,
+      threadId: first.threadId,
+      source: "project_task",
+      taskId: randomUUID(),
+    });
+    await store.createPendingRun(first);
+    await store.createPendingRun(second);
+
+    const [left, right] = await Promise.all([
+      store.claimRun({
+        runId: first.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        sources: ["project_task"],
+      }),
+      store.claimRun({
+        runId: second.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+        sources: ["project_task"],
+      }),
+    ]);
+
+    expect([left, right].filter((run) => run !== undefined)).toHaveLength(1);
+    const statuses = await Promise.all([
+      store.getRun(first.tenantId, first.runId),
+      store.getRun(second.tenantId, second.runId),
+    ]);
+    expect(statuses.map((run) => run?.status).sort()).toEqual([
+      "pending",
+      "running",
+    ]);
+  });
+
+  it("allows only one concurrent cross-source claim on a thread", async () => {
+    const apiRun = createInput();
+    const taskRun = createInput({
+      tenantId: apiRun.tenantId,
+      threadId: apiRun.threadId,
+      source: "project_task",
+      taskId: randomUUID(),
+    });
+    await store.createPendingRun(apiRun);
+    await store.createPendingRun(taskRun);
+
+    const [apiClaim, taskClaim] = await Promise.all([
+      store.claimRun({
+        runId: apiRun.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        sources: ["run_api", "project_task"],
+      }),
+      store.claimRun({
+        runId: taskRun.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+        sources: ["run_api", "project_task"],
+      }),
+    ]);
+
+    expect([apiClaim, taskClaim].filter((run) => run !== undefined)).toHaveLength(1);
+    const statuses = await Promise.all([
+      store.getRun(apiRun.tenantId, apiRun.runId),
+      store.getRun(taskRun.tenantId, taskRun.runId),
+    ]);
+    expect(statuses.map((run) => run?.status).sort()).toEqual([
+      "pending",
+      "running",
+    ]);
+  });
+
+  it("claims once, increments attempts and keeps one running row per thread", async () => {
+    const first = createInput();
+    const second = createInput({
+      tenantId: first.tenantId,
+      threadId: first.threadId,
+      input: { message: "second", id: randomUUID() },
+    });
+    await store.createPendingRun(first);
+    await store.createPendingRun(second);
+
+    const claimed = await store.claimRun({
+      runId: first.runId,
+      owner: "worker-1",
+      leaseMs: 60_000,
+    });
+    expect(claimed).toMatchObject({
+      status: "running",
+      owner: "worker-1",
+      attempts: 1,
+    });
+    expect(claimed?.leaseExpiresAt).toBeGreaterThan(Date.now());
+    expect(claimed?.startedAt).toBeGreaterThan(0);
+
+    await expect(
+      store.claimRun({
+        runId: second.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+      }),
+    ).resolves.toBeUndefined();
+    await expect(
+      store.claimRun({
+        runId: first.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+      }),
+    ).resolves.toBeUndefined();
+
+    await expect(store.findRunning(first.tenantId, first.threadId)).resolves.toMatchObject({
+      runId: first.runId,
+    });
+    await store.settleRun(first.runId, "worker-1", "success");
+    await expect(
+      store.claimRun({
+        runId: second.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+      }),
+    ).resolves.toMatchObject({ status: "running", runId: second.runId });
+  });
+
+  it("does not claim a run before its availableAt time", async () => {
+    const input = createInput({ availableAt: Date.now() + 60_000 });
+    await store.createPendingRun(input);
+
+    await expect(
+      store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 }),
+    ).resolves.toBeUndefined();
+
+    await cleanupPool.query(
+      `UPDATE lattice_agent_runs SET available_at = now() - interval '1 second'
+        WHERE run_id = $1`,
+      [input.runId],
+    );
+    await expect(
+      store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 }),
+    ).resolves.toMatchObject({ status: "running", owner: "worker-1" });
+  });
+
+  it("lets only one of two concurrent claims win", async () => {
+    const input = createInput();
+    await store.createPendingRun(input);
+
+    const claimed = await Promise.all([
+      store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 }),
+      store.claimRun({ runId: input.runId, owner: "worker-2", leaseMs: 60_000 }),
+    ]);
+
+    expect(claimed.filter((run) => run !== undefined)).toHaveLength(1);
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "running",
+      attempts: 1,
+    });
+  });
+
+  it("paginates pending runs with a keyset cursor", async () => {
+    const base = Date.now() - 3_000;
+    const first = createInput({ availableAt: base });
+    const second = createInput({ availableAt: base + 1_000 });
+    const third = createInput({ availableAt: base + 2_000 });
+    await store.createPendingRun(first);
+    await store.createPendingRun(second);
+    await store.createPendingRun(third);
+
+    const page = await store.listPendingRuns({ limit: 2 });
+    expect(page.map((run) => run.runId)).toEqual([first.runId, second.runId]);
+
+    const cursorRun = page[page.length - 1];
+    const next = await store.listPendingRuns({
+      limit: 2,
+      after: {
+        availableAt: cursorRun.availableAt ?? 0,
+        createdAt: cursorRun.createdAt,
+        runId: cursorRun.runId,
+      },
+    });
+    expect(next.map((run) => run.runId)).toEqual([third.runId]);
+  });
+
+  it("interleaves pending runs across tenants when fairTenants is enabled", async () => {
+    const tenantA = testTenantId();
+    const tenantB = testTenantId();
+    const base = Date.now() - 30_000;
+    const a1 = createInput({
+      tenantId: tenantA,
+      availableAt: base + 1_000,
+    });
+    const a2 = createInput({
+      tenantId: tenantA,
+      availableAt: base + 2_000,
+    });
+    const a3 = createInput({
+      tenantId: tenantA,
+      availableAt: base + 3_000,
+    });
+    const b1 = createInput({
+      tenantId: tenantB,
+      availableAt: base + 4_000,
+    });
+    const b2 = createInput({
+      tenantId: tenantB,
+      availableAt: base + 5_000,
+    });
+    for (const input of [a1, a2, a3, b1, b2]) {
+      await store.createPendingRun(input);
+    }
+
+    const normal = await store.listPendingRuns({ limit: 50 });
+    expect(normal.map((run) => run.runId)).toEqual(
+      [a1, a2, a3, b1, b2].map((input) => input.runId),
+    );
+
+    const fair = await store.listPendingRuns({
+      limit: 50,
+      fairTenants: true,
+    });
+    expect(fair.map((run) => run.runId)).toEqual(
+      [a1, b1, a2, b2, a3].map((input) => input.runId),
+    );
+  });
+
+  it("rejects a keyset cursor when fairTenants is enabled", async () => {
+    await expect(
+      store.listPendingRuns({
+        fairTenants: true,
+        after: { availableAt: Date.now(), createdAt: Date.now(), runId: randomUUID() },
+      }),
+    ).rejects.toBeInstanceOf(RangeError);
+  });
+
+  it("skips pending runs whose thread already has a running row", async () => {
+    const first = createInput();
+    const other = createInput();
+    await store.createPendingRun(first);
+    await store.createPendingRun(other);
+    await store.claimRun({ runId: first.runId, owner: "worker-1", leaseMs: 60_000 });
+    await store.createPendingRun(
+      createInput({
+        tenantId: first.tenantId,
+        threadId: first.threadId,
+        input: { message: "queued behind", id: randomUUID() },
+      }),
+    );
+
+    const pending = await store.listPendingRuns();
+    expect(pending.map((run) => run.runId)).toEqual([other.runId]);
+  });
+
+  it("blocks later runs on a thread with an unfinished rollback barrier", async () => {
+    const old = createInput();
+    await store.createPendingRun(old);
+    await store.claimRun({ runId: old.runId, owner: "worker-1", leaseMs: 60_000 });
+    await expect(store.requestRollback(old.tenantId, old.runId)).resolves.toBe("requested");
+    await store.settleRun(old.runId, "worker-1", "interrupted", "Rollback requested");
+
+    const later = createInput({
+      tenantId: old.tenantId,
+      threadId: old.threadId,
+      input: { message: "later", id: randomUUID() },
+    });
+    await store.createPendingRun(later);
+
+    await expect(store.listPendingRuns()).resolves.toEqual([]);
+    await expect(
+      store.claimRun({ runId: later.runId, owner: "worker-2", leaseMs: 60_000 }),
+    ).resolves.toBeUndefined();
+
+    await store.deleteRun(old.tenantId, old.runId, { status: "interrupted" });
+    await expect(
+      store.claimRun({ runId: later.runId, owner: "worker-2", leaseMs: 60_000 }),
+    ).resolves.toMatchObject({ status: "running", runId: later.runId });
+  });
+
+  it("rejects rollback of a legacy row without flagging it", async () => {
+    const input = createInput();
+    await store.beginRun({
+      runId: randomUUID(),
+      tenantId: input.tenantId,
+      assistantId: input.assistantId,
+      threadId: input.threadId,
+      owner: "legacy-worker",
+      leaseMs: 60_000,
+      inputMessageIds: [],
+    });
+    const legacy = (await store.listRuns(input.tenantId, input.threadId))[0];
+
+    await expect(store.requestRollback(input.tenantId, legacy.runId)).resolves.toBe(
+      "terminal",
+    );
+    await expect(store.getRun(input.tenantId, legacy.runId)).resolves.toMatchObject({
+      status: "running",
+      source: "legacy",
+      rollbackRequested: false,
+    });
+  });
+
+  it("rejects a pending run that provides both input and command", async () => {
+    const input = createInput({ command: { resume: "answer" } });
+    await expect(store.createPendingRun(input)).rejects.toBeInstanceOf(
+      AgentRunInputCommandConflictError,
+    );
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toBeUndefined();
+  });
+
+  it("rejects a duplicate input message id on the same thread", async () => {
+    const input = createInput({ input: { message: "hello", id: "duplicate-id" } });
+    await store.createPendingRun(input);
+
+    await expect(
+      store.createPendingRun(
+        createInput({
+          tenantId: input.tenantId,
+          threadId: input.threadId,
+          input: { message: "again", id: "duplicate-id" },
+        }),
+      ),
+    ).rejects.toBeInstanceOf(AgentRunDuplicateMessageIdError);
+
+    await expect(
+      store.createPendingRun(
+        createInput({
+          tenantId: input.tenantId,
+          threadId: randomUUID(),
+          input: { message: "other", id: "duplicate-id" },
+        }),
+      ),
+    ).resolves.toMatchObject({ status: "pending" });
+  });
+
+  it("isolates legacy rows from the pending queue", async () => {
+    const input = createInput();
+    const legacyThread = randomUUID();
+    await store.beginRun({
+      runId: randomUUID(),
+      tenantId: input.tenantId,
+      assistantId: input.assistantId,
+      threadId: legacyThread,
+      owner: "legacy-worker",
+      leaseMs: 60_000,
+      inputMessageIds: [randomUUID()],
+    });
+    await store.createPendingRun(input);
+
+    const pending = await store.listPendingRuns();
+    expect(pending.map((run) => run.runId)).toEqual([input.runId]);
+    const legacy = await store.listRuns(input.tenantId, legacyThread);
+    expect(legacy[0]).toMatchObject({ source: "legacy", status: "running" });
+  });
+
+  it("does not let a running legacy row block a run_api pending run on the same thread", async () => {
+    const input = createInput();
+    await store.beginRun({
+      runId: randomUUID(),
+      tenantId: input.tenantId,
+      assistantId: input.assistantId,
+      threadId: input.threadId,
+      owner: "legacy-worker",
+      leaseMs: 60_000,
+      inputMessageIds: [randomUUID()],
+    });
+    await store.createPendingRun(input);
+
+    await expect(store.listPendingRuns()).resolves.toEqual([
+      expect.objectContaining({ runId: input.runId, source: "run_api" }),
+    ]);
+    await expect(
+      store.claimRun({
+        runId: input.runId,
+        owner: "worker-1",
+        leaseMs: 5_000,
+        inputMessageIds: [input.input!.id as string],
+      }),
+    ).resolves.toMatchObject({ runId: input.runId, status: "running", owner: "worker-1" });
+  });
+
+  it("cancels pending and expired running rows as interrupted", async () => {
+    const pending = createInput();
+    await store.createPendingRun(pending);
+    await expect(store.requestCancel(pending.tenantId, pending.runId)).resolves.toBe(
+      "cancelled",
+    );
+    await expect(store.getRun(pending.tenantId, pending.runId)).resolves.toMatchObject({
+      status: "interrupted",
+      cancelRequested: true,
+    });
+
+    const running = createInput();
+    await store.createPendingRun(running);
+    await store.claimRun({ runId: running.runId, owner: "worker-1", leaseMs: 60_000 });
+    await expire(running.runId);
+    await expect(store.requestCancel(running.tenantId, running.runId)).resolves.toBe(
+      "cancelled",
+    );
+    await expect(store.getRun(running.tenantId, running.runId)).resolves.toMatchObject({
+      status: "interrupted",
+      cancelRequested: true,
+    });
+  });
+
+  it("deletes pending rollback rows and flags running rollback rows", async () => {
+    const pending = createInput();
+    await store.createPendingRun(pending);
+    await expect(store.requestRollback(pending.tenantId, pending.runId)).resolves.toBe(
+      "deleted",
+    );
+    await expect(store.getRun(pending.tenantId, pending.runId)).resolves.toBeUndefined();
+
+    const running = createInput();
+    await store.createPendingRun(running);
+    await store.claimRun({ runId: running.runId, owner: "worker-1", leaseMs: 60_000 });
+    await expect(store.requestRollback(running.tenantId, running.runId)).resolves.toBe(
+      "requested",
+    );
+    await expect(store.getRun(running.tenantId, running.runId)).resolves.toMatchObject({
+      status: "running",
+      rollbackRequested: true,
+      cancelRequested: true,
+    });
+  });
+
+  it("guards checkpoint and delete mutations by tenant, owner and state", async () => {
+    const input = createInput();
+    await store.createPendingRun(input);
+    await store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 });
+
+    await expect(
+      store.setBaseCheckpoint(input.tenantId, input.runId, "worker-2", "checkpoint-1"),
+    ).resolves.toBe(false);
+    await expect(
+      store.setBaseCheckpoint(randomUUID(), input.runId, "worker-1", "checkpoint-1"),
+    ).resolves.toBe(false);
+    await expect(
+      store.setBaseCheckpoint(input.tenantId, input.runId, "worker-1", "checkpoint-1"),
+    ).resolves.toBe(true);
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      baseCheckpointId: "checkpoint-1",
+    });
+
+    await expect(store.deleteRun(randomUUID(), input.runId)).resolves.toBe(false);
+    await expect(
+      store.deleteRun(input.tenantId, input.runId, { status: "pending" }),
+    ).resolves.toBe(false);
+    await expect(
+      store.deleteRun(input.tenantId, input.runId, { owner: "worker-2" }),
+    ).resolves.toBe(false);
+    await expect(
+      store.deleteRun(input.tenantId, input.runId, {
+        status: "running",
+        owner: "worker-1",
+      }),
+    ).resolves.toBe(true);
+  });
+
+  it("requeues an expired run atomically with settleExpiredRun", async () => {
+    const input = createInput();
+    await store.createPendingRun(input);
+    await store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 });
+    await store.setBaseCheckpoint(input.tenantId, input.runId, "worker-1", "checkpoint-1");
+    await expire(input.runId);
+    const availableAt = Date.now() + 500;
+
+    await expect(
+      store.settleExpiredRun(input.runId, "pending", { availableAt }),
+    ).resolves.toBe(true);
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "pending",
+      owner: null,
+      leaseExpiresAt: null,
+      startedAt: null,
+      availableAt,
+      baseCheckpointId: "checkpoint-1",
+      attempts: 1,
+      finishedAt: null,
+    });
+  });
+
+  it("reclaims an expired running run without requeueing it", async () => {
+    const input = createInput();
+    await store.createPendingRun(input);
+    await store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 });
+    await store.setBaseCheckpoint(input.tenantId, input.runId, "worker-1", "cp-base");
+    await expire(input.runId);
+
+    const recovered = await store.recoverExpiredRun(input.runId, "worker-2", 60_000);
+
+    expect(recovered).toMatchObject({
+      runId: input.runId,
+      status: "running",
+      source: "run_api",
+      owner: "worker-2",
+      attempts: 2,
+      baseCheckpointId: "cp-base",
+      baseCheckpointCaptured: true,
+    });
+    await expect(
+      store.recoverExpiredRun(input.runId, "worker-3", 60_000),
+    ).resolves.toBeUndefined();
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "running",
+      owner: "worker-2",
+      attempts: 2,
+    });
+  });
+
+  it("interrupts an expired run atomically with settleExpiredRun", async () => {
+    const input = createInput();
+    await store.createPendingRun(input);
+    await store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 });
+    await expire(input.runId);
+
+    await expect(
+      store.settleExpiredRun(input.runId, "interrupted", { error: "lease lost" }),
+    ).resolves.toBe(true);
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "interrupted",
+      error: "lease lost",
+      attempts: 1,
+      rollbackRequested: false,
+    });
+    await expect(
+      store.settleExpiredRun(input.runId, "interrupted"),
+    ).resolves.toBe(false);
+  });
+
+  it("only settleExpiredRun matches expired run_api running rows", async () => {
+    const live = createInput();
+    await store.createPendingRun(live);
+    await store.claimRun({ runId: live.runId, owner: "worker-1", leaseMs: 60_000 });
+
+    const legacy = createInput();
+    const legacyRunId = randomUUID();
+    await store.beginRun({
+      runId: legacyRunId,
+      tenantId: legacy.tenantId,
+      assistantId: legacy.assistantId,
+      threadId: legacy.threadId,
+      owner: "legacy-worker",
+      leaseMs: 60_000,
+      inputMessageIds: [],
+    });
+    await cleanupPool.query(
+      `UPDATE lattice_agent_runs
+          SET lease_expires_at = now() - interval '1 second'
+        WHERE run_id = $1`,
+      [legacyRunId],
+    );
+
+    await expect(store.settleExpiredRun(randomUUID(), "pending")).resolves.toBe(false);
+    await expect(store.settleExpiredRun(live.runId, "pending")).resolves.toBe(false);
+    await expect(store.settleExpiredRun(legacyRunId, "pending")).resolves.toBe(false);
+  });
+
+  it("reports only expired run_api rows without mutating them or legacy rows", async () => {
+    const input = createInput();
+    await store.createPendingRun(input);
+    await store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 });
+    await expire(input.runId);
+
+    const legacy = createInput();
+    await store.beginRun({
+      runId: randomUUID(),
+      tenantId: legacy.tenantId,
+      assistantId: legacy.assistantId,
+      threadId: legacy.threadId,
+      owner: "legacy-worker",
+      leaseMs: 60_000,
+      inputMessageIds: [],
+    });
+    await cleanupPool.query(
+      `UPDATE lattice_agent_runs SET lease_expires_at = now() - interval '1 second'
+        WHERE tenant_id = $1`,
+      [legacy.tenantId],
+    );
+
+    const reaped = await store.reapExpired();
+
+    expect(reaped).toEqual([
+      expect.objectContaining({
+        runId: input.runId,
+        source: "run_api",
+        attempts: 1,
+        cancelRequested: false,
+        rollbackRequested: false,
+      }),
+    ]);
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "running",
+    });
+    await expect(store.getRun(legacy.tenantId, (await store.listRuns(legacy.tenantId, legacy.threadId))[0].runId))
+      .resolves.toMatchObject({ status: "running" });
+
+    const legacyReaped = await store.reapLegacyExpired?.();
+    expect(legacyReaped).toEqual([
+      expect.objectContaining({ tenantId: legacy.tenantId }),
+    ]);
+  });
+
+  it("persists the capability revision and lists rollback work after interruption", async () => {
+    const input = createInput({ capabilityRuntimeRevision: "inherited-rev" });
+    const created = await store.createPendingRun(input);
+    expect(created.capabilityRuntimeRevision).toBe("inherited-rev");
+
+    await store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 });
+    await store.setBaseCheckpoint(input.tenantId, input.runId, "worker-1", "cp-base");
+    await expect(
+      store.setCapabilityRevision(input.tenantId, input.runId, "worker-1", "resolved-rev"),
+    ).resolves.toBe(true);
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      capabilityRuntimeRevision: "resolved-rev",
+      baseCheckpointId: "cp-base",
+      baseCheckpointCaptured: true,
+    });
+
+    await expect(store.requestRollback(input.tenantId, input.runId)).resolves.toBe("requested");
+    await expire(input.runId);
+
+    await expect(store.listRollbackRuns?.()).resolves.toEqual([
+      expect.objectContaining({
+        runId: input.runId,
+        rollbackRequested: true,
+        baseCheckpointId: "cp-base",
+      }),
+    ]);
+
+    await store.settleExpiredRun(input.runId, "interrupted", { error: "rewind failed" });
+    await expect(store.listRollbackRuns?.()).resolves.toEqual([
+      expect.objectContaining({
+        runId: input.runId,
+        rollbackRequested: true,
+        baseCheckpointId: "cp-base",
+      }),
+    ]);
+  });
+
+  it("records bounded rollback failures and converges terminally", async () => {
+    const input = createInput();
+    await store.createPendingRun(input);
+    await store.claimRun({ runId: input.runId, owner: "worker-1", leaseMs: 60_000 });
+    await expect(store.requestRollback(input.tenantId, input.runId)).resolves.toBe("requested");
+
+    await expect(
+      store.recordRollbackFailure(
+        input.tenantId,
+        input.runId,
+        "ROLLBACK_FAILED: nope",
+        2,
+      ),
+    ).resolves.toEqual({ attempts: 1, terminal: false });
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "running",
+      rollbackRequested: true,
+      rollbackAttempts: 1,
+      rollbackError: "ROLLBACK_FAILED: nope",
+    });
+
+    await expect(
+      store.recordRollbackFailure(
+        input.tenantId,
+        input.runId,
+        "ROLLBACK_FAILED: nope",
+        2,
+      ),
+    ).resolves.toEqual({ attempts: 2, terminal: true });
+    await expect(store.getRun(input.tenantId, input.runId)).resolves.toMatchObject({
+      status: "error",
+      rollbackRequested: false,
+      rollbackAttempts: 2,
+      error: "ROLLBACK_FAILED: nope",
+      rollbackError: "ROLLBACK_FAILED: nope",
+    });
+  });
+
+  it("guards success settlement against rollback and repairs uncaptured rollback", async () => {
+    const guarded = createInput();
+    await store.createPendingRun(guarded);
+    await store.claimRun({ runId: guarded.runId, owner: "worker-1", leaseMs: 60_000 });
+    await store.setBaseCheckpoint(guarded.tenantId, guarded.runId, "worker-1", "cp-base");
+    await expect(store.requestRollback(guarded.tenantId, guarded.runId)).resolves.toBe("requested");
+
+    await expect(
+      store.settleRunUnlessRollbackRequested(guarded.runId, "worker-1", "success"),
+    ).resolves.toBe(false);
+    await expect(store.getRun(guarded.tenantId, guarded.runId)).resolves.toMatchObject({
+      status: "running",
+      rollbackRequested: true,
+    });
+
+    const uncaptured = createInput();
+    await store.createPendingRun(uncaptured);
+    await store.claimRun({ runId: uncaptured.runId, owner: "worker-1", leaseMs: 60_000 });
+    await expect(store.requestRollback(uncaptured.tenantId, uncaptured.runId)).resolves.toBe("requested");
+
+    await expect(
+      store.settleRollbackFailure(
+        uncaptured.tenantId,
+        uncaptured.runId,
+        "Rollback not applied: base checkpoint was never captured",
+      ),
+    ).resolves.toBe(true);
+    await expect(store.getRun(uncaptured.tenantId, uncaptured.runId)).resolves.toMatchObject({
+      status: "error",
+      rollbackRequested: false,
+      error: expect.stringContaining("never captured"),
+    });
+  });
+
+  it("admits only one of two concurrent creates at the tenant limit", async () => {
+    const tenantId = testTenantId();
+    const first = createInput({ tenantId, maxActiveRunsForTenant: 1 });
+    const second = createInput({ tenantId, maxActiveRunsForTenant: 1 });
+
+    const results = await Promise.allSettled([
+      store.createPendingRun(first),
+      store.createPendingRun(second),
+    ]);
+
+    const fulfilled = results.filter((result) => result.status === "fulfilled");
+    const rejected = results.filter(
+      (result): result is PromiseRejectedResult => result.status === "rejected",
+    );
+    expect(fulfilled).toHaveLength(1);
+    expect(rejected).toHaveLength(1);
+    expect(rejected[0].reason).toBeInstanceOf(AgentRunTenantQuotaExceededError);
+    expect(rejected[0].reason).toMatchObject({
+      code: "RUN_TENANT_QUOTA_EXCEEDED",
+      tenantId,
+      limit: 1,
+    });
+
+    const rows = await cleanupPool.query<{ count: string }>(
+      `SELECT count(*)::text AS count
+         FROM lattice_agent_runs
+        WHERE tenant_id = $1 AND source = 'run_api'
+          AND status IN ('pending', 'running')`,
+      [tenantId],
+    );
+    expect(rows.rows[0].count).toBe("1");
+  });
+
+  it("frees tenant capacity when an active run reaches a terminal state", async () => {
+    const tenantId = testTenantId();
+    const first = createInput({ tenantId, maxActiveRunsForTenant: 1 });
+    await store.createPendingRun(first);
+
+    await expect(
+      store.createPendingRun(createInput({ tenantId, maxActiveRunsForTenant: 1 })),
+    ).rejects.toBeInstanceOf(AgentRunTenantQuotaExceededError);
+
+    await expect(store.requestCancel(tenantId, first.runId)).resolves.toBe("cancelled");
+
+    await expect(
+      store.createPendingRun(createInput({ tenantId, maxActiveRunsForTenant: 1 })),
+    ).resolves.toMatchObject({ tenantId, status: "pending" });
+  });
+
+  it("excludes pending rows for a tenant at the running cap and lists other tenants", async () => {
+    const tenantId = testTenantId();
+    const running = createInput({ tenantId });
+    await store.createPendingRun(running);
+    await expect(
+      store.claimRun({
+        runId: running.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+      }),
+    ).resolves.toMatchObject({ status: "running" });
+
+    const capped = createInput({ tenantId });
+    await store.createPendingRun(capped);
+    const other = createInput({ tenantId: testTenantId() });
+    await store.createPendingRun(other);
+
+    const pending = await store.listPendingRuns({
+      limit: 50,
+      maxRunningPerTenant: 1,
+      fairTenants: true,
+    });
+    const ids = pending.map((run) => run.runId);
+    expect(ids).toContain(other.runId);
+    expect(ids).not.toContain(capped.runId);
+    expect(ids).not.toContain(running.runId);
+
+    await expect(
+      store.claimRun({
+        runId: capped.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toBeUndefined();
+    await expect(
+      store.claimRun({
+        runId: other.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toMatchObject({ runId: other.runId, status: "running" });
+  });
+
+  it("never exceeds the running cap under concurrent claims", async () => {
+    const tenantId = testTenantId();
+    const first = createInput({ tenantId });
+    const second = createInput({ tenantId });
+    await store.createPendingRun(first);
+    await store.createPendingRun(second);
+
+    const results = await Promise.all([
+      store.claimRun({
+        runId: first.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        maxRunningPerTenant: 1,
+      }),
+      store.claimRun({
+        runId: second.runId,
+        owner: "worker-2",
+        leaseMs: 60_000,
+        maxRunningPerTenant: 1,
+      }),
+    ]);
+
+    expect(results.filter((run) => run !== undefined)).toHaveLength(1);
+    const rows = await cleanupPool.query<{ count: string }>(
+      `SELECT count(*)::text AS count
+         FROM lattice_agent_runs
+        WHERE tenant_id = $1 AND source = 'run_api' AND status = 'running'`,
+      [tenantId],
+    );
+    expect(rows.rows[0].count).toBe("1");
+  });
+
+  it("frees running capacity when a run reaches a terminal state", async () => {
+    const tenantId = testTenantId();
+    const first = createInput({ tenantId });
+    const second = createInput({ tenantId });
+    await store.createPendingRun(first);
+    await store.createPendingRun(second);
+    await expect(
+      store.claimRun({
+        runId: first.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toMatchObject({ status: "running" });
+    await expect(
+      store.claimRun({
+        runId: second.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toBeUndefined();
+
+    await expect(store.settleRun(first.runId, "worker-1", "success")).resolves.toBe(true);
+
+    await expect(
+      store.claimRun({
+        runId: second.runId,
+        owner: "worker-1",
+        leaseMs: 60_000,
+        maxRunningPerTenant: 1,
+      }),
+    ).resolves.toMatchObject({ runId: second.runId, status: "running" });
+  });
+});