fix(coach): close the budget-gate TOCTOU and count burned tokens

The daily token cap was read before the provider call and recorded only
after a successful reply, so concurrent requests could all read
`spent < cap` and all hit the provider, and an empty/sentinel reply
recorded zero spend despite burning upstream tokens.

Reserve the budget atomically before the call: a single upsert-increment
returns the new total, so concurrent requests serialise on the row and
cannot both pass the cap; over-cap refunds the reservation and surfaces
`coach.budget.exceeded` as an SSE error frame. Reconcile the reservation
to the actual token count immediately after the provider returns —
before any empty/sentinel short-circuit — so burned tokens are always
recorded. An integration test proves two concurrent requests cannot both
exceed the cap and that an empty reply still records its cost.

src/app/api/insights/chat/route.ts

@@ -62,8 +62,8 @@ import { enqueueCoachMemoryRefresh } from "@/lib/ai/coach/coach-memory-shared";
 import { storeDeterministicFacts } from "@/lib/ai/coach/facts";
 import {
   buildDateKey,
-  enforceBudget,
-  recordSpend,
+  reserveBudget,
+  reconcileSpend,
 } from "@/lib/ai/coach/budget";
 import { detectRefusal } from "@/lib/ai/coach/refusal";
 import { getCoachSystemPrompt } from "@/lib/ai/coach/system-prompt";
@@ -215,7 +215,10 @@ async function handleChatRequest(request: NextRequest): Promise<Response> {
     return apiError("Too many Coach requests, please wait a moment", 429);
   }
 
