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 } },