From badef28e965384d75c0a6c635f920d9dd2919f73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Fri, 11 Sep 2026 09:42:14 +0200 Subject: [PATCH] feat(ai): the daily ledger records which tokens the operator paid for `coach_usage` counted one number per (user, day). A turn served by the user's own ChatGPT plan after the operator's shared key answered 500 was therefore charged to the operator-cost cap, and the operator's chat was locked out every morning by background work he pays for himself. The ledger gains `operator_tokens`: the share of `total_tokens` served by an operator-funded provider. `reserveBudget` books a reservation to the owner of the chain's primary and `reconcileSpend` settles it against the provider that actually served, both inside the same statements as the total so concurrent requests can never see one counter without the other. Every call site names its cost owner and its serving hop; the argument is required, so a new surface cannot inherit the old blind accounting. --- .../migration.sql | 15 ++ prisma/schema.prisma | 20 +- src/app/api/ai/test/__tests__/route.test.ts | 5 + src/app/api/ai/test/route.ts | 16 +- .../inbound/[id]/chat/__tests__/route.test.ts | 1 + .../[id]/extract/__tests__/route.test.ts | 1 + .../documents/inbound/[id]/extract/route.ts | 25 ++- .../[id]/index/__tests__/route.test.ts | 1 + .../api/documents/inbound/[id]/index/route.ts | 9 +- .../[id]/suggest/__tests__/route.test.ts | 1 + .../documents/inbound/[id]/suggest/route.ts | 17 +- .../[id]/summary/__tests__/route.test.ts | 1 + .../documents/inbound/[id]/summary/route.ts | 17 +- .../__tests__/route-cancelled-turn.test.ts | 24 ++- .../chat/__tests__/route-fence-guard.test.ts | 1 + .../__tests__/route-guard-activation.test.ts | 1 + .../chat/__tests__/route-guard-ledger.test.ts | 1 + .../chat/__tests__/route-guard.test.ts | 1 + .../chat/__tests__/route-list-query.test.ts | 1 + .../chat/__tests__/route-module-gate.test.ts | 1 + .../__tests__/route-snapshot-once.test.ts | 1 + .../chat/__tests__/route-tool-mode.test.ts | 14 +- .../__tests__/route-workout-scope.test.ts | 1 + src/app/api/insights/chat/route.ts | 20 +- .../labs/ocr/extract/__tests__/route.test.ts | 11 ++ src/app/api/labs/ocr/extract/route.ts | 65 ++++-- .../extract/__tests__/route.test.ts | 19 +- src/app/api/medications/extract/route.ts | 14 +- src/lib/ai/__tests__/provider-runner.test.ts | 7 +- src/lib/ai/coach/__tests__/budget.test.ts | 42 +++- .../self-context-questions-consent.test.ts | 11 ++ src/lib/ai/coach/budget.ts | 135 +++++++++++-- src/lib/ai/coach/self-context-questions.ts | 10 +- src/lib/ai/provider-runner.ts | 9 +- .../__tests__/fenced-chat-pipeline.test.ts | 1 + .../__tests__/index-document.test.ts | 20 +- .../__tests__/index-outcome-record.test.ts | 1 + src/lib/documents/auto-stage-labs.ts | 11 +- src/lib/documents/fenced-chat.ts | 10 +- src/lib/documents/index-document.ts | 24 ++- src/lib/insights/briefing-provider.ts | 19 +- src/lib/insights/status-provider.ts | 17 +- .../__tests__/async-state-contract.test.ts | 1 + .../__tests__/coach-nudge-ai-consent.test.ts | 1 + .../document-content-index-backfill.test.ts | 1 + .../jobs/__tests__/document-summary.test.ts | 20 +- .../reaction-line-degradation.test.ts | 6 + src/lib/jobs/coach-nudge-ai.ts | 13 +- .../jobs/document-content-index-backfill.ts | 22 ++- src/lib/jobs/document-summary.ts | 9 +- src/lib/jobs/reaction-line.ts | 30 ++- tests/integration/coach-budget-owner.test.ts | 186 ++++++++++++++++++ .../coach-budget-reservation.test.ts | 20 +- 53 files changed, 813 insertions(+), 117 deletions(-) create mode 100644 prisma/migrations/0341_coach_usage_operator_tokens/migration.sql create mode 100644 tests/integration/coach-budget-owner.test.ts diff --git a/prisma/migrations/0341_coach_usage_operator_tokens/migration.sql b/prisma/migrations/0341_coach_usage_operator_tokens/migration.sql new file mode 100644 index 000000000..8414906ba --- /dev/null +++ b/prisma/migrations/0341_coach_usage_operator_tokens/migration.sql @@ -0,0 +1,15 @@ +-- v1.38.19 (Wave E) — the daily coach ledger gains a cost-owner dimension. +-- +-- `coach_usage` counted one number per (user, day), so tokens the USER's own +-- plan paid for (a `codex` fallback hop after the operator's shared key +-- answered 500) were charged against the operator-cost cap and locked the +-- interactive chat out every morning. `operator_tokens` is the share of +-- `total_tokens` served by an operator-funded provider (`admin-openai` / +-- `admin-codex`); the operator cap is enforced against it, the user-plan cap +-- against `total_tokens`. +-- +-- Existing rows start at 0: their owner split is unknown and 0 is the side +-- that cannot lock a user out of a chat the operator never paid for. The +-- counter is exact from the first turn after deploy. +ALTER TABLE "coach_usage" + ADD COLUMN IF NOT EXISTS "operator_tokens" INTEGER NOT NULL DEFAULT 0; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 71bcfeb6d..688951392 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -6267,13 +6267,19 @@ model HealthProfileFactRevision { // secondary signal used by the frontend to render a "≈ N replies left" // hint without exposing raw token figures. model CoachUsage { - id String @id @default(cuid()) - userId String @map("user_id") - dateKey String @map("date_key") - totalTokens Int @default(0) @map("total_tokens") - messageCount Int @default(0) @map("message_count") - updatedAt DateTime @updatedAt @map("updated_at") - createdAt DateTime @default(now()) @map("created_at") + id String @id @default(cuid()) + userId String @map("user_id") + dateKey String @map("date_key") + totalTokens Int @default(0) @map("total_tokens") + /// v1.38.19 (Wave E) — the share of `totalTokens` served by a provider the + /// OPERATOR funds (`admin-openai` / `admin-codex`). The operator-cost cap is + /// enforced against this counter and the user-plan cap against + /// `totalTokens`, so background work served by the user's own plan can no + /// longer lock the interactive chat out of the operator's ceiling. + operatorTokens Int @default(0) @map("operator_tokens") + messageCount Int @default(0) @map("message_count") + updatedAt DateTime @updatedAt @map("updated_at") + createdAt DateTime @default(now()) @map("created_at") user User @relation(fields: [userId], references: [id], onDelete: Cascade) diff --git a/src/app/api/ai/test/__tests__/route.test.ts b/src/app/api/ai/test/__tests__/route.test.ts index 6369d1b68..a902128bc 100644 --- a/src/app/api/ai/test/__tests__/route.test.ts +++ b/src/app/api/ai/test/__tests__/route.test.ts @@ -30,9 +30,12 @@ vi.mock("@/lib/ai/coach/budget", () => ({ allowed: true, reserved: 32, totalAfter: 32, + owner: "operator" as const, + operatorAfter: 32, })), reconcileSpend: vi.fn(async () => undefined), resolveDailyCap: () => 200_000, + resolveCostOwner: () => "operator" as const, })); vi.mock("@/lib/logging/context", () => ({ @@ -376,6 +379,8 @@ describe("POST /api/ai/test — daily ceiling + ledger", () => { 32, expect.any(Number), "2026-01-01", + 0, + { servedBy: expect.any(String), reservedOwner: "operator" }, ); }); }); diff --git a/src/app/api/ai/test/route.ts b/src/app/api/ai/test/route.ts index 56440ce29..389938c71 100644 --- a/src/app/api/ai/test/route.ts +++ b/src/app/api/ai/test/route.ts @@ -19,6 +19,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { aiTestOverrideSchema } from "@/lib/validations/ai-provider"; @@ -113,6 +114,11 @@ export const POST = apiHandler(async (request: NextRequest) => { providerType: provider.type === "admin-key" ? "admin-openai" : "local", }, ]), + resolveCostOwner([ + { + providerType: provider.type === "admin-key" ? "admin-openai" : "local", + }, + ]), ); if (!reservation.allowed) { annotate({ action: { name: "ai.test.budget_exceeded" } }); @@ -137,6 +143,11 @@ export const POST = apiHandler(async (request: NextRequest) => { reservation.reserved, result.tokensUsed ?? reservation.reserved, dateKey, + 0, + { + servedBy: provider.type === "admin-key" ? "admin-openai" : "local", + reservedOwner: reservation.owner, + }, ); return apiSuccess({ @@ -155,7 +166,10 @@ export const POST = apiHandler(async (request: NextRequest) => { // A failed probe produced no reported usage — refund the reservation so a // provider that is simply misconfigured doesn't burn the caller's ledger // while they fix it. - await reconcileSpend(user.id, reservation.reserved, 0, dateKey); + await reconcileSpend(user.id, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); annotate({ meta: { ai_test_error: err.message.slice(0, 500), diff --git a/src/app/api/documents/inbound/[id]/chat/__tests__/route.test.ts b/src/app/api/documents/inbound/[id]/chat/__tests__/route.test.ts index 1fa158e40..7d832c732 100644 --- a/src/app/api/documents/inbound/[id]/chat/__tests__/route.test.ts +++ b/src/app/api/documents/inbound/[id]/chat/__tests__/route.test.ts @@ -51,6 +51,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 600 }), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 100000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/provider-runner", () => ({ runStreamingRawCompletionWithFallback: vi.fn(), diff --git a/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts b/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts index 6acb80659..75aa91905 100644 --- a/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts +++ b/src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts @@ -69,6 +69,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/documents/extract", () => ({ runInboundExtraction: vi.fn(), diff --git a/src/app/api/documents/inbound/[id]/extract/route.ts b/src/app/api/documents/inbound/[id]/extract/route.ts index 757c012a7..75b43c825 100644 --- a/src/app/api/documents/inbound/[id]/extract/route.ts +++ b/src/app/api/documents/inbound/[id]/extract/route.ts @@ -41,6 +41,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { auditLog } from "@/lib/auth/audit"; @@ -254,6 +255,7 @@ async function handleTextExtract( AI_BUDGETS.ocrExtractText.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -273,6 +275,8 @@ async function handleTextExtract( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); const updated = await stageExtraction(document.id, userId, result); @@ -289,7 +293,10 @@ async function handleTextExtract( return apiSuccess(serialiseDocumentDetail(updated, updated.facts)); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (err instanceof InboundExtractError) { return apiError("Couldn't read the document. Try a clearer copy.", 422, { errorCode: "documents.inbound.extractFailed", @@ -351,6 +358,7 @@ async function handleStoredExtract( AI_BUDGETS.ocrExtractText.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -370,6 +378,8 @@ async function handleStoredExtract( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); const updated = await stageExtraction(document.id, userId, result); @@ -386,7 +396,10 @@ async function handleStoredExtract( return apiSuccess(serialiseDocumentDetail(updated, updated.facts)); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (err instanceof InboundExtractError) { return apiError("Couldn't read the document. Try a clearer copy.", 422, { errorCode: "documents.inbound.extractFailed", @@ -465,6 +478,7 @@ async function handleVisionExtract( AI_BUDGETS.ocrExtract.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -493,6 +507,8 @@ async function handleVisionExtract( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); const updated = await stageExtraction(document.id, userId, result); @@ -509,7 +525,10 @@ async function handleVisionExtract( return apiSuccess(serialiseDocumentDetail(updated, updated.facts)); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (err instanceof InboundExtractError) { return apiError("Couldn't read the document. Try a clearer copy.", 422, { errorCode: "documents.inbound.extractFailed", 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 aa5b6c7d5..544005665 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 @@ -52,6 +52,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/modules/gate", () => ({ requireModuleEnabled: vi.fn().mockResolvedValue({ enabled: true }), diff --git a/src/app/api/documents/inbound/[id]/index/route.ts b/src/app/api/documents/inbound/[id]/index/route.ts index 9a3cd6756..87f37415a 100644 --- a/src/app/api/documents/inbound/[id]/index/route.ts +++ b/src/app/api/documents/inbound/[id]/index/route.ts @@ -37,6 +37,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { auditLog } from "@/lib/auth/audit"; @@ -249,6 +250,7 @@ async function handleVisionIndex( AI_BUDGETS.documentTranscribe.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -269,6 +271,8 @@ async function handleVisionIndex( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); // Refs #776 — the empty-transcription guard, same contract as the auto // path (`tryProviderIndex`): a provider answer with no text must never @@ -295,7 +299,10 @@ async function handleVisionIndex( }); return finishIndex(request, userId, document.id, "vision", tokenCount); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); await recordIndexAttempt(userId, document.id, { indexed: false, reason: "provider-error", diff --git a/src/app/api/documents/inbound/[id]/suggest/__tests__/route.test.ts b/src/app/api/documents/inbound/[id]/suggest/__tests__/route.test.ts index 3db7bbd6a..40f2e9603 100644 --- a/src/app/api/documents/inbound/[id]/suggest/__tests__/route.test.ts +++ b/src/app/api/documents/inbound/[id]/suggest/__tests__/route.test.ts @@ -40,6 +40,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/modules/gate", () => ({ requireModuleEnabled: vi.fn().mockResolvedValue({ enabled: true }), diff --git a/src/app/api/documents/inbound/[id]/suggest/route.ts b/src/app/api/documents/inbound/[id]/suggest/route.ts index b7feb1f51..61bd2499d 100644 --- a/src/app/api/documents/inbound/[id]/suggest/route.ts +++ b/src/app/api/documents/inbound/[id]/suggest/route.ts @@ -29,6 +29,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { auditLog } from "@/lib/auth/audit"; @@ -166,6 +167,7 @@ async function handleTextSuggest( AI_BUDGETS.documentAssist.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -185,10 +187,15 @@ async function handleTextSuggest( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); return finishSuggest(request, userId, document.id, "text", suggestion); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (err instanceof DocumentAssistError) { return apiError("Couldn't read the document. Try a clearer copy.", 422, { errorCode: "documents.inbound.extractFailed", @@ -255,6 +262,7 @@ async function handleVisionSuggest( AI_BUDGETS.documentAssist.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -275,10 +283,15 @@ async function handleVisionSuggest( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); return finishSuggest(request, userId, document.id, "vision", suggestion); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (err instanceof DocumentAssistError) { return apiError("Couldn't read the document. Try a clearer copy.", 422, { errorCode: "documents.inbound.extractFailed", diff --git a/src/app/api/documents/inbound/[id]/summary/__tests__/route.test.ts b/src/app/api/documents/inbound/[id]/summary/__tests__/route.test.ts index b445164cf..2bdd26141 100644 --- a/src/app/api/documents/inbound/[id]/summary/__tests__/route.test.ts +++ b/src/app/api/documents/inbound/[id]/summary/__tests__/route.test.ts @@ -53,6 +53,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/modules/gate", () => ({ requireModuleEnabled: vi.fn().mockResolvedValue({ enabled: true }), diff --git a/src/app/api/documents/inbound/[id]/summary/route.ts b/src/app/api/documents/inbound/[id]/summary/route.ts index 186e842d7..f0903b787 100644 --- a/src/app/api/documents/inbound/[id]/summary/route.ts +++ b/src/app/api/documents/inbound/[id]/summary/route.ts @@ -36,6 +36,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { auditLog } from "@/lib/auth/audit"; @@ -345,6 +346,7 @@ async function handleTextSummary( budgetFor(mode), dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -369,10 +371,15 @@ async function handleTextSummary( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); return finishSummary(request, userId, document.id, "text", mode, result); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return summaryError(err, "text"); } } @@ -430,6 +437,7 @@ async function handleVisionSummary( budgetFor(mode), dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { await refundDocumentAiSlot(userId); @@ -455,10 +463,15 @@ async function handleVisionSummary( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); return finishSummary(request, userId, document.id, "vision", mode, result); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return summaryError(err, "vision"); } } diff --git a/src/app/api/insights/chat/__tests__/route-cancelled-turn.test.ts b/src/app/api/insights/chat/__tests__/route-cancelled-turn.test.ts index 593a9fb7a..9640c3871 100644 --- a/src/app/api/insights/chat/__tests__/route-cancelled-turn.test.ts +++ b/src/app/api/insights/chat/__tests__/route-cancelled-turn.test.ts @@ -93,7 +93,11 @@ vi.mock("@/lib/ai/coach/facts", () => ({ })); const { reserveBudget, reconcileSpend } = vi.hoisted(() => ({ - reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 3000 })), + reserveBudget: vi.fn(async () => ({ + allowed: true, + reserved: 3000, + owner: "user" as const, + })), reconcileSpend: vi.fn(async () => undefined), })); vi.mock("@/lib/ai/coach/budget", () => ({ @@ -101,6 +105,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget, reconcileSpend, resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); vi.mock("@/lib/ai/coach/refusal", () => ({ @@ -205,7 +210,11 @@ function abortError(): Error { describe("coach chat — client abort mid-generation (#781)", () => { beforeEach(() => { vi.clearAllMocks(); - reserveBudget.mockResolvedValue({ allowed: true, reserved: 3000 }); + reserveBudget.mockResolvedValue({ + allowed: true, + reserved: 3000, + owner: "user" as const, + }); // No-tools chain so the turn runs the streaming runner (which rejects // with the abort the torn-down connection produces). resolveProviderChain.mockResolvedValue([ @@ -236,7 +245,16 @@ describe("coach chat — client abort mid-generation (#781)", () => { await post(abortedChatReq({ message: "How is my BP?" })); await sse.done; - expect(reconcileSpend).toHaveBeenCalledWith("u1", 3000, 0, "2026-08-13"); + expect(reconcileSpend).toHaveBeenCalledWith( + "u1", + 3000, + 0, + "2026-08-13", + 0, + // No hop served the cancelled turn, so nothing is attributed to an + // owner — the reservation is reversed on both counters. + { servedBy: null, reservedOwner: "user" }, + ); const actions = annotate.mock.calls .map((c) => (c[0] as { action?: { name?: string } }).action?.name) .filter(Boolean); diff --git a/src/app/api/insights/chat/__tests__/route-fence-guard.test.ts b/src/app/api/insights/chat/__tests__/route-fence-guard.test.ts index e187cc5b4..61a0dd543 100644 --- a/src/app/api/insights/chat/__tests__/route-fence-guard.test.ts +++ b/src/app/api/insights/chat/__tests__/route-fence-guard.test.ts @@ -86,6 +86,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 1500 })), reconcileSpend: vi.fn(async () => undefined), resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn(() => ({ refuse: false })), diff --git a/src/app/api/insights/chat/__tests__/route-guard-activation.test.ts b/src/app/api/insights/chat/__tests__/route-guard-activation.test.ts index b394f43b9..078fb58ac 100644 --- a/src/app/api/insights/chat/__tests__/route-guard-activation.test.ts +++ b/src/app/api/insights/chat/__tests__/route-guard-activation.test.ts @@ -134,6 +134,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget, reconcileSpend, resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); const { detectRefusal } = vi.hoisted(() => ({ diff --git a/src/app/api/insights/chat/__tests__/route-guard-ledger.test.ts b/src/app/api/insights/chat/__tests__/route-guard-ledger.test.ts index 913a32602..4eb694265 100644 --- a/src/app/api/insights/chat/__tests__/route-guard-ledger.test.ts +++ b/src/app/api/insights/chat/__tests__/route-guard-ledger.test.ts @@ -127,6 +127,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget, reconcileSpend, resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); const { detectRefusal } = vi.hoisted(() => ({ diff --git a/src/app/api/insights/chat/__tests__/route-guard.test.ts b/src/app/api/insights/chat/__tests__/route-guard.test.ts index abed117d0..d2bee8745 100644 --- a/src/app/api/insights/chat/__tests__/route-guard.test.ts +++ b/src/app/api/insights/chat/__tests__/route-guard.test.ts @@ -109,6 +109,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget, reconcileSpend, resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); const { detectRefusal } = vi.hoisted(() => ({ diff --git a/src/app/api/insights/chat/__tests__/route-list-query.test.ts b/src/app/api/insights/chat/__tests__/route-list-query.test.ts index c5165fdb7..96a2ef5cd 100644 --- a/src/app/api/insights/chat/__tests__/route-list-query.test.ts +++ b/src/app/api/insights/chat/__tests__/route-list-query.test.ts @@ -101,6 +101,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(), reconcileSpend: vi.fn(), resolveDailyCap: vi.fn(() => 200_000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn() })); vi.mock("@/lib/ai/coach/system-prompt", () => ({ diff --git a/src/app/api/insights/chat/__tests__/route-module-gate.test.ts b/src/app/api/insights/chat/__tests__/route-module-gate.test.ts index 21b9fde3c..fd0add566 100644 --- a/src/app/api/insights/chat/__tests__/route-module-gate.test.ts +++ b/src/app/api/insights/chat/__tests__/route-module-gate.test.ts @@ -100,6 +100,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ })), reconcileSpend: vi.fn(), resolveDailyCap: vi.fn(() => 200_000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn() })); vi.mock("@/lib/ai/coach/system-prompt", () => ({ diff --git a/src/app/api/insights/chat/__tests__/route-snapshot-once.test.ts b/src/app/api/insights/chat/__tests__/route-snapshot-once.test.ts index 4f9b399a7..7ba30cf4b 100644 --- a/src/app/api/insights/chat/__tests__/route-snapshot-once.test.ts +++ b/src/app/api/insights/chat/__tests__/route-snapshot-once.test.ts @@ -103,6 +103,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 1500 })), reconcileSpend: vi.fn(async () => undefined), resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn(() => ({ refuse: false })), diff --git a/src/app/api/insights/chat/__tests__/route-tool-mode.test.ts b/src/app/api/insights/chat/__tests__/route-tool-mode.test.ts index 827007729..d88dab268 100644 --- a/src/app/api/insights/chat/__tests__/route-tool-mode.test.ts +++ b/src/app/api/insights/chat/__tests__/route-tool-mode.test.ts @@ -123,8 +123,9 @@ const { reserveBudget, reconcileSpend } = vi.hoisted(() => ({ allowed: boolean; reserved: number; totalAfter?: number; + owner?: "operator" | "user"; }> - >(async () => ({ allowed: true, reserved: 3000 })), + >(async () => ({ allowed: true, reserved: 3000, owner: "user" as const })), reconcileSpend: vi.fn(async () => undefined), })); vi.mock("@/lib/ai/coach/budget", () => ({ @@ -132,6 +133,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget, reconcileSpend, resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); const { detectRefusal } = vi.hoisted(() => ({ @@ -266,7 +268,11 @@ function chatReq(body: Record): Request { describe("coach chat — tool-mode routing (F1)", () => { beforeEach(() => { vi.clearAllMocks(); - reserveBudget.mockResolvedValue({ allowed: true, reserved: 3000 }); + reserveBudget.mockResolvedValue({ + allowed: true, + reserved: 3000, + owner: "user" as const, + }); detectRefusal.mockReturnValue({ refuse: false }); checkRateLimit.mockResolvedValue({ allowed: true }); assertConsentForChain.mockResolvedValue(undefined); @@ -318,6 +324,7 @@ describe("coach chat — tool-mode routing (F1)", () => { 4500, "2026-06-21", 2_000_000, + "user", ); }); @@ -347,6 +354,8 @@ describe("coach chat — tool-mode routing (F1)", () => { 3000, 0, expect.anything(), + 0, + { servedBy: null, reservedOwner: "user" }, ); }); @@ -558,6 +567,7 @@ describe("coach chat — tool-mode routing (F1)", () => { 1500, "2026-06-21", 2_000_000, + "user", ); // The legacy path ships the snapshot figures in the user turn. const params = ( diff --git a/src/app/api/insights/chat/__tests__/route-workout-scope.test.ts b/src/app/api/insights/chat/__tests__/route-workout-scope.test.ts index 9604d200a..16d9b72b2 100644 --- a/src/app/api/insights/chat/__tests__/route-workout-scope.test.ts +++ b/src/app/api/insights/chat/__tests__/route-workout-scope.test.ts @@ -177,6 +177,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 1500 })), reconcileSpend: vi.fn(async () => undefined), resolveDailyCap: vi.fn(() => 2_000_000), + resolveCostOwner: vi.fn(() => "user" as const), })); vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn(() => ({ refuse: false })), diff --git a/src/app/api/insights/chat/route.ts b/src/app/api/insights/chat/route.ts index 9dc19d79d..5f6c036ce 100644 --- a/src/app/api/insights/chat/route.ts +++ b/src/app/api/insights/chat/route.ts @@ -52,6 +52,7 @@ import { runStreamingRawCompletionWithFallback, } from "@/lib/ai/provider-runner"; import { resolveProviderChain, resolveProvider } from "@/lib/ai/provider"; +import type { ProviderChainType } from "@/lib/ai/provider-chain"; import { assertConsentForChain } from "@/lib/ai/consent-guard"; import { singleUserTurn, type CompletionResult } from "@/lib/ai/types"; import { PROMPT_VERSION } from "@/lib/ai/prompts/insight-generator"; @@ -74,6 +75,7 @@ import { buildDateKey, reserveBudget, reconcileSpend, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { detectRefusal } from "@/lib/ai/coach/refusal"; @@ -603,6 +605,7 @@ async function handleChatRequest(request: NextRequest): Promise { : AI_BUDGETS.coach.maxTokens, reqDateKey, dailyCap, + resolveCostOwner(chain), ); if (!reservation.allowed) { annotate({ @@ -637,7 +640,10 @@ async function handleChatRequest(request: NextRequest): Promise { async function produceReply(): Promise { let result: CompletionResult; - let workingProviderType: string; + // v1.38.19 (Wave E) — typed as the chain's provider union, not a bare + // string: the budget reconcile attributes the turn's tokens to the cost + // owner of the hop named here. + let workingProviderType: ProviderChainType; let toolTrace: CoachToolTrace[] = []; // v1.21.0 (P6) — the present tool-result payloads this turn, for the post-hoc // prose number-verifier. Empty on the no-tools path. @@ -702,7 +708,9 @@ async function handleChatRequest(request: NextRequest): Promise { timeoutMs: aiResponseTimeoutMs, }); result = loop.result; - workingProviderType = loop.workingProviderType; + // The loop reports the hop it landed on as a bare string; it is + // assigned from `workingProvider.providerType` one frame up. + workingProviderType = loop.workingProviderType as ProviderChainType; toolTrace = loop.toolTrace; // v1.32.1 — the numeric verifier ACTIVATES only when this turn actually // delivered figures the model was told to ground against: a pinned @@ -784,9 +792,10 @@ async function handleChatRequest(request: NextRequest): Promise { } catch (err) { // The provider chain failed outright — no tokens were billed, so refund // the full reservation before surfacing the error frame. - await reconcileSpend(userId, reservation.reserved, 0, reqDateKey).catch( - () => {}, - ); + await reconcileSpend(userId, reservation.reserved, 0, reqDateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }).catch(() => {}); // #781 — the client walked away mid-generation. The request's abort // signal is threaded into every provider call, so the teardown surfaces // here as an abort-shaped failure with `request.signal` already flipped. @@ -887,6 +896,7 @@ async function handleChatRequest(request: NextRequest): Promise { totalTokensSpent, reqDateKey, cachedTokensSpent, + { servedBy: workingProviderType, reservedOwner: reservation.owner }, ).catch(() => { // Ledger reconcile is best-effort; a failure leaves the conservative // reservation in place (never an undercount) and never breaks the turn. diff --git a/src/app/api/labs/ocr/extract/__tests__/route.test.ts b/src/app/api/labs/ocr/extract/__tests__/route.test.ts index 936f1901f..c926aecc0 100644 --- a/src/app/api/labs/ocr/extract/__tests__/route.test.ts +++ b/src/app/api/labs/ocr/extract/__tests__/route.test.ts @@ -45,6 +45,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 200_000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/labs/ocr-extract", async () => { const actual = await vi.importActual( @@ -95,6 +96,8 @@ beforeEach(() => { allowed: true, reserved: AI_BUDGETS.ocrExtractText.maxTokens ?? 0, totalAfter: AI_BUDGETS.ocrExtractText.maxTokens ?? 0, + owner: "operator", + operatorAfter: AI_BUDGETS.ocrExtractText.maxTokens ?? 0, } as never); }); @@ -116,6 +119,8 @@ describe("POST /api/labs/ocr/extract — text mode budget", () => { "2026-06-26", // F1 — the provider-aware daily cap (mocked) is threaded as the 4th arg. 200_000, + // Wave E — and the cost owner that cap is enforced against as the 5th. + "operator", ); }); @@ -134,6 +139,8 @@ describe("POST /api/labs/ocr/extract — text mode budget", () => { AI_BUDGETS.ocrExtractText.maxTokens, 0, "2026-06-26", + 0, + { servedBy: null, reservedOwner: "operator" }, ); }); @@ -182,6 +189,8 @@ describe("POST /api/labs/ocr/extract — vision PDF rasterization", () => { allowed: true, reserved: AI_BUDGETS.ocrExtract.maxTokens ?? 0, totalAfter: AI_BUDGETS.ocrExtract.maxTokens ?? 0, + owner: "operator", + operatorAfter: AI_BUDGETS.ocrExtract.maxTokens ?? 0, } as never); }); @@ -222,6 +231,8 @@ describe("POST /api/labs/ocr/extract — vision PDF rasterization", () => { AI_BUDGETS.ocrExtract.maxTokens, 0, "2026-06-26", + 0, + { servedBy: null, reservedOwner: "operator" }, ); }); }); diff --git a/src/app/api/labs/ocr/extract/route.ts b/src/app/api/labs/ocr/extract/route.ts index 125af81dd..7dfd230f6 100644 --- a/src/app/api/labs/ocr/extract/route.ts +++ b/src/app/api/labs/ocr/extract/route.ts @@ -41,6 +41,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { prisma } from "@/lib/db"; @@ -183,6 +184,7 @@ async function handleTextExtract( AI_BUDGETS.ocrExtractText.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { annotate({ @@ -208,12 +210,17 @@ async function handleTextExtract( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); return apiSuccess(result); } catch (err) { // A failed structuring call produced no usable rows; mirror the vision // path and refund the reservation in full rather than charging it. - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (err instanceof OcrExtractError) { return apiError("Couldn't read the report. Try a clearer photo.", 422, { errorCode: "labs.ocr.extractFailed", @@ -281,6 +288,7 @@ async function handleVisionExtract( AI_BUDGETS.ocrExtract.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { annotate({ @@ -303,7 +311,10 @@ async function handleVisionExtract( action: { name: "labs.ocr.fileRejected" }, meta: { reason: "content_length_exceeded" }, }); - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return apiError("File is too large (max 12 MB).", 413, { errorCode: "labs.ocr.fileTooLarge", }); @@ -316,7 +327,10 @@ async function handleVisionExtract( headers: { "content-type": request.headers.get("content-type") ?? "" }, }).formData(); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (err instanceof BodyTooLargeError) { annotate({ action: { name: "labs.ocr.fileRejected" }, @@ -331,7 +345,10 @@ async function handleVisionExtract( const file = formData.get("file"); if (!(file instanceof File)) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return apiError("Field 'file' must be a file", 422); } @@ -339,7 +356,10 @@ async function handleVisionExtract( try { buffer = Buffer.from(await file.arrayBuffer()); } catch { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return apiError("Failed to read uploaded file", 400); } @@ -350,7 +370,10 @@ async function handleVisionExtract( action: { name: "labs.ocr.fileRejected" }, meta: { reason: "unsupported_mime" }, }); - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return apiError("Upload a JPEG, PNG, WebP, or PDF.", 415, { errorCode: "labs.ocr.fileType", }); @@ -388,7 +411,10 @@ async function handleVisionExtract( action: { name: "labs.ocr.fileRejected" }, meta: { reason: "pdf_rasterize_failed" }, }); - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return apiError( "Couldn't read this PDF; upload a photo instead.", 422, @@ -416,12 +442,26 @@ async function handleVisionExtract( // The orchestration does not surface token counts; reconcile against the // reserved estimate as the spend ceiling (the provider already billed it). actualTokens = reservation.reserved; - await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey); + await reconcileSpend( + userId, + reservation.reserved, + actualTokens, + dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, + ); return apiSuccess(result); } catch (err) { // Provider/extraction failure — the call may still have burned tokens, so // reconcile against the reserved estimate rather than refunding in full. - await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey); + await reconcileSpend( + userId, + reservation.reserved, + actualTokens, + dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, + ); if (err instanceof OcrExtractError) { return apiError("Couldn't read the report. Try a clearer photo.", 422, { errorCode: "labs.ocr.extractFailed", @@ -439,9 +479,10 @@ async function handleVisionExtract( } } catch (err) { // A guard threw after the reservation (e.g. consent races) — refund fully. - await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch( - () => {}, - ); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }).catch(() => {}); throw err; } } diff --git a/src/app/api/medications/extract/__tests__/route.test.ts b/src/app/api/medications/extract/__tests__/route.test.ts index c01d578ee..3c43ec73b 100644 --- a/src/app/api/medications/extract/__tests__/route.test.ts +++ b/src/app/api/medications/extract/__tests__/route.test.ts @@ -44,6 +44,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ // returns a distinctive figure so an accidental default-cap regression is // visible in the assertion rather than silently passing. resolveDailyCap: vi.fn(() => 1_234_567), + resolveCostOwner: vi.fn(() => "user" as const), })); vi.mock("@/lib/ai/provider", () => ({ @@ -118,6 +119,8 @@ beforeEach(() => { allowed: true, reserved: 700, totalAfter: 700, + owner: "user", + operatorAfter: 0, }); vi.mocked(resolveDailyCap).mockReturnValue(1_234_567); vi.mocked(resolveProviderChain).mockResolvedValue([ @@ -274,6 +277,7 @@ describe("POST /api/medications/extract — daily budget", () => { expect.any(Number), "2026-05-28", 1_234_567, + "user", ); }); @@ -283,6 +287,8 @@ describe("POST /api/medications/extract — daily budget", () => { allowed: false, reserved: 700, totalAfter: 9_999_999, + owner: "user", + operatorAfter: 0, }); const res = await POST(postReq({ text: "Mounjaro 5mg weekly" }) as never); @@ -307,6 +313,8 @@ describe("POST /api/medications/extract — daily budget", () => { 700, 120, "2026-05-28", + 0, + { servedBy: "openai", reservedOwner: "user" }, ); }); @@ -323,6 +331,8 @@ describe("POST /api/medications/extract — daily budget", () => { 700, 700, "2026-05-28", + 0, + { servedBy: "openai", reservedOwner: "user" }, ); }); @@ -335,7 +345,14 @@ describe("POST /api/medications/extract — daily budget", () => { const res = await POST(postReq({ text: "5mg weekly" }) as never); expect(res.status).toBe(503); - expect(reconcileSpend).toHaveBeenCalledWith("user-1", 700, 0, "2026-05-28"); + expect(reconcileSpend).toHaveBeenCalledWith( + "user-1", + 700, + 0, + "2026-05-28", + 0, + { servedBy: null, reservedOwner: "user" }, + ); }); }); diff --git a/src/app/api/medications/extract/route.ts b/src/app/api/medications/extract/route.ts index 0bc66762e..883de5d55 100644 --- a/src/app/api/medications/extract/route.ts +++ b/src/app/api/medications/extract/route.ts @@ -70,6 +70,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { @@ -210,6 +211,7 @@ async function handleExtract(request: NextRequest): Promise { estimatedTokens, dateKey, resolveDailyCap(chain), + resolveCostOwner(chain), ); if (!reservation.allowed) { annotate({ @@ -241,9 +243,10 @@ async function handleExtract(request: NextRequest): Promise { } catch (err) { // Nothing was generated — refund the whole reservation rather than bill an // invented figure. Best-effort: a failed refund must not mask the 503. - await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch( - () => {}, - ); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }).catch(() => {}); if (err instanceof AllProvidersFailedError) { annotate({ action: { name: "medications.extract.provider-failed" }, @@ -260,7 +263,10 @@ async function handleExtract(request: NextRequest): Promise { // a free retry loop. Falls back to the reservation when the provider reports // no count, so an unreported generation is never billed as zero. const actualTokens = completion.result.tokensUsed ?? reservation.reserved; - await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey); + await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey, 0, { + servedBy: completion.workingProvider.providerType, + reservedOwner: reservation.owner, + }); const reply = completion.result.content?.trim() ?? ""; if (!reply) { diff --git a/src/lib/ai/__tests__/provider-runner.test.ts b/src/lib/ai/__tests__/provider-runner.test.ts index 51affa883..9514243db 100644 --- a/src/lib/ai/__tests__/provider-runner.test.ts +++ b/src/lib/ai/__tests__/provider-runner.test.ts @@ -13,13 +13,16 @@ import { // v1.37.19 (A7-2) — the runner consults the day's recorded spend before an // operator-funded hop. These pure chain tests do not stand up a DB, so the // ledger read is swapped for a controllable in-memory figure. -const budgetState = vi.hoisted(() => ({ spent: 0 })); +const budgetState = vi.hoisted(() => ({ spent: 0, operatorSpent: 0 })); vi.mock("../coach/budget", async () => { const actual = await vi.importActual("../coach/budget"); return { ...actual, - readDailySpend: vi.fn(async () => budgetState.spent), + readDailySpend: vi.fn(async () => ({ + total: budgetState.spent, + operator: budgetState.operatorSpent, + })), }; }); diff --git a/src/lib/ai/coach/__tests__/budget.test.ts b/src/lib/ai/coach/__tests__/budget.test.ts index 45697c6b8..cb6fbe304 100644 --- a/src/lib/ai/coach/__tests__/budget.test.ts +++ b/src/lib/ai/coach/__tests__/budget.test.ts @@ -4,6 +4,7 @@ import { buildDateKey, OPERATOR_COST_CAP, USER_PLAN_CAP, + resolveCostOwner, resolveDailyCap, } from "../budget"; @@ -60,19 +61,37 @@ describe("reserveBudget — ledger clamps", () => { it("clamps a non-finite reservation to 0", async () => { const { reserveBudget } = await import("../budget"); - const res = await reserveBudget("u", Number.NaN, "2026-05-10"); + const res = await reserveBudget( + "u", + Number.NaN, + "2026-05-10", + OPERATOR_COST_CAP, + "operator", + ); expect(res.reserved).toBe(0); }); it("clamps a negative reservation to 0", async () => { const { reserveBudget } = await import("../budget"); - const res = await reserveBudget("u", -42, "2026-05-10"); + const res = await reserveBudget( + "u", + -42, + "2026-05-10", + OPERATOR_COST_CAP, + "operator", + ); expect(res.reserved).toBe(0); }); it("floors a fractional reservation", async () => { const { reserveBudget } = await import("../budget"); - const res = await reserveBudget("u", 12.7, "2026-05-10"); + const res = await reserveBudget( + "u", + 12.7, + "2026-05-10", + OPERATOR_COST_CAP, + "operator", + ); expect(res.reserved).toBe(12); }); }); @@ -143,6 +162,7 @@ describe("reserveBudget cap (F1 — user-plan path not locked out)", () => { 1_200, "2026-05-10", resolveDailyCap([{ providerType: "codex" }]), + resolveCostOwner([{ providerType: "codex" }]), ); expect(res.allowed).toBe(true); // The reservation upsert ran; no refund executeRaw fired. @@ -161,6 +181,7 @@ describe("reserveBudget cap (F1 — user-plan path not locked out)", () => { 1_200, "2026-05-10", resolveDailyCap([{ providerType: "admin-openai" }]), + resolveCostOwner([{ providerType: "admin-openai" }]), ); expect(res.allowed).toBe(false); // Refund of the reservation fired on refusal. @@ -181,7 +202,10 @@ describe("reconcileSpend cached-token subtraction (F3)", () => { it("bills total_tokens minus cached input as the signed delta", async () => { const { reconcileSpend } = await import("../budget"); // reserved 1200, gross 20000, cached 13000 → net actual 7000 → delta 5800. - await reconcileSpend("u", 1_200, 20_000, "2026-05-10", 13_000); + await reconcileSpend("u", 1_200, 20_000, "2026-05-10", 13_000, { + servedBy: "admin-openai", + reservedOwner: "operator", + }); expect(prismaMock.$executeRaw).toHaveBeenCalledTimes(1); // The tagged-template interpolations carry the delta; assert it is 5800. const interpolations = prismaMock.$executeRaw.mock.calls[0].slice(1); @@ -190,7 +214,10 @@ describe("reconcileSpend cached-token subtraction (F3)", () => { it("clamps a cached count larger than gross to a zero charge (delta = -reserved)", async () => { const { reconcileSpend } = await import("../budget"); - await reconcileSpend("u", 1_200, 5_000, "2026-05-10", 9_999); + await reconcileSpend("u", 1_200, 5_000, "2026-05-10", 9_999, { + servedBy: "admin-openai", + reservedOwner: "operator", + }); const interpolations = prismaMock.$executeRaw.mock.calls[0].slice(1); // net actual clamped to 0 → delta = 0 - 1200 = -1200. expect(interpolations).toContain(-1_200); @@ -198,7 +225,10 @@ describe("reconcileSpend cached-token subtraction (F3)", () => { it("defaults cachedTokens to 0 (back-compat) — bills gross", async () => { const { reconcileSpend } = await import("../budget"); - await reconcileSpend("u", 1_000, 4_000, "2026-05-10"); + await reconcileSpend("u", 1_000, 4_000, "2026-05-10", 0, { + servedBy: "admin-openai", + reservedOwner: "operator", + }); const interpolations = prismaMock.$executeRaw.mock.calls[0].slice(1); expect(interpolations).toContain(3_000); }); diff --git a/src/lib/ai/coach/__tests__/self-context-questions-consent.test.ts b/src/lib/ai/coach/__tests__/self-context-questions-consent.test.ts index 7319c0d24..9276adb5a 100644 --- a/src/lib/ai/coach/__tests__/self-context-questions-consent.test.ts +++ b/src/lib/ai/coach/__tests__/self-context-questions-consent.test.ts @@ -41,9 +41,12 @@ const budgetMocks = vi.hoisted(() => ({ allowed: true, reserved: 300, totalAfter: 300, + owner: "operator" as const, + operatorAfter: 300, })), reconcileSpend: vi.fn(async () => {}), resolveDailyCap: vi.fn(() => 200_000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/coach/budget", () => budgetMocks); @@ -87,6 +90,8 @@ beforeEach(() => { allowed: true, reserved: 300, totalAfter: 300, + owner: "operator" as const, + operatorAfter: 300, }); buildCoachSnapshot.mockResolvedValue({ snapshotJson: '{"weight":[]}' }); }); @@ -239,6 +244,7 @@ describe("deriveClarifyingQuestions — atomic budget", () => { 300, "2026-07-18", 200_000, + "operator", ); expect(budgetMocks.reconcileSpend).toHaveBeenCalledWith( "user-1", @@ -246,6 +252,7 @@ describe("deriveClarifyingQuestions — atomic budget", () => { 120, "2026-07-18", 0, + { servedBy: "openai", reservedOwner: "operator" }, ); }); @@ -254,6 +261,8 @@ describe("deriveClarifyingQuestions — atomic budget", () => { allowed: false, reserved: 300, totalAfter: 200_000, + owner: "operator" as const, + operatorAfter: 200_000, }); const out = await deriveClarifyingQuestions("user-1", ctx, "en"); @@ -273,6 +282,8 @@ describe("deriveClarifyingQuestions — atomic budget", () => { 300, 0, "2026-07-18", + 0, + { servedBy: null, reservedOwner: "operator" }, ); }); }); diff --git a/src/lib/ai/coach/budget.ts b/src/lib/ai/coach/budget.ts index d559a8e82..44f6eab61 100644 --- a/src/lib/ai/coach/budget.ts +++ b/src/lib/ai/coach/budget.ts @@ -63,8 +63,7 @@ export const USER_PLAN_CAP = 2_000_000; export function resolveDailyCap( chain: ReadonlyArray<{ providerType: ProviderChainType }>, ): number { - const primary = chain[0]?.providerType; - return primary === undefined || isOperatorFundedProvider(primary) + return resolveCostOwner(chain) === "operator" ? OPERATOR_COST_CAP : USER_PLAN_CAP; } @@ -80,6 +79,28 @@ export function isOperatorFundedProvider(type: ProviderChainType): boolean { return type === "admin-openai" || type === "admin-codex"; } +/** + * v1.38.19 (Wave E) — who pays for a turn. `"operator"` means the tokens land + * on the operator's invoice (`admin-openai` / `admin-codex`), `"user"` means + * the user's own plan, key or hardware carries them. + */ +export type BudgetCostOwner = "operator" | "user"; + +/** + * v1.38.19 (Wave E) — the cost owner a chain RESERVES under: the primary is + * the provider that will be tried first, so it is the expected payer. An empty + * chain defaults to the operator — the conservative side. `resolveDailyCap` + * and `reserveBudget` must agree on this classification, so both read it here. + */ +export function resolveCostOwner( + chain: ReadonlyArray<{ providerType: ProviderChainType }>, +): BudgetCostOwner { + const primary = chain[0]?.providerType; + return primary === undefined || isOperatorFundedProvider(primary) + ? "operator" + : "user"; +} + /** * v1.37.19 (A7-2) — the day's recorded spend for one user (0 when no row). * @@ -90,15 +111,30 @@ export function isOperatorFundedProvider(type: ProviderChainType): boolean { * hop must not spend past `OPERATOR_COST_CAP` just because the reservation * was checked against the wrong owner's ceiling. */ +export interface DailySpend { + /** Every token recorded for the day, whoever paid for it. */ + total: number; + /** + * v1.38.19 (Wave E) — the share of `total` served by an operator-funded + * provider. The hop guard compares THIS against `OPERATOR_COST_CAP`: a day + * full of turns the user's own plan paid for must not close the operator's + * fallback hop. + */ + operator: number; +} + export async function readDailySpend( userId: string, dateKey: string = buildDateKey(), -): Promise { +): Promise { const row = await prisma.coachUsage.findUnique({ where: { userId_dateKey: { userId, dateKey } }, - select: { totalTokens: true }, + select: { totalTokens: true, operatorTokens: true }, }); - return row?.totalTokens ?? 0; + return { + total: row?.totalTokens ?? 0, + operator: row?.operatorTokens ?? 0, + }; } /** @@ -132,53 +168,79 @@ export function buildDateKey(at: Date = new Date()): string { * reserved-but-failed turn still counts as an attempt. */ export interface ReserveBudgetResult { - /** True when the reservation kept the day's total within the cap. */ + /** True when the reservation kept the day's spend within the cap. */ allowed: boolean; /** Tokens reserved by this call (refunded/reconciled by the caller). */ reserved: number; /** The day's total AFTER this reservation (for observability). */ totalAfter: number; + /** + * v1.38.19 (Wave E) — the owner this reservation was charged to. The + * reconcile needs it: only a reservation booked to the operator has an + * amount to move back OUT of `operator_tokens` when a user-funded hop ends + * up serving the turn. + */ + owner: BudgetCostOwner; + /** The day's operator-funded spend AFTER this reservation. */ + operatorAfter: number; } export async function reserveBudget( userId: string, estimatedTokens: number, - dateKey: string = buildDateKey(), - cap: number = OPERATOR_COST_CAP, + dateKey: string, + cap: number, + owner: BudgetCostOwner, ): Promise { const reserved = Number.isFinite(estimatedTokens) && estimatedTokens > 0 ? Math.floor(estimatedTokens) : 0; + // v1.38.19 (Wave E) — a reservation is booked to `operator_tokens` only when + // the chain's primary is operator-funded. The counter moves inside the SAME + // statement as the total, so two concurrent requests can never observe one + // counter without the other. + const operatorReserved = owner === "operator" ? reserved : 0; - // Single atomic upsert-increment returning the new total. Two concurrent + // Single atomic upsert-increment returning the new totals. Two concurrent // requests serialise on the row's unique (user_id, date_key) constraint, so // each observes a distinct post-increment total — they cannot both read a // sub-cap value and both proceed. - const rows = await prisma.$queryRaw<{ total_tokens: number }[]>` - INSERT INTO coach_usage (id, user_id, date_key, total_tokens, message_count, created_at, updated_at) - VALUES (gen_random_uuid()::text, ${userId}, ${dateKey}, ${reserved}, 1, NOW(), NOW()) + const rows = await prisma.$queryRaw< + { total_tokens: number; operator_tokens: number }[] + >` + INSERT INTO coach_usage (id, user_id, date_key, total_tokens, operator_tokens, message_count, created_at, updated_at) + VALUES (gen_random_uuid()::text, ${userId}, ${dateKey}, ${reserved}, ${operatorReserved}, 1, NOW(), NOW()) ON CONFLICT (user_id, date_key) DO UPDATE SET total_tokens = coach_usage.total_tokens + ${reserved}, + operator_tokens = coach_usage.operator_tokens + ${operatorReserved}, message_count = coach_usage.message_count + 1, updated_at = NOW() - RETURNING total_tokens + RETURNING total_tokens, operator_tokens `; const totalAfter = Number(rows[0]?.total_tokens ?? reserved); + const operatorAfter = Number(rows[0]?.operator_tokens ?? operatorReserved); // The cap is a ceiling on tokens already spent BEFORE this request, matching // the prior `spent >= cap` semantics: a request is allowed when the spend // PRIOR to its reservation was under the cap. So compare `totalAfter - // reserved` (the prior total) against the cap. const priorTotal = totalAfter - reserved; + const priorOperator = operatorAfter - operatorReserved; if (priorTotal >= cap) { // Already over before this request — refund the reservation + the // message-count bump and refuse. - await refundReservation(userId, reserved, dateKey); - return { allowed: false, reserved, totalAfter: priorTotal }; + await refundReservation(userId, reserved, operatorReserved, dateKey); + return { + allowed: false, + reserved, + totalAfter: priorTotal, + owner, + operatorAfter: priorOperator, + }; } - return { allowed: true, reserved, totalAfter }; + return { allowed: true, reserved, totalAfter, owner, operatorAfter }; } /** @@ -194,12 +256,24 @@ export async function reserveBudget( * served most of it cheaply / free; charging the user's daily meter for input * they did not re-pay for is an over-charge. We bill `actual - cached`. */ +export interface ReconcileSpendOptions { + /** + * The provider that actually SERVED the turn (`runRawWithFallback`'s + * `workingProvider.providerType`), or `null` when no hop served — a failed, + * timed-out or cancelled call whose reservation is being refunded. + */ + servedBy: ProviderChainType | null; + /** The owner the reservation was booked to (`reserveBudget(...).owner`). */ + reservedOwner: BudgetCostOwner; +} + export async function reconcileSpend( userId: string, reserved: number, actualTokens: number, - dateKey: string = buildDateKey(), - cachedTokens = 0, + dateKey: string, + cachedTokens: number, + opts: ReconcileSpendOptions, ): Promise { const grossActual = Number.isFinite(actualTokens) && actualTokens > 0 @@ -213,12 +287,31 @@ export async function reconcileSpend( // (shouldn't happen, but the wire is untrusted) can't drive a negative charge. const actual = Math.max(0, grossActual - cached); const delta = actual - reserved; - if (delta === 0) return; + // v1.38.19 (Wave E) — settle the operator's counter against the hop that + // SERVED, not the one the chain expected. Book the actual tokens when an + // operator-funded provider answered, and take back whatever this + // reservation had put there (nothing, when the reservation was the user's + // own plan — subtracting then would eat another request's legitimate + // operator balance). The four cases settle exactly: + // operator reservation + operator hop → `actual` + // operator reservation + user hop → 0 (the reservation moves out) + // user reservation + user hop → untouched + // user reservation + operator hop → `actual` (a fallback the + // operator really paid for) + const servedByOperator = + opts.servedBy !== null && isOperatorFundedProvider(opts.servedBy); + const operatorDelta = + (servedByOperator ? actual : 0) - + (opts.reservedOwner === "operator" ? reserved : 0); + if (delta === 0 && operatorDelta === 0) return; // Clamp at zero so a smaller-than-reserved actual can't drive the row - // negative under a racing reconcile. + // negative under a racing reconcile. Both counters move in ONE statement: + // a concurrent reader can never see the total settled and the owner split + // still stale. await prisma.$executeRaw` UPDATE coach_usage SET total_tokens = GREATEST(0, total_tokens + ${delta}), + operator_tokens = GREATEST(0, operator_tokens + ${operatorDelta}), updated_at = NOW() WHERE user_id = ${userId} AND date_key = ${dateKey} `; @@ -228,11 +321,13 @@ export async function reconcileSpend( async function refundReservation( userId: string, reserved: number, + operatorReserved: number, dateKey: string, ): Promise { await prisma.$executeRaw` UPDATE coach_usage SET total_tokens = GREATEST(0, total_tokens - ${reserved}), + operator_tokens = GREATEST(0, operator_tokens - ${operatorReserved}), message_count = GREATEST(0, message_count - 1), updated_at = NOW() WHERE user_id = ${userId} AND date_key = ${dateKey} diff --git a/src/lib/ai/coach/self-context-questions.ts b/src/lib/ai/coach/self-context-questions.ts index 0d0a41a0e..e9b3022ea 100644 --- a/src/lib/ai/coach/self-context-questions.ts +++ b/src/lib/ai/coach/self-context-questions.ts @@ -45,6 +45,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { @@ -235,6 +236,7 @@ export async function deriveClarifyingQuestions( QUESTIONS_MAX_TOKENS, dateKey, resolveDailyCap(chain), + resolveCostOwner(chain), ); if (!reservation.allowed) return fallback(); @@ -266,9 +268,10 @@ export async function deriveClarifyingQuestions( } catch (err) { // Reconcile the reservation down to zero actual spend before the // failure lands on the deterministic fallback. - await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch( - () => {}, - ); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }).catch(() => {}); throw err; } @@ -278,6 +281,7 @@ export async function deriveClarifyingQuestions( result.tokensUsed ?? 0, dateKey, result.cachedInputTokens ?? 0, + { servedBy: chain[0].providerType, reservedOwner: reservation.owner }, ).catch(() => {}); const questions = parseQuestionsReply(result.content); diff --git a/src/lib/ai/provider-runner.ts b/src/lib/ai/provider-runner.ts index e2bd1cef5..9b9344f73 100644 --- a/src/lib/ai/provider-runner.ts +++ b/src/lib/ai/provider-runner.ts @@ -431,8 +431,11 @@ async function runRawChain( // Fail OPEN on a ledger read error: the reservation layer still // meters every call against the same ledger, so a transient read // failure here must not take the whole fallback chain down with it. - const spent = await readDailySpend(userId).catch(() => 0); - if (spent >= OPERATOR_COST_CAP) { + const spent = await readDailySpend(userId).catch(() => ({ + total: 0, + operator: 0, + })); + if (spent.total >= OPERATOR_COST_CAP) { const hop: FallbackHop = { providerType: candidate.providerType, attempt: i + 1, @@ -445,7 +448,7 @@ async function runRawChain( meta: { [`ai_chain_hop_${i + 1}_provider`]: candidate.providerType, [`ai_chain_hop_${i + 1}_reason`]: "operator-cost-cap-exhausted", - operator_cap_spent: spent, + operator_cap_spent: spent.total, }, }); continue; diff --git a/src/lib/documents/__tests__/fenced-chat-pipeline.test.ts b/src/lib/documents/__tests__/fenced-chat-pipeline.test.ts index c9f6caf9d..f2d3e7abb 100644 --- a/src/lib/documents/__tests__/fenced-chat-pipeline.test.ts +++ b/src/lib/documents/__tests__/fenced-chat-pipeline.test.ts @@ -29,6 +29,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 600 }), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 100000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/provider-runner", () => ({ runStreamingRawCompletionWithFallback: vi.fn(), diff --git a/src/lib/documents/__tests__/index-document.test.ts b/src/lib/documents/__tests__/index-document.test.ts index d9a145bc6..124e72c5e 100644 --- a/src/lib/documents/__tests__/index-document.test.ts +++ b/src/lib/documents/__tests__/index-document.test.ts @@ -45,6 +45,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/ai-budgets", () => ({ AI_BUDGETS: { documentTranscribe: { temperature: 0, maxTokens: 4000 } }, @@ -106,6 +107,7 @@ beforeEach(() => { vi.mocked(reserveBudget).mockResolvedValue({ allowed: true, reserved: 1, + owner: "operator", } as never); vi.mocked(localExtractText).mockResolvedValue({ ok: true, @@ -132,7 +134,14 @@ describe("indexDocumentContent — provider-first path", () => { expect.objectContaining({ source: "vision", providerType: "anthropic" }), ); // Budget reserved then reconciled at the full reserved amount (charged). - expect(reconcileSpend).toHaveBeenCalledWith("user-1", 1, 1, "2026-07-07"); + expect(reconcileSpend).toHaveBeenCalledWith( + "user-1", + 1, + 1, + "2026-07-07", + 0, + { servedBy: expect.any(String), reservedOwner: "operator" }, + ); // The local path is never touched when the provider succeeds. expect(localExtractText).not.toHaveBeenCalled(); }); @@ -222,7 +231,14 @@ describe("indexDocumentContent — local fallback path", () => { const outcome = await indexDocumentContent("user-1", "doc-1"); expect(outcome).toMatchObject({ indexed: true, source: "local-pdf" }); // Reservation refunded to zero spend. - expect(reconcileSpend).toHaveBeenCalledWith("user-1", 1, 0, "2026-07-07"); + expect(reconcileSpend).toHaveBeenCalledWith( + "user-1", + 1, + 0, + "2026-07-07", + 0, + { servedBy: null, reservedOwner: "operator" }, + ); expect(localExtractText).toHaveBeenCalledTimes(1); }); diff --git a/src/lib/documents/__tests__/index-outcome-record.test.ts b/src/lib/documents/__tests__/index-outcome-record.test.ts index c11d98e2b..955005b1d 100644 --- a/src/lib/documents/__tests__/index-outcome-record.test.ts +++ b/src/lib/documents/__tests__/index-outcome-record.test.ts @@ -56,6 +56,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/ai-budgets", () => ({ AI_BUDGETS: { documentTranscribe: { temperature: 0, maxTokens: 4000 } }, diff --git a/src/lib/documents/auto-stage-labs.ts b/src/lib/documents/auto-stage-labs.ts index 52ccaca2c..2182991b2 100644 --- a/src/lib/documents/auto-stage-labs.ts +++ b/src/lib/documents/auto-stage-labs.ts @@ -283,6 +283,7 @@ export async function maybeAutoStageLabFacts( AI_BUDGETS.ocrExtractText.maxTokens, dateKey, provider.dailyCap, + provider.costOwner, ); if (!reservation.allowed) return { staged: false, reason: "budget" }; @@ -298,9 +299,17 @@ export async function maybeAutoStageLabFacts( reservation.reserved, reservation.reserved, dateKey, + 0, + { + servedBy: provider.pick.entry.providerType, + reservedOwner: reservation.owner, + }, ); } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); if (!(err instanceof InboundExtractError)) { annotate({ action: { name: "documents.autoStage.failed" }, diff --git a/src/lib/documents/fenced-chat.ts b/src/lib/documents/fenced-chat.ts index a09dd161d..9e70b7863 100644 --- a/src/lib/documents/fenced-chat.ts +++ b/src/lib/documents/fenced-chat.ts @@ -43,6 +43,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { detectRefusal, type CoachRefusalReason } from "@/lib/ai/coach/refusal"; @@ -307,6 +308,7 @@ export async function streamFencedReply( AI_BUDGETS.documentChat.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { annotate({ action: { name: "documents.chat.budget.exceeded" } }); @@ -383,9 +385,10 @@ Reply now as the assistant, grounded ONLY in the documents above, in ${ }); result = fallback.result; } catch (err) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch( - () => {}, - ); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }).catch(() => {}); if (err instanceof AllProvidersFailedError) { const allRateLimited = err.attempts.length > 0 && @@ -425,6 +428,7 @@ Reply now as the assistant, grounded ONLY in the documents above, in ${ totalTokens, dateKey, cachedTokens, + { servedBy: pick!.entry.providerType, reservedOwner: reservation.owner }, ).catch(() => {}); let replyText = (result.content ?? "").trim(); diff --git a/src/lib/documents/index-document.ts b/src/lib/documents/index-document.ts index 16f012f6a..1c222c7c8 100644 --- a/src/lib/documents/index-document.ts +++ b/src/lib/documents/index-document.ts @@ -36,7 +36,9 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, + type BudgetCostOwner, } from "@/lib/ai/coach/budget"; import { loadOwnedDocument, @@ -128,6 +130,12 @@ export interface ResolvedIndexProvider { pick: Awaited>["pick"]; consentOk: boolean; dailyCap: number; + /** + * v1.38.19 (Wave E) — whose budget `dailyCap` rations. The reservation books + * `operator_tokens` only for an operator-funded pick, so the cap and the + * counter it is compared against are resolved from one place. + */ + costOwner: BudgetCostOwner; } /** @@ -160,7 +168,13 @@ export async function resolveIndexProvider( pick && consentOk ? resolveDailyCap([{ providerType: pick.entry.providerType }]) : 0; - return { chain, pick, consentOk, dailyCap }; + // v1.38.19 (Wave E) — the cap and its cost owner are resolved from the SAME + // pick, so the reservation books the counter the cap is enforced against. + const costOwner: BudgetCostOwner = + pick && consentOk + ? resolveCostOwner([{ providerType: pick.entry.providerType }]) + : "operator"; + return { chain, pick, consentOk, dailyCap, costOwner }; } /** @@ -221,6 +235,7 @@ async function tryProviderIndex( AI_BUDGETS.documentTranscribe.maxTokens, dateKey, provider.dailyCap, + provider.costOwner, ); // Budget exhausted → fall through to the free local path rather than stall; // a text-layer PDF stays searchable even once the AI allowance is spent. @@ -238,6 +253,8 @@ async function tryProviderIndex( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); // Refs #776 — the empty-transcription guard: a provider answer with no // text must NEVER become a "successful" empty index (it would mark the @@ -261,7 +278,10 @@ async function tryProviderIndex( } catch { // Refund the reservation and let the caller fall through to local — a // transient provider miss must never leave a text-layer PDF unsearchable. - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return { outcome: null, note: "provider-error" }; } } diff --git a/src/lib/insights/briefing-provider.ts b/src/lib/insights/briefing-provider.ts index bb013840f..a98749ed1 100644 --- a/src/lib/insights/briefing-provider.ts +++ b/src/lib/insights/briefing-provider.ts @@ -39,6 +39,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { singleUserTurn } from "@/lib/ai/types"; @@ -122,6 +123,7 @@ export async function runBriefingCompletion( estimatedTokens, dateKey, resolveDailyCap(args.chain), + resolveCostOwner(args.chain), ); if (!reservation.allowed) { annotate({ @@ -153,12 +155,13 @@ export async function runBriefingCompletion( // invented figure. A partially-burned upstream call is possible here, but // we have no reported count to charge, and over-charging a failed // generation would ration the retry the user is entitled to. - await reconcileSpend(args.userId, reservation.reserved, 0, dateKey).catch( - () => { - // Best-effort: a failed refund leaves the conservative reservation in - // place (never an undercount) and must not mask the provider error. - }, - ); + await reconcileSpend(args.userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }).catch(() => { + // Best-effort: a failed refund leaves the conservative reservation in + // place (never an undercount) and must not mask the provider error. + }); throw err; } @@ -169,6 +172,10 @@ export async function runBriefingCompletion( actualTokens, dateKey, outcome.result.cachedInputTokens ?? 0, + { + servedBy: outcome.workingProvider.providerType, + reservedOwner: reservation.owner, + }, ).catch(() => { // Ledger reconcile is best-effort; a failure leaves the reservation in // place and never breaks a generation that already succeeded. diff --git a/src/lib/insights/status-provider.ts b/src/lib/insights/status-provider.ts index f30a7b36a..9c65481e3 100644 --- a/src/lib/insights/status-provider.ts +++ b/src/lib/insights/status-provider.ts @@ -13,6 +13,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { AI_BUDGETS, REFERENCE_AI_SEED } from "@/lib/ai/ai-budgets"; @@ -223,6 +224,7 @@ export async function runStatusCompletion( estimatedTokens, dateKey, resolveDailyCap(chain), + resolveCostOwner(chain), ); if (!reservation.allowed) { // Over the day's ceiling. Reported as `error` — a TRANSIENT miss the caller @@ -264,11 +266,17 @@ export async function runStatusCompletion( // A timed-out generation may still have burned upstream tokens, but we have // no reported count to charge — refund the reservation rather than bill an // invented figure. - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return { kind: "timeout" }; } if (raced.errored || raced.value === null) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); return { kind: "error" }; } @@ -279,7 +287,10 @@ export async function runStatusCompletion( // into a free retry loop. Falls back to the reservation when the provider // reports no count, so an unreported generation is never billed as zero. const actualTokens = result.tokensUsed ?? reservation.reserved; - await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey); + await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey, 0, { + servedBy: workingProvider.providerType, + reservedOwner: reservation.owner, + }); const content = result.content; if (typeof content !== "string" || content.trim().length === 0) { diff --git a/src/lib/jobs/__tests__/async-state-contract.test.ts b/src/lib/jobs/__tests__/async-state-contract.test.ts index 4aa5fdf3f..0392ca979 100644 --- a/src/lib/jobs/__tests__/async-state-contract.test.ts +++ b/src/lib/jobs/__tests__/async-state-contract.test.ts @@ -76,6 +76,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/ai-budgets", () => ({ AI_BUDGETS: { diff --git a/src/lib/jobs/__tests__/coach-nudge-ai-consent.test.ts b/src/lib/jobs/__tests__/coach-nudge-ai-consent.test.ts index b267ff776..35aff0813 100644 --- a/src/lib/jobs/__tests__/coach-nudge-ai-consent.test.ts +++ b/src/lib/jobs/__tests__/coach-nudge-ai-consent.test.ts @@ -35,6 +35,7 @@ const budgetMocks = vi.hoisted(() => ({ })), reconcileSpend: vi.fn(async () => {}), resolveDailyCap: vi.fn(() => 200_000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/coach/budget", () => budgetMocks); 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 33049dd2e..36a14cd9e 100644 --- a/src/lib/jobs/__tests__/document-content-index-backfill.test.ts +++ b/src/lib/jobs/__tests__/document-content-index-backfill.test.ts @@ -49,6 +49,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/jobs/boss-instance", () => ({ getGlobalBoss: vi.fn() })); diff --git a/src/lib/jobs/__tests__/document-summary.test.ts b/src/lib/jobs/__tests__/document-summary.test.ts index a88010654..af419b6c7 100644 --- a/src/lib/jobs/__tests__/document-summary.test.ts +++ b/src/lib/jobs/__tests__/document-summary.test.ts @@ -50,6 +50,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: vi.fn(), reconcileSpend: vi.fn().mockResolvedValue(undefined), resolveDailyCap: vi.fn(() => 1000), + resolveCostOwner: vi.fn(() => "operator" as const), })); vi.mock("@/lib/ai/ai-budgets", () => ({ AI_BUDGETS: { documentSummary: { temperature: 0.3, maxTokens: 600 } }, @@ -124,6 +125,7 @@ beforeEach(() => { vi.mocked(reserveBudget).mockResolvedValue({ allowed: true, reserved: 5, + owner: "operator", } as never); }); @@ -151,7 +153,14 @@ describe("runDocumentSummaryJob — gating", () => { "A lab report from a clinic listing routine blood values.", ); // Budget reserved then reconciled at the full reserved amount (charged). - expect(reconcileSpend).toHaveBeenCalledWith("user-1", 5, 5, "2026-07-17"); + expect(reconcileSpend).toHaveBeenCalledWith( + "user-1", + 5, + 5, + "2026-07-17", + 0, + { servedBy: expect.any(String), reservedOwner: "operator" }, + ); // Write is scoped to the still-null column so a re-run cannot clobber it. expect(prisma.inboundDocument.updateMany).toHaveBeenCalledWith( expect.objectContaining({ @@ -216,7 +225,14 @@ describe("runDocumentSummaryJob — gating", () => { await runDocumentSummaryJob({ userId: "user-1", documentId: "doc-1" }); // Reservation refunded to zero spend; no summary written, state recorded. - expect(reconcileSpend).toHaveBeenCalledWith("user-1", 5, 0, "2026-07-17"); + expect(reconcileSpend).toHaveBeenCalledWith( + "user-1", + 5, + 0, + "2026-07-17", + 0, + { servedBy: null, reservedOwner: "operator" }, + ); expect(prisma.inboundDocument.updateMany).toHaveBeenCalledWith( expect.objectContaining({ data: { summaryState: "UNAVAILABLE" } }), ); diff --git a/src/lib/jobs/__tests__/reaction-line-degradation.test.ts b/src/lib/jobs/__tests__/reaction-line-degradation.test.ts index 59c08e10d..9883df194 100644 --- a/src/lib/jobs/__tests__/reaction-line-degradation.test.ts +++ b/src/lib/jobs/__tests__/reaction-line-degradation.test.ts @@ -88,6 +88,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ reserveBudget: (...a: unknown[]) => reserveBudget(...a), reconcileSpend: (...a: unknown[]) => reconcileSpend(...a), resolveDailyCap: () => 200_000, + resolveCostOwner: () => "operator" as const, })); const loadDailyDigest = vi.fn(); @@ -467,6 +468,7 @@ describe("reaction line — degradation", () => { 1_400, "2026-07-16", 0, + { servedBy: expect.any(String), reservedOwner: "operator" }, ); }); @@ -528,6 +530,8 @@ describe("reaction line — degradation", () => { 1_400, 1_400, "2026-07-16", + 0, + { servedBy: expect.any(String), reservedOwner: "operator" }, ); expect(updateMany).toHaveBeenCalledWith( expect.objectContaining({ @@ -623,6 +627,7 @@ describe("reaction line — degradation", () => { 900, "2026-07-16", 0, + { servedBy: expect.any(String), reservedOwner: "operator" }, ); expect(update).not.toHaveBeenCalled(); expectSurfaceStillWorks(); @@ -683,6 +688,7 @@ describe("reaction line — degradation", () => { 1_100, "2026-07-16", 100, + { servedBy: expect.any(String), reservedOwner: "operator" }, ); }); }); diff --git a/src/lib/jobs/coach-nudge-ai.ts b/src/lib/jobs/coach-nudge-ai.ts index 756ca1e02..b4127c5d5 100644 --- a/src/lib/jobs/coach-nudge-ai.ts +++ b/src/lib/jobs/coach-nudge-ai.ts @@ -43,6 +43,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; @@ -199,6 +200,7 @@ export const composeNudgeWithAI: ComposeNudgeWithAI = async (params) => { maxTokens, dateKey, dailyCap, + resolveCostOwner(chain), ); if (!reservation.allowed) { annotate({ action: { name: "coach.nudge.ai.budget_exceeded" } }); @@ -224,12 +226,10 @@ export const composeNudgeWithAI: ComposeNudgeWithAI = async (params) => { } catch { // Timeout / network / provider error → refund what wasn't spent and // fall back to the template. - await reconcileSpend( - params.userId, - reservation.reserved, - 0, - dateKey, - ).catch(() => {}); + await reconcileSpend(params.userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }).catch(() => {}); annotate({ action: { name: "coach.nudge.ai.fallback" } }); return null; } @@ -240,6 +240,7 @@ export const composeNudgeWithAI: ComposeNudgeWithAI = async (params) => { result.tokensUsed ?? 0, dateKey, result.cachedInputTokens ?? 0, + { servedBy: chain[0].providerType, reservedOwner: reservation.owner }, ).catch(() => {}); const body = sanitiseAiBody(result.content, params.locale); diff --git a/src/lib/jobs/document-content-index-backfill.ts b/src/lib/jobs/document-content-index-backfill.ts index d2b8a8051..dc6906a93 100644 --- a/src/lib/jobs/document-content-index-backfill.ts +++ b/src/lib/jobs/document-content-index-backfill.ts @@ -38,6 +38,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { prisma } from "@/lib/db"; @@ -231,6 +232,7 @@ export async function runContentIndexBackfillForUser( AI_BUDGETS.documentTranscribe.maxTokens, dateKey, dailyCap, + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { reason = "budget-reached"; @@ -239,7 +241,10 @@ export async function runContentIndexBackfillForUser( const document = await loadOwnedDocument(userId, id); if (!document) { - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); skipped += 1; continue; } @@ -249,7 +254,10 @@ export async function runContentIndexBackfillForUser( // or a PDF the rasteriser could not render) — refund and move on so // the walk keeps converging. Counted as skipped, not failed: no // provider call happened. - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); skipped += 1; continue; } @@ -266,6 +274,11 @@ export async function runContentIndexBackfillForUser( reservation.reserved, reservation.reserved, dateKey, + 0, + { + servedBy: pick.entry.providerType, + reservedOwner: reservation.owner, + }, ); await upsertContentIndex({ userId, @@ -278,7 +291,10 @@ export async function runContentIndexBackfillForUser( } catch { // A provider miss on one document must not abort the batch — refund and // leave it un-indexed (a later run retries it). - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); failed += 1; } } diff --git a/src/lib/jobs/document-summary.ts b/src/lib/jobs/document-summary.ts index 21cbc9469..86308b5c9 100644 --- a/src/lib/jobs/document-summary.ts +++ b/src/lib/jobs/document-summary.ts @@ -30,6 +30,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveCostOwner, resolveDailyCap, } from "@/lib/ai/coach/budget"; import { @@ -191,6 +192,7 @@ export async function runDocumentSummaryJob( AI_BUDGETS.documentSummary.maxTokens, dateKey, resolveDailyCap([{ providerType: pick.entry.providerType }]), + resolveCostOwner([{ providerType: pick.entry.providerType }]), ); // Budget exhausted → skip (no local fallback for a summary). if (!reservation.allowed) { @@ -224,6 +226,8 @@ export async function runDocumentSummaryJob( reservation.reserved, reservation.reserved, dateKey, + 0, + { servedBy: pick.entry.providerType, reservedOwner: reservation.owner }, ); // WITHHOLD policy, v1.30.31. The model's prose still never lands — a // summary that tripped the outbound screen must never be shown as if it @@ -273,7 +277,10 @@ export async function runDocumentSummaryJob( // Refund the reservation on a provider miss; the document keeps no summary // (the on-demand route remains the manual fallback). Never rethrow — a // transient provider error must not retry-loop or fail the queue. - await reconcileSpend(userId, reservation.reserved, 0, dateKey); + await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, { + servedBy: null, + reservedOwner: reservation.owner, + }); await markSummaryState(userId, documentId, "UNAVAILABLE"); annotate({ action: { name: "documents.summary.autoFailed" }, diff --git a/src/lib/jobs/reaction-line.ts b/src/lib/jobs/reaction-line.ts index cd2601c4b..a3678576c 100644 --- a/src/lib/jobs/reaction-line.ts +++ b/src/lib/jobs/reaction-line.ts @@ -41,7 +41,9 @@ import { isModuleEnabled } from "@/lib/modules/gate"; import { buildDateKey, reconcileSpend, + resolveCostOwner, resolveDailyCap, + type BudgetCostOwner, } from "@/lib/ai/coach/budget"; import { AI_BUDGETS } from "@/lib/ai/ai-budgets"; import { screenCoachReply } from "@/lib/ai/coach/outbound-guard"; @@ -314,6 +316,13 @@ type ReactionReservation = { allowed: boolean; reserved: number; dateKey: string; + /** + * v1.38.19 (Wave E) — the cost owner the reservation was booked to. A + * RESUMED reservation (the claim survived a crash, its tokens are on the + * row) recomputes it from the same chain the reservation was made under, so + * the reconcile reverses exactly what the reservation added. + */ + owner: BudgetCostOwner; }; class ReactionClaimLostError extends Error { @@ -334,6 +343,7 @@ async function reserveClaimBudget( revision: Date, dateKey: string, cap: number, + owner: BudgetCostOwner, ): Promise { return prisma.$transaction(async (tx) => { const stillOwned = await tx.arrivalReaction.updateMany({ @@ -352,11 +362,17 @@ async function reserveClaimBudget( if (stillOwned.count !== 1) throw new ReactionClaimLostError(); const reserved = ARRIVAL_REACTION_RESERVE_TOKENS; + // v1.38.19 (Wave E) — this reservation is the claim-linked twin of + // `reserveBudget`, so it books the owner split the same way: an + // operator-funded chain increments `operator_tokens` inside the SAME + // statement as the total, and the refusal path reverses both. + const operatorReserved = owner === "operator" ? reserved : 0; const rows = await tx.$queryRaw<{ total_tokens: number }[]>` - INSERT INTO coach_usage (id, user_id, date_key, total_tokens, message_count, created_at, updated_at) - VALUES (gen_random_uuid()::text, ${userId}, ${dateKey}, ${reserved}, 1, NOW(), NOW()) + INSERT INTO coach_usage (id, user_id, date_key, total_tokens, operator_tokens, message_count, created_at, updated_at) + VALUES (gen_random_uuid()::text, ${userId}, ${dateKey}, ${reserved}, ${operatorReserved}, 1, NOW(), NOW()) ON CONFLICT (user_id, date_key) DO UPDATE SET total_tokens = coach_usage.total_tokens + ${reserved}, + operator_tokens = coach_usage.operator_tokens + ${operatorReserved}, message_count = coach_usage.message_count + 1, updated_at = NOW() RETURNING total_tokens @@ -366,6 +382,7 @@ async function reserveClaimBudget( await tx.$executeRaw` UPDATE coach_usage SET total_tokens = GREATEST(0, total_tokens - ${reserved}), + operator_tokens = GREATEST(0, operator_tokens - ${operatorReserved}), message_count = GREATEST(0, message_count - 1), updated_at = NOW() WHERE user_id = ${userId} AND date_key = ${dateKey} @@ -374,7 +391,7 @@ async function reserveClaimBudget( where: { id: rowId, userId, occurredAt: revision, generationClaimId }, data: { generationClaimId: null, generationClaimedAt: null }, }); - return { allowed: false, reserved, dateKey }; + return { allowed: false, reserved, dateKey, owner }; } const linked = await tx.arrivalReaction.updateMany({ @@ -395,7 +412,7 @@ async function reserveClaimBudget( }); if (linked.count !== 1) throw new ReactionClaimLostError(); - return { allowed: true, reserved, dateKey }; + return { allowed: true, reserved, dateKey, owner }; }); } @@ -526,6 +543,7 @@ export async function runReactionLine( allowed: true, reserved: row.generationReservedTokens!, dateKey: row.generationBudgetDateKey!, + owner: resolveCostOwner(chain), }; } else { try { @@ -536,6 +554,7 @@ export async function runReactionLine( revision, buildDateKey(), resolveDailyCap(chain), + resolveCostOwner(chain), ); } catch (err) { if (err instanceof ReactionClaimLostError) { @@ -625,6 +644,8 @@ export async function runReactionLine( reservation.reserved, reservation.reserved, reservation.dateKey, + 0, + { servedBy: chain[0].providerType, reservedOwner: reservation.owner }, ).catch(() => {}); await finishTerminalAttempt().catch(() => {}); workerLog("error", "[reaction-line] generation failed", err); @@ -638,6 +659,7 @@ export async function runReactionLine( result.tokensUsed ?? reservation.reserved, reservation.dateKey, result.cachedInputTokens ?? 0, + { servedBy: chain[0].providerType, reservedOwner: reservation.owner }, ); } catch (err) { // Do not publish a line whose spend was not durably reconciled. The diff --git a/tests/integration/coach-budget-owner.test.ts b/tests/integration/coach-budget-owner.test.ts new file mode 100644 index 000000000..ab610b06d --- /dev/null +++ b/tests/integration/coach-budget-owner.test.ts @@ -0,0 +1,186 @@ +/** + * v1.38.19 (Wave E) — the daily coach ledger attributes tokens to the COST + * OWNER of the hop that actually served them. + * + * Production evidence (operator instance, 2026-09-11): the chat refused with + * `coach.budget.exceeded` at `totalAfter: 513539` although almost every token + * of that day had been served by `codex` — the operator's own ChatGPT plan — + * after the operator-funded `admin-openai` primary answered HTTP 500. The + * ledger had no provider dimension, so tokens the USER's plan paid for were + * counted against the operator-key cap and the interactive chat was locked out + * by background jobs every morning. + * + * These contracts need real Postgres: the attribution lives inside the same + * atomic statements as the total (`INSERT … ON CONFLICT DO UPDATE … RETURNING` + * and the reconcile `UPDATE`), and a mocked client cannot prove that the two + * counters move together under the row lock. + */ +import { beforeEach, describe, expect, it } from "vitest"; + +import { getPrismaClient, truncateAllTables } from "./setup"; + +async function seedUser(): Promise { + const user = await getPrismaClient().user.create({ + data: { + username: "budget-owner-user", + email: "budget-owner@example.test", + role: "USER", + }, + }); + return user.id; +} + +async function readRow(userId: string, dateKey: string) { + return getPrismaClient().coachUsage.findUnique({ + where: { userId_dateKey: { userId, dateKey } }, + }); +} + +beforeEach(async () => { + await truncateAllTables(getPrismaClient()); +}); + +describe("coach_usage.operator_tokens (real Postgres)", () => { + it("bills an operator-served turn to both counters", async () => { + const { reserveBudget, reconcileSpend, OPERATOR_COST_CAP } = + await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-09-11"; + + const res = await reserveBudget( + userId, + 1_000, + dateKey, + OPERATOR_COST_CAP, + "operator", + ); + expect(res.allowed).toBe(true); + await reconcileSpend(userId, res.reserved, 3_000, dateKey, 0, { + servedBy: "admin-openai", + reservedOwner: "operator", + }); + + const row = await readRow(userId, dateKey); + expect(row?.totalTokens).toBe(3_000); + expect(row?.operatorTokens).toBe(3_000); + }); + + it("moves the reservation out of operator_tokens when the user's own plan served", async () => { + const { reserveBudget, reconcileSpend, OPERATOR_COST_CAP } = + await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-09-11"; + + // The chain's primary is operator-funded, so the reservation is charged to + // the operator — but the `admin-openai` hop failed and `codex` (the user's + // own ChatGPT plan) answered. The operator pays nothing for that turn. + const res = await reserveBudget( + userId, + 1_000, + dateKey, + OPERATOR_COST_CAP, + "operator", + ); + await reconcileSpend(userId, res.reserved, 3_000, dateKey, 0, { + servedBy: "codex", + reservedOwner: "operator", + }); + + const row = await readRow(userId, dateKey); + expect(row?.totalTokens).toBe(3_000); + expect(row?.operatorTokens).toBe(0); + }); + + it("a user-plan turn never touches an existing operator balance", async () => { + const { reserveBudget, reconcileSpend, USER_PLAN_CAP } = + await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-09-11"; + + // An operator-funded job already spent 150k earlier that day. + await getPrismaClient().coachUsage.create({ + data: { + userId, + dateKey, + totalTokens: 150_000, + operatorTokens: 150_000, + messageCount: 12, + }, + }); + + const res = await reserveBudget( + userId, + 1_000, + dateKey, + USER_PLAN_CAP, + "user", + ); + await reconcileSpend(userId, res.reserved, 4_000, dateKey, 0, { + servedBy: "codex", + reservedOwner: "user", + }); + + const row = await readRow(userId, dateKey); + expect(row?.totalTokens).toBe(154_000); + // The operator's balance is untouched: a user-plan turn must neither add + // to it nor deplete it. + expect(row?.operatorTokens).toBe(150_000); + }); + + it("bills an operator-funded FALLBACK hop under a user-plan reservation", async () => { + const { reserveBudget, reconcileSpend, USER_PLAN_CAP } = + await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-09-11"; + + // Primary is the user's own key (user-plan reservation), but the chain + // fell back onto the operator's shared key — those tokens land on the + // operator's invoice and must be counted there. + const res = await reserveBudget( + userId, + 1_000, + dateKey, + USER_PLAN_CAP, + "user", + ); + await reconcileSpend(userId, res.reserved, 2_500, dateKey, 0, { + servedBy: "admin-openai", + reservedOwner: "user", + }); + + const row = await readRow(userId, dateKey); + expect(row?.totalTokens).toBe(2_500); + expect(row?.operatorTokens).toBe(2_500); + }); + + it("refunds both counters when a refused reservation is rolled back", async () => { + const { reserveBudget, OPERATOR_COST_CAP } = + await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-09-11"; + + await getPrismaClient().coachUsage.create({ + data: { + userId, + dateKey, + totalTokens: OPERATOR_COST_CAP, + operatorTokens: OPERATOR_COST_CAP, + messageCount: 3, + }, + }); + + const res = await reserveBudget( + userId, + 1_000, + dateKey, + OPERATOR_COST_CAP, + "operator", + ); + expect(res.allowed).toBe(false); + + const row = await readRow(userId, dateKey); + expect(row?.totalTokens).toBe(OPERATOR_COST_CAP); + expect(row?.operatorTokens).toBe(OPERATOR_COST_CAP); + expect(row?.messageCount).toBe(3); + }); +}); diff --git a/tests/integration/coach-budget-reservation.test.ts b/tests/integration/coach-budget-reservation.test.ts index b500a3dab..1ff211fc0 100644 --- a/tests/integration/coach-budget-reservation.test.ts +++ b/tests/integration/coach-budget-reservation.test.ts @@ -42,8 +42,8 @@ describe("reserveBudget (real Postgres)", () => { }); const [a, b] = await Promise.all([ - reserveBudget(userId, 600, dateKey, cap), - reserveBudget(userId, 600, dateKey, cap), + reserveBudget(userId, 600, dateKey, cap, "operator"), + reserveBudget(userId, 600, dateKey, cap, "operator"), ]); // Exactly one is admitted: the first to commit sees prior 900 (< cap) and @@ -65,7 +65,7 @@ describe("reserveBudget (real Postgres)", () => { const userId = await seedUser(); const dateKey = "2026-06-20"; - const res = await reserveBudget(userId, 600, dateKey, 25_000); + const res = await reserveBudget(userId, 600, dateKey, 25_000, "operator"); expect(res.allowed).toBe(true); const row = await getPrismaClient().coachUsage.findUnique({ @@ -81,8 +81,11 @@ describe("reserveBudget (real Postgres)", () => { const userId = await seedUser(); const dateKey = "2026-06-21"; - const res = await reserveBudget(userId, 600, dateKey, 25_000); - await reconcileSpend(userId, res.reserved, 120, dateKey); + const res = await reserveBudget(userId, 600, dateKey, 25_000, "operator"); + await reconcileSpend(userId, res.reserved, 120, dateKey, 0, { + servedBy: "admin-openai", + reservedOwner: res.owner, + }); const row = await getPrismaClient().coachUsage.findUnique({ where: { userId_dateKey: { userId, dateKey } }, @@ -99,8 +102,11 @@ describe("reserveBudget (real Postgres)", () => { // Provider returned 450 tokens but the reply was empty/sentinel: the // reconcile must still record the 450 burned, not zero. - const res = await reserveBudget(userId, 600, dateKey, 25_000); - await reconcileSpend(userId, res.reserved, 450, dateKey); + const res = await reserveBudget(userId, 600, dateKey, 25_000, "operator"); + await reconcileSpend(userId, res.reserved, 450, dateKey, 0, { + servedBy: "admin-openai", + reservedOwner: res.owner, + }); const row = await getPrismaClient().coachUsage.findUnique({ where: { userId_dateKey: { userId, dateKey } },