@checkstack/ai-backend 0.1.0

package/src/propose-apply/store.ts

@@ -0,0 +1,224 @@
+import { and, eq, lt, gt } from "drizzle-orm";
+import type { SafeDatabase } from "@checkstack/backend-api";
+import type { AiToolEffect } from "@checkstack/ai-common";
+import * as schema from "../schema";
+import type { AiToolCallRow } from "../schema";
+import { generateProposalNonce } from "./token";
+
+type AiDatabase = SafeDatabase<typeof schema>;
+
+/** Default proposal-token TTL — 10 minutes (LOCKED, §13.4). */
+export const PROPOSAL_TTL_MS = 10 * 60 * 1000;
+
+export interface AuditPrincipal {
+  kind: "user" | "application";
+  id: string;
+}
+
+/**
+ * The durable audit log + propose/apply token store. A `proposed` row IS the
+ * token (decision §4, §13.4); there is no separate ephemeral table. All state
+ * is in shared Postgres, so a token proposed on one pod is consumable on any
+ * other (state-and-scale).
+ */
+export interface AiToolCallStore {
+  /** Record a directly-executed tool (read, or an automation-run mutate). */
+  recordExecuted(args: {
+    principal: AuditPrincipal;
+    transport: "chat" | "mcp" | "automation";
+    conversationId?: string;
+    toolName: string;
+    argsHash: string;
+    resultSnapshot?: Record<string, unknown>;
+  }): Promise<AiToolCallRow>;
+
+  /** Record a failed execute (audit only). */
+  recordFailed(args: {
+    principal: AuditPrincipal;
+    transport: "chat" | "mcp" | "automation";
+    conversationId?: string;
+    toolName: string;
+    effect: AiToolEffect;
+    argsHash: string;
+    error: string;
+  }): Promise<AiToolCallRow>;
+
+  /**
+   * Persist a `proposed` row (the token store). Returns the row plus the fresh
+   * nonce; the caller formats `propose:<id>.<nonce>`.
+   */
+  createProposal(args: {
+    principal: AuditPrincipal;
+    transport: "chat" | "mcp";
+    conversationId?: string;
+    toolName: string;
+    effect: AiToolEffect;
+    argsHash: string;
+    proposedPayload: Record<string, unknown>;
+    resultSnapshot?: Record<string, unknown>;
+    now?: Date;
+  }): Promise<{ row: AiToolCallRow; nonce: string }>;
+
+  /**
+   * Atomically consume a `proposed` row: a single
+   * `UPDATE ... WHERE id = ? AND status = 'proposed'` flips it to `applied` and
+   * RETURNS the row. A second apply finds `status != 'proposed'` and gets
+   * `undefined` (single-use even under concurrent calls). Returns `undefined`
+   * when the row is missing, already consumed, or the conditions don't hold.
+   *
+   * `applier` is the principal that actually called `apply` — stamped into
+   * `appliedByKind`/`appliedById` so the audit log records WHO applied, not just
+   * who proposed (P3 review item 1). Usually identical to the proposer.
+   */
+  consumeProposal(args: {
+    rowId: string;
+    applier: AuditPrincipal;
+    now?: Date;
+  }): Promise<AiToolCallRow | undefined>;
+
+  /** Fetch a proposal row by id without consuming it (for nonce/TTL checks). */
+  getProposal(rowId: string): Promise<AiToolCallRow | undefined>;
+
+  /**
+   * Sweep: flip expired `proposed` rows to `expired`, retaining them as audit
+   * history. Returns the number of rows expired.
+   */
+  expireStaleProposals(now?: Date): Promise<number>;
+}
+
+export function createAiToolCallStore({
+  db,
+}: {
+  db: AiDatabase;
+}): AiToolCallStore {
+  return {
+    async recordExecuted({
+      principal,
+      transport,
+      conversationId,
+      toolName,
+      argsHash,
+      resultSnapshot,
+    }) {
+      const [row] = await db
+        .insert(schema.aiToolCalls)
+        .values({
+          principalKind: principal.kind,
+          principalId: principal.id,
+          transport,
+          conversationId,
+          toolName,
+          effect: "read",
+          argsHash,
+          status: "executed",
+          resultSnapshot,
+        })
+        .returning();
+      return row;
+    },
+
+    async recordFailed({
+      principal,
+      transport,
+      conversationId,
+      toolName,
+      effect,
+      argsHash,
+      error,
+    }) {
+      const [row] = await db
+        .insert(schema.aiToolCalls)
+        .values({
+          principalKind: principal.kind,
+          principalId: principal.id,
+          transport,
+          conversationId,
+          toolName,
+          effect,
+          argsHash,
+          status: "failed",
+          error,
+        })
+        .returning();
+      return row;
+    },
+
+    async createProposal({
+      principal,
+      transport,
+      conversationId,
+      toolName,
+      effect,
+      argsHash,
+      proposedPayload,
+      resultSnapshot,
+      now = new Date(),
+    }) {
+      const nonce = generateProposalNonce();
+      const [row] = await db
+        .insert(schema.aiToolCalls)
+        .values({
+          principalKind: principal.kind,
+          principalId: principal.id,
+          transport,
+          conversationId,
+          toolName,
+          effect,
+          argsHash,
+          status: "proposed",
+          proposalNonce: nonce,
+          proposalExpiresAt: new Date(now.getTime() + PROPOSAL_TTL_MS),
+          proposedPayload,
+          resultSnapshot,
+          proposedAt: now,
+        })
+        .returning();
+      return { row, nonce };
+    },
+
+    async consumeProposal({ rowId, applier, now = new Date() }) {
+      // Atomic single-use: only a still-`proposed`, non-expired row transitions
+      // to `applied`. A concurrent second apply matches zero rows.
+      const [row] = await db
+        .update(schema.aiToolCalls)
+        .set({
+          status: "applied",
+          appliedAt: now,
+          appliedByKind: applier.kind,
+          appliedById: applier.id,
+        })
+        .where(
+          and(
+            eq(schema.aiToolCalls.id, rowId),
+            eq(schema.aiToolCalls.status, "proposed"),
+            gt(schema.aiToolCalls.proposalExpiresAt, now),
+          ),
+        )
+        .returning();
+      return row;
+    },
+
+    async getProposal(rowId) {
+      const rows = await db
+        .select()
+        .from(schema.aiToolCalls)
+        .where(eq(schema.aiToolCalls.id, rowId))
+        .limit(1);
+      return rows[0];
+    },
+
+    async expireStaleProposals(now = new Date()) {
+      const expired = await db
+        .update(schema.aiToolCalls)
+        .set({ status: "expired" })
+        .where(
+          and(
+            eq(schema.aiToolCalls.status, "proposed"),
+            lt(schema.aiToolCalls.proposalExpiresAt, now),
+          ),
+        )
+        .returning({ id: schema.aiToolCalls.id });
+      return expired.length;
+    },
+  };
+}

