diff --git a/.env.production.example b/.env.production.example index dfc0a7497..c90fb5377 100644 --- a/.env.production.example +++ b/.env.production.example @@ -230,6 +230,18 @@ API_TOKEN_HMAC_KEY="" # ("AI response was cut off"). See docs/integrations/ai-providers.md. # INSIGHTS_MAX_TOKENS="" +# Per-user hourly ceiling for POST /api/insights/generate (optional). +# Default 10. Lower it on a tight LLM budget; the 24h cache already +# short-circuits read traffic. See docs/integrations/ai-providers.md. +# INSIGHTS_RATE_LIMIT_PER_HOUR="" + +# Per-user hourly ceiling for the document-AI actions (optional). Read/index, +# extract, suggest and summarise share one bucket per user. Default 6, +# clamped to 1-1000. Raise it when users batch-read many documents against a +# generous provider budget; lower it on a tight one. See +# docs/integrations/ai-providers.md. +# DOCUMENT_AI_LIMIT_PER_HOUR="" + # ----------------------------------------------------------------------------- # Off-host backups -- all-or-none diff --git a/docker-compose.yml b/docker-compose.yml index 7add1e61d..e304d1cab 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -199,6 +199,18 @@ x-healthlog-env: &healthlog_env # Listed here so an operator value in .env reaches the container (the # `environment:` block is a whitelist). See docs/integrations/ai-providers.md. INSIGHTS_MAX_TOKENS: "${INSIGHTS_MAX_TOKENS:-}" + # Per-user hourly ceiling for POST /api/insights/generate (optional; + # default 10). Was documented but missing from this whitelist, so an + # operator value never reached the container. See + # docs/integrations/ai-providers.md. + INSIGHTS_RATE_LIMIT_PER_HOUR: "${INSIGHTS_RATE_LIMIT_PER_HOUR:-}" + # Per-user hourly ceiling for the document-AI actions (read/index, extract, + # suggest, summarise — one shared bucket; optional; default 6, clamped + # 1-1000). Raise on a self-host with a generous provider budget, lower on a + # tight one. Listed here so an operator value in .env reaches the container + # (the `environment:` block is a whitelist). See + # docs/integrations/ai-providers.md. + DOCUMENT_AI_LIMIT_PER_HOUR: "${DOCUMENT_AI_LIMIT_PER_HOUR:-}" # v1.17.1 — SMTP transport for the email notification channel (optional, # all-or-none: HOST + PORT + FROM enable it; USER/PASS optional; SECURE # toggles implicit TLS). Listed here so operator values in .env reach the diff --git a/docs/api/openapi.yaml b/docs/api/openapi.yaml index 8072fcec3..47fc1f345 100644 --- a/docs/api/openapi.yaml +++ b/docs/api/openapi.yaml @@ -12495,9 +12495,11 @@ paths: description: "Optional AI assist on an already-stored document. Runs ONE provider call over the stored original (VISION, empty body) or browser-OCR'd text (TEXT, `application/json` `{ mode: \"text\", text }`, opt-in local OCR) and returns a `{ title, kind, documentDate }` DRAFT for the edit form. AI-consent / rate / budget gated (shares the - 6/hour extract bucket). WRITES NOTHING — never stages facts, never flips status; the user reviews and saves. 422 - (`documents.inbound.providerUnsupported`) with no provider configured. Never interprets or diagnoses; the title - is a neutral filing label." + per-user document-AI bucket with extract/summary/index — default 6/hour, operator-tunable via + `DOCUMENT_AI_LIMIT_PER_HOUR`; a slot is only consumed when the request reaches the provider, and the 429 carries + the reset instant in `X-RateLimit-Reset` and `meta.retryAt`). WRITES NOTHING — never stages facts, never flips + status; the user reviews and saves. 422 (`documents.inbound.providerUnsupported`) with no provider configured. + Never interprets or diagnoses; the title is a neutral filing label." parameters: - name: id in: path diff --git a/docs/integrations/ai-providers.md b/docs/integrations/ai-providers.md index 66f03496b..80cc2abc3 100644 --- a/docs/integrations/ai-providers.md +++ b/docs/integrations/ai-providers.md @@ -393,6 +393,20 @@ verbose model gets its briefing cut off mid-JSON — the API reports that case as "AI response was cut off" rather than the generic invalid-JSON error. +`DOCUMENT_AI_LIMIT_PER_HOUR` (default `6`, clamped to 1–1000) bounds +the document-AI actions per user per hour. Read/index, extract, +suggest and summarise share one bucket, and a slot is only consumed +when a request actually reaches the provider — a request that fails +during input preparation (unreadable file, exhausted daily budget, +malformed body) hands its slot back. Raise it when users batch-read +many documents against a generous provider budget; lower it on a +tight one. The 429 names the actual wait (the window is an hour) and +carries the reset instant in `X-RateLimit-Reset`. + +All three variables must also be listed in the compose +`environment:` whitelist to reach the container; the shipped +`docker-compose.yml` lists them. + ## Connection test `POST /api/ai/test` runs a one-shot probe against the resolved diff --git a/messages/de.json b/messages/de.json index 9386be584..437a2d668 100644 --- a/messages/de.json +++ b/messages/de.json @@ -9560,7 +9560,10 @@ "unavailableAction": "KI-Einstellungen öffnen", "errorGeneric": "Dokument konnte nicht gelesen werden. Bitte erneut versuchen.", "errorOcr": "Aus diesem Bild konnte kein Text gelesen werden.", - "errorRateLimited": "Zu viele KI-Anfragen — bitte in ein paar Minuten erneut versuchen.", + "errorRateLimited": "Das stündliche KI-Leselimit ist erreicht. Bitte später erneut versuchen.", + "errorRateLimitedWaitOne": "Das stündliche KI-Leselimit ist erreicht. Bitte in etwa einer Minute erneut versuchen.", + "errorRateLimitedWaitFew": "Das stündliche KI-Leselimit ist erreicht. Bitte in etwa {count} Minuten erneut versuchen.", + "errorRateLimitedWaitOther": "Das stündliche KI-Leselimit ist erreicht. Bitte in etwa {count} Minuten erneut versuchen.", "errorBudget": "Dein KI-Budget für heute ist aufgebraucht.", "errorPdfNeedsVision": "Zum Scannen von PDFs wird ein Claude-Vision-Anbieter benötigt.", "errorFileType": "Dieser Dateityp kann nicht gelesen werden.", diff --git a/messages/en.json b/messages/en.json index 18df41bfc..3c98234b1 100644 --- a/messages/en.json +++ b/messages/en.json @@ -9560,7 +9560,10 @@ "unavailableAction": "Open AI settings", "errorGeneric": "Couldn't read the document. Try again.", "errorOcr": "Couldn't read text from this image.", - "errorRateLimited": "Too many AI requests — try again in a few minutes.", + "errorRateLimited": "The hourly AI reading limit is reached. Try again later.", + "errorRateLimitedWaitOne": "The hourly AI reading limit is reached. Try again in about a minute.", + "errorRateLimitedWaitFew": "The hourly AI reading limit is reached. Try again in about {count} minutes.", + "errorRateLimitedWaitOther": "The hourly AI reading limit is reached. Try again in about {count} minutes.", "errorBudget": "Your AI budget for today is used up.", "errorPdfNeedsVision": "Scanning PDFs needs a Claude vision provider.", "errorFileType": "This file type can't be read.", diff --git a/messages/es.json b/messages/es.json index 5bc566a52..2446eaffa 100644 --- a/messages/es.json +++ b/messages/es.json @@ -9560,7 +9560,10 @@ "unavailableAction": "Abrir ajustes de IA", "errorGeneric": "No se pudo leer el documento. Inténtalo de nuevo.", "errorOcr": "No se pudo leer texto de esta imagen.", - "errorRateLimited": "Demasiadas solicitudes de IA: inténtalo en unos minutos.", + "errorRateLimited": "Se alcanzó el límite de lecturas de IA por hora. Inténtalo más tarde.", + "errorRateLimitedWaitOne": "Se alcanzó el límite de lecturas de IA por hora. Inténtalo en aproximadamente un minuto.", + "errorRateLimitedWaitFew": "Se alcanzó el límite de lecturas de IA por hora. Inténtalo en aproximadamente {count} minutos.", + "errorRateLimitedWaitOther": "Se alcanzó el límite de lecturas de IA por hora. Inténtalo en aproximadamente {count} minutos.", "errorBudget": "Tu presupuesto de IA de hoy se ha agotado.", "errorPdfNeedsVision": "Escanear PDF requiere un proveedor de visión Claude.", "errorFileType": "Este tipo de archivo no se puede leer.", diff --git a/messages/fr.json b/messages/fr.json index def066936..419eaa2dc 100644 --- a/messages/fr.json +++ b/messages/fr.json @@ -9560,7 +9560,10 @@ "unavailableAction": "Ouvrir les paramètres d'IA", "errorGeneric": "Impossible de lire le document. Réessayez.", "errorOcr": "Impossible de lire le texte de cette image.", - "errorRateLimited": "Trop de requêtes IA — réessayez dans quelques minutes.", + "errorRateLimited": "La limite horaire de lectures IA est atteinte. Réessayez plus tard.", + "errorRateLimitedWaitOne": "La limite horaire de lectures IA est atteinte. Réessayez dans environ une minute.", + "errorRateLimitedWaitFew": "La limite horaire de lectures IA est atteinte. Réessayez dans environ {count} minutes.", + "errorRateLimitedWaitOther": "La limite horaire de lectures IA est atteinte. Réessayez dans environ {count} minutes.", "errorBudget": "Votre budget IA du jour est épuisé.", "errorPdfNeedsVision": "La lecture des PDF nécessite un fournisseur de vision Claude.", "errorFileType": "Ce type de fichier ne peut pas être lu.", diff --git a/messages/it.json b/messages/it.json index c3fe2cfe2..acde06dc0 100644 --- a/messages/it.json +++ b/messages/it.json @@ -9560,7 +9560,10 @@ "unavailableAction": "Apri impostazioni IA", "errorGeneric": "Impossibile leggere il documento. Riprova.", "errorOcr": "Impossibile leggere il testo da questa immagine.", - "errorRateLimited": "Troppe richieste IA — riprova tra qualche minuto.", + "errorRateLimited": "Il limite orario di letture IA è stato raggiunto. Riprova più tardi.", + "errorRateLimitedWaitOne": "Il limite orario di letture IA è stato raggiunto. Riprova tra circa un minuto.", + "errorRateLimitedWaitFew": "Il limite orario di letture IA è stato raggiunto. Riprova tra circa {count} minuti.", + "errorRateLimitedWaitOther": "Il limite orario di letture IA è stato raggiunto. Riprova tra circa {count} minuti.", "errorBudget": "Il tuo budget IA di oggi è esaurito.", "errorPdfNeedsVision": "La scansione dei PDF richiede un provider di visione Claude.", "errorFileType": "Questo tipo di file non può essere letto.", diff --git a/messages/pl.json b/messages/pl.json index c8bf4c825..098808b07 100644 --- a/messages/pl.json +++ b/messages/pl.json @@ -9560,7 +9560,10 @@ "unavailableAction": "Otwórz ustawienia AI", "errorGeneric": "Nie udało się odczytać dokumentu. Spróbuj ponownie.", "errorOcr": "Nie udało się odczytać tekstu z tego obrazu.", - "errorRateLimited": "Zbyt wiele żądań AI — spróbuj za kilka minut.", + "errorRateLimited": "Godzinowy limit odczytów AI został osiągnięty. Spróbuj ponownie później.", + "errorRateLimitedWaitOne": "Godzinowy limit odczytów AI został osiągnięty. Spróbuj ponownie za około minutę.", + "errorRateLimitedWaitFew": "Godzinowy limit odczytów AI został osiągnięty. Spróbuj ponownie za około {count} minuty.", + "errorRateLimitedWaitOther": "Godzinowy limit odczytów AI został osiągnięty. Spróbuj ponownie za około {count} minut.", "errorBudget": "Twój dzisiejszy budżet AI został wyczerpany.", "errorPdfNeedsVision": "Skanowanie plików PDF wymaga dostawcy wizji Claude.", "errorFileType": "Tego typu pliku nie można odczytać.", diff --git a/scripts/env-manifest.json b/scripts/env-manifest.json index 2a7ab0c02..3bf61318d 100644 --- a/scripts/env-manifest.json +++ b/scripts/env-manifest.json @@ -129,6 +129,14 @@ { "name": "INSIGHTS_MAX_TOKENS", "purpose": "Daily-briefing output-token ceiling. Default 2500, clamped to 500-8000." + }, + { + "name": "INSIGHTS_RATE_LIMIT_PER_HOUR", + "purpose": "Per-user hourly ceiling for POST /api/insights/generate. Default 10; values below 1 fall back to the default." + }, + { + "name": "DOCUMENT_AI_LIMIT_PER_HOUR", + "purpose": "Per-user hourly ceiling for the document-AI actions (read/index, extract, suggest, summarise share one bucket). Default 6, clamped to 1-1000." } ] }, diff --git a/src/app/api/documents/inbound/[id]/extract/route.ts b/src/app/api/documents/inbound/[id]/extract/route.ts index c5c124c4d..05ea711eb 100644 --- a/src/app/api/documents/inbound/[id]/extract/route.ts +++ b/src/app/api/documents/inbound/[id]/extract/route.ts @@ -55,21 +55,21 @@ import { resolveDocumentTextProvider, resolveDocumentVisionProvider, } from "@/lib/documents/provider-order"; +import { + checkDocumentAiRateLimit, + documentAiRateLimited, + DOCUMENT_AI_TEXT_BODY_MAX_BYTES, + refundDocumentAiSlot, +} from "@/lib/documents/ai-route-support"; import { detectOcrMimeType } from "@/lib/labs/ocr-upload"; import { annotate } from "@/lib/logging/context"; import { requireModuleEnabled } from "@/lib/modules/gate"; -import { checkRateLimit, rateLimitHeaders } from "@/lib/rate-limit"; import { inboundTextExtractSchema } from "@/lib/validations/inbound-documents"; export const dynamic = "force-dynamic"; type RouteParams = { params: Promise<{ id: string }> }; -/** Inbound extractions are expensive (vision) / metered (text). 6/hour. */ -const EXTRACT_LIMIT_PER_HOUR = 6; -const EXTRACT_WINDOW_MS = 60 * 60 * 1000; -const TEXT_BODY_MAX_BYTES = 512 * 1024; - /** A loaded, owner-scoped, still-extractable document row. */ type LoadedDocument = { id: string; @@ -213,28 +213,20 @@ async function handleTextExtract( surface: "insights", }); - const rl = await checkRateLimit( - `documents-inbound:${userId}`, - EXTRACT_LIMIT_PER_HOUR, - EXTRACT_WINDOW_MS, - ); - if (!rl.allowed) { - const response = apiError("Too many extractions. Try again later.", 429, { - errorCode: "documents.inbound.rateLimited", - }); - for (const [k, v] of Object.entries(rateLimitHeaders(rl))) { - response.headers.set(k, v); - } - return response; - } + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); const { data: body, error: jsonError } = await safeJson(request, { - maxBytes: TEXT_BODY_MAX_BYTES, + maxBytes: DOCUMENT_AI_TEXT_BODY_MAX_BYTES, }); - if (jsonError) return jsonError; + if (jsonError) { + await refundDocumentAiSlot(userId); + return jsonError; + } const parsed = inboundTextExtractSchema.safeParse(body); if (!parsed.success) { + await refundDocumentAiSlot(userId); return apiError("Invalid document text payload", 422, { errorCode: "documents.inbound.extractFailed", }); @@ -248,6 +240,7 @@ async function handleTextExtract( resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { + await refundDocumentAiSlot(userId); return apiError("Your AI usage budget for today is reached.", 429, { errorCode: "documents.inbound.budgetExceeded", }); @@ -315,23 +308,12 @@ async function handleVisionExtract( surface: "insights", }); - const rl = await checkRateLimit( - `documents-inbound:${userId}`, - EXTRACT_LIMIT_PER_HOUR, - EXTRACT_WINDOW_MS, - ); - if (!rl.allowed) { - const response = apiError("Too many extractions. Try again later.", 429, { - errorCode: "documents.inbound.rateLimited", - }); - for (const [k, v] of Object.entries(rateLimitHeaders(rl))) { - response.headers.set(k, v); - } - return response; - } + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); // Decrypt the stored original and re-derive its MIME from the bytes (never - // trust a stored label for the provider call). + // trust a stored label for the provider call). Every miss below happens + // before the provider dispatch, so the charged slot goes back. let buffer: Buffer; try { buffer = decryptDocumentContent( @@ -339,6 +321,7 @@ async function handleVisionExtract( document.contentCodec, ); } catch { + await refundDocumentAiSlot(userId); return apiError("Couldn't read the stored document.", 422, { errorCode: "documents.inbound.extractFailed", }); @@ -346,6 +329,7 @@ async function handleVisionExtract( const mime = detectOcrMimeType(buffer); if (!mime) { + await refundDocumentAiSlot(userId); return apiError( "This document can't be scanned. Use local OCR (text mode).", 422, @@ -354,6 +338,7 @@ async function handleVisionExtract( } if (mime === "application/pdf" && !pick.pdfSupported) { + await refundDocumentAiSlot(userId); return apiError( "PDF scanning needs a Claude vision provider; use local OCR instead.", 422, @@ -369,6 +354,7 @@ async function handleVisionExtract( resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { + await refundDocumentAiSlot(userId); return apiError("Your AI usage budget for today is reached.", 429, { errorCode: "documents.inbound.budgetExceeded", }); diff --git a/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts b/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts index c132d5d6b..0ec7a3fa9 100644 --- a/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts +++ b/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts @@ -5,6 +5,12 @@ import { NextRequest } from "next/server"; * Content-index route (Document vault P2). Pins: vision path transcribes then * upserts the index; text path indexes posted OCR text with no provider; * provider-gated on vision; only the index sibling is written (never facts). + * + * Bucket honesty: the route charges the shared document-AI bucket up front + * (so a 429 stays cheap) but a slot must reflect work actually performed — + * a 429 names the reset instant in `meta.retryAt`, a preparation failure + * before any provider dispatch refunds the charged slot, and a successful + * dispatch keeps it. */ vi.mock("@/lib/db", () => ({ @@ -47,7 +53,11 @@ vi.mock("@/lib/modules/gate", () => ({ })); vi.mock("@/lib/rate-limit", () => ({ checkRateLimit: vi.fn().mockResolvedValue({ allowed: true }), - rateLimitHeaders: () => ({}), + refundRateLimit: vi.fn().mockResolvedValue(undefined), + rateLimitHeaders: (result: { remaining?: number; resetAt?: number }) => ({ + "X-RateLimit-Remaining": String(result.remaining ?? 0), + "X-RateLimit-Reset": new Date(result.resetAt ?? 0).toISOString(), + }), })); vi.mock("@/lib/auth/session", () => ({ getSession: vi.fn() })); vi.mock("@/lib/auth/audit", () => ({ @@ -69,7 +79,10 @@ vi.mock("next/headers", () => ({ import { POST } from "../route"; import { prisma } from "@/lib/db"; import { getSession } from "@/lib/auth/session"; +import { checkRateLimit, refundRateLimit } from "@/lib/rate-limit"; +import { reserveBudget } from "@/lib/ai/coach/budget"; import { resolveDocumentVisionProvider } from "@/lib/documents/provider-order"; +import { decryptDocumentContent } from "@/lib/documents/store"; import { transcribeDocument } from "@/lib/documents/describe"; import { upsertContentIndex } from "@/lib/documents/content-index"; @@ -95,9 +108,30 @@ const textReq = (id: string, text: string) => }, ); +const PICK = { + chain: [{ providerType: "anthropic", instance: {} }], + pick: { + entry: { providerType: "anthropic", instance: {} }, + providerType: "anthropic", + pdfSupported: true, + }, +}; + beforeEach(() => { vi.clearAllMocks(); vi.mocked(getSession).mockResolvedValue(SESSION_OK as never); + vi.mocked(checkRateLimit).mockResolvedValue({ + allowed: true, + remaining: 5, + resetAt: Date.now() + 60 * 60 * 1000, + }); + vi.mocked(reserveBudget).mockResolvedValue({ + allowed: true, + reserved: 1, + } as never); + vi.mocked(decryptDocumentContent).mockReturnValue( + Buffer.from([1, 2, 3]) as never, + ); vi.mocked(prisma.inboundDocument.findFirst).mockResolvedValue({ id: "doc-1", kind: "OTHER", @@ -110,14 +144,7 @@ beforeEach(() => { describe("POST /api/documents/inbound/[id]/index", () => { it("vision: transcribes then upserts the index, never facts", async () => { - vi.mocked(resolveDocumentVisionProvider).mockResolvedValue({ - chain: [{ providerType: "anthropic", instance: {} }], - pick: { - entry: { providerType: "anthropic", instance: {} }, - providerType: "anthropic", - pdfSupported: true, - }, - } as never); + vi.mocked(resolveDocumentVisionProvider).mockResolvedValue(PICK as never); vi.mocked(transcribeDocument).mockResolvedValue({ text: "hemoglobin 14.2 leukocytes normal", } as never); @@ -176,3 +203,60 @@ describe("POST /api/documents/inbound/[id]/index", () => { expect(upsertContentIndex).not.toHaveBeenCalled(); }); }); + +describe("POST /api/documents/inbound/[id]/index — bucket honesty", () => { + beforeEach(() => { + vi.mocked(resolveDocumentVisionProvider).mockResolvedValue(PICK as never); + vi.mocked(transcribeDocument).mockResolvedValue({ + text: "hemoglobin 14.2 leukocytes normal", + } as never); + }); + + it("returns 429 with the reset instant in meta.retryAt when the bucket is exhausted", async () => { + const resetAt = Date.now() + 47 * 60_000; + vi.mocked(checkRateLimit).mockResolvedValue({ + allowed: false, + remaining: 0, + resetAt, + }); + const res = await POST(visionReq("doc-1") as never, ctx("doc-1") as never); + expect(res.status).toBe(429); + const body = (await res.json()) as { + meta?: { errorCode?: string; retryAt?: string }; + }; + expect(body.meta?.errorCode).toBe("documents.inbound.rateLimited"); + expect(body.meta?.retryAt).toBe(new Date(resetAt).toISOString()); + expect(refundRateLimit).not.toHaveBeenCalled(); + }); + + it("refunds the charged slot when input preparation fails before dispatch", async () => { + vi.mocked(decryptDocumentContent).mockImplementation(() => { + throw new Error("bad codec"); + }); + const res = await POST(visionReq("doc-1") as never, ctx("doc-1") as never); + expect(res.status).toBe(422); + expect(refundRateLimit).toHaveBeenCalledTimes(1); + expect(refundRateLimit).toHaveBeenCalledWith("documents-inbound:user-1"); + expect(transcribeDocument).not.toHaveBeenCalled(); + }); + + it("refunds the charged slot when the daily budget refuses before dispatch", async () => { + vi.mocked(reserveBudget).mockResolvedValue({ + allowed: false, + reserved: 0, + } as never); + const res = await POST(visionReq("doc-1") as never, ctx("doc-1") as never); + expect(res.status).toBe(429); + const body = (await res.json()) as { meta?: { errorCode?: string } }; + expect(body.meta?.errorCode).toBe("documents.inbound.budgetExceeded"); + expect(refundRateLimit).toHaveBeenCalledTimes(1); + expect(transcribeDocument).not.toHaveBeenCalled(); + }); + + it("keeps the slot once the provider was actually dispatched", async () => { + const res = await POST(visionReq("doc-1") as never, ctx("doc-1") as never); + expect(res.status).toBe(200); + expect(transcribeDocument).toHaveBeenCalledTimes(1); + expect(refundRateLimit).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/documents/inbound/[id]/index/route.ts b/src/app/api/documents/inbound/[id]/index/route.ts index 665a93d05..3d47dfea9 100644 --- a/src/app/api/documents/inbound/[id]/index/route.ts +++ b/src/app/api/documents/inbound/[id]/index/route.ts @@ -39,12 +39,12 @@ import { } from "@/lib/ai/coach/budget"; import { auditLog } from "@/lib/auth/audit"; import { - DOCUMENT_AI_BUCKET, - DOCUMENT_AI_LIMIT_PER_HOUR, + checkDocumentAiRateLimit, + documentAiRateLimited, DOCUMENT_AI_TEXT_BODY_MAX_BYTES, - DOCUMENT_AI_WINDOW_MS, loadOwnedDocument, prepareVisionInput, + refundDocumentAiSlot, type LoadedDocument, } from "@/lib/documents/ai-route-support"; import { upsertContentIndex } from "@/lib/documents/content-index"; @@ -56,23 +56,12 @@ import { resolveDocumentVisionProvider } from "@/lib/documents/provider-order"; import { annotate } from "@/lib/logging/context"; import { requireModuleEnabled } from "@/lib/modules/gate"; import { prisma } from "@/lib/db"; -import { checkRateLimit, rateLimitHeaders } from "@/lib/rate-limit"; import { inboundTextExtractSchema } from "@/lib/validations/inbound-documents"; export const dynamic = "force-dynamic"; type RouteParams = { params: Promise<{ id: string }> }; -function rateLimited(rl: Awaited>): Response { - const response = apiError("Too many requests. Try again later.", 429, { - errorCode: "documents.inbound.rateLimited", - }); - for (const [k, v] of Object.entries(rateLimitHeaders(rl))) { - response.headers.set(k, v); - } - return response; -} - export const POST = apiHandler( async (request: NextRequest, { params }: RouteParams) => { const { user } = await requireAuth(); @@ -131,19 +120,19 @@ async function handleTextIndex( }); } - const rl = await checkRateLimit( - `${DOCUMENT_AI_BUCKET}:${userId}`, - DOCUMENT_AI_LIMIT_PER_HOUR, - DOCUMENT_AI_WINDOW_MS, - ); - if (!rl.allowed) return rateLimited(rl); + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); const { data: body, error: jsonError } = await safeJson(request, { maxBytes: DOCUMENT_AI_TEXT_BODY_MAX_BYTES, }); - if (jsonError) return jsonError; + if (jsonError) { + await refundDocumentAiSlot(userId); + return jsonError; + } const parsed = inboundTextExtractSchema.safeParse(body); if (!parsed.success) { + await refundDocumentAiSlot(userId); return apiError("Invalid document text payload", 422, { errorCode: "documents.inbound.extractFailed", }); @@ -177,15 +166,13 @@ async function handleVisionIndex( surface: "insights", }); - const rl = await checkRateLimit( - `${DOCUMENT_AI_BUCKET}:${userId}`, - DOCUMENT_AI_LIMIT_PER_HOUR, - DOCUMENT_AI_WINDOW_MS, - ); - if (!rl.allowed) return rateLimited(rl); + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); const vision = await prepareVisionInput(document, pick.pdfSupported); if (!vision.ok) { + // Preparation failed before any provider dispatch — the slot goes back. + await refundDocumentAiSlot(userId); if (vision.reason === "pdfNeedsAnthropic") { return apiError( "PDF scanning needs a Claude vision provider; use local OCR instead.", @@ -213,6 +200,7 @@ async function handleVisionIndex( resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { + await refundDocumentAiSlot(userId); return apiError("Your AI usage budget for today is reached.", 429, { errorCode: "documents.inbound.budgetExceeded", }); diff --git a/src/app/api/documents/inbound/[id]/suggest/route.ts b/src/app/api/documents/inbound/[id]/suggest/route.ts index 79a939fb2..94f435641 100644 --- a/src/app/api/documents/inbound/[id]/suggest/route.ts +++ b/src/app/api/documents/inbound/[id]/suggest/route.ts @@ -36,12 +36,12 @@ import { type DocumentSuggestion, } from "@/lib/documents/assist"; import { - DOCUMENT_AI_BUCKET, - DOCUMENT_AI_LIMIT_PER_HOUR, + checkDocumentAiRateLimit, + documentAiRateLimited, DOCUMENT_AI_TEXT_BODY_MAX_BYTES, - DOCUMENT_AI_WINDOW_MS, loadOwnedDocument, prepareVisionInput, + refundDocumentAiSlot, type LoadedDocument, } from "@/lib/documents/ai-route-support"; import { @@ -51,23 +51,12 @@ import { import { annotate } from "@/lib/logging/context"; import { requireModuleEnabled } from "@/lib/modules/gate"; import { prisma } from "@/lib/db"; -import { checkRateLimit, rateLimitHeaders } from "@/lib/rate-limit"; import { inboundTextExtractSchema } from "@/lib/validations/inbound-documents"; export const dynamic = "force-dynamic"; type RouteParams = { params: Promise<{ id: string }> }; -function rateLimited(rl: Awaited>): Response { - const response = apiError("Too many requests. Try again later.", 429, { - errorCode: "documents.inbound.rateLimited", - }); - for (const [k, v] of Object.entries(rateLimitHeaders(rl))) { - response.headers.set(k, v); - } - return response; -} - export const POST = apiHandler( async (request: NextRequest, { params }: RouteParams) => { const { user } = await requireAuth(); @@ -146,19 +135,19 @@ async function handleTextSuggest( surface: "insights", }); - const rl = await checkRateLimit( - `${DOCUMENT_AI_BUCKET}:${userId}`, - DOCUMENT_AI_LIMIT_PER_HOUR, - DOCUMENT_AI_WINDOW_MS, - ); - if (!rl.allowed) return rateLimited(rl); + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); const { data: body, error: jsonError } = await safeJson(request, { maxBytes: DOCUMENT_AI_TEXT_BODY_MAX_BYTES, }); - if (jsonError) return jsonError; + if (jsonError) { + await refundDocumentAiSlot(userId); + return jsonError; + } const parsed = inboundTextExtractSchema.safeParse(body); if (!parsed.success) { + await refundDocumentAiSlot(userId); return apiError("Invalid document text payload", 422, { errorCode: "documents.inbound.extractFailed", }); @@ -172,6 +161,7 @@ async function handleTextSuggest( resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { + await refundDocumentAiSlot(userId); return apiError("Your AI usage budget for today is reached.", 429, { errorCode: "documents.inbound.budgetExceeded", }); @@ -226,15 +216,13 @@ async function handleVisionSuggest( surface: "insights", }); - const rl = await checkRateLimit( - `${DOCUMENT_AI_BUCKET}:${userId}`, - DOCUMENT_AI_LIMIT_PER_HOUR, - DOCUMENT_AI_WINDOW_MS, - ); - if (!rl.allowed) return rateLimited(rl); + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); const vision = await prepareVisionInput(document, pick.pdfSupported); if (!vision.ok) { + // Preparation failed before any provider dispatch — the slot goes back. + await refundDocumentAiSlot(userId); if (vision.reason === "pdfNeedsAnthropic") { return apiError( "PDF scanning needs a Claude vision provider; use local OCR instead.", @@ -262,6 +250,7 @@ async function handleVisionSuggest( resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { + await refundDocumentAiSlot(userId); return apiError("Your AI usage budget for today is reached.", 429, { errorCode: "documents.inbound.budgetExceeded", }); diff --git a/src/app/api/documents/inbound/[id]/summary/route.ts b/src/app/api/documents/inbound/[id]/summary/route.ts index 9003afbe4..55fbea7ec 100644 --- a/src/app/api/documents/inbound/[id]/summary/route.ts +++ b/src/app/api/documents/inbound/[id]/summary/route.ts @@ -38,12 +38,12 @@ import { } from "@/lib/ai/coach/budget"; import { auditLog } from "@/lib/auth/audit"; import { - DOCUMENT_AI_BUCKET, - DOCUMENT_AI_LIMIT_PER_HOUR, + checkDocumentAiRateLimit, + documentAiRateLimited, DOCUMENT_AI_TEXT_BODY_MAX_BYTES, - DOCUMENT_AI_WINDOW_MS, loadOwnedDocument, prepareVisionInput, + refundDocumentAiSlot, type LoadedDocument, } from "@/lib/documents/ai-route-support"; import { @@ -62,7 +62,6 @@ import { import { annotate } from "@/lib/logging/context"; import { requireModuleEnabled } from "@/lib/modules/gate"; import { prisma } from "@/lib/db"; -import { checkRateLimit, rateLimitHeaders } from "@/lib/rate-limit"; import { DOCUMENT_SUMMARY_MODES, inboundTextExtractSchema, @@ -76,16 +75,6 @@ export const dynamic = "force-dynamic"; type RouteParams = { params: Promise<{ id: string }> }; -function rateLimited(rl: Awaited>): Response { - const response = apiError("Too many requests. Try again later.", 429, { - errorCode: "documents.inbound.rateLimited", - }); - for (const [k, v] of Object.entries(rateLimitHeaders(rl))) { - response.headers.set(k, v); - } - return response; -} - type SummaryPersistenceTarget = { userId: string; documentId: string; @@ -295,19 +284,19 @@ async function handleTextSummary( }); } - const rl = await checkRateLimit( - `${DOCUMENT_AI_BUCKET}:${userId}`, - DOCUMENT_AI_LIMIT_PER_HOUR, - DOCUMENT_AI_WINDOW_MS, - ); - if (!rl.allowed) return rateLimited(rl); + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); const { data: body, error: jsonError } = await safeJson(request, { maxBytes: DOCUMENT_AI_TEXT_BODY_MAX_BYTES, }); - if (jsonError) return jsonError; + if (jsonError) { + await refundDocumentAiSlot(userId); + return jsonError; + } const parsed = inboundTextExtractSchema.safeParse(body); if (!parsed.success) { + await refundDocumentAiSlot(userId); return apiError("Invalid document text payload", 422, { errorCode: "documents.inbound.extractFailed", }); @@ -326,15 +315,22 @@ async function handleTextSummary( const { pick } = await resolveDocumentTextProvider(userId); if (!pick) { + await refundDocumentAiSlot(userId); return apiError("No AI provider is configured", 422, { errorCode: "documents.inbound.providerUnsupported", }); } - await assertDocumentEgressConsent({ - userId, - providerType: pick.providerType, - surface: "insights", - }); + try { + await assertDocumentEgressConsent({ + userId, + providerType: pick.providerType, + surface: "insights", + }); + } catch (err) { + // Refused before any dispatch — the slot goes back with the refusal. + await refundDocumentAiSlot(userId); + throw err; + } const dateKey = buildDateKey(); const reservation = await reserveBudget( @@ -344,6 +340,7 @@ async function handleTextSummary( resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { + await refundDocumentAiSlot(userId); return apiError("Your AI usage budget for today is reached.", 429, { errorCode: "documents.inbound.budgetExceeded", }); @@ -394,15 +391,13 @@ async function handleVisionSummary( surface: "insights", }); - const rl = await checkRateLimit( - `${DOCUMENT_AI_BUCKET}:${userId}`, - DOCUMENT_AI_LIMIT_PER_HOUR, - DOCUMENT_AI_WINDOW_MS, - ); - if (!rl.allowed) return rateLimited(rl); + const rl = await checkDocumentAiRateLimit(userId); + if (!rl.allowed) return documentAiRateLimited(rl); const vision = await prepareVisionInput(document, pick.pdfSupported); if (!vision.ok) { + // Preparation failed before any provider dispatch — the slot goes back. + await refundDocumentAiSlot(userId); if (vision.reason === "pdfNeedsAnthropic") { return apiError( "PDF scanning needs a Claude vision provider; use local OCR instead.", @@ -430,6 +425,7 @@ async function handleVisionSummary( resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { + await refundDocumentAiSlot(userId); return apiError("Your AI usage budget for today is reached.", 429, { errorCode: "documents.inbound.budgetExceeded", }); diff --git a/src/components/documents/__tests__/document-ai-panels.test.tsx b/src/components/documents/__tests__/document-ai-panels.test.tsx index ea24d4aa4..53748f937 100644 --- a/src/components/documents/__tests__/document-ai-panels.test.tsx +++ b/src/components/documents/__tests__/document-ai-panels.test.tsx @@ -153,7 +153,7 @@ describe("", () => { output="summary" result={null} isPending - errorKey={null} + errorText={null} onClose={noop} />, ); @@ -167,7 +167,7 @@ describe("", () => { output="summary" result={{ summary: "A lipid panel dated 1 March 2026." }} isPending={false} - errorKey={null} + errorText={null} onClose={noop} />, ); @@ -180,7 +180,7 @@ describe("", () => { output="text" result={{ text: "LDL 3.1 mmol/L" }} isPending={false} - errorKey={null} + errorText={null} onClose={noop} />, ); @@ -195,7 +195,7 @@ describe("", () => { output="summary" result={null} isPending={false} - errorKey="documents.assist.errorGeneric" + errorText="Couldn't read the document. Try again." onClose={noop} />, ); @@ -209,7 +209,7 @@ describe("", () => { output="summary" result={{ summary: "A lipid panel dated 1 March 2026." }} isPending={false} - errorKey={null} + errorText={null} onClose={noop} />, ); diff --git a/src/components/documents/__tests__/use-document-assist.test.ts b/src/components/documents/__tests__/use-document-assist.test.ts new file mode 100644 index 000000000..0a661bc7f --- /dev/null +++ b/src/components/documents/__tests__/use-document-assist.test.ts @@ -0,0 +1,59 @@ +/** + * The honest-wait derivation for a document-AI 429. The routes mirror the + * bucket's reset instant into `meta.retryAt`; the client turns it into "try + * again in about N minutes". This pins the mapping: only a document-AI + * rate-limit error yields minutes, the floor is one minute, and anything + * malformed falls back to null (the generic copy). + */ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { ApiError } from "@/lib/api/api-fetch"; + +import { documentAiRetryMinutes } from "../use-document-assist"; + +const NOW = new Date("2026-08-12T12:00:00.000Z").getTime(); + +beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(NOW); +}); + +afterEach(() => { + vi.useRealTimers(); +}); + +function rateLimited(retryAt?: unknown): ApiError { + return new ApiError("Too many document AI requests this hour.", 429, { + errorCode: "documents.inbound.rateLimited", + ...(retryAt !== undefined ? { retryAt } : {}), + }); +} + +describe("documentAiRetryMinutes", () => { + it("derives whole minutes from the reset instant, rounding up", () => { + const err = rateLimited(new Date(NOW + 42.4 * 60_000).toISOString()); + expect(documentAiRetryMinutes(err)).toBe(43); + }); + + it("floors at one minute for a reset only seconds away", () => { + const err = rateLimited(new Date(NOW + 20_000).toISOString()); + expect(documentAiRetryMinutes(err)).toBe(1); + }); + + it("returns null when the error carries no usable instant", () => { + expect(documentAiRetryMinutes(rateLimited())).toBeNull(); + expect(documentAiRetryMinutes(rateLimited("soon"))).toBeNull(); + expect(documentAiRetryMinutes(rateLimited(42))).toBeNull(); + }); + + it("returns null for any other error class", () => { + expect( + documentAiRetryMinutes( + new ApiError("Budget", 429, { + errorCode: "documents.inbound.budgetExceeded", + }), + ), + ).toBeNull(); + expect(documentAiRetryMinutes(new Error("network"))).toBeNull(); + }); +}); diff --git a/src/components/documents/document-ai-panels.tsx b/src/components/documents/document-ai-panels.tsx index 43f08e1a0..919711f4c 100644 --- a/src/components/documents/document-ai-panels.tsx +++ b/src/components/documents/document-ai-panels.tsx @@ -202,13 +202,14 @@ export function DocumentSummaryPanel({ output, result, isPending, - errorKey, + errorText, onClose, }: { output: DocumentSummaryMode; result: DocumentDescribeResult | null; isPending: boolean; - errorKey: string | null; + /** Resolved user-facing error sentence (already translated), or null. */ + errorText: string | null; onClose: () => void; }) { const { t } = useTranslations(); @@ -247,9 +248,9 @@ export function DocumentSummaryPanel({ - ) : errorKey ? ( + ) : errorText ? (

