From de088febe6b8423c830519a6b5daa816a2c73c0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Sun, 19 Jul 2026 05:31:38 +0200 Subject: [PATCH] Cap the medication-extract budget by the paying credential MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../extract/__tests__/route.test.ts | 131 +++++++++++++++++- src/app/api/medications/extract/route.ts | 72 ++++++++-- 2 files changed, 183 insertions(+), 20 deletions(-) diff --git a/src/app/api/medications/extract/__tests__/route.test.ts b/src/app/api/medications/extract/__tests__/route.test.ts index 09d7fbcf3..b80464eb3 100644 --- a/src/app/api/medications/extract/__tests__/route.test.ts +++ b/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); diff --git a/src/app/api/medications/extract/route.ts b/src/app/api/medications/extract/route.ts index 4119ee4bc..5b84598a8 100644 --- a/src/app/api/medications/extract/route.ts +++ b/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 { }); } - 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 { 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 { 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 { 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 { 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: {