package/src/propose-apply/token.test.ts

@@ -0,0 +1,52 @@
+import { describe, expect, test } from "bun:test";
+import {
+  formatProposalToken,
+  generateProposalNonce,
+  nonceMatches,
+  parseProposalToken,
+} from "./token";
+
+describe("proposal token codec", () => {
+  test("format -> parse round-trips id and nonce", () => {
+    const rowId = "11111111-2222-3333-4444-555555555555";
+    const nonce = generateProposalNonce();
+    const token = formatProposalToken({ rowId, nonce });
+    expect(token.startsWith("propose:")).toBe(true);
+    const parsed = parseProposalToken(token);
+    expect(parsed).toEqual({ rowId, nonce });
+  });
+
+  test("nonce is 64 hex chars (32 random bytes)", () => {
+    const nonce = generateProposalNonce();
+    expect(nonce).toMatch(/^[0-9a-f]{64}$/);
+  });
+
+  test("rejects a token without the propose: prefix", () => {
+    expect(parseProposalToken("apply:abc.def")).toBeUndefined();
+  });
+
+  test("rejects a token with no separator", () => {
+    expect(parseProposalToken("propose:abcdef")).toBeUndefined();
+  });
+
+  test("rejects a token with an empty id", () => {
+    expect(parseProposalToken("propose:.abc")).toBeUndefined();
+  });
+
+  test("splits on the FIRST dot so a hex nonce is preserved whole", () => {
+    // Defensive: nonces are hex (no dots), but parsing must not lose data even
+    // if a future id format contained a dot in the nonce position.
+    const parsed = parseProposalToken("propose:row1.aa.bb");
+    expect(parsed).toEqual({ rowId: "row1", nonce: "aa.bb" });
+  });
+
+  test("nonceMatches is true for equal nonces, false otherwise", () => {
+    const nonce = generateProposalNonce();
+    expect(nonceMatches({ candidate: nonce, stored: nonce })).toBe(true);
+    expect(nonceMatches({ candidate: "deadbeef", stored: nonce })).toBe(false);
+  });
+
+  test("nonceMatches returns false on length mismatch without throwing", () => {
+    expect(nonceMatches({ candidate: "ab", stored: "abcd" })).toBe(false);
+  });
+});

