From bb58b40b2aa79d81b13cd9c22d92a2a55bf4d90c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Sat, 3 Oct 2026 21:14:43 +0200 Subject: [PATCH] Run fact extraction in the background, re-checking under a row lock The extract route accepts `Prefer: respond-async` for all three inputs (text, stored, vision) and queues a DOCUMENT_EXTRACT run (added to migration 0370). The run's result is the count of staged facts and the document status; the facts themselves are read from the document. Without the header the route answers with the document detail as before; both paths run one shared body. A background read can take minutes, and a review the person finishes in the meantime must not have its approved facts replaced. Staging now locks the document row and asks again whether it may stage, refusing with the route's own alreadyConfirmed / alreadyPartlyConfirmed codes. The synchronous path gets the same re-check. The vault's stored-text extract uses the background path and shows the shared waiting line in place of its hint while the read runs. --- docs/api/openapi.yaml | 95 ++-- .../migration.sql | 5 +- prisma/schema.prisma | 3 + .../ai-capability-route-inventory.test.ts | 1 + .../[id]/extract/__tests__/route.test.ts | 32 +- .../documents/inbound/[id]/extract/route.ts | 462 +++++++----------- .../__tests__/document-facts-review.test.tsx | 2 + .../documents/document-facts-review.tsx | 25 +- .../documents/use-document-facts.ts | 39 +- src/lib/documents/ai-runs/extract-run.ts | 287 +++++++++++ src/lib/documents/ai-runs/types.ts | 23 +- .../jobs/__tests__/document-ai-run.test.ts | 116 +++++ src/lib/jobs/document-ai-run.ts | 62 +++ src/lib/openapi/routes/ai-runs.ts | 14 +- src/lib/openapi/routes/documents.ts | 21 +- tests/integration/document-ai-runs.test.ts | 168 +++++++ 16 files changed, 1003 insertions(+), 352 deletions(-) create mode 100644 src/lib/documents/ai-runs/extract-run.ts diff --git a/docs/api/openapi.yaml b/docs/api/openapi.yaml index 35847ef8a..46f4d3b65 100644 --- a/docs/api/openapi.yaml +++ b/docs/api/openapi.yaml @@ -16701,13 +16701,32 @@ paths: structures the document's own stored extracted text (its content index) — the manual recovery for a skipped/failed automatic staging run, 422 `documents.inbound.notIndexed` when no stored text exists. Extraction reproduces what the document states — it never interprets. Nothing reaches the structured stores here; the - confirm route is the only write path." + confirm route is the only write path. With `Prefer: respond-async` (v1.40) the extraction runs in the + background: after the same quick refusals the route answers 202 with a run id, a second request while an + extraction of this document from the same input is queued or running answers with that run, and `GET + /api/ai-runs/{id}` serves `DocumentExtractRunResult` (`{ documentId, factsStaged, status }`, not the document + detail) as the run's `result`. Staging re-checks the document under a row lock, so a review finished while the + read ran fails the run with `documents.inbound.alreadyPartlyConfirmed` (409) or + `documents.inbound.alreadyConfirmed` (422) and leaves the review untouched; the synchronous form answers the + same in that race." parameters: - name: id in: path required: true schema: type: string + - &a37 + name: Prefer + in: header + required: false + schema: + type: string + example: respond-async + description: "Send `respond-async` (RFC 7240) to run the read in the background: the route answers 202 with a run id + once every quick refusal has passed, and `GET /api/ai-runs/{id}` serves the result. Without it the route + answers synchronously with the result, as it always has; that form is retained for clients that have not + adopted runs, and its retirement will be announced one minor release ahead. A preference the route cannot + apply is ignored and the response then carries no `Preference-Applied`." requestBody: required: false content: @@ -16721,8 +16740,9 @@ paths: application/json: schema: $ref: "#/components/schemas/InboundExtractEnvelope" + "202": *a36 "401": *a1 - "403": &a37 + "403": &a38 description: "AI refusal: the `documentAi` capability is unavailable for this record. The envelope carries `meta.errorCode`, `meta.capability` and `meta.reason` (one of `AiUnavailableReason`). `assistant.disabled.documentAi` — the operator turned off reading documents, or every AI feature (the master @@ -16746,7 +16766,7 @@ paths: $ref: "#/components/schemas/ErrorEnvelope" "422": *a4 "429": *a5 - "503": &a38 + "503": &a39 description: "`ai.unavailable` — the inputs behind the AI capability could not be read, and the answer failed closed. Nothing was sent. Retry later." content: @@ -16776,18 +16796,7 @@ paths: required: true schema: type: string - - &a39 - name: Prefer - in: header - required: false - schema: - type: string - example: respond-async - description: "Send `respond-async` (RFC 7240) to run the read in the background: the route answers 202 with a run id - once every quick refusal has passed, and `GET /api/ai-runs/{id}` serves the result. Without it the route - answers synchronously with the result, as it always has; that form is retained for clients that have not - adopted runs, and its retirement will be announced one minor release ahead. A preference the route cannot - apply is ignored and the response then carries no `Preference-Applied`." + - *a37 requestBody: required: false content: @@ -16824,10 +16833,10 @@ paths: $ref: "#/components/schemas/DocumentSuggestEnvelope" "202": *a36 "401": *a1 - "403": *a37 + "403": *a38 "422": *a4 "429": *a5 - "503": *a38 + "503": *a39 /api/documents/inbound/{id}/summary: post: tags: @@ -16877,7 +16886,7 @@ paths: type: boolean default: false description: With `persist=true`, explicitly replace an existing stored summary. - - *a39 + - *a37 requestBody: required: false content: @@ -16914,10 +16923,10 @@ paths: $ref: "#/components/schemas/DocumentSummaryEnvelope" "202": *a36 "401": *a1 - "403": *a37 + "403": *a38 "422": *a4 "429": *a5 - "503": *a38 + "503": *a39 /api/documents/inbound/{id}/index: post: tags: @@ -16943,7 +16952,7 @@ paths: required: true schema: type: string - - *a39 + - *a37 requestBody: required: false content: @@ -16980,10 +16989,10 @@ paths: $ref: "#/components/schemas/DocumentIndexEnvelope" "202": *a36 "401": *a1 - "403": *a37 + "403": *a38 "422": *a4 "429": *a5 - "503": *a38 + "503": *a39 /api/documents/inbound/{id}/chat: post: tags: @@ -17023,10 +17032,10 @@ paths: type: string description: "SSE frames: `data: \\n\\n`. See the operation description for the per-`type` frame shapes." "401": *a1 - "403": *a37 + "403": *a38 "422": *a4 "429": *a5 - "503": *a38 + "503": *a39 get: tags: - Documents @@ -17089,10 +17098,10 @@ paths: schema: $ref: "#/components/schemas/DocumentReindexEnvelope" "401": *a1 - "403": *a37 + "403": *a38 "422": *a4 "429": *a5 - "503": *a38 + "503": *a39 /api/documents/inbound/{id}/facts/{factId}: patch: tags: @@ -63669,6 +63678,7 @@ components: - LABS_OCR_EXTRACT - DOCUMENT_SUMMARY - DOCUMENT_SUGGEST + - DOCUMENT_EXTRACT documentId: anyOf: - type: string @@ -63688,11 +63698,13 @@ components: - $ref: "#/components/schemas/OcrExtractResponse" - $ref: "#/components/schemas/DocumentSummaryResponse" - $ref: "#/components/schemas/DocumentSuggestResponse" + - $ref: "#/components/schemas/DocumentExtractRunResult" - type: "null" description: "Present once `status` is `SUCCEEDED`: the body the queuing route answers with synchronously. `DOCUMENT_INDEX` carries `DocumentIndexResponse`, `LABS_OCR_EXTRACT` carries `OcrExtractResponse`, - `DOCUMENT_SUMMARY` carries `DocumentSummaryResponse`, `DOCUMENT_SUGGEST` carries `DocumentSuggestResponse`. - Decode by `kind`; treat an unknown `kind` as a run this client did not start." + `DOCUMENT_SUMMARY` carries `DocumentSummaryResponse`, `DOCUMENT_SUGGEST` carries `DocumentSuggestResponse`, + `DOCUMENT_EXTRACT` carries `DocumentExtractRunResult` (not the extract route's detail body). Decode by + `kind`; treat an unknown `kind` as a run this client did not start." error: anyOf: - type: object @@ -63951,6 +63963,31 @@ components: required: - suggestions additionalProperties: false + DocumentExtractRunResult: + type: object + properties: + documentId: + type: string + factsStaged: + type: integer + minimum: -9007199254740991 + maximum: 9007199254740991 + description: How many facts were staged PENDING for review. Read them from `GET /api/documents/inbound/{id}`; the + confirm route stays the only write into the record. + status: + type: string + enum: + - STORED + - EXTRACTING + - EXTRACTED + - FAILED + - CONFIRMED + - DISCARDED + required: + - documentId + - factsStaged + - status + additionalProperties: false ListBiomarkersEnvelope: type: object properties: diff --git a/prisma/migrations/0370_document_ai_run_kinds_summary_suggest_extract/migration.sql b/prisma/migrations/0370_document_ai_run_kinds_summary_suggest_extract/migration.sql index 6d1ede60a..a64f219f2 100644 --- a/prisma/migrations/0370_document_ai_run_kinds_summary_suggest_extract/migration.sql +++ b/prisma/migrations/0370_document_ai_run_kinds_summary_suggest_extract/migration.sql @@ -1,6 +1,6 @@ --- Two more kinds of background document AI run. +-- Three more kinds of background document AI run. -- --- Summaries and filing suggestions can now run in the same background worker +-- Summaries, filing suggestions and fact extraction can now run in the same background worker -- as "Read with AI" and the lab scan (0366): the request queues the read and -- answers 202, and the result is polled from the run. The synchronous routes -- are unchanged; the new kinds only name the run. @@ -13,3 +13,4 @@ -- (inert with no rows that carry them). ALTER TYPE "document_ai_run_kind" ADD VALUE IF NOT EXISTS 'DOCUMENT_SUMMARY'; ALTER TYPE "document_ai_run_kind" ADD VALUE IF NOT EXISTS 'DOCUMENT_SUGGEST'; +ALTER TYPE "document_ai_run_kind" ADD VALUE IF NOT EXISTS 'DOCUMENT_EXTRACT'; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 473f7c987..69cd6245f 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -9042,6 +9042,9 @@ enum DocumentAiRunKind { /// Filing suggestions (title, kind, date) for a stored document. Drafts /// only: nothing is written to the document. DOCUMENT_SUGGEST + /// Structured facts read out of a stored document, staged PENDING for the + /// review step. The confirm route stays the only write into the record. + DOCUMENT_EXTRACT @@map("document_ai_run_kind") } diff --git a/src/app/api/__tests__/ai-capability-route-inventory.test.ts b/src/app/api/__tests__/ai-capability-route-inventory.test.ts index 91c3c29c1..2bb7afd8c 100644 --- a/src/app/api/__tests__/ai-capability-route-inventory.test.ts +++ b/src/app/api/__tests__/ai-capability-route-inventory.test.ts @@ -171,6 +171,7 @@ const RUN_QUEUED_BY: Record = { LABS_OCR_EXTRACT: "src/app/api/labs/ocr/extract/route.ts", DOCUMENT_SUMMARY: "src/app/api/documents/inbound/[id]/summary/route.ts", DOCUMENT_SUGGEST: "src/app/api/documents/inbound/[id]/suggest/route.ts", + DOCUMENT_EXTRACT: "src/app/api/documents/inbound/[id]/extract/route.ts", }; describe("background run kinds", () => { diff --git a/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts b/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts index ccb39338b..a03e75fe6 100644 --- a/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts +++ b/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts @@ -19,8 +19,11 @@ import { NextRequest } from "next/server"; */ vi.mock("@/lib/db", () => { + // The staging transaction locks the document row and asks again whether + // it may stage (a review can finish while a background read runs). const tx = { - extractedFact: { deleteMany: vi.fn() }, + $queryRaw: vi.fn(async () => [{ status: "STORED" }]), + extractedFact: { deleteMany: vi.fn(), count: vi.fn(async () => 0) }, inboundDocument: { update: vi.fn(), findUniqueOrThrow: vi.fn() }, }; return { @@ -123,7 +126,11 @@ import { AiUnavailableError } from "@/lib/ai/capabilities/refusal"; const tx = ( prisma as unknown as { __tx: { - extractedFact: { deleteMany: ReturnType }; + $queryRaw: ReturnType; + extractedFact: { + deleteMany: ReturnType; + count: ReturnType; + }; inboundDocument: { update: ReturnType; findUniqueOrThrow: ReturnType; @@ -304,6 +311,27 @@ describe("POST /api/documents/inbound/[id]/extract — stored mode", () => { ); }); + it("refuses to stage when a fact was approved while the read ran, keeping the review", async () => { + tx.extractedFact.count.mockResolvedValueOnce(1); + const res = await POST(storedReq("doc-1") as never, ctx("doc-1") as never); + expect(res.status).toBe(409); + const body = await res.json(); + expect(body.meta.errorCode).toBe( + "documents.inbound.alreadyPartlyConfirmed", + ); + expect(tx.extractedFact.deleteMany).not.toHaveBeenCalled(); + expect(tx.inboundDocument.update).not.toHaveBeenCalled(); + }); + + it("refuses to stage when the document was confirmed while the read ran", async () => { + tx.$queryRaw.mockResolvedValueOnce([{ status: "CONFIRMED" }]); + const res = await POST(storedReq("doc-1") as never, ctx("doc-1") as never); + expect(res.status).toBe(422); + const body = await res.json(); + expect(body.meta.errorCode).toBe("documents.inbound.alreadyConfirmed"); + expect(tx.extractedFact.deleteMany).not.toHaveBeenCalled(); + }); + it("422s (notIndexed) without stored text, before any bucket charge", async () => { vi.mocked(loadDocumentChatText).mockResolvedValue(null as never); const res = await POST(storedReq("doc-1") as never, ctx("doc-1") as never); diff --git a/src/app/api/documents/inbound/[id]/extract/route.ts b/src/app/api/documents/inbound/[id]/extract/route.ts index 930d8c895..1fc8338c5 100644 --- a/src/app/api/documents/inbound/[id]/extract/route.ts +++ b/src/app/api/documents/inbound/[id]/extract/route.ts @@ -22,6 +22,15 @@ * The document is UNTRUSTED (prompt-injection): the server never acts on an * instruction inside it. The staged facts land PENDING for the mandatory * review-then-confirm screen; nothing reaches the structured stores here. + * + * v1.40 — `Prefer: respond-async` (RFC 7240) runs the extraction in the + * background worker: every refusal that can be answered quickly is still + * answered here, then the route answers 202 with a run id. The run's result is + * `{ documentId, factsStaged, status }` rather than the whole detail (the + * facts are staged on the document; the client reads them from there). + * Without the header the route behaves exactly as before; the iPhone app + * relies on that. Both paths run one body (`executeDocumentExtract`), which + * re-checks the document under a row lock before it stages anything. */ import { Buffer } from "node:buffer"; @@ -40,24 +49,13 @@ import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; import { requireAiCapability } from "@/lib/ai/capabilities/gate"; import { buildDateKey, - reconcileSpend, reserveBudget, resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; -import { auditLog } from "@/lib/auth/audit"; +import type { ProviderChainType } from "@/lib/ai/provider-chain"; import { prisma } from "@/lib/db"; -import { - InboundExtractError, - runInboundExtraction, - type InboundExtractionResult, -} from "@/lib/documents/extract"; -import { - decryptDocumentContent, - encryptFactData, - encryptFactProvenance, - serialiseDocumentDetail, -} from "@/lib/documents/store"; +import { serialiseDocumentDetail } from "@/lib/documents/store"; import { requireDocumentTextProvider, requireDocumentVisionProvider, @@ -66,73 +64,34 @@ import { checkDocumentAiRateLimit, documentAiRateLimited, DOCUMENT_AI_TEXT_BODY_MAX_BYTES, + loadOwnedDocument, refundDocumentAiSlot, + type LoadedDocument, } from "@/lib/documents/ai-route-support"; +import { + executeDocumentExtract, + type DocumentExtractInput, +} from "@/lib/documents/ai-runs/extract-run"; +import { + acceptedRunResponse, + outcomeResponse, + prefersRespondAsync, +} from "@/lib/documents/ai-runs/http"; +import { findLiveDocumentRun, startAiRun } from "@/lib/documents/ai-runs/start"; +import type { AiRunBudget } from "@/lib/documents/ai-runs/types"; import { loadDocumentChatText } from "@/lib/documents/content-index"; -import { detectOcrMimeType } from "@/lib/labs/ocr-upload"; import { annotate } from "@/lib/logging/context"; import { requireModuleEnabled } from "@/lib/modules/gate"; import { inboundStoredExtractSchema, inboundTextExtractSchema, } from "@/lib/validations/inbound-documents"; -import { dateOnlyAtNoonUtc } from "@/lib/tz/date-only"; export const dynamic = "force-dynamic"; type RouteParams = { params: Promise<{ id: string }> }; -/** A loaded, owner-scoped, still-extractable document row. */ -type LoadedDocument = { - id: string; - kind: string; - contentEncrypted: Uint8Array; - contentCodec: string; - mimeType: string; - status: string; -}; - -/** - * Replace the document's staged facts and flip it to EXTRACTED in one - * transaction. Re-extraction is allowed only when NO fact is APPROVED yet (the - * POST handler refuses it otherwise), so clearing every prior staged fact here - * can never drop an approved row or sever a committed-record provenance link — - * the only facts present are PENDING / REJECTED leftovers from an earlier run. - */ -async function stageExtraction( - documentId: string, - userId: string, - result: InboundExtractionResult, -) { - return prisma.$transaction(async (tx) => { - await tx.extractedFact.deleteMany({ where: { documentId, userId } }); - await tx.inboundDocument.update({ - where: { id: documentId }, - data: { - status: "EXTRACTED", - providerType: result.providerType, - reportDate: result.reportDate - ? dateOnlyAtNoonUtc(result.reportDate) - : null, - facts: { - create: result.facts.map((f) => ({ - userId, - factType: f.factType, - status: "PENDING" as const, - confidence: f.confidence, - needsReview: f.needsReview, - dataEncrypted: encryptFactData(f.data), - provenanceEncrypted: encryptFactProvenance(f.provenance), - })), - }, - }, - }); - return tx.inboundDocument.findUniqueOrThrow({ - where: { id: documentId }, - include: { facts: { orderBy: { createdAt: "asc" } } }, - }); - }); -} +type ExtractInputMode = "text" | "stored" | "vision"; export const POST = apiHandler( async (request: NextRequest, { params }: RouteParams) => { @@ -146,17 +105,7 @@ export const POST = apiHandler( await requireAiCapability("documentAi", { pickDecides: true }); const { id } = await params; - const document = await prisma.inboundDocument.findFirst({ - where: { id, userId: user.id, deletedAt: null }, - select: { - id: true, - kind: true, - contentEncrypted: true, - contentCodec: true, - mimeType: true, - status: true, - }, - }); + const document = await loadOwnedDocument(user.id, id); if (!document) { return apiError("Document not found", 404, { errorCode: "documents.inbound.notFound", @@ -175,7 +124,8 @@ export const POST = apiHandler( // committed-record provenance link — and re-stage the same facts as PENDING, // letting the user approve them a second time and duplicate the committed // lab / condition / medication. Block it: the user must finish reviewing or - // discard the document first. + // discard the document first. Staging asks again under a row lock, because + // a review can finish while the read runs. const approvedCount = await prisma.extractedFact.count({ where: { documentId: document.id, userId: user.id, status: "APPROVED" }, }); @@ -191,6 +141,7 @@ export const POST = apiHandler( ); } + const background = prefersRespondAsync(request); const contentType = request.headers.get("content-type") ?? ""; if (contentType.includes("application/json")) { // One body read for both JSON modes, BEFORE any bucket charge — a @@ -200,20 +151,106 @@ export const POST = apiHandler( }); if (jsonError) return jsonError; if (inboundStoredExtractSchema.safeParse(body).success) { - return handleStoredExtract(request, user.id, document); + return handleStoredExtract(request, user.id, document, background); } - return handleTextExtract(request, user.id, document, body); + return handleTextExtract(request, user.id, document, body, background); } - return handleVisionExtract(request, user.id, document); + return handleVisionExtract(request, user.id, document, background); }, ); +/** + * An extraction of this document from the same input, already queued or + * running: attach to it rather than charge a second one. Background only. + */ +async function liveRunResponse( + userId: string, + documentId: string, + input: ExtractInputMode, +): Promise { + const live = await findLiveDocumentRun( + userId, + documentId, + "DOCUMENT_EXTRACT", + (params) => params.extract?.input === input, + ); + return live ? acceptedRunResponse(live) : null; +} + +/** Reserve the read's budget; null (slot refunded) when the day is spent. */ +async function reserveFor( + userId: string, + providerType: ProviderChainType, + tokens: number, +): Promise { + const dateKey = buildDateKey(); + const reservation = await reserveBudget( + userId, + tokens, + dateKey, + resolveDailyCap([{ providerType }]), + resolveCostOwner([{ providerType }]), + "coach", + ); + if (!reservation.allowed) { + await refundDocumentAiSlot(userId); + return null; + } + return { reserved: reservation.reserved, owner: reservation.owner, dateKey }; +} + +function budgetExceeded(): Response { + return apiError("Your AI usage budget for today is reached.", 429, { + errorCode: "documents.inbound.budgetExceeded", + }); +} + +/** Queue the extraction for the background worker. */ +function queueExtract( + userId: string, + document: LoadedDocument, + input: ExtractInputMode, + budget: AiRunBudget, + text: string | null, +): Promise { + return startAiRun({ + userId, + kind: "DOCUMENT_EXTRACT", + documentId: document.id, + params: { + mode: input === "vision" ? "vision" : "text", + budget, + extract: { input }, + }, + input: text === null ? null : Buffer.from(text, "utf8"), + refundSlot: () => refundDocumentAiSlot(userId), + }); +} + +/** The synchronous path: run the shared body and answer with the detail. */ +async function runNow( + request: NextRequest, + userId: string, + document: LoadedDocument, + input: DocumentExtractInput, +): Promise { + const outcome = await executeDocumentExtract({ + userId, + document, + input, + origin: { ipAddress: getClientIp(request), worker: false }, + }); + if (!outcome.ok) return outcomeResponse(outcome); + return apiSuccess(serialiseDocumentDetail(outcome.data, outcome.data.facts)); +} + /** TEXT mode — structure in-browser-OCR'd text against the stored row. */ async function handleTextExtract( request: NextRequest, userId: string, document: LoadedDocument, body: unknown, + background: boolean, ): Promise { const row = await prisma.user.findUnique({ where: { id: userId }, @@ -228,6 +265,11 @@ async function handleTextExtract( const pick = await requireDocumentTextProvider(userId); + if (background) { + const live = await liveRunResponse(userId, document.id, "text"); + if (live) return live; + } + const rl = await checkDocumentAiRateLimit(userId); if (!rl.allowed) return documentAiRateLimited(rl); @@ -244,68 +286,22 @@ async function handleTextExtract( ); } - const dateKey = buildDateKey(); - const reservation = await reserveBudget( + const budget = await reserveFor( userId, + pick.entry.providerType, AI_BUDGETS.ocrExtractText.maxTokens, - dateKey, - resolveDailyCap([{ providerType: pick.entry.providerType }]), - resolveCostOwner([{ providerType: pick.entry.providerType }]), - "coach", ); - if (!reservation.allowed) { - await refundDocumentAiSlot(userId); - return apiError("Your AI usage budget for today is reached.", 429, { - errorCode: "documents.inbound.budgetExceeded", - }); - } + if (!budget) return budgetExceeded(); - try { - const result = await runInboundExtraction({ - provider: pick.entry.instance, - providerType: pick.providerType, - ocrText: parsed.data.text, - }); - await reconcileSpend( - userId, - reservation.reserved, - reservation.reserved, - dateKey, - 0, - { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, - ); - - const updated = await stageExtraction(document.id, userId, result); - - await auditLog("documents.inbound.extract", { - userId, - ipAddress: getClientIp(request), - details: { - documentId: document.id, - facts: updated.facts.length, - mode: "text", - }, - }); - - return apiSuccess(serialiseDocumentDetail(updated, updated.facts)); - } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { - servedBy: null, - reservedOwner: reservation.owner, - }); - if (err instanceof InboundExtractError) { - return apiError("Couldn't read the document. Try a clearer copy.", 422, { - errorCode: "documents.inbound.extractFailed", - }); - } - annotate({ - action: { name: "documents.inbound.extractFailed" }, - meta: { reason: "provider_error", mode: "text" }, - }); - return apiError("Couldn't read the document. Try a clearer copy.", 502, { - errorCode: "documents.inbound.extractFailed", - }); + if (background) { + return queueExtract(userId, document, "text", budget, parsed.data.text); } + return runNow(request, userId, document, { + mode: "text", + text: parsed.data.text, + pick, + budget, + }); } /** @@ -324,6 +320,7 @@ async function handleStoredExtract( request: NextRequest, userId: string, document: LoadedDocument, + background: boolean, ): Promise { const chat = await loadDocumentChatText(userId, document.id); if (!chat || !chat.text.trim()) { @@ -334,71 +331,31 @@ async function handleStoredExtract( const pick = await requireDocumentTextProvider(userId); + if (background) { + const live = await liveRunResponse(userId, document.id, "stored"); + if (live) return live; + } + const rl = await checkDocumentAiRateLimit(userId); if (!rl.allowed) return documentAiRateLimited(rl); - const dateKey = buildDateKey(); - const reservation = await reserveBudget( + const budget = await reserveFor( userId, + pick.entry.providerType, AI_BUDGETS.ocrExtractText.maxTokens, - dateKey, - resolveDailyCap([{ providerType: pick.entry.providerType }]), - resolveCostOwner([{ providerType: pick.entry.providerType }]), - "coach", ); - if (!reservation.allowed) { - await refundDocumentAiSlot(userId); - return apiError("Your AI usage budget for today is reached.", 429, { - errorCode: "documents.inbound.budgetExceeded", - }); - } + if (!budget) return budgetExceeded(); - try { - const result = await runInboundExtraction({ - provider: pick.entry.instance, - providerType: pick.providerType, - ocrText: chat.text, - }); - await reconcileSpend( - userId, - reservation.reserved, - reservation.reserved, - dateKey, - 0, - { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, - ); - - const updated = await stageExtraction(document.id, userId, result); - - await auditLog("documents.inbound.extract", { - userId, - ipAddress: getClientIp(request), - details: { - documentId: document.id, - facts: updated.facts.length, - mode: "stored", - }, - }); - - return apiSuccess(serialiseDocumentDetail(updated, updated.facts)); - } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { - servedBy: null, - reservedOwner: reservation.owner, - }); - if (err instanceof InboundExtractError) { - return apiError("Couldn't read the document. Try a clearer copy.", 422, { - errorCode: "documents.inbound.extractFailed", - }); - } - annotate({ - action: { name: "documents.inbound.extractFailed" }, - meta: { reason: "provider_error", mode: "stored" }, - }); - return apiError("Couldn't read the document. Try a clearer copy.", 502, { - errorCode: "documents.inbound.extractFailed", - }); + if (background) { + // The worker loads the stored text itself; nothing is sealed into the run. + return queueExtract(userId, document, "stored", budget, null); } + return runNow(request, userId, document, { + mode: "stored", + text: chat.text, + pick, + budget, + }); } /** VISION mode — run the stored original through the vision provider. */ @@ -406,116 +363,29 @@ async function handleVisionExtract( request: NextRequest, userId: string, document: LoadedDocument, + background: boolean, ): Promise { const pick = await requireDocumentVisionProvider(userId); + if (background) { + const live = await liveRunResponse(userId, document.id, "vision"); + if (live) return live; + } + 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). Every miss below happens - // before the provider dispatch, so the charged slot goes back. - let buffer: Buffer; - try { - buffer = decryptDocumentContent( - document.contentEncrypted, - document.contentCodec, - ); - } catch { - await refundDocumentAiSlot(userId); - return apiError("Couldn't read the stored document.", 422, { - errorCode: "documents.inbound.extractFailed", - }); - } - - const mime = detectOcrMimeType(buffer); - if (!mime) { - await refundDocumentAiSlot(userId); - return apiError( - "This document can't be scanned. Use local OCR (text mode).", - 422, - { errorCode: "documents.inbound.fileType" }, - ); - } - - if (mime === "application/pdf" && !pick.pdfSupported) { - await refundDocumentAiSlot(userId); - return apiError( - "PDF scanning needs a Claude vision provider; use local OCR instead.", - 422, - { errorCode: "documents.inbound.pdfNeedsAnthropic" }, - ); - } - - const dateKey = buildDateKey(); - const reservation = await reserveBudget( + const budget = await reserveFor( userId, + pick.entry.providerType, AI_BUDGETS.ocrExtract.maxTokens, - dateKey, - resolveDailyCap([{ providerType: pick.entry.providerType }]), - resolveCostOwner([{ providerType: pick.entry.providerType }]), - "coach", ); - if (!reservation.allowed) { - await refundDocumentAiSlot(userId); - return apiError("Your AI usage budget for today is reached.", 429, { - errorCode: "documents.inbound.budgetExceeded", - }); - } + if (!budget) return budgetExceeded(); - const dataBase64 = buffer.toString("base64"); - const images = - mime === "application/pdf" ? [] : [{ mediaType: mime, dataBase64 }]; - const documents = - mime === "application/pdf" - ? [{ mediaType: "application/pdf" as const, dataBase64 }] - : []; - - try { - const result = await runInboundExtraction({ - provider: pick.entry.instance, - providerType: pick.providerType, - images, - documents, - }); - await reconcileSpend( - userId, - reservation.reserved, - reservation.reserved, - dateKey, - 0, - { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, - ); - - const updated = await stageExtraction(document.id, userId, result); - - await auditLog("documents.inbound.extract", { - userId, - ipAddress: getClientIp(request), - details: { - documentId: document.id, - facts: updated.facts.length, - mode: "vision", - }, - }); - - return apiSuccess(serialiseDocumentDetail(updated, updated.facts)); - } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { - servedBy: null, - reservedOwner: reservation.owner, - }); - if (err instanceof InboundExtractError) { - return apiError("Couldn't read the document. Try a clearer copy.", 422, { - errorCode: "documents.inbound.extractFailed", - }); - } - annotate({ - action: { name: "documents.inbound.extractFailed" }, - meta: { reason: "provider_error", mode: "vision" }, - }); - return apiError("Couldn't read the document. Try a clearer copy.", 502, { - errorCode: "documents.inbound.extractFailed", - }); + if (background) { + // The worker decrypts the stored original itself, after it has picked + // the provider again and re-checked the wire for it. + return queueExtract(userId, document, "vision", budget, null); } + return runNow(request, userId, document, { mode: "vision", pick, budget }); } diff --git a/src/components/documents/__tests__/document-facts-review.test.tsx b/src/components/documents/__tests__/document-facts-review.test.tsx index c2c64d314..33c2b7b47 100644 --- a/src/components/documents/__tests__/document-facts-review.test.tsx +++ b/src/components/documents/__tests__/document-facts-review.test.tsx @@ -125,6 +125,8 @@ describe("", () => { expect(html).toContain("Extract lab values"); // Honest about the mechanism: stored text, review-first, no auto-save. expect(html).toContain("nothing is saved automatically"); + // No background read is running, so no run line takes the hint's place. + expect(html).not.toContain("document-facts-extract-run-phase"); }); it("withdraws the extract offer once facts were approved", () => { diff --git a/src/components/documents/document-facts-review.tsx b/src/components/documents/document-facts-review.tsx index ae20eedff..f3ac38128 100644 --- a/src/components/documents/document-facts-review.tsx +++ b/src/components/documents/document-facts-review.tsx @@ -46,6 +46,7 @@ import type { ObservationFactData, } from "@/lib/validations/inbound-documents"; +import { AiRunPhaseNote } from "./document-ai-panels"; import { useDocumentAiErrorText } from "./use-document-assist"; import { useConfirmFacts, @@ -625,12 +626,12 @@ export function DocumentFactsSection({ storedExtract.mutate( { documentId: doc.id }, { - onSuccess: (detail) => { + onSuccess: (result) => { setFailures({}); setIncluded({}); - if ( - detail.facts.some((fact) => fact.status === "PENDING") - ) { + // The staged facts arrive with the document refetch the + // hook starts; the review opens onto them. + if (result.factsStaged > 0) { setOpen(true); } else { toast.info(t("documents.review.extractEmpty")); @@ -654,9 +655,19 @@ export function DocumentFactsSection({ ? t("documents.review.extracting") : t("documents.review.extract")} -

