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.
This commit is contained in:
Marc-André Bombeck committed 2026-06-26 09:26:47 +02:00
1 parent a20e140a2f
commit fb54b46034
3 files changed
+151 -9

No files matched your search

@@ -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<typeof import("@/lib/api-handler")>(
"@/lib/api-handler",
);
return {
...actual,
apiHandler: <T extends (...args: unknown[]) => Promise<Response>>(
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<typeof import("@/lib/labs/ocr-extract")>(
"@/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",
);
});
});
+10 -9
View File
@@ -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",
+12
View File
@@ -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<string, AiBudget>;