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); + }); +});