From 2c46712e13f342efa35a818ade03b805b13be573 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Fri, 11 Sep 2026 09:56:44 +0200 Subject: [PATCH] fix(ai): the operator cap counts only tokens the operator paid for The reservation compared the day's TOTAL against whichever ceiling the chain resolved, and the chain walker's hop guard did the same. On an account whose background jobs run on the user's own ChatGPT plan, both saw a total the operator had not paid for: the chat refused at 06:42Z with 513 539 tokens against the 200 000 operator ceiling, and the fallback hop to the shared key was closed on the same arithmetic. An operator-funded chain is now measured against `operator_tokens` and a user-plan chain keeps its abuse ceiling on the total; the hop guard reads the operator share too. A day of turns the user's own plan paid for can no longer lock anyone out of the operator's ceiling. --- .../insights/generate/__tests__/route.test.ts | 8 +- src/lib/ai/__tests__/provider-runner.test.ts | 60 ++++++++++++++ src/lib/ai/coach/__tests__/budget.test.ts | 78 ++++++++++++++++++- src/lib/ai/coach/budget.ts | 12 ++- src/lib/ai/provider-runner.ts | 7 +- .../briefing-provider-budget.test.ts | 25 +++++- ...ensive-briefing-grounding-disposal.test.ts | 2 +- .../comprehensive-briefing-reroll.test.ts | 2 +- .../comprehensive-budget-refusal.test.ts | 22 +++++- .../comprehensive-generate-force.test.ts | 2 +- .../comprehensive-json-retry.test.ts | 2 +- .../comprehensive-privacy-flip.test.ts | 2 +- .../comprehensive-timeout-honored.test.ts | 2 +- .../comprehensive-warm-timeout.test.ts | 2 +- .../__tests__/status-provider-budget.test.ts | 25 +++++- tests/integration/coach-budget-owner.test.ts | 60 ++++++++++++++ 16 files changed, 285 insertions(+), 26 deletions(-) diff --git a/src/app/api/insights/generate/__tests__/route.test.ts b/src/app/api/insights/generate/__tests__/route.test.ts index 5289fa0d1..fe320041c 100644 --- a/src/app/api/insights/generate/__tests__/route.test.ts +++ b/src/app/api/insights/generate/__tests__/route.test.ts @@ -39,7 +39,7 @@ vi.mock("@/lib/db", () => ({ // egress and reconciles after (`reserveBudget` / `reconcileSpend`), both // over raw SQL. A zero prior total keeps every generation under the cap, // so these suites keep testing what they were written to test. - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: vi.fn(async () => ({ @@ -183,7 +183,7 @@ beforeEach(() => { // `clearAllMocks` does not undo a `mockImplementation`, so the ledger stub is // re-seeded per test — otherwise the over-cap case below leaks into the rest. vi.mocked(prisma.$queryRaw).mockImplementation((async () => [ - { total_tokens: 0 }, + { total_tokens: 0, operator_tokens: 0 }, ]) as never); selfContextText.mockResolvedValue(null); vi.mocked(checkRateLimit).mockResolvedValue({ @@ -254,8 +254,8 @@ describe("POST /api/insights/generate — daily token ceiling", () => { strings: TemplateStringsArray, ) => String(strings[0]).includes("coach_usage") - ? [{ total_tokens: 10_000_000 }] - : [{ total_tokens: 0 }]) as never); + ? [{ total_tokens: 10_000_000, operator_tokens: 10_000_000 }] + : [{ total_tokens: 0, operator_tokens: 0 }]) as never); makeWorkingProvider(); const res = await POST(jsonRequest({ force: true }) as never); diff --git a/src/lib/ai/__tests__/provider-runner.test.ts b/src/lib/ai/__tests__/provider-runner.test.ts index 9514243db..bb5509aed 100644 --- a/src/lib/ai/__tests__/provider-runner.test.ts +++ b/src/lib/ai/__tests__/provider-runner.test.ts @@ -114,6 +114,7 @@ function err(status: number, msg = "boom"): Error & { httpStatus: number } { beforeEach(() => { clearLastWorkingProviderCache(); budgetState.spent = 0; + budgetState.operatorSpent = 0; }); afterEach(() => { @@ -700,6 +701,7 @@ describe("runRawCompletionWithFallback — operator-cost cap at hop time", () => it("refuses the admin-* fallback hop once the operator cap is exhausted", async () => { budgetState.spent = 200_000; // OPERATOR_COST_CAP + budgetState.operatorSpent = 200_000; const codex = new ScriptedProvider({ type: "codex", script: [{ ok: false, error: err(500) }], @@ -730,6 +732,7 @@ describe("runRawCompletionWithFallback — operator-cost cap at hop time", () => it("lets the admin-* fallback run while the operator cap has headroom", async () => { budgetState.spent = 100; + budgetState.operatorSpent = 100; const codex = new ScriptedProvider({ type: "codex", script: [{ ok: false, error: err(500) }], @@ -748,6 +751,63 @@ describe("runRawCompletionWithFallback — operator-cost cap at hop time", () => expect(adminOpenai.callCount).toBe(1); }); + // v1.38.19 (Wave E) — the guard reads the OPERATOR-funded share. + // + // Production evidence (2026-09-11): the operator's day held ~1.2 M tokens, + // nearly all served by `codex` on his own ChatGPT plan. Comparing the day's + // TOTAL against the operator ceiling closed the operator's own fallback hop + // on money he had never spent. + it("lets the admin-* hop run on a big day the user's own plan paid for", async () => { + budgetState.spent = 1_200_000; + budgetState.operatorSpent = 150_000; + const codex = new ScriptedProvider({ + type: "codex", + script: [{ ok: false, error: err(500) }], + }); + const adminOpenai = new ScriptedProvider({ script: [{ ok: true }] }); + + const outcome = await runRawCompletionWithFallback({ + userId: "u-owner-split", + providers: [ + { providerType: "codex", instance: codex }, + { providerType: "admin-openai", instance: adminOpenai }, + ], + params, + }); + expect(outcome.workingProvider.providerType).toBe("admin-openai"); + expect(adminOpenai.callCount).toBe(1); + }); + + it("refuses the admin-* hop once the OPERATOR-funded share is exhausted", async () => { + budgetState.spent = 1_200_000; + budgetState.operatorSpent = 250_000; + const codex = new ScriptedProvider({ + type: "codex", + script: [{ ok: false, error: err(500) }], + }); + const adminOpenai = new ScriptedProvider({ script: [{ ok: true }] }); + + await expect( + runRawCompletionWithFallback({ + userId: "u-owner-exhausted", + providers: [ + { providerType: "codex", instance: codex }, + { providerType: "admin-openai", instance: adminOpenai }, + ], + params, + }), + ).rejects.toMatchObject({ + attempts: [ + expect.objectContaining({ providerType: "codex" }), + expect.objectContaining({ + providerType: "admin-openai", + failureReason: "operator-cost-cap-exhausted", + }), + ], + }); + expect(adminOpenai.callCount).toBe(0); + }); + it("never consults the ledger for a user-funded chain", async () => { budgetState.spent = 999_999_999; const openai = new ScriptedProvider({ script: [{ ok: true }] }); diff --git a/src/lib/ai/coach/__tests__/budget.test.ts b/src/lib/ai/coach/__tests__/budget.test.ts index cb6fbe304..5b8f012c7 100644 --- a/src/lib/ai/coach/__tests__/budget.test.ts +++ b/src/lib/ai/coach/__tests__/budget.test.ts @@ -172,7 +172,7 @@ describe("reserveBudget cap (F1 — user-plan path not locked out)", () => { 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 }, + { total_tokens: priorSpend + 1_200, operator_tokens: priorSpend + 1_200 }, ]); prismaMock.$executeRaw.mockResolvedValue(0); const { reserveBudget } = await import("../budget"); @@ -189,6 +189,82 @@ describe("reserveBudget cap (F1 — user-plan path not locked out)", () => { }); }); +/** + * v1.38.19 (Wave E) — the cap is enforced against the counter it is ABOUT. + * + * Production evidence (2026-09-11): the operator's day held 1.0–1.45 M tokens, + * almost all served by `codex` on his own ChatGPT plan after the shared + * `admin-openai` key answered 500 — and the chat refused at + * `totalAfter: 513539` against the 200 k operator ceiling. Comparing the total + * is the defect; the operator ceiling may only see operator-funded tokens. + */ +describe("reserveBudget cap by cost owner (Wave E)", () => { + 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(); + prismaMock.$executeRaw.mockResolvedValue(0); + }); + + it("admits an operator turn on a day the user's own plan filled", async () => { + // 1.2 M tokens on the day, but only 150 k of them on the operator's key. + prismaMock.$queryRaw.mockResolvedValue([ + { total_tokens: 1_200_000 + 1_200, operator_tokens: 150_000 + 1_200 }, + ]); + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget( + "u", + 1_200, + "2026-09-11", + OPERATOR_COST_CAP, + "operator", + ); + expect(res.allowed).toBe(true); + expect(res.operatorAfter).toBe(151_200); + expect(prismaMock.$executeRaw).not.toHaveBeenCalled(); + }); + + it("refuses an operator turn once the OPERATOR-funded share reaches the cap", async () => { + prismaMock.$queryRaw.mockResolvedValue([ + { + total_tokens: 1_200_000 + 1_200, + operator_tokens: OPERATOR_COST_CAP + 1_200, + }, + ]); + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget( + "u", + 1_200, + "2026-09-11", + OPERATOR_COST_CAP, + "operator", + ); + expect(res.allowed).toBe(false); + expect(prismaMock.$executeRaw).toHaveBeenCalled(); + }); + + it("keeps the user-plan ceiling on the TOTAL, not on the operator share", async () => { + prismaMock.$queryRaw.mockResolvedValue([ + { total_tokens: USER_PLAN_CAP + 1_200, operator_tokens: 0 }, + ]); + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget( + "u", + 1_200, + "2026-09-11", + USER_PLAN_CAP, + "user", + ); + expect(res.allowed).toBe(false); + }); +}); + describe("reconcileSpend cached-token subtraction (F3)", () => { let prismaMock: { $executeRaw: ReturnType }; diff --git a/src/lib/ai/coach/budget.ts b/src/lib/ai/coach/budget.ts index 44f6eab61..ab4c3f506 100644 --- a/src/lib/ai/coach/budget.ts +++ b/src/lib/ai/coach/budget.ts @@ -223,11 +223,17 @@ export async function reserveBudget( // 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. + // PRIOR to its reservation was under the cap. So compare the post-increment + // figure minus this reservation against the cap. + // + // v1.38.19 (Wave E) — and compare the counter the cap is ABOUT. An + // operator-funded chain is measured against the day's operator-funded spend, + // never against a total that the user's own plan inflated; a user-plan chain + // keeps its abuse ceiling on the total. const priorTotal = totalAfter - reserved; const priorOperator = operatorAfter - operatorReserved; - if (priorTotal >= cap) { + const prior = owner === "operator" ? priorOperator : priorTotal; + if (prior >= cap) { // Already over before this request — refund the reservation + the // message-count bump and refuse. await refundReservation(userId, reserved, operatorReserved, dateKey); diff --git a/src/lib/ai/provider-runner.ts b/src/lib/ai/provider-runner.ts index 9b9344f73..e414423fd 100644 --- a/src/lib/ai/provider-runner.ts +++ b/src/lib/ai/provider-runner.ts @@ -435,7 +435,10 @@ async function runRawChain( total: 0, operator: 0, })); - if (spent.total >= OPERATOR_COST_CAP) { + // v1.38.19 (Wave E) — the OPERATOR-funded share, not the day's total: a + // day filled by the user's own plan must not close the operator's + // fallback hop on money the operator never spent. + if (spent.operator >= OPERATOR_COST_CAP) { const hop: FallbackHop = { providerType: candidate.providerType, attempt: i + 1, @@ -448,7 +451,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.total, + operator_cap_spent: spent.operator, }, }); continue; diff --git a/src/lib/insights/__tests__/briefing-provider-budget.test.ts b/src/lib/insights/__tests__/briefing-provider-budget.test.ts index b79a4a7ff..b51df1541 100644 --- a/src/lib/insights/__tests__/briefing-provider-budget.test.ts +++ b/src/lib/insights/__tests__/briefing-provider-budget.test.ts @@ -21,6 +21,8 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; /** In-memory stand-in for the day's `coach_usage.total_tokens`. */ let ledgerTotal = 0; +/** …and of `coach_usage.operator_tokens` (Wave E). */ +let ledgerOperator = 0; vi.mock("@/lib/db", () => ({ prisma: { @@ -28,8 +30,10 @@ vi.mock("@/lib/db", () => ({ // (strings, userId, dateKey, reserved, reserved). $queryRaw: vi.fn( async (_strings: TemplateStringsArray, ...v: unknown[]) => { + // Wave E — (userId, dateKey, reserved, operatorReserved, …). ledgerTotal += Number(v[2] ?? 0); - return [{ total_tokens: ledgerTotal }]; + ledgerOperator += Number(v[3] ?? 0); + return [{ total_tokens: ledgerTotal, operator_tokens: ledgerOperator }]; }, ), // `reconcileSpend` (+delta) and `refundReservation` (-reserved). @@ -37,9 +41,17 @@ vi.mock("@/lib/db", () => ({ async (strings: TemplateStringsArray, ...v: unknown[]) => { const sql = strings.join("?"); const amount = Number(v[0] ?? 0); - if (sql.includes("total_tokens + ")) ledgerTotal += amount; - else if (sql.includes("total_tokens - ")) ledgerTotal -= amount; + // Wave E — the operator-funded share moves in the same statement. + const operatorAmount = Number(v[1] ?? 0); + if (sql.includes("total_tokens + ")) { + ledgerTotal += amount; + ledgerOperator += operatorAmount; + } else if (sql.includes("total_tokens - ")) { + ledgerTotal -= amount; + ledgerOperator -= operatorAmount; + } ledgerTotal = Math.max(0, ledgerTotal); + ledgerOperator = Math.max(0, ledgerOperator); return 1; }, ), @@ -104,6 +116,7 @@ function mockProviderReply( beforeEach(() => { vi.clearAllMocks(); ledgerTotal = 0; + ledgerOperator = 0; }); describe("runBriefingCompletion — ledger accounting", () => { @@ -142,6 +155,7 @@ describe("runBriefingCompletion — ledger accounting", () => { it("refuses a generation once the day's cap is already spent", async () => { ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; mockProviderReply(500); await expect( @@ -201,6 +215,7 @@ describe("runBriefingCompletion — ledger accounting", () => { describe("runBriefingCompletion — cost owner decides the ceiling", () => { it("measures an operator-key chain against the operator ceiling", async () => { ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; mockProviderReply(100); await expect( @@ -214,6 +229,7 @@ describe("runBriefingCompletion — cost owner decides the ceiling", () => { // their own key pays their own bill, so the operator's ceiling is a // category error for them — they must still be served. ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; mockProviderReply(100); const outcome = await runBriefingCompletion( @@ -226,6 +242,9 @@ describe("runBriefingCompletion — cost owner decides the ceiling", () => { it("still bounds a BYOK chain at the user-plan ceiling", async () => { ledgerTotal = USER_PLAN_CAP; + // Not one of those tokens was the operator's: the abuse ceiling on the + // total is what must still refuse this generation. + ledgerOperator = 0; mockProviderReply(100); await expect( diff --git a/src/lib/insights/__tests__/comprehensive-briefing-grounding-disposal.test.ts b/src/lib/insights/__tests__/comprehensive-briefing-grounding-disposal.test.ts index 1ad03b0e3..53a38a769 100644 --- a/src/lib/insights/__tests__/comprehensive-briefing-grounding-disposal.test.ts +++ b/src/lib/insights/__tests__/comprehensive-briefing-grounding-disposal.test.ts @@ -24,7 +24,7 @@ const extractFeatures = vi.fn(); vi.mock("@/lib/db", () => ({ prisma: { - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: (...a: unknown[]) => findUnique(...a), diff --git a/src/lib/insights/__tests__/comprehensive-briefing-reroll.test.ts b/src/lib/insights/__tests__/comprehensive-briefing-reroll.test.ts index 9b8e680b7..7517e6f66 100644 --- a/src/lib/insights/__tests__/comprehensive-briefing-reroll.test.ts +++ b/src/lib/insights/__tests__/comprehensive-briefing-reroll.test.ts @@ -24,7 +24,7 @@ vi.mock("@/lib/db", () => ({ // egress and reconciles after (`reserveBudget` / `reconcileSpend`), both // over raw SQL. A zero prior total keeps every generation under the cap, // so these suites keep testing what they were written to test. - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: (...a: unknown[]) => findUnique(...a), diff --git a/src/lib/insights/__tests__/comprehensive-budget-refusal.test.ts b/src/lib/insights/__tests__/comprehensive-budget-refusal.test.ts index 7b439de55..8d214b0b6 100644 --- a/src/lib/insights/__tests__/comprehensive-budget-refusal.test.ts +++ b/src/lib/insights/__tests__/comprehensive-budget-refusal.test.ts @@ -15,6 +15,8 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; let ledgerTotal = 0; +/** …and of `coach_usage.operator_tokens` (Wave E). */ +let ledgerOperator = 0; const findUnique = vi.fn(); const userUpdate = vi.fn(); @@ -29,17 +31,27 @@ vi.mock("@/lib/db", () => ({ prisma: { $queryRaw: vi.fn( async (_strings: TemplateStringsArray, ...v: unknown[]) => { + // Wave E — (userId, dateKey, reserved, operatorReserved, …). ledgerTotal += Number(v[2] ?? 0); - return [{ total_tokens: ledgerTotal }]; + ledgerOperator += Number(v[3] ?? 0); + return [{ total_tokens: ledgerTotal, operator_tokens: ledgerOperator }]; }, ), $executeRaw: vi.fn( async (strings: TemplateStringsArray, ...v: unknown[]) => { const sql = strings.join("?"); const amount = Number(v[0] ?? 0); - if (sql.includes("total_tokens + ")) ledgerTotal += amount; - else if (sql.includes("total_tokens - ")) ledgerTotal -= amount; + // Wave E — the operator-funded share moves in the same statement. + const operatorAmount = Number(v[1] ?? 0); + if (sql.includes("total_tokens + ")) { + ledgerTotal += amount; + ledgerOperator += operatorAmount; + } else if (sql.includes("total_tokens - ")) { + ledgerTotal -= amount; + ledgerOperator -= operatorAmount; + } ledgerTotal = Math.max(0, ledgerTotal); + ledgerOperator = Math.max(0, ledgerOperator); return 1; }, ), @@ -110,6 +122,7 @@ const VALID = JSON.stringify({ dailyBriefing: { paragraph: "ok" } }); beforeEach(() => { vi.clearAllMocks(); ledgerTotal = 0; + ledgerOperator = 0; resolveProviderChain.mockResolvedValue([ { providerType: "admin-openai", instance: {} }, ]); @@ -145,6 +158,7 @@ describe("generateComprehensiveInsight — daily token ceiling", () => { it("refuses over-cap as `skipped: budget` without contacting a provider", async () => { ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; const outcome = await generateComprehensiveInsight("u1", { locale: "de" }); @@ -166,6 +180,7 @@ describe("generateComprehensiveInsight — daily token ceiling", () => { { providerType: "anthropic", instance: {} }, ]); ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; const outcome = await generateComprehensiveInsight("u1", { locale: "de" }); @@ -178,6 +193,7 @@ describe("generateComprehensiveInsight — daily token ceiling", () => { { providerType: "anthropic", instance: {} }, ]); ledgerTotal = USER_PLAN_CAP; + ledgerOperator = USER_PLAN_CAP; const outcome = await generateComprehensiveInsight("u1", { locale: "de" }); diff --git a/src/lib/insights/__tests__/comprehensive-generate-force.test.ts b/src/lib/insights/__tests__/comprehensive-generate-force.test.ts index 612d1d265..8ab32fb81 100644 --- a/src/lib/insights/__tests__/comprehensive-generate-force.test.ts +++ b/src/lib/insights/__tests__/comprehensive-generate-force.test.ts @@ -31,7 +31,7 @@ vi.mock("@/lib/db", () => ({ // egress and reconciles after (`reserveBudget` / `reconcileSpend`), both // over raw SQL. A zero prior total keeps every generation under the cap, // so these suites keep testing what they were written to test. - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: (...a: unknown[]) => findUnique(...a), diff --git a/src/lib/insights/__tests__/comprehensive-json-retry.test.ts b/src/lib/insights/__tests__/comprehensive-json-retry.test.ts index 7f4a3a8f2..f64cfdc86 100644 --- a/src/lib/insights/__tests__/comprehensive-json-retry.test.ts +++ b/src/lib/insights/__tests__/comprehensive-json-retry.test.ts @@ -22,7 +22,7 @@ vi.mock("@/lib/db", () => ({ // egress and reconciles after (`reserveBudget` / `reconcileSpend`), both // over raw SQL. A zero prior total keeps every generation under the cap, // so these suites keep testing what they were written to test. - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: (...a: unknown[]) => findUnique(...a), diff --git a/src/lib/insights/__tests__/comprehensive-privacy-flip.test.ts b/src/lib/insights/__tests__/comprehensive-privacy-flip.test.ts index 5b30f25e1..4ee3bc274 100644 --- a/src/lib/insights/__tests__/comprehensive-privacy-flip.test.ts +++ b/src/lib/insights/__tests__/comprehensive-privacy-flip.test.ts @@ -25,7 +25,7 @@ const getSelfContextTextForUser = vi.fn(); vi.mock("@/lib/db", () => ({ prisma: { - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: (...a: unknown[]) => findUnique(...a), diff --git a/src/lib/insights/__tests__/comprehensive-timeout-honored.test.ts b/src/lib/insights/__tests__/comprehensive-timeout-honored.test.ts index 2c8692e50..73bc1fc6c 100644 --- a/src/lib/insights/__tests__/comprehensive-timeout-honored.test.ts +++ b/src/lib/insights/__tests__/comprehensive-timeout-honored.test.ts @@ -26,7 +26,7 @@ vi.mock("@/lib/db", () => ({ // egress and reconciles after (`reserveBudget` / `reconcileSpend`), both // over raw SQL. A zero prior total keeps every generation under the cap, // so these suites keep testing what they were written to test. - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: (...a: unknown[]) => findUnique(...a), diff --git a/src/lib/insights/__tests__/comprehensive-warm-timeout.test.ts b/src/lib/insights/__tests__/comprehensive-warm-timeout.test.ts index 138aff034..24ee58f6d 100644 --- a/src/lib/insights/__tests__/comprehensive-warm-timeout.test.ts +++ b/src/lib/insights/__tests__/comprehensive-warm-timeout.test.ts @@ -28,7 +28,7 @@ vi.mock("@/lib/db", () => ({ // egress and reconciles after (`reserveBudget` / `reconcileSpend`), both // over raw SQL. A zero prior total keeps every generation under the cap, // so these suites keep testing what they were written to test. - $queryRaw: vi.fn(async () => [{ total_tokens: 0 }]), + $queryRaw: vi.fn(async () => [{ total_tokens: 0, operator_tokens: 0 }]), $executeRaw: vi.fn(async () => 0), user: { findUnique: (...a: unknown[]) => findUnique(...a), diff --git a/src/lib/insights/__tests__/status-provider-budget.test.ts b/src/lib/insights/__tests__/status-provider-budget.test.ts index 9e5bd760d..0593bdf7e 100644 --- a/src/lib/insights/__tests__/status-provider-budget.test.ts +++ b/src/lib/insights/__tests__/status-provider-budget.test.ts @@ -18,6 +18,8 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; /** In-memory stand-in for the day's `coach_usage.total_tokens`. */ let ledgerTotal = 0; +/** …and of `coach_usage.operator_tokens` (Wave E). */ +let ledgerOperator = 0; vi.mock("@/lib/db", () => ({ prisma: { @@ -25,8 +27,10 @@ vi.mock("@/lib/db", () => ({ // (strings, userId, dateKey, reserved, reserved). $queryRaw: vi.fn( async (_strings: TemplateStringsArray, ...v: unknown[]) => { + // Wave E — (userId, dateKey, reserved, operatorReserved, …). ledgerTotal += Number(v[2] ?? 0); - return [{ total_tokens: ledgerTotal }]; + ledgerOperator += Number(v[3] ?? 0); + return [{ total_tokens: ledgerTotal, operator_tokens: ledgerOperator }]; }, ), // `reconcileSpend` (+delta) and `refundReservation` (-reserved). @@ -34,9 +38,17 @@ vi.mock("@/lib/db", () => ({ async (strings: TemplateStringsArray, ...v: unknown[]) => { const sql = strings.join("?"); const amount = Number(v[0] ?? 0); - if (sql.includes("total_tokens + ")) ledgerTotal += amount; - else if (sql.includes("total_tokens - ")) ledgerTotal -= amount; + // Wave E — the operator-funded share moves in the same statement. + const operatorAmount = Number(v[1] ?? 0); + if (sql.includes("total_tokens + ")) { + ledgerTotal += amount; + ledgerOperator += operatorAmount; + } else if (sql.includes("total_tokens - ")) { + ledgerTotal -= amount; + ledgerOperator -= operatorAmount; + } ledgerTotal = Math.max(0, ledgerTotal); + ledgerOperator = Math.max(0, ledgerOperator); return 1; }, ), @@ -98,6 +110,7 @@ function mockProviderReply(tokensUsed: number | null) { beforeEach(() => { vi.clearAllMocks(); ledgerTotal = 0; + ledgerOperator = 0; resolveProviderChain.mockResolvedValue(OPERATOR_CHAIN); resolveProvider.mockResolvedValue({ type: "none" }); }); @@ -126,6 +139,7 @@ describe("runStatusCompletion — ledger accounting", () => { it("refuses a generation once the day's cap is already spent", async () => { ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; mockProviderReply(500); const result = await runStatusCompletion(completionArgs()); @@ -158,6 +172,7 @@ describe("runStatusCompletion — cost owner decides the ceiling", () => { it("measures an operator-key chain against the operator ceiling", async () => { // Just under the operator ceiling: an operator-key user is refused here. ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; mockProviderReply(100); const result = await runStatusCompletion(completionArgs()); @@ -172,6 +187,7 @@ describe("runStatusCompletion — cost owner decides the ceiling", () => { // category error for them — they must still be served. resolveProviderChain.mockResolvedValue(BYOK_CHAIN); ledgerTotal = OPERATOR_COST_CAP; + ledgerOperator = OPERATOR_COST_CAP; runRawCompletionWithFallback.mockResolvedValue({ result: { content: '{"summary":"ok"}', model: "m", tokensUsed: 100 }, workingProvider: { providerType: "anthropic" }, @@ -186,6 +202,9 @@ describe("runStatusCompletion — cost owner decides the ceiling", () => { it("still bounds a BYOK chain at the user-plan ceiling", async () => { resolveProviderChain.mockResolvedValue(BYOK_CHAIN); ledgerTotal = USER_PLAN_CAP; + // Not one of those tokens was the operator's: the abuse ceiling on the + // total is what must still refuse this generation. + ledgerOperator = 0; mockProviderReply(100); const result = await runStatusCompletion(completionArgs()); diff --git a/tests/integration/coach-budget-owner.test.ts b/tests/integration/coach-budget-owner.test.ts index ab610b06d..95f81f8a5 100644 --- a/tests/integration/coach-budget-owner.test.ts +++ b/tests/integration/coach-budget-owner.test.ts @@ -153,6 +153,66 @@ describe("coach_usage.operator_tokens (real Postgres)", () => { expect(row?.operatorTokens).toBe(2_500); }); + it("admits the operator's chat on a day his own plan filled", async () => { + const { reserveBudget, OPERATOR_COST_CAP } = + await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-09-11"; + + // The operator's 06:42Z lockout, reproduced: 1.2 M tokens on the day, of + // which only 150 k were served by the operator-funded key — the rest by + // `codex` on his own ChatGPT plan. + await getPrismaClient().coachUsage.create({ + data: { + userId, + dateKey, + totalTokens: 1_200_000, + operatorTokens: 150_000, + messageCount: 240, + }, + }); + + const res = await reserveBudget( + userId, + 3_000, + dateKey, + OPERATOR_COST_CAP, + "operator", + ); + expect(res.allowed).toBe(true); + expect(res.operatorAfter).toBe(153_000); + }); + + it("refuses once the operator-funded share reaches the operator cap", 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: 1_200_000, + operatorTokens: OPERATOR_COST_CAP, + messageCount: 240, + }, + }); + + const res = await reserveBudget( + userId, + 3_000, + dateKey, + OPERATOR_COST_CAP, + "operator", + ); + expect(res.allowed).toBe(false); + + const row = await readRow(userId, dateKey); + expect(row?.totalTokens).toBe(1_200_000); + expect(row?.operatorTokens).toBe(OPERATOR_COST_CAP); + }); + it("refunds both counters when a refused reservation is rolled back", async () => { const { reserveBudget, OPERATOR_COST_CAP } = await import("@/lib/ai/coach/budget");