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(