diff --git a/src/app/api/insights/chat/route.ts b/src/app/api/insights/chat/route.ts index 09ecf7d5d..89ead0497 100644 --- a/src/app/api/insights/chat/route.ts +++ b/src/app/api/insights/chat/route.ts @@ -62,8 +62,8 @@ import { enqueueCoachMemoryRefresh } from "@/lib/ai/coach/coach-memory-shared"; import { storeDeterministicFacts } from "@/lib/ai/coach/facts"; import { buildDateKey, - enforceBudget, - recordSpend, + reserveBudget, + reconcileSpend, } from "@/lib/ai/coach/budget"; import { detectRefusal } from "@/lib/ai/coach/refusal"; import { getCoachSystemPrompt } from "@/lib/ai/coach/system-prompt"; @@ -215,7 +215,10 @@ async function handleChatRequest(request: NextRequest): Promise { return apiError("Too many Coach requests, please wait a moment", 429); } - await enforceBudget(userId); + // v1.18.7 (SENIOR-DEV HIGH) — the daily token cap is enforced atomically + // by reserving budget BEFORE the provider call (see the reservation right + // before `runRawCompletionWithFallback`), not by a read-then-write here. + // The old read-before-call gate let concurrent requests all pass the cap. const locale = await resolveServerLocale({ request, @@ -431,6 +434,27 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`; // ConsentRequiredError → apiHandler returns 403 + `consent.ai.required`. await assertConsentForChain({ userId, chain, surface: "coach" }); + // v1.18.7 (SENIOR-DEV HIGH) — atomically RESERVE the day's budget before + // the provider call. The reservation increments the day's total by the + // per-call ceiling (`maxTokens`) in one upsert and returns the new total; + // concurrent requests serialise on the row so they cannot all pass the cap. + // Over-cap → 429 refusal frame, reservation already refunded. The actual + // token count is reconciled against this reservation after the call, + // including on empty / sentinel replies (their tokens were still burned). + const reqDateKey = buildDateKey(); + const reservation = await reserveBudget( + userId, + AI_BUDGETS.coach.maxTokens, + reqDateKey, + ); + if (!reservation.allowed) { + annotate({ + action: { name: "coach.budget.exceeded" }, + meta: { totalAfter: reservation.totalAfter }, + }); + return streamProviderError({ code: "coach.budget.exceeded" }); + } + let result: CompletionResult; let workingProviderType: string; try { @@ -447,6 +471,11 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`; result = fallback.result; workingProviderType = fallback.workingProvider.providerType; } 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( + () => {}, + ); if (err instanceof AllProvidersFailedError) { annotate({ action: { name: "insights.coach.providerFailed" }, @@ -484,6 +513,21 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`; throw err; } + // v1.18.7 (SENIOR-DEV MEDIUM) — the provider call returned, so its tokens + // were billed regardless of reply quality. Reconcile the reservation + // against the actual count NOW, before any empty/sentinel short-circuit, so + // an empty or sentinel-only reply still records its burned cost (the old + // post-hoc `recordSpend` ran only on the happy path, undercounting these). + await reconcileSpend( + userId, + reservation.reserved, + result.tokensUsed ?? 0, + reqDateKey, + ).catch(() => { + // Ledger reconcile is best-effort; a failure leaves the conservative + // reservation in place (never an undercount) and never breaks the turn. + }); + const rawReply = (result.content ?? "").trim(); if (!rawReply) { return streamProviderError({ code: "coach.provider.empty" }); @@ -611,13 +655,10 @@ Reply now as the assistant, in ${locale === "de" ? "German" : "English"}.`; promptVersion: PROMPT_VERSION, }); - // Bump the day's spend ledger AFTER persistence so a retried request - // doesn't double-count when the persistence layer rolled back. - await recordSpend({ - userId, - tokens: result.tokensUsed ?? 0, - dateKey: buildDateKey(), - }); + // v1.18.7 — the day's spend was already reconciled against the + // reservation immediately after the provider returned (above), so there is + // no post-persistence ledger bump here. The reservation guarantees the + // tokens are counted even if persistence or streaming fails afterwards. annotate({ action: { name: "insights.coach.replied" }, diff --git a/src/lib/ai/coach/budget.ts b/src/lib/ai/coach/budget.ts index 897056628..02df866e7 100644 --- a/src/lib/ai/coach/budget.ts +++ b/src/lib/ai/coach/budget.ts @@ -61,6 +61,119 @@ export async function enforceBudget( } } +/** + * v1.18.7 (SENIOR-DEV HIGH) — atomically reserve budget BEFORE the provider + * call, closing the read-then-write TOCTOU window. + * + * `enforceBudget` reads spend, then `recordSpend` bumps it only AFTER a + * successful reply, so up to `rate-limit` (20/min) concurrent requests could + * each read `spent < cap` before any spend landed and all hit the provider. + * This instead does a single atomic SQL upsert that increments the day's + * total by an ESTIMATED reservation and returns the new total. The caller + * checks the returned total against the cap: if the reservation pushed it + * over, the request is refused and the reservation refunded; otherwise the + * call proceeds and the actual token count is reconciled afterwards. Mirrors + * the `rate-limit.ts` atomic-upsert pattern — multi-instance correctness is + * structural, not a hopeful read-then-write. + * + * The first request of a fresh day still lands cheaply via the upsert + * `create` branch; `messageCount` is bumped at reservation time so a + * reserved-but-failed turn still counts as an attempt. + */ +export interface ReserveBudgetResult { + /** True when the reservation kept the day's total 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; +} + +export async function reserveBudget( + userId: string, + estimatedTokens: number, + dateKey: string = buildDateKey(), + cap: number = MAX_TOKENS_PER_USER_PER_DAY, +): Promise { + const reserved = + Number.isFinite(estimatedTokens) && estimatedTokens > 0 + ? Math.floor(estimatedTokens) + : 0; + + // Single atomic upsert-increment returning the new total. 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()) + ON CONFLICT (user_id, date_key) DO UPDATE SET + total_tokens = coach_usage.total_tokens + ${reserved}, + message_count = coach_usage.message_count + 1, + updated_at = NOW() + RETURNING total_tokens + `; + const totalAfter = Number(rows[0]?.total_tokens ?? reserved); + + // 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; + 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 }; + } + + return { allowed: true, reserved, totalAfter }; +} + +/** + * Reconcile a reservation against the actual tokens the provider reported. + * `actual - reserved` is applied as a signed delta (clamped so the row never + * goes negative). Called after every provider call — including empty / + * sentinel / refusal replies, whose upstream tokens were still burned + * (SENIOR-DEV MEDIUM: spend undercount). + */ +export async function reconcileSpend( + userId: string, + reserved: number, + actualTokens: number, + dateKey: string = buildDateKey(), +): Promise { + const actual = + Number.isFinite(actualTokens) && actualTokens > 0 + ? Math.floor(actualTokens) + : 0; + const delta = actual - reserved; + if (delta === 0) return; + // Clamp at zero so a smaller-than-reserved actual can't drive the row + // negative under a racing reconcile. + await prisma.$executeRaw` + UPDATE coach_usage + SET total_tokens = GREATEST(0, total_tokens + ${delta}), + updated_at = NOW() + WHERE user_id = ${userId} AND date_key = ${dateKey} + `; +} + +/** Refund a reservation (tokens + the message-count bump) on a refusal. */ +async function refundReservation( + userId: string, + reserved: number, + dateKey: string, +): Promise { + await prisma.$executeRaw` + UPDATE coach_usage + SET total_tokens = GREATEST(0, total_tokens - ${reserved}), + message_count = GREATEST(0, message_count - 1), + updated_at = NOW() + WHERE user_id = ${userId} AND date_key = ${dateKey} + `; +} + export interface RecordSpendParams { userId: string; tokens: number; diff --git a/tests/integration/coach-budget-reservation.test.ts b/tests/integration/coach-budget-reservation.test.ts new file mode 100644 index 000000000..86a7dcbd5 --- /dev/null +++ b/tests/integration/coach-budget-reservation.test.ts @@ -0,0 +1,109 @@ +/** + * v1.18.7 (SENIOR-DEV HIGH + MEDIUM) — atomic Coach budget reservation. + * + * `reserveBudget` does a single atomic upsert-increment so concurrent + * requests cannot all pass the daily cap (the old read-before-call gate + * could). `reconcileSpend` then adjusts the reservation to the actual token + * count — including on empty / sentinel replies whose tokens were still + * burned (the prior post-hoc `recordSpend` undercounted these). These tests + * prove both contracts against a real Postgres; a mocked client could not + * detect a missing atomic increment. + */ +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-user", email: "budget@example.test", role: "USER" }, + }); + return user.id; +} + +beforeEach(async () => { + await truncateAllTables(getPrismaClient()); +}); + +describe("reserveBudget (real Postgres)", () => { + it("lets at most one of two concurrent requests push past the cap", async () => { + const { reserveBudget } = await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-06-19"; + const cap = 1000; + + // Pre-spend 900 so a single 600-token reservation stays under the cap + // but two concurrent ones cannot both be admitted. + await getPrismaClient().coachUsage.create({ + data: { userId, dateKey, totalTokens: 900, messageCount: 1 }, + }); + + const [a, b] = await Promise.all([ + reserveBudget(userId, 600, dateKey, cap), + reserveBudget(userId, 600, dateKey, cap), + ]); + + // Exactly one is admitted: the first to commit sees prior 900 (< cap) and + // reserves to 1500; the second sees prior 1500 (>= cap) and is refused + + // refunded. Order is non-deterministic, so assert the SET of outcomes. + const allowed = [a, b].filter((r) => r.allowed).length; + expect(allowed).toBe(1); + + // The refused reservation was refunded, so the row reflects exactly one + // admitted 600-token reservation on top of the 900 baseline. + const row = await getPrismaClient().coachUsage.findUnique({ + where: { userId_dateKey: { userId, dateKey } }, + }); + expect(row?.totalTokens).toBe(1500); + }); + + it("admits the first request of a fresh day via the upsert create branch", async () => { + const { reserveBudget } = await import("@/lib/ai/coach/budget"); + const userId = await seedUser(); + const dateKey = "2026-06-20"; + + const res = await reserveBudget(userId, 600, dateKey, 25_000); + expect(res.allowed).toBe(true); + + const row = await getPrismaClient().coachUsage.findUnique({ + where: { userId_dateKey: { userId, dateKey } }, + }); + expect(row?.totalTokens).toBe(600); + expect(row?.messageCount).toBe(1); + }); + + it("reconciles the reservation down to a smaller actual count", async () => { + const { reserveBudget, reconcileSpend } = await import( + "@/lib/ai/coach/budget" + ); + 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 row = await getPrismaClient().coachUsage.findUnique({ + where: { userId_dateKey: { userId, dateKey } }, + }); + // 600 reserved → reconciled to the actual 120 burned. + expect(row?.totalTokens).toBe(120); + }); + + it("still records burned tokens on an empty/sentinel reply (no undercount)", async () => { + const { reserveBudget, reconcileSpend } = await import( + "@/lib/ai/coach/budget" + ); + const userId = await seedUser(); + const dateKey = "2026-06-22"; + + // 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 row = await getPrismaClient().coachUsage.findUnique({ + where: { userId_dateKey: { userId, dateKey } }, + }); + expect(row?.totalTokens).toBe(450); + expect(row?.totalTokens).toBeGreaterThan(0); + }); +}); diff --git a/tests/integration/coach-chat.test.ts b/tests/integration/coach-chat.test.ts index 8ea8e62e5..508b4e78d 100644 --- a/tests/integration/coach-chat.test.ts +++ b/tests/integration/coach-chat.test.ts @@ -214,7 +214,12 @@ describe("POST /api/insights/chat — integration", () => { expect(usage?.messageCount).toBe(1); }); - it("returns 429 when daily token budget is exhausted", async () => { + it("refuses with a budget-exceeded error frame when the daily cap is spent", async () => { + // v1.18.7 — the cap is now enforced by an atomic reservation right + // before the provider call. An over-budget request surfaces the + // structured `coach.budget.exceeded` code as an SSE `error` frame + // (HTTP 200, like every other coach error) rather than a JSON 429, and + // the reservation is refunded so the ledger stays at the cap. const { userId } = await seedUserWithSession(); const prisma = getPrismaClient(); const dateKey = new Date().toISOString().slice(0, 10); @@ -237,10 +242,16 @@ describe("POST /api/insights/chat — integration", () => { }), }), ); - expect(res.status).toBe(429); - const env = (await res.json()) as { error: string }; - expect(env.error).toBe("coach.budget.exceeded"); + expect(res.status).toBe(200); + const body = await readStream(res); + expect(body).toContain("coach.budget.exceeded"); expect(runProviderMock).not.toHaveBeenCalled(); + + // The refused reservation was refunded — the ledger holds at the cap. + const usage = await prisma.coachUsage.findUnique({ + where: { userId_dateKey: { userId, dateKey } }, + }); + expect(usage?.totalTokens).toBe(MAX_TOKENS_PER_USER_PER_DAY); }); it("returns 404 (not 403) when conversationId belongs to another user", async () => {