- {t(errorKey)} + {errorText}

) : body !== null ? (

void; suggestPending: boolean; + /** Resolved user-facing error sentence (already translated), or null. */ suggestErrorKey: string | null; onSummarise: (output: DocumentSummaryMode) => void; summaryPending: boolean; @@ -105,6 +106,7 @@ export function DocumentAiSection({ onDismissSuggestion: () => void; summaryOutput: DocumentSummaryMode | null; summaryResult: DocumentDescribeResult | null; + /** Resolved user-facing error sentence (already translated), or null. */ summaryErrorKey: string | null; onCloseSummary: () => void; hasContentIndex: boolean; @@ -233,7 +235,7 @@ export function DocumentAiSection({ {suggestErrorKey ? (

- {t(suggestErrorKey)} + {suggestErrorKey}

) : null} @@ -255,7 +257,7 @@ export function DocumentAiSection({ output={summaryOutput} result={summaryResult} isPending={summaryPending} - errorKey={summaryErrorKey} + errorText={summaryErrorKey} onClose={onCloseSummary} /> ) : null} diff --git a/src/components/documents/document-detail-sheet.tsx b/src/components/documents/document-detail-sheet.tsx index 2d947998b..4d1707585 100644 --- a/src/components/documents/document-detail-sheet.tsx +++ b/src/components/documents/document-detail-sheet.tsx @@ -87,8 +87,8 @@ import { DocumentShareSheet } from "./document-share-sheet"; import { DOCUMENT_KIND_ICONS } from "./document-kind-meta"; import { useIndexDocument } from "./use-content-index"; import { - documentAiErrorKey, useDocumentAiCapability, + useDocumentAiErrorText, useDocumentsAutoAiRead, useDocumentSummary, useSuggestDetails, @@ -393,6 +393,9 @@ export function DocumentDetailSheet({ ); const suggest = useSuggestDetails(); const summary = useDocumentSummary(); + // Resolves an AI error to final copy — a 429 names the actual wait, derived + // from the reset instant the route mirrors into the error meta. + const aiErrorText = useDocumentAiErrorText(); // Stored generation has an independent mutation state so transient // Summarise/Show text requests cannot change this block's spinner or result. const storedSummary = useDocumentSummary(); @@ -593,7 +596,7 @@ export function DocumentDetailSheet({ queryKey: queryKeys.inboundDocument(documentId), }); }, - onError: (error) => toast.error(t(documentAiErrorKey(error))), + onError: (error) => toast.error(aiErrorText(error)), }, ); }; @@ -604,7 +607,7 @@ export function DocumentDetailSheet({ { mode: aiMode, target: aiTarget }, { onSuccess: () => toast.success(t("documents.ai.readDone")), - onError: (error) => toast.error(t(documentAiErrorKey(error))), + onError: (error) => toast.error(aiErrorText(error)), }, ); }; @@ -836,7 +839,7 @@ export function DocumentDetailSheet({ onSuggest={runSuggest} suggestPending={suggest.isPending} suggestErrorKey={ - suggest.isError ? documentAiErrorKey(suggest.error) : null + suggest.isError ? aiErrorText(suggest.error) : null } onSummarise={runSummary} summaryPending={summary.isPending} @@ -851,7 +854,7 @@ export function DocumentDetailSheet({ summaryOutput={summaryOutput} summaryResult={summary.data ?? null} summaryErrorKey={ - summary.isError ? documentAiErrorKey(summary.error) : null + summary.isError ? aiErrorText(summary.error) : null } onCloseSummary={() => { setSummaryOutput(null); diff --git a/src/components/documents/use-document-assist.ts b/src/components/documents/use-document-assist.ts index da5ae7814..e1da97b87 100644 --- a/src/components/documents/use-document-assist.ts +++ b/src/components/documents/use-document-assist.ts @@ -13,9 +13,12 @@ * prefills; the summary is transient. The only mutation is the caller applying * a reviewed draft through the existing edit-on-commit machinery. */ +import { useCallback } from "react"; + import { useMutation, useQuery } from "@tanstack/react-query"; import { apiGet, ApiError } from "@/lib/api/api-fetch"; +import { useTranslations } from "@/lib/i18n/context"; import { queryKeys } from "@/lib/query-keys"; import type { DocumentAiCapabilityDto, @@ -129,6 +132,45 @@ export function useDocumentSummary() { * `meta.errorCode`; client-side precondition failures carry a `reason`. The * default is the calm "try again" message — never a raw provider string. */ +/** + * The actual wait, in whole minutes, behind a document-AI 429. The routes + * mirror the bucket's reset instant into the error meta (`retryAt`, ISO 8601); + * the window is an hour, so a fixed "try again in a few minutes" understated + * the wait for most of it. Null when the error is not a document-AI rate-limit + * response or carries no usable instant — the caller falls back to the generic + * copy. Floored at one minute: a reset seconds away still reads as "about + * 1 minute", never "0 minutes". + */ +export function documentAiRetryMinutes(err: unknown): number | null { + if (!(err instanceof ApiError)) return null; + if (err.meta?.errorCode !== "documents.inbound.rateLimited") return null; + const raw = err.meta?.retryAt; + if (typeof raw !== "string") return null; + const at = Date.parse(raw); + if (!Number.isFinite(at)) return null; + return Math.max(1, Math.ceil((at - Date.now()) / 60_000)); +} + +/** + * Resolve a document-AI error to the user-facing sentence. Rate-limited + * responses derive their copy from the ACTUAL reset instant ("try again in + * about N minutes"); everything else maps through `documentAiErrorKey`. The + * result is final text — callers render or toast it as-is. + */ +export function useDocumentAiErrorText(): (err: unknown) => string { + const { t, tCount } = useTranslations(); + return useCallback( + (err: unknown) => { + const minutes = documentAiRetryMinutes(err); + if (minutes !== null) { + return tCount("documents.assist.errorRateLimitedWait", minutes); + } + return t(documentAiErrorKey(err)); + }, + [t, tCount], + ); +} + export function documentAiErrorKey(err: unknown): string { if (err instanceof DocumentAssistClientError) { switch (err.reason) { diff --git a/src/lib/documents/__tests__/ai-route-support-limit.test.ts b/src/lib/documents/__tests__/ai-route-support-limit.test.ts new file mode 100644 index 000000000..d2a4b1110 --- /dev/null +++ b/src/lib/documents/__tests__/ai-route-support-limit.test.ts @@ -0,0 +1,94 @@ +/** + * The document-AI hourly ceiling is an operator setting + * (`DOCUMENT_AI_LIMIT_PER_HOUR`), not a constant. This pins the resolver's + * contract: unset / non-numeric falls back to the default 6, a numeric value + * takes effect, and the clamp keeps a typo from zeroing the surface out + * (floor 1) or removing the cap entirely (ceiling 1000). + * + * Also pins the shared 429 builder: the response mirrors the bucket's reset + * instant into `meta.retryAt` so the client can name the ACTUAL wait instead + * of a fixed "a few minutes". + */ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@/lib/db", () => ({ + prisma: { + inboundDocument: { findFirst: vi.fn() }, + $queryRaw: vi.fn(), + $executeRaw: vi.fn(), + }, +})); +vi.mock("@/lib/documents/store", () => ({ + decryptDocumentContent: vi.fn(), +})); +vi.mock("@/lib/documents/rasterize-pdf", () => ({ + rasterizePdf: vi.fn(), +})); +vi.mock("@/lib/labs/ocr-upload", () => ({ + detectOcrMimeType: vi.fn(), +})); + +import { + documentAiRateLimited, + resolveDocumentAiLimitPerHour, +} from "../ai-route-support"; + +const ORIGINAL = process.env.DOCUMENT_AI_LIMIT_PER_HOUR; + +beforeEach(() => { + delete process.env.DOCUMENT_AI_LIMIT_PER_HOUR; +}); + +afterEach(() => { + if (ORIGINAL === undefined) delete process.env.DOCUMENT_AI_LIMIT_PER_HOUR; + else process.env.DOCUMENT_AI_LIMIT_PER_HOUR = ORIGINAL; +}); + +describe("resolveDocumentAiLimitPerHour", () => { + it("defaults to 6 when the variable is unset", () => { + expect(resolveDocumentAiLimitPerHour()).toBe(6); + }); + + it("takes an operator value from the environment", () => { + process.env.DOCUMENT_AI_LIMIT_PER_HOUR = "12"; + expect(resolveDocumentAiLimitPerHour()).toBe(12); + }); + + it("falls back to the default on a non-numeric value", () => { + process.env.DOCUMENT_AI_LIMIT_PER_HOUR = "plenty"; + expect(resolveDocumentAiLimitPerHour()).toBe(6); + }); + + it("floors at 1 so a zero/negative typo cannot zero the surface out", () => { + process.env.DOCUMENT_AI_LIMIT_PER_HOUR = "0"; + expect(resolveDocumentAiLimitPerHour()).toBe(1); + process.env.DOCUMENT_AI_LIMIT_PER_HOUR = "-4"; + expect(resolveDocumentAiLimitPerHour()).toBe(1); + }); + + it("caps at 1000 so a typo cannot remove the ceiling entirely", () => { + process.env.DOCUMENT_AI_LIMIT_PER_HOUR = "500000"; + expect(resolveDocumentAiLimitPerHour()).toBe(1000); + }); +}); + +describe("documentAiRateLimited", () => { + it("mirrors the reset instant into meta.retryAt and the headers", async () => { + const resetAt = Date.now() + 42 * 60_000; + const response = documentAiRateLimited({ + allowed: false, + remaining: 0, + resetAt, + }); + expect(response.status).toBe(429); + expect(response.headers.get("X-RateLimit-Reset")).toBe( + new Date(resetAt).toISOString(), + ); + const body = (await response.json()) as { + error: string | null; + meta?: { errorCode?: string; retryAt?: string }; + }; + expect(body.meta?.errorCode).toBe("documents.inbound.rateLimited"); + expect(body.meta?.retryAt).toBe(new Date(resetAt).toISOString()); + }); +}); diff --git a/src/lib/documents/ai-route-support.ts b/src/lib/documents/ai-route-support.ts index a974ad5ab..e294000fe 100644 --- a/src/lib/documents/ai-route-support.ts +++ b/src/lib/documents/ai-route-support.ts @@ -9,16 +9,88 @@ */ import { Buffer } from "node:buffer"; +import { apiError } from "@/lib/api-response"; import { prisma } from "@/lib/db"; import { rasterizePdf } from "@/lib/documents/rasterize-pdf"; import { decryptDocumentContent } from "@/lib/documents/store"; import { detectOcrMimeType } from "@/lib/labs/ocr-upload"; +import { + checkRateLimit, + rateLimitHeaders, + refundRateLimit, + type RateLimitResult, +} from "@/lib/rate-limit"; + +/** + * The document-AI actions (extract / suggest / summary / index) share one + * per-user hourly egress bucket. The ceiling is operator-tunable via + * `DOCUMENT_AI_LIMIT_PER_HOUR` (default 6, clamped to 1–1000 so a typo can + * neither zero the surface out nor remove the cap entirely) — same posture as + * `INSIGHTS_RATE_LIMIT_PER_HOUR` on the briefing route. Unset / non-numeric + * falls back to the default. + */ +const DOCUMENT_AI_LIMIT_DEFAULT = 6; +const DOCUMENT_AI_LIMIT_MIN = 1; +const DOCUMENT_AI_LIMIT_MAX = 1000; + +export function resolveDocumentAiLimitPerHour(): number { + const raw = process.env.DOCUMENT_AI_LIMIT_PER_HOUR; + if (!raw) return DOCUMENT_AI_LIMIT_DEFAULT; + const parsed = Number.parseInt(raw, 10); + if (!Number.isFinite(parsed)) return DOCUMENT_AI_LIMIT_DEFAULT; + return Math.min( + DOCUMENT_AI_LIMIT_MAX, + Math.max(DOCUMENT_AI_LIMIT_MIN, parsed), + ); +} -/** The three AI actions share the extract route's 6/hour egress bucket. */ -export const DOCUMENT_AI_LIMIT_PER_HOUR = 6; export const DOCUMENT_AI_WINDOW_MS = 60 * 60 * 1000; /** The shared rate-limit bucket key prefix (identical to the extract route). */ export const DOCUMENT_AI_BUCKET = "documents-inbound"; + +/** Charge one slot of the shared per-user document-AI bucket. */ +export async function checkDocumentAiRateLimit( + userId: string, +): Promise { + return checkRateLimit( + `${DOCUMENT_AI_BUCKET}:${userId}`, + resolveDocumentAiLimitPerHour(), + DOCUMENT_AI_WINDOW_MS, + ); +} + +/** + * Give a charged slot back when the request failed BEFORE the metered work + * (provider dispatch, or the local index/echo work on the no-egress paths). + * The bucket is charged early so a 429 stays cheap, but a slot must reflect + * work actually performed — a malformed body, an unreadable original, or an + * exhausted budget dispatched nothing and burns no allowance. Best-effort: + * a failed refund never masks the response the caller is about to return. + */ +export async function refundDocumentAiSlot(userId: string): Promise { + try { + await refundRateLimit(`${DOCUMENT_AI_BUCKET}:${userId}`); + } catch { + // The original (error) response matters more than the bookkeeping. + } +} + +/** + * The shared 429 for the document-AI bucket. Carries the standard + * `X-RateLimit-*` headers AND mirrors the reset instant into the error meta + * (`retryAt`, ISO 8601) so the client can tell the user the ACTUAL wait — + * the window is an hour, not "a few minutes". + */ +export function documentAiRateLimited(rl: RateLimitResult): Response { + const response = apiError("Too many document AI requests this hour.", 429, { + errorCode: "documents.inbound.rateLimited", + retryAt: new Date(rl.resetAt).toISOString(), + }); + for (const [k, v] of Object.entries(rateLimitHeaders(rl))) { + response.headers.set(k, v); + } + return response; +} /** Cap on a browser-OCR text body (mirrors the extract text mode). */ export const DOCUMENT_AI_TEXT_BODY_MAX_BYTES = 512 * 1024; diff --git a/src/lib/openapi/routes/documents.ts b/src/lib/openapi/routes/documents.ts index 3d4afe2be..2462f681f 100644 --- a/src/lib/openapi/routes/documents.ts +++ b/src/lib/openapi/routes/documents.ts @@ -769,7 +769,7 @@ export const inboundDocumentPaths: NonNullable = { tags: ["Documents"], summary: "Suggest filing metadata (drafts only)", description: - 'Optional AI assist on an already-stored document. Runs ONE provider call over the stored original (VISION, empty body) or browser-OCR\'d text (TEXT, `application/json` `{ mode: "text", text }`, opt-in local OCR) and returns a `{ title, kind, documentDate }` DRAFT for the edit form. AI-consent / rate / budget gated (shares the 6/hour extract bucket). WRITES NOTHING — never stages facts, never flips status; the user reviews and saves. 422 (`documents.inbound.providerUnsupported`) with no provider configured. Never interprets or diagnoses; the title is a neutral filing label.', + 'Optional AI assist on an already-stored document. Runs ONE provider call over the stored original (VISION, empty body) or browser-OCR\'d text (TEXT, `application/json` `{ mode: "text", text }`, opt-in local OCR) and returns a `{ title, kind, documentDate }` DRAFT for the edit form. AI-consent / rate / budget gated (shares the per-user document-AI bucket with extract/summary/index — default 6/hour, operator-tunable via `DOCUMENT_AI_LIMIT_PER_HOUR`; a slot is only consumed when the request reaches the provider, and the 429 carries the reset instant in `X-RateLimit-Reset` and `meta.retryAt`). WRITES NOTHING — never stages facts, never flips status; the user reviews and saves. 422 (`documents.inbound.providerUnsupported`) with no provider configured. Never interprets or diagnoses; the title is a neutral filing label.', parameters: [ { name: "id", in: "path", required: true, schema: { type: "string" } }, ], diff --git a/src/lib/rate-limit.ts b/src/lib/rate-limit.ts index d72de6da5..cad590f3b 100644 --- a/src/lib/rate-limit.ts +++ b/src/lib/rate-limit.ts @@ -52,6 +52,24 @@ export async function checkRateLimit( }; } +/** + * Give one slot back to a bucket whose charge turned out to be premature. + * + * A caller that charges the bucket up front (to keep the 429 cheap) but then + * fails BEFORE doing the metered work refunds the slot so the user's allowance + * reflects work actually performed, not attempts. Scoped to the live window: + * once `reset_at` has passed the counter is stale and a decrement would bleed + * into the NEXT window's budget, so an expired row is left untouched. Never + * drops below zero (concurrent refunds after a window flip are a no-op). + */ +export async function refundRateLimit(key: string): Promise { + await prisma.$executeRaw` + UPDATE rate_limits + SET count = GREATEST(count - 1, 0) + WHERE key = ${key} AND reset_at >= NOW() + `; +} + /** * Rate limit response headers. */