From fb54b4603458797aa75de542175a9921815474f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Fri, 26 Jun 2026 09:26:47 +0200 Subject: [PATCH] fix(labs): reserve a proportionate budget for text-mode OCR Text-mode OCR structures browser-extracted text rather than running a vision call, so it reserved and reconciled the full vision budget even though the spend is far lower. Add an ocrExtractText ceiling and reserve that in text mode, and refund the reservation in full on a clean extract failure instead of charging the reserved estimate. --- .../labs/ocr/extract/__tests__/route.test.ts | 129 ++++++++++++++++++ src/app/api/labs/ocr/extract/route.ts | 19 +-- src/lib/ai/ai-budgets.ts | 12 ++ 3 files changed, 151 insertions(+), 9 deletions(-) create mode 100644 src/app/api/labs/ocr/extract/__tests__/route.test.ts diff --git a/src/app/api/labs/ocr/extract/__tests__/route.test.ts b/src/app/api/labs/ocr/extract/__tests__/route.test.ts new file mode 100644 index 000000000..055a56a1f --- /dev/null +++ b/src/app/api/labs/ocr/extract/__tests__/route.test.ts @@ -0,0 +1,129 @@ +/** + * v1.20.1 — POST /api/labs/ocr/extract, TEXT mode. + * + * Focus: the text-mode structuring pass reserves the proportionate text budget + * ceiling (`AI_BUDGETS.ocrExtractText`), not the far larger vision ceiling, and + * on a clean extraction failure it refunds the reservation in full rather than + * charging it. + */ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +vi.mock("@/lib/api-handler", async () => { + const actual = + await vi.importActual( + "@/lib/api-handler", + ); + return { + ...actual, + apiHandler: Promise>( + h: T, + ): T => h, + requireAuth: vi.fn(), + }; +}); + +vi.mock("@/lib/db", () => ({ + prisma: { user: { findUnique: vi.fn() } }, +})); +vi.mock("@/lib/logging/context", () => ({ annotate: vi.fn() })); +vi.mock("@/lib/labs/ocr-capability", () => ({ + resolveTextProvider: vi.fn(), + resolveVisionProvider: vi.fn(), +})); +vi.mock("@/lib/ai/consent-guard", () => ({ + assertConsentForChain: vi.fn().mockResolvedValue(undefined), +})); +vi.mock("@/lib/rate-limit", () => ({ + checkRateLimit: vi.fn().mockResolvedValue({ allowed: true }), + rateLimitHeaders: vi.fn(() => ({})), +})); +vi.mock("@/lib/ai/coach/budget", () => ({ + buildDateKey: vi.fn(() => "2026-06-26"), + reserveBudget: vi.fn(), + reconcileSpend: vi.fn().mockResolvedValue(undefined), +})); +vi.mock("@/lib/labs/ocr-extract", async () => { + const actual = await vi.importActual( + "@/lib/labs/ocr-extract", + ); + return { ...actual, runOcrExtraction: vi.fn() }; +}); + +import { POST } from "../route"; +import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; +import { requireAuth } from "@/lib/api-handler"; +import { prisma } from "@/lib/db"; +import { resolveTextProvider } from "@/lib/labs/ocr-capability"; +import { reserveBudget, reconcileSpend } from "@/lib/ai/coach/budget"; +import { OcrExtractError, runOcrExtraction } from "@/lib/labs/ocr-extract"; + +const SESSION_OK = { + session: { id: "sess-1", expiresAt: new Date(Date.now() + 3_600_000) }, + user: { id: "user-1", username: "tester", role: "USER" as const }, +}; + +function textReq(text = "Glucose 95 mg/dL"): Request { + return new Request("http://localhost/api/labs/ocr/extract", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ mode: "text", text }), + }); +} + +beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(requireAuth).mockResolvedValue(SESSION_OK as never); + vi.mocked(prisma.user.findUnique).mockResolvedValue({ + labsLocalOcrEnabled: true, + } as never); + vi.mocked(resolveTextProvider).mockResolvedValue({ + chain: [{ providerType: "openai", instance: {} as never }], + pick: { + entry: { providerType: "openai", instance: {} as never }, + providerType: "openai", + }, + } as never); + vi.mocked(reserveBudget).mockResolvedValue({ + allowed: true, + reserved: AI_BUDGETS.ocrExtractText.maxTokens ?? 0, + totalAfter: AI_BUDGETS.ocrExtractText.maxTokens ?? 0, + } as never); +}); + +describe("POST /api/labs/ocr/extract — text mode budget", () => { + it("reserves the cheaper text ceiling, not the vision ceiling", async () => { + vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never); + + const res = await POST(textReq()); + expect(res.status).toBe(200); + + // The reservation must use the text ceiling — proven distinct from the + // vision ceiling so a future merge can't silently re-point it. + expect(AI_BUDGETS.ocrExtractText.maxTokens).toBeLessThan( + AI_BUDGETS.ocrExtract.maxTokens ?? Infinity, + ); + expect(reserveBudget).toHaveBeenCalledWith( + "user-1", + AI_BUDGETS.ocrExtractText.maxTokens, + "2026-06-26", + ); + }); + + it("refunds the reservation in full on a clean extract failure", async () => { + vi.mocked(runOcrExtraction).mockRejectedValue( + new OcrExtractError("unreadable"), + ); + + const res = await POST(textReq()); + expect(res.status).toBe(422); + + // actual spend reconciles to 0 — the failed structuring pass is refunded, + // not charged at the reserved estimate. + expect(reconcileSpend).toHaveBeenCalledWith( + "user-1", + AI_BUDGETS.ocrExtractText.maxTokens, + 0, + "2026-06-26", + ); + }); +}); diff --git a/src/app/api/labs/ocr/extract/route.ts b/src/app/api/labs/ocr/extract/route.ts index 201ee7211..e35b710d6 100644 --- a/src/app/api/labs/ocr/extract/route.ts +++ b/src/app/api/labs/ocr/extract/route.ts @@ -138,12 +138,14 @@ async function handleTextExtract( }); } - // Budget (text is cheaper than vision; reuse the same ceiling as the upper - // bound — a text-only structuring call always spends at most this). + // Budget — text mode is a plain text→JSON structuring pass, far cheaper than + // a vision call, so it reserves the proportionate text ceiling rather than + // the vision budget. Over-charging the vision rate for a text call would + // exhaust the day budget against spend that never happened. const dateKey = buildDateKey(); const reservation = await reserveBudget( userId, - AI_BUDGETS.ocrExtract.maxTokens, + AI_BUDGETS.ocrExtractText.maxTokens, dateKey, ); if (!reservation.allowed) { @@ -163,6 +165,8 @@ async function handleTextExtract( providerType: pick.providerType, ocrText: parsed.data.text, }); + // A clean structuring pass spent (at most) the reservation; the provider + // already billed it, so reconcile against the reserved estimate. await reconcileSpend( userId, reservation.reserved, @@ -171,12 +175,9 @@ async function handleTextExtract( ); return apiSuccess(result); } catch (err) { - await reconcileSpend( - userId, - reservation.reserved, - reservation.reserved, - dateKey, - ); + // A failed structuring call produced no usable rows; mirror the vision + // path and refund the reservation in full rather than charging it. + await reconcileSpend(userId, reservation.reserved, 0, dateKey); if (err instanceof OcrExtractError) { return apiError("Couldn't read the report. Try a clearer photo.", 422, { errorCode: "labs.ocr.extractFailed", diff --git a/src/lib/ai/ai-budgets.ts b/src/lib/ai/ai-budgets.ts index e0e67d14e..c1b4538fe 100644 --- a/src/lib/ai/ai-budgets.ts +++ b/src/lib/ai/ai-budgets.ts @@ -109,4 +109,16 @@ export const AI_BUDGETS = { * gates with a tight 6/hour rate bucket + reserveBudget + upload-size cap. */ ocrExtract: { temperature: 0, maxTokens: 4000 }, + + /** + * v1.20.1 — Lab-OCR text-mode structuring. The browser OCR's the image + * (tesseract.js) and POSTs only the extracted TEXT, so the provider does a + * plain text→JSON structuring pass — no expensive vision input tokens. The + * output contract is identical (the same dozen-plus analyte rows), but the + * spend ceiling is far lower than the vision path's 4000, so the day-budget + * reservation should reflect that rather than over-charging a text call at + * the vision rate. 1500 tokens comfortably covers a dense report's JSON + * envelope. Temperature 0 — faithful structuring, not generation. + */ + ocrExtractText: { temperature: 0, maxTokens: 1500 }, } as const satisfies Record;