Cap the medication-extract budget by the paying credential

The route used the enforceBudget + recordSpend read-then-write pair and
passed no cap, so every call measured against the operator ceiling — a
self-hoster on their own key was rationed by the operator's bill while
paying for their own usage. It also left the TOCTOU window the rest of
the codebase closed: concurrent calls could each read a sub-cap total
and all proceed.

Move the gate after the chain resolves and reserve atomically against
resolveDailyCap(chain), matching runStatusCompletion and the document
paths. Reconcile on every exit: the provider's reported count on a
completed call, a full refund when nothing was generated.

src/app/api/medications/extract/__tests__/route.test.ts

@@ -34,8 +34,16 @@ vi.mock("@/lib/rate-limit", () => ({
 
 vi.mock("@/lib/ai/coach/budget", () => ({
   buildDateKey: vi.fn(() => "2026-05-28"),
-  enforceBudget: vi.fn(async () => undefined),
-  recordSpend: vi.fn(async () => undefined),
+  reserveBudget: vi.fn(async () => ({
+    allowed: true,
+    reserved: 700,
+    totalAfter: 700,
+  })),
+  reconcileSpend: vi.fn(async () => undefined),
+  // The REAL cap resolver is used in the cap test below; the default mock
+  // returns a distinctive figure so an accidental default-cap regression is
+  // visible in the assertion rather than silently passing.
+  resolveDailyCap: vi.fn(() => 1_234_567),
 }));
 
 vi.mock("@/lib/ai/provider", () => ({
@@ -73,9 +81,16 @@ vi.mock("@/lib/logging/context", () => ({
 import { POST } from "../route";
 import { requireAuth, HttpError } from "@/lib/api-handler";
 import { checkRateLimit } from "@/lib/rate-limit";
-import { enforceBudget } from "@/lib/ai/coach/budget";
+import {
+  reserveBudget,
+  reconcileSpend,
+  resolveDailyCap,
+} from "@/lib/ai/coach/budget";
 import { resolveProvider, resolveProviderChain } from "@/lib/ai/provider";
-import { runRawCompletionWithFallback } from "@/lib/ai/provider-runner";
+import {
+  runRawCompletionWithFallback,
+  AllProvidersFailedError,
+} from "@/lib/ai/provider-runner";
 import { assertConsentForChain } from "@/lib/ai/consent-guard";
 
 function postReq(body: unknown): NextRequest {
@@ -98,7 +113,12 @@ beforeEach(() => {
     remaining: 9,
     resetAt: Date.now() + 60_000,
   });
-  vi.mocked(enforceBudget).mockResolvedValue(undefined);
+  vi.mocked(reserveBudget).mockResolvedValue({
+    allowed: true,
+    reserved: 700,
+    totalAfter: 700,
+  });
+  vi.mocked(resolveDailyCap).mockReturnValue(1_234_567);
   vi.mocked(resolveProviderChain).mockResolvedValue([
     {
       providerType: "openai",
@@ -216,6 +236,107 @@ describe("POST /api/medications/extract — happy path", () => {
   });
 });
 
+describe("POST /api/medications/extract — daily budget", () => {
+  function okCompletion(tokensUsed: number | null) {
+    return {
+      result: {
+        content: JSON.stringify({ dose: "5", doseUnit: "mg" }),
+        tokensUsed,
+        model: "gpt-4",
+        providerType: "anthropic",
+      },
+      workingProvider: { providerType: "openai", instance: {} as never },
+      fallbackHops: [],
+    };
+  }
+
+  it("caps against the COST OWNER — the resolved chain, not the operator ceiling", async () => {
+    vi.mocked(requireAuth).mockResolvedValue(AUTH_OK as never);
+    vi.mocked(runRawCompletionWithFallback).mockResolvedValueOnce(
+      okCompletion(120) as never,
+    );
+
+    await POST(postReq({ text: "5mg weekly" }) as never);
+
+    // The cap is derived from the chain that would actually pay for the call.
+    // A user-key chain ("openai") therefore never rides the operator ceiling —
+    // `budget.test.ts` pins that resolveDailyCap([{openai}]) === USER_PLAN_CAP.
+    expect(resolveDailyCap).toHaveBeenCalledWith(
+      expect.arrayContaining([
+        expect.objectContaining({ providerType: "openai" }),
+      ]),
+    );
+    // ...and the reservation is made against exactly that cap, not a default.
+    expect(reserveBudget).toHaveBeenCalledWith(
+      "user-1",
+      expect.any(Number),
+      "2026-05-28",
+      1_234_567,
+    );
+  });
+
+  it("returns 429 without calling the provider when the day's cap is spent", async () => {
+    vi.mocked(requireAuth).mockResolvedValue(AUTH_OK as never);
+    vi.mocked(reserveBudget).mockResolvedValueOnce({
+      allowed: false,
+      reserved: 700,
+      totalAfter: 9_999_999,
+    });
+
+    const res = await POST(postReq({ text: "Mounjaro 5mg weekly" }) as never);
+
+    expect(res.status).toBe(429);
+    const body = await res.json();
+    expect(body.error).toMatch(/budget/i);
+    // The refusal happens BEFORE any upstream spend.
+    expect(runRawCompletionWithFallback).not.toHaveBeenCalled();
+  });
+
+  it("reconciles the reservation against the provider's reported count", async () => {
+    vi.mocked(requireAuth).mockResolvedValue(AUTH_OK as never);
+    vi.mocked(runRawCompletionWithFallback).mockResolvedValueOnce(
+      okCompletion(120) as never,
+    );
+
+    await POST(postReq({ text: "5mg weekly" }) as never);
+
+    expect(reconcileSpend).toHaveBeenCalledWith(
+      "user-1",
+      700,
+      120,
+      "2026-05-28",
+    );
+  });
+
+  it("bills the reservation, not zero, when the provider reports no count", async () => {
+    vi.mocked(requireAuth).mockResolvedValue(AUTH_OK as never);
+    vi.mocked(runRawCompletionWithFallback).mockResolvedValueOnce(
+      okCompletion(null) as never,
+    );
+
+    await POST(postReq({ text: "5mg weekly" }) as never);
+
+    expect(reconcileSpend).toHaveBeenCalledWith(
+      "user-1",
+      700,
+      700,
+      "2026-05-28",
+    );
+  });
+
+  it("refunds the whole reservation when the provider chain fails", async () => {
+    vi.mocked(requireAuth).mockResolvedValue(AUTH_OK as never);
+    vi.mocked(runRawCompletionWithFallback).mockRejectedValueOnce(
+      new AllProvidersFailedError([]),
+    );
+
+    const res = await POST(postReq({ text: "5mg weekly" }) as never);
+
+    expect(res.status).toBe(503);
+    expect(reconcileSpend).toHaveBeenCalledWith("user-1", 700, 0, "2026-05-28");
+  });
+});
+
 describe("POST /api/medications/extract — validation + provider errors", () => {
   it("returns 422 on a malformed body", async () => {
     vi.mocked(requireAuth).mockResolvedValue(AUTH_OK as never);

src/app/api/medications/extract/route.ts

@@ -26,13 +26,19 @@
  *   1. requireAuth()                       — cookie OR Bearer.
  *   2. requireAssistantSurface("coach")    — operator can disable.
  *   3. checkRateLimit(...)                 — 10 requests / 5 min / user.
- *   4. enforceBudget()                     — daily Coach token ceiling.
- *   5. resolveProviderChain()              — fall back through providers.
+ *   4. resolveProviderChain()              — fall back through providers.
+ *   5. reserveBudget(resolveDailyCap(chain)) — atomic daily token ceiling,
+ *      capped against the COST OWNER: a self-hoster on their own key or a
+ *      local model is measured against the user-plan ceiling, not the
+ *      operator's. Runs AFTER the chain resolves because the cap depends
+ *      on which credential would actually pay for the call.
  *   6. runRawCompletionWithFallback()      — same machinery the Coach
  *      uses; structured-output is enforced via JSON.parse + Zod here
  *      rather than via the strict-schema wrapper.
  *   7. applyCitationGuard()                — drop hallucinated name/dose.
- *   8. recordSpend()                       — bump the day's ledger.
+ *   8. reconcileSpend()                    — settle the reservation against
+ *      the provider's reported count on EVERY exit path, refunding when
+ *      nothing was generated.
  *
  * Response: `{ data: MedicationExtractionResult, error: null }`. Every
  * field on the result is optional; the wizard is the merge target.
@@ -62,8 +68,9 @@ import { assertConsentForChain } from "@/lib/ai/consent-guard";
 
 import {
   buildDateKey,
-  enforceBudget,
-  recordSpend,
+  reconcileSpend,
+  reserveBudget,
+  resolveDailyCap,
 } from "@/lib/ai/coach/budget";
 import {
   applyCitationGuard,
@@ -75,6 +82,12 @@ import {
 /** Hard cap on the free-text payload; mirrors `coachChatRequestSchema.message`. */
 const MAX_TEXT_LENGTH = 2000;
 
+/**
+ * Output ceiling for the extraction reply — a single small JSON object. Also
+ * the output half of the budget reservation estimate.
+ */
+const MAX_OUTPUT_TOKENS = 600;
+
 const requestSchema = z.object({
   text: z.string().min(1).max(MAX_TEXT_LENGTH),
   /** Optional UI locale; informational hint for the model. */
@@ -154,8 +167,6 @@ async function handleExtract(request: NextRequest): Promise<Response> {
     });
   }
 
-  await enforceBudget(userId);
-
   // Resolve the same provider chain the Coach uses. If the user has no
   // provider configured AND the operator has not seeded an admin key,
   // surface a structured 503 — the wizard hides the overlay.
@@ -181,6 +192,31 @@ async function handleExtract(request: NextRequest): Promise<Response> {
     locale,
   });
 
+  // Atomically reserve the day's budget BEFORE the provider call. The old
+  // `enforceBudget` read-then-write left a TOCTOU window (concurrent calls
+  // could each observe a sub-cap total and all proceed) AND passed no cap, so
+  // it always measured against `OPERATOR_COST_CAP` — rationing a self-hoster
+  // paying for their own key by the operator's ceiling. `resolveDailyCap`
+  // charges the operator ceiling only when the chain's primary is the
+  // operator's own credential; every other chain gets the user-plan ceiling.
+  const dateKey = buildDateKey();
+  const estimatedTokens =
+    MAX_OUTPUT_TOKENS +
+    Math.ceil((systemPrompt.length + userPrompt.length) / 4);
+  const reservation = await reserveBudget(
+    userId,
+    estimatedTokens,
+    dateKey,
+    resolveDailyCap(chain),
+  );
+  if (!reservation.allowed) {
+    annotate({
+      action: { name: "medications.extract.budget-exceeded" },
+      meta: { totalAfter: reservation.totalAfter },
+    });
+    return apiError("coach.budget.exceeded", 429);
+  }
+
   let completion;
   try {
     completion = await runRawCompletionWithFallback({
@@ -194,13 +230,18 @@ async function handleExtract(request: NextRequest): Promise<Response> {
         temperature: 0.1,
         // The reply body is a single small JSON object; 600 is well
         // over the empirical max we have seen on a 2k-char input.
-        maxTokens: 600,
+        maxTokens: MAX_OUTPUT_TOKENS,
         // Parsed with `parseModelJson` (JSON.parse) below — opt the OpenAI /
         // Codex chains into their strict JSON-object mode (gated on this flag).
         responseFormat: "json",
       }),
     });
   } catch (err) {
+    // Nothing was generated — refund the whole reservation rather than bill an
+    // invented figure. Best-effort: a failed refund must not mask the 503.
+    await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch(
+      () => {},
+    );
     if (err instanceof AllProvidersFailedError) {
       annotate({
         action: { name: "medications.extract.provider-failed" },
@@ -211,6 +252,14 @@ async function handleExtract(request: NextRequest): Promise<Response> {
     throw err;
   }
 
+  // Reconcile against what the provider actually reported. This runs for the
+  // unusable-reply branches too: those tokens were burned upstream even though
+  // the body could not be used, so they stay on the ledger rather than funding
+  // a free retry loop. Falls back to the reservation when the provider reports
+  // no count, so an unreported generation is never billed as zero.
+  const actualTokens = completion.result.tokensUsed ?? reservation.reserved;
+  await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey);
+
   const reply = completion.result.content?.trim() ?? "";
   if (!reply) {
     annotate({ action: { name: "medications.extract.empty-reply" } });
@@ -242,13 +291,6 @@ async function handleExtract(request: NextRequest): Promise<Response> {
     text,
   );
 
-  // Bump the day's spend ledger so the daily budget gate stays in sync.
-  await recordSpend({
-    userId,
-    tokens: completion.result.tokensUsed ?? 0,
-    dateKey: buildDateKey(),
-  });
-
   annotate({
     action: { name: "medications.extract.ok" },
     meta: {