From a1ee2ac003be911441bcbfddff1162ac680dc655 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Sat, 24 Oct 2026 11:28:57 +0200 Subject: [PATCH] fix(documents): keep document reads inside their job's expiry Since the lab staging reads a long report in parts, one auto-index could make the transcription plus up to twenty reads of two calls each at the person's response timeout, hours past pg-boss's fifteen-minute default, which failed the job and retried it beside the reads still running. The staging now takes the run's deadline as a schedule and plans only the reads that fit before it. The index job is long, with an expiry above the calls that run before the plan can say no, a stop on jobDeadline and a lock per document; the document run hands the same staging its own deadline, which moves the run's expiry as the reads are planned. The document summary, one call and its reasoning repeat, gets an expiry above that worst case. ocrReadCallMs in the read core is the one per-call timeout both deadlines divide by. The expiry guard now names every importer of the document and lab read entry points with the queues that reach it. The runtime guard's planted cases stripped every file's comments again on each rescan, which took over a second alone and timed out under a full parallel run; the stripped text is now memoised. Refs #1274 --- src/__tests__/ai-job-expiry-guard.test.ts | 96 ++++++++++++++--- src/__tests__/baselines/test-as-never.json | 4 +- src/__tests__/job-queue-runtime-guard.test.ts | 19 +++- .../transcription-read-core-guard.test.ts | 14 ++- .../[id]/index/__tests__/route.test.ts | 5 +- .../__tests__/auto-stage-labs.test.ts | 21 ++++ src/lib/documents/ai-runs/index-run.ts | 12 ++- src/lib/documents/auto-stage-labs.ts | 17 ++- src/lib/documents/ocr-read.ts | 15 +++ .../__tests__/async-state-contract.test.ts | 3 +- src/lib/jobs/__tests__/document-index.test.ts | 102 +++++++++++++++--- .../jobs/__tests__/document-summary.test.ts | 3 +- src/lib/jobs/document-ai-run.ts | 43 ++++++-- src/lib/jobs/document-index-queue.ts | 26 +++++ src/lib/jobs/document-index.ts | 40 ++++++- src/lib/jobs/document-summary-queue.ts | 14 +++ src/lib/jobs/document-summary.ts | 6 +- src/lib/jobs/model-time.ts | 12 ++- src/lib/jobs/queue-runtime.ts | 32 ++++-- src/lib/jobs/registry/maintenance.ts | 4 + src/lib/jobs/reminder/document-handlers.ts | 22 ++-- 21 files changed, 434 insertions(+), 76 deletions(-) diff --git a/src/__tests__/ai-job-expiry-guard.test.ts b/src/__tests__/ai-job-expiry-guard.test.ts index f48d4739a..5e17d801c 100644 --- a/src/__tests__/ai-job-expiry-guard.test.ts +++ b/src/__tests__/ai-job-expiry-guard.test.ts @@ -11,13 +11,17 @@ * workout line and the Coach memory pass could each do the same once a status * run waited for its reasoning repeat. * - * Two entry points carry every background model run at the person's - * timeout: `runStatusCompletion` (a status run, `STATUS_MODEL_CALLS` calls) - * and `generateComprehensiveInsight` (`COMPREHENSIVE_MODEL_CALLS` calls). The - * guard holds: + * The entry points that carry a background model run at the person's + * timeout: `runStatusCompletion` (a status run, `STATUS_MODEL_CALLS` calls), + * `generateComprehensiveInsight` (`COMPREHENSIVE_MODEL_CALLS` calls), and the + * document and lab reads: `runOcrReads` (a report's reads, as many as the + * run's deadline allows, or up to `OCR_MAX_PAGES` without one), + * `transcribeDocument`, `runDocumentSummary`, `runDocumentAssist` and + * `executeSupplementLabelExtraction`. The guard holds: * * 1. every file that imports one of them is named below with the queues - * whose jobs reach it, so a new caller has to be attributed; + * whose jobs reach it, or as a request surface with its reason, so a new + * caller has to be attributed; * 2. each of those queues is `bounded` in `queue-runtime.ts`, with model * time at least the runs it makes at the largest setting and an expiry * above that, or `long`, where the run's lock makes the retry after an @@ -26,13 +30,16 @@ * * Its limit: the run counts per queue are declared here, not derived; the * Coach memory refresh's three runs, for instance, are three calls in - * `runCoachMemoryRefresh` that a reader has to count. Model calls outside the - * two entry points (the document and lab runs, the reaction line, the nudge) - * keep their own deadlines and are not covered here. + * `runCoachMemoryRefresh` that a reader has to count. A `long` queue is held + * to its lock, not to arithmetic: its reads are planned against the job's + * deadline (`jobDeadline`), which `document-index.test.ts` and the document + * run's own tests prove. Model calls outside these entry points (the + * reaction line and the nudge, both under a fixed surface ceiling) are not + * covered here. * - * Mutation check: declare `morning-digest-refresh` short again, drop a run - * from a queue's model time, or import `runStatusCompletion` into a new file, - * and this goes red. The "plants" case runs the check over the table as it + * Mutation check: declare `morning-digest-refresh` or `document-index` short + * again, drop a run from a queue's model time, or import `runStatusCompletion` + * or `runOcrReads` into a new file, and this goes red. The "plants" case runs the check over the table as it * stood before #1274. */ import { join } from "node:path"; @@ -46,6 +53,7 @@ import { } from "./helpers/source-files"; import { QUEUE_RUNTIME, type QueueRuntime } from "@/lib/jobs/queue-runtime"; import { + COMPLETION_CEILING_MS, COMPREHENSIVE_WARM_CEILING_MS, STATUS_RUN_CEILING_MS, } from "@/lib/jobs/model-time"; @@ -56,7 +64,15 @@ import { import { AI_RESPONSE_TIMEOUT_MAX_SECONDS } from "@/lib/ai/effective-timeout"; const SRC = join(process.cwd(), "src"); -const ENTRY_POINTS = ["runStatusCompletion", "generateComprehensiveInsight"]; +const ENTRY_POINTS = [ + "runStatusCompletion", + "generateComprehensiveInsight", + "runOcrReads", + "transcribeDocument", + "runDocumentSummary", + "runDocumentAssist", + "executeSupplementLabelExtraction", +]; /** Each importer of an entry point, and the queues whose jobs reach it. */ const MODEL_RUN_SITES: Readonly> = { @@ -79,11 +95,39 @@ const MODEL_RUN_SITES: Readonly> = { "lib/jobs/workout-insight-generate.ts": ["workout-insight-generate"], "lib/jobs/insight-pregenerate.ts": ["insight-pregenerate"], "lib/jobs/morning-digest-refresh.ts": ["morning-digest-refresh"], + "lib/documents/extract.ts": ["document-ai-run", "document-index"], + "lib/documents/index-document.ts": ["document-index", "document-ai-run"], + "lib/documents/ai-runs/index-run.ts": ["document-ai-run"], + "lib/documents/ai-runs/summary-run.ts": ["document-ai-run"], + "lib/documents/ai-runs/suggest-run.ts": ["document-ai-run"], + "lib/labs/ocr-extract.ts": ["document-ai-run"], + "lib/jobs/document-ai-run.ts": ["document-ai-run"], + "lib/jobs/document-content-index-backfill.ts": [ + "document-content-index-backfill", + ], + "lib/jobs/document-summary.ts": ["document-summary"], +}; + +/** + * Importers no queue reaches: a request waits on them, and the browser's + * abort follows the published setting (`clientAbortMs`). + */ +const REQUEST_SITES: Readonly> = { + "app/api/documents/inbound/[id]/summary/route.ts": + "The on-demand transcription a person asked for; the background path runs through document-ai-run.", }; /** What one job of a bounded queue runs, at most, one after another. */ const RUNS: Readonly< - Record + Record< + string, + { + status?: number; + comprehensive?: number; + completions?: number; + why: string; + } + > > = { "insight-status-generate": { status: 1, @@ -101,6 +145,10 @@ const RUNS: Readonly< comprehensive: 1, why: "One forced comprehensive generation.", }, + "document-summary": { + completions: 1, + why: "One summary call (DOCUMENT_SUMMARY_MODEL_CALLS) and its reasoning repeat.", + }, }; /** Value specifiers of every `import { … } from` in `text`. */ @@ -126,10 +174,15 @@ function importers(): string[] { } /** Seconds of model work `runs` makes at the largest setting. */ -function modelSecondsOf(runs: { status?: number; comprehensive?: number }) { +function modelSecondsOf(runs: { + status?: number; + comprehensive?: number; + completions?: number; +}) { return ( ((runs.status ?? 0) * STATUS_RUN_CEILING_MS + - (runs.comprehensive ?? 0) * COMPREHENSIVE_WARM_CEILING_MS) / + (runs.comprehensive ?? 0) * COMPREHENSIVE_WARM_CEILING_MS + + (runs.completions ?? 0) * COMPLETION_CEILING_MS) / 1000 ); } @@ -177,8 +230,10 @@ function findings(table: Readonly>): string[] { describe("jobs that wait on the response timeout outlast it", () => { it("names every importer of a model run (an empty set is a failure)", () => { const found = importers(); - expect(found.length).toBeGreaterThanOrEqual(10); - expect(found).toEqual(Object.keys(MODEL_RUN_SITES).sort()); + expect(found.length).toBeGreaterThanOrEqual(20); + expect(found).toEqual( + [...Object.keys(MODEL_RUN_SITES), ...Object.keys(REQUEST_SITES)].sort(), + ); }); it("sizes the run ceilings from the largest setting and the call counts", () => { @@ -200,6 +255,13 @@ describe("jobs that wait on the response timeout outlast it", () => { ).toEqual([]); }); + it("names every request surface with a reason, and none twice", () => { + for (const [rel, why] of Object.entries(REQUEST_SITES)) { + expect(why.length, rel).toBeGreaterThan(30); + expect(rel in MODEL_RUN_SITES, rel).toBe(false); + } + }); + it("plants: the table before #1274 is caught", () => { const before: Record = { ...QUEUE_RUNTIME, diff --git a/src/__tests__/baselines/test-as-never.json b/src/__tests__/baselines/test-as-never.json index a13953429..64018c1a6 100644 --- a/src/__tests__/baselines/test-as-never.json +++ b/src/__tests__/baselines/test-as-never.json @@ -1,7 +1,7 @@ { "$comment": "as never casts in tests, per test file under src/. Maintained by src/__tests__/test-realism-ratchet.test.ts.", "count": 524, - "total": 5490, + "total": 5489, "entries": { "__tests__/api/auth/login/mfa-branch.test.ts": 10, "__tests__/api/auth/mfa/verify.test.ts": 9, @@ -433,7 +433,7 @@ "lib/jobs/__tests__/data-arrival-worker.test.ts": 18, "lib/jobs/__tests__/document-ai-run.test.ts": 1, "lib/jobs/__tests__/document-content-index-backfill.test.ts": 29, - "lib/jobs/__tests__/document-index.test.ts": 6, + "lib/jobs/__tests__/document-index.test.ts": 5, "lib/jobs/__tests__/document-summary-reaper.test.ts": 1, "lib/jobs/__tests__/document-summary.test.ts": 19, "lib/jobs/__tests__/geo-backfill.test.ts": 1, diff --git a/src/__tests__/job-queue-runtime-guard.test.ts b/src/__tests__/job-queue-runtime-guard.test.ts index 27a15294c..cfc4de732 100644 --- a/src/__tests__/job-queue-runtime-guard.test.ts +++ b/src/__tests__/job-queue-runtime-guard.test.ts @@ -60,9 +60,20 @@ const source = new Map( files.map((f) => [f, readFileSync(join(SRC, f), "utf8")]), ); -/** Strip comments so documentation that quotes code never matches. */ +/** + * Strip comments so documentation that quotes code never matches. Memoised + * per text: the planted cases rescan the whole tree once per finding, and + * stripping every file again each time made that one case take seconds, + * past the default timeout under a full parallel run. + */ +const stripped = new Map(); function code(text: string): string { - return stripComments(text); + let out = stripped.get(text); + if (out === undefined) { + out = stripComments(text); + stripped.set(text, out); + } + return out; } /** Index just past the bracket that closes the one at `open`. */ @@ -326,9 +337,7 @@ describe("every queue declares its runtime", () => { { definitions: [planted({})], ...scan(texts(budgeted)) }, ), ).toEqual(["planted-long: expiry 900s out of range"]); - // Each finding rescans every send in the tree; under a full parallel run - // that took six to nine seconds against the five-second default. - }, 30_000); + }); }); /** diff --git a/src/__tests__/transcription-read-core-guard.test.ts b/src/__tests__/transcription-read-core-guard.test.ts index 430f00cb0..9c6a64327 100644 --- a/src/__tests__/transcription-read-core-guard.test.ts +++ b/src/__tests__/transcription-read-core-guard.test.ts @@ -70,12 +70,18 @@ describe("transcription read core", () => { expect(setters).toEqual(["documents/ocr-read.ts"]); }); - it("reads the per-read timeout in the core and in the deadline that plans the reads, nowhere else", () => { + it("reads the per-read timeout in the core only, and every deadline that plans reads asks it", () => { const readers = filesMatching(/\bAI_BUDGETS\.ocrExtract\.timeoutMs\b/); - requireFloor("files reading the transcription timeout", readers, 2); - expect(readers).toEqual([ - "documents/ocr-read.ts", + requireFloor("files reading the transcription timeout", readers, 1); + expect(readers).toEqual(["documents/ocr-read.ts"]); + // The document run and the auto-index job divide their deadlines by the + // same per-call timeout the reads send (`ocrReadCallMs`). + const planners = filesMatching(/\bocrReadCallMs\s*\(/).filter( + (rel) => rel !== "documents/ocr-read.ts", + ); + expect(planners).toEqual([ "jobs/document-ai-run.ts", + "jobs/document-index.ts", ]); }); diff --git a/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts b/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts index 74097db12..8307653b2 100644 --- a/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts +++ b/src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts @@ -291,7 +291,10 @@ describe("POST /api/documents/inbound/[id]/index — staging continuation", () = } as never); const res = await POST(visionReq("doc-1") as never, ctx("doc-1") as never); expect(res.status).toBe(200); - expect(maybeAutoStageLabFacts).toHaveBeenCalledWith("user-1", "doc-1"); + // A request has no deadline: the staging reads as it always did. + expect(maybeAutoStageLabFacts).toHaveBeenCalledWith("user-1", "doc-1", { + schedule: undefined, + }); const body = await res.json(); expect(body.data.labFactsStaged).toBe(4); // Read with AI ends an import's hold on automatic AI reading (#1038), diff --git a/src/lib/documents/__tests__/auto-stage-labs.test.ts b/src/lib/documents/__tests__/auto-stage-labs.test.ts index 85c903be5..09525392b 100644 --- a/src/lib/documents/__tests__/auto-stage-labs.test.ts +++ b/src/lib/documents/__tests__/auto-stage-labs.test.ts @@ -173,6 +173,27 @@ describe("maybeAutoStageLabFacts", () => { expect(created.data.every((f) => f.status === "PENDING")).toBe(true); }); + it("plans a long report's reads against the run's deadline (#1274)", async () => { + // Four hundred lines of a panel, many reads of text. + const longReport = Array.from( + { length: 400 }, + (_, i) => `Hämoglobin ${i} g/dl Referenzbereich 13-17 LDL 120 mg/dl`, + ).join("\n"); + mockText.mockResolvedValue({ text: longReport, source: "verbatim" }); + + await maybeAutoStageLabFacts("u1", "d1"); + const unbounded = mockExtract.mock.calls[0]![0].chunks.length; + expect(unbounded).toBeGreaterThan(1); + + const arm = vi.fn(async () => {}); + await maybeAutoStageLabFacts("u1", "d1", { + schedule: { maxModelCalls: () => 2, arm }, + }); + // Two calls left before the deadline: one read (a call and its retry). + expect(mockExtract.mock.calls[1]![0].chunks).toHaveLength(1); + expect(arm).toHaveBeenCalledWith(2); + }); + it("skips (manual path only) when a module is off", async () => { mockModule.mockImplementation(async (_u, key) => key !== "labs"); const res = await maybeAutoStageLabFacts("u1", "d1"); diff --git a/src/lib/documents/ai-runs/index-run.ts b/src/lib/documents/ai-runs/index-run.ts index e8e955683..99f67b1fa 100644 --- a/src/lib/documents/ai-runs/index-run.ts +++ b/src/lib/documents/ai-runs/index-run.ts @@ -12,6 +12,7 @@ * (status, message, `documents.inbound.*` code); the caller renders it as a * response or stores it on the run. */ +import type { OcrSchedule } from "@/lib/documents/ocr-plan"; import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; import type { ProviderChainType } from "@/lib/ai/provider-chain"; import { auditLog } from "@/lib/auth/audit"; @@ -56,6 +57,11 @@ export const DOCUMENT_INDEX_RESERVE_TOKENS = export interface IndexOrigin { ipAddress: string | null; worker: boolean; + /** + * A background run's deadline, which the lab staging after the index plans + * its reads against. A request has none: it reads as before. + */ + schedule?: OcrSchedule; } export type DocumentIndexInput = @@ -109,9 +115,9 @@ async function finishIndex( // 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 staging = await maybeAutoStageLabFacts(userId, documentId, { + schedule: origin.schedule, + }).catch(() => null); const labFactsStaged = staging?.staged === true ? staging.facts : 0; return { ok: true, diff --git a/src/lib/documents/auto-stage-labs.ts b/src/lib/documents/auto-stage-labs.ts index c9a1d25ef..07efe063b 100644 --- a/src/lib/documents/auto-stage-labs.ts +++ b/src/lib/documents/auto-stage-labs.ts @@ -41,6 +41,10 @@ import { runInboundExtraction, type InboundExtractionResult, } from "@/lib/documents/extract"; +import { + OCR_EXTRACT_MODEL_CALLS, + type OcrSchedule, +} from "@/lib/documents/ocr-plan"; import { OcrExtractError } from "@/lib/documents/ocr-read"; import { encryptFactData, encryptFactProvenance } from "@/lib/documents/store"; import { prisma } from "@/lib/db"; @@ -251,10 +255,16 @@ async function stageFactsIfStored( * Auto-stage lab facts for one freshly indexed document (owner-scoped). Called * fire-and-forget by the auto-index job after a successful index; every guard * miss returns a tagged no-op so the manual extract button stays the fallback. + * + * `schedule` is the run's deadline when a job reads (the auto-index job, the + * document run): the report's text is then split into as many reads as fit + * before it, instead of up to `OCR_MAX_PAGES` reads of two calls each that a + * slow model could stretch past the job's expiry (#1274). */ export async function maybeAutoStageLabFacts( userId: string, documentId: string, + options: { schedule?: OcrSchedule } = {}, ): Promise { const [inboundOn, labsOn] = await Promise.all([ isModuleEnabled(userId, "inboundDocuments"), @@ -321,9 +331,14 @@ export async function maybeAutoStageLabFacts( { cap, surface: "job" }, ); - const plan = planInboundTextReads(chat.text, undefined); + const plan = planInboundTextReads(chat.text, options.schedule); let result: InboundExtractionResult; try { + // A run that cannot fit even one read before its deadline refuses here; + // the reservation is settled below like any other failed read. + await options.schedule?.arm( + plan.args.chunks.length * OCR_EXTRACT_MODEL_CALLS, + ); result = await runInboundExtraction({ userId, provider: provider.pick.entry.instance, diff --git a/src/lib/documents/ocr-read.ts b/src/lib/documents/ocr-read.ts index 15445dcb4..9e647710a 100644 --- a/src/lib/documents/ocr-read.ts +++ b/src/lib/documents/ocr-read.ts @@ -25,6 +25,7 @@ * The raw image / PDF bytes live in memory only; nothing here logs them. */ import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; +import { callTimeoutMs } from "@/lib/ai/effective-timeout"; import { aiEgressRefusal } from "@/lib/ai/capabilities/egress"; import type { AiUnavailableError } from "@/lib/ai/capabilities/refusal"; import { @@ -69,6 +70,20 @@ export function ocrReadParams(prompt: { }; } +/** + * The timeout one call of a read runs under for a record whose setting is + * `responseTimeoutSeconds`: the read's own default, lifted by the setting. + * What a deadline that plans reads divides its time by. + */ +export function ocrReadCallMs( + responseTimeoutSeconds: number | null | undefined, +): number { + return callTimeoutMs( + { timeoutMs: AI_BUDGETS.ocrExtract.timeoutMs }, + responseTimeoutSeconds, + ); +} + /** Appended to the user turn when the first answer was not the schema. */ const RETRY_SUFFIX = "\n\nYour previous response was not valid JSON matching the schema. Return ONLY the JSON object described in the system prompt, with no prose or markdown."; diff --git a/src/lib/jobs/__tests__/async-state-contract.test.ts b/src/lib/jobs/__tests__/async-state-contract.test.ts index 76f7a55f0..2c0f30af8 100644 --- a/src/lib/jobs/__tests__/async-state-contract.test.ts +++ b/src/lib/jobs/__tests__/async-state-contract.test.ts @@ -77,7 +77,8 @@ vi.mock("@/lib/ai/capabilities/gate", () => ({ aiCapabilityForJob: vi.fn(), aiCapabilityForRecord: vi.fn(), })); -vi.mock("@/lib/ai/ai-budgets", () => ({ +vi.mock("@/lib/ai/ai-budgets", async (importOriginal) => ({ + ...(await importOriginal()), AI_BUDGETS: { documentSummary: { temperature: 0.3, maxTokens: 600 }, documentTranscribe: { temperature: 0, maxTokens: 4000 }, diff --git a/src/lib/jobs/__tests__/document-index.test.ts b/src/lib/jobs/__tests__/document-index.test.ts index 456005ac2..80cf09f67 100644 --- a/src/lib/jobs/__tests__/document-index.test.ts +++ b/src/lib/jobs/__tests__/document-index.test.ts @@ -14,13 +14,38 @@ vi.mock("@/lib/logging/context", () => ({ annotate: vi.fn() })); vi.mock("@/lib/documents/auto-stage-labs", () => ({ maybeAutoStageLabFacts: vi.fn().mockResolvedValue(undefined), })); +const { findUnique } = vi.hoisted(() => ({ findUnique: vi.fn() })); +vi.mock("@/lib/db", () => ({ prisma: { user: { findUnique } } })); -import { enqueueDocumentIndex, runDocumentIndex } from "../document-index"; -import { DOCUMENT_INDEX_QUEUE } from "../document-index-queue"; +import { + enqueueDocumentIndex, + runDocumentIndex, + stagingSchedule, +} from "../document-index"; +import { + DOCUMENT_INDEX_EXPIRE_SECONDS, + DOCUMENT_INDEX_QUEUE, + DOCUMENT_INDEX_UNPLANNED_COMPLETIONS, +} from "../document-index-queue"; +import { JOB_BUDGET_SHARE } from "../job-budget"; +import { COMPLETION_CEILING_MS } from "../model-time"; +import { AI_RESPONSE_TIMEOUT_MAX_SECONDS } from "@/lib/ai/effective-timeout"; +import { + OCR_EXTRACT_MODEL_CALLS, + OCR_MAX_PAGES, + maxReadsFor, +} from "@/lib/documents/ocr-plan"; import { indexDocumentContent } from "@/lib/documents/index-document"; import { getGlobalBoss } from "@/lib/jobs/boss-instance"; import { maybeAutoStageLabFacts } from "@/lib/documents/auto-stage-labs"; +/** A document the local text layer indexed. */ +const INDEXED_LOCALLY = { + indexed: true, + source: "local-pdf", + tokenCount: 3, +} as never; + beforeEach(() => { vi.clearAllMocks(); }); @@ -35,7 +60,10 @@ describe("enqueueDocumentIndex", () => { expect(send).toHaveBeenCalledWith( DOCUMENT_INDEX_QUEUE, expect.objectContaining({ userId: "user-1", documentId: "doc-1" }), - expect.objectContaining({ singletonKey: "document-index|doc-1" }), + expect.objectContaining({ + singletonKey: "document-index|doc-1", + expireInSeconds: DOCUMENT_INDEX_EXPIRE_SECONDS, + }), ); }); @@ -66,24 +94,34 @@ describe("enqueueDocumentIndex", () => { describe("runDocumentIndex", () => { it("delegates to the shared decision tree", async () => { - vi.mocked(indexDocumentContent).mockResolvedValue({ - indexed: true, - source: "local-pdf", - tokenCount: 3, - } as never); + vi.mocked(indexDocumentContent).mockResolvedValue(INDEXED_LOCALLY); await runDocumentIndex({ userId: "user-1", documentId: "doc-1" }); expect(indexDocumentContent).toHaveBeenCalledWith("user-1", "doc-1", { localOnly: false, }); - expect(maybeAutoStageLabFacts).toHaveBeenCalledWith("user-1", "doc-1"); + expect(maybeAutoStageLabFacts).toHaveBeenCalledWith("user-1", "doc-1", { + schedule: undefined, + }); + }); + + it("hands the lab staging the job's deadline", async () => { + vi.mocked(indexDocumentContent).mockResolvedValue(INDEXED_LOCALLY); + findUnique.mockResolvedValue({ aiResponseTimeoutSeconds: 300 }); + let clock = 0; + await runDocumentIndex( + { userId: "user-1", documentId: "doc-1" }, + { deadlineAt: 3_000_000, now: () => clock }, + ); + const schedule = vi.mocked(maybeAutoStageLabFacts).mock.calls[0][2] + ?.schedule; + // 300 s per call, twice for the reasoning repeat: five calls fit. + expect(schedule?.maxModelCalls()).toBe(5); + clock = 2_000_000; + expect(schedule?.maxModelCalls()).toBe(1); }); it("a deferred import indexes locally and stages no lab facts", async () => { - vi.mocked(indexDocumentContent).mockResolvedValue({ - indexed: true, - source: "local-pdf", - tokenCount: 3, - } as never); + vi.mocked(indexDocumentContent).mockResolvedValue(INDEXED_LOCALLY); await runDocumentIndex({ userId: "user-1", documentId: "doc-1", @@ -100,3 +138,39 @@ describe("runDocumentIndex", () => { expect(indexDocumentContent).not.toHaveBeenCalled(); }); }); + +describe("document-index — a slow model cannot outlive the job (#1274)", () => { + it("ends below the expiry at the largest setting, however long the report", async () => { + findUnique.mockResolvedValue({ + aiResponseTimeoutSeconds: AI_RESPONSE_TIMEOUT_MAX_SECONDS, + }); + const expiryMs = DOCUMENT_INDEX_EXPIRE_SECONDS * 1000; + const deadlineAt = expiryMs * JOB_BUDGET_SHARE; + // The transcription takes its whole worst case before the staging plans. + const stagingStarts = COMPLETION_CEILING_MS; + const schedule = await stagingSchedule( + "user-1", + deadlineAt, + () => stagingStarts, + ); + const reads = Math.min(maxReadsFor(schedule), OCR_MAX_PAGES); + const ends = + stagingStarts + reads * OCR_EXTRACT_MODEL_CALLS * COMPLETION_CEILING_MS; + expect(ends).toBeLessThan(expiryMs); + // Without the deadline the staging would plan every part of a long + // report: twenty reads of two calls, hours past the expiry. + expect( + stagingStarts + + OCR_MAX_PAGES * OCR_EXTRACT_MODEL_CALLS * COMPLETION_CEILING_MS, + ).toBeGreaterThan(expiryMs); + }); + + it("counts the transcription and the first read the deadline cannot refuse", () => { + expect(DOCUMENT_INDEX_UNPLANNED_COMPLETIONS).toBe( + 1 + OCR_EXTRACT_MODEL_CALLS, + ); + expect(DOCUMENT_INDEX_EXPIRE_SECONDS * 1000).toBeGreaterThan( + DOCUMENT_INDEX_UNPLANNED_COMPLETIONS * COMPLETION_CEILING_MS, + ); + }); +}); diff --git a/src/lib/jobs/__tests__/document-summary.test.ts b/src/lib/jobs/__tests__/document-summary.test.ts index edf3846fc..30680185c 100644 --- a/src/lib/jobs/__tests__/document-summary.test.ts +++ b/src/lib/jobs/__tests__/document-summary.test.ts @@ -50,7 +50,8 @@ vi.mock("@/lib/ai/capabilities/gate", () => ({ aiCapabilityForJob: vi.fn(), aiCapabilityForRecord: vi.fn(), })); -vi.mock("@/lib/ai/ai-budgets", () => ({ +vi.mock("@/lib/ai/ai-budgets", async (importOriginal) => ({ + ...(await importOriginal()), AI_BUDGETS: { documentSummary: { temperature: 0.3, maxTokens: 600 } }, })); diff --git a/src/lib/jobs/document-ai-run.ts b/src/lib/jobs/document-ai-run.ts index f24b70dc4..4df270a6a 100644 --- a/src/lib/jobs/document-ai-run.ts +++ b/src/lib/jobs/document-ai-run.ts @@ -36,7 +36,6 @@ import { PICK_DECIDED_REASONS, type AiCapabilityKey, } from "@/lib/ai/capabilities/types"; -import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; import { callTimeoutMs } from "@/lib/ai/effective-timeout"; import { prisma } from "@/lib/db"; import { loadOwnedDocument } from "@/lib/documents/ai-route-support"; @@ -77,6 +76,7 @@ import { OCR_EXTRACT_MODEL_CALLS, type OcrSchedule, } from "@/lib/documents/ocr-plan"; +import { ocrReadCallMs } from "@/lib/documents/ocr-read"; import { requireDocumentTextProvider, requireDocumentVisionProvider, @@ -220,7 +220,16 @@ async function runDocumentIndex( userId: run.userId, document, input: { mode: "text", text: run.input?.toString("utf8") ?? "" }, - origin: { ipAddress: null, worker: true }, + origin: { + ipAddress: null, + worker: true, + schedule: deadlineSchedule( + run.id, + ocrReadCallMs(owner?.aiResponseTimeoutSeconds ?? null), + deadlineAt, + now, + ), + }, }); } const pick = await requireDocumentVisionProvider(run.userId); @@ -237,7 +246,16 @@ async function runDocumentIndex( userId: run.userId, document, input: { mode: "vision", pick, budget: run.params.budget }, - origin: { ipAddress: null, worker: true }, + origin: { + ipAddress: null, + worker: true, + schedule: deadlineSchedule( + run.id, + ocrReadCallMs(pick.entry.instance.responseTimeoutSeconds), + deadlineAt, + now, + ), + }, }); } @@ -255,11 +273,22 @@ async function ocrReadSchedule( deadlineAt: number | undefined, now: () => number, ): Promise { - const perCallMs = callTimeoutMs( - { timeoutMs: AI_BUDGETS.ocrExtract.timeoutMs }, - responseTimeoutSeconds, - ); + const perCallMs = ocrReadCallMs(responseTimeoutSeconds); await armDeadline(runId, perCallMs, OCR_EXTRACT_MODEL_CALLS, deadlineAt, now); + return deadlineSchedule(runId, perCallMs, deadlineAt, now); +} + +/** + * The schedule of reads still to come before the job's deadline, armed only + * when the reads are planned: the lab staging after an index plans its reads + * once the transcription is done (#1274). + */ +function deadlineSchedule( + runId: string, + perCallMs: number, + deadlineAt: number | undefined, + now: () => number, +): OcrSchedule { return { maxModelCalls: () => deadlineAt === undefined diff --git a/src/lib/jobs/document-index-queue.ts b/src/lib/jobs/document-index-queue.ts index caf1d936d..6f311d6fc 100644 --- a/src/lib/jobs/document-index-queue.ts +++ b/src/lib/jobs/document-index-queue.ts @@ -3,7 +3,33 @@ * schedule and concurrency. They live apart from the handler so the worker * registry and every sender can read them without loading the job. */ +import { + COMPLETION_CEILING_MS, + expireAboveModelTime, +} from "@/lib/jobs/model-time"; + export const DOCUMENT_INDEX_QUEUE = "document-index"; +/** + * The completions that run before the job's deadline can turn a read away: + * the transcription, and the first lab-staging read with its corrective + * retry (`1 + OCR_EXTRACT_MODEL_CALLS`; the staging plans at least one read). + * Every further read is planned against the deadline. + */ +export const DOCUMENT_INDEX_UNPLANNED_COMPLETIONS = 3; + +/** + * The job's expiry. Since the lab staging reads a long report in parts, one + * index could make the transcription plus up to twenty reads of two calls + * each at the person's timeout, hours past pg-boss's fifteen-minute default, + * which then failed the job and retried it beside the reads still running + * (#1274). The staging now plans its reads to end before three quarters of + * this (`jobDeadline`); what runs before that plan, at the largest setting, + * still ends below it. + */ +export const DOCUMENT_INDEX_EXPIRE_SECONDS = expireAboveModelTime( + DOCUMENT_INDEX_UNPLANNED_COMPLETIONS * COMPLETION_CEILING_MS, +); + /** Serial concurrency — provider calls + PDF parse, kept off the request pool. */ export const DOCUMENT_INDEX_CONCURRENCY = 1; diff --git a/src/lib/jobs/document-index.ts b/src/lib/jobs/document-index.ts index 4cfacec37..710e3552c 100644 --- a/src/lib/jobs/document-index.ts +++ b/src/lib/jobs/document-index.ts @@ -12,12 +12,19 @@ * The queue MUST be registered in the maintenance registrar * (`src/lib/jobs/reminder/register-maintenance.ts`) so pg-boss provisions it. */ +import { STATUS_MODEL_CALLS } from "@/lib/ai/ai-budgets"; import { maybeAutoStageLabFacts } from "@/lib/documents/auto-stage-labs"; import { indexDocumentContent } from "@/lib/documents/index-document"; +import type { OcrSchedule } from "@/lib/documents/ocr-plan"; +import { ocrReadCallMs } from "@/lib/documents/ocr-read"; +import { prisma } from "@/lib/db"; import { getGlobalBoss } from "@/lib/jobs/boss-instance"; import { annotate } from "@/lib/logging/context"; import { caughtAs, logCaught } from "@/lib/logging/signal"; -import { DOCUMENT_INDEX_QUEUE } from "./document-index-queue"; +import { + DOCUMENT_INDEX_EXPIRE_SECONDS, + DOCUMENT_INDEX_QUEUE, +} from "./document-index-queue"; export interface DocumentIndexPayload { userId: string; @@ -38,6 +45,7 @@ export interface DocumentIndexPayload { */ export async function runDocumentIndex( payload: DocumentIndexPayload, + options: { deadlineAt?: number; now?: () => number } = {}, ): Promise { const { userId, documentId } = payload; if (!userId || !documentId) return; @@ -60,12 +68,39 @@ export async function runDocumentIndex( // A deferred import stages nothing: lab staging is an AI read, and the // person asked for those to wait. if (outcome.indexed && !localOnly) { - await maybeAutoStageLabFacts(userId, documentId).catch( + const schedule = + options.deadlineAt === undefined + ? undefined + : await stagingSchedule(userId, options.deadlineAt, options.now); + await maybeAutoStageLabFacts(userId, documentId, { schedule }).catch( caughtAs("documents.auto_index.lab_staging_failed"), ); } } +/** + * The lab staging's reads, planned against the job's deadline. One call of a + * read may take the person's timeout twice (the call and the one repeat a + * thinking model gets), so that is the unit the deadline is divided by. + * There is no run row to move: the deadline is the job's own, fixed. + */ +export async function stagingSchedule( + userId: string, + deadlineAt: number, + now: () => number = Date.now, +): Promise { + const owner = await prisma.user.findUnique({ + where: { id: userId }, + select: { aiResponseTimeoutSeconds: true }, + }); + const perCallMs = + ocrReadCallMs(owner?.aiResponseTimeoutSeconds ?? null) * STATUS_MODEL_CALLS; + return { + maxModelCalls: () => Math.floor((deadlineAt - now()) / perCallMs), + arm: async () => {}, + }; +} + /** * Enqueue a per-document index job. Coalesced by `singletonKey` per document so * a duplicate/idempotent re-upload cannot pile up parallel runs. Fire-and- @@ -93,6 +128,7 @@ export async function enqueueDocumentIndex( retryDelay: 30, retryBackoff: true, singletonKey: `document-index|${documentId}`, + expireInSeconds: DOCUMENT_INDEX_EXPIRE_SECONDS, }); return { enqueued: Boolean(jobId) }; } catch (err) { diff --git a/src/lib/jobs/document-summary-queue.ts b/src/lib/jobs/document-summary-queue.ts index c509924d7..f0cb6387e 100644 --- a/src/lib/jobs/document-summary-queue.ts +++ b/src/lib/jobs/document-summary-queue.ts @@ -3,7 +3,21 @@ * schedule and concurrency. They live apart from the handler so the worker * registry and every sender can read them without loading the job. */ +import { + COMPLETION_CEILING_MS, + expireAboveModelTime, +} from "@/lib/jobs/model-time"; + export const DOCUMENT_SUMMARY_QUEUE = "document-summary"; +/** + * The job's expiry: one summary call (`DOCUMENT_SUMMARY_MODEL_CALLS`) and its + * reasoning repeat at the largest response timeout, plus headroom. Two + * ten-minute attempts outlived pg-boss's fifteen-minute default (#1274). + */ +export const DOCUMENT_SUMMARY_EXPIRE_SECONDS = expireAboveModelTime( + COMPLETION_CEILING_MS, +); + /** Serial concurrency — provider calls, kept off the request pool. */ export const DOCUMENT_SUMMARY_CONCURRENCY = 1; diff --git a/src/lib/jobs/document-summary.ts b/src/lib/jobs/document-summary.ts index dfebd34e6..a1a9bbb84 100644 --- a/src/lib/jobs/document-summary.ts +++ b/src/lib/jobs/document-summary.ts @@ -49,7 +49,10 @@ import { getGlobalBoss } from "@/lib/jobs/boss-instance"; import { annotate } from "@/lib/logging/context"; import type { DocumentSummaryState } from "@/generated/prisma/client"; import { logCaught } from "@/lib/logging/signal"; -import { DOCUMENT_SUMMARY_QUEUE } from "./document-summary-queue"; +import { + DOCUMENT_SUMMARY_EXPIRE_SECONDS, + DOCUMENT_SUMMARY_QUEUE, +} from "./document-summary-queue"; /** * Record what became of this document's summary. Never touches a READY row — @@ -380,6 +383,7 @@ export async function enqueueDocumentSummary( retryDelay: 30, retryBackoff: true, singletonKey: `document-summary|${documentId}`, + expireInSeconds: DOCUMENT_SUMMARY_EXPIRE_SECONDS, }); // PENDING is claimed only once a job genuinely exists. That is the whole // point: the detail view may say "being generated" for this state, so diff --git a/src/lib/jobs/model-time.ts b/src/lib/jobs/model-time.ts index b71a8fa43..e6df15444 100644 --- a/src/lib/jobs/model-time.ts +++ b/src/lib/jobs/model-time.ts @@ -59,12 +59,18 @@ export function comprehensiveWarmBudgetMs( } /** - * The longest one status run (`runStatusCompletion`) waits at the largest - * setting: each of its `STATUS_MODEL_CALLS` calls at the full timeout. + * The longest one completion takes at the largest setting, as the clients run + * it: the call and the one repeat the Local and gateway clients send when a + * thinking model answered with reasoning alone (`STATUS_MODEL_CALLS`), each + * at the full timeout. A document read's call and its corrective retry are + * two completions. */ -export const STATUS_RUN_CEILING_MS = +export const COMPLETION_CEILING_MS = AI_RESPONSE_TIMEOUT_MAX_MS * STATUS_MODEL_CALLS; +/** The longest one status run (`runStatusCompletion`, one completion) waits. */ +export const STATUS_RUN_CEILING_MS = COMPLETION_CEILING_MS; + /** * Room a job's expiry leaves on top of its model work for what the handler * does around it: reading the record, building the prompt, persisting. diff --git a/src/lib/jobs/queue-runtime.ts b/src/lib/jobs/queue-runtime.ts index 230cfbc87..94e966bfe 100644 --- a/src/lib/jobs/queue-runtime.ts +++ b/src/lib/jobs/queue-runtime.ts @@ -15,10 +15,11 @@ * quarters of the default expiry (`jobBudget`) and hands the rest to its * next run, so it ends before fifteen minutes by construction. * - `bounded`: one unit of model work (one account's note, refresh or - * memory pass) whose calls wait the person's response timeout each, up - * to ten minutes, which can outlast fifteen minutes. Every call sits - * behind an outer bound, so the job's longest run is known - * (`modelSeconds`, from `model-time.ts`); it carries an explicit + * memory pass, a document summary) whose calls wait the person's + * response timeout each, up to ten minutes, which can outlast fifteen + * minutes. Every call carries its own abort and most sit behind an outer + * bound, so the job's longest run is known (`modelSeconds`, from + * `model-time.ts`); it carries an explicit * `expireInSeconds` above that on every send, and no retry can start * beside a run that is still going (#1274). * - `long`: a job that may run for hours. It carries an explicit @@ -58,9 +59,12 @@ import { COACH_MEMORY_REFRESH_EXPIRE_SECONDS, } from "@/lib/ai/coach/coach-memory-shared"; import { + COMPLETION_CEILING_MS, COMPREHENSIVE_WARM_CEILING_MS, STATUS_RUN_CEILING_MS, } from "@/lib/jobs/model-time"; +import { DOCUMENT_INDEX_EXPIRE_SECONDS } from "@/lib/jobs/document-index-queue"; +import { DOCUMENT_SUMMARY_EXPIRE_SECONDS } from "@/lib/jobs/document-summary-queue"; /** pg-boss's default job expiry, in seconds. */ export const DEFAULT_JOB_EXPIRE_SECONDS = 15 * 60; @@ -547,7 +551,17 @@ export const QUEUE_RUNTIME: Readonly> = { exclusive: "lockedPass", why: "Up to 200 documents of one account, a provider transcription each.", }, - "document-index": short("One document."), + "document-index": { + runtime: "long", + expireInSeconds: DOCUMENT_INDEX_EXPIRE_SECONDS, + expiryVia: ["DOCUMENT_INDEX_EXPIRE_SECONDS"], + stop: { + file: "lib/jobs/reminder/document-handlers.ts", + fn: "handleDocumentIndex", + }, + exclusive: "lockedPass", + why: "One document: the transcription, then a lab report's staging reads, as many as fit before the job's deadline at the person's timeout.", + }, "document-ai-run": { runtime: "long", expireInSeconds: DOCUMENT_AI_RUN_EXPIRE_SECONDS, @@ -563,7 +577,13 @@ export const QUEUE_RUNTIME: Readonly> = { "Up to 500 conditional updates and one deleteMany.", ), "document-thumbnail": short("One document."), - "document-summary": short("One document."), + "document-summary": { + runtime: "bounded", + modelSeconds: COMPLETION_CEILING_MS / 1000, + expireInSeconds: DOCUMENT_SUMMARY_EXPIRE_SECONDS, + expiryVia: ["DOCUMENT_SUMMARY_EXPIRE_SECONDS"], + why: "One document's summary: one call and its reasoning repeat at the person's timeout.", + }, "document-summary-catchup": short("Enqueues at most 200 documents."), "document-thumbnail-backfill": short("Enqueues at most 1 000 documents."), "environment-fetch": short( diff --git a/src/lib/jobs/registry/maintenance.ts b/src/lib/jobs/registry/maintenance.ts index d0f2ec23e..146a43aa6 100644 --- a/src/lib/jobs/registry/maintenance.ts +++ b/src/lib/jobs/registry/maintenance.ts @@ -667,6 +667,10 @@ const jobs: readonly JobDefinition[] = [ "Per-document, no re-converging discovery pass. Collapse queued duplicates only, so a re-index after a content change is never dropped.", }, work: { localConcurrency: DOCUMENT_INDEX_CONCURRENCY }, + // One index per document at a time: a delivery after the expiry finds the + // reads still running and reports so instead of reading beside them. + lock: (job: Job<{ documentId?: string }>) => + `document:${job.data.documentId ?? ""}`, load: async () => (await documents()).handleDocumentIndex, }), // One run per job under its run id's lock, so a redelivery never reads diff --git a/src/lib/jobs/reminder/document-handlers.ts b/src/lib/jobs/reminder/document-handlers.ts index fa7c2d55c..6bdfa2fc2 100644 --- a/src/lib/jobs/reminder/document-handlers.ts +++ b/src/lib/jobs/reminder/document-handlers.ts @@ -31,7 +31,7 @@ import { runThumbnailBackfillForUser, type ThumbnailBackfillPayload, } from "@/lib/jobs/document-thumbnail-backfill"; -import { jobBudget } from "@/lib/jobs/job-budget"; +import { jobBudget, jobDeadline } from "@/lib/jobs/job-budget"; import { jobDone, type JobOutcome } from "@/lib/jobs/job-outcome"; import { recordError } from "@/lib/jobs/worker-status"; @@ -102,7 +102,7 @@ export async function handleContentIndexBackfill( /** One stored document's job, by the step that runs it. */ function perDocument

( name: string, - run: (payload: P) => Promise, + run: (payload: P, job: Job

) => Promise, ): (jobs: Job

[]) => Promise { return async (jobs) => { let documents = 0; @@ -114,7 +114,7 @@ function perDocument

( continue; } try { - await run(job.data); + await run(job.data, job); documents++; } catch (err) { recordError(); @@ -132,12 +132,18 @@ function perDocument

( /** * Content indexing on upload: the vision provider when configured and - * consented, the local text layer otherwise. + * consented, the local text layer otherwise. The lab staging after it plans + * its reads against the job's own deadline (`jobDeadline`), so a long report + * read by a slow model stops before pg-boss's expiry instead of being retried + * beside itself (#1274). */ -export const handleDocumentIndex = perDocument( - "document-index", - runDocumentIndex, -); +export function handleDocumentIndex( + jobs: Job[], +): Promise { + return perDocument("document-index", (payload, job) => + runDocumentIndex(payload, { deadlineAt: jobDeadline(job) }), + )(jobs); +} /** The preview thumbnail on upload: local decode and downscale, no egress. */ export const handleDocumentThumbnail = perDocument(