fix(ai-runs): do not count time in the document reader against a run's deadline

A background AI run's deadline is armed in model time plus a fixed margin.
Time spent waiting in the document reader's queue, and the read itself, came
out of that margin, so a person's lab scan or document read could be timed
out by the reaper because other work held the reader. When the reader grants
the slot it now moves the run's deadline by the wait plus the most the read
may take, and gives back the unused part when the read ends, never past the
job's own expiry and only while the run is running. The interactive queue
wait stays within the margin, which covers the run until the slot is
granted.

src/lib/documents/__tests__/pdf-isolation.test.ts

@@ -12,6 +12,7 @@ vi.mock("@/lib/logging/context", () => ({ annotate }));
 import { gif } from "../../__tests__/gif-fixture";
 import {
   MAX_QUEUED_PER_OWNER,
+  READER_WAIT_MS,
   runOnce,
   runPdfChild,
   withPdfWork,
@@ -301,3 +302,76 @@ describe("the reader's queue", () => {
     }
   });
 });
+
+describe("the caller's clock while it waits for the reader", () => {
+  const request = { op: "text", data: Buffer.from("x"), maxPages: 1 } as const;
+
+  it("moves by the wait plus the read's allowance when the slot is granted, then gives back what the read did not use", async () => {
+    let open: () => void = () => {};
+    const gate = new Promise<void>((resolve) => (open = resolve));
+    const holder = vi.fn(async (): Promise<PdfChildResult<PdfChildRequest>> => {
+      await gate;
+      return { ok: true, value: "first" };
+    });
+    const quick = vi.fn(async (): Promise<PdfChildResult<PdfChildRequest>> => ({
+      ok: true,
+      value: "second",
+    }));
+    const moves: number[] = [];
+    const first = withPdfWork({ owner: "a", lane: "interactive" }, () =>
+      runPdfChild(request, { timeoutMs: 1_000, run: holder }),
+    );
+    const second = withPdfWork(
+      {
+        owner: "b",
+        lane: "interactive",
+        extendDeadline: async (ms) => {
+          moves.push(ms);
+        },
+      },
+      () => runPdfChild(request, { timeoutMs: 1_000, run: quick }),
+    );
+    await new Promise((resolve) => setTimeout(resolve, 150));
+    open();
+    await expect(first).resolves.toMatchObject({ ok: true });
+    await expect(second).resolves.toMatchObject({ ok: true, value: "second" });
+
+    // Granted after about 150 ms of waiting, with two attempts' allowance.
+    expect(moves).toHaveLength(2);
+    expect(moves[0]!).toBeGreaterThanOrEqual(2_000 + 140);
+    expect(moves[0]!).toBeLessThan(2_000 + 1_000);
+    // The quick read gives back nearly all of its allowance.
+    expect(moves[1]!).toBeLessThan(0);
+    expect(moves[1]!).toBeGreaterThan(-2_000 - 1);
+    // Net: the wait and the read's own time, nothing more.
+    expect(moves[0]! + moves[1]!).toBeGreaterThanOrEqual(140);
+    expect(moves[0]! + moves[1]!).toBeLessThan(1_000);
+  });
+
+  it("reads on when moving the clock fails", async () => {
+    const run = vi.fn(async (): Promise<PdfChildResult<PdfChildRequest>> => ({
+      ok: true,
+      value: "text",
+    }));
+    await expect(
+      withPdfWork(
+        {
+          owner: "c",
+          lane: "interactive",
+          extendDeadline: async () => {
+            throw new Error("database gone");
+          },
+        },
+        () => runPdfChild(request, { timeoutMs: 1_000, run }),
+      ),
+    ).resolves.toEqual({ ok: true, value: "text" });
+  });
+
+  it("waits no longer for the reader than a run's margin can absorb before the slot moves its deadline", async () => {
+    // The deadline moves when the slot is granted; until then the run's
+    // armed margin is all that stands between a waiting read and the
+    // reaper.
+    const { AI_RUN_MARGIN_MS } = await import("../ai-runs/types");
+    expect(READER_WAIT_MS.interactive).toBeLessThanOrEqual(AI_RUN_MARGIN_MS);
+  });
+});

src/lib/documents/pdf-isolation.ts

@@ -143,6 +143,16 @@ export interface PdfWork {
   /** The account the work is done for; the queue takes accounts in turn. */
   owner: string;
   lane: ReaderLane;
+  /**
+   * Moves the caller's own clock by `ms`. A background AI run has a deadline
+   * measured in model calls, and time spent waiting for the reader or
+   * reading is not model time: when the slot is granted the clock moves by
+   * the wait plus the most the read may take, and when the read ends it moves
+   * back by the part of that allowance the read did not use. So the run's
+   * model-time budget starts once the reader is done with it, however long
+   * the queue was. Failures are the caller's to swallow; the read goes on.
+   */
+  extendDeadline?: (ms: number) => Promise<void>;
 }
 
 const workContext = new AsyncLocalStorage<PdfWork>();
@@ -265,10 +275,15 @@ export async function runPdfChild<R extends PdfChildRequest>(
   }
 
   const work = currentPdfWork();
+  const queuedAt = Date.now();
   const ticket = await queue.acquire(work.owner, work.lane);
   if (!ticket.ok) {
     return { ok: false, reason: "busy", retryAfterMs: ticket.retryAfterMs };
   }