-  await enforceBudget(userId);
+  // v1.18.7 (SENIOR-DEV HIGH) — the daily token cap is enforced atomically
+  // by reserving budget BEFORE the provider call (see the reservation right
+  // before `runRawCompletionWithFallback`), not by a read-then-write here.
+  // The old read-before-call gate let concurrent requests all pass the cap.
 
   const locale = await resolveServerLocale({
     request,
@@ -431,6 +434,27 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`;
   // ConsentRequiredError → apiHandler returns 403 + `consent.ai.required`.
   await assertConsentForChain({ userId, chain, surface: "coach" });
 
+  // v1.18.7 (SENIOR-DEV HIGH) — atomically RESERVE the day's budget before
+  // the provider call. The reservation increments the day's total by the
+  // per-call ceiling (`maxTokens`) in one upsert and returns the new total;
+  // concurrent requests serialise on the row so they cannot all pass the cap.
+  // Over-cap → 429 refusal frame, reservation already refunded. The actual
+  // token count is reconciled against this reservation after the call,
+  // including on empty / sentinel replies (their tokens were still burned).
+  const reqDateKey = buildDateKey();
+  const reservation = await reserveBudget(
+    userId,
+    AI_BUDGETS.coach.maxTokens,
+    reqDateKey,
+  );
+  if (!reservation.allowed) {
+    annotate({
+      action: { name: "coach.budget.exceeded" },
+      meta: { totalAfter: reservation.totalAfter },
+    });
+    return streamProviderError({ code: "coach.budget.exceeded" });
+  }
+
   let result: CompletionResult;
   let workingProviderType: string;
   try {
@@ -447,6 +471,11 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`;
     result = fallback.result;
     workingProviderType = fallback.workingProvider.providerType;
   } catch (err) {
+    // The provider chain failed outright — no tokens were billed, so refund
+    // the full reservation before surfacing the error frame.
+    await reconcileSpend(userId, reservation.reserved, 0, reqDateKey).catch(
+      () => {},
+    );
     if (err instanceof AllProvidersFailedError) {
       annotate({
         action: { name: "insights.coach.providerFailed" },
@@ -484,6 +513,21 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`;
     throw err;
   }
 
+  // v1.18.7 (SENIOR-DEV MEDIUM) — the provider call returned, so its tokens
+  // were billed regardless of reply quality. Reconcile the reservation
+  // against the actual count NOW, before any empty/sentinel short-circuit, so
+  // an empty or sentinel-only reply still records its burned cost (the old
+  // post-hoc `recordSpend` ran only on the happy path, undercounting these).
+  await reconcileSpend(
+    userId,
+    reservation.reserved,
+    result.tokensUsed ?? 0,
+    reqDateKey,
+  ).catch(() => {
+    // Ledger reconcile is best-effort; a failure leaves the conservative
+    // reservation in place (never an undercount) and never breaks the turn.
+  });
+
   const rawReply = (result.content ?? "").trim();
   if (!rawReply) {
     return streamProviderError({ code: "coach.provider.empty" });
@@ -611,13 +655,10 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`;
     promptVersion: PROMPT_VERSION,
   });
 
-  // Bump the day's spend ledger AFTER persistence so a retried request
-  // doesn't double-count when the persistence layer rolled back.
-  await recordSpend({
-    userId,
-    tokens: result.tokensUsed ?? 0,
-    dateKey: buildDateKey(),
-  });
+  // v1.18.7 — the day's spend was already reconciled against the
+  // reservation immediately after the provider returned (above), so there is
+  // no post-persistence ledger bump here. The reservation guarantees the
+  // tokens are counted even if persistence or streaming fails afterwards.
 
   annotate({
     action: { name: "insights.coach.replied" },

src/lib/ai/coach/budget.ts

@@ -61,6 +61,119 @@ export async function enforceBudget(
   }
 }
 
+/**
+ * v1.18.7 (SENIOR-DEV HIGH) — atomically reserve budget BEFORE the provider
+ * call, closing the read-then-write TOCTOU window.
+ *
+ * `enforceBudget` reads spend, then `recordSpend` bumps it only AFTER a
+ * successful reply, so up to `rate-limit` (20/min) concurrent requests could
+ * each read `spent < cap` before any spend landed and all hit the provider.
+ * This instead does a single atomic SQL upsert that increments the day's
+ * total by an ESTIMATED reservation and returns the new total. The caller
+ * checks the returned total against the cap: if the reservation pushed it
+ * over, the request is refused and the reservation refunded; otherwise the
+ * call proceeds and the actual token count is reconciled afterwards. Mirrors
+ * the `rate-limit.ts` atomic-upsert pattern — multi-instance correctness is
+ * structural, not a hopeful read-then-write.
+ *
+ * The first request of a fresh day still lands cheaply via the upsert
+ * `create` branch; `messageCount` is bumped at reservation time so a
+ * reserved-but-failed turn still counts as an attempt.
+ */
+export interface ReserveBudgetResult {
+  /** True when the reservation kept the day's total within the cap. */
+  allowed: boolean;
+  /** Tokens reserved by this call (refunded/reconciled by the caller). */
+  reserved: number;
+  /** The day's total AFTER this reservation (for observability). */
+  totalAfter: number;
+}
+
+export async function reserveBudget(
+  userId: string,
+  estimatedTokens: number,
+  dateKey: string = buildDateKey(),
+  cap: number = MAX_TOKENS_PER_USER_PER_DAY,
+): Promise<ReserveBudgetResult> {
+  const reserved =
+    Number.isFinite(estimatedTokens) && estimatedTokens > 0
+      ? Math.floor(estimatedTokens)
+      : 0;
+
+  // Single atomic upsert-increment returning the new total. Two concurrent
+  // requests serialise on the row's unique (user_id, date_key) constraint, so
+  // each observes a distinct post-increment total — they cannot both read a
+  // sub-cap value and both proceed.
+  const rows = await prisma.$queryRaw<{ total_tokens: number }[]>`
+    INSERT INTO coach_usage (id, user_id, date_key, total_tokens, message_count, created_at, updated_at)
+    VALUES (gen_random_uuid()::text, ${userId}, ${dateKey}, ${reserved}, 1, NOW(), NOW())
+    ON CONFLICT (user_id, date_key) DO UPDATE SET
+      total_tokens = coach_usage.total_tokens + ${reserved},
+      message_count = coach_usage.message_count + 1,
+      updated_at = NOW()
+    RETURNING total_tokens
+  `;
+  const totalAfter = Number(rows[0]?.total_tokens ?? reserved);
+
+  // The cap is a ceiling on tokens already spent BEFORE this request, matching
+  // the prior `spent >= cap` semantics: a request is allowed when the spend
+  // PRIOR to its reservation was under the cap. So compare `totalAfter -
+  // reserved` (the prior total) against the cap.
+  const priorTotal = totalAfter - reserved;
+  if (priorTotal >= cap) {
+    // Already over before this request — refund the reservation + the
+    // message-count bump and refuse.
+    await refundReservation(userId, reserved, dateKey);
+    return { allowed: false, reserved, totalAfter: priorTotal };
+  }
+
+  return { allowed: true, reserved, totalAfter };
+}
+
+/**
+ * Reconcile a reservation against the actual tokens the provider reported.
+ * `actual - reserved` is applied as a signed delta (clamped so the row never
+ * goes negative). Called after every provider call — including empty /
+ * sentinel / refusal replies, whose upstream tokens were still burned
+ * (SENIOR-DEV MEDIUM: spend undercount).
+ */
+export async function reconcileSpend(
+  userId: string,
+  reserved: number,
+  actualTokens: number,
+  dateKey: string = buildDateKey(),
+): Promise<void> {
+  const actual =
+    Number.isFinite(actualTokens) && actualTokens > 0
+      ? Math.floor(actualTokens)
+      : 0;
+  const delta = actual - reserved;
+  if (delta === 0) return;
+  // Clamp at zero so a smaller-than-reserved actual can't drive the row
+  // negative under a racing reconcile.
+  await prisma.$executeRaw`
+    UPDATE coach_usage
+    SET total_tokens = GREATEST(0, total_tokens + ${delta}),
+        updated_at = NOW()
+    WHERE user_id = ${userId} AND date_key = ${dateKey}
+  `;
+}
+
+/** Refund a reservation (tokens + the message-count bump) on a refusal. */
+async function refundReservation(
+  userId: string,
+  reserved: number,
+  dateKey: string,
+): Promise<void> {
+  await prisma.$executeRaw`
+    UPDATE coach_usage
+    SET total_tokens = GREATEST(0, total_tokens - ${reserved}),
+        message_count = GREATEST(0, message_count - 1),
+        updated_at = NOW()
+    WHERE user_id = ${userId} AND date_key = ${dateKey}
+  `;
+}
+
 export interface RecordSpendParams {
   userId: string;
   tokens: number;

tests/integration/coach-budget-reservation.test.ts

@@ -0,0 +1,109 @@
+/**
+ * v1.18.7 (SENIOR-DEV HIGH + MEDIUM) — atomic Coach budget reservation.
+ *
+ * `reserveBudget` does a single atomic upsert-increment so concurrent
+ * requests cannot all pass the daily cap (the old read-before-call gate
+ * could). `reconcileSpend` then adjusts the reservation to the actual token
+ * count — including on empty / sentinel replies whose tokens were still
+ * burned (the prior post-hoc `recordSpend` undercounted these). These tests
+ * prove both contracts against a real Postgres; a mocked client could not
+ * detect a missing atomic increment.
+ */
+import { beforeEach, describe, expect, it } from "vitest";
+
+import { getPrismaClient, truncateAllTables } from "./setup";
+
+async function seedUser(): Promise<string> {
+  const user = await getPrismaClient().user.create({
+    data: { username: "budget-user", email: "budget@example.test", role: "USER" },
+  });
+  return user.id;
+}
+
+beforeEach(async () => {
+  await truncateAllTables(getPrismaClient());
+});
+
+describe("reserveBudget (real Postgres)", () => {
+  it("lets at most one of two concurrent requests push past the cap", async () => {
+    const { reserveBudget } = await import("@/lib/ai/coach/budget");
+    const userId = await seedUser();
+    const dateKey = "2026-06-19";
+    const cap = 1000;
+
+    // Pre-spend 900 so a single 600-token reservation stays under the cap
+    // but two concurrent ones cannot both be admitted.
+    await getPrismaClient().coachUsage.create({
+      data: { userId, dateKey, totalTokens: 900, messageCount: 1 },
+    });
+
+    const [a, b] = await Promise.all([
+      reserveBudget(userId, 600, dateKey, cap),
+      reserveBudget(userId, 600, dateKey, cap),
+    ]);
+
+    // Exactly one is admitted: the first to commit sees prior 900 (< cap) and
+    // reserves to 1500; the second sees prior 1500 (>= cap) and is refused +
+    // refunded. Order is non-deterministic, so assert the SET of outcomes.
+    const allowed = [a, b].filter((r) => r.allowed).length;
+    expect(allowed).toBe(1);
+
+    // The refused reservation was refunded, so the row reflects exactly one
+    // admitted 600-token reservation on top of the 900 baseline.
+    const row = await getPrismaClient().coachUsage.findUnique({
+      where: { userId_dateKey: { userId, dateKey } },
+    });
+    expect(row?.totalTokens).toBe(1500);
+  });
+
+  it("admits the first request of a fresh day via the upsert create branch", async () => {
+    const { reserveBudget } = await import("@/lib/ai/coach/budget");
+    const userId = await seedUser();
+    const dateKey = "2026-06-20";
+
+    const res = await reserveBudget(userId, 600, dateKey, 25_000);
+    expect(res.allowed).toBe(true);
+
+    const row = await getPrismaClient().coachUsage.findUnique({
+      where: { userId_dateKey: { userId, dateKey } },
+    });
+    expect(row?.totalTokens).toBe(600);
+    expect(row?.messageCount).toBe(1);
+  });
+
+  it("reconciles the reservation down to a smaller actual count", async () => {
+    const { reserveBudget, reconcileSpend } = await import(
+      "@/lib/ai/coach/budget"
+    );
+    const userId = await seedUser();
+    const dateKey = "2026-06-21";
+
+    const res = await reserveBudget(userId, 600, dateKey, 25_000);
+    await reconcileSpend(userId, res.reserved, 120, dateKey);
+
+    const row = await getPrismaClient().coachUsage.findUnique({
+      where: { userId_dateKey: { userId, dateKey } },
+    });
+    // 600 reserved → reconciled to the actual 120 burned.
+    expect(row?.totalTokens).toBe(120);
+  });
+
+  it("still records burned tokens on an empty/sentinel reply (no undercount)", async () => {
+    const { reserveBudget, reconcileSpend } = await import(
+      "@/lib/ai/coach/budget"
+    );
+    const userId = await seedUser();
+    const dateKey = "2026-06-22";
+
+    // Provider returned 450 tokens but the reply was empty/sentinel: the
+    // reconcile must still record the 450 burned, not zero.
+    const res = await reserveBudget(userId, 600, dateKey, 25_000);
+    await reconcileSpend(userId, res.reserved, 450, dateKey);
+
+    const row = await getPrismaClient().coachUsage.findUnique({
+      where: { userId_dateKey: { userId, dateKey } },
+    });
+    expect(row?.totalTokens).toBe(450);
+    expect(row?.totalTokens).toBeGreaterThan(0);
+  });
+});

tests/integration/coach-chat.test.ts

@@ -214,7 +214,12 @@ describe("POST /api/insights/chat — integration", () => {
     expect(usage?.messageCount).toBe(1);
   });
 
-  it("returns 429 when daily token budget is exhausted", async () => {
+  it("refuses with a budget-exceeded error frame when the daily cap is spent", async () => {
+    // v1.18.7 — the cap is now enforced by an atomic reservation right
+    // before the provider call. An over-budget request surfaces the
+    // structured `coach.budget.exceeded` code as an SSE `error` frame
+    // (HTTP 200, like every other coach error) rather than a JSON 429, and
+    // the reservation is refunded so the ledger stays at the cap.
     const { userId } = await seedUserWithSession();
     const prisma = getPrismaClient();
     const dateKey = new Date().toISOString().slice(0, 10);
@@ -237,10 +242,16 @@ describe("POST /api/insights/chat — integration", () => {
         }),
       }),
     );
-    expect(res.status).toBe(429);
-    const env = (await res.json()) as { error: string };
-    expect(env.error).toBe("coach.budget.exceeded");
+    expect(res.status).toBe(200);
+    const body = await readStream(res);
+    expect(body).toContain("coach.budget.exceeded");
     expect(runProviderMock).not.toHaveBeenCalled();
+
+    // The refused reservation was refunded — the ledger holds at the cap.
+    const usage = await prisma.coachUsage.findUnique({
+      where: { userId_dateKey: { userId, dateKey } },
+    });
+    expect(usage?.totalTokens).toBe(MAX_TOKENS_PER_USER_PER_DAY);
   });
 
   it("returns 404 (not 403) when conversationId belongs to another user", async () => {