diff --git a/src/lib/jobs/__tests__/document-content-index-backfill-failing.test.ts b/src/lib/jobs/__tests__/document-content-index-backfill-failing.test.ts index dcdd6126e..98ea5f7ab 100644 --- a/src/lib/jobs/__tests__/document-content-index-backfill-failing.test.ts +++ b/src/lib/jobs/__tests__/document-content-index-backfill-failing.test.ts @@ -77,6 +77,9 @@ vi.mock("@/lib/ai/coach/budget", () => ({ resolveCostOwner: () => "user", })); vi.mock("@/lib/jobs/boss-instance", () => ({ getGlobalBoss: vi.fn() })); +vi.mock("@/lib/ai/capabilities/egress", () => ({ + aiEgressRefusal: vi.fn().mockResolvedValue(null), +})); import { BACKFILL_FAILURES_IN_A_ROW, diff --git a/src/lib/jobs/__tests__/document-content-index-backfill.test.ts b/src/lib/jobs/__tests__/document-content-index-backfill.test.ts index 189bcc9b9..3a668daf3 100644 --- a/src/lib/jobs/__tests__/document-content-index-backfill.test.ts +++ b/src/lib/jobs/__tests__/document-content-index-backfill.test.ts @@ -54,6 +54,9 @@ vi.mock("@/lib/ai/coach/budget", () => ({ resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/jobs/boss-instance", () => ({ getGlobalBoss: vi.fn() })); +vi.mock("@/lib/ai/capabilities/egress", () => ({ + aiEgressRefusal: vi.fn().mockResolvedValue(null), +})); import { runContentIndexBackfillForUser } from "../document-content-index-backfill"; import { prisma } from "@/lib/db"; @@ -62,7 +65,8 @@ import { rasterizePdf } from "@/lib/documents/rasterize-pdf"; import { transcribeDocument } from "@/lib/documents/describe"; import { resolveDocumentVisionProvider } from "@/lib/documents/provider-order"; import { AiUnavailableError } from "@/lib/ai/capabilities/refusal"; -import { reserveBudget } from "@/lib/ai/coach/budget"; +import { reconcileSpend, reserveBudget } from "@/lib/ai/coach/budget"; +import { aiEgressRefusal } from "@/lib/ai/capabilities/egress"; import { aiCapabilityForJob } from "@/lib/ai/capabilities/gate"; import { decryptIndexText, @@ -137,6 +141,14 @@ function seedIndexRows( }) as never); } +/** A pick, and one page of candidates with these ids. */ +function walk(ids: string[]) { + vi.mocked(resolveDocumentVisionProvider).mockResolvedValue(PICK as never); + vi.mocked(prisma.inboundDocument.findMany).mockResolvedValue( + ids.map((id) => ({ id })) as never, + ); +} + describe("runContentIndexBackfillForUser", () => { it("no-ops with no provider", async () => { vi.mocked(resolveDocumentVisionProvider).mockResolvedValue({ @@ -253,13 +265,9 @@ describe("runContentIndexBackfillForUser", () => { ); it("indexes the not-yet-indexed documents", async () => { - vi.mocked(resolveDocumentVisionProvider).mockResolvedValue(PICK as never); // One short page (< PAGE_SIZE) — the walk breaks after it, so a single // findMany return is enough (a trailing once would leak to the next test). - vi.mocked(prisma.inboundDocument.findMany).mockResolvedValue([ - { id: "d1" }, - { id: "d2" }, - ] as never); + walk(["d1", "d2"]); const result = await runContentIndexBackfillForUser("user-1"); expect(result).toEqual({ @@ -276,6 +284,36 @@ describe("runContentIndexBackfillForUser", () => { ).toMatchObject({ aiReadDeferred: false }); }); + it("settles a document's reservation once when its index write fails after the read", async () => { + walk(["d1"]); + vi.mocked(upsertContentIndex).mockRejectedValueOnce(new Error("write")); + + const result = await runContentIndexBackfillForUser("user-1"); + expect(result).toMatchObject({ indexed: 0, failed: 1 }); + // The read happened and is charged; nothing is handed back a second time. + expect(reconcileSpend).toHaveBeenCalledTimes(1); + expect(vi.mocked(reconcileSpend).mock.calls[0]![2]).toBe(1); + }); + + it("asks the wire again before every document's read", async () => { + walk(["d1", "d2"]); + // Consent withdrawn while the walk is between the two documents. + vi.mocked(aiEgressRefusal) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce( + new AiUnavailableError("documentAi", "consent_required"), + ); + + const result = await runContentIndexBackfillForUser("user-1"); + expect(transcribeDocument).toHaveBeenCalledTimes(1); + expect(aiEgressRefusal).toHaveBeenCalledWith("documentAi", "user-1", [ + "anthropic", + ]); + expect(result).toMatchObject({ indexed: 1, reason: "no-consent" }); + // The second document's reservation went back untouched. + expect(vi.mocked(reconcileSpend).mock.calls.at(-1)![2]).toBe(0); + }); + it("includes PDFs as candidates for an image-input provider and rasterises them", async () => { // The per-document authority (`prepareVisionInput`) renders a PDF to page // images when the provider cannot take PDFs natively. The candidate MIME diff --git a/src/lib/jobs/document-content-index-backfill.ts b/src/lib/jobs/document-content-index-backfill.ts index e6e59197a..45fc3d078 100644 --- a/src/lib/jobs/document-content-index-backfill.ts +++ b/src/lib/jobs/document-content-index-backfill.ts @@ -31,6 +31,7 @@ * (`src/lib/jobs/reminder/register-maintenance.ts`) so pg-boss provisions it. */ import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; +import { aiEgressRefusal } from "@/lib/ai/capabilities/egress"; import { aiCapabilityForJob } from "@/lib/ai/capabilities/gate"; import { PICK_DECIDED_REASONS } from "@/lib/ai/capabilities/types"; import { @@ -333,6 +334,22 @@ export async function runContentIndexBackfillForUser( continue; } + // The wire again, before every read: a consent withdrawn or a switch + // turned off while the walk runs stops it at the next document, not + // at the next run. + const refusal = await aiEgressRefusal("documentAi", userId, [ + pick.providerType, + ]); + if (refusal) { + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); + reason = + refusal.reason === "consent_required" ? "no-consent" : "unavailable"; + break outer; + } + calls += 1; const read = await transcribeAndIndex(userId, document.id, pick, vision, { reserved: reservation.reserved, @@ -392,14 +409,24 @@ async function transcribeAndIndex( servedBy, reservedOwner: budget.owner, }); + let text: string; try { - const { text } = await transcribeDocument({ + ({ text } = await transcribeDocument({ provider: pick.entry.instance, providerType: pick.providerType, images: vision.images, documents: vision.documents, - }); - await settle(budget.reserved, pick.entry.providerType); + })); + } catch (err) { + await settle(spentTokensOf(err) ?? 0, null); + // Counted into the run's `failed` tally by the caller. + return { error: err }; + } + // Settled once, for the read that was made. The index write below is ours, + // not the provider's: a failure there charges nothing back, or the + // reservation would be handed back a second time. + await settle(budget.reserved, pick.entry.providerType); + try { await upsertContentIndex({ userId, documentId, @@ -407,12 +434,10 @@ async function transcribeAndIndex( source: "vision", providerType: pick.providerType, }); - return "indexed"; } catch (err) { - await settle(spentTokensOf(err) ?? 0, null); - // Counted into the run's `failed` tally by the caller. return { error: err }; } + return "indexed"; } /**