package/src/propose-apply/token.ts

@@ -0,0 +1,71 @@
+import { randomBytes, timingSafeEqual } from "node:crypto";
+
+/**
+ * Proposal token codec (decision §13.4).
+ *
+ * A proposal is a row in `ai_tool_calls` with `status = "proposed"`. The opaque
+ * token handed to the caller is `propose:<rowId>.<nonce>`. `apply` parses it,
+ * fetches the row by id, and accepts only if the row is still `proposed`, the
+ * nonce matches in CONSTANT TIME, and the TTL has not elapsed.
+ *
+ * This module is pure (no DB) so the token grammar + constant-time compare are
+ * unit-tested in isolation.
+ */
+
+const TOKEN_PREFIX = "propose:";
+const NONCE_BYTES = 32;
+
+export interface ParsedProposalToken {
+  rowId: string;
+  nonce: string;
+}
+
+/** A fresh, cryptographically-random nonce (hex). */
+export function generateProposalNonce(): string {
+  return randomBytes(NONCE_BYTES).toString("hex");
+}
+
+/** Build the opaque token for a proposal row. */
+export function formatProposalToken({
+  rowId,
+  nonce,
+}: ParsedProposalToken): string {
+  return `${TOKEN_PREFIX}${rowId}.${nonce}`;
+}
+
+/**
+ * Parse an opaque proposal token. Returns `undefined` for any malformed token
+ * (wrong prefix, missing separator, empty id/nonce) so callers reject without
+ * branching on a thrown error.
+ */
+export function parseProposalToken(
+  token: string,
+): ParsedProposalToken | undefined {
+  if (!token.startsWith(TOKEN_PREFIX)) return undefined;
+  const body = token.slice(TOKEN_PREFIX.length);
+  // The rowId is a UUID (no dots); the nonce is hex (no dots). Split on the
+  // FIRST dot so a row id is never confused with the nonce.
+  const sep = body.indexOf(".");
+  if (sep <= 0) return undefined;
+  const rowId = body.slice(0, sep);
+  const nonce = body.slice(sep + 1);
+  if (!rowId || !nonce) return undefined;
+  return { rowId, nonce };
+}
+
+/**
+ * Constant-time nonce comparison. Length-mismatched / non-hex inputs return
+ * false without leaking timing about the stored nonce.
+ */
+export function nonceMatches({
+  candidate,
+  stored,
+}: {
+  candidate: string;
+  stored: string;
+}): boolean {
+  const a = Buffer.from(candidate, "utf8");
+  const b = Buffer.from(stored, "utf8");
+  if (a.length !== b.length) return false;
+  return timingSafeEqual(a, b);
+}

package/src/rate-limit/spend-ledger.it.test.ts

