From b77a7214745a59e45fb72df6399c20d96a74dc82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Sun, 19 Jul 2026 05:35:13 +0200 Subject: [PATCH] Remove the superseded budget read-then-write pair MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit enforceBudget and recordSpend have no callers left. Both defects they carried recurred every time a surface copied them: a TOCTOU window between the read and the bump, and a cap parameter defaulting to the operator ceiling so an omitted argument rationed a self-hoster paying for their own key. Deleting them means a future surface cannot reach for the older, simpler-looking helper and reintroduce either. getDailyTokenSpend goes with them — it existed to feed the read half. The ledger-poisoning clamps recordSpend carried (NaN, negative, fractional token counts) move to reserveBudget, where the property now lives. Refresh the doc comments that still narrated the old flow. --- .../chat/__tests__/route-module-gate.test.ts | 8 +- src/app/api/insights/chat/route.ts | 5 +- src/lib/ai/coach/__tests__/budget.test.ts | 90 ++++++------------- src/lib/ai/coach/budget.ts | 74 +++------------ src/lib/ai/coach/conversation-summary.ts | 2 +- 5 files changed, 47 insertions(+), 132 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 7b8e75db2..21b9fde3c 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 @@ -93,8 +93,12 @@ vi.mock("@/lib/ai/coach/coach-memory-shared", () => ({ vi.mock("@/lib/ai/coach/facts", () => ({ storeDeterministicFacts: vi.fn() })); vi.mock("@/lib/ai/coach/budget", () => ({ buildDateKey: vi.fn(), - enforceBudget: vi.fn(), - recordSpend: vi.fn(), + reserveBudget: vi.fn(async () => ({ + allowed: true, + reserved: 0, + totalAfter: 0, + })), + reconcileSpend: vi.fn(), resolveDailyCap: vi.fn(() => 200_000), })); vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn() })); diff --git a/src/app/api/insights/chat/route.ts b/src/app/api/insights/chat/route.ts index c8b22d84e..a48a829db 100644 --- a/src/app/api/insights/chat/route.ts +++ b/src/app/api/insights/chat/route.ts @@ -9,8 +9,9 @@ * Behaviour: * 1. requireAuth() — cookie session OR bearer token (iOS app). * 2. Validate body with `coachChatRequestSchema`. - * 3. enforceBudget() — 429 with `coach.budget.exceeded` when the - * user has already burned the day's token cap. + * 3. reserveBudget(resolveDailyCap(chain)) — 429 with + * `coach.budget.exceeded` when the day's token cap for the + * credential that would pay for the call is already burned. * 4. detectRefusal() — pattern-based prompt-injection + * off-topic guard. Refusal emits a single `token` frame with * the localised refusal copy and a `done` frame; never hits a diff --git a/src/lib/ai/coach/__tests__/budget.test.ts b/src/lib/ai/coach/__tests__/budget.test.ts index 63ec53a3f..45697c6b8 100644 --- a/src/lib/ai/coach/__tests__/budget.test.ts +++ b/src/lib/ai/coach/__tests__/budget.test.ts @@ -1,6 +1,5 @@ import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; -import { HttpError } from "@/lib/api-handler"; import { buildDateKey, OPERATOR_COST_CAP, @@ -32,84 +31,49 @@ describe("buildDateKey", () => { }); }); -describe("budget", () => { +/** + * Ledger-poisoning guard. These clamps used to live on `recordSpend`, which is + * gone along with `enforceBudget` — the read-then-write pair that carried both + * a TOCTOU window and an `OPERATOR_COST_CAP` default that rationed a + * self-hoster's own key. The PROPERTY they protected still matters and now + * belongs to the reservation path, so it is asserted there: a provider that + * reports `tokensUsed: NaN` or a negative count must not poison the meter. + */ +describe("reserveBudget — ledger clamps", () => { let prismaMock: { - coachUsage: { - findUnique: ReturnType; - upsert: ReturnType; - }; + $queryRaw: ReturnType; + $executeRaw: ReturnType; }; beforeEach(async () => { const dbModule = await import("@/lib/db"); prismaMock = dbModule.prisma as unknown as typeof prismaMock; - prismaMock.coachUsage.findUnique.mockReset(); - prismaMock.coachUsage.upsert.mockReset(); + prismaMock.$queryRaw.mockReset(); + prismaMock.$queryRaw.mockResolvedValue([{ total_tokens: 0 }]); + prismaMock.$executeRaw.mockReset(); + prismaMock.$executeRaw.mockResolvedValue(0); }); afterEach(() => { vi.restoreAllMocks(); }); - it("returns 0 when no usage row exists", async () => { - prismaMock.coachUsage.findUnique.mockResolvedValue(null); - const { getDailyTokenSpend } = await import("../budget"); - const spent = await getDailyTokenSpend("user-1", "2026-05-10"); - expect(spent).toBe(0); + it("clamps a non-finite reservation to 0", async () => { + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget("u", Number.NaN, "2026-05-10"); + expect(res.reserved).toBe(0); }); - it("returns the persisted token total", async () => { - prismaMock.coachUsage.findUnique.mockResolvedValue({ totalTokens: 4_200 }); - const { getDailyTokenSpend } = await import("../budget"); - const spent = await getDailyTokenSpend("user-1", "2026-05-10"); - expect(spent).toBe(4_200); + it("clamps a negative reservation to 0", async () => { + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget("u", -42, "2026-05-10"); + expect(res.reserved).toBe(0); }); - it("under-budget passes enforceBudget without throwing", async () => { - prismaMock.coachUsage.findUnique.mockResolvedValue({ totalTokens: 5_000 }); - const { enforceBudget } = await import("../budget"); - await expect( - enforceBudget("user-1", "2026-05-10"), - ).resolves.toBeUndefined(); - }); - - it("over-budget throws HttpError(429)", async () => { - prismaMock.coachUsage.findUnique.mockResolvedValue({ - totalTokens: OPERATOR_COST_CAP, - }); - const { enforceBudget } = await import("../budget"); - await expect(enforceBudget("user-1", "2026-05-10")).rejects.toMatchObject({ - statusCode: 429, - message: "coach.budget.exceeded", - }); - await expect(enforceBudget("user-1", "2026-05-10")).rejects.toBeInstanceOf( - HttpError, - ); - }); - - it("recordSpend clamps non-finite tokens to 0", async () => { - prismaMock.coachUsage.upsert.mockResolvedValue({}); - const { recordSpend } = await import("../budget"); - await recordSpend({ userId: "u", tokens: NaN }); - const args = prismaMock.coachUsage.upsert.mock.calls[0][0]; - expect(args.create.totalTokens).toBe(0); - expect(args.update.totalTokens.increment).toBe(0); - }); - - it("recordSpend clamps negative tokens to 0", async () => { - prismaMock.coachUsage.upsert.mockResolvedValue({}); - const { recordSpend } = await import("../budget"); - await recordSpend({ userId: "u", tokens: -42 }); - const args = prismaMock.coachUsage.upsert.mock.calls[0][0]; - expect(args.create.totalTokens).toBe(0); - }); - - it("recordSpend floors fractional tokens", async () => { - prismaMock.coachUsage.upsert.mockResolvedValue({}); - const { recordSpend } = await import("../budget"); - await recordSpend({ userId: "u", tokens: 12.7 }); - const args = prismaMock.coachUsage.upsert.mock.calls[0][0]; - expect(args.create.totalTokens).toBe(12); + it("floors a fractional reservation", async () => { + const { reserveBudget } = await import("../budget"); + const res = await reserveBudget("u", 12.7, "2026-05-10"); + expect(res.reserved).toBe(12); }); }); diff --git a/src/lib/ai/coach/budget.ts b/src/lib/ai/coach/budget.ts index 875c169f6..21aa7d2c5 100644 --- a/src/lib/ai/coach/budget.ts +++ b/src/lib/ai/coach/budget.ts @@ -15,7 +15,6 @@ * the user's timezone. */ import { prisma } from "@/lib/db"; -import { HttpError } from "@/lib/api-handler"; import type { ProviderChainType } from "@/lib/ai/provider-chain"; /** @@ -82,37 +81,6 @@ export function buildDateKey(at: Date = new Date()): string { return at.toISOString().slice(0, 10); } -/** - * Read the current day's token spend for `userId`. Returns 0 when no - * row exists yet (first request of the day). - */ -export async function getDailyTokenSpend( - userId: string, - dateKey: string = buildDateKey(), -): Promise { - const row = await prisma.coachUsage.findUnique({ - where: { userId_dateKey: { userId, dateKey } }, - select: { totalTokens: true }, - }); - return row?.totalTokens ?? 0; -} - -/** - * Throw 429 when the user has already burned the day's budget. Called - * BEFORE the provider chain runs — the route emits a refusal SSE - * frame instead of hitting any upstream LLM. - */ -export async function enforceBudget( - userId: string, - dateKey: string = buildDateKey(), - cap: number = OPERATOR_COST_CAP, -): Promise { - const spent = await getDailyTokenSpend(userId, dateKey); - if (spent >= cap) { - throw new HttpError(429, "coach.budget.exceeded"); - } -} - /** * v1.18.7 (SENIOR-DEV HIGH) — atomically reserve budget BEFORE the provider * call, closing the read-then-write TOCTOU window. @@ -240,37 +208,15 @@ async function refundReservation( `; } -export interface RecordSpendParams { - userId: string; - tokens: number; - dateKey?: string; -} - /** - * Bump the day's `totalTokens` and `messageCount` after a successful - * assistant reply. Upsert keeps the first request of the day cheap; - * subsequent calls land on the unique-index path. - * - * Negative or non-finite token figures are clamped to zero so a - * provider that returns `tokensUsed: NaN` cannot poison the ledger. + * The `enforceBudget` (read-then-write check) + `recordSpend` (post-hoc bump) + * pair that preceded the reservation model is deliberately GONE, not + * deprecated. It carried two defects that recurred every time a new surface + * copied it: a TOCTOU window between the read and the write, and a `cap` + * parameter that defaulted to `OPERATOR_COST_CAP`, so any caller that omitted + * it rationed a self-hoster's own key by the operator's ceiling. Removing the + * functions means a future surface cannot reintroduce either defect by + * reaching for the older, simpler-looking helper — + * `reserveBudget` / `reconcileSpend` with an explicit `resolveDailyCap(chain)` + * is the only path left. */ -export async function recordSpend(params: RecordSpendParams): Promise { - const dateKey = params.dateKey ?? buildDateKey(); - const tokens = - Number.isFinite(params.tokens) && params.tokens > 0 - ? Math.floor(params.tokens) - : 0; - await prisma.coachUsage.upsert({ - where: { userId_dateKey: { userId: params.userId, dateKey } }, - create: { - userId: params.userId, - dateKey, - totalTokens: tokens, - messageCount: 1, - }, - update: { - totalTokens: { increment: tokens }, - messageCount: { increment: 1 }, - }, - }); -} diff --git a/src/lib/ai/coach/conversation-summary.ts b/src/lib/ai/coach/conversation-summary.ts index a9e0f90f1..10ea0620f 100644 --- a/src/lib/ai/coach/conversation-summary.ts +++ b/src/lib/ai/coach/conversation-summary.ts @@ -15,7 +15,7 @@ * - `annotate` carries counts + ids only, never the summary content. * * Background generation is intentionally UNBILLED — like the period-narrative - * warm, it does not call `recordSpend`. The operator's per-user daily token + * warm, it books no spend of its own. The operator's per-user daily token * ceiling is the cost backstop; the interactive reply meter is not charged for * background memory upkeep. *