@checkstack/ai-backend 0.1.0

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

@@ -0,0 +1,176 @@
+import { describe, expect, test, mock } from "bun:test";
+import type { AiSpendCap } from "@checkstack/ai-common";
+import type { AuditPrincipal } from "../propose-apply/store";
+import {
+  checkSpendCap,
+  enforceSpendCap,
+  recordSpend,
+  SpendCapExceededError,
+} from "./spend-ledger";
+
+/**
+ * Per-integration LLM spend cap (Phase 6, the locked-off knob, now wired).
+ *
+ * Unit-level proof that the cap is a shared-Postgres ROLLING-WINDOW SUM over
+ * `ai_spend`: the check reads the SUM through the injected `db` (no pod-local
+ * counter), enforcement throws once `used >= tokenBudget`, the cap is OFF when
+ * unconfigured, and a turn's usage is recorded as input + output tokens. The
+ * cross-pod property (the sum counts EVERY pod's writes) is exercised against a
+ * live Postgres in `spend-ledger.it.test.ts`.
+ */
+
+const principal: AuditPrincipal = { kind: "user", id: "u1" };
+const cap: AiSpendCap = { tokenBudget: 1000, windowMinutes: 60 };
+
+/** A fake db whose SUM query resolves to `total`, simulating the ledger read. */
+function sumDb(total: number) {
+  const where = mock(() => Promise.resolve([{ total }]));
+  const from = mock(() => ({ where }));
+  const select = mock(() => ({ from }));
+  return { select } as never;
+}
+
+/** A fake db capturing the inserted spend row. */
+function insertDb() {
+  const captured: Array<Record<string, unknown>> = [];
+  const values = mock((v: Record<string, unknown>) => {
+    captured.push(v);
+    return Promise.resolve([]);
+  });
+  const insert = mock(() => ({ values }));
+  return { db: { insert } as never, captured };
+}
+
+describe("checkSpendCap (rolling-window SUM over ai_spend)", () => {
+  test("under budget: a new turn is allowed", async () => {
+    const result = await checkSpendCap({
+      db: sumDb(400),
+      integrationId: "ai.openai-compatible.c1",
+      principal,
+      cap,
+    });
+    expect(result.used).toBe(400);
+    expect(result.allowed).toBe(true);
+    expect(result.tokenBudget).toBe(1000);
+  });
+
+  test("at/over budget: a new turn is refused (used >= tokenBudget)", async () => {
+    const atCap = await checkSpendCap({
+      db: sumDb(1000),
+      integrationId: "c1",
+      principal,
+      cap,
+    });
+    expect(atCap.allowed).toBe(false);
+    const over = await checkSpendCap({
+      db: sumDb(1500),
+      integrationId: "c1",
+      principal,
+      cap,
+    });
+    expect(over.allowed).toBe(false);
+  });
+
+  test("a numeric-string SUM (as pg returns it) coerces to a number", async () => {
+    // pg returns sum() over an integer column as a string; the ledger coerces.
+    const where = mock(() => Promise.resolve([{ total: "750" }]));
+    const from = mock(() => ({ where }));
+    const db = { select: mock(() => ({ from })) } as never;
+    const result = await checkSpendCap({
+      db,
+      integrationId: "c1",
+      principal,
+      cap,
+    });
+    expect(result.used).toBe(750);
+    expect(result.allowed).toBe(true);
+  });
+
+  test("the window start is now - windowMinutes (the rolling window)", async () => {
+    const now = new Date("2026-06-02T12:00:00Z");
+    const result = await checkSpendCap({
+      db: sumDb(0),
+      integrationId: "c1",
+      principal,
+      cap: { tokenBudget: 100, windowMinutes: 30 },
+      now,
+    });
+    expect(result.windowStart).toEqual(new Date("2026-06-02T11:30:00Z"));
+  });
+});
+
+describe("enforceSpendCap", () => {
+  test("OFF by default: no cap configured = a no-op (never throws, never reads)", async () => {
+    const select = mock(() => ({ from: () => ({ where: () => [] }) }));
+    const db = { select } as never;
+    await expect(
+      enforceSpendCap({ db, integrationId: "c1", principal, cap: undefined }),
+    ).resolves.toBeUndefined();
+    // Default OFF means we never even hit the ledger.
+    expect(select).not.toHaveBeenCalled();
+  });
+
+  test("throws SpendCapExceededError when over the configured budget", async () => {
+    await expect(
+      enforceSpendCap({
+        db: sumDb(1200),
+        integrationId: "c1",
+        principal,
+        cap,
+      }),
+    ).rejects.toBeInstanceOf(SpendCapExceededError);
+  });
+
+  test("passes when under the configured budget", async () => {
+    await expect(
+      enforceSpendCap({ db: sumDb(10), integrationId: "c1", principal, cap }),
+    ).resolves.toBeUndefined();
+  });
+});
+
+describe("recordSpend (append to the shared ledger)", () => {
+  test("persists input + output + total tokens for the turn", async () => {
+    const { db, captured } = insertDb();
+    await recordSpend({
+      db,
+      integrationId: "ai.openai-compatible.c1",
+      principal,
+      conversationId: "conv-1",
+      model: "gpt-4o-mini",
+      usage: { inputTokens: 120, outputTokens: 80 },
+    });
+    expect(captured).toHaveLength(1);
+    const row = captured[0];
+    expect(row?.integrationId).toBe("ai.openai-compatible.c1");
+    expect(row?.principalKind).toBe("user");
+    expect(row?.inputTokens).toBe(120);
+    expect(row?.outputTokens).toBe(80);
+    expect(row?.totalTokens).toBe(200);
+    expect(row?.conversationId).toBe("conv-1");
+    expect(row?.model).toBe("gpt-4o-mini");
+  });
+
+  test("missing / negative usage coerces to 0 (a provider omitting usage never corrupts the ledger)", async () => {
+    const { db, captured } = insertDb();
+    await recordSpend({
+      db,
+      integrationId: "c1",
+      principal,
+      usage: { inputTokens: Number.NaN, outputTokens: -5 },
+    });
+    const row = captured[0];
+    expect(row?.inputTokens).toBe(0);
+    expect(row?.outputTokens).toBe(0);
+    expect(row?.totalTokens).toBe(0);
+  });
+});
+
+describe("SpendCapExceededError", () => {
+  test("carries used + cap and a clear budget-exceeded message", () => {
+    const err = new SpendCapExceededError(1200, 1000);
+    expect(err.used).toBe(1200);
+    expect(err.tokenBudget).toBe(1000);
+    expect(err.message).toContain("spend/budget exceeded");
+    expect(err.name).toBe("SpendCapExceededError");
+  });
+});

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