+  const grantedAt = Date.now();
+  // The most this read may take: every attempt to its time limit.
+  const allowance = limits.timeoutMs * MAX_ATTEMPTS;
+  await moveCallerClock(work, grantedAt - queuedAt + allowance);
   try {
     for (let attempt = 1; ; attempt++) {
       const result = limits.run
@@ -279,6 +294,19 @@ export async function runPdfChild<R extends PdfChildRequest>(
     }
   } finally {
     ticket.release();
+    await moveCallerClock(
+      work,
+      -Math.max(0, allowance - (Date.now() - grantedAt)),
+    );
+  }
+}
+
+async function moveCallerClock(work: PdfWork, ms: number): Promise<void> {
+  if (!work.extendDeadline || ms === 0) return;
+  try {
+    await work.extendDeadline(ms);
+  } catch {
+    // The caller's clock is the caller's; the read is not failed for it.
   }
 }
 

src/lib/jobs/__tests__/document-ai-run.test.ts

@@ -293,7 +293,7 @@ describe("runDocumentAiRun", () => {
       return { ok: true, data: {} };
     });
     await runDocumentAiRun("run1", NOW + 2_700_000, now);
-    expect(seen).toEqual({ owner: "u1", lane: "interactive" });
+    expect(seen).toMatchObject({ owner: "u1", lane: "interactive" });
   });
 
   it("answers a busy document reader with a retryable code and hands the reservation back", async () => {

src/lib/jobs/document-ai-run.ts

@@ -187,6 +187,34 @@ async function armDeadline(
   });
 }
 
+/**
+ * Move a running run's deadline by `ms`, never past the job's own expiry.
+ * The document reader calls it (`PdfWork.extendDeadline`): the deadline is
+ * model time, and a wait in the reader's queue or the read itself is not, so
+ * a person's read is not timed out because other work held the reader.
+ */
+export async function moveRunDeadline(
+  runId: string,
+  ms: number,
+  deadlineAt: number | undefined,
+): Promise<void> {
+  const delta = Math.round(ms);
+  if (deadlineAt === undefined) {
+    await prisma.$executeRaw`
+      UPDATE document_ai_runs
+      SET expires_at = expires_at + (${delta}::double precision * interval '1 millisecond')
+      WHERE id = ${runId} AND status = 'RUNNING'`;
+    return;
+  }
+  await prisma.$executeRaw`
+    UPDATE document_ai_runs
+    SET expires_at = LEAST(
+      expires_at + (${delta}::double precision * interval '1 millisecond'),
+      ${new Date(deadlineAt)}
+    )
+    WHERE id = ${runId} AND status = 'RUNNING'`;
+}
+
 async function runDocumentIndex(
   run: ClaimedAiRun,
   deadlineAt: number | undefined,
@@ -619,7 +647,11 @@ export async function runDocumentAiRun(
     // A person asked for this read and is waiting on it: the document
     // reader serves it ahead of background work, in turn with other people.
     outcome = await withPdfWork(
-      { owner: run.userId, lane: "interactive" },
+      {
+        owner: run.userId,
+        lane: "interactive",
+        extendDeadline: (ms) => moveRunDeadline(run.id, ms, deadlineAt),
+      },
       () => RUN_DISPATCH[run.kind](run, deadlineAt, now),
     );
   } catch (err) {

tests/integration/document-ai-runs.test.ts

@@ -540,6 +540,68 @@ describe("the reaper", () => {
   });
 });
 
+/**
+ * A run's deadline is model time. The document reader moves it by the time
+ * a read waited for the reader and took, so a person's lab scan is not
+ * timed out because other work held the reader; never past the job's own
+ * expiry, and only while the run is running.
+ */
+describe("the deadline while the document reader works", () => {
+  async function runningRun(expiresAt: Date) {
+    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 };
+    await getPrismaClient().documentAiRun.update({
+      where: { id: runId },
+      data: { status: "RUNNING", expiresAt },
+    });
+    return runId;
+  }
+  const expiry = async (runId: string) =>
+    (
+      await getPrismaClient().documentAiRun.findUniqueOrThrow({
+        where: { id: runId },
+      })
+    ).expiresAt.getTime();
+
+  it("moves a running run's deadline by the reader's time, so the reaper leaves it alone", async () => {
+    const { moveRunDeadline } = await import("@/lib/jobs/document-ai-run");
+    const { reapAiRuns } = await import("@/lib/documents/ai-runs/store");
+    const armed = new Date(Date.now() + 5_000);
+    const runId = await runningRun(armed);
+
+    // The slot is granted after a 60 s wait, with 70 s allowed for the read.
+    await moveRunDeadline(runId, 130_000, undefined);
+    expect(await expiry(runId)).toBe(armed.getTime() + 130_000);
+    // Without the move the reaper would have taken it a minute later.
+    expect(
+      (await reapAiRuns(new Date(armed.getTime() + 60_000))).timedOut,
+    ).toBe(0);
+
+    // The read took 10 s of its 70: the unused 60 s go back.
+    await moveRunDeadline(runId, -60_000, undefined);
+    expect(await expiry(runId)).toBe(armed.getTime() + 70_000);
+  });
+
+  it("never moves it past the job's own expiry, and leaves a run that is not running alone", async () => {
+    const { moveRunDeadline } = await import("@/lib/jobs/document-ai-run");
+    const armed = new Date(Date.now() + 5_000);
+    const runId = await runningRun(armed);
+    const jobEnds = armed.getTime() + 20_000;
+    await moveRunDeadline(runId, 130_000, jobEnds);
+    expect(await expiry(runId)).toBe(jobEnds);
+
+    await getPrismaClient().documentAiRun.update({
+      where: { id: runId },
+      data: { status: "FAILED" },
+    });
+    await moveRunDeadline(runId, 130_000, undefined);
+    expect(await expiry(runId)).toBe(jobEnds);
+  });
+});
+
 describe("a summary in the background", () => {
   it("marks a stored summary PENDING while queued and READY when it lands", async () => {
     const user = await makeUser();