@@ -0,0 +1,224 @@
+/**
+ * Per-integration LLM spend-cap cross-pod enforcement (real Postgres) — Phase 6,
+ * state-and-scale §14.5.
+ *
+ * The spend cap is a shared-Postgres ROLLING-WINDOW SUM over `ai_spend`: the
+ * spend is the total tokens a principal has incurred against an integration in
+ * the trailing window. Because the sum is read from the SAME shared table that
+ * every pod writes to, the cap holds across ALL pods. An in-memory per-pod token
+ * counter would let N pods EACH allow the cap = N x the intended spend — a leak a
+ * single-process unit test can never catch. This test simulates TWO pods (two
+ * independent pools to the SAME schema) and asserts the combined token usage is
+ * counted against ONE shared cap (mirrors the tool-budget cross-pod IT).
+ *
+ * Gated behind `CHECKSTACK_IT=1`; connection from `CHECKSTACK_IT_PG_URL`. Runs in
+ * a freshly created, self-cleaning schema.
+ */
+import { afterAll, beforeAll, describe, expect, it } from "bun:test";
+import { drizzle } from "drizzle-orm/node-postgres";
+import { Pool } from "pg";
+import type { SafeDatabase } from "@checkstack/backend-api";
+import type { AiSpendCap } from "@checkstack/ai-common";
+import * as schema from "../schema";
+import type { AuditPrincipal } from "../propose-apply/store";
+import {
+  SpendCapExceededError,
+  checkSpendCap,
+  enforceSpendCap,
+  recordSpend,
+} from "./spend-ledger";
+
+const PG_URL =
+  process.env.CHECKSTACK_IT_PG_URL ??
+  "postgres://postgres:postgres@localhost:5432/postgres";
+
+const SCHEMA = `it_ai_spend_${crypto.randomUUID().replace(/-/g, "")}`;
+const INTEGRATION = "ai.openai-compatible.c1";
+
+interface Pod {
+  pool: Pool;
+  db: SafeDatabase<typeof schema>;
+  end(): Promise<void>;
+}
+
+function makePod(): Pod {
+  const pool = new Pool({
+    connectionString: PG_URL,
+    options: `-c search_path=${SCHEMA}`,
+  });
+  const db = drizzle(pool, { schema }) as unknown as SafeDatabase<typeof schema>;
+  return { pool, db, end: () => pool.end() };
+}
+
+describe.skipIf(!process.env.CHECKSTACK_IT)(
+  "per-integration spend cap (shared Postgres, cross-pod)",
+  () => {
+    let admin: Pool;
+    let podA: Pod;
+    let podB: Pod;
+
+    beforeAll(async () => {
+      admin = new Pool({ connectionString: PG_URL });
+      await admin.query(`CREATE SCHEMA "${SCHEMA}"`);
+      await admin.query(`
+        CREATE TABLE "${SCHEMA}".ai_spend (
+          id text PRIMARY KEY,
+          integration_id text NOT NULL,
+          principal_kind text NOT NULL,
+          principal_id text NOT NULL,
+          conversation_id text,
+          model text,
+          input_tokens integer NOT NULL DEFAULT 0,
+          output_tokens integer NOT NULL DEFAULT 0,
+          total_tokens integer NOT NULL DEFAULT 0,
+          created_at timestamp NOT NULL DEFAULT now()
+        )
+      `);
+      podA = makePod();
+      podB = makePod();
+    });
+
+    afterAll(async () => {
+      await podA?.end();
+      await podB?.end();
+      await admin.query(`DROP SCHEMA IF EXISTS "${SCHEMA}" CASCADE`);
+      await admin.end();
+    });
+
+    it("sums token usage from BOTH pods against ONE shared cap (no N x leak)", async () => {
+      const principal: AuditPrincipal = { kind: "user", id: "spend-u1" };
+      const cap: AiSpendCap = { tokenBudget: 1000, windowMinutes: 60 };
+
+      // Record turns on ALTERNATING pods. Whichever pod later checks the cap
+      // sees the COMBINED total, because the ledger is the shared table.
+      await recordSpend({
+        db: podA.db,
+        integrationId: INTEGRATION,
+        principal,
+        usage: { inputTokens: 200, outputTokens: 100 }, // 300
+      });
+      await recordSpend({
+        db: podB.db,
+        integrationId: INTEGRATION,
+        principal,
+        usage: { inputTokens: 250, outputTokens: 150 }, // 400 -> 700
+      });
+
+      // 700 < 1000: a new turn is within budget, read from EITHER pod.
+      const onA = await checkSpendCap({
+        db: podA.db,
+        integrationId: INTEGRATION,
+        principal,
+        cap,
+      });
+      const onB = await checkSpendCap({
+        db: podB.db,
+        integrationId: INTEGRATION,
+        principal,
+        cap,
+      });
+      expect(onA.used).toBe(700);
+      expect(onB.used).toBe(700);
+      expect(onA.allowed).toBe(true);
+      expect(onB.allowed).toBe(true);
+
+      // One more turn (on pod B) crosses the cap.
+      await recordSpend({
+        db: podB.db,
+        integrationId: INTEGRATION,
+        principal,
+        usage: { inputTokens: 200, outputTokens: 200 }, // 400 -> 1100
+      });
+
+      // Now BOTH pods agree the principal is over the SHARED cap — not 1000/pod.
+      const overA = await checkSpendCap({
+        db: podA.db,
+        integrationId: INTEGRATION,
+        principal,
+        cap,
+      });
+      const overB = await checkSpendCap({
+        db: podB.db,
+        integrationId: INTEGRATION,
+        principal,
+        cap,
+      });
+      expect(overA.used).toBe(1100);
+      expect(overB.used).toBe(1100);
+      expect(overA.allowed).toBe(false);
+      expect(overB.allowed).toBe(false);
+
+      await expect(
+        enforceSpendCap({
+          db: podA.db,
+          integrationId: INTEGRATION,
+          principal,
+          cap,
+        }),
+      ).rejects.toBeInstanceOf(SpendCapExceededError);
+      await expect(
+        enforceSpendCap({
+          db: podB.db,
+          integrationId: INTEGRATION,
+          principal,
+          cap,
+        }),
+      ).rejects.toBeInstanceOf(SpendCapExceededError);
+    });
+
+    it("the cap is per-integration: spend on one integration does not block another", async () => {
+      const principal: AuditPrincipal = { kind: "user", id: "spend-u2" };
+      const cap: AiSpendCap = { tokenBudget: 100, windowMinutes: 60 };
+      await recordSpend({
+        db: podA.db,
+        integrationId: "ai.openai-compatible.c1",
+        principal,
+        usage: { inputTokens: 200, outputTokens: 0 }, // over on c1
+      });
+      // c1 is over; c2 has no spend, so a turn there is still allowed.
+      await expect(
+        enforceSpendCap({
+          db: podB.db,
+          integrationId: "ai.openai-compatible.c1",
+          principal,
+          cap,
+        }),
+      ).rejects.toBeInstanceOf(SpendCapExceededError);
+      await expect(
+        enforceSpendCap({
+          db: podB.db,
+          integrationId: "ai.openai-compatible.c2",
+          principal,
+          cap,
+        }),
+      ).resolves.toBeUndefined();
+    });
+
+    it("the window slides: only recent spend counts (older turns drop out)", async () => {
+      const principal: AuditPrincipal = { kind: "user", id: "spend-window" };
+      // An OLD turn outside the window.
+      await podA.pool.query(
+        `INSERT INTO "${SCHEMA}".ai_spend
+           (id, integration_id, principal_kind, principal_id, total_tokens, created_at)
+         VALUES ($1, $2, 'user', $3, 5000, now() - interval '120 minutes')`,
+        [crypto.randomUUID(), INTEGRATION, principal.id],
+      );
+      // A recent one.
+      await recordSpend({
+        db: podB.db,
+        integrationId: INTEGRATION,
+        principal,
+        usage: { inputTokens: 10, outputTokens: 10 },
+      });
+      // With a 60-minute window the old 5000-token turn is excluded.
+      const result = await checkSpendCap({
+        db: podB.db,
+        integrationId: INTEGRATION,
+        principal,
+        cap: { tokenBudget: 1000, windowMinutes: 60 },
+      });
+      expect(result.used).toBe(20);
+      expect(result.allowed).toBe(true);
+    });
+  },
+);