@@ -0,0 +1,162 @@
+import { and, eq, gte, sql } from "drizzle-orm";
+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";
+
+type AiDatabase = SafeDatabase<typeof schema>;
+
+/**
+ * Per-integration LLM SPEND CAP (Phase 6, the locked-off knob, now wired).
+ *
+ * A shared-Postgres ROLLING-WINDOW SUM over `ai_spend`: the spend is the total
+ * number of tokens this principal has incurred against THIS 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 (state-and-scale §14.5) — an
+ * in-memory per-pod token counter would let N pods each allow the cap = N x the
+ * intended spend, which a single-process test can never catch. Mirrors the
+ * per-principal tool rate-limit budget exactly, reusing the
+ * `ai_spend_integration_principal_created_idx` index.
+ *
+ * Token-count (not USD) is deliberate: deterministic + provider-agnostic. Every
+ * OpenAI-compatible provider reports tokens via the AI SDK `usage`; only some
+ * publish a price table (self-hosted Ollama/vLLM have none), so a USD cap would
+ * require a per-model pricing table that drifts and is meaningless for local
+ * models. Tokens are the one unit every provider reports.
+ *
+ * The cap is OFF by default: enforcement only runs when a connection configures
+ * `spendCap`. With no cap configured, `enforceSpendCap` is a no-op.
+ */
+
+export interface SpendUsageInput {
+  inputTokens: number;
+  outputTokens: number;
+}
+
+export interface SpendCheckResult {
+  /** Total tokens already used by this principal+integration in the window. */
+  used: number;
+  /** Whether a NEW turn is within budget (`used < tokenBudget`). */
+  allowed: boolean;
+  /** The configured token budget (the cap). */
+  tokenBudget: number;
+  /** Trailing window start the sum was taken against. */
+  windowStart: Date;
+}
+
+/**
+ * Sum the principal's token usage against an integration over the cap's trailing
+ * window and decide whether a new turn is within budget. Returns
+ * `allowed: false` once `used >= tokenBudget`.
+ */
+export async function checkSpendCap({
+  db,
+  integrationId,
+  principal,
+  cap,
+  now = new Date(),
+}: {
+  db: AiDatabase;
+  integrationId: string;
+  principal: AuditPrincipal;
+  cap: AiSpendCap;
+  now?: Date;
+}): Promise<SpendCheckResult> {
+  const windowStart = new Date(now.getTime() - cap.windowMinutes * 60_000);
+
+  const rows = await db
+    .select({
+      total: sql<number>`coalesce(sum(${schema.aiSpend.totalTokens}), 0)`,
+    })
+    .from(schema.aiSpend)
+    .where(
+      and(
+        eq(schema.aiSpend.integrationId, integrationId),
+        eq(schema.aiSpend.principalKind, principal.kind),
+        eq(schema.aiSpend.principalId, principal.id),
+        gte(schema.aiSpend.createdAt, windowStart),
+      ),
+    );
+
+  // `sum` over an integer column comes back as a numeric string from pg; coerce.
+  const used = Number(rows[0]?.total ?? 0);
+  return {
+    used,
+    allowed: used < cap.tokenBudget,
+    tokenBudget: cap.tokenBudget,
+    windowStart,
+  };
+}
+
+/** Error thrown when a principal is over an integration's LLM spend cap. */
+export class SpendCapExceededError extends Error {
+  constructor(
+    public readonly used: number,
+    public readonly tokenBudget: number,
+  ) {
+    super(
+      `LLM spend/budget exceeded: ${used} tokens used against this integration in the current window (cap ${tokenBudget}). Try again later or raise the integration's spend cap.`,
+    );
+    this.name = "SpendCapExceededError";
+  }
+}
+
+/**
+ * Enforce the cap BEFORE a chat turn: throws {@link SpendCapExceededError} when
+ * the principal is over the integration's configured budget. A no-op when no
+ * `cap` is configured (default OFF — no cap unless set on the connection).
+ */
+export async function enforceSpendCap({
+  db,
+  integrationId,
+  principal,
+  cap,
+  now,
+}: {
+  db: AiDatabase;
+  integrationId: string;
+  principal: AuditPrincipal;
+  cap?: AiSpendCap;
+  now?: Date;
+}): Promise<void> {
+  if (!cap) return; // default OFF: no cap unless configured.
+  const result = await checkSpendCap({ db, integrationId, principal, cap, now });
+  if (!result.allowed) {
+    throw new SpendCapExceededError(result.used, result.tokenBudget);
+  }
+}
+
+/**
+ * Record one turn's token usage into the shared `ai_spend` ledger (append-only).
+ * Called from the agent loop's `onFinish` with the AI-SDK `totalUsage`. Negative
+ * / undefined token counts coerce to 0 so a provider that omits usage never
+ * corrupts the ledger.
+ */
+export async function recordSpend({
+  db,
+  integrationId,
+  principal,
+  conversationId,
+  model,
+  usage,
+}: {
+  db: AiDatabase;
+  integrationId: string;
+  principal: AuditPrincipal;
+  conversationId?: string;
+  model?: string;
+  usage: SpendUsageInput;
+}): Promise<void> {
+  const inputTokens = Math.max(0, Math.trunc(usage.inputTokens || 0));
+  const outputTokens = Math.max(0, Math.trunc(usage.outputTokens || 0));
+  await db.insert(schema.aiSpend).values({
+    integrationId,
+    principalKind: principal.kind,
+    principalId: principal.id,
+    conversationId,
+    model,
+    inputTokens,
+    outputTokens,
+    totalTokens: inputTokens + outputTokens,
+  });
+}

