From 378334850610f88a9adb93a810c64c688dcbdccc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Fri, 26 Jun 2026 12:53:04 +0200 Subject: [PATCH] fix(coach): make the daily token budget provider-aware The daily AI-token gate was a flat 25,000/day for every provider. That cap exists to bound the OPERATOR's API bill, but a ChatGPT-OAuth (Codex) or BYOK chain egresses on the USER's own plan/key and costs the operator nothing, so gating it on the operator-cost ceiling is wrong. A single gpt-5.x reasoning turn legitimately reports 20k-40k gross total_tokens (re-sent system prompt + inventory + tool defs per round + hidden reasoning, summed across rounds), so one or two turns exhausted the flat cap and locked the user out of a plan they pay for. Classify the chain's cost owner by its primary provider via resolveDailyCap: an admin-openai primary (operator pays) keeps the operator-cost cap; every user-egress primary (codex/openai/anthropic/local) gets a generous abuse-only ceiling. Threaded into the Coach chat gate and both OCR-extract modes. Raise the operator-cost cap 25,000 -> 200,000 so the operator-key path also survives a normal day of reasoning turns, keeping gross-token accounting rather than reworking the reconcile math. Subtract cached input tokens at reconcile: the Responses-API gross total still includes prompt-cached input the user did not re-pay for. The codex client already parsed cached_tokens; sum it through the tool loop and bill total_tokens minus cached. --- .../chat/__tests__/route-module-gate.test.ts | 1 + .../__tests__/route-snapshot-once.test.ts | 1 + .../chat/__tests__/route-tool-mode.test.ts | 21 ++- src/app/api/insights/chat/route.ts | 14 ++ .../labs/ocr/extract/__tests__/route.test.ts | 3 + src/app/api/labs/ocr/extract/route.ts | 9 ++ src/lib/ai/coach/__tests__/budget.test.ts | 138 +++++++++++++++++- src/lib/ai/coach/budget.ts | 73 ++++++++- src/lib/ai/coach/tools/loop.ts | 9 ++ 9 files changed, 262 insertions(+), 7 deletions(-) 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 8b6c82fd3..6c92359fb 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 @@ -94,6 +94,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ buildDateKey: vi.fn(), enforceBudget: vi.fn(), recordSpend: vi.fn(), + resolveDailyCap: vi.fn(() => 200_000), })); 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 ef6332dde..eeaea14f3 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 @@ -100,6 +100,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ buildDateKey: vi.fn(() => "2026-06-21"), reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 1500 })), reconcileSpend: vi.fn(async () => undefined), + resolveDailyCap: vi.fn(() => 2_000_000), })); 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 3f4a1e17d..c38816b39 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 @@ -111,6 +111,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ buildDateKey: vi.fn(() => "2026-06-21"), reserveBudget, reconcileSpend, + resolveDailyCap: vi.fn(() => 2_000_000), })); const { detectRefusal } = vi.hoisted(() => ({ @@ -245,8 +246,14 @@ describe("coach chat — tool-mode routing (F1)", () => { { providerType: "anthropic", instance: {} }, ]); await post(chatReq({ message: "How is my BP?" })); - // maxTokens (1500) × MAX_ROUNDS (2). - expect(reserveBudget).toHaveBeenCalledWith("u1", 3000, "2026-06-21"); + // maxTokens (1500) × MAX_ROUNDS (2). The 4th arg is the provider-aware + // daily cap (F1) — mocked to the user-plan ceiling for this BYOK chain. + expect(reserveBudget).toHaveBeenCalledWith( + "u1", + 3000, + "2026-06-21", + 2_000_000, + ); }); it("persists the tool trace onto provenance", async () => { @@ -321,8 +328,14 @@ describe("coach chat — tool-mode routing (F1)", () => { expect(runCoachToolLoop).not.toHaveBeenCalled(); expect(runRawCompletionWithFallback).toHaveBeenCalledTimes(1); - // Single-round budget reservation in the fallback path. - expect(reserveBudget).toHaveBeenCalledWith("u1", 1500, "2026-06-21"); + // Single-round budget reservation in the fallback path; the 4th arg is the + // provider-aware daily cap (F1). + expect(reserveBudget).toHaveBeenCalledWith( + "u1", + 1500, + "2026-06-21", + 2_000_000, + ); // The legacy path ships the snapshot figures in the user turn. const params = ( (runRawCompletionWithFallback.mock.calls[0] as unknown[])[0] as { diff --git a/src/app/api/insights/chat/route.ts b/src/app/api/insights/chat/route.ts index 9fca7a8d8..13d02b2f4 100644 --- a/src/app/api/insights/chat/route.ts +++ b/src/app/api/insights/chat/route.ts @@ -64,6 +64,7 @@ import { buildDateKey, reserveBudget, reconcileSpend, + resolveDailyCap, } from "@/lib/ai/coach/budget"; import { detectRefusal } from "@/lib/ai/coach/refusal"; import { @@ -503,6 +504,12 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`; // so reserve the per-call ceiling × the round count up front and reconcile // the SUMMED actual tokens afterwards. The atomic reserve/reconcile // primitives are unchanged; only the reserved amount scales. + // v1.21.0 (F1) — the daily ceiling is the OPERATOR's cost cap only when the + // chain egresses via the operator's own key (`admin-openai` primary). A + // ChatGPT-OAuth/Codex or BYOK chain runs on the user's OWN plan/key and costs + // the operator nothing, so it gets the generous user-plan ceiling — gating it + // on the operator-cost cap would lock the user out of a plan they pay for. + const dailyCap = resolveDailyCap(chain); const reqDateKey = buildDateKey(); const reservation = await reserveBudget( userId, @@ -510,6 +517,7 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`; ? AI_BUDGETS.coach.maxTokens * MAX_ROUNDS : AI_BUDGETS.coach.maxTokens, reqDateKey, + dailyCap, ); if (!reservation.allowed) { annotate({ @@ -526,6 +534,9 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`; // prose number-verifier. Empty on the no-tools path. let toolResultPayloads: unknown[] = []; let totalTokensSpent: number; + // v1.21.0 (F3) — cached-input tokens to subtract at reconcile (prompt-cached + // input the user did not re-pay for must not be billed to the daily meter). + let cachedTokensSpent = 0; try { if (toolMode) { // v1.20.0 (F1) — base context: the full system prompt + a tool-mode @@ -565,6 +576,7 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}. Fetch toolTrace = loop.toolTrace; toolResultPayloads = (loop.toolResults ?? []).map((r) => r.data); totalTokensSpent = loop.totalTokens; + cachedTokensSpent = loop.cachedTokens; } else { const fallback = await runRawCompletionWithFallback({ userId, @@ -587,6 +599,7 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}. Fetch result = fallback.result; workingProviderType = fallback.workingProvider.providerType; totalTokensSpent = result.tokensUsed ?? 0; + cachedTokensSpent = result.cachedInputTokens ?? 0; } } catch (err) { // The provider chain failed outright — no tokens were billed, so refund @@ -643,6 +656,7 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}. Fetch reservation.reserved, totalTokensSpent, reqDateKey, + cachedTokensSpent, ).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 055a56a1f..f4617a05c 100644 --- a/src/app/api/labs/ocr/extract/__tests__/route.test.ts +++ b/src/app/api/labs/ocr/extract/__tests__/route.test.ts @@ -41,6 +41,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({ buildDateKey: vi.fn(() => "2026-06-26"), reserveBudget: vi.fn(), reconcileSpend: vi.fn().mockResolvedValue(undefined), + resolveDailyCap: vi.fn(() => 200_000), })); vi.mock("@/lib/labs/ocr-extract", async () => { const actual = await vi.importActual( @@ -106,6 +107,8 @@ describe("POST /api/labs/ocr/extract — text mode budget", () => { "user-1", AI_BUDGETS.ocrExtractText.maxTokens, "2026-06-26", + // F1 — the provider-aware daily cap (mocked) is threaded as the 4th arg. + 200_000, ); }); diff --git a/src/app/api/labs/ocr/extract/route.ts b/src/app/api/labs/ocr/extract/route.ts index e35b710d6..ebddbf99c 100644 --- a/src/app/api/labs/ocr/extract/route.ts +++ b/src/app/api/labs/ocr/extract/route.ts @@ -35,6 +35,7 @@ import { buildDateKey, reconcileSpend, reserveBudget, + resolveDailyCap, } from "@/lib/ai/coach/budget"; import { prisma } from "@/lib/db"; import { @@ -142,11 +143,15 @@ async function handleTextExtract( // a vision call, so it reserves the proportionate text ceiling rather than // the vision budget. Over-charging the vision rate for a text call would // exhaust the day budget against spend that never happened. + // v1.21.0 (F1) — the operator-cost cap applies only when the picked provider + // egresses on the operator's own key; a BYOK / Codex / local pick runs on the + // user's own plan and gets the generous user-plan ceiling. const dateKey = buildDateKey(); const reservation = await reserveBudget( userId, AI_BUDGETS.ocrExtractText.maxTokens, dateKey, + resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { annotate({ @@ -234,11 +239,15 @@ async function handleVisionExtract( } // 4. Reserve the day's budget BEFORE the provider call (atomic, TOCTOU-safe). + // v1.21.0 (F1) — the operator-cost cap applies only when the picked vision + // provider egresses on the operator's own key; a BYOK / Codex pick runs on + // the user's own plan and gets the generous user-plan ceiling. const dateKey = buildDateKey(); const reservation = await reserveBudget( userId, AI_BUDGETS.ocrExtract.maxTokens, dateKey, + resolveDailyCap([{ providerType: pick.entry.providerType }]), ); if (!reservation.allowed) { annotate({ diff --git a/src/lib/ai/coach/__tests__/budget.test.ts b/src/lib/ai/coach/__tests__/budget.test.ts index aca47c40b..267cb1dcf 100644 --- a/src/lib/ai/coach/__tests__/budget.test.ts +++ b/src/lib/ai/coach/__tests__/budget.test.ts @@ -1,7 +1,13 @@ import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; import { HttpError } from "@/lib/api-handler"; -import { buildDateKey, MAX_TOKENS_PER_USER_PER_DAY } from "../budget"; +import { + buildDateKey, + MAX_TOKENS_PER_USER_PER_DAY, + OPERATOR_COST_CAP, + USER_PLAN_CAP, + resolveDailyCap, +} from "../budget"; vi.mock("@/lib/db", () => ({ prisma: { @@ -9,6 +15,8 @@ vi.mock("@/lib/db", () => ({ findUnique: vi.fn(), upsert: vi.fn(), }, + $executeRaw: vi.fn(), + $queryRaw: vi.fn(), }, })); @@ -105,3 +113,131 @@ describe("budget", () => { expect(args.create.totalTokens).toBe(12); }); }); + +describe("resolveDailyCap (F1 — provider-aware cap)", () => { + it("applies the operator-cost cap to an operator-key (admin-openai) primary", () => { + expect(resolveDailyCap([{ providerType: "admin-openai" }])).toBe( + OPERATOR_COST_CAP, + ); + }); + + it("applies the generous user-plan cap to a ChatGPT-OAuth (codex) primary", () => { + expect(resolveDailyCap([{ providerType: "codex" }])).toBe(USER_PLAN_CAP); + }); + + it("applies the user-plan cap to BYOK openai / anthropic / local primaries", () => { + expect(resolveDailyCap([{ providerType: "openai" }])).toBe(USER_PLAN_CAP); + expect(resolveDailyCap([{ providerType: "anthropic" }])).toBe( + USER_PLAN_CAP, + ); + expect(resolveDailyCap([{ providerType: "local" }])).toBe(USER_PLAN_CAP); + }); + + it("classifies on the PRIMARY entry — a user-egress chain with an admin-openai fallback stays user-plan", () => { + expect( + resolveDailyCap([ + { providerType: "codex" }, + { providerType: "admin-openai" }, + ]), + ).toBe(USER_PLAN_CAP); + }); + + it("defaults an empty chain to the conservative operator cap", () => { + expect(resolveDailyCap([])).toBe(OPERATOR_COST_CAP); + }); +}); + +describe("reserveBudget cap (F1 — user-plan path not locked out)", () => { + let prismaMock: { + $queryRaw: ReturnType; + $executeRaw: ReturnType; + }; + + beforeEach(async () => { + const dbModule = await import("@/lib/db"); + prismaMock = dbModule.prisma as unknown as typeof prismaMock; + prismaMock.$queryRaw.mockReset(); + prismaMock.$executeRaw.mockReset(); + }); + + it("a user-plan chain does NOT trip after a spend that exceeds the operator cap", async () => { + // Prior spend well past the 200k operator cap, but under the user-plan cap. + const priorSpend = OPERATOR_COST_CAP + 50_000; + prismaMock.$queryRaw.mockResolvedValue([ + { total_tokens: priorSpend + 1_200 }, + ]); + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget( + "u", + 1_200, + "2026-05-10", + resolveDailyCap([{ providerType: "codex" }]), + ); + expect(res.allowed).toBe(true); + // The reservation upsert ran; no refund executeRaw fired. + expect(prismaMock.$executeRaw).not.toHaveBeenCalled(); + }); + + it("the operator-key path STILL trips once prior spend reaches the operator cap", async () => { + const priorSpend = OPERATOR_COST_CAP; + prismaMock.$queryRaw.mockResolvedValue([ + { total_tokens: priorSpend + 1_200 }, + ]); + prismaMock.$executeRaw.mockResolvedValue(0); + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget( + "u", + 1_200, + "2026-05-10", + resolveDailyCap([{ providerType: "admin-openai" }]), + ); + expect(res.allowed).toBe(false); + // Refund of the reservation fired on refusal. + expect(prismaMock.$executeRaw).toHaveBeenCalled(); + }); +}); + +describe("reconcileSpend cached-token subtraction (F3)", () => { + let prismaMock: { $executeRaw: ReturnType }; + + beforeEach(async () => { + const dbModule = await import("@/lib/db"); + prismaMock = dbModule.prisma as unknown as typeof prismaMock; + prismaMock.$executeRaw.mockReset(); + prismaMock.$executeRaw.mockResolvedValue(0); + }); + + 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); + 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); + expect(interpolations).toContain(5_800); + }); + + 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); + const interpolations = prismaMock.$executeRaw.mock.calls[0].slice(1); + // net actual clamped to 0 → delta = 0 - 1200 = -1200. + expect(interpolations).toContain(-1_200); + }); + + 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"); + const interpolations = prismaMock.$executeRaw.mock.calls[0].slice(1); + expect(interpolations).toContain(3_000); + }); +}); + +describe("MAX_TOKENS_PER_USER_PER_DAY (F2 — reasoning-aware operator cap)", () => { + it("is sized for reasoning turns, not a single non-reasoning reply", () => { + // Was 25_000 (≈ one gpt-5.x reasoning turn). Raised so the operator-key + // path survives a normal day of reasoning turns. + expect(MAX_TOKENS_PER_USER_PER_DAY).toBeGreaterThanOrEqual(150_000); + expect(OPERATOR_COST_CAP).toBe(MAX_TOKENS_PER_USER_PER_DAY); + }); +}); diff --git a/src/lib/ai/coach/budget.ts b/src/lib/ai/coach/budget.ts index 02df866e7..f52f52340 100644 --- a/src/lib/ai/coach/budget.ts +++ b/src/lib/ai/coach/budget.ts @@ -16,8 +16,63 @@ */ import { prisma } from "@/lib/db"; import { HttpError } from "@/lib/api-handler"; +import type { ProviderChainType } from "@/lib/ai/provider-chain"; -export const MAX_TOKENS_PER_USER_PER_DAY = 25_000; +/** + * Operator-cost daily ceiling — the cap that protects the OPERATOR's LLM bill. + * + * It applies ONLY to the operator-managed-key path (`admin-openai`): the + * server's own OpenAI key, where every token the user spends lands on the + * operator's invoice. + * + * v1.21.0 (F1/F2) — raised from the historical 25_000. That figure was sized + * for a single non-reasoning ~600-token chat reply (~20 turns/day). With the + * v1.20 tool loop charging the Responses-API gross `total_tokens` of a `gpt-5.x` + * reasoning turn (re-sent system prompt + inventory + 7 tool defs per round + + * hidden reasoning, summed across rounds ≈ 20k–40k), 25k locked the user out + * after one turn. 200k keeps gross accounting (no per-round output-token + * surgery) while leaving room for a normal day of reasoning turns on the + * operator's key. The cap stays a real ceiling on the operator's exposure. + */ +export const MAX_TOKENS_PER_USER_PER_DAY = 200_000; + +/** + * v1.21.0 (F1) — historical name retained for any external reference; the + * operator-cost cap is the only cap the ledger enforces. + */ +export const OPERATOR_COST_CAP = MAX_TOKENS_PER_USER_PER_DAY; + +/** + * v1.21.0 (F1) — the daily ceiling for a chain whose egress runs on the + * USER's own plan / key (ChatGPT-OAuth/Codex, BYOK OpenAI/Anthropic, or a + * self-hosted local model). The operator pays nothing for these, so the + * operator-cost cap is a category error here — gating them on it locks a user + * out of a plan they pay for. We keep only a generous abuse ceiling so a + * runaway client loop can't write unbounded rows; a normal user never reaches + * it. + */ +export const USER_PLAN_CAP = 2_000_000; + +/** + * v1.21.0 (F1) — classify a resolved provider chain's cost owner and return + * the daily cap that applies. + * + * The chain is a fallback list; the FIRST entry is the provider that will be + * tried first and is the expected cost owner. Only when that primary provider + * is `admin-openai` (the user has no personal provider and falls back to the + * operator's key) does the operator pay — so the operator-cost cap applies. + * Every other primary (`codex` / `openai` / `anthropic` / `local`) is the + * user's own egress, so the generous user-plan cap applies. An empty chain + * (no provider resolved) defaults to the operator cap — the conservative side. + */ +export function resolveDailyCap( + chain: ReadonlyArray<{ providerType: ProviderChainType }>, +): number { + const primary = chain[0]?.providerType; + return primary === "admin-openai" || primary === undefined + ? OPERATOR_COST_CAP + : USER_PLAN_CAP; +} /** * Build the UTC day-key for a given clock. Defaults to "now". @@ -136,17 +191,31 @@ export async function reserveBudget( * goes negative). Called after every provider call — including empty / * sentinel / refusal replies, whose upstream tokens were still burned * (SENIOR-DEV MEDIUM: spend undercount). + * + * v1.21.0 (F3) — `cachedTokens` (the Responses-API `cached_tokens` count) is + * subtracted from the charged amount. The gross `total_tokens` a reasoning + * provider reports still includes the full input even when prompt-caching + * 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 async function reconcileSpend( userId: string, reserved: number, actualTokens: number, dateKey: string = buildDateKey(), + cachedTokens = 0, ): Promise { - const actual = + const grossActual = Number.isFinite(actualTokens) && actualTokens > 0 ? Math.floor(actualTokens) : 0; + const cached = + Number.isFinite(cachedTokens) && cachedTokens > 0 + ? Math.floor(cachedTokens) + : 0; + // Bill net of cached input; clamp so a cached count larger than the gross + // (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; // Clamp at zero so a smaller-than-reserved actual can't drive the row diff --git a/src/lib/ai/coach/tools/loop.ts b/src/lib/ai/coach/tools/loop.ts index 5e8208e4a..852813e67 100644 --- a/src/lib/ai/coach/tools/loop.ts +++ b/src/lib/ai/coach/tools/loop.ts @@ -45,6 +45,12 @@ export interface CoachToolLoopResult { workingProviderType: string; /** Summed tokens across every round (for budget reconcile). */ totalTokens: number; + /** + * v1.21.0 (F3) — summed cached-input tokens across every round. Subtracted + * from the charged amount at reconcile so prompt-cached input the user did + * not re-pay for is not billed to their daily meter. + */ + cachedTokens: number; /** Number of model round-trips made. */ rounds: number; /** Which tools ran + whether each found data (persisted onto provenance). */ @@ -86,6 +92,7 @@ export async function runCoachToolLoop(args: { const messages: AiMessage[] = [...args.messages]; let totalTokens = 0; + let cachedTokens = 0; let rounds = 0; let workingProviderType = ""; const toolTrace: CoachToolTrace[] = []; @@ -116,6 +123,7 @@ export async function runCoachToolLoop(args: { const result = fallback.result; workingProviderType = fallback.workingProvider.providerType; totalTokens += result.tokensUsed ?? 0; + cachedTokens += result.cachedInputTokens ?? 0; const calls = result.toolCalls ?? []; const wantsTools = @@ -136,6 +144,7 @@ export async function runCoachToolLoop(args: { result, workingProviderType, totalTokens, + cachedTokens, rounds, toolTrace, toolResults,