- {t("documents.review.extractHint")} -

+ {storedExtract.runPhase === "idle" ? ( +

+ {t("documents.review.extractHint")} +

+ ) : ( + // While the read runs, the hint gives way to where it stands: + // one muted line, the same one every document read shows. + + )} )} diff --git a/src/components/documents/use-document-facts.ts b/src/components/documents/use-document-facts.ts index 497724761..469155b9b 100644 --- a/src/components/documents/use-document-facts.ts +++ b/src/components/documents/use-document-facts.ts @@ -15,11 +15,15 @@ */ import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { resolveAiRun, useAiRunPhase } from "@/hooks/use-ai-run"; import { apiPatch, apiPost } from "@/lib/api/api-fetch"; +import type { + AiRunAccepted, + DocumentExtractRunResult, +} from "@/lib/documents/ai-runs/types"; import { invalidateKeys, queryKeys } from "@/lib/query-keys"; import type { ExtractedFactDto, - InboundDocumentDetailDto, InboundFactEdit, } from "@/lib/validations/inbound-documents"; @@ -35,17 +39,35 @@ export interface ConfirmFactsResult { * Structure the document's own stored extracted text into staged facts — * the manual recovery when the automatic staging run was skipped or failed. * Stages PENDING facts only; the review + confirm steps stay mandatory. - * A provider call is slow, so the request opts out of the 15 s default. + * + * The extraction runs in the background (v1.40): the request only queues it, + * so no proxy can cut a slow model, and `runPhase` says so while it runs. The + * run reports how many facts it staged; the facts themselves are read from + * the document, which is refetched when the run ends. */ export function useStoredExtract() { const queryClient = useQueryClient(); - return useMutation({ - mutationFn: ({ documentId }) => - apiPost( - `/api/documents/inbound/${documentId}/extract`, - { mode: "stored" }, - { signal: AbortSignal.timeout(120_000) }, + const { phase, onQueued, onProgress, reset } = useAiRunPhase(); + const mutation = useMutation< + DocumentExtractRunResult, + Error, + { documentId: string } + >({ + mutationFn: async ({ documentId }) => + resolveAiRun( + queryClient, + await apiPost( + `/api/documents/inbound/${documentId}/extract`, + { mode: "stored" }, + { + // The request only queues the read: the enqueue, no model. + signal: AbortSignal.timeout(60_000), + headers: { Prefer: "respond-async" }, + }, + ), + { onQueued, onProgress }, ), + onSettled: reset, onSuccess: (_data, { documentId }) => { void invalidateKeys(queryClient, [ queryKeys.inboundDocument(documentId), @@ -53,6 +75,7 @@ export function useStoredExtract() { ]); }, }); + return { ...mutation, runPhase: mutation.isPending ? phase : "idle" }; } /** diff --git a/src/lib/documents/ai-runs/extract-run.ts b/src/lib/documents/ai-runs/extract-run.ts new file mode 100644 index 000000000..226a6c98f --- /dev/null +++ b/src/lib/documents/ai-runs/extract-run.ts @@ -0,0 +1,287 @@ +/** + * Fact extraction from a stored document: the body both the synchronous + * extract route and the background run execute, so the two paths cannot + * drift. + * + * One extraction (a read and its one corrective retry) over text the browser + * read, the document's own stored text, or the stored original, against a + * reservation the caller already took and which is settled here. The facts + * land PENDING for the mandatory review step; the confirm route stays the only + * write into the person's record. + * + * Staging re-checks the document under a row lock. The route refuses a + * confirmed or partly confirmed document before it charges anything, but a + * background read can take minutes, and a review the person finishes in the + * meantime must never have its approved facts replaced by fresh PENDING ones + * (which could then be approved a second time). So the staging transaction + * locks the document row, asks again, and refuses with the route's own code. + */ +import { Buffer } from "node:buffer"; + +import type { ProviderChainType } from "@/lib/ai/provider-chain"; +import { auditLog } from "@/lib/auth/audit"; +import { prisma } from "@/lib/db"; +import { + refundDocumentAiSlot, + type LoadedDocument, +} from "@/lib/documents/ai-route-support"; +import { + InboundExtractError, + runInboundExtraction, + type InboundExtractionResult, + type RunInboundExtractionArgs, +} from "@/lib/documents/extract"; +import { + decryptDocumentContent, + encryptFactData, + encryptFactProvenance, +} from "@/lib/documents/store"; +import { detectOcrMimeType } from "@/lib/labs/ocr-upload"; +import { annotate } from "@/lib/logging/context"; +import { dateOnlyAtNoonUtc } from "@/lib/tz/date-only"; + +import { + runFailure, + unreadableDocument, + type DocumentTextPick, + type DocumentVisionPick, + type RunOrigin, +} from "./run-support"; +import { settleAiRunBudget } from "./store"; +import type { AiRunBudget, AiRunOutcome } from "./types"; + +/** The extraction and its one corrective retry. */ +export const DOCUMENT_EXTRACT_MODEL_CALLS = 2; + +/** + * What the extraction reads: text the browser read (`text`), the document's + * own content index (`stored`, loaded by the caller), or the stored original + * (`vision`). + */ +export type DocumentExtractInput = + | { + mode: "text" | "stored"; + text: string; + pick: DocumentTextPick; + budget: AiRunBudget; + } + | { mode: "vision"; pick: DocumentVisionPick; budget: AiRunBudget }; + +/** The staged document with its facts, for the route's detail body. */ +export type StagedDocument = Awaited>; + +class StagingRefused extends Error { + constructor(readonly outcome: AiRunOutcome) { + super("staging refused"); + } +} + +/** + * Replace the document's staged facts and flip it to EXTRACTED in one + * transaction, after locking the row and asking again whether it may be. + * + * Re-extraction is allowed only when NO fact is APPROVED and the document is + * not CONFIRMED, so clearing every prior staged fact here can never drop an + * approved row or sever a committed-record provenance link: the only facts + * present are PENDING / REJECTED leftovers from an earlier run. + */ +async function stageExtraction( + documentId: string, + userId: string, + result: InboundExtractionResult, +) { + return prisma.$transaction(async (tx) => { + const locked = await tx.$queryRaw<{ status: string }[]>` + SELECT status::text AS status + FROM inbound_documents + WHERE id = ${documentId} AND user_id = ${userId} AND deleted_at IS NULL + FOR UPDATE`; + const status = locked[0]?.status; + if (!status) { + throw new StagingRefused( + runFailure({ + status: 404, + message: "Document not found", + errorCode: "documents.inbound.notFound", + }), + ); + } + if (status === "CONFIRMED") { + throw new StagingRefused( + runFailure({ + status: 422, + message: "This document has already been confirmed.", + errorCode: "documents.inbound.alreadyConfirmed", + }), + ); + } + const approved = await tx.extractedFact.count({ + where: { documentId, userId, status: "APPROVED" }, + }); + if (approved > 0) { + throw new StagingRefused( + runFailure({ + status: 409, + message: + "Some facts from this document are already confirmed. Finish reviewing or discard it before extracting again.", + errorCode: "documents.inbound.alreadyPartlyConfirmed", + }), + ); + } + await tx.extractedFact.deleteMany({ where: { documentId, userId } }); + await tx.inboundDocument.update({ + where: { id: documentId }, + data: { + status: "EXTRACTED", + providerType: result.providerType, + reportDate: result.reportDate + ? dateOnlyAtNoonUtc(result.reportDate) + : null, + facts: { + create: result.facts.map((f) => ({ + userId, + factType: f.factType, + status: "PENDING" as const, + confidence: f.confidence, + needsReview: f.needsReview, + dataEncrypted: encryptFactData(f.data), + provenanceEncrypted: encryptFactProvenance(f.provenance), + })), + }, + }, + }); + return tx.inboundDocument.findUniqueOrThrow({ + where: { id: documentId }, + include: { facts: { orderBy: { createdAt: "asc" } } }, + }); + }); +} + +/** + * The provider input for the stored original: decrypted, its MIME re-derived + * from the bytes (never the stored label), a PDF only for a provider that + * reads one. Every miss happens before the dispatch. + */ +function visionArgs( + document: LoadedDocument, + pick: DocumentVisionPick, +): + Pick | AiRunOutcome { + let buffer: Buffer; + try { + buffer = decryptDocumentContent( + document.contentEncrypted, + document.contentCodec, + ); + } catch { + return runFailure({ + status: 422, + message: "Couldn't read the stored document.", + errorCode: "documents.inbound.extractFailed", + }); + } + const mime = detectOcrMimeType(buffer); + if (!mime) { + return runFailure({ + status: 422, + message: "This document can't be scanned. Use local OCR (text mode).", + errorCode: "documents.inbound.fileType", + }); + } + if (mime === "application/pdf" && !pick.pdfSupported) { + return runFailure({ + status: 422, + message: + "PDF scanning needs a Claude vision provider; use local OCR instead.", + errorCode: "documents.inbound.pdfNeedsAnthropic", + }); + } + const dataBase64 = buffer.toString("base64"); + return mime === "application/pdf" + ? { images: [], documents: [{ mediaType: "application/pdf", dataBase64 }] } + : { images: [{ mediaType: mime, dataBase64 }], documents: [] }; +} + +/** Extract and stage facts from one owned document. The caller authorised and charged it. */ +export async function executeDocumentExtract(args: { + userId: string; + document: LoadedDocument; + input: DocumentExtractInput; + origin: RunOrigin; +}): Promise> { + const { userId, document, input, origin } = args; + const { pick, budget } = input; + const servedBy = pick.entry.providerType as ProviderChainType; + + let source: Pick< + RunInboundExtractionArgs, + "images" | "documents" | "ocrText" + >; + if (input.mode === "vision") { + const prepared = visionArgs(document, input.pick); + if ("ok" in prepared) { + // Preparation failed before any provider dispatch: the reservation and + // the slot go back. + await settleAiRunBudget(userId, budget, 0, null); + await refundDocumentAiSlot(userId); + return prepared; + } + source = prepared; + } else { + source = { ocrText: input.text }; + } + + let result: InboundExtractionResult; + try { + result = await runInboundExtraction({ + provider: pick.entry.instance, + providerType: pick.providerType, + ...source, + }); + } catch (err) { + await settleAiRunBudget(userId, budget, 0, null); + if (!(err instanceof InboundExtractError)) { + annotate({ + action: { name: "documents.inbound.extractFailed" }, + meta: { reason: "provider_error", mode: input.mode }, + }); + } + return unreadableDocument(err instanceof InboundExtractError); + } + await settleAiRunBudget(userId, budget, budget.reserved, servedBy); + + let staged: StagedDocument; + try { + staged = await stageExtraction(document.id, userId, result); + } catch (err) { + if (!(err instanceof StagingRefused)) { + // The reservation is already settled for the call that was made; a + // write that failed is answered as the route always answered it. + annotate({ + action: { name: "documents.inbound.extractFailed" }, + meta: { reason: "staging_error", mode: input.mode }, + }); + return unreadableDocument(false); + } + annotate({ + action: { name: "documents.inbound.extractRefusedAtStaging" }, + meta: { + documentId: document.id, + errorCode: err.outcome.ok ? null : err.outcome.errorCode, + }, + }); + return err.outcome; + } + + await auditLog("documents.inbound.extract", { + userId, + ...(origin.worker ? { actorUserId: userId } : {}), + ipAddress: origin.ipAddress, + details: { + documentId: document.id, + facts: staged.facts.length, + mode: input.mode, + }, + }); + return { ok: true, data: staged }; +} diff --git a/src/lib/documents/ai-runs/types.ts b/src/lib/documents/ai-runs/types.ts index 74986eb77..2c63c7d92 100644 --- a/src/lib/documents/ai-runs/types.ts +++ b/src/lib/documents/ai-runs/types.ts @@ -10,7 +10,10 @@ */ import type { AiCapabilityKey } from "@/lib/ai/capabilities/types"; import type { Locale } from "@/lib/i18n/config"; -import type { InboundDocumentKindValue } from "@/lib/validations/inbound-documents"; +import type { + InboundDocumentKindValue, + InboundDocumentStatusValue, +} from "@/lib/validations/inbound-documents"; import type { OcrExtractResponseDto } from "@/lib/validations/labs-ocr"; export const DOCUMENT_AI_RUN_KINDS = [ @@ -18,6 +21,7 @@ export const DOCUMENT_AI_RUN_KINDS = [ "LABS_OCR_EXTRACT", "DOCUMENT_SUMMARY", "DOCUMENT_SUGGEST", + "DOCUMENT_EXTRACT", ] as const; export type DocumentAiRunKindValue = (typeof DOCUMENT_AI_RUN_KINDS)[number]; @@ -33,6 +37,7 @@ export const DOCUMENT_AI_RUN_CAPABILITY: Record< LABS_OCR_EXTRACT: "labsOcr", DOCUMENT_SUMMARY: "documentAi", DOCUMENT_SUGGEST: "documentAi", + DOCUMENT_EXTRACT: "documentAi", }; export const DOCUMENT_AI_RUN_STATUSES = [ @@ -90,6 +95,8 @@ export interface AiRunParams { replace: boolean; locale: Locale; }; + /** An extract run: which text it structures. `stored` reads the content index. */ + extract?: { input: "text" | "stored" | "vision" }; } /** "Read with AI": the same body the synchronous index route answers with. */ @@ -114,11 +121,23 @@ export interface DocumentSuggestRunResult { }; } +/** + * An extract run. Smaller than the synchronous route's body (the whole + * document detail): the facts are staged on the document, and the client + * reads them from there. + */ +export interface DocumentExtractRunResult { + documentId: string; + factsStaged: number; + status: InboundDocumentStatusValue; +} + export type AiRunResult = | DocumentIndexRunResult | OcrExtractResponseDto | DocumentSummaryRunResult - | DocumentSuggestRunResult; + | DocumentSuggestRunResult + | DocumentExtractRunResult; /** * What one run body produced: the success body, or the failure the diff --git a/src/lib/jobs/__tests__/document-ai-run.test.ts b/src/lib/jobs/__tests__/document-ai-run.test.ts index 752157713..d66bdd9c2 100644 --- a/src/lib/jobs/__tests__/document-ai-run.test.ts +++ b/src/lib/jobs/__tests__/document-ai-run.test.ts @@ -79,6 +79,16 @@ vi.mock("@/lib/documents/ai-runs/suggest-run", () => ({ executeDocumentSuggest: (...a: unknown[]) => executeDocumentSuggest(...a), })); +const executeDocumentExtract = vi.fn(); +vi.mock("@/lib/documents/ai-runs/extract-run", () => ({ + DOCUMENT_EXTRACT_MODEL_CALLS: 2, + executeDocumentExtract: (...a: unknown[]) => executeDocumentExtract(...a), +})); +const loadDocumentChatText = vi.fn(); +vi.mock("@/lib/documents/content-index", () => ({ + loadDocumentChatText: (...a: unknown[]) => loadDocumentChatText(...a), +})); + const loadOwnedDocument = vi.fn(); vi.mock("@/lib/documents/ai-route-support", () => ({ loadOwnedDocument: (...a: unknown[]) => loadOwnedDocument(...a), @@ -554,3 +564,109 @@ describe("summary and suggestion runs", () => { expect(markQueuedSummaryUnavailable).not.toHaveBeenCalled(); }); }); + +describe("extract runs", () => { + const extractRun = (input: "stored" | "vision" | "text") => + claimed({ + kind: "DOCUMENT_EXTRACT", + params: { + mode: input === "vision" ? "vision" : "text", + budget, + extract: { input }, + }, + input: input === "text" ? Buffer.from("Hb 14 g/dL") : null, + }); + + beforeEach(() => { + executeDocumentExtract.mockResolvedValue({ + ok: true, + data: { + id: "d1", + status: "EXTRACTED", + facts: [{ id: "f1" }, { id: "f2" }], + }, + }); + loadDocumentChatText.mockResolvedValue({ text: "Hb 14 g/dL" }); + requireDocumentTextProvider.mockResolvedValue({ + entry: { + providerType: "local", + instance: { responseTimeoutSeconds: 100 }, + }, + providerType: "local", + }); + }); + + it("stores only the count and the status, with two calls' deadline", async () => { + claimAiRun.mockResolvedValue(extractRun("vision")); + requireDocumentVisionProvider.mockResolvedValue(visionPick(300)); + expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe( + "succeeded", + ); + expect(updateMany).toHaveBeenCalledWith({ + where: { id: "run1", status: "RUNNING" }, + data: { expiresAt: new Date(NOW + 2 * 300_000 + 60_000) }, + }); + expect(completeAiRun).toHaveBeenCalledWith( + "run1", + { documentId: "d1", factsStaged: 2, status: "EXTRACTED" }, + new Date(NOW), + ); + }); + + it("reads the stored text for a stored extract and re-checks the text provider's wire", async () => { + claimAiRun.mockResolvedValue(extractRun("stored")); + expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe( + "succeeded", + ); + expect(loadDocumentChatText).toHaveBeenCalledWith("u1", "d1"); + expect(aiEgressRefusal).toHaveBeenCalledWith("documentAi", "u1", ["local"]); + expect(executeDocumentExtract).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ mode: "stored", text: "Hb 14 g/dL" }), + }), + ); + }); + + it("fails a stored extract whose text is gone, before any provider pick", async () => { + claimAiRun.mockResolvedValue(extractRun("stored")); + loadDocumentChatText.mockResolvedValue(null); + expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe("failed"); + expect(requireDocumentTextProvider).not.toHaveBeenCalled(); + expect(executeDocumentExtract).not.toHaveBeenCalled(); + expect(settleAiRunBudget).toHaveBeenCalledOnce(); + expect(failAiRun).toHaveBeenCalledWith( + "run1", + expect.objectContaining({ + status: 422, + errorCode: "documents.inbound.notIndexed", + }), + "RUNNING", + new Date(NOW), + ); + }); + + it("stores the staging refusal when a review finished while the read ran", async () => { + claimAiRun.mockResolvedValue(extractRun("text")); + executeDocumentExtract.mockResolvedValue({ + ok: false, + status: 409, + message: "Some facts from this document are already confirmed.", + errorCode: "documents.inbound.alreadyPartlyConfirmed", + }); + expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe("failed"); + expect(executeDocumentExtract).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ mode: "text", text: "Hb 14 g/dL" }), + }), + ); + expect(failAiRun).toHaveBeenCalledWith( + "run1", + expect.objectContaining({ + status: 409, + errorCode: "documents.inbound.alreadyPartlyConfirmed", + }), + "RUNNING", + new Date(NOW), + ); + }); +}); diff --git a/src/lib/jobs/document-ai-run.ts b/src/lib/jobs/document-ai-run.ts index c8244de7e..40388d7e0 100644 --- a/src/lib/jobs/document-ai-run.ts +++ b/src/lib/jobs/document-ai-run.ts @@ -39,7 +39,12 @@ import { import { callTimeoutMs } from "@/lib/ai/effective-timeout"; import { prisma } from "@/lib/db"; import { loadOwnedDocument } from "@/lib/documents/ai-route-support"; +import { loadDocumentChatText } from "@/lib/documents/content-index"; import { refusalAsFailure } from "@/lib/documents/ai-runs/http"; +import { + DOCUMENT_EXTRACT_MODEL_CALLS, + executeDocumentExtract, +} from "@/lib/documents/ai-runs/extract-run"; import { DOCUMENT_INDEX_MODEL_CALLS, executeDocumentIndex, @@ -389,6 +394,62 @@ async function runDocumentSuggestRead( }); } +/** The text an extract run structures: the sealed browser text, or the content index. */ +async function extractText(run: ClaimedAiRun): Promise { + if (run.params.extract?.input !== "stored") return runText(run); + const chat = run.documentId + ? await loadDocumentChatText(run.userId, run.documentId) + : null; + if (!chat || !chat.text.trim()) { + throw new RunRefused({ + status: 422, + message: "Read the document first, then extract.", + errorCode: "documents.inbound.notIndexed", + }); + } + return chat.text; +} + +async function runDocumentExtractRead( + run: ClaimedAiRun, + deadlineAt: number | undefined, + now: () => number, +): Promise> { + const document = await loadRunDocument(run); + const text = run.params.mode === "text" ? await extractText(run) : null; + const read = await pickForDocumentRead( + run, + DOCUMENT_EXTRACT_MODEL_CALLS, + deadlineAt, + now, + ); + const budget = runBudget(run); + const outcome = await executeDocumentExtract({ + userId: run.userId, + document, + input: + read.mode === "text" + ? { + mode: run.params.extract?.input === "stored" ? "stored" : "text", + text: text ?? "", + pick: read.pick, + budget, + } + : { mode: "vision", pick: read.pick, budget }, + origin: { ipAddress: null, worker: true }, + }); + if (!outcome.ok) return outcome; + // The facts are staged on the document; the run carries only the count. + return { + ok: true, + data: { + documentId: outcome.data.id, + factsStaged: outcome.data.facts.length, + status: outcome.data.status, + }, + }; +} + type RunDispatcher = ( run: ClaimedAiRun, deadlineAt: number | undefined, @@ -405,6 +466,7 @@ const RUN_DISPATCH: Record = { LABS_OCR_EXTRACT: runLabsOcr, DOCUMENT_SUMMARY: runDocumentSummaryRead, DOCUMENT_SUGGEST: runDocumentSuggestRead, + DOCUMENT_EXTRACT: runDocumentExtractRead, }; /** diff --git a/src/lib/openapi/routes/ai-runs.ts b/src/lib/openapi/routes/ai-runs.ts index e65d7c1e0..d2eeefba6 100644 --- a/src/lib/openapi/routes/ai-runs.ts +++ b/src/lib/openapi/routes/ai-runs.ts @@ -7,16 +7,18 @@ * Part of the OpenAPI route table; aggregated in `./index.ts`. * * iOS coordination (v1.40): `POST /api/documents/inbound/{id}/index`, - * `…/summary` and `…/suggest` keep their 200 bodies when the request carries - * no `Prefer` header. A client that - * sends `Prefer: respond-async` gets 202 with a run id and polls the run here; - * the result is the same body the route would have answered with. + * `…/summary`, `…/suggest` and `…/extract` keep their 200 bodies when the + * request carries no `Prefer` header. A client that sends + * `Prefer: respond-async` gets 202 with a run id and polls the run here; the + * result is the body the route would have answered with, except for extract, + * whose run carries `DocumentExtractRunResult`. * `POST /api/labs/ocr/extract` (web only) always answers 202. */ import type { ZodOpenApiObject } from "zod-openapi"; import { z } from "zod/v4"; import { + documentExtractRunResult, documentIndexResponse, documentSuggestResponse, documentSummaryResponse, @@ -32,6 +34,7 @@ const aiRun = z "LABS_OCR_EXTRACT", "DOCUMENT_SUMMARY", "DOCUMENT_SUGGEST", + "DOCUMENT_EXTRACT", ]), documentId: z .string() @@ -44,10 +47,11 @@ const aiRun = z ocrExtractResponse, documentSummaryResponse, documentSuggestResponse, + documentExtractRunResult, ]) .nullable() .describe( - "Present once `status` is `SUCCEEDED`: the body the queuing route answers with synchronously. `DOCUMENT_INDEX` carries `DocumentIndexResponse`, `LABS_OCR_EXTRACT` carries `OcrExtractResponse`, `DOCUMENT_SUMMARY` carries `DocumentSummaryResponse`, `DOCUMENT_SUGGEST` carries `DocumentSuggestResponse`. Decode by `kind`; treat an unknown `kind` as a run this client did not start.", + "Present once `status` is `SUCCEEDED`: the body the queuing route answers with synchronously. `DOCUMENT_INDEX` carries `DocumentIndexResponse`, `LABS_OCR_EXTRACT` carries `OcrExtractResponse`, `DOCUMENT_SUMMARY` carries `DocumentSummaryResponse`, `DOCUMENT_SUGGEST` carries `DocumentSuggestResponse`, `DOCUMENT_EXTRACT` carries `DocumentExtractRunResult` (not the extract route's detail body). Decode by `kind`; treat an unknown `kind` as a run this client did not start.", ), error: z .object({ diff --git a/src/lib/openapi/routes/documents.ts b/src/lib/openapi/routes/documents.ts index bc63f1941..4d413fd4f 100644 --- a/src/lib/openapi/routes/documents.ts +++ b/src/lib/openapi/routes/documents.ts @@ -287,6 +287,23 @@ const documentUploadEnvelope = dataEnvelope( "DocumentUploadEnvelope", ); +/** + * A `DOCUMENT_EXTRACT` run's result. Smaller than the synchronous extract + * body: the facts are staged on the document and read from it. + */ +export const documentExtractRunResult = z + .object({ + documentId: z.string(), + factsStaged: z + .number() + .int() + .describe( + "How many facts were staged PENDING for review. Read them from `GET /api/documents/inbound/{id}`; the confirm route stays the only write into the record.", + ), + status: statusEnum, + }) + .meta({ id: "DocumentExtractRunResult" }); + /** The suggest route's body; also a `DOCUMENT_SUGGEST` run's result. */ export const documentSuggestResponse = z .object({ @@ -919,7 +936,7 @@ export const inboundDocumentPaths: NonNullable = { tags: ["Documents"], summary: "Extract facts from a stored document", description: - 'Optional AI enhancement on an already-stored document. Runs the dedicated OCR/vision provider over the stored original and stages STRUCTURED FACTS for review. Answers under the `documentAi` capability (403 with the capability envelope when it is closed), then rate / budget gated. With no provider configured this returns 422 (`documents.inbound.providerUnsupported`, with `meta.capability` and `meta.reason = "no_provider"`) — the stored document is untouched, only the enhancement fails. Three modes: VISION (empty body) decrypts and scans the stored original (PDF needs an Anthropic vision provider); TEXT (`application/json`, opt-in local OCR) `{ mode: "text", text }` structures browser-OCR\'d text; STORED (`application/json`) `{ mode: "stored" }` structures the document\'s own stored extracted text (its content index) — the manual recovery for a skipped/failed automatic staging run, 422 `documents.inbound.notIndexed` when no stored text exists. Extraction reproduces what the document states — it never interprets. Nothing reaches the structured stores here; the confirm route is the only write path.', + 'Optional AI enhancement on an already-stored document. Runs the dedicated OCR/vision provider over the stored original and stages STRUCTURED FACTS for review. Answers under the `documentAi` capability (403 with the capability envelope when it is closed), then rate / budget gated. With no provider configured this returns 422 (`documents.inbound.providerUnsupported`, with `meta.capability` and `meta.reason = "no_provider"`) — the stored document is untouched, only the enhancement fails. Three modes: VISION (empty body) decrypts and scans the stored original (PDF needs an Anthropic vision provider); TEXT (`application/json`, opt-in local OCR) `{ mode: "text", text }` structures browser-OCR\'d text; STORED (`application/json`) `{ mode: "stored" }` structures the document\'s own stored extracted text (its content index) — the manual recovery for a skipped/failed automatic staging run, 422 `documents.inbound.notIndexed` when no stored text exists. Extraction reproduces what the document states — it never interprets. Nothing reaches the structured stores here; the confirm route is the only write path. With `Prefer: respond-async` (v1.40) the extraction runs in the background: after the same quick refusals the route answers 202 with a run id, a second request while an extraction of this document from the same input is queued or running answers with that run, and `GET /api/ai-runs/{id}` serves `DocumentExtractRunResult` (`{ documentId, factsStaged, status }`, not the document detail) as the run\'s `result`. Staging re-checks the document under a row lock, so a review finished while the read ran fails the run with `documents.inbound.alreadyPartlyConfirmed` (409) or `documents.inbound.alreadyConfirmed` (422) and leaves the review untouched; the synchronous form answers the same in that race.', parameters: [ { name: "id", @@ -927,6 +944,7 @@ export const inboundDocumentPaths: NonNullable = { required: true, schema: { type: "string" }, }, + preferRespondAsyncParameter, ], requestBody: { required: false, @@ -964,6 +982,7 @@ export const inboundDocumentPaths: NonNullable = { "Re-extraction refused: at least one fact on this document is already APPROVED. `meta.errorCode` = `documents.inbound.alreadyPartlyConfirmed`. Re-extracting would sever committed-record provenance and duplicate committed records, so the user must finish reviewing or discard the document first.", content: { "application/json": { schema: errorEnvelope } }, }, + ...aiRunAccepted, ...aiExtractionRefusals("documentAi"), ...stdResponses, }, diff --git a/tests/integration/document-ai-runs.test.ts b/tests/integration/document-ai-runs.test.ts index ee54eb1e3..2b82d6575 100644 --- a/tests/integration/document-ai-runs.test.ts +++ b/tests/integration/document-ai-runs.test.ts @@ -16,6 +16,9 @@ * anything reaches a model. * 6. The reaper fails a run no worker took (and hands its reservation back) * and deletes a finished run an hour after it ended. + * 8. A stored-text extraction stages its facts from the worker, and a + * review finished while the read ran is never overwritten: staging + * re-checks the document under a row lock and fails the run instead. * 7. Summary and suggestion runs (0370) round-trip the routes' own bodies; a * summary that is to be stored marks the document PENDING while queued, * READY when it lands and UNAVAILABLE when the run fails; a transient one @@ -93,6 +96,29 @@ vi.mock("@/lib/labs/ocr-extract", async (importOriginal) => ({ rows: [], })), })); +vi.mock("@/lib/documents/extract", async (importOriginal) => ({ + ...(await importOriginal()), + runInboundExtraction: vi.fn(async () => ({ + reportDate: "2027-01-10", + kind: "OTHER", + providerType: "anthropic", + facts: [ + { + factType: "OBSERVATION", + confidence: 0.95, + needsReview: false, + data: { label: "Hemoglobin", value: 13.9, unit: "g/dL" }, + provenance: { + sourceText: "Hemoglobin 13.9 g/dL", + anchored: true, + sourceOffset: 0, + page: null, + confidence: 0.95, + }, + }, + ], + })), +})); vi.mock("@/lib/documents/auto-stage-labs", () => ({ maybeAutoStageLabFacts: vi.fn(async () => ({ staged: false, @@ -105,6 +131,7 @@ import { runDocumentSummary, transcribeDocument, } from "@/lib/documents/describe"; +import { runInboundExtraction } from "@/lib/documents/extract"; import { runOcrExtraction } from "@/lib/labs/ocr-extract"; const PNG = Buffer.from( @@ -267,6 +294,7 @@ beforeEach(async () => { vi.mocked(runOcrExtraction).mockClear(); vi.mocked(runDocumentSummary).mockClear(); vi.mocked(runDocumentAssist).mockClear(); + vi.mocked(runInboundExtraction).mockClear(); }); describe("Read with AI in the background", () => { @@ -680,3 +708,143 @@ describe("filing suggestions in the background", () => { expect(foreignPoll.body.meta?.errorCode).toBe("aiRuns.notFound"); }); }); + +async function postStoredExtract(documentId: string, prefer: boolean) { + const { POST } = + await import("@/app/api/documents/inbound/[id]/extract/route"); + const { NextRequest } = await import("next/server"); + const headers: Record = { + "content-type": "application/json", + }; + if (prefer) headers.prefer = "respond-async"; + const res = await POST( + new NextRequest( + `http://localhost/api/documents/inbound/${documentId}/extract`, + { method: "POST", headers, body: JSON.stringify({ mode: "stored" }) }, + ) as never, + { params: Promise.resolve({ id: documentId }) } as never, + ); + return { status: res.status, body: (await res.json()) as Envelope }; +} + +async function indexedDocument(userId: string) { + const document = await storeDocument(userId); + const { upsertContentIndex } = await import("@/lib/documents/content-index"); + await upsertContentIndex({ + userId, + documentId: document.id, + text: "Hemoglobin 13.9 g/dL reference range 12-16", + source: "text-ocr", + providerType: null, + }); + return document; +} + +describe("fact extraction in the background", () => { + it("stages the facts in the worker and serves the count", async () => { + const user = await makeUser(); + await grantConsent(user.id); + const document = await indexedDocument(user.id); + await signIn(user.id); + + const queued = await postStoredExtract(document.id, true); + expect(queued.status).toBe(202); + const { runId } = queued.body.data as { runId: string }; + const stored = await getPrismaClient().documentAiRun.findUniqueOrThrow({ + where: { id: runId }, + }); + // The worker reads the content index itself; nothing is sealed. + expect(stored.inputEncrypted).toBeNull(); + expect(runInboundExtraction).not.toHaveBeenCalled(); + + expect(await work(runId)).toBe("succeeded"); + expect(runInboundExtraction).toHaveBeenCalledWith( + expect.objectContaining({ + ocrText: expect.stringContaining("Hemoglobin 13.9"), + }), + ); + const done = await poll(runId); + expect(done.body.data).toMatchObject({ + status: "SUCCEEDED", + kind: "DOCUMENT_EXTRACT", + result: { + documentId: document.id, + factsStaged: 1, + status: "EXTRACTED", + }, + }); + const facts = await getPrismaClient().extractedFact.findMany({ + where: { documentId: document.id }, + }); + expect(facts.map((f) => f.status)).toEqual(["PENDING"]); + }); + + it("never replaces a review that was finished while the read ran", async () => { + const user = await makeUser(); + await grantConsent(user.id); + const document = await indexedDocument(user.id); + await signIn(user.id); + + const queued = await postStoredExtract(document.id, true); + const { runId } = queued.body.data as { runId: string }; + + // The person approves a fact from an earlier staging before the worker + // gets to the run. + const { encryptFactData, encryptFactProvenance } = + await import("@/lib/documents/store"); + const approved = await getPrismaClient().extractedFact.create({ + data: { + documentId: document.id, + userId: user.id, + factType: "OBSERVATION", + status: "APPROVED", + confidence: 0.9, + needsReview: false, + dataEncrypted: encryptFactData({ + label: "Glucose", + value: 95, + unit: "mg/dL", + } as never), + provenanceEncrypted: encryptFactProvenance({ + sourceText: "Glucose 95", + anchored: true, + sourceOffset: 0, + page: null, + confidence: 0.9, + } as never), + }, + }); + + expect(await work(runId)).toBe("failed"); + const done = await poll(runId); + expect(done.body.data).toMatchObject({ + status: "FAILED", + error: { + code: "documents.inbound.alreadyPartlyConfirmed", + status: 409, + }, + }); + const facts = await getPrismaClient().extractedFact.findMany({ + where: { documentId: document.id }, + }); + expect(facts.map((f) => [f.id, f.status])).toEqual([ + [approved.id, "APPROVED"], + ]); + }); + + it("answers with the document detail, exactly as before, without the header", async () => { + const user = await makeUser(); + await grantConsent(user.id); + const document = await indexedDocument(user.id); + await signIn(user.id); + + const res = await postStoredExtract(document.id, false); + expect(res.status).toBe(200); + expect(res.body.data).toMatchObject({ + id: document.id, + status: "EXTRACTED", + }); + expect((res.body.data as { facts: unknown[] }).facts).toHaveLength(1); + expect(await getPrismaClient().documentAiRun.count()).toBe(0); + }); +});