package/src/rate-limit/tool-budget.it.test.ts

@@ -0,0 +1,173 @@
+/**
+ * Per-principal tool-budget cross-pod enforcement (real Postgres) — matrix #17,
+ * plan §14.5, state-and-scale §9.
+ *
+ * The per-principal tool budget is a shared-Postgres ROLLING-WINDOW counter over
+ * `ai_tool_calls`: the budget is the number of rows the principal has written in
+ * the trailing window. Because the count is read from the SAME shared table that
+ * every pod writes to, the cap holds across ALL pods. An in-memory per-pod
+ * limiter would let N pods EACH allow the cap = N x the intended limit — 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 activity is
+ * counted against ONE shared budget (mirrors the P2 DCR 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 * as schema from "../schema";
+import type { AuditPrincipal } from "../propose-apply/store";
+import {
+  ToolBudgetExceededError,
+  checkToolBudget,
+  enforceToolBudget,
+} from "./tool-budget";
+
+const PG_URL =
+  process.env.CHECKSTACK_IT_PG_URL ??
+  "postgres://postgres:postgres@localhost:5432/postgres";
+
+const SCHEMA = `it_ai_budget_${crypto.randomUUID().replace(/-/g, "")}`;
+
+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() };
+}
+
+/** Insert one `ai_tool_calls` row (a recorded tool invocation) on a given pod. */
+async function recordCall(pod: Pod, principal: AuditPrincipal): Promise<void> {
+  await pod.pool.query(
+    `INSERT INTO "${SCHEMA}".ai_tool_calls
+       (id, principal_kind, principal_id, transport, tool_name, effect, args_hash, status)
+     VALUES ($1, $2, $3, 'mcp', 'incident.list', 'read', 'h', 'executed')`,
+    [crypto.randomUUID(), principal.kind, principal.id],
+  );
+}
+
+describe.skipIf(!process.env.CHECKSTACK_IT)(
+  "per-principal tool budget (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}"`);
+      // The minimal columns checkToolBudget counts over (principal + createdAt).
+      await admin.query(`
+        CREATE TABLE "${SCHEMA}".ai_tool_calls (
+          id text PRIMARY KEY,
+          principal_kind text NOT NULL,
+          principal_id text NOT NULL,
+          transport text NOT NULL,
+          conversation_id text,
+          tool_name text NOT NULL,
+          effect text NOT NULL,
+          args_hash text NOT NULL,
+          status text NOT NULL,
+          proposal_nonce text,
+          proposal_expires_at timestamp,
+          applied_by_kind text,
+          applied_by_id text,
+          result_snapshot jsonb,
+          proposed_payload jsonb,
+          error text,
+          proposed_at timestamp,
+          applied_at timestamp,
+          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("counts activity from BOTH pods against ONE shared budget (no N x leak)", async () => {
+      const principal: AuditPrincipal = { kind: "user", id: "budget-u1" };
+      const max = 4;
+
+      // Alternate the pod that records each call. Whichever pod later checks the
+      // budget sees the COMBINED count, because the counter is the shared table.
+      await recordCall(podA, principal);
+      await recordCall(podB, principal);
+      await recordCall(podA, principal);
+
+      // 3 recorded < 4: a new call is within budget, read from EITHER pod.
+      const onA = await checkToolBudget({ db: podA.db, principal, max });
+      const onB = await checkToolBudget({ db: podB.db, principal, max });
+      expect(onA.used).toBe(3);
+      expect(onB.used).toBe(3);
+      expect(onA.allowed).toBe(true);
+      expect(onB.allowed).toBe(true);
+
+      // One more recorded call (on pod B) reaches the cap.
+      await recordCall(podB, principal);
+
+      // Now BOTH pods agree the principal is at/over budget — not 4-per-pod.
+      const overA = await checkToolBudget({ db: podA.db, principal, max });
+      const overB = await checkToolBudget({ db: podB.db, principal, max });
+      expect(overA.used).toBe(4);
+      expect(overB.used).toBe(4);
+      expect(overA.allowed).toBe(false);
+      expect(overB.allowed).toBe(false);
+    });
+
+    it("enforceToolBudget throws cross-pod once the shared cap is met", async () => {
+      const principal: AuditPrincipal = { kind: "application", id: "budget-app" };
+      const max = 2;
+      await recordCall(podA, principal);
+      await recordCall(podB, principal);
+
+      // Within-budget check on a fresh principal stays under: assert the throw is
+      // driven by the SHARED count, not pod-local state.
+      await expect(
+        enforceToolBudget({ db: podB.db, principal, max }),
+      ).rejects.toBeInstanceOf(ToolBudgetExceededError);
+      await expect(
+        enforceToolBudget({ db: podA.db, principal, max }),
+      ).rejects.toBeInstanceOf(ToolBudgetExceededError);
+    });
+
+    it("the window slides: only recent calls count (older ones drop out)", async () => {
+      const principal: AuditPrincipal = { kind: "user", id: "budget-window" };
+      // Insert a row with an OLD created_at (outside the window).
+      await podA.pool.query(
+        `INSERT INTO "${SCHEMA}".ai_tool_calls
+           (id, principal_kind, principal_id, transport, tool_name, effect, args_hash, status, created_at)
+         VALUES ($1, 'user', $2, 'mcp', 'incident.list', 'read', 'h', 'executed', now() - interval '10 minutes')`,
+        [crypto.randomUUID(), principal.id],
+      );
+      // A recent one.
+      await recordCall(podB, principal);
+
+      // With a 1-minute window the old row is excluded; only the recent counts.
+      const result = await checkToolBudget({
+        db: podB.db,
+        principal,
+        max: 60,
+        windowMs: 60_000,
+      });
+      expect(result.used).toBe(1);
+    });
+  },
+);

