mirror of
https://github.com/MBombeck/HealthLog.git
synced 2026-11-08 12:32:43 +00:00
feat(documents): read documents and lab scans in the background
"Read with AI" and the lab scan used to run their model calls inside the
request. At the default response time a read of three calls outlasted the
100 s a proxy in front of the server commonly allows, and the person saw a
cut connection for a read the server was still doing.
Both now run as a background run (`DocumentAiRun`, migration 0366):
- `POST /api/documents/inbound/{id}/index` with `Prefer: respond-async`
answers 202 with a run id once every quick refusal has passed. Without
the header it answers exactly as before, which the iPhone app relies on.
A second request while a read of the same document is queued or running
attaches to that run.
- `POST /api/labs/ocr/extract` always answers 202. Only the web client
calls it.
- `GET /api/ai-runs/{id}` serves the run: the same body the synchronous
route answers with, or the failure code and status it would have given.
Scoped to the caller; another account's run is a 404.
The worker (`document-ai-run`, a long queue with a one-hour expiry, one
job per run under `lockedPass`) claims the run, asks the capability again
for the record and re-checks the wire for the provider it picks, sets the
run's deadline from the person's AI response time and refuses a read that
cannot fit the job's expiry. Sync route and worker run one body per kind,
so the two paths cannot drift. Budget is reserved at enqueue and settled
once by whoever ends the run.
Input and result are sealed with labelled AES-256-GCM, registered for key
rotation; the input is dropped when the run ends. A five-minute reaper
fails a run no worker took within fifteen minutes (and hands its
reservation back), fails a run past its deadline, and deletes a finished
run after an hour. Runs are outside backups and inside the data wipe.
The web client polls through the query-key factory and says, calmly, that
the read continues in the background, and when no worker has picked it up.
This commit is contained in:
1 parent
556bd6d6cd
commit
0acf3da575
49 files changed
+4533
-870
No files matched your search
+479
-282
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,54 @@
|
||||
-- Document AI as background runs. A read that used to run inside the request
|
||||
-- now runs in the worker; the request answers 202 with this row's id and the
|
||||
-- client polls it. The row is a ticket: the input is cleared when the run
|
||||
-- finishes and the row is deleted an hour later. Idempotent, so a re-run of a
|
||||
-- half-applied migration converges.
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE "document_ai_run_kind" AS ENUM ('DOCUMENT_INDEX', 'LABS_OCR_EXTRACT');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE "document_ai_run_status" AS ENUM ('QUEUED', 'RUNNING', 'SUCCEEDED', 'FAILED');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "document_ai_runs" (
|
||||
"id" TEXT NOT NULL,
|
||||
"user_id" TEXT NOT NULL,
|
||||
"kind" "document_ai_run_kind" NOT NULL,
|
||||
"document_id" TEXT,
|
||||
"status" "document_ai_run_status" NOT NULL DEFAULT 'QUEUED',
|
||||
"params_json" JSONB NOT NULL,
|
||||
"input_encrypted" BYTEA,
|
||||
"result_encrypted" BYTEA,
|
||||
"error_code" TEXT,
|
||||
"error_status" INTEGER,
|
||||
"error_message" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"started_at" TIMESTAMP(3),
|
||||
"finished_at" TIMESTAMP(3),
|
||||
"expires_at" TIMESTAMP(3) NOT NULL,
|
||||
CONSTRAINT "document_ai_runs_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "document_ai_runs_user_id_created_at_idx"
|
||||
ON "document_ai_runs" ("user_id", "created_at");
|
||||
CREATE INDEX IF NOT EXISTS "document_ai_runs_document_id_kind_status_idx"
|
||||
ON "document_ai_runs" ("document_id", "kind", "status");
|
||||
CREATE INDEX IF NOT EXISTS "document_ai_runs_status_expires_at_idx"
|
||||
ON "document_ai_runs" ("status", "expires_at");
|
||||
|
||||
DO $$ BEGIN
|
||||
ALTER TABLE "document_ai_runs"
|
||||
ADD CONSTRAINT "document_ai_runs_user_id_fkey"
|
||||
FOREIGN KEY ("user_id") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE;
|
||||
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
|
||||
DO $$ BEGIN
|
||||
ALTER TABLE "document_ai_runs"
|
||||
ADD CONSTRAINT "document_ai_runs_document_id_fkey"
|
||||
FOREIGN KEY ("document_id") REFERENCES "inbound_documents"("id") ON DELETE CASCADE ON UPDATE CASCADE;
|
||||
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
@@ -942,6 +942,7 @@ model User {
|
||||
// v1.39.3 — the person's own connection to a Paperless-ngx or Papra
|
||||
// instance for the document picker. See DocumentSourceConnection.
|
||||
documentSourceConnections DocumentSourceConnection[]
|
||||
documentAiRuns DocumentAiRun[]
|
||||
extractedFacts ExtractedFact[]
|
||||
documentConditionLinks DocumentConditionLink[]
|
||||
documentContentIndexes DocumentContentIndex[]
|
||||
@@ -8190,6 +8191,8 @@ model InboundDocument {
|
||||
shareLinks ClinicianShareLinkDocument[]
|
||||
/// Further import source keys that resolved to this document by content.
|
||||
sourceAliases DocumentSourceAlias[]
|
||||
/// Background AI reads of this document (v1.40). Transient; cascades.
|
||||
aiRuns DocumentAiRun[]
|
||||
|
||||
createdAt DateTime @default(now()) @map("created_at")
|
||||
updatedAt DateTime @updatedAt @map("updated_at")
|
||||
@@ -8987,3 +8990,75 @@ model OnboardingRecord {
|
||||
|
||||
@@map("onboarding_records")
|
||||
}
|
||||
|
||||
/// v1.40 — what a background document AI run is doing.
|
||||
enum DocumentAiRunKind {
|
||||
/// "Read with AI": transcribe the stored original into the content index,
|
||||
/// then continue into lab staging.
|
||||
DOCUMENT_INDEX
|
||||
/// A lab report scan (photo, PDF or browser-read text) into proposed rows
|
||||
/// for the review screen. Nothing is written until the person commits.
|
||||
LABS_OCR_EXTRACT
|
||||
|
||||
@@map("document_ai_run_kind")
|
||||
}
|
||||
|
||||
/// v1.40 — where a background document AI run stands.
|
||||
enum DocumentAiRunStatus {
|
||||
QUEUED
|
||||
RUNNING
|
||||
SUCCEEDED
|
||||
FAILED
|
||||
|
||||
@@map("document_ai_run_status")
|
||||
}
|
||||
|
||||
/// v1.40 — one model read that runs in the background instead of inside the
|
||||
/// request that asked for it, so a proxy that cuts long requests cannot cut
|
||||
/// the read. The request answers 202 with the run id; the client polls
|
||||
/// `GET /api/ai-runs/{id}`. A ticket, not a record: the input is dropped the
|
||||
/// moment the run finishes and the whole row is deleted an hour later
|
||||
/// (`document-ai-run-reaper`). Never in a backup.
|
||||
model DocumentAiRun {
|
||||
id String @id @default(cuid())
|
||||
userId String @map("user_id")
|
||||
kind DocumentAiRunKind
|
||||
/// The document a document run reads. Null for a lab scan, whose input is
|
||||
/// the upload itself.
|
||||
documentId String? @map("document_id")
|
||||
status DocumentAiRunStatus @default(QUEUED)
|
||||
/// How to run it: the input mode, the upload's sniffed MIME, the budget
|
||||
/// reservation taken at enqueue. Never health data.
|
||||
paramsJson Json @map("params_json")
|
||||
/// AES-256-GCM (`encryptBytes`, labelled) of the input the worker needs and
|
||||
/// cannot read anywhere else: the uploaded scan, or the text the browser
|
||||
/// read off it. Null for a run that reads a stored document, and set to
|
||||
/// null when the run finishes.
|
||||
inputEncrypted Bytes? @map("input_encrypted")
|
||||
/// AES-256-GCM (`encryptBytes`, labelled) of the result as JSON. It holds
|
||||
/// model text and lab values. Served by the poll route until the row goes.
|
||||
resultEncrypted Bytes? @map("result_encrypted")
|
||||
/// The failure code a client branches on, the same code the synchronous
|
||||
/// route would have answered with.
|
||||
errorCode String? @map("error_code")
|
||||
/// The HTTP status the synchronous route would have answered with, so a
|
||||
/// client handles a failed run as it handled the failed request.
|
||||
errorStatus Int? @map("error_status")
|
||||
/// The failure in one sanitised sentence. Never upstream text.
|
||||
errorMessage String? @map("error_message")
|
||||
createdAt DateTime @default(now()) @map("created_at")
|
||||
startedAt DateTime? @map("started_at")
|
||||
finishedAt DateTime? @map("finished_at")
|
||||
/// The next deadline the reaper acts on: queued runs fail as
|
||||
/// `aiRuns.workerUnavailable`, running ones as `aiRuns.timedOut`, finished
|
||||
/// ones are deleted.
|
||||
expiresAt DateTime @map("expires_at")
|
||||
|
||||
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
|
||||
document InboundDocument? @relation(fields: [documentId], references: [id], onDelete: Cascade)
|
||||
|
||||
@@index([userId, createdAt])
|
||||
@@index([documentId, kind, status])
|
||||
@@index([status, expiresAt])
|
||||
@@map("document_ai_runs")
|
||||
}
|
||||
@@ -924,6 +924,20 @@ async function main() {
|
||||
),
|
||||
);
|
||||
|
||||
// ───── v1.40 background document AI runs (Bytes, binary2, batched) ─────
|
||||
// A run's sealed input and result. Transient (deleted an hour after the run
|
||||
// finishes), but a run in flight during a rotation must still open under the
|
||||
// new key. Walked through the registry so the binary codec, the labels and
|
||||
// the disposable rule are honoured; NULL once a run has finished (input) or
|
||||
// before it has (result), which the walk skips.
|
||||
for (const field of ["inputEncrypted", "resultEncrypted"] as const) {
|
||||
results.push(
|
||||
await rotateRegistryColumn("DocumentAiRun", field, {
|
||||
documentAiRun: prisma.documentAiRun,
|
||||
} as unknown as CorpusClient),
|
||||
);
|
||||
}
|
||||
|
||||
// ───── Whole-account backup blob (String, batched) ─────
|
||||
// `DataBackup.data` holds `packBackupBlob()` output: AES-256-GCM ciphertext
|
||||
// under a column that does NOT carry the `*Encrypted` suffix, which is why
|
||||
|
||||
@@ -107,7 +107,7 @@ const EGRESS_SITES: Record<string, EgressSite> = {
|
||||
"lib/documents/describe.ts": {
|
||||
recheck: false,
|
||||
reason:
|
||||
"Document summary; reached from the summary route and the summary job, which resolve `documentAi` first.",
|
||||
"Document summary and transcription; reached from the summary and index routes, the summary job, and the background run worker (`lib/jobs/document-ai-run.ts`), which resolve `documentAi` first.",
|
||||
},
|
||||
"lib/documents/extract.ts": {
|
||||
recheck: false,
|
||||
@@ -117,7 +117,7 @@ const EGRESS_SITES: Record<string, EgressSite> = {
|
||||
"lib/labs/ocr-extract.ts": {
|
||||
recheck: false,
|
||||
reason:
|
||||
"Lab report scan; reached from the labs OCR route, which gates `labsOcr`.",
|
||||
"Lab report scan; reached from the background run worker (`lib/jobs/document-ai-run.ts`) for a scan the labs OCR route queued under `labsOcr`.",
|
||||
},
|
||||
"app/api/medications/extract/route.ts": {
|
||||
recheck: false,
|
||||
@@ -154,6 +154,42 @@ const EGRESS_SITES: Record<string, EgressSite> = {
|
||||
},
|
||||
};
|
||||
|
||||
/**
|
||||
* Background workers that reach one of the admitted helpers above with no
|
||||
* request around them. The route that queued the work asked the capability
|
||||
* once; the worker runs later, after a switch may have been turned off or
|
||||
* consent withdrawn, so it must ask again for the record and re-check the wire
|
||||
* for the provider it picked.
|
||||
*/
|
||||
const WORKER_REACHERS: Record<string, string> = {
|
||||
"lib/jobs/document-ai-run.ts":
|
||||
"Background document reads and lab scans (v1.40): `transcribeDocument` and `runOcrExtraction` through the shared run bodies.",
|
||||
};
|
||||
|
||||
const JOB_CAPABILITY = /\baiCapabilityForJob\s*\(/;
|
||||
const WIRE_RECHECK = /\baiEgressRefusal\s*\(/;
|
||||
|
||||
describe("workers that reach an admitted helper", () => {
|
||||
it.each(Object.keys(WORKER_REACHERS))(
|
||||
"%s asks the capability for the record and re-checks the wire",
|
||||
(rel) => {
|
||||
const text = code(rel);
|
||||
expect(JOB_CAPABILITY.test(text), `${rel}: no aiCapabilityForJob`).toBe(
|
||||
true,
|
||||
);
|
||||
expect(WIRE_RECHECK.test(text), `${rel}: no aiEgressRefusal`).toBe(true);
|
||||
},
|
||||
);
|
||||
|
||||
it("a worker whose re-check is only in a comment fails the matcher", () => {
|
||||
const planted = stripComments(
|
||||
"// aiEgressRefusal(key, userId, types)\nawait aiCapabilityForJob(userId, key);",
|
||||
);
|
||||
expect(JOB_CAPABILITY.test(planted)).toBe(true);
|
||||
expect(WIRE_RECHECK.test(planted)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("AI egress call sites", () => {
|
||||
const found = sourceFiles().filter((rel) => EGRESS_CALL.test(code(rel)));
|
||||
|
||||
|
||||
@@ -194,6 +194,14 @@ export const AI_ROUTES: Readonly<Record<string, AiRouteEntry>> = {
|
||||
* claim the guard checks against the file.
|
||||
*/
|
||||
export const DATA_ROUTES: Readonly<Record<string, string>> = {
|
||||
// v1.40 — the poll for a background document AI run. A plain read of a row
|
||||
// the caller created: every capability question was asked when the run was
|
||||
// queued and again in the worker before anything left, and serving the
|
||||
// answer to a read the person already asked for calls no model. Not
|
||||
// `aiCapabilityToServe` either: a run result is the caller's own request
|
||||
// result, not stored model text shown later.
|
||||
"src/app/api/ai-runs/[id]/route.ts":
|
||||
"The caller's own background run: status, result, failure.",
|
||||
// The design's list of formerly AI-gated data reads and writes.
|
||||
"src/app/api/insights/derived/batch/route.ts":
|
||||
"Deterministic scores and assessments for the whole grid.",
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
/**
|
||||
* v1.40 — GET /api/ai-runs/{id}: where one background document AI run stands.
|
||||
*
|
||||
* A plain read of a row the caller created: the run id came back in the 202
|
||||
* of the route that queued it, and the row is scoped to the caller here, so
|
||||
* another account's run (or one that never existed, or one an hour past its
|
||||
* end) is the same 404. No AI gate: serving the answer to a read the person
|
||||
* already asked for calls no model, and every capability question was asked
|
||||
* when the run was queued and again in the worker before anything left.
|
||||
*
|
||||
* `result` is the body the synchronous route answers with (by `kind`), and
|
||||
* `error` carries the code and status it would have answered with, so a
|
||||
* client handles both paths through one table.
|
||||
*/
|
||||
import { apiHandler, requireAuth } from "@/lib/api-handler";
|
||||
import { apiError, apiSuccess } from "@/lib/api-response";
|
||||
import { readAiRunForUser } from "@/lib/documents/ai-runs/store";
|
||||
import { AI_RUN_ERROR_CODES } from "@/lib/documents/ai-runs/types";
|
||||
|
||||
export const dynamic = "force-dynamic";
|
||||
|
||||
type RouteParams = { params: Promise<{ id: string }> };
|
||||
|
||||
export const GET = apiHandler(
|
||||
async (_request: Request, { params }: RouteParams) => {
|
||||
const { user } = await requireAuth();
|
||||
const { id } = await params;
|
||||
const run = await readAiRunForUser(user.id, id);
|
||||
if (!run) {
|
||||
return apiError("Run not found", 404, {
|
||||
errorCode: AI_RUN_ERROR_CODES.notFound,
|
||||
});
|
||||
}
|
||||
const response = apiSuccess(run);
|
||||
// A poll answer is for one person at one instant.
|
||||
response.headers.set("Cache-Control", "no-store");
|
||||
if (run.retryAfterMs !== null) {
|
||||
response.headers.set(
|
||||
"Retry-After",
|
||||
String(Math.max(1, Math.ceil(run.retryAfterMs / 1000))),
|
||||
);
|
||||
}
|
||||
return response;
|
||||
},
|
||||
);
|
||||
@@ -19,46 +19,52 @@
|
||||
* the lab / extract text mode already uses — no provider egress at all.
|
||||
*
|
||||
* Persists ONLY AES-256-GCM ciphertext text + opaque HMAC token hashes (A4).
|
||||
*
|
||||
* v1.40 — `Prefer: respond-async` (RFC 7240) runs the read in the background
|
||||
* worker: every refusal that can be answered quickly is still answered here,
|
||||
* then the route answers 202 with a run id and `GET /api/ai-runs/{id}` serves
|
||||
* the same body this route answers with synchronously. Without the header the
|
||||
* route behaves exactly as before; the iPhone app relies on that. Both paths
|
||||
* run one body (`executeDocumentIndex`).
|
||||
*/
|
||||
import { Buffer } from "node:buffer";
|
||||
|
||||
import { NextRequest } from "next/server";
|
||||
|
||||
import { apiHandler, requireAuth } from "@/lib/api-handler";
|
||||
import {
|
||||
apiError,
|
||||
apiSuccess,
|
||||
apiValidationError,
|
||||
getClientIp,
|
||||
safeJson,
|
||||
sanitiseZodIssues,
|
||||
} from "@/lib/api-response";
|
||||
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 {
|
||||
checkDocumentAiRateLimit,
|
||||
documentAiRateLimited,
|
||||
DOCUMENT_AI_TEXT_BODY_MAX_BYTES,
|
||||
loadOwnedDocument,
|
||||
prepareVisionInput,
|
||||
refundDocumentAiSlot,
|
||||
type LoadedDocument,
|
||||
} from "@/lib/documents/ai-route-support";
|
||||
import { maybeAutoStageLabFacts } from "@/lib/documents/auto-stage-labs";
|
||||
import { upsertContentIndex } from "@/lib/documents/content-index";
|
||||
import { recordIndexAttempt } from "@/lib/documents/index-document";
|
||||
import {
|
||||
DocumentDescribeError,
|
||||
transcribeDocument,
|
||||
} from "@/lib/documents/describe";
|
||||
acceptedRunResponse,
|
||||
outcomeResponse,
|
||||
prefersRespondAsync,
|
||||
} from "@/lib/documents/ai-runs/http";
|
||||
import {
|
||||
DOCUMENT_INDEX_RESERVE_TOKENS,
|
||||
executeDocumentIndex,
|
||||
} from "@/lib/documents/ai-runs/index-run";
|
||||
import { findLiveDocumentRun, startAiRun } from "@/lib/documents/ai-runs/start";
|
||||
import { requireDocumentVisionProvider } from "@/lib/documents/provider-order";
|
||||
import { annotate } from "@/lib/logging/context";
|
||||
import { requireModuleEnabled } from "@/lib/modules/gate";
|
||||
import { prisma } from "@/lib/db";
|
||||
import { inboundTextExtractSchema } from "@/lib/validations/inbound-documents";
|
||||
@@ -82,57 +88,26 @@ export const POST = apiHandler(
|
||||
});
|
||||
}
|
||||
|
||||
const background = prefersRespondAsync(request);
|
||||
const contentType = request.headers.get("content-type") ?? "";
|
||||
if (contentType.includes("application/json")) {
|
||||
return handleTextIndex(request, user.id, document);
|
||||
return handleTextIndex(request, user.id, document, background);
|
||||
}
|
||||
return handleVisionIndex(request, user.id, document);
|
||||
return handleVisionIndex(request, user.id, document, background);
|
||||
},
|
||||
);
|
||||
|
||||
async function finishIndex(
|
||||
request: NextRequest,
|
||||
/**
|
||||
* A read of this document already queued or running: attach to it rather than
|
||||
* charge a second one. Background path only; the synchronous path keeps its
|
||||
* behaviour of one request, one read.
|
||||
*/
|
||||
async function liveRunResponse(
|
||||
userId: string,
|
||||
documentId: string,
|
||||
source: "vision" | "text-ocr",
|
||||
tokenCount: number,
|
||||
): Promise<Response> {
|
||||
// Refs #776 — the manual route is the second writer of the attempt record
|
||||
// (the auto job + backfill share `indexLoadedDocument`): a success clears
|
||||
// any stored failure reason so the detail view stops explaining a problem
|
||||
// that no longer exists.
|
||||
await recordIndexAttempt(userId, documentId, {
|
||||
indexed: true,
|
||||
source,
|
||||
tokenCount,
|
||||
});
|
||||
// Read with AI is the person reading the document on purpose, so an
|
||||
// import's hold on automatic AI reading (`aiRead=defer`) ends here, before
|
||||
// the lab staging below that would otherwise refuse it.
|
||||
await prisma.inboundDocument.updateMany({
|
||||
where: { id: documentId, userId, aiReadDeferred: true },
|
||||
data: { aiReadDeferred: false },
|
||||
});
|
||||
await auditLog("documents.inbound.index", {
|
||||
userId,
|
||||
ipAddress: getClientIp(request),
|
||||
details: { documentId, source, tokens: tokenCount },
|
||||
});
|
||||
annotate({
|
||||
action: { name: "documents.contentIndex.upsert" },
|
||||
meta: { documentId, source, tokens: tokenCount },
|
||||
});
|
||||
// A manual read continues into the SAME lab staging the automatic index
|
||||
// worker performs, so a skipped or failed auto run is recoverable per
|
||||
// document without a re-upload. Every guard lives inside the helper (both
|
||||
// modules on, still STORED with no facts, provider + consent, looks like a
|
||||
// lab report) — a non-lab document is a tagged no-op, and a staging failure
|
||||
// never fails the index that just succeeded.
|
||||
const staging = await maybeAutoStageLabFacts(userId, documentId).catch(
|
||||
() => null,
|
||||
);
|
||||
const labFactsStaged = staging?.staged === true ? staging.facts : 0;
|
||||
return apiSuccess({ documentId, indexed: true, tokenCount, labFactsStaged });
|
||||
): Promise<Response | null> {
|
||||
const live = await findLiveDocumentRun(userId, documentId, "DOCUMENT_INDEX");
|
||||
return live ? acceptedRunResponse(live) : null;
|
||||
}
|
||||
|
||||
/** TEXT mode — index browser-OCR'd text (no provider egress). */
|
||||
@@ -140,6 +115,7 @@ async function handleTextIndex(
|
||||
request: NextRequest,
|
||||
userId: string,
|
||||
document: LoadedDocument,
|
||||
background: boolean,
|
||||
): Promise<Response> {
|
||||
const row = await prisma.user.findUnique({
|
||||
where: { id: userId },
|
||||
@@ -151,6 +127,11 @@ async function handleTextIndex(
|
||||
});
|
||||
}
|
||||
|
||||
if (background) {
|
||||
const live = await liveRunResponse(userId, document.id);
|
||||
if (live) return live;
|
||||
}
|
||||
|
||||
const rl = await checkDocumentAiRateLimit(userId);
|
||||
if (!rl.allowed) return documentAiRateLimited(rl);
|
||||
|
||||
@@ -174,14 +155,27 @@ async function handleTextIndex(
|
||||
);
|
||||
}
|
||||
|
||||
const { tokenCount } = await upsertContentIndex({
|
||||
userId,
|
||||
documentId: document.id,
|
||||
text: parsed.data.text,
|
||||
source: "text-ocr",
|
||||
providerType: null,
|
||||
});
|
||||
return finishIndex(request, userId, document.id, "text-ocr", tokenCount);
|
||||
if (background) {
|
||||
// The text is indexed in the worker, which then continues into the same
|
||||
// lab staging (model calls) the synchronous path makes.
|
||||
return startAiRun({
|
||||
userId,
|
||||
kind: "DOCUMENT_INDEX",
|
||||
documentId: document.id,
|
||||
params: { mode: "text" },
|
||||
input: Buffer.from(parsed.data.text, "utf8"),
|
||||
refundSlot: () => refundDocumentAiSlot(userId),
|
||||
});
|
||||
}
|
||||
|
||||
return outcomeResponse(
|
||||
await executeDocumentIndex({
|
||||
userId,
|
||||
document,
|
||||
input: { mode: "text", text: parsed.data.text },
|
||||
origin: { ipAddress: getClientIp(request), worker: false },
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
/** VISION mode — transcribe the stored original, then index the text. */
|
||||
@@ -189,6 +183,7 @@ async function handleVisionIndex(
|
||||
request: NextRequest,
|
||||
userId: string,
|
||||
document: LoadedDocument,
|
||||
background: boolean,
|
||||
): Promise<Response> {
|
||||
// Transcribing the stored original is model work; indexing text the browser
|
||||
// already read (the text mode above) is not, and stays open with AI off so
|
||||
@@ -197,59 +192,18 @@ async function handleVisionIndex(
|
||||
await requireAiCapability("documentAi", { pickDecides: true });
|
||||
const pick = await requireDocumentVisionProvider(userId);
|
||||
|
||||
if (background) {
|
||||
const live = await liveRunResponse(userId, document.id);
|
||||
if (live) return live;
|
||||
}
|
||||
|
||||
const rl = await checkDocumentAiRateLimit(userId);
|
||||
if (!rl.allowed) return documentAiRateLimited(rl);
|
||||
|
||||
const vision = await prepareVisionInput(document, pick.pdfSupported);
|
||||
if (!vision.ok) {
|
||||
// Preparation failed before any provider dispatch — the slot goes back.
|
||||
// Refs #776 — each failure is also recorded on the row so the detail view
|
||||
// can explain the missing index after the toast is gone.
|
||||
await refundDocumentAiSlot(userId);
|
||||
if (vision.reason === "pdfNeedsAnthropic") {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "pdf-needs-anthropic",
|
||||
});
|
||||
return apiError(
|
||||
"PDF scanning needs a Claude vision provider; use local OCR instead.",
|
||||
422,
|
||||
{ errorCode: "documents.inbound.pdfNeedsAnthropic" },
|
||||
);
|
||||
}
|
||||
if (vision.reason === "rasterFailed") {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "raster-failed",
|
||||
});
|
||||
return apiError("The PDF pages couldn't be rendered for scanning.", 422, {
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
if (vision.reason === "fileType") {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "local-unsupported",
|
||||
});
|
||||
return apiError(
|
||||
"This document can't be scanned. Use local OCR (text mode).",
|
||||
422,
|
||||
{ errorCode: "documents.inbound.fileType" },
|
||||
);
|
||||
}
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "decrypt-error",
|
||||
});
|
||||
return apiError("Couldn't read the stored document.", 422, {
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
|
||||
const dateKey = buildDateKey();
|
||||
const reservation = await reserveBudget(
|
||||
userId,
|
||||
AI_BUDGETS.documentTranscribe.maxTokens,
|
||||
DOCUMENT_INDEX_RESERVE_TOKENS,
|
||||
dateKey,
|
||||
resolveDailyCap([{ providerType: pick.entry.providerType }]),
|
||||
resolveCostOwner([{ providerType: pick.entry.providerType }]),
|
||||
@@ -261,66 +215,30 @@ async function handleVisionIndex(
|
||||
errorCode: "documents.inbound.budgetExceeded",
|
||||
});
|
||||
}
|
||||
const budget = {
|
||||
reserved: reservation.reserved,
|
||||
owner: reservation.owner,
|
||||
dateKey,
|
||||
};
|
||||
|
||||
try {
|
||||
const { text } = await transcribeDocument({
|
||||
provider: pick.entry.instance,
|
||||
providerType: pick.providerType,
|
||||
images: vision.images,
|
||||
documents: vision.documents,
|
||||
});
|
||||
await reconcileSpend(
|
||||
userId,
|
||||
reservation.reserved,
|
||||
reservation.reserved,
|
||||
dateKey,
|
||||
0,
|
||||
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
|
||||
);
|
||||
// Refs #776 — the empty-transcription guard, same contract as the auto
|
||||
// path (`tryProviderIndex`): a provider answer with no text must never
|
||||
// become a "successful" empty index. The spend stays charged (the
|
||||
// provider was called); the honest answer is an error plus the recorded
|
||||
// reason.
|
||||
if (text.trim().length === 0) {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "empty-transcription",
|
||||
});
|
||||
return apiError(
|
||||
"The provider returned no text for this document. Try a clearer copy.",
|
||||
422,
|
||||
{ errorCode: "documents.inbound.extractFailed" },
|
||||
);
|
||||
}
|
||||
const { tokenCount } = await upsertContentIndex({
|
||||
if (background) {
|
||||
// The worker picks the provider again and re-checks the wire for it right
|
||||
// before the stored original leaves.
|
||||
return startAiRun({
|
||||
userId,
|
||||
kind: "DOCUMENT_INDEX",
|
||||
documentId: document.id,
|
||||
text,
|
||||
source: "vision",
|
||||
providerType: pick.providerType,
|
||||
});
|
||||
return finishIndex(request, userId, document.id, "vision", tokenCount);
|
||||
} catch (err) {
|
||||
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
|
||||
servedBy: null,
|
||||
reservedOwner: reservation.owner,
|
||||
});
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "provider-error",
|
||||
});
|
||||
if (err instanceof DocumentDescribeError) {
|
||||
return apiError("Couldn't read the document. Try a clearer copy.", 422, {
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
annotate({
|
||||
action: { name: "documents.contentIndex.failed" },
|
||||
meta: { reason: "provider_error", mode: "vision" },
|
||||
});
|
||||
return apiError("Couldn't read the document. Try a clearer copy.", 502, {
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
params: { mode: "vision", budget },
|
||||
refundSlot: () => refundDocumentAiSlot(userId),
|
||||
});
|
||||
}
|
||||
|
||||
return outcomeResponse(
|
||||
await executeDocumentIndex({
|
||||
userId,
|
||||
document,
|
||||
input: { mode: "vision", pick, budget },
|
||||
origin: { ipAddress: getClientIp(request), worker: false },
|
||||
}),
|
||||
);
|
||||
}
|
||||
@@ -1,10 +1,13 @@
|
||||
/**
|
||||
* v1.20.1 — POST /api/labs/ocr/extract, TEXT mode.
|
||||
* v1.20.1 — POST /api/labs/ocr/extract.
|
||||
*
|
||||
* Focus: the text-mode structuring pass reserves the proportionate text budget
|
||||
* ceiling (`AI_BUDGETS.ocrExtractText`), not the far larger vision ceiling, and
|
||||
* on a clean extraction failure it refunds the reservation in full rather than
|
||||
* charging it.
|
||||
* ceiling (`AI_BUDGETS.ocrExtractText`), not the far larger vision ceiling.
|
||||
*
|
||||
* v1.40 — the route queues the read and answers 202; the read itself, its
|
||||
* budget settlement and its failures are pinned in
|
||||
* `src/lib/labs/__tests__/ocr-run.test.ts`. What stays here is everything the
|
||||
* route still refuses before anything is queued.
|
||||
*/
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
|
||||
@@ -48,6 +51,18 @@ vi.mock("@/lib/ai/coach/budget", () => ({
|
||||
resolveDailyCapFor: vi.fn(() => 200_000),
|
||||
resolveCostOwner: vi.fn(() => "operator" as const),
|
||||
}));
|
||||
vi.mock("@/lib/documents/ai-runs/start", () => ({
|
||||
startAiRun: vi.fn(
|
||||
async () =>
|
||||
new Response(
|
||||
JSON.stringify({
|
||||
data: { runId: "run-1", status: "QUEUED", pollAfterMs: 1500 },
|
||||
error: null,
|
||||
}),
|
||||
{ status: 202 },
|
||||
),
|
||||
),
|
||||
}));
|
||||
vi.mock("@/lib/labs/ocr-extract", async () => {
|
||||
const actual = await vi.importActual<typeof import("@/lib/labs/ocr-extract")>(
|
||||
"@/lib/labs/ocr-extract",
|
||||
@@ -56,6 +71,7 @@ vi.mock("@/lib/labs/ocr-extract", async () => {
|
||||
});
|
||||
|
||||
import { POST } from "../route";
|
||||
import { startAiRun } from "@/lib/documents/ai-runs/start";
|
||||
import { AI_BUDGETS } from "@/lib/ai/ai-budgets";
|
||||
import { requireAuth } from "@/lib/api-handler";
|
||||
import { prisma } from "@/lib/db";
|
||||
@@ -63,10 +79,9 @@ import { requireLabsOcrProvider } from "@/lib/labs/ocr-capability";
|
||||
import { requireAiCapability } from "@/lib/ai/capabilities/gate";
|
||||
import { AiUnavailableError } from "@/lib/ai/capabilities/refusal";
|
||||
import { rasterizePdf } from "@/lib/documents/rasterize-pdf";
|
||||
import { reserveBudget, reconcileSpend } from "@/lib/ai/coach/budget";
|
||||
import { reserveBudget } from "@/lib/ai/coach/budget";
|
||||
import { checkRateLimit, refundRateLimit } from "@/lib/rate-limit";
|
||||
import { OcrExtractError, runOcrExtraction } from "@/lib/labs/ocr-extract";
|
||||
import { annotate } from "@/lib/logging/context";
|
||||
import { runOcrExtraction } from "@/lib/labs/ocr-extract";
|
||||
|
||||
const SESSION_OK = {
|
||||
session: { id: "sess-1", expiresAt: new Date(Date.now() + 3_600_000) },
|
||||
@@ -113,10 +128,8 @@ beforeEach(() => {
|
||||
|
||||
describe("POST /api/labs/ocr/extract — text mode budget", () => {
|
||||
it("reserves the cheaper text ceiling, not the vision ceiling", async () => {
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
|
||||
const res = await POST(textReq());
|
||||
expect(res.status).toBe(200);
|
||||
expect(res.status).toBe(202);
|
||||
|
||||
// The reservation must use the text ceiling — proven distinct from the
|
||||
// vision ceiling so a future merge can't silently re-point it.
|
||||
@@ -136,61 +149,29 @@ describe("POST /api/labs/ocr/extract — text mode budget", () => {
|
||||
);
|
||||
});
|
||||
|
||||
it("refunds the reservation in full on a clean extract failure", async () => {
|
||||
vi.mocked(runOcrExtraction).mockRejectedValue(
|
||||
new OcrExtractError("unreadable"),
|
||||
);
|
||||
|
||||
const res = await POST(textReq());
|
||||
expect(res.status).toBe(422);
|
||||
|
||||
// actual spend reconciles to 0 — the failed structuring pass is refunded,
|
||||
// not charged at the reserved estimate.
|
||||
expect(reconcileSpend).toHaveBeenCalledWith(
|
||||
"user-1",
|
||||
AI_BUDGETS.ocrExtractText.maxTokens,
|
||||
0,
|
||||
"2026-06-26",
|
||||
0,
|
||||
{ servedBy: null, reservedOwner: "operator" },
|
||||
);
|
||||
});
|
||||
|
||||
it("identifies provider failures instead of blaming image quality", async () => {
|
||||
vi.mocked(runOcrExtraction).mockRejectedValue(new Error("provider down"));
|
||||
|
||||
const res = await POST(textReq());
|
||||
const body = (await res.json()) as { error: string };
|
||||
|
||||
expect(res.status).toBe(502);
|
||||
expect(body.error).toContain("configured AI provider");
|
||||
});
|
||||
|
||||
it("never passes the upstream body on, to the caller or to the event", async () => {
|
||||
const upstreamBody = "ami-id instance-id iam/security-credentials/admin";
|
||||
vi.mocked(runOcrExtraction).mockRejectedValue(
|
||||
Object.assign(new Error("Local AI request failed (403)"), {
|
||||
httpStatus: 403,
|
||||
model: "llama3",
|
||||
bodyExcerpt: upstreamBody,
|
||||
}),
|
||||
);
|
||||
|
||||
const res = await POST(textReq());
|
||||
expect(res.status).toBe(502);
|
||||
expect(await res.text()).not.toContain("iam/security-credentials");
|
||||
const events = JSON.stringify(vi.mocked(annotate).mock.calls);
|
||||
expect(events).not.toContain("iam/security-credentials");
|
||||
// The status and the model stay, so an operator can still tell why.
|
||||
expect(annotate).toHaveBeenCalledWith(
|
||||
it("queues the text, sealed in the run, with the reservation it took", async () => {
|
||||
const res = await POST(textReq("Glucose 95 mg/dL"));
|
||||
expect(res.status).toBe(202);
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
expect(startAiRun).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
meta: expect.objectContaining({ upstreamStatus: 403, model: "llama3" }),
|
||||
userId: "user-1",
|
||||
kind: "LABS_OCR_EXTRACT",
|
||||
params: {
|
||||
mode: "text",
|
||||
budget: {
|
||||
reserved: AI_BUDGETS.ocrExtractText.maxTokens,
|
||||
owner: "operator",
|
||||
dateKey: "2026-06-26",
|
||||
},
|
||||
},
|
||||
input: Buffer.from("Glucose 95 mg/dL", "utf8"),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("POST /api/labs/ocr/extract — vision PDF rasterization", () => {
|
||||
describe("POST /api/labs/ocr/extract — vision uploads", () => {
|
||||
function pdfReq(): Request {
|
||||
// A minimal `%PDF-` header is all `detectOcrMimeType` needs to sniff a PDF.
|
||||
const bytes = new Uint8Array([
|
||||
@@ -225,50 +206,23 @@ describe("POST /api/labs/ocr/extract — vision PDF rasterization", () => {
|
||||
} as never);
|
||||
});
|
||||
|
||||
it("rasterizes a PDF for a non-Anthropic vision provider and sends the page images", async () => {
|
||||
vi.mocked(rasterizePdf).mockResolvedValue({
|
||||
ok: true,
|
||||
images: [{ mediaType: "image/jpeg", dataBase64: "cGFnZQ==" }],
|
||||
pageCount: 1,
|
||||
});
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
|
||||
it("queues a PDF with its sniffed type; the worker renders it", async () => {
|
||||
const res = await POST(pdfReq());
|
||||
expect(res.status).toBe(200);
|
||||
// The whole PDF was read: no coverage note on the response.
|
||||
const whole = (await res.json()) as { data: Record<string, unknown> };
|
||||
expect(whole.data).not.toHaveProperty("pageCoverage");
|
||||
expect(rasterizePdf).toHaveBeenCalledOnce();
|
||||
// The rendered page images flow through as `input_image`s; no native PDF
|
||||
// document block is sent for a non-Anthropic provider.
|
||||
expect(runOcrExtraction).toHaveBeenCalledWith(
|
||||
expect(res.status).toBe(202);
|
||||
// Rendering is the worker's job now; nothing is read in the request.
|
||||
expect(rasterizePdf).not.toHaveBeenCalled();
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
expect(startAiRun).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
images: [{ mediaType: "image/jpeg", dataBase64: "cGFnZQ==" }],
|
||||
documents: [],
|
||||
kind: "LABS_OCR_EXTRACT",
|
||||
params: expect.objectContaining({
|
||||
mode: "vision",
|
||||
mime: "application/pdf",
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("says on the response when only the first pages of a long PDF were read", async () => {
|
||||
const tenPages = Array.from({ length: 10 }, () => ({
|
||||
mediaType: "image/jpeg" as const,
|
||||
dataBase64: "cGFnZQ==",
|
||||
}));
|
||||
vi.mocked(rasterizePdf).mockResolvedValue({
|
||||
ok: true,
|
||||
images: tenPages,
|
||||
pageCount: 23,
|
||||
});
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
|
||||
const res = await POST(pdfReq());
|
||||
expect(res.status).toBe(200);
|
||||
const body = (await res.json()) as {
|
||||
data: { pageCoverage?: { read: number; total: number } };
|
||||
};
|
||||
expect(body.data.pageCoverage).toEqual({ read: 10, total: 23 });
|
||||
});
|
||||
|
||||
it("hands the slot back when the upload is not an image or PDF", async () => {
|
||||
const form = new FormData();
|
||||
form.append(
|
||||
@@ -287,36 +241,13 @@ describe("POST /api/labs/ocr/extract — vision PDF rasterization", () => {
|
||||
expect(refundRateLimit).toHaveBeenCalledWith("labs-ocr:user-1");
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("falls back to pdfNeedsAnthropic when rasterization fails", async () => {
|
||||
vi.mocked(rasterizePdf).mockResolvedValue({
|
||||
ok: false,
|
||||
reason: "render-failed",
|
||||
});
|
||||
|
||||
const res = await POST(pdfReq());
|
||||
expect(res.status).toBe(422);
|
||||
const body = (await res.json()) as { error: string | null };
|
||||
expect(body.error).toBeTruthy();
|
||||
// A failed render never reaches the provider — and refunds the reservation.
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
expect(reconcileSpend).toHaveBeenCalledWith(
|
||||
"user-1",
|
||||
AI_BUDGETS.ocrExtract.maxTokens,
|
||||
0,
|
||||
"2026-06-26",
|
||||
0,
|
||||
{ servedBy: null, reservedOwner: "operator" },
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("POST /api/labs/ocr/extract — the hourly scan bucket", () => {
|
||||
it("charges the ceiling the operator set", async () => {
|
||||
vi.stubEnv("LABS_OCR_LIMIT_PER_HOUR", "40");
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
const res = await POST(textReq() as never);
|
||||
expect(res.status).toBe(200);
|
||||
expect(res.status).toBe(202);
|
||||
expect(checkRateLimit).toHaveBeenCalledWith(
|
||||
"labs-ocr:user-1",
|
||||
40,
|
||||
@@ -366,15 +297,9 @@ describe("POST /api/labs/ocr/extract — the hourly scan bucket", () => {
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("keeps the slot once the provider was actually called", async () => {
|
||||
vi.mocked(runOcrExtraction).mockRejectedValue(new Error("provider down"));
|
||||
const failed = await POST(textReq() as never);
|
||||
expect(failed.status).toBe(502);
|
||||
expect(refundRateLimit).not.toHaveBeenCalled();
|
||||
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
const ok = await POST(textReq() as never);
|
||||
expect(ok.status).toBe(200);
|
||||
it("keeps the slot once the scan is queued", async () => {
|
||||
const queued = await POST(textReq() as never);
|
||||
expect(queued.status).toBe(202);
|
||||
expect(refundRateLimit).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -3,7 +3,15 @@
|
||||
*
|
||||
* Read-only (NOT idempotent) extraction of a paper lab report into STRUCTURED
|
||||
* proposed rows for the mandatory human-review screen. Nothing is written to
|
||||
* the database here; the raw upload is never persisted or logged.
|
||||
* the person's record here; the raw upload is never logged.
|
||||
*
|
||||
* v1.40 — the read runs in the background worker. This route answers every
|
||||
* refusal it can answer quickly (capability, provider, rate bucket, budget, a
|
||||
* malformed or oversized upload, an unknown file type), seals the upload or
|
||||
* the text into a `DocumentAiRun` (AES-256-GCM, dropped the moment the run
|
||||
* finishes, the row deleted an hour later) and answers 202 with the run id.
|
||||
* `GET /api/ai-runs/{id}` serves the proposed rows. Only the web client calls
|
||||
* this route, so there is no synchronous form to keep.
|
||||
*
|
||||
* Two modes, dispatched on the request content-type:
|
||||
*
|
||||
@@ -22,9 +30,11 @@
|
||||
* re-check the capability for that pick (`requireLabsOcrProvider`: the
|
||||
* operator's switch, the labs module, and a document consent receipt for a
|
||||
* pick that leaves the machine) → rate-limit
|
||||
* (`LABS_OCR_LIMIT_PER_HOUR`, default 6/h) → reserveBudget → run extraction
|
||||
* → reconcile budget. A slot is charged early so a 429 stays cheap, and it is
|
||||
* handed back when the scan fails before the provider is called.
|
||||
* (`LABS_OCR_LIMIT_PER_HOUR`, default 6/h) → reserveBudget → queue the run.
|
||||
* The worker resolves the provider again, re-checks the wire, runs the
|
||||
* extraction and reconciles the budget (`src/lib/labs/ocr-run.ts`). A slot is
|
||||
* charged early so a 429 stays cheap, and it is handed back when the scan
|
||||
* fails before the provider is called.
|
||||
*
|
||||
* Extracted text is UNTRUSTED (prompt-injection): the server never acts on an
|
||||
* instruction inside the document — the human review step is the safety
|
||||
@@ -35,12 +45,10 @@ import { Buffer } from "node:buffer";
|
||||
import { apiHandler, requireAuth } from "@/lib/api-handler";
|
||||
import {
|
||||
apiError,
|
||||
apiSuccess,
|
||||
apiValidationError,
|
||||
safeJson,
|
||||
sanitiseZodIssues,
|
||||
} from "@/lib/api-response";
|
||||
import { AI_BUDGETS } from "@/lib/ai/ai-budgets";
|
||||
import { requireAiCapability } from "@/lib/ai/capabilities/gate";
|
||||
import {
|
||||
buildDateKey,
|
||||
@@ -49,10 +57,10 @@ import {
|
||||
resolveCostOwner,
|
||||
resolveDailyCap,
|
||||
} from "@/lib/ai/coach/budget";
|
||||
import { startAiRun } from "@/lib/documents/ai-runs/start";
|
||||
import { prisma } from "@/lib/db";
|
||||
import { rasterizePdf } from "@/lib/documents/rasterize-pdf";
|
||||
import { requireLabsOcrProvider } from "@/lib/labs/ocr-capability";
|
||||
import { OcrExtractError, runOcrExtraction } from "@/lib/labs/ocr-extract";
|
||||
import { OCR_RESERVE_TOKENS } from "@/lib/labs/ocr-run";
|
||||
import {
|
||||
BodyTooLargeError,
|
||||
detectOcrMimeType,
|
||||
@@ -72,27 +80,6 @@ export const dynamic = "force-dynamic";
|
||||
/** OCR'd text is bounded in the schema; cap the JSON body proportionally. */
|
||||
const TEXT_BODY_MAX_BYTES = 512 * 1024;
|
||||
|
||||
/**
|
||||
* The wide-event facts of a failed provider call: the status and the model,
|
||||
* never the upstream body. The base URL can be one a user typed, and what an
|
||||
* arbitrary endpoint answers has no business in the event stream (v1.39.3);
|
||||
* the response to the caller carries neither.
|
||||
*/
|
||||
function providerFailureMeta(error: unknown, mode?: "text") {
|
||||
const err = error as {
|
||||
httpStatus?: unknown;
|
||||
model?: unknown;
|
||||
};
|
||||
return {
|
||||
reason: "provider_error",
|
||||
...(mode ? { mode } : {}),
|
||||
...(typeof err.httpStatus === "number"
|
||||
? { upstreamStatus: err.httpStatus }
|
||||
: {}),
|
||||
...(typeof err.model === "string" ? { model: err.model } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
export const POST = apiHandler(async (request: Request) => {
|
||||
const { user } = await requireAuth();
|
||||
|
||||
@@ -111,7 +98,7 @@ export const POST = apiHandler(async (request: Request) => {
|
||||
});
|
||||
|
||||
/**
|
||||
* TEXT mode (v1.18.10) — structure in-browser-OCR'd text via any configured
|
||||
* TEXT mode (v1.18.10) — queue in-browser-OCR'd text for any configured
|
||||
* provider. No image bytes reach the server; the raw image stayed on-device.
|
||||
*/
|
||||
async function handleTextExtract(
|
||||
@@ -173,7 +160,7 @@ async function handleTextExtract(
|
||||
const dateKey = buildDateKey();
|
||||
const reservation = await reserveBudget(
|
||||
userId,
|
||||
AI_BUDGETS.ocrExtractText.maxTokens,
|
||||
OCR_RESERVE_TOKENS.text,
|
||||
dateKey,
|
||||
resolveDailyCap([{ providerType: pick.entry.providerType }]),
|
||||
resolveCostOwner([{ providerType: pick.entry.providerType }]),
|
||||
@@ -190,51 +177,27 @@ async function handleTextExtract(
|
||||
});
|
||||
}
|
||||
|
||||
try {
|
||||
const result = await runOcrExtraction({
|
||||
userId,
|
||||
provider: pick.entry.instance,
|
||||
providerType: pick.providerType,
|
||||
ocrText: parsed.data.text,
|
||||
});
|
||||
// A clean structuring pass spent (at most) the reservation; the provider
|
||||
// already billed it, so reconcile against the reserved estimate.
|
||||
await reconcileSpend(
|
||||
userId,
|
||||
reservation.reserved,
|
||||
reservation.reserved,
|
||||
dateKey,
|
||||
0,
|
||||
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
|
||||
);
|
||||
return apiSuccess(result);
|
||||
} catch (err) {
|
||||
// A failed structuring call produced no usable rows; mirror the vision
|
||||
// path and refund the reservation in full rather than charging it.
|
||||
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
|
||||
servedBy: null,
|
||||
reservedOwner: reservation.owner,
|
||||
});
|
||||
if (err instanceof OcrExtractError) {
|
||||
return apiError("Couldn't read the report. Try a clearer photo.", 422, {
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
});
|
||||
}
|
||||
annotate({
|
||||
action: { name: "labs.ocr.extractFailed" },
|
||||
meta: providerFailureMeta(err, "text"),
|
||||
});
|
||||
return apiError(
|
||||
"The configured AI provider could not process this report. Check the provider configuration and retry.",
|
||||
502,
|
||||
{ errorCode: "labs.ocr.extractFailed" },
|
||||
);
|
||||
}
|
||||
// The structuring pass runs in the worker, which picks the provider again
|
||||
// and re-checks the wire for it before the text leaves.
|
||||
return startAiRun({
|
||||
userId,
|
||||
kind: "LABS_OCR_EXTRACT",
|
||||
params: {
|
||||
mode: "text",
|
||||
budget: {
|
||||
reserved: reservation.reserved,
|
||||
owner: reservation.owner,
|
||||
dateKey,
|
||||
},
|
||||
},
|
||||
input: Buffer.from(parsed.data.text, "utf8"),
|
||||
refundSlot: () => refundLabsOcrSlot(userId),
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* VISION mode — run a multipart photo / PDF through the user's vision-capable
|
||||
* provider. The image transits server memory ephemerally and is never stored.
|
||||
* VISION mode — queue a multipart photo / PDF for the user's vision-capable
|
||||
* provider. The upload is sealed into the run until the read finishes.
|
||||
*/
|
||||
async function handleVisionExtract(
|
||||
request: Request,
|
||||
@@ -258,7 +221,7 @@ async function handleVisionExtract(
|
||||
const dateKey = buildDateKey();
|
||||
const reservation = await reserveBudget(
|
||||
userId,
|
||||
AI_BUDGETS.ocrExtract.maxTokens,
|
||||
OCR_RESERVE_TOKENS.vision,
|
||||
dateKey,
|
||||
resolveDailyCap([{ providerType: pick.entry.providerType }]),
|
||||
resolveCostOwner([{ providerType: pick.entry.providerType }]),
|
||||
@@ -359,120 +322,23 @@ async function handleVisionExtract(
|
||||
});
|
||||
}
|
||||
|
||||
// 7. PDF handling — mirror the document vault's `prepareVisionInput`
|
||||
// (`src/lib/documents/ai-route-support.ts`) exactly:
|
||||
// - Anthropic (`pick.pdfSupported`) reads the native `document` block at
|
||||
// highest fidelity;
|
||||
// - every other vision provider (codex / any image-only wire) reads the
|
||||
// PDF via `rasterizePdf`, which renders the pages to JPEG `input_image`
|
||||
// parts the provider already handles. Rendering is pure local compute
|
||||
// (no egress). `rasterizePdf` never throws: a malformed / encrypted /
|
||||
// unrenderable PDF returns `{ ok: false }`, and we keep the shipped
|
||||
// `pdfNeedsAnthropic` reject as the graceful fallback.
|
||||
let images: {
|
||||
mediaType: "image/jpeg" | "image/png" | "image/webp";
|
||||
dataBase64: string;
|
||||
}[];
|
||||
let documents: { mediaType: "application/pdf"; dataBase64: string }[];
|
||||
let pageCoverage: { read: number; total: number } | undefined;
|
||||
|
||||
if (mime === "application/pdf") {
|
||||
if (pick.pdfSupported) {
|
||||
images = [];
|
||||
documents = [
|
||||
{
|
||||
mediaType: "application/pdf",
|
||||
dataBase64: buffer.toString("base64"),
|
||||
},
|
||||
];
|
||||
} else {
|
||||
const raster = await rasterizePdf(buffer);
|
||||
if (!raster.ok) {
|
||||
annotate({
|
||||
action: { name: "labs.ocr.fileRejected" },
|
||||
meta: { reason: "pdf_rasterize_failed" },
|
||||
});
|
||||
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
|
||||
servedBy: null,
|
||||
reservedOwner: reservation.owner,
|
||||
});
|
||||
await refundLabsOcrSlot(userId);
|
||||
return apiError(
|
||||
"Couldn't read this PDF; upload a photo instead.",
|
||||
422,
|
||||
{ errorCode: "labs.ocr.pdfNeedsAnthropic" },
|
||||
);
|
||||
}
|
||||
images = raster.images;
|
||||
documents = [];
|
||||
// A PDF longer than the page cap is read from its first pages only.
|
||||
// That used to live in the wide event alone; the person reviewing the
|
||||
// rows is the one who needs to know the rest of the report was not
|
||||
// read.
|
||||
if (raster.pageCount > raster.images.length) {
|
||||
pageCoverage = {
|
||||
read: raster.images.length,
|
||||
total: raster.pageCount,
|
||||
};
|
||||
annotate({
|
||||
action: { name: "labs.ocr.pagesCapped" },
|
||||
meta: { read: raster.images.length, total: raster.pageCount },
|
||||
});
|
||||
}
|
||||
}
|
||||
} else {
|
||||
images = [{ mediaType: mime, dataBase64: buffer.toString("base64") }];
|
||||
documents = [];
|
||||
}
|
||||
|
||||
// 8. Run the extraction. The actual token spend reconciles the reservation.
|
||||
let actualTokens = 0;
|
||||
try {
|
||||
const result = await runOcrExtraction({
|
||||
userId,
|
||||
provider: pick.entry.instance,
|
||||
providerType: pick.providerType,
|
||||
images,
|
||||
documents,
|
||||
});
|
||||
// The orchestration does not surface token counts; reconcile against the
|
||||
// reserved estimate as the spend ceiling (the provider already billed it).
|
||||
actualTokens = reservation.reserved;
|
||||
await reconcileSpend(
|
||||
userId,
|
||||
reservation.reserved,
|
||||
actualTokens,
|
||||
dateKey,
|
||||
0,
|
||||
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
|
||||
);
|
||||
return apiSuccess(pageCoverage ? { ...result, pageCoverage } : result);
|
||||
} catch (err) {
|
||||
// Provider/extraction failure — the call may still have burned tokens, so
|
||||
// reconcile against the reserved estimate rather than refunding in full.
|
||||
await reconcileSpend(
|
||||
userId,
|
||||
reservation.reserved,
|
||||
actualTokens,
|
||||
dateKey,
|
||||
0,
|
||||
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
|
||||
);
|
||||
if (err instanceof OcrExtractError) {
|
||||
return apiError("Couldn't read the report. Try a clearer photo.", 422, {
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
});
|
||||
}
|
||||
annotate({
|
||||
action: { name: "labs.ocr.extractFailed" },
|
||||
meta: providerFailureMeta(err),
|
||||
});
|
||||
return apiError(
|
||||
"The configured AI provider could not process this report. Check the provider configuration and retry.",
|
||||
502,
|
||||
{ errorCode: "labs.ocr.extractFailed" },
|
||||
);
|
||||
}
|
||||
// 7. Queue the read. The worker renders a PDF for a provider that cannot
|
||||
// read one natively, runs the extraction, and settles the reservation.
|
||||
return await startAiRun({
|
||||
userId,
|
||||
kind: "LABS_OCR_EXTRACT",
|
||||
params: {
|
||||
mode: "vision",
|
||||
mime,
|
||||
budget: {
|
||||
reserved: reservation.reserved,
|
||||
owner: reservation.owner,
|
||||
dateKey,
|
||||
},
|
||||
},
|
||||
input: buffer,
|
||||
refundSlot: () => refundLabsOcrSlot(userId),
|
||||
});
|
||||
} catch (err) {
|
||||
// A guard threw after the reservation (e.g. consent races) — refund fully.
|
||||
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
|
||||
|
||||
@@ -1,21 +1,28 @@
|
||||
/**
|
||||
* The browser waits as long as the server may.
|
||||
* The browser waits as long as the server may, and no longer has to wait at
|
||||
* all for a read that runs in the background.
|
||||
*
|
||||
* Document reads and lab OCR used to abort on fixed windows (120 s, 90 s, and
|
||||
* the 15 s default on the lab text path) while the server, for a slow local
|
||||
* model, was allowed far longer. The windows now derive from
|
||||
* `ai.provider.responseTimeoutMs` on the account payload: the per-call value
|
||||
* times the model calls one request can make, plus a 30 s margin.
|
||||
* model, was allowed far longer. A synchronous route still derives its window
|
||||
* from `ai.provider.responseTimeoutMs` (per call, times the calls, plus a
|
||||
* 30 s margin). Since v1.40 "Read with AI" and the lab scan only queue the
|
||||
* read: the request gets a short ceiling for the upload and the enqueue, and
|
||||
* the result comes from the run.
|
||||
*/
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { QueryClient } from "@tanstack/react-query";
|
||||
|
||||
const apiPost = vi.fn();
|
||||
const apiFetch = vi.fn();
|
||||
vi.mock("@/lib/api/api-fetch", () => ({
|
||||
const apiGet = vi.fn();
|
||||
vi.mock("@/lib/api/api-fetch", async (importOriginal) => ({
|
||||
ApiError: (await importOriginal<typeof import("@/lib/api/api-fetch")>())
|
||||
.ApiError,
|
||||
apiPost: (...a: unknown[]) => apiPost(...a),
|
||||
apiFetch: (...a: unknown[]) => apiFetch(...a),
|
||||
apiFetchRaw: vi.fn(),
|
||||
apiGet: vi.fn(),
|
||||
apiGet: (...a: unknown[]) => apiGet(...a),
|
||||
apiPatch: vi.fn(),
|
||||
}));
|
||||
vi.mock("@/lib/labs/local-ocr", () => ({
|
||||
@@ -36,29 +43,55 @@ const target = {
|
||||
servingClass: "attachment" as const,
|
||||
};
|
||||
|
||||
const accepted = { runId: "run1", status: "QUEUED", pollAfterMs: 0 };
|
||||
const indexResult = {
|
||||
documentId: "d1",
|
||||
indexed: true,
|
||||
tokenCount: 12,
|
||||
labFactsStaged: 0,
|
||||
};
|
||||
|
||||
let timeout: ReturnType<typeof vi.spyOn>;
|
||||
let queryClient: QueryClient;
|
||||
|
||||
beforeEach(() => {
|
||||
apiPost.mockReset().mockResolvedValue({});
|
||||
apiFetch.mockReset().mockResolvedValue({});
|
||||
apiPost.mockReset().mockResolvedValue(accepted);
|
||||
apiFetch.mockReset().mockResolvedValue(accepted);
|
||||
apiGet.mockReset().mockResolvedValue({
|
||||
id: "run1",
|
||||
status: "SUCCEEDED",
|
||||
result: indexResult,
|
||||
error: null,
|
||||
retryAfterMs: null,
|
||||
});
|
||||
timeout = vi.spyOn(AbortSignal, "timeout");
|
||||
queryClient = new QueryClient();
|
||||
vi.useFakeTimers({ shouldAdvanceTime: true });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
timeout.mockRestore();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
describe("document AI requests", () => {
|
||||
it("give Read with AI three model calls under the published setting", async () => {
|
||||
await runDocumentIndex({
|
||||
it("queue Read with AI and resolve the run's result", async () => {
|
||||
const result = await runDocumentIndex({
|
||||
mode: "vision",
|
||||
target,
|
||||
responseTimeoutMs: 300_000,
|
||||
background: { queryClient },
|
||||
});
|
||||
expect(timeout).toHaveBeenCalledWith(3 * 300_000 + 30_000);
|
||||
expect(result).toEqual(indexResult);
|
||||
expect(timeout).toHaveBeenCalledWith(60_000);
|
||||
expect(timeout).not.toHaveBeenCalledWith(3 * 300_000 + 30_000);
|
||||
const init = apiPost.mock.calls[0][2] as RequestInit;
|
||||
expect(new Headers(init.headers).get("Prefer")).toBe("respond-async");
|
||||
expect(apiGet).toHaveBeenCalledWith("/api/ai-runs/run1");
|
||||
});
|
||||
|
||||
it("give a one-call route one call plus the margin", async () => {
|
||||
it("give a one-call synchronous route one call plus the margin", async () => {
|
||||
apiPost.mockResolvedValue({ title: "x" });
|
||||
await runDocumentAi({
|
||||
path: "/api/documents/inbound/d1/summary?mode=summary",
|
||||
mode: "vision",
|
||||
@@ -67,28 +100,54 @@ describe("document AI requests", () => {
|
||||
modelCalls: 1,
|
||||
});
|
||||
expect(timeout).toHaveBeenCalledWith(330_000);
|
||||
});
|
||||
|
||||
it("follow the server default when nothing is set", async () => {
|
||||
await runDocumentIndex({
|
||||
mode: "vision",
|
||||
target,
|
||||
responseTimeoutMs: 60_000,
|
||||
});
|
||||
expect(timeout).toHaveBeenCalledWith(210_000);
|
||||
const init = apiPost.mock.calls[0][2] as RequestInit;
|
||||
expect(new Headers(init.headers ?? {}).get("Prefer")).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("lab OCR requests", () => {
|
||||
it("vision upload waits for the read and its retry", async () => {
|
||||
await postOcrExtract({ file: new File(["x"], "r.png") }, 300_000);
|
||||
expect(timeout).toHaveBeenCalledWith(630_000);
|
||||
it("vision upload queues the read and polls the rows", async () => {
|
||||
const rows = { reportDate: null, providerType: "anthropic", rows: [] };
|
||||
apiGet.mockResolvedValue({
|
||||
id: "run1",
|
||||
status: "SUCCEEDED",
|
||||
result: rows,
|
||||
error: null,
|
||||
retryAfterMs: null,
|
||||
});
|
||||
const result = await postOcrExtract(
|
||||
{ file: new File(["x"], "r.png") },
|
||||
queryClient,
|
||||
);
|
||||
expect(result).toEqual(rows);
|
||||
expect(timeout).toHaveBeenCalledWith(60_000);
|
||||
expect(apiFetch.mock.calls[0][1].signal).toBeInstanceOf(AbortSignal);
|
||||
});
|
||||
|
||||
it("text mode no longer rides the 15 s default", async () => {
|
||||
await postOcrExtract({ text: "Hb 14" }, 300_000);
|
||||
expect(timeout).toHaveBeenCalledWith(630_000);
|
||||
await postOcrExtract({ text: "Hb 14" }, queryClient);
|
||||
expect(timeout).toHaveBeenCalledWith(60_000);
|
||||
expect(apiPost.mock.calls[0][2].signal).toBeInstanceOf(AbortSignal);
|
||||
});
|
||||
|
||||
it("a failed run throws what the synchronous route would have", async () => {
|
||||
apiGet.mockResolvedValue({
|
||||
id: "run1",
|
||||
status: "FAILED",
|
||||
result: null,
|
||||
error: {
|
||||
code: "labs.ocr.extractFailed",
|
||||
message: "Couldn't read the report.",
|
||||
status: 422,
|
||||
},
|
||||
retryAfterMs: null,
|
||||
});
|
||||
await expect(
|
||||
postOcrExtract({ text: "Hb 14" }, queryClient),
|
||||
).rejects.toMatchObject({
|
||||
name: "ApiError",
|
||||
status: 422,
|
||||
meta: { errorCode: "labs.ocr.extractFailed" },
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -33,6 +33,7 @@
|
||||
import { FileText, ScanText, WandSparkles } from "lucide-react";
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import type { AiRunPhase } from "@/hooks/use-ai-run";
|
||||
import { useTranslations } from "@/lib/i18n/context";
|
||||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
@@ -77,6 +78,7 @@ export function DocumentAiSection({
|
||||
contentIndexSource,
|
||||
lastIndexOutcome = null,
|
||||
indexPending,
|
||||
indexRunPhase = "idle",
|
||||
onIndex,
|
||||
}: {
|
||||
aiEnabled: boolean;
|
||||
@@ -120,6 +122,8 @@ export function DocumentAiSection({
|
||||
*/
|
||||
lastIndexOutcome?: DocumentIndexOutcomeValue | null;
|
||||
indexPending: boolean;
|
||||
/** v1.40 — where a background read stands, for the calm line under it. */
|
||||
indexRunPhase?: AiRunPhase;
|
||||
onIndex: () => void;
|
||||
}) {
|
||||
const { t } = useTranslations();
|
||||
@@ -248,6 +252,18 @@ export function DocumentAiSection({
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
{indexRunPhase !== "idle" ? (
|
||||
<p
|
||||
role="status"
|
||||
data-slot="document-index-run-phase"
|
||||
className="text-muted-foreground text-xs"
|
||||
>
|
||||
{indexRunPhase === "waitingForWorker"
|
||||
? t("aiRuns.waitingForWorker")
|
||||
: t("aiRuns.backgroundDocument")}
|
||||
</p>
|
||||
) : null}
|
||||
|
||||
{indexOutcomeText ? (
|
||||
<p
|
||||
data-slot="document-index-outcome"
|
||||
|
||||
@@ -15,9 +15,20 @@
|
||||
* payload publishes (the person's response-timeout setting, resolved on the
|
||||
* server) times the model calls the route can make, plus a margin, so the
|
||||
* browser never abandons a read the server is still allowed to finish.
|
||||
*
|
||||
* v1.40 — a route that offers the background path is called with
|
||||
* `Prefer: respond-async` (`background`). The request then only queues the
|
||||
* read and answers 202 quickly, and the result is polled from the run, so no
|
||||
* proxy in front of the server can cut a long read. Such a call carries no
|
||||
* model-time ceiling at all; the run's own deadline (the same setting) ends
|
||||
* it on the server.
|
||||
*/
|
||||
import type { QueryClient } from "@tanstack/react-query";
|
||||
|
||||
import { resolveAiRun, type WaitForAiRunOptions } from "@/hooks/use-ai-run";
|
||||
import { apiFetchRaw, apiPost } from "@/lib/api/api-fetch";
|
||||
import { clientAbortMs } from "@/lib/ai/effective-timeout";
|
||||
import type { AiRunAccepted } from "@/lib/documents/ai-runs/types";
|
||||
import { ocrImageToText, LocalOcrError } from "@/lib/labs/local-ocr";
|
||||
|
||||
/** The transport an AI call uses, resolved from the OCR capability probe. */
|
||||
@@ -67,6 +78,15 @@ async function fetchOriginalAsFile(target: DocumentAiTarget): Promise<File> {
|
||||
});
|
||||
}
|
||||
|
||||
/** How long a queuing request may take: the upload and the enqueue, no model. */
|
||||
const QUEUE_REQUEST_TIMEOUT_MS = 60_000;
|
||||
|
||||
/** The background path of a document AI call. */
|
||||
export interface DocumentAiBackground extends WaitForAiRunOptions {
|
||||
queryClient: QueryClient;
|
||||
onQueued?: () => void;
|
||||
}
|
||||
|
||||
/** Run one document AI call over the chosen transport. */
|
||||
export async function runDocumentAi<T>(opts: {
|
||||
path: string;
|
||||
@@ -80,8 +100,20 @@ export async function runDocumentAi<T>(opts: {
|
||||
* staging read with its one corrective retry).
|
||||
*/
|
||||
modelCalls: number;
|
||||
/** Queue the read and poll it (the route must offer `Prefer: respond-async`). */
|
||||
background?: DocumentAiBackground;
|
||||
}): Promise<T> {
|
||||
const timeoutMs = clientAbortMs(opts.responseTimeoutMs, opts.modelCalls);
|
||||
const background = opts.background;
|
||||
const timeoutMs = background
|
||||
? QUEUE_REQUEST_TIMEOUT_MS
|
||||
: clientAbortMs(opts.responseTimeoutMs, opts.modelCalls);
|
||||
const headers: HeadersInit | undefined = background
|
||||
? { Prefer: "respond-async" }
|
||||
: undefined;
|
||||
const settle = async (data: T | AiRunAccepted): Promise<T> =>
|
||||
background
|
||||
? resolveAiRun<T>(background.queryClient, data, background)
|
||||
: (data as T);
|
||||
if (opts.mode === "text") {
|
||||
if (
|
||||
opts.target.servingClass !== "inline" ||
|
||||
@@ -99,16 +131,21 @@ export async function runDocumentAi<T>(opts: {
|
||||
? new DocumentAssistClientError("ocr")
|
||||
: err;
|
||||
}
|
||||
return apiPost<T>(
|
||||
opts.path,
|
||||
{ mode: "text", text },
|
||||
{ signal: AbortSignal.timeout(timeoutMs) },
|
||||
return settle(
|
||||
await apiPost<T | AiRunAccepted>(
|
||||
opts.path,
|
||||
{ mode: "text", text },
|
||||
{ signal: AbortSignal.timeout(timeoutMs), headers },
|
||||
),
|
||||
);
|
||||
}
|
||||
// Vision: empty body → the route dispatches to its vision handler.
|
||||
return apiPost<T>(opts.path, undefined, {
|
||||
signal: AbortSignal.timeout(timeoutMs),
|
||||
});
|
||||
return settle(
|
||||
await apiPost<T | AiRunAccepted>(opts.path, undefined, {
|
||||
signal: AbortSignal.timeout(timeoutMs),
|
||||
headers,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
/** The index route's result, including the lab-staging continuation. */
|
||||
@@ -123,11 +160,16 @@ export interface DocumentIndexResult {
|
||||
labFactsStaged: number;
|
||||
}
|
||||
|
||||
/** Populate / refresh one document's content index over the chosen transport. */
|
||||
/**
|
||||
* Populate / refresh one document's content index over the chosen transport.
|
||||
* Always in the background: the read and the lab staging after it are up to
|
||||
* three model calls, which no proxy timeout should be able to cut.
|
||||
*/
|
||||
export function runDocumentIndex(opts: {
|
||||
mode: DocumentAiMode;
|
||||
target: DocumentAiTarget;
|
||||
responseTimeoutMs: number;
|
||||
background: DocumentAiBackground;
|
||||
}): Promise<DocumentIndexResult> {
|
||||
return runDocumentAi<DocumentIndexResult>({
|
||||
path: `/api/documents/inbound/${opts.target.documentId}/index`,
|
||||
@@ -135,5 +177,6 @@ export function runDocumentIndex(opts: {
|
||||
target: opts.target,
|
||||
responseTimeoutMs: opts.responseTimeoutMs,
|
||||
modelCalls: 3,
|
||||
background: opts.background,
|
||||
});
|
||||
}
|
||||
@@ -937,6 +937,7 @@ export function DocumentDetailSheet({
|
||||
contentIndexSource={doc.contentIndexSource}
|
||||
lastIndexOutcome={doc.lastIndexOutcome}
|
||||
indexPending={indexDoc.isPending}
|
||||
indexRunPhase={indexDoc.runPhase}
|
||||
onIndex={runIndex}
|
||||
/>
|
||||
)}
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
import { useMutation, useQueryClient } from "@tanstack/react-query";
|
||||
|
||||
import { useAiProviderState } from "@/hooks/use-ai-capability";
|
||||
import { useAiRunPhase } from "@/hooks/use-ai-run";
|
||||
import { apiPost } from "@/lib/api/api-fetch";
|
||||
import { invalidateKeys, queryKeys } from "@/lib/query-keys";
|
||||
|
||||
@@ -27,21 +28,34 @@ import {
|
||||
type DocumentIndexResult,
|
||||
} from "./document-ai-transport";
|
||||
|
||||
/** Index / re-index one document for content search. */
|
||||
/**
|
||||
* Index / re-index one document for content search. The read runs in the
|
||||
* background (v1.40); `runPhase` says so while it does, so the screen can tell
|
||||
* the person they may leave: the index and any staged lab values land on the
|
||||
* document whether or not anyone is still watching.
|
||||
*/
|
||||
export function useIndexDocument() {
|
||||
const queryClient = useQueryClient();
|
||||
const { responseTimeoutMs } = useAiProviderState();
|
||||
return useMutation<
|
||||
const { phase, onQueued, onProgress, reset } = useAiRunPhase();
|
||||
const mutation = useMutation<
|
||||
DocumentIndexResult,
|
||||
Error,
|
||||
{ mode: DocumentAiMode; target: DocumentAiTarget }
|
||||
>({
|
||||
mutationFn: ({ mode, target }) =>
|
||||
runDocumentIndex({ mode, target, responseTimeoutMs }),
|
||||
runDocumentIndex({
|
||||
mode,
|
||||
target,
|
||||
responseTimeoutMs,
|
||||
background: { queryClient, onQueued, onProgress },
|
||||
}),
|
||||
onSettled: reset,
|
||||
onSuccess: () => {
|
||||
void invalidateKeys(queryClient, [queryKeys.documents()]);
|
||||
},
|
||||
});
|
||||
return { ...mutation, runPhase: mutation.isPending ? phase : "idle" };
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -207,6 +207,10 @@ export function documentAiErrorKey(err: unknown): string {
|
||||
return "documents.assist.errorProvider";
|
||||
case "documents.inbound.notIndexed":
|
||||
return "documents.assist.errorNotIndexed";
|
||||
case "aiRuns.workerUnavailable":
|
||||
return "aiRuns.errorWorkerUnavailable";
|
||||
case "aiRuns.timedOut":
|
||||
return "aiRuns.errorTimedOut";
|
||||
case "consent.ai.required":
|
||||
return "documents.assist.errorConsent";
|
||||
case "ai.provider.none":
|
||||
|
||||
@@ -82,6 +82,10 @@ export function extractErrorMessage(
|
||||
return t("labs.ocr.consentRequired");
|
||||
case "ai.provider.none":
|
||||
return t("labs.ocr.providerUnsupported");
|
||||
case "aiRuns.workerUnavailable":
|
||||
return t("aiRuns.errorWorkerUnavailable");
|
||||
case "aiRuns.timedOut":
|
||||
return t("aiRuns.errorTimedOut");
|
||||
default:
|
||||
break;
|
||||
}
|
||||
@@ -332,6 +336,17 @@ export function OcrReviewDialog({
|
||||
? t("labs.ocr.readingOnDevice")
|
||||
: t("labs.ocr.extracting")}
|
||||
</span>
|
||||
{extract.runPhase !== "idle" ? (
|
||||
<span
|
||||
role="status"
|
||||
data-slot="ocr-run-phase"
|
||||
className="text-muted-foreground text-xs"
|
||||
>
|
||||
{extract.runPhase === "waitingForWorker"
|
||||
? t("aiRuns.waitingForWorker")
|
||||
: t("aiRuns.backgroundScan")}
|
||||
</span>
|
||||
) : null}
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
|
||||
@@ -6,17 +6,26 @@
|
||||
* - `useOcrCapability()` — the cheap capability probe that decides whether the
|
||||
* "Scan a report" affordance shows.
|
||||
* - `useOcrExtract()` — uploads the photo / PDF and returns the proposed rows.
|
||||
* A vision call is slow, so it opts out of the default 15 s fetch timeout in
|
||||
* favour of a generous 90 s window.
|
||||
* Since v1.40 the route only queues the read (202) and the rows come from
|
||||
* the background run, so no proxy timeout can cut a slow model.
|
||||
* - `useOcrCommit()` — writes the user-confirmed rows and invalidates the
|
||||
* labs + biomarker query keys.
|
||||
*/
|
||||
import type { AiCapabilityState } from "@/lib/ai/capabilities/types";
|
||||
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
|
||||
import {
|
||||
useMutation,
|
||||
useQuery,
|
||||
useQueryClient,
|
||||
type QueryClient,
|
||||
} from "@tanstack/react-query";
|
||||
|
||||
import { useAiProviderState } from "@/hooks/use-ai-capability";
|
||||
import { clientAbortMs } from "@/lib/ai/effective-timeout";
|
||||
import {
|
||||
resolveAiRun,
|
||||
useAiRunPhase,
|
||||
type WaitForAiRunOptions,
|
||||
} from "@/hooks/use-ai-run";
|
||||
import { apiFetch, apiGet, apiPatch, apiPost } from "@/lib/api/api-fetch";
|
||||
import type { AiRunAccepted } from "@/lib/documents/ai-runs/types";
|
||||
import { ocrImageToText } from "@/lib/labs/local-ocr";
|
||||
import {
|
||||
aiInputDependentKeys,
|
||||
@@ -94,49 +103,60 @@ export function useOcrCapability(enabled: boolean) {
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* The most sequential model calls one extract request makes: the read, and one
|
||||
* corrective retry when the reply does not parse.
|
||||
*/
|
||||
const OCR_EXTRACT_MODEL_CALLS = 2;
|
||||
/** How long the queuing request may take: the upload and the enqueue, no model. */
|
||||
const OCR_QUEUE_REQUEST_TIMEOUT_MS = 60_000;
|
||||
|
||||
/**
|
||||
* POST one lab-report extraction. Vision mode uploads the file; text mode posts
|
||||
* only the text the browser OCR'd. Either way a model reads it, so the request
|
||||
* waits as long as the server may: the person's response-timeout setting as
|
||||
* the account payload publishes it (`ai.provider.responseTimeoutMs`), times
|
||||
* the calls one request can make, plus a margin. The 15 s default used to cut
|
||||
* the text path off long before any model answered.
|
||||
* POST one lab-report extraction and follow the background run to its rows.
|
||||
* Vision mode uploads the file; text mode posts only the text the browser
|
||||
* OCR'd. The request itself only queues the read, so it carries a short
|
||||
* ceiling; the run's own deadline (the person's AI response time) ends a slow
|
||||
* read on the server.
|
||||
*/
|
||||
export function postOcrExtract(
|
||||
export async function postOcrExtract(
|
||||
input: { file: File } | { text: string },
|
||||
responseTimeoutMs: number,
|
||||
queryClient: QueryClient,
|
||||
progress: WaitForAiRunOptions & { onQueued?: () => void } = {},
|
||||
): Promise<OcrExtractResponseDto> {
|
||||
const signal = AbortSignal.timeout(
|
||||
clientAbortMs(responseTimeoutMs, OCR_EXTRACT_MODEL_CALLS),
|
||||
);
|
||||
const signal = AbortSignal.timeout(OCR_QUEUE_REQUEST_TIMEOUT_MS);
|
||||
let accepted: OcrExtractResponseDto | AiRunAccepted;
|
||||
if ("text" in input) {
|
||||
return apiPost<OcrExtractResponseDto>(
|
||||
accepted = await apiPost<OcrExtractResponseDto | AiRunAccepted>(
|
||||
"/api/labs/ocr/extract",
|
||||
{ mode: "text", text: input.text },
|
||||
{ signal },
|
||||
);
|
||||
} else {
|
||||
const form = new FormData();
|
||||
form.append("file", input.file);
|
||||
accepted = await apiFetch<OcrExtractResponseDto | AiRunAccepted>(
|
||||
"/api/labs/ocr/extract",
|
||||
{ method: "POST", body: form, signal },
|
||||
);
|
||||
}
|
||||
const form = new FormData();
|
||||
form.append("file", input.file);
|
||||
return apiFetch<OcrExtractResponseDto>("/api/labs/ocr/extract", {
|
||||
method: "POST",
|
||||
body: form,
|
||||
signal,
|
||||
return resolveAiRun<OcrExtractResponseDto>(queryClient, accepted, progress);
|
||||
}
|
||||
|
||||
/**
|
||||
* The extract mutation for one mode, plus the run phase the dialog shows while
|
||||
* the read runs in the background.
|
||||
*/
|
||||
function useOcrExtractMutation(
|
||||
read: (file: File) => Promise<{ file: File } | { text: string }>,
|
||||
) {
|
||||
const queryClient = useQueryClient();
|
||||
const { phase, onQueued, onProgress, reset } = useAiRunPhase();
|
||||
const mutation = useMutation<OcrExtractResponseDto, Error, File>({
|
||||
mutationFn: async (file: File) =>
|
||||
postOcrExtract(await read(file), queryClient, { onQueued, onProgress }),
|
||||
onSettled: reset,
|
||||
});
|
||||
return { ...mutation, runPhase: mutation.isPending ? phase : "idle" };
|
||||
}
|
||||
|
||||
/** Upload + extract (VISION mode). Resolves with the proposed review rows. */
|
||||
export function useOcrExtract() {
|
||||
const { responseTimeoutMs } = useAiProviderState();
|
||||
return useMutation<OcrExtractResponseDto, Error, File>({
|
||||
mutationFn: (file: File) => postOcrExtract({ file }, responseTimeoutMs),
|
||||
});
|
||||
return useOcrExtractMutation(async (file) => ({ file }));
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -146,11 +166,9 @@ export function useOcrExtract() {
|
||||
* is shared verbatim.
|
||||
*/
|
||||
export function useOcrTextExtract() {
|
||||
const { responseTimeoutMs } = useAiProviderState();
|
||||
return useMutation<OcrExtractResponseDto, Error, File>({
|
||||
mutationFn: async (file: File) =>
|
||||
postOcrExtract({ text: await ocrImageToText(file) }, responseTimeoutMs),
|
||||
});
|
||||
return useOcrExtractMutation(async (file) => ({
|
||||
text: await ocrImageToText(file),
|
||||
}));
|
||||
}
|
||||
|
||||
/** The local-OCR opt-in preference (read + toggle). */
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
/**
|
||||
* Following a background run from the browser (v1.40): a synchronous body
|
||||
* passes through untouched, a queued run is polled through the key factory
|
||||
* until it ends, the poll interval follows the server's `retryAfterMs`, the
|
||||
* progress callback sees the queue wait (the "waiting for the background
|
||||
* worker" hint), a failure throws the synchronous route's `ApiError`, and the
|
||||
* answer does not stay in the cache.
|
||||
*/
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { QueryClient } from "@tanstack/react-query";
|
||||
|
||||
const apiGet = vi.fn();
|
||||
vi.mock("@/lib/api/api-fetch", async (importOriginal) => ({
|
||||
ApiError: (await importOriginal<typeof import("@/lib/api/api-fetch")>())
|
||||
.ApiError,
|
||||
apiGet: (...a: unknown[]) => apiGet(...a),
|
||||
}));
|
||||
|
||||
import {
|
||||
AI_RUN_WORKER_HINT_MS,
|
||||
isAiRunAccepted,
|
||||
resolveAiRun,
|
||||
waitForAiRun,
|
||||
} from "@/hooks/use-ai-run";
|
||||
import { queryKeys } from "@/lib/query-keys";
|
||||
|
||||
const accepted = { runId: "r1", status: "QUEUED" as const, pollAfterMs: 1500 };
|
||||
|
||||
function run(overrides: Record<string, unknown>) {
|
||||
return {
|
||||
id: "r1",
|
||||
kind: "DOCUMENT_INDEX",
|
||||
documentId: "d1",
|
||||
status: "QUEUED",
|
||||
result: null,
|
||||
error: null,
|
||||
createdAt: "2027-01-15T10:00:00.000Z",
|
||||
startedAt: null,
|
||||
finishedAt: null,
|
||||
queuedForMs: 0,
|
||||
retryAfterMs: 1500,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
let queryClient: QueryClient;
|
||||
const waits: number[] = [];
|
||||
const wait = async (ms: number) => {
|
||||
waits.push(ms);
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
apiGet.mockReset();
|
||||
waits.length = 0;
|
||||
queryClient = new QueryClient({
|
||||
defaultOptions: { queries: { retry: false } },
|
||||
});
|
||||
});
|
||||
|
||||
describe("isAiRunAccepted", () => {
|
||||
it("tells a queued run from a result", () => {
|
||||
expect(isAiRunAccepted(accepted)).toBe(true);
|
||||
expect(isAiRunAccepted({ indexed: true, tokenCount: 1 })).toBe(false);
|
||||
expect(isAiRunAccepted(null)).toBe(false);
|
||||
expect(isAiRunAccepted({ runId: "r1", status: "SUCCEEDED" })).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("waitForAiRun", () => {
|
||||
it("polls until the run succeeds and backs off as the server says", async () => {
|
||||
apiGet
|
||||
.mockResolvedValueOnce(run({ status: "QUEUED", queuedForMs: 1000 }))
|
||||
.mockResolvedValueOnce(
|
||||
run({ status: "RUNNING", queuedForMs: null, retryAfterMs: 4000 }),
|
||||
)
|
||||
.mockResolvedValueOnce(
|
||||
run({
|
||||
status: "SUCCEEDED",
|
||||
result: { indexed: true },
|
||||
retryAfterMs: null,
|
||||
}),
|
||||
);
|
||||
const progress = vi.fn();
|
||||
const result = await waitForAiRun(queryClient, accepted, {
|
||||
wait,
|
||||
onProgress: progress,
|
||||
});
|
||||
expect(result).toEqual({ indexed: true });
|
||||
expect(waits).toEqual([1500, 1500, 4000]);
|
||||
expect(apiGet).toHaveBeenCalledWith("/api/ai-runs/r1");
|
||||
expect(progress).toHaveBeenCalledTimes(2);
|
||||
expect(queryClient.getQueryData(queryKeys.aiRun("r1"))).toBeUndefined();
|
||||
});
|
||||
|
||||
it("throws the failure the synchronous route would have answered with", async () => {
|
||||
apiGet.mockResolvedValueOnce(
|
||||
run({
|
||||
status: "FAILED",
|
||||
error: {
|
||||
code: "consent.ai.required",
|
||||
message: "AI consent is required for this feature",
|
||||
status: 403,
|
||||
},
|
||||
retryAfterMs: null,
|
||||
}),
|
||||
);
|
||||
await expect(
|
||||
waitForAiRun(queryClient, accepted, { wait }),
|
||||
).rejects.toMatchObject({
|
||||
name: "ApiError",
|
||||
status: 403,
|
||||
meta: { errorCode: "consent.ai.required" },
|
||||
});
|
||||
expect(queryClient.getQueryData(queryKeys.aiRun("r1"))).toBeUndefined();
|
||||
});
|
||||
|
||||
it("stops when the caller aborts", async () => {
|
||||
const controller = new AbortController();
|
||||
controller.abort(new Error("left"));
|
||||
await expect(
|
||||
waitForAiRun(queryClient, accepted, { signal: controller.signal }),
|
||||
).rejects.toThrow("left");
|
||||
expect(apiGet).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("resolveAiRun", () => {
|
||||
it("passes a synchronous body straight through", async () => {
|
||||
const onQueued = vi.fn();
|
||||
expect(
|
||||
await resolveAiRun(queryClient, { indexed: true }, { wait, onQueued }),
|
||||
).toEqual({ indexed: true });
|
||||
expect(onQueued).not.toHaveBeenCalled();
|
||||
expect(apiGet).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("reports the queue and follows the run", async () => {
|
||||
apiGet.mockResolvedValueOnce(
|
||||
run({ status: "SUCCEEDED", result: { rows: [] }, retryAfterMs: null }),
|
||||
);
|
||||
const onQueued = vi.fn();
|
||||
expect(
|
||||
await resolveAiRun(queryClient, accepted, { wait, onQueued }),
|
||||
).toEqual({ rows: [] });
|
||||
expect(onQueued).toHaveBeenCalledOnce();
|
||||
});
|
||||
});
|
||||
|
||||
describe("the worker hint threshold", () => {
|
||||
it("is long enough not to flash on an ordinary queue wait", () => {
|
||||
expect(AI_RUN_WORKER_HINT_MS).toBeGreaterThanOrEqual(30_000);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,136 @@
|
||||
"use client";
|
||||
|
||||
/**
|
||||
* v1.40 — follow a background document AI run to its end.
|
||||
*
|
||||
* A route that queues a read answers 202 with `{ runId, status, pollAfterMs }`.
|
||||
* `resolveAiRun` takes whatever the route answered: a synchronous body passes
|
||||
* straight through, a queued run is polled through the query cache
|
||||
* (`queryKeys.aiRun`) until it ends. The result resolves exactly like the
|
||||
* synchronous body did; a failed run throws the same `ApiError` (message,
|
||||
* status, `meta.errorCode`) the synchronous route would have, so every
|
||||
* existing error mapping keeps working.
|
||||
*
|
||||
* `useAiRunPhase` is the calm state a screen shows while a read runs in the
|
||||
* background, and the hint when no worker has picked it up yet.
|
||||
*/
|
||||
import { useCallback, useState } from "react";
|
||||
import type { QueryClient } from "@tanstack/react-query";
|
||||
|
||||
import { ApiError, apiGet } from "@/lib/api/api-fetch";
|
||||
import type { AiRunAccepted, AiRunDto } from "@/lib/documents/ai-runs/types";
|
||||
import { queryKeys } from "@/lib/query-keys";
|
||||
|
||||
/** A queued run has waited this long: say the background worker may be down. */
|
||||
export const AI_RUN_WORKER_HINT_MS = 30_000;
|
||||
|
||||
const DEFAULT_POLL_MS = 1500;
|
||||
|
||||
/** Whether a route's 2xx body is a queued run rather than its result. */
|
||||
export function isAiRunAccepted(data: unknown): data is AiRunAccepted {
|
||||
return (
|
||||
typeof data === "object" &&
|
||||
data !== null &&
|
||||
typeof (data as { runId?: unknown }).runId === "string" &&
|
||||
(data as { status?: unknown }).status === "QUEUED"
|
||||
);
|
||||
}
|
||||
|
||||
function sleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (signal?.aborted) {
|
||||
reject(signal.reason);
|
||||
return;
|
||||
}
|
||||
const timer = setTimeout(() => {
|
||||
signal?.removeEventListener("abort", onAbort);
|
||||
resolve();
|
||||
}, ms);
|
||||
const onAbort = () => {
|
||||
clearTimeout(timer);
|
||||
reject(signal?.reason);
|
||||
};
|
||||
signal?.addEventListener("abort", onAbort, { once: true });
|
||||
});
|
||||
}
|
||||
|
||||
export interface WaitForAiRunOptions {
|
||||
/** Called with every poll answer while the run has not ended. */
|
||||
onProgress?: (run: AiRunDto) => void;
|
||||
signal?: AbortSignal;
|
||||
/** Injectable for tests. */
|
||||
wait?: (ms: number, signal?: AbortSignal) => Promise<void>;
|
||||
}
|
||||
|
||||
/** Poll one run until it ends; resolve its result or throw its failure. */
|
||||
export async function waitForAiRun<T>(
|
||||
queryClient: QueryClient,
|
||||
accepted: AiRunAccepted,
|
||||
options: WaitForAiRunOptions = {},
|
||||
): Promise<T> {
|
||||
const wait = options.wait ?? sleep;
|
||||
const key = queryKeys.aiRun(accepted.runId);
|
||||
let delay = accepted.pollAfterMs > 0 ? accepted.pollAfterMs : DEFAULT_POLL_MS;
|
||||
try {
|
||||
for (;;) {
|
||||
await wait(delay, options.signal);
|
||||
const run = await queryClient.fetchQuery({
|
||||
queryKey: key,
|
||||
queryFn: () =>
|
||||
apiGet<AiRunDto>(
|
||||
`/api/ai-runs/${encodeURIComponent(accepted.runId)}`,
|
||||
),
|
||||
staleTime: 0,
|
||||
retry: 2,
|
||||
});
|
||||
if (run.status === "SUCCEEDED") return run.result as T;
|
||||
if (run.status === "FAILED") {
|
||||
const error = run.error ?? {
|
||||
code: "aiRuns.failed",
|
||||
message: "The background read failed.",
|
||||
status: 500,
|
||||
};
|
||||
throw new ApiError(error.message, error.status, {
|
||||
errorCode: error.code,
|
||||
});
|
||||
}
|
||||
options.onProgress?.(run);
|
||||
delay = run.retryAfterMs ?? DEFAULT_POLL_MS;
|
||||
}
|
||||
} finally {
|
||||
// The answer has been handed to the caller; it does not stay in memory.
|
||||
queryClient.removeQueries({ queryKey: key });
|
||||
}
|
||||
}
|
||||
|
||||
/** A route's 2xx body, followed to its result when it is a queued run. */
|
||||
export async function resolveAiRun<T>(
|
||||
queryClient: QueryClient,
|
||||
data: T | AiRunAccepted,
|
||||
options: WaitForAiRunOptions & { onQueued?: () => void } = {},
|
||||
): Promise<T> {
|
||||
if (!isAiRunAccepted(data)) return data;
|
||||
options.onQueued?.();
|
||||
return waitForAiRun<T>(queryClient, data, options);
|
||||
}
|
||||
|
||||
export type AiRunPhase = "idle" | "background" | "waitingForWorker";
|
||||
|
||||
/**
|
||||
* The phase a screen renders while a read runs in the background: `background`
|
||||
* once the route queued it, `waitingForWorker` once it has sat in the queue
|
||||
* past {@link AI_RUN_WORKER_HINT_MS}.
|
||||
*/
|
||||
export function useAiRunPhase() {
|
||||
const [phase, setPhase] = useState<AiRunPhase>("idle");
|
||||
const onQueued = useCallback(() => setPhase("background"), []);
|
||||
const onProgress = useCallback((run: AiRunDto) => {
|
||||
setPhase(
|
||||
run.status === "QUEUED" && (run.queuedForMs ?? 0) >= AI_RUN_WORKER_HINT_MS
|
||||
? "waitingForWorker"
|
||||
: "background",
|
||||
);
|
||||
}, []);
|
||||
const reset = useCallback(() => setPhase("idle"), []);
|
||||
return { phase, onQueued, onProgress, reset };
|
||||
}
|
||||
@@ -53,6 +53,13 @@ export const BACKUP_CHUNK_AAD = "healthlog/data-backup-chunk/v1";
|
||||
*/
|
||||
export const WORKOUT_ROUTE_GEOMETRY_AAD = "healthlog/workout-route-geometry/v1";
|
||||
|
||||
/**
|
||||
* The associated-data labels a background document AI run's input and result
|
||||
* are sealed with, so neither opens as the other or as anything else.
|
||||
*/
|
||||
export const DOCUMENT_AI_RUN_INPUT_AAD = "healthlog/document-ai-run-input/v1";
|
||||
export const DOCUMENT_AI_RUN_RESULT_AAD = "healthlog/document-ai-run-result/v1";
|
||||
|
||||
export interface EncryptedColumn {
|
||||
/** Prisma model name (PascalCase, as declared in schema.prisma). */
|
||||
readonly model: string;
|
||||
@@ -439,6 +446,30 @@ export const ENCRYPTED_COLUMNS: readonly EncryptedColumn[] = [
|
||||
// is PHI; rotation re-encrypts it on the standard Bytes walk.
|
||||
{ model: "DocumentThumbnail", field: "thumbnailEncrypted", kind: "bytes" },
|
||||
|
||||
// ───── v1.40 background document AI runs (Bytes, binary2) ─────
|
||||
// The input a run reads (an uploaded lab scan, or the text the browser read
|
||||
// off it) and the result it produced (model text, lab values). Both live at
|
||||
// most an hour past the run, but a rotation that skipped them would leave a
|
||||
// run unreadable under the retired key. Disposable: a row whose value cannot
|
||||
// be opened is a dead ticket, and deleting it loses nothing the person can
|
||||
// not ask for again.
|
||||
{
|
||||
model: "DocumentAiRun",
|
||||
field: "inputEncrypted",
|
||||
kind: "bytes",
|
||||
codec: "binary2",
|
||||
aad: DOCUMENT_AI_RUN_INPUT_AAD,
|
||||
disposable: true,
|
||||
},
|
||||
{
|
||||
model: "DocumentAiRun",
|
||||
field: "resultEncrypted",
|
||||
kind: "bytes",
|
||||
codec: "binary2",
|
||||
aad: DOCUMENT_AI_RUN_RESULT_AAD,
|
||||
disposable: true,
|
||||
},
|
||||
|
||||
// ───── Whole-account backup blob (String, batched) ─────
|
||||
// The weekly disaster-recovery snapshot and every manually uploaded pack.
|
||||
// Named `data`, so the `*Encrypted` scan never saw it; the value is
|
||||
|
||||
@@ -160,6 +160,10 @@ export const WIPE_MODELS = [
|
||||
// credential to another system behind an account that erased everything.
|
||||
"DocumentSourceConnection",
|
||||
"ExtractedFact",
|
||||
// v1.40 — background AI runs over a document or a lab scan. A queued scan
|
||||
// holds the uploaded report until it is read, and a finished run holds what
|
||||
// was read off it; a child of the document, so it goes before it.
|
||||
"DocumentAiRun",
|
||||
"ImportJob",
|
||||
|
||||
// ── AI ──────────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
/**
|
||||
* The HTTP edges of a background run: which requests asked for it, and the
|
||||
* 202 they get.
|
||||
*/
|
||||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import { AiUnavailableError } from "@/lib/ai/capabilities/refusal";
|
||||
import {
|
||||
acceptedRunResponse,
|
||||
outcomeResponse,
|
||||
prefersRespondAsync,
|
||||
refusalAsFailure,
|
||||
workerUnavailableResponse,
|
||||
} from "@/lib/documents/ai-runs/http";
|
||||
|
||||
const req = (prefer?: string) =>
|
||||
new Request("http://localhost/x", {
|
||||
method: "POST",
|
||||
headers: prefer === undefined ? {} : { Prefer: prefer },
|
||||
});
|
||||
|
||||
describe("prefersRespondAsync", () => {
|
||||
it("reads RFC 7240 preferences, case-insensitively and among others", () => {
|
||||
expect(prefersRespondAsync(req("respond-async"))).toBe(true);
|
||||
expect(prefersRespondAsync(req("Respond-Async"))).toBe(true);
|
||||
expect(prefersRespondAsync(req("return=minimal, respond-async"))).toBe(
|
||||
true,
|
||||
);
|
||||
expect(prefersRespondAsync(req("respond-async; wait=10"))).toBe(true);
|
||||
});
|
||||
|
||||
it("keeps the synchronous path for every other request", () => {
|
||||
expect(prefersRespondAsync(req())).toBe(false);
|
||||
expect(prefersRespondAsync(req("return=minimal"))).toBe(false);
|
||||
expect(prefersRespondAsync(req("respond-asynchronously"))).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("acceptedRunResponse", () => {
|
||||
it("answers 202 with the run, the poll address and the applied preference", async () => {
|
||||
const res = acceptedRunResponse("run1");
|
||||
expect(res.status).toBe(202);
|
||||
expect(res.headers.get("Preference-Applied")).toBe("respond-async");
|
||||
expect(res.headers.get("Location")).toBe("/api/ai-runs/run1");
|
||||
expect(await res.json()).toEqual({
|
||||
data: { runId: "run1", status: "QUEUED", pollAfterMs: 1500 },
|
||||
error: null,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("outcomes", () => {
|
||||
it("renders a failure as the route's own envelope", async () => {
|
||||
const res = outcomeResponse({
|
||||
ok: false,
|
||||
status: 422,
|
||||
message: "Couldn't read the stored document.",
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
expect(res.status).toBe(422);
|
||||
expect((await res.json()).meta).toEqual({
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
});
|
||||
|
||||
it("keeps a capability refusal's status and code", () => {
|
||||
expect(
|
||||
refusalAsFailure(new AiUnavailableError("labsOcr", "consent_required")),
|
||||
).toEqual({
|
||||
status: 403,
|
||||
message: "AI consent is required for this feature",
|
||||
errorCode: "consent.ai.required",
|
||||
});
|
||||
expect(refusalAsFailure(new Error("x"))).toBeNull();
|
||||
});
|
||||
|
||||
it("names the missing worker", async () => {
|
||||
const res = workerUnavailableResponse();
|
||||
expect(res.status).toBe(503);
|
||||
expect((await res.json()).meta.errorCode).toBe("aiRuns.workerUnavailable");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,219 @@
|
||||
/**
|
||||
* The run service's reads and the reaper's decisions, with Prisma stubbed.
|
||||
* The integration file (`tests/integration/document-ai-runs.test.ts`) proves
|
||||
* the same against Postgres; this one pins the arithmetic and the branches.
|
||||
*/
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const findFirst = vi.fn();
|
||||
const findMany = vi.fn();
|
||||
const updateMany = vi.fn();
|
||||
const deleteMany = vi.fn();
|
||||
vi.mock("@/lib/db", () => ({
|
||||
prisma: {
|
||||
documentAiRun: {
|
||||
findFirst: (...a: unknown[]) => findFirst(...a),
|
||||
findMany: (...a: unknown[]) => findMany(...a),
|
||||
updateMany: (...a: unknown[]) => updateMany(...a),
|
||||
deleteMany: (...a: unknown[]) => deleteMany(...a),
|
||||
},
|
||||
},
|
||||
}));
|
||||
const reconcileSpend = vi.fn();
|
||||
vi.mock("@/lib/ai/coach/budget", () => ({
|
||||
reconcileSpend: (...a: unknown[]) => reconcileSpend(...a),
|
||||
}));
|
||||
|
||||
import { _resetCryptoCacheForTests, encryptBytes } from "@/lib/crypto";
|
||||
import { DOCUMENT_AI_RUN_RESULT_AAD } from "@/lib/crypto/encrypted-columns";
|
||||
import { readAiRunForUser, reapAiRuns } from "@/lib/documents/ai-runs/store";
|
||||
|
||||
const NOW = new Date("2027-01-15T10:00:00.000Z");
|
||||
const at = (msAgo: number) => new Date(NOW.getTime() - msAgo);
|
||||
|
||||
function row(overrides: Record<string, unknown>) {
|
||||
return {
|
||||
id: "r1",
|
||||
kind: "LABS_OCR_EXTRACT",
|
||||
documentId: null,
|
||||
status: "QUEUED",
|
||||
resultEncrypted: null,
|
||||
errorCode: null,
|
||||
errorStatus: null,
|
||||
errorMessage: null,
|
||||
createdAt: at(10_000),
|
||||
startedAt: null,
|
||||
finishedAt: null,
|
||||
expiresAt: new Date(NOW.getTime() + 60_000),
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.stubEnv("ENCRYPTION_KEYS", "");
|
||||
vi.stubEnv("ENCRYPTION_ACTIVE_KEY_ID", "");
|
||||
vi.stubEnv("ENCRYPTION_KEY", "ab".repeat(32));
|
||||
_resetCryptoCacheForTests();
|
||||
vi.clearAllMocks();
|
||||
updateMany.mockResolvedValue({ count: 1 });
|
||||
deleteMany.mockResolvedValue({ count: 0 });
|
||||
});
|
||||
|
||||
describe("readAiRunForUser", () => {
|
||||
it("scopes the read to the caller", async () => {
|
||||
findFirst.mockResolvedValue(null);
|
||||
expect(await readAiRunForUser("u1", "r1", NOW)).toBeNull();
|
||||
expect(findFirst).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ where: { id: "r1", userId: "u1" } }),
|
||||
);
|
||||
});
|
||||
|
||||
it("reports the queue wait and a poll interval that backs off", async () => {
|
||||
findFirst.mockResolvedValue(row({ createdAt: at(10_000) }));
|
||||
const early = await readAiRunForUser("u1", "r1", NOW);
|
||||
expect(early).toMatchObject({ status: "QUEUED", queuedForMs: 10_000 });
|
||||
expect(early?.retryAfterMs).toBe(1500 + 500);
|
||||
|
||||
findFirst.mockResolvedValue(row({ createdAt: at(10 * 60_000) }));
|
||||
const late = await readAiRunForUser("u1", "r1", NOW);
|
||||
expect(late?.retryAfterMs).toBe(5000);
|
||||
});
|
||||
|
||||
it("opens a result and stops the poll", async () => {
|
||||
const sealed = encryptBytes(
|
||||
Buffer.from(JSON.stringify({ rows: [], reportDate: null })),
|
||||
DOCUMENT_AI_RUN_RESULT_AAD,
|
||||
);
|
||||
findFirst.mockResolvedValue(
|
||||
row({
|
||||
status: "SUCCEEDED",
|
||||
resultEncrypted: new Uint8Array(sealed),
|
||||
finishedAt: at(1000),
|
||||
}),
|
||||
);
|
||||
expect(await readAiRunForUser("u1", "r1", NOW)).toMatchObject({
|
||||
status: "SUCCEEDED",
|
||||
result: { rows: [], reportDate: null },
|
||||
error: null,
|
||||
retryAfterMs: null,
|
||||
queuedForMs: null,
|
||||
});
|
||||
});
|
||||
|
||||
it("reports a result that no longer opens as a failure, never an empty success", async () => {
|
||||
findFirst.mockResolvedValue(
|
||||
row({
|
||||
status: "SUCCEEDED",
|
||||
resultEncrypted: new Uint8Array([1, 2, 3]),
|
||||
}),
|
||||
);
|
||||
expect(await readAiRunForUser("u1", "r1", NOW)).toMatchObject({
|
||||
status: "FAILED",
|
||||
result: null,
|
||||
error: { code: "aiRuns.failed", status: 500 },
|
||||
});
|
||||
});
|
||||
|
||||
it("hands back the route's own failure", async () => {
|
||||
findFirst.mockResolvedValue(
|
||||
row({
|
||||
status: "FAILED",
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
errorStatus: 422,
|
||||
errorMessage: "Couldn't read the report. Try a clearer photo.",
|
||||
}),
|
||||
);
|
||||
expect((await readAiRunForUser("u1", "r1", NOW))?.error).toEqual({
|
||||
code: "labs.ocr.extractFailed",
|
||||
status: 422,
|
||||
message: "Couldn't read the report. Try a clearer photo.",
|
||||
});
|
||||
});
|
||||
|
||||
it("treats a finished run past its retention as gone", async () => {
|
||||
findFirst.mockResolvedValue(row({ status: "FAILED", expiresAt: at(1) }));
|
||||
expect(await readAiRunForUser("u1", "r1", NOW)).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("reapAiRuns", () => {
|
||||
it("fails a waiting run as worker-unavailable and refunds it, a running one as timed out", async () => {
|
||||
findMany.mockResolvedValue([
|
||||
{
|
||||
id: "q1",
|
||||
userId: "u1",
|
||||
status: "QUEUED",
|
||||
paramsJson: {
|
||||
mode: "vision",
|
||||
budget: { reserved: 8000, owner: "user", dateKey: "2027-01-15" },
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "x1",
|
||||
userId: "u2",
|
||||
status: "RUNNING",
|
||||
paramsJson: { mode: "text" },
|
||||
},
|
||||
]);
|
||||
deleteMany.mockResolvedValue({ count: 4 });
|
||||
|
||||
expect(await reapAiRuns(NOW)).toEqual({
|
||||
workerUnavailable: 1,
|
||||
timedOut: 1,
|
||||
deleted: 4,
|
||||
});
|
||||
expect(updateMany).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
where: { id: "q1", status: "QUEUED" },
|
||||
data: expect.objectContaining({
|
||||
status: "FAILED",
|
||||
errorCode: "aiRuns.workerUnavailable",
|
||||
inputEncrypted: null,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expect(updateMany).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
where: { id: "x1", status: "RUNNING" },
|
||||
data: expect.objectContaining({ errorCode: "aiRuns.timedOut" }),
|
||||
}),
|
||||
);
|
||||
// Only the run nobody claimed is refunded; a running one's reservation
|
||||
// belongs to the worker that claimed it.
|
||||
expect(reconcileSpend).toHaveBeenCalledOnce();
|
||||
expect(reconcileSpend).toHaveBeenCalledWith(
|
||||
"u1",
|
||||
8000,
|
||||
0,
|
||||
"2027-01-15",
|
||||
0,
|
||||
{
|
||||
servedBy: null,
|
||||
reservedOwner: "user",
|
||||
},
|
||||
);
|
||||
expect(deleteMany).toHaveBeenCalledWith({
|
||||
where: {
|
||||
status: { in: ["SUCCEEDED", "FAILED"] },
|
||||
expiresAt: { lte: NOW },
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it("does not refund a run another party ended first", async () => {
|
||||
findMany.mockResolvedValue([
|
||||
{
|
||||
id: "q1",
|
||||
userId: "u1",
|
||||
status: "QUEUED",
|
||||
paramsJson: {
|
||||
mode: "vision",
|
||||
budget: { reserved: 8000, owner: "user", dateKey: "2027-01-15" },
|
||||
},
|
||||
},
|
||||
]);
|
||||
updateMany.mockResolvedValue({ count: 0 });
|
||||
expect((await reapAiRuns(NOW)).workerUnavailable).toBe(0);
|
||||
expect(reconcileSpend).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,75 @@
|
||||
/**
|
||||
* The HTTP side of a background run: reading `Prefer: respond-async`, the 202
|
||||
* a route answers with, and turning a run body's outcome back into the
|
||||
* response the synchronous route has always given.
|
||||
*/
|
||||
import { apiError, apiSuccess } from "@/lib/api-response";
|
||||
import { AiUnavailableError } from "@/lib/ai/capabilities/refusal";
|
||||
|
||||
import type { AiRunFailure } from "./store";
|
||||
import {
|
||||
AI_RUN_ERROR_CODES,
|
||||
AI_RUN_POLL_AFTER_MS,
|
||||
type AiRunAccepted,
|
||||
type AiRunOutcome,
|
||||
} from "./types";
|
||||
|
||||
/**
|
||||
* Whether the request asked for the background path (RFC 7240). Without the
|
||||
* header the four document routes answer exactly as they always have, which
|
||||
* is what a client that has not adopted runs relies on.
|
||||
*/
|
||||
export function prefersRespondAsync(request: Request): boolean {
|
||||
const header = request.headers.get("prefer");
|
||||
if (!header) return false;
|
||||
return header
|
||||
.split(",")
|
||||
.some(
|
||||
(preference) =>
|
||||
preference.split(";")[0]?.trim().toLowerCase() === "respond-async",
|
||||
);
|
||||
}
|
||||
|
||||
/** The 202 for a queued run, with the poll address and the applied preference. */
|
||||
export function acceptedRunResponse(runId: string): Response {
|
||||
const body: AiRunAccepted = {
|
||||
runId,
|
||||
status: "QUEUED",
|
||||
pollAfterMs: AI_RUN_POLL_AFTER_MS,
|
||||
};
|
||||
const response = apiSuccess(body, 202);
|
||||
response.headers.set("Preference-Applied", "respond-async");
|
||||
response.headers.set("Location", `/api/ai-runs/${runId}`);
|
||||
return response;
|
||||
}
|
||||
|
||||
/** The answer when the run could not be handed to the background worker. */
|
||||
export function workerUnavailableResponse(): Response {
|
||||
return apiError(
|
||||
"The background worker is not available. Try again in a moment.",
|
||||
503,
|
||||
{ errorCode: AI_RUN_ERROR_CODES.workerUnavailable },
|
||||
);
|
||||
}
|
||||
|
||||
/** A run body's outcome as the synchronous route's response. */
|
||||
export function outcomeResponse<T>(outcome: AiRunOutcome<T>): Response {
|
||||
return outcome.ok
|
||||
? apiSuccess(outcome.data)
|
||||
: apiError(outcome.message, outcome.status, {
|
||||
errorCode: outcome.errorCode,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* A capability refusal thrown inside a run, as the failure the route would
|
||||
* have rendered for it. Null for anything else.
|
||||
*/
|
||||
export function refusalAsFailure(error: unknown): AiRunFailure | null {
|
||||
if (!(error instanceof AiUnavailableError)) return null;
|
||||
return {
|
||||
status: error.status,
|
||||
message: error.message,
|
||||
errorCode: error.meta.errorCode,
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
/**
|
||||
* "Read with AI" for one stored document: the body both the synchronous index
|
||||
* route and the background run execute, so the two paths cannot drift.
|
||||
*
|
||||
* Vision mode decrypts the stored original, makes one provider transcription
|
||||
* call against a reservation the caller already took, and indexes the text.
|
||||
* Text mode indexes text the browser read on the device (no provider call).
|
||||
* Both then record the attempt, lift an import's hold on automatic reading,
|
||||
* and continue into the same lab staging the automatic index worker runs.
|
||||
*
|
||||
* Every failure comes back as the outcome the route has always answered with
|
||||
* (status, message, `documents.inbound.*` code); the caller renders it as a
|
||||
* response or stores it on the run.
|
||||
*/
|
||||
import { AI_BUDGETS } from "@/lib/ai/ai-budgets";
|
||||
import type { ProviderChainType } from "@/lib/ai/provider-chain";
|
||||
import { auditLog } from "@/lib/auth/audit";
|
||||
import { prisma } from "@/lib/db";
|
||||
import {
|
||||
prepareVisionInput,
|
||||
refundDocumentAiSlot,
|
||||
type LoadedDocument,
|
||||
} from "@/lib/documents/ai-route-support";
|
||||
import { maybeAutoStageLabFacts } from "@/lib/documents/auto-stage-labs";
|
||||
import { upsertContentIndex } from "@/lib/documents/content-index";
|
||||
import {
|
||||
DocumentDescribeError,
|
||||
transcribeDocument,
|
||||
} from "@/lib/documents/describe";
|
||||
import { recordIndexAttempt } from "@/lib/documents/index-document";
|
||||
import type { requireDocumentVisionProvider } from "@/lib/documents/provider-order";
|
||||
import { annotate } from "@/lib/logging/context";
|
||||
|
||||
import { settleAiRunBudget } from "./store";
|
||||
import type {
|
||||
AiRunBudget,
|
||||
AiRunOutcome,
|
||||
DocumentIndexRunResult,
|
||||
} from "./types";
|
||||
|
||||
export type DocumentVisionPick = Awaited<
|
||||
ReturnType<typeof requireDocumentVisionProvider>
|
||||
>;
|
||||
|
||||
/**
|
||||
* The most sequential model calls one read makes: the transcription, then the
|
||||
* lab staging read with its one corrective retry.
|
||||
*/
|
||||
export const DOCUMENT_INDEX_MODEL_CALLS = 3;
|
||||
|
||||
/** The tokens a vision read reserves before it is dispatched. */
|
||||
export const DOCUMENT_INDEX_RESERVE_TOKENS =
|
||||
AI_BUDGETS.documentTranscribe.maxTokens;
|
||||
|
||||
/** Who asked: a request (with its client address) or a background run. */
|
||||
export interface IndexOrigin {
|
||||
ipAddress: string | null;
|
||||
worker: boolean;
|
||||
}
|
||||
|
||||
export type DocumentIndexInput =
|
||||
| { mode: "text"; text: string }
|
||||
| { mode: "vision"; pick: DocumentVisionPick; budget: AiRunBudget };
|
||||
|
||||
/** A failure in the shape the route would answer with (`errorCode` literal, so the catalogue guard reads it). */
|
||||
const failure = (f: {
|
||||
status: number;
|
||||
message: string;
|
||||
errorCode: string;
|
||||
}): AiRunOutcome<never> => ({ ok: false, ...f });
|
||||
|
||||
async function finishIndex(
|
||||
userId: string,
|
||||
documentId: string,
|
||||
source: "vision" | "text-ocr",
|
||||
tokenCount: number,
|
||||
origin: IndexOrigin,
|
||||
): Promise<AiRunOutcome<DocumentIndexRunResult>> {
|
||||
// Refs #776 — the manual read is the second writer of the attempt record
|
||||
// (the auto job + backfill share `indexLoadedDocument`): a success clears
|
||||
// any stored failure reason so the detail view stops explaining a problem
|
||||
// that no longer exists.
|
||||
await recordIndexAttempt(userId, documentId, {
|
||||
indexed: true,
|
||||
source,
|
||||
tokenCount,
|
||||
});
|
||||
// Read with AI is the person reading the document on purpose, so an
|
||||
// import's hold on automatic AI reading (`aiRead=defer`) ends here, before
|
||||
// the lab staging below that would otherwise refuse it.
|
||||
await prisma.inboundDocument.updateMany({
|
||||
where: { id: documentId, userId, aiReadDeferred: true },
|
||||
data: { aiReadDeferred: false },
|
||||
});
|
||||
await auditLog("documents.inbound.index", {
|
||||
userId,
|
||||
// A background run has no request left to attribute the action from.
|
||||
...(origin.worker ? { actorUserId: userId } : {}),
|
||||
ipAddress: origin.ipAddress,
|
||||
details: { documentId, source, tokens: tokenCount },
|
||||
});
|
||||
annotate({
|
||||
action: { name: "documents.contentIndex.upsert" },
|
||||
meta: { documentId, source, tokens: tokenCount },
|
||||
});
|
||||
// A manual read continues into the SAME lab staging the automatic index
|
||||
// worker performs, so a skipped or failed auto run is recoverable per
|
||||
// document without a re-upload. Every guard lives inside the helper (both
|
||||
// modules on, still STORED with no facts, provider + consent, looks like a
|
||||
// lab report) — a non-lab document is a tagged no-op, and a staging failure
|
||||
// never fails the index that just succeeded.
|
||||
const staging = await maybeAutoStageLabFacts(userId, documentId).catch(
|
||||
() => null,
|
||||
);
|
||||
const labFactsStaged = staging?.staged === true ? staging.facts : 0;
|
||||
return {
|
||||
ok: true,
|
||||
data: { documentId, indexed: true, tokenCount, labFactsStaged },
|
||||
};
|
||||
}
|
||||
|
||||
/** Index one owned document. The caller has authorised and charged the read. */
|
||||
export async function executeDocumentIndex(args: {
|
||||
userId: string;
|
||||
document: LoadedDocument;
|
||||
input: DocumentIndexInput;
|
||||
origin: IndexOrigin;
|
||||
}): Promise<AiRunOutcome<DocumentIndexRunResult>> {
|
||||
const { userId, document, input, origin } = args;
|
||||
|
||||
if (input.mode === "text") {
|
||||
const { tokenCount } = await upsertContentIndex({
|
||||
userId,
|
||||
documentId: document.id,
|
||||
text: input.text,
|
||||
source: "text-ocr",
|
||||
providerType: null,
|
||||
});
|
||||
return finishIndex(userId, document.id, "text-ocr", tokenCount, origin);
|
||||
}
|
||||
|
||||
const { pick, budget } = input;
|
||||
const servedBy = pick.entry.providerType as ProviderChainType;
|
||||
|
||||
const vision = await prepareVisionInput(document, pick.pdfSupported);
|
||||
if (!vision.ok) {
|
||||
// Preparation failed before any provider dispatch — the reservation and
|
||||
// the slot go back. Refs #776 — each failure is also recorded on the row
|
||||
// so the detail view can explain the missing index after the toast is
|
||||
// gone.
|
||||
await settleAiRunBudget(userId, budget, 0, null);
|
||||
await refundDocumentAiSlot(userId);
|
||||
if (vision.reason === "pdfNeedsAnthropic") {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "pdf-needs-anthropic",
|
||||
});
|
||||
return failure({
|
||||
status: 422,
|
||||
message:
|
||||
"PDF scanning needs a Claude vision provider; use local OCR instead.",
|
||||
errorCode: "documents.inbound.pdfNeedsAnthropic",
|
||||
});
|
||||
}
|
||||
if (vision.reason === "rasterFailed") {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "raster-failed",
|
||||
});
|
||||
return failure({
|
||||
status: 422,
|
||||
message: "The PDF pages couldn't be rendered for scanning.",
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
if (vision.reason === "fileType") {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "local-unsupported",
|
||||
});
|
||||
return failure({
|
||||
status: 422,
|
||||
message: "This document can't be scanned. Use local OCR (text mode).",
|
||||
errorCode: "documents.inbound.fileType",
|
||||
});
|
||||
}
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "decrypt-error",
|
||||
});
|
||||
return failure({
|
||||
status: 422,
|
||||
message: "Couldn't read the stored document.",
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
|
||||
let text: string;
|
||||
try {
|
||||
({ text } = await transcribeDocument({
|
||||
provider: pick.entry.instance,
|
||||
providerType: pick.providerType,
|
||||
images: vision.images,
|
||||
documents: vision.documents,
|
||||
}));
|
||||
} catch (err) {
|
||||
await settleAiRunBudget(userId, budget, 0, null);
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "provider-error",
|
||||
});
|
||||
if (err instanceof DocumentDescribeError) {
|
||||
return failure({
|
||||
status: 422,
|
||||
message: "Couldn't read the document. Try a clearer copy.",
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
annotate({
|
||||
action: { name: "documents.contentIndex.failed" },
|
||||
meta: { reason: "provider_error", mode: "vision" },
|
||||
});
|
||||
return failure({
|
||||
status: 502,
|
||||
message: "Couldn't read the document. Try a clearer copy.",
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
await settleAiRunBudget(userId, budget, budget.reserved, servedBy);
|
||||
|
||||
// Refs #776 — the empty-transcription guard, same contract as the auto
|
||||
// path (`tryProviderIndex`): a provider answer with no text must never
|
||||
// become a "successful" empty index. The spend stays charged (the
|
||||
// provider was called); the honest answer is an error plus the recorded
|
||||
// reason.
|
||||
if (text.trim().length === 0) {
|
||||
await recordIndexAttempt(userId, document.id, {
|
||||
indexed: false,
|
||||
reason: "empty-transcription",
|
||||
});
|
||||
return failure({
|
||||
status: 422,
|
||||
message:
|
||||
"The provider returned no text for this document. Try a clearer copy.",
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
}
|
||||
const { tokenCount } = await upsertContentIndex({
|
||||
userId,
|
||||
documentId: document.id,
|
||||
text,
|
||||
source: "vision",
|
||||
providerType: pick.providerType,
|
||||
});
|
||||
return finishIndex(userId, document.id, "vision", tokenCount, origin);
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
/**
|
||||
* Start a background run from a route that has already authorised and
|
||||
* charged it: create the row, hand it to the worker, answer 202. When the
|
||||
* worker cannot be reached the run is failed on the spot, the reservation and
|
||||
* the rate slot go back, and the answer is 503, so nobody polls a run that
|
||||
* nothing will ever pick up.
|
||||
*/
|
||||
import type { Buffer } from "node:buffer";
|
||||
|
||||
import { prisma } from "@/lib/db";
|
||||
import { enqueueDocumentAiRun } from "@/lib/jobs/document-ai-run";
|
||||
import { annotate } from "@/lib/logging/context";
|
||||
|
||||
import { acceptedRunResponse, workerUnavailableResponse } from "./http";
|
||||
import { createAiRun, failAiRun, settleAiRunBudget } from "./store";
|
||||
import {
|
||||
AI_RUN_ERROR_CODES,
|
||||
type AiRunParams,
|
||||
type DocumentAiRunKindValue,
|
||||
} from "./types";
|
||||
|
||||
/**
|
||||
* The live run of this kind for this document, if one is queued or running:
|
||||
* a second press of "Read with AI" (another tab, a double click) attaches to
|
||||
* it instead of paying for a second read.
|
||||
*/
|
||||
export async function findLiveDocumentRun(
|
||||
userId: string,
|
||||
documentId: string,
|
||||
kind: DocumentAiRunKindValue,
|
||||
): Promise<string | null> {
|
||||
const row = await prisma.documentAiRun.findFirst({
|
||||
where: {
|
||||
userId,
|
||||
documentId,
|
||||
kind,
|
||||
status: { in: ["QUEUED", "RUNNING"] },
|
||||
expiresAt: { gt: new Date() },
|
||||
},
|
||||
select: { id: true },
|
||||
orderBy: { createdAt: "desc" },
|
||||
});
|
||||
return row?.id ?? null;
|
||||
}
|
||||
|
||||
export async function startAiRun(args: {
|
||||
userId: string;
|
||||
kind: DocumentAiRunKindValue;
|
||||
documentId?: string | null;
|
||||
params: AiRunParams;
|
||||
input?: Buffer | null;
|
||||
/** Hand the route's rate slot back when the run cannot start. */
|
||||
refundSlot: () => Promise<void>;
|
||||
}): Promise<Response> {
|
||||
const runId = await createAiRun({
|
||||
userId: args.userId,
|
||||
kind: args.kind,
|
||||
documentId: args.documentId ?? null,
|
||||
params: args.params,
|
||||
input: args.input ?? null,
|
||||
});
|
||||
const enqueued = await enqueueDocumentAiRun(runId, args.userId);
|
||||
if (!enqueued) {
|
||||
await failAiRun(
|
||||
runId,
|
||||
{
|
||||
status: 503,
|
||||
errorCode: AI_RUN_ERROR_CODES.workerUnavailable,
|
||||
message: "The background worker is not available.",
|
||||
},
|
||||
"QUEUED",
|
||||
);
|
||||
await settleAiRunBudget(args.userId, args.params.budget, 0, null);
|
||||
await args.refundSlot();
|
||||
annotate({
|
||||
action: { name: "ai_runs.worker_unavailable" },
|
||||
meta: { kind: args.kind },
|
||||
});
|
||||
return workerUnavailableResponse();
|
||||
}
|
||||
annotate({
|
||||
action: { name: "ai_runs.queued" },
|
||||
meta: { kind: args.kind, mode: args.params.mode },
|
||||
});
|
||||
return acceptedRunResponse(runId);
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
/**
|
||||
* The `DocumentAiRun` row: create, claim, finish, fail, read, reap.
|
||||
*
|
||||
* Every state change is a conditional `updateMany` on the state it leaves, so
|
||||
* the worker, the reaper and a second delivery of the same job can never both
|
||||
* move one run: whoever matched the row did it, the rest see a count of zero.
|
||||
* The budget reservation taken at enqueue is settled exactly once, by the
|
||||
* party that ends the run (the worker after its provider call, or the reaper
|
||||
* for a run no worker ever picked up).
|
||||
*/
|
||||
import { Buffer } from "node:buffer";
|
||||
|
||||
import { Prisma } from "@/generated/prisma/client";
|
||||
import { reconcileSpend } from "@/lib/ai/coach/budget";
|
||||
import type { ProviderChainType } from "@/lib/ai/provider-chain";
|
||||
import { decryptBytes, encryptBytes } from "@/lib/crypto";
|
||||
import {
|
||||
DOCUMENT_AI_RUN_INPUT_AAD,
|
||||
DOCUMENT_AI_RUN_RESULT_AAD,
|
||||
} from "@/lib/crypto/encrypted-columns";
|
||||
import { prisma } from "@/lib/db";
|
||||
|
||||
import {
|
||||
AI_RUN_ERROR_CODES,
|
||||
AI_RUN_POLL_AFTER_MS,
|
||||
AI_RUN_QUEUE_WAIT_MS,
|
||||
AI_RUN_RETENTION_MS,
|
||||
type AiRunBudget,
|
||||
type AiRunDto,
|
||||
type AiRunParams,
|
||||
type AiRunResult,
|
||||
type DocumentAiRunKindValue,
|
||||
} from "./types";
|
||||
|
||||
/** The longest a client is told to wait between polls. */
|
||||
const MAX_POLL_INTERVAL_MS = 5000;
|
||||
|
||||
export interface CreateAiRunArgs {
|
||||
userId: string;
|
||||
kind: DocumentAiRunKindValue;
|
||||
documentId?: string | null;
|
||||
params: AiRunParams;
|
||||
/** The input only the run can read later: an upload, or browser-read text. */
|
||||
input?: Buffer | null;
|
||||
now?: Date;
|
||||
}
|
||||
|
||||
/** Insert a queued run. The caller enqueues the job right after. */
|
||||
export async function createAiRun(args: CreateAiRunArgs): Promise<string> {
|
||||
const now = args.now ?? new Date();
|
||||
const row = await prisma.documentAiRun.create({
|
||||
data: {
|
||||
userId: args.userId,
|
||||
kind: args.kind,
|
||||
documentId: args.documentId ?? null,
|
||||
status: "QUEUED",
|
||||
paramsJson: args.params as unknown as Prisma.InputJsonValue,
|
||||
inputEncrypted: args.input
|
||||
? new Uint8Array(encryptBytes(args.input, DOCUMENT_AI_RUN_INPUT_AAD))
|
||||
: null,
|
||||
createdAt: now,
|
||||
expiresAt: new Date(now.getTime() + AI_RUN_QUEUE_WAIT_MS),
|
||||
},
|
||||
select: { id: true },
|
||||
});
|
||||
return row.id;
|
||||
}
|
||||
|
||||
/** A run the worker has claimed, with its input opened. */
|
||||
export interface ClaimedAiRun {
|
||||
id: string;
|
||||
userId: string;
|
||||
kind: DocumentAiRunKindValue;
|
||||
documentId: string | null;
|
||||
params: AiRunParams;
|
||||
input: Buffer | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Move a queued run to RUNNING with the deadline the reaper enforces. Null
|
||||
* when the run is gone or no longer queued (a second delivery, or the reaper
|
||||
* failed it while it waited).
|
||||
*/
|
||||
export async function claimAiRun(
|
||||
id: string,
|
||||
runDeadline: Date,
|
||||
now: Date = new Date(),
|
||||
): Promise<ClaimedAiRun | null> {
|
||||
const moved = await prisma.documentAiRun.updateMany({
|
||||
where: { id, status: "QUEUED" },
|
||||
data: { status: "RUNNING", startedAt: now, expiresAt: runDeadline },
|
||||
});
|
||||
if (moved.count !== 1) return null;
|
||||
const row = await prisma.documentAiRun.findUnique({
|
||||
where: { id },
|
||||
select: {
|
||||
id: true,
|
||||
userId: true,
|
||||
kind: true,
|
||||
documentId: true,
|
||||
paramsJson: true,
|
||||
inputEncrypted: true,
|
||||
},
|
||||
});
|
||||
if (!row) return null;
|
||||
let input: Buffer | null = null;
|
||||
if (row.inputEncrypted) {
|
||||
input = decryptBytes(
|
||||
Buffer.from(row.inputEncrypted),
|
||||
DOCUMENT_AI_RUN_INPUT_AAD,
|
||||
);
|
||||
}
|
||||
return {
|
||||
id: row.id,
|
||||
userId: row.userId,
|
||||
kind: row.kind,
|
||||
documentId: row.documentId,
|
||||
params: row.paramsJson as unknown as AiRunParams,
|
||||
input,
|
||||
};
|
||||
}
|
||||
|
||||
/** Store the result and drop the input. False when the run was already ended. */
|
||||
export async function completeAiRun(
|
||||
id: string,
|
||||
result: AiRunResult,
|
||||
now: Date = new Date(),
|
||||
): Promise<boolean> {
|
||||
const sealed = encryptBytes(
|
||||
Buffer.from(JSON.stringify(result), "utf8"),
|
||||
DOCUMENT_AI_RUN_RESULT_AAD,
|
||||
);
|
||||
const moved = await prisma.documentAiRun.updateMany({
|
||||
where: { id, status: "RUNNING" },
|
||||
data: {
|
||||
status: "SUCCEEDED",
|
||||
resultEncrypted: new Uint8Array(sealed),
|
||||
inputEncrypted: null,
|
||||
finishedAt: now,
|
||||
expiresAt: new Date(now.getTime() + AI_RUN_RETENTION_MS),
|
||||
},
|
||||
});
|
||||
return moved.count === 1;
|
||||
}
|
||||
|
||||
export interface AiRunFailure {
|
||||
status: number;
|
||||
message: string;
|
||||
errorCode: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Fail a run that is still queued or running and drop its input. False when
|
||||
* another party ended it first.
|
||||
*/
|
||||
export async function failAiRun(
|
||||
id: string,
|
||||
failure: AiRunFailure,
|
||||
from: "QUEUED" | "RUNNING" | "ANY" = "ANY",
|
||||
now: Date = new Date(),
|
||||
): Promise<boolean> {
|
||||
const moved = await prisma.documentAiRun.updateMany({
|
||||
where: {
|
||||
id,
|
||||
status: from === "ANY" ? { in: ["QUEUED", "RUNNING"] } : from,
|
||||
},
|
||||
data: {
|
||||
status: "FAILED",
|
||||
errorCode: failure.errorCode,
|
||||
errorStatus: failure.status,
|
||||
errorMessage: failure.message,
|
||||
inputEncrypted: null,
|
||||
finishedAt: now,
|
||||
expiresAt: new Date(now.getTime() + AI_RUN_RETENTION_MS),
|
||||
},
|
||||
});
|
||||
return moved.count === 1;
|
||||
}
|
||||
|
||||
/**
|
||||
* Settle the reservation taken at enqueue. `servedBy` null with `actual` 0
|
||||
* hands it back in full. Best effort: the run's own state matters more than
|
||||
* the meter.
|
||||
*/
|
||||
export async function settleAiRunBudget(
|
||||
userId: string,
|
||||
budget: AiRunBudget | undefined,
|
||||
actual: number,
|
||||
servedBy: ProviderChainType | null,
|
||||
): Promise<void> {
|
||||
if (!budget || budget.reserved <= 0) return;
|
||||
try {
|
||||
await reconcileSpend(userId, budget.reserved, actual, budget.dateKey, 0, {
|
||||
servedBy,
|
||||
reservedOwner: budget.owner,
|
||||
});
|
||||
} catch {
|
||||
// The run's outcome is already decided; a failed refund is not a reason
|
||||
// to change it.
|
||||
}
|
||||
}
|
||||
|
||||
function openResult(sealed: Uint8Array | null): AiRunResult | null {
|
||||
if (!sealed) return null;
|
||||
try {
|
||||
const plain = decryptBytes(Buffer.from(sealed), DOCUMENT_AI_RUN_RESULT_AAD);
|
||||
return JSON.parse(plain.toString("utf8")) as AiRunResult;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* One run as its owner sees it, or null when it does not exist, belongs to
|
||||
* somebody else, or has outlived its retention (the reaper may not have
|
||||
* deleted it yet; it is gone either way).
|
||||
*/
|
||||
export async function readAiRunForUser(
|
||||
userId: string,
|
||||
id: string,
|
||||
now: Date = new Date(),
|
||||
): Promise<AiRunDto | null> {
|
||||
const row = await prisma.documentAiRun.findFirst({
|
||||
where: { id, userId },
|
||||
select: {
|
||||
id: true,
|
||||
kind: true,
|
||||
documentId: true,
|
||||
status: true,
|
||||
resultEncrypted: true,
|
||||
errorCode: true,
|
||||
errorStatus: true,
|
||||
errorMessage: true,
|
||||
createdAt: true,
|
||||
startedAt: true,
|
||||
finishedAt: true,
|
||||
expiresAt: true,
|
||||
},
|
||||
});
|
||||
if (!row) return null;
|
||||
const terminal = row.status === "SUCCEEDED" || row.status === "FAILED";
|
||||
if (terminal && row.expiresAt.getTime() <= now.getTime()) return null;
|
||||
|
||||
const result =
|
||||
row.status === "SUCCEEDED" ? openResult(row.resultEncrypted) : null;
|
||||
// A result that no longer opens (a key rotated away under it) is reported
|
||||
// as a failure rather than as a success with nothing in it.
|
||||
const unreadable = row.status === "SUCCEEDED" && result === null;
|
||||
const age = now.getTime() - row.createdAt.getTime();
|
||||
const sinceStart = row.startedAt
|
||||
? now.getTime() - row.startedAt.getTime()
|
||||
: 0;
|
||||
return {
|
||||
id: row.id,
|
||||
kind: row.kind,
|
||||
documentId: row.documentId,
|
||||
status: unreadable ? "FAILED" : row.status,
|
||||
result,
|
||||
error:
|
||||
row.status === "FAILED" || unreadable
|
||||
? {
|
||||
code: unreadable
|
||||
? AI_RUN_ERROR_CODES.failed
|
||||
: (row.errorCode ?? AI_RUN_ERROR_CODES.failed),
|
||||
message: unreadable
|
||||
? "The result of this run can no longer be read."
|
||||
: (row.errorMessage ?? "The run failed."),
|
||||
status: unreadable ? 500 : (row.errorStatus ?? 500),
|
||||
}
|
||||
: null,
|
||||
createdAt: row.createdAt.toISOString(),
|
||||
startedAt: row.startedAt?.toISOString() ?? null,
|
||||
finishedAt: row.finishedAt?.toISOString() ?? null,
|
||||
queuedForMs: row.status === "QUEUED" ? Math.max(0, age) : null,
|
||||
retryAfterMs: terminal
|
||||
? null
|
||||
: Math.min(
|
||||
MAX_POLL_INTERVAL_MS,
|
||||
// Poll briskly at first, then back off: most reads finish within
|
||||
// seconds, and a slow local model can take minutes.
|
||||
AI_RUN_POLL_AFTER_MS + Math.floor(Math.max(age, sinceStart) / 20),
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
export interface ReapSummary {
|
||||
workerUnavailable: number;
|
||||
timedOut: number;
|
||||
deleted: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* The reaper's one pass, keyed on `expiresAt`, which always holds the next
|
||||
* deadline of a run:
|
||||
*
|
||||
* - QUEUED past it: no worker picked the run up. Failed as
|
||||
* `aiRuns.workerUnavailable`, and its reservation is handed back.
|
||||
* - RUNNING past it: the run outlived every model call it was allowed.
|
||||
* Failed as `aiRuns.timedOut`. The reservation stays with the worker that
|
||||
* claimed it; a worker that died with it leaves it charged, the cautious
|
||||
* direction for a meter.
|
||||
* - SUCCEEDED / FAILED past it: an hour after it finished. Deleted.
|
||||
*/
|
||||
export async function reapAiRuns(
|
||||
now: Date = new Date(),
|
||||
batch = 500,
|
||||
): Promise<ReapSummary> {
|
||||
const summary: ReapSummary = {
|
||||
workerUnavailable: 0,
|
||||
timedOut: 0,
|
||||
deleted: 0,
|
||||
};
|
||||
|
||||
const stuck = await prisma.documentAiRun.findMany({
|
||||
where: { status: { in: ["QUEUED", "RUNNING"] }, expiresAt: { lte: now } },
|
||||
select: { id: true, userId: true, status: true, paramsJson: true },
|
||||
orderBy: { expiresAt: "asc" },
|
||||
take: batch,
|
||||
});
|
||||
for (const run of stuck) {
|
||||
if (run.status === "QUEUED") {
|
||||
const failed = await failAiRun(
|
||||
run.id,
|
||||
{
|
||||
status: 503,
|
||||
errorCode: AI_RUN_ERROR_CODES.workerUnavailable,
|
||||
message:
|
||||
"No background worker picked this up. Ask the server operator to check the worker.",
|
||||
},
|
||||
"QUEUED",
|
||||
now,
|
||||
);
|
||||
if (failed) {
|
||||
summary.workerUnavailable += 1;
|
||||
const params = run.paramsJson as unknown as AiRunParams | null;
|
||||
await settleAiRunBudget(run.userId, params?.budget, 0, null);
|
||||
}
|
||||
} else {
|
||||
const failed = await failAiRun(
|
||||
run.id,
|
||||
{
|
||||
status: 504,
|
||||
errorCode: AI_RUN_ERROR_CODES.timedOut,
|
||||
message: "Reading took longer than the AI response time allows.",
|
||||
},
|
||||
"RUNNING",
|
||||
now,
|
||||
);
|
||||
if (failed) summary.timedOut += 1;
|
||||
}
|
||||
}
|
||||
|
||||
const gone = await prisma.documentAiRun.deleteMany({
|
||||
where: { status: { in: ["SUCCEEDED", "FAILED"] }, expiresAt: { lte: now } },
|
||||
});
|
||||
summary.deleted = gone.count;
|
||||
return summary;
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
/**
|
||||
* The shapes a background document AI run carries: its kinds and states, the
|
||||
* outcome a run body returns, the error codes only a run can produce, and the
|
||||
* DTO the poll route publishes.
|
||||
*
|
||||
* A run is a ticket. The request that asked for the read answers 202 with the
|
||||
* run id, the worker does the read, and the client polls
|
||||
* `GET /api/ai-runs/{id}` until the run is terminal. The result lives in the
|
||||
* row for an hour after it finishes and then the row is deleted.
|
||||
*/
|
||||
import type { OcrExtractResponseDto } from "@/lib/validations/labs-ocr";
|
||||
|
||||
export const DOCUMENT_AI_RUN_KINDS = [
|
||||
"DOCUMENT_INDEX",
|
||||
"LABS_OCR_EXTRACT",
|
||||
] as const;
|
||||
export type DocumentAiRunKindValue = (typeof DOCUMENT_AI_RUN_KINDS)[number];
|
||||
|
||||
export const DOCUMENT_AI_RUN_STATUSES = [
|
||||
"QUEUED",
|
||||
"RUNNING",
|
||||
"SUCCEEDED",
|
||||
"FAILED",
|
||||
] as const;
|
||||
export type DocumentAiRunStatusValue =
|
||||
(typeof DOCUMENT_AI_RUN_STATUSES)[number];
|
||||
|
||||
/**
|
||||
* The codes only a background run produces. Every other failure a run
|
||||
* reports carries the code its synchronous route answers with, so a client
|
||||
* maps both paths through one table.
|
||||
*/
|
||||
export const AI_RUN_ERROR_CODES = {
|
||||
notFound: "aiRuns.notFound",
|
||||
workerUnavailable: "aiRuns.workerUnavailable",
|
||||
timedOut: "aiRuns.timedOut",
|
||||
failed: "aiRuns.failed",
|
||||
} as const;
|
||||
|
||||
/** How long a queued run waits for a worker before it is failed. */
|
||||
export const AI_RUN_QUEUE_WAIT_MS = 15 * 60 * 1000;
|
||||
|
||||
/** How long a finished run stays readable before the reaper deletes it. */
|
||||
export const AI_RUN_RETENTION_MS = 60 * 60 * 1000;
|
||||
|
||||
/** Room a running run gets on top of its model calls (decrypt, raster, write). */
|
||||
export const AI_RUN_MARGIN_MS = 60 * 1000;
|
||||
|
||||
/** The interval the 202 tells a client to wait before its first poll. */
|
||||
export const AI_RUN_POLL_AFTER_MS = 1500;
|
||||
|
||||
/** The budget reservation taken at enqueue, settled by whoever ends the run. */
|
||||
export interface AiRunBudget {
|
||||
reserved: number;
|
||||
owner: "operator" | "user";
|
||||
dateKey: string;
|
||||
}
|
||||
|
||||
/** What a run needs to know to run. Never health data. */
|
||||
export interface AiRunParams {
|
||||
/** `vision` reads the stored original or the uploaded scan; `text` reads text the browser read. */
|
||||
mode: "vision" | "text";
|
||||
/** The sniffed MIME of an uploaded scan (lab scan, vision mode only). */
|
||||
mime?: "image/jpeg" | "image/png" | "image/webp" | "application/pdf";
|
||||
/** Present when the enqueue reserved budget for a provider call. */
|
||||
budget?: AiRunBudget;
|
||||
}
|
||||
|
||||
/** "Read with AI": the same body the synchronous index route answers with. */
|
||||
export interface DocumentIndexRunResult {
|
||||
documentId: string;
|
||||
indexed: true;
|
||||
tokenCount: number;
|
||||
labFactsStaged: number;
|
||||
}
|
||||
|
||||
export type AiRunResult = DocumentIndexRunResult | OcrExtractResponseDto;
|
||||
|
||||
/**
|
||||
* What one run body produced: the success body, or the failure the
|
||||
* synchronous route would have answered with.
|
||||
*/
|
||||
export type AiRunOutcome<T> =
|
||||
| { ok: true; data: T }
|
||||
| { ok: false; status: number; message: string; errorCode: string };
|
||||
|
||||
/** The poll route's body. */
|
||||
export interface AiRunDto {
|
||||
id: string;
|
||||
kind: DocumentAiRunKindValue;
|
||||
documentId: string | null;
|
||||
status: DocumentAiRunStatusValue;
|
||||
result: AiRunResult | null;
|
||||
error: { code: string; message: string; status: number } | null;
|
||||
createdAt: string;
|
||||
startedAt: string | null;
|
||||
finishedAt: string | null;
|
||||
/** How long the run has waited for a worker, while it is still queued. */
|
||||
queuedForMs: number | null;
|
||||
/** When to ask again, while the run is not terminal. */
|
||||
retryAfterMs: number | null;
|
||||
}
|
||||
|
||||
/** The 202 body. */
|
||||
export interface AiRunAccepted {
|
||||
runId: string;
|
||||
status: "QUEUED";
|
||||
pollAfterMs: number;
|
||||
}
|
||||
@@ -676,6 +676,8 @@ export const NOT_IN_BACKUP_MODELS: Readonly<Record<string, string>> = {
|
||||
"When this host last put this account's copy in this operator's bucket, and how big it was. It describes one deployment's relationship with one bucket, so restoring it elsewhere would assert an off-host copy that host has never written.",
|
||||
DocumentImportKey:
|
||||
"The source keys of imported documents the person deleted and the purge then removed, kept so a re-run of the importer does not store them again. The backup carries only live documents, so the tombstones that hold the same memory before the purge do not travel either; carrying the ledger alone would restore half of one decision. After a restore, a deliberate re-import can bring such a document back, and the person deletes it again.",
|
||||
DocumentAiRun:
|
||||
"A ticket for one background read (a document read with AI, a lab scan), deleted an hour after it finishes. The document and the committed rows travel; the ticket does not, and a restored run would point at a worker job that does not exist on the receiving host.",
|
||||
DocumentSourceAlias:
|
||||
"Further import source keys that were answered with an already stored document because the bytes matched. The next import run writes the same alias again the moment it meets the same bytes, so a restore loses nothing a re-run does not rebuild; and for a deleted document it restores no less than the tombstones and ledger beside it do, which do not travel either.",
|
||||
ImportJob:
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
/**
|
||||
* The background document AI run worker (v1.40).
|
||||
*
|
||||
* What only the worker decides, with every collaborator stubbed: a run is
|
||||
* claimed once; the capability is asked again for the record and the wire is
|
||||
* re-checked for the provider actually picked, so a switch turned off or a
|
||||
* consent withdrawn while the run waited stops it before anything leaves; the
|
||||
* deadline follows the person's AI response time and a read that cannot fit
|
||||
* the job's expiry is refused; and whatever ends the run settles the budget
|
||||
* exactly once.
|
||||
*
|
||||
* Mutation checks, run against this file: deleting the `aiEgressRefusal` call
|
||||
* in `assertRunMayEgress` turns "refuses when consent was withdrawn" red;
|
||||
* deleting the `aiCapabilityForJob` check turns "refuses when the operator
|
||||
* turned reading off" red; dropping the `ends > deadlineAt` comparison turns
|
||||
* "refuses a read that cannot fit" red.
|
||||
*/
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const claimAiRun = vi.fn();
|
||||
const completeAiRun = vi.fn();
|
||||
const failAiRun = vi.fn();
|
||||
const settleAiRunBudget = vi.fn();
|
||||
vi.mock("@/lib/documents/ai-runs/store", () => ({
|
||||
claimAiRun: (...a: unknown[]) => claimAiRun(...a),
|
||||
completeAiRun: (...a: unknown[]) => completeAiRun(...a),
|
||||
failAiRun: (...a: unknown[]) => failAiRun(...a),
|
||||
settleAiRunBudget: (...a: unknown[]) => settleAiRunBudget(...a),
|
||||
}));
|
||||
|
||||
const aiCapabilityForJob = vi.fn();
|
||||
vi.mock("@/lib/ai/capabilities/gate", () => ({
|
||||
aiCapabilityForJob: (...a: unknown[]) => aiCapabilityForJob(...a),
|
||||
}));
|
||||
const aiEgressRefusal = vi.fn();
|
||||
vi.mock("@/lib/ai/capabilities/egress", () => ({
|
||||
aiEgressRefusal: (...a: unknown[]) => aiEgressRefusal(...a),
|
||||
}));
|
||||
|
||||
const requireDocumentVisionProvider = vi.fn();
|
||||
vi.mock("@/lib/documents/provider-order", () => ({
|
||||
requireDocumentVisionProvider: (...a: unknown[]) =>
|
||||
requireDocumentVisionProvider(...a),
|
||||
}));
|
||||
const requireLabsOcrProvider = vi.fn();
|
||||
vi.mock("@/lib/labs/ocr-capability", () => ({
|
||||
requireLabsOcrProvider: (...a: unknown[]) => requireLabsOcrProvider(...a),
|
||||
}));
|
||||
|
||||
const executeDocumentIndex = vi.fn();
|
||||
vi.mock("@/lib/documents/ai-runs/index-run", () => ({
|
||||
DOCUMENT_INDEX_MODEL_CALLS: 3,
|
||||
executeDocumentIndex: (...a: unknown[]) => executeDocumentIndex(...a),
|
||||
}));
|
||||
const executeOcrExtraction = vi.fn();
|
||||
vi.mock("@/lib/labs/ocr-run", () => ({
|
||||
OCR_EXTRACT_MODEL_CALLS: 2,
|
||||
executeOcrExtraction: (...a: unknown[]) => executeOcrExtraction(...a),
|
||||
}));
|
||||
|
||||
const loadOwnedDocument = vi.fn();
|
||||
vi.mock("@/lib/documents/ai-route-support", () => ({
|
||||
loadOwnedDocument: (...a: unknown[]) => loadOwnedDocument(...a),
|
||||
}));
|
||||
|
||||
const updateMany = vi.fn();
|
||||
const findUser = vi.fn();
|
||||
vi.mock("@/lib/db", () => ({
|
||||
prisma: {
|
||||
documentAiRun: { updateMany: (...a: unknown[]) => updateMany(...a) },
|
||||
user: { findUnique: (...a: unknown[]) => findUser(...a) },
|
||||
},
|
||||
}));
|
||||
|
||||
const send = vi.fn();
|
||||
let boss: { send: typeof send } | null = { send };
|
||||
vi.mock("@/lib/jobs/boss-instance", () => ({
|
||||
getGlobalBoss: () => boss,
|
||||
}));
|
||||
|
||||
import { AiUnavailableError } from "@/lib/ai/capabilities/refusal";
|
||||
import {
|
||||
DOCUMENT_AI_RUN_EXPIRE_SECONDS,
|
||||
DOCUMENT_AI_RUN_QUEUE,
|
||||
enqueueDocumentAiRun,
|
||||
handleDocumentAiRunJobs,
|
||||
runDocumentAiRun,
|
||||
} from "@/lib/jobs/document-ai-run";
|
||||
|
||||
const NOW = 1_800_000_000_000;
|
||||
const now = () => NOW;
|
||||
const budget = { reserved: 4000, owner: "user", dateKey: "2027-01-15" };
|
||||
|
||||
function claimed(overrides: Record<string, unknown> = {}) {
|
||||
return {
|
||||
id: "run1",
|
||||
userId: "u1",
|
||||
kind: "DOCUMENT_INDEX",
|
||||
documentId: "d1",
|
||||
params: { mode: "vision", budget },
|
||||
input: null,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
const visionPick = (seconds: number | null) => ({
|
||||
entry: {
|
||||
providerType: "anthropic",
|
||||
instance: { responseTimeoutSeconds: seconds },
|
||||
},
|
||||
providerType: "anthropic",
|
||||
pdfSupported: true,
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
boss = { send };
|
||||
claimAiRun.mockResolvedValue(claimed());
|
||||
aiCapabilityForJob.mockResolvedValue({ available: true, reason: null });
|
||||
aiEgressRefusal.mockResolvedValue(null);
|
||||
requireDocumentVisionProvider.mockResolvedValue(visionPick(null));
|
||||
loadOwnedDocument.mockResolvedValue({ id: "d1" });
|
||||
executeDocumentIndex.mockResolvedValue({
|
||||
ok: true,
|
||||
data: { documentId: "d1", indexed: true, tokenCount: 3, labFactsStaged: 0 },
|
||||
});
|
||||
updateMany.mockResolvedValue({ count: 1 });
|
||||
completeAiRun.mockResolvedValue(true);
|
||||
failAiRun.mockResolvedValue(true);
|
||||
});
|
||||
|
||||
describe("enqueueDocumentAiRun", () => {
|
||||
it("sends one run with its expiry and no retry", async () => {
|
||||
send.mockResolvedValue("job1");
|
||||
expect(await enqueueDocumentAiRun("run1", "u1")).toBe(true);
|
||||
expect(send).toHaveBeenCalledWith(
|
||||
DOCUMENT_AI_RUN_QUEUE,
|
||||
{ runId: "run1", userId: "u1" },
|
||||
{ retryLimit: 0, expireInSeconds: DOCUMENT_AI_RUN_EXPIRE_SECONDS },
|
||||
);
|
||||
});
|
||||
|
||||
it("reports no queue as not enqueued", async () => {
|
||||
boss = null;
|
||||
expect(await enqueueDocumentAiRun("run1", "u1")).toBe(false);
|
||||
});
|
||||
|
||||
it("reports a failed send as not enqueued", async () => {
|
||||
send.mockRejectedValue(new Error("db down"));
|
||||
expect(await enqueueDocumentAiRun("run1", "u1")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("runDocumentAiRun", () => {
|
||||
it("skips a run another delivery already claimed", async () => {
|
||||
claimAiRun.mockResolvedValue(null);
|
||||
expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe(
|
||||
"skipped",
|
||||
);
|
||||
expect(executeDocumentIndex).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("reads, stores the result, and arms the deadline from the setting", async () => {
|
||||
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 + 3 * 300_000 + 60_000) },
|
||||
});
|
||||
expect(aiEgressRefusal).toHaveBeenCalledWith("documentAi", "u1", [
|
||||
"anthropic",
|
||||
]);
|
||||
expect(completeAiRun).toHaveBeenCalledWith(
|
||||
"run1",
|
||||
{ documentId: "d1", indexed: true, tokenCount: 3, labFactsStaged: 0 },
|
||||
new Date(NOW),
|
||||
);
|
||||
expect(settleAiRunBudget).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("refuses when the operator turned reading off while the run waited", async () => {
|
||||
aiCapabilityForJob.mockResolvedValue({
|
||||
available: false,
|
||||
reason: "operator_disabled",
|
||||
});
|
||||
expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe("failed");
|
||||
expect(executeDocumentIndex).not.toHaveBeenCalled();
|
||||
expect(settleAiRunBudget).toHaveBeenCalledWith("u1", budget, 0, null);
|
||||
expect(failAiRun).toHaveBeenCalledWith(
|
||||
"run1",
|
||||
expect.objectContaining({
|
||||
status: 403,
|
||||
errorCode: "assistant.disabled.documentAi",
|
||||
}),
|
||||
"RUNNING",
|
||||
new Date(NOW),
|
||||
);
|
||||
});
|
||||
|
||||
it("refuses when consent was withdrawn for the picked provider", async () => {
|
||||
aiEgressRefusal.mockResolvedValue(
|
||||
new AiUnavailableError("documentAi", "consent_required"),
|
||||
);
|
||||
expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe("failed");
|
||||
expect(executeDocumentIndex).not.toHaveBeenCalled();
|
||||
expect(failAiRun).toHaveBeenCalledWith(
|
||||
"run1",
|
||||
expect.objectContaining({
|
||||
status: 403,
|
||||
errorCode: "consent.ai.required",
|
||||
}),
|
||||
"RUNNING",
|
||||
new Date(NOW),
|
||||
);
|
||||
});
|
||||
|
||||
it("refuses a read that cannot fit the job's expiry", async () => {
|
||||
requireDocumentVisionProvider.mockResolvedValue(visionPick(600));
|
||||
// Three calls at ten minutes plus the margin do not fit in twenty minutes.
|
||||
expect(await runDocumentAiRun("run1", NOW + 20 * 60_000, now)).toBe(
|
||||
"failed",
|
||||
);
|
||||
expect(executeDocumentIndex).not.toHaveBeenCalled();
|
||||
expect(settleAiRunBudget).toHaveBeenCalledWith("u1", budget, 0, null);
|
||||
expect(failAiRun).toHaveBeenCalledWith(
|
||||
"run1",
|
||||
expect.objectContaining({ errorCode: "aiRuns.timedOut", status: 504 }),
|
||||
"RUNNING",
|
||||
new Date(NOW),
|
||||
);
|
||||
});
|
||||
|
||||
it("stores the route's own failure when the read fails", async () => {
|
||||
executeDocumentIndex.mockResolvedValue({
|
||||
ok: false,
|
||||
status: 422,
|
||||
message: "Couldn't read the document. Try a clearer copy.",
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
});
|
||||
expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe("failed");
|
||||
// The body settled its own reservation; the worker must not settle again.
|
||||
expect(settleAiRunBudget).not.toHaveBeenCalled();
|
||||
expect(failAiRun).toHaveBeenCalledWith(
|
||||
"run1",
|
||||
{
|
||||
status: 422,
|
||||
errorCode: "documents.inbound.extractFailed",
|
||||
message: "Couldn't read the document. Try a clearer copy.",
|
||||
},
|
||||
"RUNNING",
|
||||
new Date(NOW),
|
||||
);
|
||||
});
|
||||
|
||||
it("fails a run whose document is gone and hands the reservation back", async () => {
|
||||
loadOwnedDocument.mockResolvedValue(null);
|
||||
expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe("failed");
|
||||
expect(settleAiRunBudget).toHaveBeenCalledWith("u1", budget, 0, null);
|
||||
expect(failAiRun).toHaveBeenCalledWith(
|
||||
"run1",
|
||||
expect.objectContaining({ errorCode: "documents.inbound.notFound" }),
|
||||
"RUNNING",
|
||||
new Date(NOW),
|
||||
);
|
||||
});
|
||||
|
||||
it("arms a text read for the lab staging that follows it", async () => {
|
||||
claimAiRun.mockResolvedValue(
|
||||
claimed({ params: { mode: "text" }, input: Buffer.from("Befund") }),
|
||||
);
|
||||
findUser.mockResolvedValue({ aiResponseTimeoutSeconds: 120 });
|
||||
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 * 120_000 + 60_000) },
|
||||
});
|
||||
expect(executeDocumentIndex).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ input: { mode: "text", text: "Befund" } }),
|
||||
);
|
||||
expect(aiEgressRefusal).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("reads a lab scan through the labs capability and its own provider", async () => {
|
||||
claimAiRun.mockResolvedValue(
|
||||
claimed({
|
||||
kind: "LABS_OCR_EXTRACT",
|
||||
documentId: null,
|
||||
params: { mode: "text", budget },
|
||||
input: Buffer.from("Hb 14 g/dL"),
|
||||
}),
|
||||
);
|
||||
requireLabsOcrProvider.mockResolvedValue({
|
||||
entry: {
|
||||
providerType: "openai",
|
||||
instance: { responseTimeoutSeconds: null },
|
||||
},
|
||||
providerType: "openai",
|
||||
});
|
||||
executeOcrExtraction.mockResolvedValue({
|
||||
ok: true,
|
||||
data: { reportDate: null, providerType: "openai", rows: [] },
|
||||
});
|
||||
expect(await runDocumentAiRun("run1", NOW + 2_700_000, now)).toBe(
|
||||
"succeeded",
|
||||
);
|
||||
expect(aiCapabilityForJob).toHaveBeenCalledWith("u1", "labsOcr");
|
||||
expect(aiEgressRefusal).toHaveBeenCalledWith("labsOcr", "u1", ["openai"]);
|
||||
expect(executeOcrExtraction).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
input: expect.objectContaining({ mode: "text", text: "Hb 14 g/dL" }),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("handleDocumentAiRunJobs", () => {
|
||||
it("runs each job under its own deadline and reports the counts", async () => {
|
||||
const outcome = await handleDocumentAiRunJobs([
|
||||
{
|
||||
id: "j1",
|
||||
data: { runId: "run1", userId: "u1" },
|
||||
expireInSeconds: DOCUMENT_AI_RUN_EXPIRE_SECONDS,
|
||||
},
|
||||
{ id: "j2", data: { runId: "", userId: "u1" }, expireInSeconds: 60 },
|
||||
] as never);
|
||||
expect(outcome).toMatchObject({
|
||||
ok: true,
|
||||
did: { jobs: 2, processed: 1, failed: 0, skipped: 1 },
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,47 @@
|
||||
/**
|
||||
* Five-minute reaper for background document AI runs (v1.40).
|
||||
*
|
||||
* A run carries its next deadline in `expiresAt`, and this pass acts on it:
|
||||
* a run no worker picked up within fifteen minutes fails as
|
||||
* `aiRuns.workerUnavailable` (its reservation handed back), a run past the
|
||||
* time its model calls were allowed fails as `aiRuns.timedOut`, and a run
|
||||
* that finished an hour ago is deleted with its sealed result. Every step is
|
||||
* a conditional write on the state it leaves, so a tick that overlaps a
|
||||
* worker, or a second tick, changes nothing twice. The next tick is the retry.
|
||||
*/
|
||||
import { type Job } from "pg-boss";
|
||||
|
||||
import { reapAiRuns } from "@/lib/documents/ai-runs/store";
|
||||
import { jobDone, jobFailed, type JobOutcome } from "@/lib/jobs/job-outcome";
|
||||
import { withBackgroundEvent } from "@/lib/logging/background";
|
||||
|
||||
export const DOCUMENT_AI_RUN_REAPER_QUEUE = "document-ai-run-reaper";
|
||||
/** Every five minutes: a queued run fails within twenty minutes at worst. */
|
||||
export const DOCUMENT_AI_RUN_REAPER_CRON = "*/5 * * * *";
|
||||
|
||||
export interface DocumentAiRunReaperPayload {
|
||||
triggeredAt?: string;
|
||||
}
|
||||
|
||||
export async function handleDocumentAiRunReaper(
|
||||
jobs: Job<DocumentAiRunReaperPayload>[],
|
||||
): Promise<JobOutcome> {
|
||||
void jobs;
|
||||
return withBackgroundEvent("job.document_ai_run_reaper", async (evt) => {
|
||||
try {
|
||||
const summary = await reapAiRuns();
|
||||
evt.setAction({ name: "ai_runs.reap" });
|
||||
evt.addMeta("ai_runs_worker_unavailable", summary.workerUnavailable);
|
||||
evt.addMeta("ai_runs_timed_out", summary.timedOut);
|
||||
evt.addMeta("ai_runs_deleted", summary.deleted);
|
||||
return jobDone({
|
||||
ai_runs_worker_unavailable: summary.workerUnavailable,
|
||||
ai_runs_timed_out: summary.timedOut,
|
||||
ai_runs_deleted: summary.deleted,
|
||||
});
|
||||
} catch (err) {
|
||||
evt.addWarning(`document-ai-run-reaper failed: ${err}`);
|
||||
return jobFailed("document ai run reaper failed", err);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,356 @@
|
||||
/**
|
||||
* The background document AI run: one model read a request handed off with
|
||||
* 202 instead of holding the connection open for it (v1.40).
|
||||
*
|
||||
* The route has already refused what it could refuse quickly (auth, module,
|
||||
* capability, provider, rate bucket, budget, a malformed upload) and sealed
|
||||
* the run's input into its `DocumentAiRun` row. This worker:
|
||||
*
|
||||
* 1. claims the row (QUEUED → RUNNING, a conditional update, so a second
|
||||
* delivery does nothing);
|
||||
* 2. asks the capability again for the record (`aiCapabilityForJob`), picks
|
||||
* the provider the way the route does and re-checks the wire for exactly
|
||||
* that provider (`aiEgressRefusal`): consent may have been withdrawn or a
|
||||
* switch turned off while the run waited;
|
||||
* 3. sets the run's deadline from the person's AI response time (one
|
||||
* `callTimeoutMs` per model call the read can make) and refuses to start
|
||||
* a read that would not fit inside the job's own expiry (`jobDeadline`);
|
||||
* 4. runs the same body the synchronous route runs, and stores the result
|
||||
* or the failure the route would have answered with.
|
||||
*
|
||||
* The queue is `long` (`queue-runtime.ts`): a read at the longest response
|
||||
* setting outlasts pg-boss's fifteen-minute default. Each job runs under the
|
||||
* lock `lockedPass` takes per run, and is sent with `retryLimit: 0`: a run
|
||||
* the worker lost is failed by the reaper and asked for again by the person,
|
||||
* never re-run behind their back.
|
||||
*
|
||||
* The queue MUST be registered in the maintenance registrar
|
||||
* (`src/lib/jobs/reminder/register-maintenance.ts`) so pg-boss provisions it.
|
||||
*/
|
||||
import type { Job } from "pg-boss";
|
||||
|
||||
import { aiEgressRefusal } from "@/lib/ai/capabilities/egress";
|
||||
import { aiCapabilityForJob } from "@/lib/ai/capabilities/gate";
|
||||
import { AiUnavailableError } from "@/lib/ai/capabilities/refusal";
|
||||
import {
|
||||
PICK_DECIDED_REASONS,
|
||||
type AiCapabilityKey,
|
||||
} from "@/lib/ai/capabilities/types";
|
||||
import { callTimeoutMs } from "@/lib/ai/effective-timeout";
|
||||
import { prisma } from "@/lib/db";
|
||||
import { loadOwnedDocument } from "@/lib/documents/ai-route-support";
|
||||
import { refusalAsFailure } from "@/lib/documents/ai-runs/http";
|
||||
import {
|
||||
DOCUMENT_INDEX_MODEL_CALLS,
|
||||
executeDocumentIndex,
|
||||
} from "@/lib/documents/ai-runs/index-run";
|
||||
import {
|
||||
claimAiRun,
|
||||
completeAiRun,
|
||||
failAiRun,
|
||||
settleAiRunBudget,
|
||||
type AiRunFailure,
|
||||
type ClaimedAiRun,
|
||||
} from "@/lib/documents/ai-runs/store";
|
||||
import {
|
||||
AI_RUN_ERROR_CODES,
|
||||
AI_RUN_MARGIN_MS,
|
||||
type AiRunOutcome,
|
||||
type AiRunResult,
|
||||
type DocumentAiRunKindValue,
|
||||
} from "@/lib/documents/ai-runs/types";
|
||||
import { requireDocumentVisionProvider } from "@/lib/documents/provider-order";
|
||||
import { getGlobalBoss } from "@/lib/jobs/boss-instance";
|
||||
import { jobDeadline } from "@/lib/jobs/job-budget";
|
||||
import { jobDone, type JobOutcome } from "@/lib/jobs/job-outcome";
|
||||
import { requireLabsOcrProvider } from "@/lib/labs/ocr-capability";
|
||||
import {
|
||||
executeOcrExtraction,
|
||||
OCR_EXTRACT_MODEL_CALLS,
|
||||
} from "@/lib/labs/ocr-run";
|
||||
import { annotate } from "@/lib/logging/context";
|
||||
|
||||
export const DOCUMENT_AI_RUN_QUEUE = "document-ai-run";
|
||||
|
||||
/** Two reads at a time: provider calls, kept off the request pool. */
|
||||
export const DOCUMENT_AI_RUN_CONCURRENCY = 2;
|
||||
|
||||
/**
|
||||
* The job's expiry: one hour. Three model calls at the longest response
|
||||
* setting (600 s) plus the document work fit inside the three quarters of it
|
||||
* `jobDeadline` allows; a read that would not is refused before it starts.
|
||||
*/
|
||||
export const DOCUMENT_AI_RUN_EXPIRE_SECONDS = 60 * 60;
|
||||
|
||||
export interface DocumentAiRunPayload {
|
||||
runId: string;
|
||||
userId: string;
|
||||
}
|
||||
|
||||
/** The capability each kind of run answers to. */
|
||||
const RUN_CAPABILITY: Record<DocumentAiRunKindValue, AiCapabilityKey> = {
|
||||
DOCUMENT_INDEX: "documentAi",
|
||||
LABS_OCR_EXTRACT: "labsOcr",
|
||||
};
|
||||
|
||||
/**
|
||||
* Hand a created run to the worker. False when no queue is reachable; the
|
||||
* caller fails the run and answers 503 rather than leaving a run nobody
|
||||
* will pick up.
|
||||
*/
|
||||
export async function enqueueDocumentAiRun(
|
||||
runId: string,
|
||||
userId: string,
|
||||
): Promise<boolean> {
|
||||
const boss = getGlobalBoss();
|
||||
if (!boss) return false;
|
||||
const payload: DocumentAiRunPayload = { runId, userId };
|
||||
try {
|
||||
const jobId = await boss.send(DOCUMENT_AI_RUN_QUEUE, payload, {
|
||||
retryLimit: 0,
|
||||
expireInSeconds: DOCUMENT_AI_RUN_EXPIRE_SECONDS,
|
||||
});
|
||||
return Boolean(jobId);
|
||||
} catch (err) {
|
||||
annotate({
|
||||
action: { name: "ai_runs.enqueue_failed" },
|
||||
meta: { reason: err instanceof Error ? err.name : "unknown" },
|
||||
});
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
class RunRefused extends Error {
|
||||
constructor(readonly failure: AiRunFailure) {
|
||||
super(failure.errorCode);
|
||||
}
|
||||
}
|
||||
|
||||
/** Refuse the run when the capability or the wire is closed for this provider. */
|
||||
async function assertRunMayEgress(
|
||||
key: AiCapabilityKey,
|
||||
userId: string,
|
||||
providerTypes: readonly string[],
|
||||
): Promise<void> {
|
||||
const state = await aiCapabilityForJob(userId, key);
|
||||
if (state.reason !== null && !PICK_DECIDED_REASONS.has(state.reason)) {
|
||||
throw new AiUnavailableError(key, state.reason);
|
||||
}
|
||||
const refusal = await aiEgressRefusal(key, userId, providerTypes);
|
||||
if (refusal) throw refusal;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the deadline the reaper enforces from the person's response time, or
|
||||
* refuse a read that cannot finish inside the job's own expiry.
|
||||
*/
|
||||
async function armDeadline(
|
||||
runId: string,
|
||||
perCallMs: number,
|
||||
modelCalls: number,
|
||||
deadlineAt: number | undefined,
|
||||
now: () => number,
|
||||
): Promise<void> {
|
||||
const ends = now() + perCallMs * modelCalls + AI_RUN_MARGIN_MS;
|
||||
if (deadlineAt !== undefined && ends > deadlineAt) {
|
||||
throw new RunRefused({
|
||||
status: 504,
|
||||
errorCode: AI_RUN_ERROR_CODES.timedOut,
|
||||
message:
|
||||
"The AI response time is set longer than a background read may run.",
|
||||
});
|
||||
}
|
||||
await prisma.documentAiRun.updateMany({
|
||||
where: { id: runId, status: "RUNNING" },
|
||||
data: { expiresAt: new Date(ends) },
|
||||
});
|
||||
}
|
||||
|
||||
async function runDocumentIndex(
|
||||
run: ClaimedAiRun,
|
||||
deadlineAt: number | undefined,
|
||||
now: () => number,
|
||||
): Promise<AiRunOutcome<AiRunResult>> {
|
||||
const document = run.documentId
|
||||
? await loadOwnedDocument(run.userId, run.documentId)
|
||||
: null;
|
||||
if (!document) {
|
||||
await settleAiRunBudget(run.userId, run.params.budget, 0, null);
|
||||
return {
|
||||
ok: false,
|
||||
status: 404,
|
||||
message: "Document not found",
|
||||
errorCode: "documents.inbound.notFound",
|
||||
};
|
||||
}
|
||||
if (run.params.mode === "text") {
|
||||
// No transcription call, but the lab staging after the index is model
|
||||
// work (a read and its corrective retry) under the person's setting.
|
||||
const owner = await prisma.user.findUnique({
|
||||
where: { id: run.userId },
|
||||
select: { aiResponseTimeoutSeconds: true },
|
||||
});
|
||||
await armDeadline(
|
||||
run.id,
|
||||
callTimeoutMs({}, owner?.aiResponseTimeoutSeconds ?? null),
|
||||
DOCUMENT_INDEX_MODEL_CALLS - 1,
|
||||
deadlineAt,
|
||||
now,
|
||||
);
|
||||
return executeDocumentIndex({
|
||||
userId: run.userId,
|
||||
document,
|
||||
input: { mode: "text", text: run.input?.toString("utf8") ?? "" },
|
||||
origin: { ipAddress: null, worker: true },
|
||||
});
|
||||
}
|
||||
const pick = await requireDocumentVisionProvider(run.userId);
|
||||
await assertRunMayEgress("documentAi", run.userId, [pick.providerType]);
|
||||
await armDeadline(
|
||||
run.id,
|
||||
callTimeoutMs({}, pick.entry.instance.responseTimeoutSeconds),
|
||||
DOCUMENT_INDEX_MODEL_CALLS,
|
||||
deadlineAt,
|
||||
now,
|
||||
);
|
||||
if (!run.params.budget) throw new Error("vision run without a reservation");
|
||||
return executeDocumentIndex({
|
||||
userId: run.userId,
|
||||
document,
|
||||
input: { mode: "vision", pick, budget: run.params.budget },
|
||||
origin: { ipAddress: null, worker: true },
|
||||
});
|
||||
}
|
||||
|
||||
async function runLabsOcr(
|
||||
run: ClaimedAiRun,
|
||||
deadlineAt: number | undefined,
|
||||
now: () => number,
|
||||
): Promise<AiRunOutcome<AiRunResult>> {
|
||||
if (!run.input) throw new Error("lab scan run without an input");
|
||||
if (run.params.mode === "text") {
|
||||
const pick = await requireLabsOcrProvider(run.userId, "text");
|
||||
await assertRunMayEgress("labsOcr", run.userId, [pick.providerType]);
|
||||
await armDeadline(
|
||||
run.id,
|
||||
callTimeoutMs({}, pick.entry.instance.responseTimeoutSeconds),
|
||||
OCR_EXTRACT_MODEL_CALLS,
|
||||
deadlineAt,
|
||||
now,
|
||||
);
|
||||
return executeOcrExtraction({
|
||||
userId: run.userId,
|
||||
budget: run.params.budget,
|
||||
input: { mode: "text", text: run.input.toString("utf8"), pick },
|
||||
});
|
||||
}
|
||||
const mime = run.params.mime;
|
||||
if (!mime) throw new Error("lab scan run without a MIME type");
|
||||
const pick = await requireLabsOcrProvider(run.userId, "vision");
|
||||
await assertRunMayEgress("labsOcr", run.userId, [pick.providerType]);
|
||||
await armDeadline(
|
||||
run.id,
|
||||
callTimeoutMs({}, pick.entry.instance.responseTimeoutSeconds),
|
||||
OCR_EXTRACT_MODEL_CALLS,
|
||||
deadlineAt,
|
||||
now,
|
||||
);
|
||||
return executeOcrExtraction({
|
||||
userId: run.userId,
|
||||
budget: run.params.budget,
|
||||
input: { mode: "vision", bytes: run.input, mime, pick },
|
||||
});
|
||||
}
|
||||
|
||||
export type DocumentAiRunResult = "succeeded" | "failed" | "skipped";
|
||||
|
||||
/**
|
||||
* Run one queued run to its end. `deadlineAt` is the instant the job's own
|
||||
* expiry allows work until (`jobDeadline`); `now` is injectable for tests.
|
||||
*/
|
||||
export async function runDocumentAiRun(
|
||||
runId: string,
|
||||
deadlineAt: number | undefined,
|
||||
now: () => number = Date.now,
|
||||
): Promise<DocumentAiRunResult> {
|
||||
const provisional = new Date(
|
||||
deadlineAt ?? now() + DOCUMENT_AI_RUN_EXPIRE_SECONDS * 1000,
|
||||
);
|
||||
const run = await claimAiRun(runId, provisional, new Date(now()));
|
||||
if (!run) return "skipped";
|
||||
|
||||
let outcome: AiRunOutcome<AiRunResult>;
|
||||
try {
|
||||
outcome =
|
||||
run.kind === "DOCUMENT_INDEX"
|
||||
? await runDocumentIndex(run, deadlineAt, now)
|
||||
: await runLabsOcr(run, deadlineAt, now);
|
||||
} catch (err) {
|
||||
// Everything that throws here threw before the body took over the
|
||||
// reservation, so it goes back here.
|
||||
await settleAiRunBudget(run.userId, run.params.budget, 0, null);
|
||||
const failure: AiRunFailure = (err instanceof RunRefused
|
||||
? err.failure
|
||||
: null) ??
|
||||
refusalAsFailure(err) ?? {
|
||||
status: 500,
|
||||
errorCode: AI_RUN_ERROR_CODES.failed,
|
||||
message: "The background read failed.",
|
||||
};
|
||||
await failAiRun(run.id, failure, "RUNNING", new Date(now()));
|
||||
annotate({
|
||||
action: { name: "ai_runs.failed" },
|
||||
meta: {
|
||||
kind: run.kind,
|
||||
capability: RUN_CAPABILITY[run.kind],
|
||||
errorCode: failure.errorCode,
|
||||
},
|
||||
});
|
||||
return "failed";
|
||||
}
|
||||
|
||||
if (outcome.ok) {
|
||||
await completeAiRun(run.id, outcome.data, new Date(now()));
|
||||
annotate({
|
||||
action: { name: "ai_runs.succeeded" },
|
||||
meta: { kind: run.kind },
|
||||
});
|
||||
return "succeeded";
|
||||
}
|
||||
await failAiRun(
|
||||
run.id,
|
||||
{
|
||||
status: outcome.status,
|
||||
errorCode: outcome.errorCode,
|
||||
message: outcome.message,
|
||||
},
|
||||
"RUNNING",
|
||||
new Date(now()),
|
||||
);
|
||||
annotate({
|
||||
action: { name: "ai_runs.failed" },
|
||||
meta: { kind: run.kind, errorCode: outcome.errorCode },
|
||||
});
|
||||
return "failed";
|
||||
}
|
||||
|
||||
/** The pg-boss handler; bound under `lockedPass` keyed by run id. */
|
||||
export async function handleDocumentAiRunJobs(
|
||||
jobs: Job<DocumentAiRunPayload>[],
|
||||
): Promise<JobOutcome> {
|
||||
let processed = 0;
|
||||
let failed = 0;
|
||||
let skipped = 0;
|
||||
for (const job of jobs) {
|
||||
const runId = job.data?.runId;
|
||||
if (!runId) {
|
||||
skipped += 1;
|
||||
continue;
|
||||
}
|
||||
const result = await runDocumentAiRun(runId, jobDeadline(job));
|
||||
if (result === "succeeded") processed += 1;
|
||||
else if (result === "failed") failed += 1;
|
||||
else skipped += 1;
|
||||
}
|
||||
return jobDone({ jobs: jobs.length, processed, failed, skipped });
|
||||
}
|
||||
@@ -43,6 +43,10 @@ export type JobFacts = Readonly<Record<string, JobFact>>;
|
||||
*/
|
||||
export const JOB_FACT_ALLOWLIST: ReadonlySet<string> = new Set([
|
||||
"access_tokens_deleted",
|
||||
// v1.40 — the document AI run reaper's three outcomes.
|
||||
"ai_runs_deleted",
|
||||
"ai_runs_timed_out",
|
||||
"ai_runs_worker_unavailable",
|
||||
// An admission delivery that found its import still running from an
|
||||
// earlier delivery of the same job and did nothing.
|
||||
"already_running",
|
||||
|
||||
@@ -32,6 +32,7 @@ import { PERIOD_NARRATIVE_EXPIRE_SECONDS } from "@/lib/jobs/period-narrative-sha
|
||||
import { PR_DETECTION_EXPIRE_SECONDS } from "@/lib/jobs/pr-detection";
|
||||
import { ENCRYPTION_KEY_ROTATE_EXPIRE_SECONDS } from "@/lib/jobs/encryption-key-rotate";
|
||||
import { CONTENT_INDEX_BACKFILL_EXPIRE_SECONDS } from "@/lib/jobs/document-content-index-backfill";
|
||||
import { DOCUMENT_AI_RUN_EXPIRE_SECONDS } from "@/lib/jobs/document-ai-run";
|
||||
import { DATA_BACKUP_SEND_OPTIONS } from "@/lib/jobs/data-backup-policy";
|
||||
import { OFFHOST_BACKUP_EXPIRE_SECONDS } from "@/lib/jobs/offhost-backup";
|
||||
import { BACKUP_RESTORE_EXPIRE_SECONDS } from "@/lib/jobs/backup-restore";
|
||||
@@ -388,6 +389,20 @@ export const QUEUE_RUNTIME: Readonly<Record<string, QueueRuntime>> = {
|
||||
why: "Up to 200 documents of one account, a provider transcription each.",
|
||||
},
|
||||
"document-index": short("One document."),
|
||||
"document-ai-run": {
|
||||
runtime: "long",
|
||||
expireInSeconds: DOCUMENT_AI_RUN_EXPIRE_SECONDS,
|
||||
expiryVia: ["DOCUMENT_AI_RUN_EXPIRE_SECONDS"],
|
||||
stop: {
|
||||
file: "lib/jobs/document-ai-run.ts",
|
||||
fn: "handleDocumentAiRunJobs",
|
||||
},
|
||||
exclusive: "lockedPass",
|
||||
why: "One person's document or lab scan, up to three model calls at their AI response time (up to ten minutes each).",
|
||||
},
|
||||
"document-ai-run-reaper": short(
|
||||
"Up to 500 conditional updates and one deleteMany.",
|
||||
),
|
||||
"document-thumbnail": short("One document."),
|
||||
"document-summary": short("One document."),
|
||||
"document-summary-catchup": short("Enqueues at most 200 documents."),
|
||||
|
||||
@@ -120,6 +120,18 @@ import {
|
||||
runDocumentIndex,
|
||||
type DocumentIndexPayload,
|
||||
} from "@/lib/jobs/document-index";
|
||||
import {
|
||||
DOCUMENT_AI_RUN_QUEUE,
|
||||
DOCUMENT_AI_RUN_CONCURRENCY,
|
||||
handleDocumentAiRunJobs,
|
||||
type DocumentAiRunPayload,
|
||||
} from "@/lib/jobs/document-ai-run";
|
||||
import {
|
||||
DOCUMENT_AI_RUN_REAPER_QUEUE,
|
||||
DOCUMENT_AI_RUN_REAPER_CRON,
|
||||
handleDocumentAiRunReaper,
|
||||
type DocumentAiRunReaperPayload,
|
||||
} from "@/lib/jobs/document-ai-run-reaper";
|
||||
import {
|
||||
DOCUMENT_THUMBNAIL_QUEUE,
|
||||
DOCUMENT_THUMBNAIL_CONCURRENCY,
|
||||
@@ -489,6 +501,14 @@ const allQueues = [
|
||||
// extraction. Without this entry pg-boss never provisions the queue and every
|
||||
// upload enqueue silently no-ops.
|
||||
DOCUMENT_INDEX_QUEUE,
|
||||
// v1.40 — background document AI runs ("Read with AI", lab scans): the
|
||||
// request answers 202 and this queue does the read. Without this entry
|
||||
// pg-boss never provisions the queue and every run waits until the reaper
|
||||
// fails it as worker-unavailable.
|
||||
DOCUMENT_AI_RUN_QUEUE,
|
||||
// v1.40 — five-minute reaper for those runs: fails a run no worker took or
|
||||
// one past its deadline, deletes a finished one after an hour.
|
||||
DOCUMENT_AI_RUN_REAPER_QUEUE,
|
||||
// Document vault — automatic per-document preview thumbnail, enqueued on
|
||||
// upload. Pure local compute (canvas/pdfjs downscale), no egress. Without
|
||||
// this entry pg-boss never provisions the queue and every upload enqueue
|
||||
@@ -646,6 +666,9 @@ const schedules: ScheduleEntry[] = [
|
||||
// tick is the retry: the updateMany is idempotent and the TTL predicate
|
||||
// re-selects anything a failed tick left behind.
|
||||
[DOCUMENT_SUMMARY_REAPER_QUEUE, DOCUMENT_SUMMARY_REAPER_CRON, cronIsTheRetry],
|
||||
// v1.40 — five-minute pass over background document AI runs. The next tick
|
||||
// is the retry: every write is conditional on the state it leaves.
|
||||
[DOCUMENT_AI_RUN_REAPER_QUEUE, DOCUMENT_AI_RUN_REAPER_CRON, cronIsTheRetry],
|
||||
// v1.32.1 (issue #588) — every-15-minute orphan-ImportJob sweep. Re-runs
|
||||
// the same reconcile the boot path uses, so a stuck "unpacking" row
|
||||
// whose worker crashed/restarted without the boot-time pass catching it
|
||||
@@ -935,6 +958,27 @@ export async function registerMaintenanceQueues(
|
||||
{ localConcurrency: 1 },
|
||||
handleDocumentPurge,
|
||||
);
|
||||
// v1.40 — background document AI runs. Each job is one run, under the lock
|
||||
// `lockedPass` takes per run id, so a redelivery never reads beside the
|
||||
// first delivery; the claim inside is conditional as well.
|
||||
await createAndWork<DocumentAiRunPayload>(
|
||||
boss,
|
||||
DOCUMENT_AI_RUN_QUEUE,
|
||||
{ localConcurrency: DOCUMENT_AI_RUN_CONCURRENCY },
|
||||
lockedPass(
|
||||
DOCUMENT_AI_RUN_QUEUE,
|
||||
(job) => `run:${job.data.runId ?? ""}`,
|
||||
handleDocumentAiRunJobs,
|
||||
),
|
||||
);
|
||||
// v1.40 — the five-minute run reaper. Single-flight; every write in it is
|
||||
// conditional, so a duplicate tick is a no-op.
|
||||
await createAndWork<DocumentAiRunReaperPayload>(
|
||||
boss,
|
||||
DOCUMENT_AI_RUN_REAPER_QUEUE,
|
||||
{ localConcurrency: 1 },
|
||||
handleDocumentAiRunReaper,
|
||||
);
|
||||
// Document vault — hourly stale-PENDING summary reaper. Single-flight;
|
||||
// the underlying updateMany is idempotent so a duplicate tick is a no-op.
|
||||
await createAndWork<DocumentSummaryReaperPayload>(
|
||||
|
||||
@@ -134,6 +134,9 @@ const subsystemSurface = {
|
||||
"document-summary-reaper": { audience: "system" },
|
||||
"document-content-index-backfill": { audience: "account" },
|
||||
"document-index": { audience: "account" },
|
||||
// v1.40 — background document AI runs and their reaper.
|
||||
"document-ai-run": { audience: "account" },
|
||||
"document-ai-run-reaper": { audience: "system" },
|
||||
"document-thumbnail": { audience: "account" },
|
||||
"document-thumbnail-backfill": { audience: "account" },
|
||||
"document-summary": { audience: "account" },
|
||||
|
||||
@@ -0,0 +1,201 @@
|
||||
/**
|
||||
* The lab scan read the background worker runs (v1.40), moved out of the
|
||||
* route with the guarantees the route's own tests used to pin: the budget is
|
||||
* settled against what the provider was asked to do, a PDF is rendered for a
|
||||
* provider that cannot read one, a long PDF says which pages were read, and a
|
||||
* provider failure is named as one without the upstream body reaching the
|
||||
* stored failure or the event stream.
|
||||
*/
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("@/lib/logging/context", () => ({ annotate: vi.fn() }));
|
||||
vi.mock("@/lib/documents/rasterize-pdf", () => ({ rasterizePdf: vi.fn() }));
|
||||
vi.mock("@/lib/ai/coach/budget", () => ({
|
||||
reconcileSpend: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
vi.mock("@/lib/labs/ocr-extract", async () => {
|
||||
const actual = await vi.importActual<typeof import("@/lib/labs/ocr-extract")>(
|
||||
"@/lib/labs/ocr-extract",
|
||||
);
|
||||
return { ...actual, runOcrExtraction: vi.fn() };
|
||||
});
|
||||
|
||||
import { reconcileSpend } from "@/lib/ai/coach/budget";
|
||||
import { rasterizePdf } from "@/lib/documents/rasterize-pdf";
|
||||
import { OcrExtractError, runOcrExtraction } from "@/lib/labs/ocr-extract";
|
||||
import { executeOcrExtraction } from "@/lib/labs/ocr-run";
|
||||
import { annotate } from "@/lib/logging/context";
|
||||
|
||||
const budget = {
|
||||
reserved: 1500,
|
||||
owner: "operator" as const,
|
||||
dateKey: "2027-01-15",
|
||||
};
|
||||
const textPick = {
|
||||
entry: { providerType: "openai", instance: {} as never },
|
||||
providerType: "openai",
|
||||
} as never;
|
||||
const codexPick = {
|
||||
entry: { providerType: "codex", instance: {} as never },
|
||||
providerType: "codex",
|
||||
pdfSupported: false,
|
||||
} as never;
|
||||
|
||||
const text = () => ({ mode: "text" as const, text: "Hb 14", pick: textPick });
|
||||
const pdf = () => ({
|
||||
mode: "vision" as const,
|
||||
bytes: Buffer.from("%PDF-1.4"),
|
||||
mime: "application/pdf" as const,
|
||||
pick: codexPick,
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe("text mode", () => {
|
||||
it("charges the reservation to the provider that served", async () => {
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
const out = await executeOcrExtraction({
|
||||
userId: "u1",
|
||||
input: text(),
|
||||
budget,
|
||||
});
|
||||
expect(out).toEqual({ ok: true, data: { rows: [] } });
|
||||
expect(reconcileSpend).toHaveBeenCalledWith(
|
||||
"u1",
|
||||
1500,
|
||||
1500,
|
||||
"2027-01-15",
|
||||
0,
|
||||
{
|
||||
servedBy: "openai",
|
||||
reservedOwner: "operator",
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("refunds the reservation in full on a clean extract failure", async () => {
|
||||
vi.mocked(runOcrExtraction).mockRejectedValue(new OcrExtractError("x"));
|
||||
const out = await executeOcrExtraction({
|
||||
userId: "u1",
|
||||
input: text(),
|
||||
budget,
|
||||
});
|
||||
expect(out).toMatchObject({
|
||||
ok: false,
|
||||
status: 422,
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
});
|
||||
expect(reconcileSpend).toHaveBeenCalledWith(
|
||||
"u1",
|
||||
1500,
|
||||
0,
|
||||
"2027-01-15",
|
||||
0,
|
||||
{
|
||||
servedBy: null,
|
||||
reservedOwner: "operator",
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("names a provider failure and never keeps the upstream body", async () => {
|
||||
vi.mocked(runOcrExtraction).mockRejectedValue(
|
||||
Object.assign(new Error("Local AI request failed (403)"), {
|
||||
httpStatus: 403,
|
||||
model: "llama3",
|
||||
bodyExcerpt: "ami-id iam/security-credentials/admin",
|
||||
}),
|
||||
);
|
||||
const out = await executeOcrExtraction({
|
||||
userId: "u1",
|
||||
input: text(),
|
||||
budget,
|
||||
});
|
||||
expect(out).toMatchObject({ ok: false, status: 502 });
|
||||
expect(JSON.stringify(out)).toContain("configured AI provider");
|
||||
expect(JSON.stringify(out)).not.toContain("iam/security-credentials");
|
||||
expect(JSON.stringify(vi.mocked(annotate).mock.calls)).not.toContain(
|
||||
"iam/security-credentials",
|
||||
);
|
||||
expect(annotate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
meta: expect.objectContaining({ upstreamStatus: 403, model: "llama3" }),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("vision mode", () => {
|
||||
it("renders a PDF for a provider without native PDF and sends the pages", async () => {
|
||||
vi.mocked(rasterizePdf).mockResolvedValue({
|
||||
ok: true,
|
||||
images: [{ mediaType: "image/jpeg", dataBase64: "cGFnZQ==" }],
|
||||
pageCount: 1,
|
||||
});
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
const out = await executeOcrExtraction({
|
||||
userId: "u1",
|
||||
input: pdf(),
|
||||
budget,
|
||||
});
|
||||
expect(out).toEqual({ ok: true, data: { rows: [] } });
|
||||
expect(runOcrExtraction).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
images: [{ mediaType: "image/jpeg", dataBase64: "cGFnZQ==" }],
|
||||
documents: [],
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("says when only the first pages of a long PDF were read", async () => {
|
||||
vi.mocked(rasterizePdf).mockResolvedValue({
|
||||
ok: true,
|
||||
images: Array.from({ length: 10 }, () => ({
|
||||
mediaType: "image/jpeg" as const,
|
||||
dataBase64: "cGFnZQ==",
|
||||
})),
|
||||
pageCount: 23,
|
||||
});
|
||||
vi.mocked(runOcrExtraction).mockResolvedValue({ rows: [] } as never);
|
||||
const out = await executeOcrExtraction({
|
||||
userId: "u1",
|
||||
input: pdf(),
|
||||
budget,
|
||||
});
|
||||
expect(out).toEqual({
|
||||
ok: true,
|
||||
data: { rows: [], pageCoverage: { read: 10, total: 23 } },
|
||||
});
|
||||
});
|
||||
|
||||
it("falls back to pdfNeedsAnthropic and refunds when rendering fails", async () => {
|
||||
vi.mocked(rasterizePdf).mockResolvedValue({
|
||||
ok: false,
|
||||
reason: "render-failed",
|
||||
});
|
||||
const out = await executeOcrExtraction({
|
||||
userId: "u1",
|
||||
input: pdf(),
|
||||
budget,
|
||||
});
|
||||
expect(out).toMatchObject({
|
||||
ok: false,
|
||||
status: 422,
|
||||
errorCode: "labs.ocr.pdfNeedsAnthropic",
|
||||
});
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
expect(reconcileSpend).toHaveBeenCalledWith(
|
||||
"u1",
|
||||
1500,
|
||||
0,
|
||||
"2027-01-15",
|
||||
0,
|
||||
{
|
||||
servedBy: null,
|
||||
reservedOwner: "operator",
|
||||
},
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,200 @@
|
||||
/**
|
||||
* A lab report scan, run in the background worker.
|
||||
*
|
||||
* `POST /api/labs/ocr/extract` checks the capability, the rate bucket and the
|
||||
* budget, validates the upload, seals it into a `DocumentAiRun` and answers
|
||||
* 202. This is the read itself: the same PDF handling the route used to do in
|
||||
* the request (a native document block on a PDF-capable provider, rendered
|
||||
* page images on any other), one extraction with its corrective retry, and the
|
||||
* reservation settled against what the provider was asked to do.
|
||||
*
|
||||
* Nothing is written to the person's record here. The proposed rows go back
|
||||
* through the run for the mandatory review screen, and the commit route stays
|
||||
* the only write path.
|
||||
*/
|
||||
import type { Buffer } from "node:buffer";
|
||||
|
||||
import { AI_BUDGETS } from "@/lib/ai/ai-budgets";
|
||||
import type { ProviderChainType } from "@/lib/ai/provider-chain";
|
||||
import { settleAiRunBudget } from "@/lib/documents/ai-runs/store";
|
||||
import type { AiRunBudget, AiRunOutcome } from "@/lib/documents/ai-runs/types";
|
||||
import { rasterizePdf } from "@/lib/documents/rasterize-pdf";
|
||||
import type { ProviderChainResolved } from "@/lib/ai/provider-runner";
|
||||
import type { VisionProviderPick } from "@/lib/labs/ocr-capability";
|
||||
import { OcrExtractError, runOcrExtraction } from "@/lib/labs/ocr-extract";
|
||||
import { annotate } from "@/lib/logging/context";
|
||||
import type { OcrExtractResponseDto } from "@/lib/validations/labs-ocr";
|
||||
|
||||
/** The most sequential model calls one scan makes: the read and one retry. */
|
||||
export const OCR_EXTRACT_MODEL_CALLS = 2;
|
||||
|
||||
/** The tokens a scan reserves at enqueue, by mode. */
|
||||
export const OCR_RESERVE_TOKENS = {
|
||||
vision: AI_BUDGETS.ocrExtract.maxTokens,
|
||||
text: AI_BUDGETS.ocrExtractText.maxTokens,
|
||||
} as const;
|
||||
|
||||
export type OcrVisionPick = NonNullable<VisionProviderPick["pick"]>;
|
||||
export interface OcrTextPick {
|
||||
entry: ProviderChainResolved;
|
||||
providerType: string;
|
||||
}
|
||||
|
||||
export type OcrRunInput =
|
||||
| { mode: "text"; text: string; pick: OcrTextPick }
|
||||
| {
|
||||
mode: "vision";
|
||||
bytes: Buffer;
|
||||
mime: "image/jpeg" | "image/png" | "image/webp" | "application/pdf";
|
||||
pick: OcrVisionPick;
|
||||
};
|
||||
|
||||
/**
|
||||
* The wide-event facts of a failed provider call: the status and the model,
|
||||
* never the upstream body. The base URL can be one a user typed, and what an
|
||||
* arbitrary endpoint answers has no business in the event stream (v1.39.3);
|
||||
* the stored failure carries neither.
|
||||
*/
|
||||
function providerFailureMeta(error: unknown, mode?: "text") {
|
||||
const err = error as { httpStatus?: unknown; model?: unknown };
|
||||
return {
|
||||
reason: "provider_error",
|
||||
...(mode ? { mode } : {}),
|
||||
...(typeof err.httpStatus === "number"
|
||||
? { upstreamStatus: err.httpStatus }
|
||||
: {}),
|
||||
...(typeof err.model === "string" ? { model: err.model } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
/** A failure in the shape the route would answer with (`errorCode` literal, so the catalogue guard reads it). */
|
||||
const failure = (f: {
|
||||
status: number;
|
||||
message: string;
|
||||
errorCode: string;
|
||||
}): AiRunOutcome<never> => ({ ok: false, ...f });
|
||||
|
||||
/** Read one lab report into proposed rows. The caller authorised the read. */
|
||||
export async function executeOcrExtraction(args: {
|
||||
userId: string;
|
||||
input: OcrRunInput;
|
||||
budget: AiRunBudget | undefined;
|
||||
}): Promise<AiRunOutcome<OcrExtractResponseDto>> {
|
||||
const { userId, input, budget } = args;
|
||||
const servedBy = input.pick.entry.providerType as ProviderChainType;
|
||||
|
||||
if (input.mode === "text") {
|
||||
try {
|
||||
const result = await runOcrExtraction({
|
||||
userId,
|
||||
provider: input.pick.entry.instance,
|
||||
providerType: input.pick.providerType,
|
||||
ocrText: input.text,
|
||||
});
|
||||
await settleAiRunBudget(userId, budget, budget?.reserved ?? 0, servedBy);
|
||||
return { ok: true, data: result };
|
||||
} catch (err) {
|
||||
// A failed structuring call produced no usable rows; the reservation
|
||||
// goes back in full rather than being charged.
|
||||
await settleAiRunBudget(userId, budget, 0, null);
|
||||
if (err instanceof OcrExtractError) {
|
||||
return failure({
|
||||
status: 422,
|
||||
message: "Couldn't read the report. Try a clearer photo.",
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
});
|
||||
}
|
||||
annotate({
|
||||
action: { name: "labs.ocr.extractFailed" },
|
||||
meta: providerFailureMeta(err, "text"),
|
||||
});
|
||||
return failure({
|
||||
status: 502,
|
||||
message:
|
||||
"The configured AI provider could not process this report. Check the provider configuration and retry.",
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
const { pick, bytes, mime } = input;
|
||||
let images: {
|
||||
mediaType: "image/jpeg" | "image/png" | "image/webp";
|
||||
dataBase64: string;
|
||||
}[];
|
||||
let documents: { mediaType: "application/pdf"; dataBase64: string }[];
|
||||
let pageCoverage: { read: number; total: number } | undefined;
|
||||
|
||||
if (mime === "application/pdf") {
|
||||
if (pick.pdfSupported) {
|
||||
images = [];
|
||||
documents = [
|
||||
{ mediaType: "application/pdf", dataBase64: bytes.toString("base64") },
|
||||
];
|
||||
} else {
|
||||
const raster = await rasterizePdf(bytes);
|
||||
if (!raster.ok) {
|
||||
annotate({
|
||||
action: { name: "labs.ocr.fileRejected" },
|
||||
meta: { reason: "pdf_rasterize_failed" },
|
||||
});
|
||||
await settleAiRunBudget(userId, budget, 0, null);
|
||||
return failure({
|
||||
status: 422,
|
||||
message: "Couldn't read this PDF; upload a photo instead.",
|
||||
errorCode: "labs.ocr.pdfNeedsAnthropic",
|
||||
});
|
||||
}
|
||||
images = raster.images;
|
||||
documents = [];
|
||||
// A PDF longer than the page cap is read from its first pages only, and
|
||||
// the person reviewing the rows is the one who needs to know.
|
||||
if (raster.pageCount > raster.images.length) {
|
||||
pageCoverage = { read: raster.images.length, total: raster.pageCount };
|
||||
annotate({
|
||||
action: { name: "labs.ocr.pagesCapped" },
|
||||
meta: { read: raster.images.length, total: raster.pageCount },
|
||||
});
|
||||
}
|
||||
}
|
||||
} else {
|
||||
images = [{ mediaType: mime, dataBase64: bytes.toString("base64") }];
|
||||
documents = [];
|
||||
}
|
||||
|
||||
try {
|
||||
const result = await runOcrExtraction({
|
||||
userId,
|
||||
provider: pick.entry.instance,
|
||||
providerType: pick.providerType,
|
||||
images,
|
||||
documents,
|
||||
});
|
||||
// The orchestration does not surface token counts; the reserved estimate
|
||||
// is the spend ceiling (the provider already billed it).
|
||||
await settleAiRunBudget(userId, budget, budget?.reserved ?? 0, servedBy);
|
||||
return {
|
||||
ok: true,
|
||||
data: pageCoverage ? { ...result, pageCoverage } : result,
|
||||
};
|
||||
} catch (err) {
|
||||
await settleAiRunBudget(userId, budget, 0, servedBy);
|
||||
if (err instanceof OcrExtractError) {
|
||||
return failure({
|
||||
status: 422,
|
||||
message: "Couldn't read the report. Try a clearer photo.",
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
});
|
||||
}
|
||||
annotate({
|
||||
action: { name: "labs.ocr.extractFailed" },
|
||||
meta: providerFailureMeta(err),
|
||||
});
|
||||
return failure({
|
||||
status: 502,
|
||||
message:
|
||||
"The configured AI provider could not process this report. Check the provider configuration and retry.",
|
||||
errorCode: "labs.ocr.extractFailed",
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,12 @@ export const ERROR_CODE_CATALOGUE: Readonly<Record<string, readonly string[]>> =
|
||||
"ai.unavailable",
|
||||
],
|
||||
ai_provider: ["ai_provider.invalid", "ai_provider.no_fields"],
|
||||
aiRuns: [
|
||||
"aiRuns.failed",
|
||||
"aiRuns.notFound",
|
||||
"aiRuns.timedOut",
|
||||
"aiRuns.workerUnavailable",
|
||||
],
|
||||
allergy: ["allergy.invalid"],
|
||||
analytics: ["analytics.invalid_query"],
|
||||
anamnesis: [
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
/**
|
||||
* The request header and the 202 response every route that queues a
|
||||
* background document AI run shares (v1.40). See `./ai-runs.ts` for the poll.
|
||||
*/
|
||||
import { z } from "zod/v4";
|
||||
|
||||
import { dataEnvelope, errorEnvelope } from "./shared";
|
||||
|
||||
/** `Prefer: respond-async` on a route that offers the background path. */
|
||||
export const preferRespondAsyncParameter = {
|
||||
name: "Prefer",
|
||||
in: "header" as const,
|
||||
required: false,
|
||||
schema: { type: "string" as const, 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`.",
|
||||
};
|
||||
|
||||
const acceptedBody = z
|
||||
.object({
|
||||
runId: z.string(),
|
||||
status: z.literal("QUEUED"),
|
||||
pollAfterMs: z
|
||||
.number()
|
||||
.int()
|
||||
.describe("How long to wait before the first poll."),
|
||||
})
|
||||
.meta({
|
||||
id: "AiRunAccepted",
|
||||
description:
|
||||
"The read was queued. Poll `GET /api/ai-runs/{runId}`. A second request for the same document while a read of it is still queued or running answers with that run instead of starting another.",
|
||||
});
|
||||
|
||||
const acceptedHeaders = z.object({
|
||||
"Preference-Applied": z.literal("respond-async").meta({
|
||||
description: "The route honoured `Prefer: respond-async`.",
|
||||
}),
|
||||
Location: z.string().meta({ description: "The run's poll address." }),
|
||||
});
|
||||
|
||||
/** The 202 a queuing route answers with. */
|
||||
export const aiRunAccepted = {
|
||||
"202": {
|
||||
description:
|
||||
"Queued for the background worker. Every refusal the route can answer quickly (capability, provider, rate limit, budget, validation) has already been answered; anything that goes wrong from here is reported on the run.",
|
||||
content: {
|
||||
"application/json": {
|
||||
schema: dataEnvelope(acceptedBody, "AiRunAcceptedEnvelope"),
|
||||
},
|
||||
},
|
||||
headers: acceptedHeaders,
|
||||
},
|
||||
"503": {
|
||||
description:
|
||||
"`aiRuns.workerUnavailable`: the run could not be handed to the background worker. Nothing was read and nothing was charged; try again.",
|
||||
content: { "application/json": { schema: errorEnvelope } },
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,98 @@
|
||||
/**
|
||||
* OpenAPI route table for background document AI runs (`/api/ai-runs/*`).
|
||||
* The request header and the 202 the queuing routes share live in
|
||||
* `./ai-run-accepted.ts`, so this module can read the result schemas of the
|
||||
* routes it serves without an import cycle.
|
||||
*
|
||||
* Part of the OpenAPI route table; aggregated in `./index.ts`.
|
||||
*
|
||||
* iOS coordination (v1.40): `POST /api/documents/inbound/{id}/index` keeps
|
||||
* its 200 body 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.
|
||||
* `POST /api/labs/ocr/extract` (web only) always answers 202.
|
||||
*/
|
||||
import type { ZodOpenApiObject } from "zod-openapi";
|
||||
import { z } from "zod/v4";
|
||||
|
||||
import { documentIndexResponse } from "./documents";
|
||||
import { ocrExtractResponse } from "./ocr";
|
||||
import { dataEnvelope, errorEnvelope, stdResponses } from "./shared";
|
||||
|
||||
const aiRun = z
|
||||
.object({
|
||||
id: z.string(),
|
||||
kind: z.enum(["DOCUMENT_INDEX", "LABS_OCR_EXTRACT"]),
|
||||
documentId: z
|
||||
.string()
|
||||
.nullable()
|
||||
.describe("The document a document run reads; null for a lab scan."),
|
||||
status: z.enum(["QUEUED", "RUNNING", "SUCCEEDED", "FAILED"]),
|
||||
result: z
|
||||
.union([documentIndexResponse, ocrExtractResponse])
|
||||
.nullable()
|
||||
.describe(
|
||||
"Present once `status` is `SUCCEEDED`: the body the queuing route answers with synchronously. `DOCUMENT_INDEX` carries `DocumentIndexResponse`, `LABS_OCR_EXTRACT` carries `OcrExtractResponse`.",
|
||||
),
|
||||
error: z
|
||||
.object({
|
||||
code: z.string(),
|
||||
message: z.string(),
|
||||
status: z.number().int(),
|
||||
})
|
||||
.nullable()
|
||||
.describe(
|
||||
"Present once `status` is `FAILED`: the `meta.errorCode`, message and HTTP status the queuing route would have answered with, or one of the run's own codes: `aiRuns.workerUnavailable` (no worker took the run within fifteen minutes), `aiRuns.timedOut` (the read outlived the calls the AI response time allows), `aiRuns.failed` (anything else).",
|
||||
),
|
||||
createdAt: z.string(),
|
||||
startedAt: z.string().nullable(),
|
||||
finishedAt: z.string().nullable(),
|
||||
queuedForMs: z
|
||||
.number()
|
||||
.int()
|
||||
.nullable()
|
||||
.describe(
|
||||
"While queued, how long the run has waited for a worker. A long wait usually means the background worker is not running.",
|
||||
),
|
||||
retryAfterMs: z
|
||||
.number()
|
||||
.int()
|
||||
.nullable()
|
||||
.describe(
|
||||
"While not terminal, when to poll again (also sent as `Retry-After`, in seconds). Null once the run has ended.",
|
||||
),
|
||||
})
|
||||
.meta({
|
||||
id: "AiRun",
|
||||
description:
|
||||
"One background document AI run. Results and failures stay readable for an hour after the run ends; the run is then deleted and answers 404.",
|
||||
});
|
||||
|
||||
export const aiRunPaths: NonNullable<ZodOpenApiObject["paths"]> = {
|
||||
"/api/ai-runs/{id}": {
|
||||
get: {
|
||||
tags: ["Documents"],
|
||||
summary: "Poll a background document AI run",
|
||||
description:
|
||||
"Where one background read stands, and its result or failure once it ends. Scoped to the caller: another account's run, a run that never existed, and a run more than an hour past its end are the same 404 `aiRuns.notFound`. A plain read of the caller's own run; it calls no model and asks no AI capability (those were asked when the run was queued, and again in the worker before anything left the server). Auth via cookie or Bearer.",
|
||||
parameters: [
|
||||
{ name: "id", in: "path", required: true, schema: { type: "string" } },
|
||||
],
|
||||
responses: {
|
||||
"200": {
|
||||
description: "The run.",
|
||||
content: {
|
||||
"application/json": {
|
||||
schema: dataEnvelope(aiRun, "AiRunEnvelope"),
|
||||
},
|
||||
},
|
||||
},
|
||||
"404": {
|
||||
description: "`aiRuns.notFound`.",
|
||||
content: { "application/json": { schema: errorEnvelope } },
|
||||
},
|
||||
...stdResponses,
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -32,6 +32,7 @@ import {
|
||||
} from "@/lib/validations/inbound-documents";
|
||||
|
||||
import { aiExtractionRefusals } from "./ai-extraction-refusals";
|
||||
import { aiRunAccepted, preferRespondAsyncParameter } from "./ai-run-accepted";
|
||||
import { aiCapabilityState } from "./profile";
|
||||
import {
|
||||
dataEnvelope,
|
||||
@@ -286,6 +287,20 @@ const documentUploadEnvelope = dataEnvelope(
|
||||
"DocumentUploadEnvelope",
|
||||
);
|
||||
|
||||
/** The index route's body; also a `DOCUMENT_INDEX` run's result. */
|
||||
export const documentIndexResponse = z
|
||||
.object({
|
||||
documentId: z.string(),
|
||||
indexed: z.boolean(),
|
||||
tokenCount: z.number(),
|
||||
labFactsStaged: z
|
||||
.number()
|
||||
.describe(
|
||||
"How many lab facts the index run auto-staged for review (0 when the document is not a lab report, staging was not eligible, or facts already exist). Staged facts are PENDING — the confirm route remains the only write into Labs.",
|
||||
),
|
||||
})
|
||||
.meta({ id: "DocumentIndexResponse" });
|
||||
|
||||
export const inboundDocumentPaths: NonNullable<ZodOpenApiObject["paths"]> = {
|
||||
"/api/documents/inbound": {
|
||||
get: {
|
||||
@@ -1060,9 +1075,10 @@ export const inboundDocumentPaths: NonNullable<ZodOpenApiObject["paths"]> = {
|
||||
tags: ["Documents"],
|
||||
summary: "Build / refresh the content-search index",
|
||||
description:
|
||||
"Populates or refreshes one document's content-search index so search matches INSIDE its body. VISION (empty body) decrypts the stored original and runs one provider transcription (answers under the `documentAi` capability — 403 with the capability envelope when it is closed, including a missing `ai_extraction` / `ai_full` receipt for a provider that leaves the machine — then rate / budget gated); TEXT (`application/json` `{ mode: \"text\", text }`, opt-in local OCR) indexes browser-OCR'd text with no provider egress and stays available with AI off, so search keeps working. Persists ONLY AES-256-GCM ciphertext of the text plus opaque HMAC token hashes — no plaintext body, no plaintext token. Idempotent; re-indexing overwrites in place. A successful index continues into the same lab auto-staging the background worker performs (still-STORED lab-looking document with no facts, both modules on, provider + consent eligible); `labFactsStaged` reports how many facts were staged PENDING for review — confirmation stays the only write into Labs.",
|
||||
"Populates or refreshes one document's content-search index so search matches INSIDE its body. VISION (empty body) decrypts the stored original and runs one provider transcription (answers under the `documentAi` capability — 403 with the capability envelope when it is closed, including a missing `ai_extraction` / `ai_full` receipt for a provider that leaves the machine — then rate / budget gated); TEXT (`application/json` `{ mode: \"text\", text }`, opt-in local OCR) indexes browser-OCR'd text with no provider egress and stays available with AI off, so search keeps working. Persists ONLY AES-256-GCM ciphertext of the text plus opaque HMAC token hashes — no plaintext body, no plaintext token. Idempotent; re-indexing overwrites in place. A successful index continues into the same lab auto-staging the background worker performs (still-STORED lab-looking document with no facts, both modules on, provider + consent eligible); `labFactsStaged` reports how many facts were staged PENDING for review — confirmation stays the only write into Labs. With `Prefer: respond-async` (v1.40) the read runs in the background: after the same quick refusals the route answers 202 with a run id, a second request while a read of this document is queued or running answers with that run, and `GET /api/ai-runs/{id}` serves this body as the run's `result`. Text mode honours the preference too, because the lab staging it continues into is model work.",
|
||||
parameters: [
|
||||
{ name: "id", in: "path", required: true, schema: { type: "string" } },
|
||||
preferRespondAsyncParameter,
|
||||
],
|
||||
requestBody: {
|
||||
required: false,
|
||||
@@ -1082,23 +1098,13 @@ export const inboundDocumentPaths: NonNullable<ZodOpenApiObject["paths"]> = {
|
||||
content: {
|
||||
"application/json": {
|
||||
schema: dataEnvelope(
|
||||
z
|
||||
.object({
|
||||
documentId: z.string(),
|
||||
indexed: z.boolean(),
|
||||
tokenCount: z.number(),
|
||||
labFactsStaged: z
|
||||
.number()
|
||||
.describe(
|
||||
"How many lab facts the index run auto-staged for review (0 when the document is not a lab report, staging was not eligible, or facts already exist). Staged facts are PENDING — the confirm route remains the only write into Labs.",
|
||||
),
|
||||
})
|
||||
.meta({ id: "DocumentIndexResponse" }),
|
||||
documentIndexResponse,
|
||||
"DocumentIndexEnvelope",
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
...aiRunAccepted,
|
||||
...aiExtractionRefusals("documentAi"),
|
||||
...stdResponses,
|
||||
},
|
||||
|
||||
@@ -64,6 +64,7 @@ import { integrationPaths } from "./integrations";
|
||||
import { insightsSignalPaths } from "./insights-signals";
|
||||
import { labsPaths } from "./labs";
|
||||
import { ocrPaths } from "./ocr";
|
||||
import { aiRunPaths } from "./ai-runs";
|
||||
import { measurementPaths } from "./measurements";
|
||||
import { measurementReminderPaths } from "./measurement-reminders";
|
||||
import { mentalHealthPaths } from "./mental-health";
|
||||
@@ -115,6 +116,7 @@ export const openApiPaths: NonNullable<ZodOpenApiObject["paths"]> = {
|
||||
...measurementReminderPaths,
|
||||
...labsPaths,
|
||||
...ocrPaths,
|
||||
...aiRunPaths,
|
||||
...biomarkerPaths,
|
||||
...encounterPaths,
|
||||
...vaccinationPaths,
|
||||
|
||||
@@ -18,6 +18,7 @@ import { z } from "zod/v4";
|
||||
import { ocrCommitSchema } from "@/lib/validations/labs-ocr";
|
||||
|
||||
import { aiExtractionRefusals } from "./ai-extraction-refusals";
|
||||
import { aiRunAccepted } from "./ai-run-accepted";
|
||||
import { aiCapabilityState } from "./profile";
|
||||
import {
|
||||
dataEnvelope,
|
||||
@@ -80,7 +81,7 @@ const extractedRow = z
|
||||
"One proposed reading transcribed from the upload. UNTRUSTED model output annotated server-side: `biomarkerMatch` flags whether the analyte links an existing catalog marker; `duplicateOf` is the id of a live reading this row likely duplicates (the review row defaults to unchecked when set); `confidence` is the model's per-field self-score (low fields are flagged for the human). Nothing is written until the commit route confirms.",
|
||||
});
|
||||
|
||||
const extractResponse = z
|
||||
export const ocrExtractResponse = z
|
||||
.object({
|
||||
reportDate: z.string().nullable(),
|
||||
providerType: z.string(),
|
||||
@@ -98,7 +99,7 @@ const extractResponse = z
|
||||
.meta({
|
||||
id: "OcrExtractResponse",
|
||||
description:
|
||||
"The proposed rows for the human review screen. NEVER written to the database — extraction is read-only and the raw upload is held in memory only.",
|
||||
"The proposed rows for the human review screen, served by `GET /api/ai-runs/{id}` for a `LABS_OCR_EXTRACT` run. Never written to the person's record; the commit route is the only write.",
|
||||
});
|
||||
|
||||
const committedRow = z
|
||||
@@ -177,7 +178,7 @@ export const ocrPaths: NonNullable<ZodOpenApiObject["paths"]> = {
|
||||
tags: ["Labs"],
|
||||
summary: "Extract lab readings from a photo, PDF, or OCR'd text",
|
||||
description:
|
||||
"Read-only (NOT idempotent) extraction. Two modes by content-type. VISION (`multipart/form-data`): a `file` (JPEG/PNG/WebP, or PDF — read natively on Anthropic, else rasterized to page images for any other vision provider; ≤ 12 MiB) is run through the user's vision-capable provider; the upload lives in memory only and is never persisted or logged. TEXT (`application/json`, opt-in local OCR): the browser OCR's the image (tesseract.js) and POSTs `{ mode: \"text\", text }` — only the extracted text reaches the server, so a text-only provider (e.g. ChatGPT-OAuth) reaches the same review/commit flow. Both modes are model work and answer under the `labsOcr` capability (the operator's reading-documents switch, the labs module, and — because a lab report is a document — an active `ai_extraction` or `ai_full` consent receipt for any provider that leaves the machine, checked for the provider actually picked: the first vision-capable entry in vision mode, the chain head in text mode). No provider is 422 `labs.ocr.providerUnsupported`; text mode without the local-OCR opt-in is 422 `labs.ocr.localOcrDisabled`. Then both pass a per-user hourly rate bucket (default 6, operator-tunable via `LABS_OCR_LIMIT_PER_HOUR`; a slot is only consumed when the scan reaches the provider, and the 429 carries `Retry-After`, the `X-RateLimit-*` triple and `meta.retryAt`) and the per-day token budget, and return proposed rows for the mandatory human review screen — nothing is written. Extracted content is treated as untrusted (prompt-injection); the review step is the safety boundary.",
|
||||
"Read-only (NOT idempotent) extraction. Two modes by content-type. VISION (`multipart/form-data`): a `file` (JPEG/PNG/WebP, or PDF — read natively on Anthropic, else rasterized to page images for any other vision provider; ≤ 12 MiB) is run through the user's vision-capable provider; the upload is never logged. TEXT (`application/json`, opt-in local OCR): the browser OCR's the image (tesseract.js) and POSTs `{ mode: \"text\", text }` — only the extracted text reaches the server, so a text-only provider (e.g. ChatGPT-OAuth) reaches the same review/commit flow. Both modes are model work and answer under the `labsOcr` capability (the operator's reading-documents switch, the labs module, and — because a lab report is a document — an active `ai_extraction` or `ai_full` consent receipt for any provider that leaves the machine, checked for the provider actually picked: the first vision-capable entry in vision mode, the chain head in text mode). No provider is 422 `labs.ocr.providerUnsupported`; text mode without the local-OCR opt-in is 422 `labs.ocr.localOcrDisabled`. Then both pass a per-user hourly rate bucket (default 6, operator-tunable via `LABS_OCR_LIMIT_PER_HOUR`; a slot is only consumed when the scan reaches the provider, and the 429 carries `Retry-After`, the `X-RateLimit-*` triple and `meta.retryAt`) and the per-day token budget, and the read is queued: since v1.40 the route answers 202 with a run id and `GET /api/ai-runs/{id}` serves the proposed rows (`OcrExtractResponse`) for the mandatory human review screen. The upload or text is held encrypted in the run until it is read, then dropped; the run is deleted an hour after it ends. Nothing is written to the record — the commit route is the only write. A file too large is 413 `labs.ocr.fileTooLarge`, an unknown type 415 `labs.ocr.fileType`, both before anything is queued. Extracted content is treated as untrusted (prompt-injection); the review step is the safety boundary.",
|
||||
requestBody: {
|
||||
required: true,
|
||||
content: {
|
||||
@@ -204,14 +205,7 @@ export const ocrPaths: NonNullable<ZodOpenApiObject["paths"]> = {
|
||||
},
|
||||
},
|
||||
responses: {
|
||||
"200": {
|
||||
description: "Proposed rows for review.",
|
||||
content: {
|
||||
"application/json": {
|
||||
schema: dataEnvelope(extractResponse, "OcrExtractEnvelope"),
|
||||
},
|
||||
},
|
||||
},
|
||||
...aiRunAccepted,
|
||||
...aiExtractionRefusals("labsOcr"),
|
||||
...stdResponses,
|
||||
},
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
/**
|
||||
* v1.40 — TanStack Query keys for background document AI runs. A run is
|
||||
* polled until it ends and then dropped from the cache, so the result (model
|
||||
* text, lab values) does not linger in memory past the screen that asked.
|
||||
*/
|
||||
export const aiRunKeys = {
|
||||
aiRuns: () => ["ai-runs"] as const,
|
||||
aiRun: (runId: string) => ["ai-runs", runId] as const,
|
||||
};
|
||||
@@ -15,6 +15,7 @@
|
||||
import type { QueryClient, QueryKey } from "@tanstack/react-query";
|
||||
|
||||
import { adminKeys } from "./admin";
|
||||
import { aiRunKeys } from "./ai-runs";
|
||||
import { allergyKeys } from "./allergies";
|
||||
import { authKeys } from "./auth";
|
||||
import { coachKeys } from "./coach";
|
||||
@@ -74,6 +75,7 @@ export const queryKeys = {
|
||||
...profileKeys,
|
||||
...sharingKeys,
|
||||
...vaccinationKeys,
|
||||
...aiRunKeys,
|
||||
};
|
||||
|
||||
export { recordSettingsKeys } from "./record-settings";
|
||||
|
||||
@@ -63,6 +63,12 @@ vi.mock("@/lib/db-compat", () => ({
|
||||
ensureDbCompatibility: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
// v1.40 — a lab scan is queued for the background worker; a queue that takes
|
||||
// the send lets the route answer 202, and the test then runs the worker.
|
||||
vi.mock("@/lib/jobs/boss-instance", () => ({
|
||||
getGlobalBoss: () => ({ send: vi.fn(async () => "job-1") }),
|
||||
}));
|
||||
|
||||
// ── The spies: every place a model would be called ─────────────────────────
|
||||
|
||||
vi.mock("@/lib/documents/assist", async (importOriginal) => ({
|
||||
@@ -344,9 +350,16 @@ describe("scanning a lab report — POST /api/labs/ocr/extract", () => {
|
||||
await grantConsent(user.id, "ai_extraction");
|
||||
await signIn(user.id);
|
||||
|
||||
const { status } = await labsScanText();
|
||||
const { status, body } = await labsScanText();
|
||||
|
||||
expect(status).toBe(200);
|
||||
// The route only queues the read; the worker is where the model is called.
|
||||
expect(status).toBe(202);
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
const { runDocumentAiRun } = await import("@/lib/jobs/document-ai-run");
|
||||
const { runId } = body.data as { runId: string };
|
||||
expect(await runDocumentAiRun(runId, Date.now() + 45 * 60_000)).toBe(
|
||||
"succeeded",
|
||||
);
|
||||
expect(runOcrExtraction).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
|
||||
@@ -0,0 +1,451 @@
|
||||
/**
|
||||
* Background document AI runs (v1.40) end to end against a real Postgres:
|
||||
* the queuing routes, the worker body, the poll route and the reaper, with a
|
||||
* spy standing in for every model call.
|
||||
*
|
||||
* What only this file can prove:
|
||||
* 1. Migration 0366 applied clean (every write below lands in the table).
|
||||
* 2. Enqueue → worker → poll round-trips the same body the synchronous
|
||||
* route answers with, for "Read with AI" and for a lab scan, and the
|
||||
* synchronous route is unchanged without `Prefer: respond-async`.
|
||||
* 3. The sealed input is ciphertext at rest and is dropped when the run
|
||||
* finishes.
|
||||
* 4. A run is its owner's: another account's poll is the same 404 as a run
|
||||
* that does not exist.
|
||||
* 5. Consent withdrawn while a run waited stops it in the worker, before
|
||||
* 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.
|
||||
*/
|
||||
import { Buffer } from "node:buffer";
|
||||
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
process.env.ENCRYPTION_KEY ??=
|
||||
"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
|
||||
|
||||
import { encrypt } from "@/lib/crypto";
|
||||
|
||||
import { cookieJar, headerJar } from "./mock-next-headers";
|
||||
import { getPrismaClient, truncateAllTables } from "./setup";
|
||||
|
||||
vi.mock("next/headers", async () => {
|
||||
const { cookieJar, headerJar } = await import("./mock-next-headers");
|
||||
return {
|
||||
headers: vi.fn(async () => ({
|
||||
get: (name: string) => headerJar.get(name.toLowerCase()) ?? null,
|
||||
})),
|
||||
cookies: vi.fn(async () => ({
|
||||
get: (name: string) => {
|
||||
const value = cookieJar.get(name);
|
||||
return value ? { name, value } : undefined;
|
||||
},
|
||||
set: (name: string, value: string) => {
|
||||
cookieJar.set(name, value);
|
||||
},
|
||||
delete: (name: string) => {
|
||||
cookieJar.delete(name);
|
||||
},
|
||||
})),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("@/lib/db-compat", () => ({
|
||||
ensureDbCompatibility: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
const queue = vi.hoisted(() => ({
|
||||
up: true,
|
||||
send: vi.fn(async () => "job-1"),
|
||||
}));
|
||||
vi.mock("@/lib/jobs/boss-instance", () => ({
|
||||
getGlobalBoss: () => (queue.up ? { send: queue.send } : null),
|
||||
}));
|
||||
const send = queue.send;
|
||||
|
||||
vi.mock("@/lib/documents/describe", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/lib/documents/describe")>()),
|
||||
transcribeDocument: vi.fn(async () => ({
|
||||
text: "Entlassbrief Leukozyten erhoeht",
|
||||
})),
|
||||
}));
|
||||
vi.mock("@/lib/labs/ocr-extract", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/lib/labs/ocr-extract")>()),
|
||||
runOcrExtraction: vi.fn(async () => ({
|
||||
reportDate: "2027-01-10",
|
||||
providerType: "anthropic",
|
||||
rows: [],
|
||||
})),
|
||||
}));
|
||||
vi.mock("@/lib/documents/auto-stage-labs", () => ({
|
||||
maybeAutoStageLabFacts: vi.fn(async () => ({
|
||||
staged: false,
|
||||
reason: "not-lab",
|
||||
})),
|
||||
}));
|
||||
|
||||
import { transcribeDocument } from "@/lib/documents/describe";
|
||||
import { runOcrExtraction } from "@/lib/labs/ocr-extract";
|
||||
|
||||
const PNG = Buffer.from(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg==",
|
||||
"base64",
|
||||
);
|
||||
|
||||
let counter = 0;
|
||||
|
||||
async function makeUser() {
|
||||
const suffix = `r-${counter++}`;
|
||||
return getPrismaClient().user.create({
|
||||
data: {
|
||||
username: `runs-${suffix}`,
|
||||
email: `runs-${suffix}@example.test`,
|
||||
role: "USER",
|
||||
timezone: "UTC",
|
||||
locale: "en",
|
||||
onboardingCompletedAt: new Date(),
|
||||
aiProvider: "ANTHROPIC",
|
||||
aiModel: "claude-sonnet-4-6",
|
||||
aiAnthropicKeyEncrypted: encrypt("sk-ant-integration-test"),
|
||||
labsLocalOcrEnabled: true,
|
||||
modulePreferencesJson: { inboundDocuments: true, labs: true },
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async function grantConsent(userId: string) {
|
||||
return getPrismaClient().consentReceipt.create({
|
||||
data: {
|
||||
userId,
|
||||
kind: "ai_extraction",
|
||||
artefact: "test",
|
||||
signedAt: new Date(),
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async function signIn(userId: string) {
|
||||
const session = await getPrismaClient().session.create({
|
||||
data: { userId, expiresAt: new Date(Date.now() + 600_000) },
|
||||
});
|
||||
cookieJar.set("healthlog_session", session.id);
|
||||
}
|
||||
|
||||
async function storeDocument(userId: string) {
|
||||
const { encryptDocumentContent } = await import("@/lib/documents/store");
|
||||
const { content, codec } = encryptDocumentContent(PNG);
|
||||
return getPrismaClient().inboundDocument.create({
|
||||
data: {
|
||||
userId,
|
||||
kind: "OTHER",
|
||||
filename: "letter.png",
|
||||
mimeType: "image/png",
|
||||
byteSize: PNG.byteLength,
|
||||
status: "STORED",
|
||||
contentEncrypted: content,
|
||||
contentCodec: codec,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
interface Envelope<T = unknown> {
|
||||
data: T;
|
||||
error: string | null;
|
||||
meta?: Record<string, unknown>;
|
||||
}
|
||||
|
||||
async function postIndex(documentId: string, prefer: boolean) {
|
||||
const { POST } = await import("@/app/api/documents/inbound/[id]/index/route");
|
||||
const { NextRequest } = await import("next/server");
|
||||
const res = await POST(
|
||||
new NextRequest(
|
||||
`http://localhost/api/documents/inbound/${documentId}/index`,
|
||||
{
|
||||
method: "POST",
|
||||
headers: prefer ? { prefer: "respond-async" } : {},
|
||||
},
|
||||
) as never,
|
||||
{ params: Promise.resolve({ id: documentId }) } as never,
|
||||
);
|
||||
return {
|
||||
status: res.status,
|
||||
headers: res.headers,
|
||||
body: (await res.json()) as Envelope,
|
||||
};
|
||||
}
|
||||
|
||||
async function postLabScanText(text: string) {
|
||||
const { POST } = await import("@/app/api/labs/ocr/extract/route");
|
||||
const res = await POST(
|
||||
new Request("http://localhost/api/labs/ocr/extract", {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({ mode: "text", text }),
|
||||
}) as never,
|
||||
);
|
||||
return { status: res.status, body: (await res.json()) as Envelope };
|
||||
}
|
||||
|
||||
async function poll(runId: string) {
|
||||
const { GET } = await import("@/app/api/ai-runs/[id]/route");
|
||||
const res = await GET(
|
||||
new Request(`http://localhost/api/ai-runs/${runId}`) as never,
|
||||
{ params: Promise.resolve({ id: runId }) } as never,
|
||||
);
|
||||
return {
|
||||
status: res.status,
|
||||
body: (await res.json()) as Envelope<Record<string, unknown>>,
|
||||
};
|
||||
}
|
||||
|
||||
async function work(runId: string) {
|
||||
const { runDocumentAiRun } = await import("@/lib/jobs/document-ai-run");
|
||||
return runDocumentAiRun(runId, Date.now() + 45 * 60_000);
|
||||
}
|
||||
|
||||
beforeEach(async () => {
|
||||
await truncateAllTables(getPrismaClient());
|
||||
cookieJar.clear();
|
||||
headerJar.clear();
|
||||
queue.up = true;
|
||||
send.mockClear();
|
||||
vi.mocked(transcribeDocument).mockClear();
|
||||
vi.mocked(runOcrExtraction).mockClear();
|
||||
});
|
||||
|
||||
describe("Read with AI in the background", () => {
|
||||
it("queues, runs in the worker, and serves the route's own body", async () => {
|
||||
const user = await makeUser();
|
||||
await grantConsent(user.id);
|
||||
const document = await storeDocument(user.id);
|
||||
await signIn(user.id);
|
||||
|
||||
const queued = await postIndex(document.id, true);
|
||||
expect(queued.status).toBe(202);
|
||||
expect(queued.headers.get("preference-applied")).toBe("respond-async");
|
||||
const { runId } = queued.body.data as { runId: string };
|
||||
expect(queued.headers.get("location")).toBe(`/api/ai-runs/${runId}`);
|
||||
expect(transcribeDocument).not.toHaveBeenCalled();
|
||||
expect(send).toHaveBeenCalledOnce();
|
||||
|
||||
const waiting = await poll(runId);
|
||||
expect(waiting.status).toBe(200);
|
||||
expect(waiting.body.data).toMatchObject({
|
||||
status: "QUEUED",
|
||||
kind: "DOCUMENT_INDEX",
|
||||
documentId: document.id,
|
||||
result: null,
|
||||
});
|
||||
|
||||
// A second press while the read is queued attaches to the same run.
|
||||
const again = await postIndex(document.id, true);
|
||||
expect(again.status).toBe(202);
|
||||
expect((again.body.data as { runId: string }).runId).toBe(runId);
|
||||
expect(send).toHaveBeenCalledOnce();
|
||||
|
||||
expect(await work(runId)).toBe("succeeded");
|
||||
expect(transcribeDocument).toHaveBeenCalledOnce();
|
||||
|
||||
const done = await poll(runId);
|
||||
expect(done.body.data).toMatchObject({
|
||||
status: "SUCCEEDED",
|
||||
error: null,
|
||||
retryAfterMs: null,
|
||||
result: {
|
||||
documentId: document.id,
|
||||
indexed: true,
|
||||
labFactsStaged: 0,
|
||||
},
|
||||
});
|
||||
const index = await getPrismaClient().documentContentIndex.findUnique({
|
||||
where: { documentId: document.id },
|
||||
});
|
||||
expect(index).not.toBeNull();
|
||||
});
|
||||
|
||||
it("answers synchronously, exactly as before, without the header", async () => {
|
||||
const user = await makeUser();
|
||||
await grantConsent(user.id);
|
||||
const document = await storeDocument(user.id);
|
||||
await signIn(user.id);
|
||||
|
||||
const res = await postIndex(document.id, false);
|
||||
expect(res.status).toBe(200);
|
||||
expect(res.headers.get("preference-applied")).toBeNull();
|
||||
expect(res.body.data).toMatchObject({
|
||||
documentId: document.id,
|
||||
indexed: true,
|
||||
labFactsStaged: 0,
|
||||
});
|
||||
expect(send).not.toHaveBeenCalled();
|
||||
expect(await getPrismaClient().documentAiRun.count()).toBe(0);
|
||||
});
|
||||
|
||||
it("stops in the worker when consent was withdrawn while the run waited", async () => {
|
||||
const user = await makeUser();
|
||||
const receipt = await grantConsent(user.id);
|
||||
const document = await storeDocument(user.id);
|
||||
await signIn(user.id);
|
||||
|
||||
const queued = await postIndex(document.id, true);
|
||||
const { runId } = queued.body.data as { runId: string };
|
||||
await getPrismaClient().consentReceipt.update({
|
||||
where: { id: receipt.id },
|
||||
data: { revokedAt: new Date() },
|
||||
});
|
||||
|
||||
expect(await work(runId)).toBe("failed");
|
||||
expect(transcribeDocument).not.toHaveBeenCalled();
|
||||
const done = await poll(runId);
|
||||
expect(done.body.data).toMatchObject({
|
||||
status: "FAILED",
|
||||
result: null,
|
||||
error: { code: "consent.ai.required", status: 403 },
|
||||
});
|
||||
});
|
||||
|
||||
it("answers 503 and leaves no live run when no worker queue is reachable", async () => {
|
||||
const user = await makeUser();
|
||||
await grantConsent(user.id);
|
||||
const document = await storeDocument(user.id);
|
||||
await signIn(user.id);
|
||||
queue.up = false;
|
||||
|
||||
const res = await postIndex(document.id, true);
|
||||
expect(res.status).toBe(503);
|
||||
expect(res.body.meta?.errorCode).toBe("aiRuns.workerUnavailable");
|
||||
const live = await getPrismaClient().documentAiRun.count({
|
||||
where: { status: { in: ["QUEUED", "RUNNING"] } },
|
||||
});
|
||||
expect(live).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("a lab scan in the background", () => {
|
||||
it("seals the text, reads it in the worker, drops the input, serves the rows", async () => {
|
||||
const user = await makeUser();
|
||||
await grantConsent(user.id);
|
||||
await signIn(user.id);
|
||||
|
||||
const queued = await postLabScanText("Hämoglobin 14,2 g/dL");
|
||||
expect(queued.status).toBe(202);
|
||||
const { runId } = queued.body.data as { runId: string };
|
||||
|
||||
const stored = await getPrismaClient().documentAiRun.findUniqueOrThrow({
|
||||
where: { id: runId },
|
||||
});
|
||||
expect(stored.kind).toBe("LABS_OCR_EXTRACT");
|
||||
expect(stored.inputEncrypted).not.toBeNull();
|
||||
expect(Buffer.from(stored.inputEncrypted!).toString("utf8")).not.toContain(
|
||||
"Hämoglobin",
|
||||
);
|
||||
expect(JSON.stringify(stored.paramsJson)).not.toContain("Hämoglobin");
|
||||
|
||||
expect(await work(runId)).toBe("succeeded");
|
||||
expect(runOcrExtraction).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ ocrText: "Hämoglobin 14,2 g/dL" }),
|
||||
);
|
||||
|
||||
const after = await getPrismaClient().documentAiRun.findUniqueOrThrow({
|
||||
where: { id: runId },
|
||||
});
|
||||
expect(after.inputEncrypted).toBeNull();
|
||||
expect(after.resultEncrypted).not.toBeNull();
|
||||
expect(Buffer.from(after.resultEncrypted!).toString("utf8")).not.toContain(
|
||||
"2027-01-10",
|
||||
);
|
||||
|
||||
const done = await poll(runId);
|
||||
expect(done.body.data).toMatchObject({
|
||||
status: "SUCCEEDED",
|
||||
kind: "LABS_OCR_EXTRACT",
|
||||
documentId: null,
|
||||
result: { reportDate: "2027-01-10", providerType: "anthropic", rows: [] },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("whose run it is", () => {
|
||||
it("is the same 404 for another account's run and for no run", async () => {
|
||||
const owner = await makeUser();
|
||||
await grantConsent(owner.id);
|
||||
await signIn(owner.id);
|
||||
const queued = await postLabScanText("Glucose 95 mg/dL");
|
||||
const { runId } = queued.body.data as { runId: string };
|
||||
|
||||
const other = await makeUser();
|
||||
await signIn(other.id);
|
||||
const foreign = await poll(runId);
|
||||
expect(foreign.status).toBe(404);
|
||||
expect(foreign.body.meta?.errorCode).toBe("aiRuns.notFound");
|
||||
|
||||
const missing = await poll("does-not-exist");
|
||||
expect(missing.status).toBe(404);
|
||||
expect(missing.body.meta?.errorCode).toBe("aiRuns.notFound");
|
||||
|
||||
await signIn(owner.id);
|
||||
expect((await poll(runId)).status).toBe(200);
|
||||
});
|
||||
});
|
||||
|
||||
describe("the reaper", () => {
|
||||
it("fails a run no worker took and hands its reservation back", async () => {
|
||||
const user = await makeUser();
|
||||
await grantConsent(user.id);
|
||||
await signIn(user.id);
|
||||
const queued = await postLabScanText("Glucose 95 mg/dL");
|
||||
const { runId } = queued.body.data as { runId: string };
|
||||
|
||||
const { readDailySpend, buildDateKey } =
|
||||
await import("@/lib/ai/coach/budget");
|
||||
expect(
|
||||
(await readDailySpend(user.id, buildDateKey())).total,
|
||||
).toBeGreaterThan(0);
|
||||
|
||||
const { reapAiRuns } = await import("@/lib/documents/ai-runs/store");
|
||||
const later = new Date(Date.now() + 16 * 60_000);
|
||||
expect(await reapAiRuns(later)).toEqual({
|
||||
workerUnavailable: 1,
|
||||
timedOut: 0,
|
||||
deleted: 0,
|
||||
});
|
||||
expect((await readDailySpend(user.id, buildDateKey())).total).toBe(0);
|
||||
|
||||
const run = await getPrismaClient().documentAiRun.findUniqueOrThrow({
|
||||
where: { id: runId },
|
||||
});
|
||||
expect(run.status).toBe("FAILED");
|
||||
expect(run.errorCode).toBe("aiRuns.workerUnavailable");
|
||||
expect(run.inputEncrypted).toBeNull();
|
||||
|
||||
// A worker that turns up late finds nothing to claim.
|
||||
expect(await work(runId)).toBe("skipped");
|
||||
expect(runOcrExtraction).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("deletes a finished run an hour after it ended, and the poll stops serving it first", async () => {
|
||||
const user = await makeUser();
|
||||
await grantConsent(user.id);
|
||||
await signIn(user.id);
|
||||
const queued = await postLabScanText("Glucose 95 mg/dL");
|
||||
const { runId } = queued.body.data as { runId: string };
|
||||
expect(await work(runId)).toBe("succeeded");
|
||||
|
||||
await getPrismaClient().documentAiRun.update({
|
||||
where: { id: runId },
|
||||
data: { expiresAt: new Date(Date.now() - 1000) },
|
||||
});
|
||||
// Past its retention the run is gone for the client even before the
|
||||
// reaper's tick.
|
||||
expect((await poll(runId)).status).toBe(404);
|
||||
|
||||
const { reapAiRuns } = await import("@/lib/documents/ai-runs/store");
|
||||
const summary = await reapAiRuns();
|
||||
expect(summary.deleted).toBe(1);
|
||||
expect(
|
||||
await getPrismaClient().documentAiRun.findUnique({
|
||||
where: { id: runId },
|
||||
}),
|
||||
).toBeNull();
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user