package/src/rate-limit/tool-budget.test.ts

@@ -0,0 +1,58 @@
+import { describe, expect, test, mock } from "bun:test";
+import {
+  checkToolBudget,
+  enforceToolBudget,
+  ToolBudgetExceededError,
+  TOOL_BUDGET_MAX_CALLS,
+} from "./tool-budget";
+
+/** Build a fake db whose count query resolves to `value`. */
+function dbReturning(value: number) {
+  const where = mock(() => Promise.resolve([{ value }]));
+  const from = mock(() => ({ where }));
+  const select = mock(() => ({ from }));
+  return { db: { select } as never, where };
+}
+
+describe("per-principal tool budget (§14.5, P3 review item 3)", () => {
+  test("under budget is allowed; the count is read over a trailing window", async () => {
+    const { db, where } = dbReturning(5);
+    const now = new Date("2026-06-01T00:01:00Z");
+    const result = await checkToolBudget({
+      db,
+      principal: { kind: "user", id: "u1" },
+      max: 10,
+      windowMs: 60_000,
+      now,
+    });
+    expect(result.used).toBe(5);
+    expect(result.allowed).toBe(true);
+    // Rolling window: start is `now - windowMs`.
+    expect(result.windowStart.getTime()).toBe(now.getTime() - 60_000);
+    expect(where).toHaveBeenCalledTimes(1);
+  });
+
+  test("at-or-over budget is refused (used >= max)", async () => {
+    const { db } = dbReturning(10);
+    const result = await checkToolBudget({
+      db,
+      principal: { kind: "user", id: "u1" },
+      max: 10,
+    });
+    expect(result.allowed).toBe(false);
+  });
+
+  test("enforceToolBudget throws ToolBudgetExceededError over budget", async () => {
+    const { db } = dbReturning(TOOL_BUDGET_MAX_CALLS);
+    await expect(
+      enforceToolBudget({ db, principal: { kind: "application", id: "a1" } }),
+    ).rejects.toBeInstanceOf(ToolBudgetExceededError);
+  });
+
+  test("enforceToolBudget passes under budget", async () => {
+    const { db } = dbReturning(0);
+    await expect(
+      enforceToolBudget({ db, principal: { kind: "user", id: "u1" } }),
+    ).resolves.toBeUndefined();
+  });
+});

package/src/rate-limit/tool-budget.ts

@@ -0,0 +1,107 @@
+import { and, count, eq, gte } from "drizzle-orm";
+import type { SafeDatabase } from "@checkstack/backend-api";
+import * as schema from "../schema";
+import type { AuditPrincipal } from "../propose-apply/store";
+
+type AiDatabase = SafeDatabase<typeof schema>;
+
+/**
+ * Per-principal tool RATE-LIMIT BUDGET (plan §14.5, P3 review item 3).
+ *
+ * A shared-Postgres ROLLING-WINDOW counter over `ai_tool_calls`: the budget is
+ * the number of rows this principal has written in the trailing window. Because
+ * the count is read from the SAME shared Postgres table that every pod writes
+ * to (state-and-scale §14.5), the cap holds across ALL pods — an in-memory
+ * per-pod limiter would let N pods each allow the cap, i.e. N x the intended
+ * limit, which a single-process test would never catch. This mirrors the P2 DCR
+ * rate-limiter pattern, reusing the existing `ai_tool_calls_principal_created_idx`
+ * index (`principalKind, principalId, createdAt`).
+ *
+ * Every tool invocation across BOTH transports (MCP `tools/call` and the chat
+ * agent loop) writes an `ai_tool_calls` row, so the count is an accurate budget
+ * of the principal's tool activity regardless of effect. The check runs BEFORE
+ * the call is executed/recorded; over-budget callers are refused.
+ */
+
+/** Default rolling window for the per-principal tool budget (1 minute). */
+export const TOOL_BUDGET_WINDOW_MS = 60_000;
+/** Default max tool calls per principal per window. */
+export const TOOL_BUDGET_MAX_CALLS = 60;
+
+export interface ToolBudgetResult {
+  /** Tool calls already recorded for this principal in the current window. */
+  used: number;
+  /** Whether a NEW call is within budget (`used < max`). */
+  allowed: boolean;
+  /** Trailing window start the count was taken against. */
+  windowStart: Date;
+}
+
+/**
+ * Read the principal's tool-call count over the trailing window and decide
+ * whether a new call is within budget. Returns `allowed: false` when the
+ * principal has already met or exceeded `max` in the window. The caller decides
+ * the response (a 429 for MCP, a friendly error for chat).
+ */
+export async function checkToolBudget({
+  db,
+  principal,
+  max = TOOL_BUDGET_MAX_CALLS,
+  windowMs = TOOL_BUDGET_WINDOW_MS,
+  now = new Date(),
+}: {
+  db: AiDatabase;
+  principal: AuditPrincipal;
+  max?: number;
+  windowMs?: number;
+  now?: Date;
+}): Promise<ToolBudgetResult> {
+  const windowStart = new Date(now.getTime() - windowMs);
+
+  const rows = await db
+    .select({ value: count() })
+    .from(schema.aiToolCalls)
+    .where(
+      and(
+        eq(schema.aiToolCalls.principalKind, principal.kind),
+        eq(schema.aiToolCalls.principalId, principal.id),
+        gte(schema.aiToolCalls.createdAt, windowStart),
+      ),
+    );
+
+  const used = rows[0]?.value ?? 0;
+  return { used, allowed: used < max, windowStart };
+}
+
+/** Error thrown / mapped to 429 when a principal is over its tool budget. */
+export class ToolBudgetExceededError extends Error {
+  constructor(
+    public readonly used: number,
+    public readonly max: number,
+  ) {
+    super(
+      `Tool rate-limit budget exceeded: ${used} calls in the current window (max ${max}). Try again shortly.`,
+    );
+    this.name = "ToolBudgetExceededError";
+  }
+}
+
+/**
+ * Enforce the budget: throws {@link ToolBudgetExceededError} when over budget.
+ * Convenience wrapper used by both transports so the enforcement is identical.
+ */
+export async function enforceToolBudget(args: {
+  db: AiDatabase;
+  principal: AuditPrincipal;
+  max?: number;
+  windowMs?: number;
+  now?: Date;
+}): Promise<void> {
+  const result = await checkToolBudget(args);
+  if (!result.allowed) {
+    throw new ToolBudgetExceededError(
+      result.used,
+      args.max ?? TOOL_BUDGET_MAX_CALLS,
+    );
+  }
+}