merge(OAUTHFIX): provider-aware AI budget

This commit is contained in:
Marc-André Bombeck committed 2026-06-26 12:53:50 +02:00
commit 8471384d01
16 files changed
+271 -14

No files matched your search

+1 -1
View File
@@ -2224,7 +2224,7 @@
"feedbackThanks": "Danke für das Signal.",
"feedbackError": "Feedback konnte nicht gespeichert werden",
"dailyLimitTitle": "Heute-Limit erreicht",
"dailyLimitBody": "Heute-Limit erreicht; Reset um 00:00 UTC.",
"dailyLimitBody": "Tageslimit erreicht. Das Budget setzt sich um Mitternacht (UTC) zurück.",
"providerRateLimitTitle": "Provider rate-limited",
"providerRateLimitBody": "Provider ist temporär überlastet; Reset in ~5 min.",
"cluster": {
+1 -1
View File
@@ -2224,7 +2224,7 @@
"feedbackThanks": "Thanks for the signal.",
"feedbackError": "Could not save feedback",
"dailyLimitTitle": "Daily limit reached",
"dailyLimitBody": "Daily limit reached; resets at 00:00 UTC.",
"dailyLimitBody": "Daily limit reached. The budget refreshes at midnight UTC.",
"providerRateLimitTitle": "Provider rate-limited",
"providerRateLimitBody": "Provider temporarily rate-limited; retry in ~5 min.",
"cluster": {
+1 -1
View File
@@ -2222,7 +2222,7 @@
"feedbackThanks": "Gracias por la indicación.",
"feedbackError": "No se pudo guardar el comentario",
"dailyLimitTitle": "Límite diario alcanzado",
"dailyLimitBody": "Límite diario alcanzado; reinicio a las 00:00 UTC.",
"dailyLimitBody": "Límite diario alcanzado. El cupo se reinicia a medianoche (UTC).",
"providerRateLimitTitle": "Proveedor con límite de tasa",
"providerRateLimitBody": "El proveedor está temporalmente saturado; reinicio en ~5 min.",
"settingsContextLabel": "Enviar contexto",
+1 -1
View File
@@ -2222,7 +2222,7 @@
"feedbackThanks": "Merci pour l’indication.",
"feedbackError": "Impossible d'enregistrer le retour",
"dailyLimitTitle": "Limite quotidienne atteinte",
"dailyLimitBody": "Limite quotidienne atteinte ; réinitialisation à 00:00 UTC.",
"dailyLimitBody": "Limite quotidienne atteinte. Le quota se réinitialise à minuit (UTC).",
"providerRateLimitTitle": "Fournisseur limité en débit",
"providerRateLimitBody": "Le fournisseur est temporairement saturé ; réinitialisation dans ~5 min.",
"settingsContextLabel": "Envoyer le contexte",
+1 -1
View File
@@ -2222,7 +2222,7 @@
"feedbackThanks": "Grazie per il segnale.",
"feedbackError": "Impossibile salvare il feedback",
"dailyLimitTitle": "Limite giornaliero raggiunto",
"dailyLimitBody": "Limite giornaliero raggiunto; reset alle 00:00 UTC.",
"dailyLimitBody": "Limite giornaliero raggiunto. Il budget si reimposta a mezzanotte (UTC).",
"providerRateLimitTitle": "Provider con rate-limit",
"providerRateLimitBody": "Il provider è temporaneamente sovraccarico; reset in ~5 min.",
"settingsContextLabel": "Invia contesto",
+1 -1
View File
@@ -2222,7 +2222,7 @@
"feedbackThanks": "Dzięki za sygnał.",
"feedbackError": "Nie udało się zapisać opinii",
"dailyLimitTitle": "Osiągnięto dzienny limit",
"dailyLimitBody": "Osiągnięto dzienny limit; reset o 00:00 UTC.",
"dailyLimitBody": "Osiągnięto dzienny limit. Limit resetuje się o północy (UTC).",
"providerRateLimitTitle": "Dostawca z ograniczeniem",
"providerRateLimitBody": "Dostawca jest tymczasowo przeciążony; reset za ~5 min.",
"settingsContextLabel": "Wysyłaj kontekst",
@@ -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", () => ({
@@ -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 })),
@@ -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 {
+14
View File
@@ -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.
@@ -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<typeof import("@/lib/labs/ocr-extract")>(
@@ -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,
);
});
+9
View File
@@ -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({
@@ -618,7 +618,9 @@ describe("<MessageThread>", () => {
})}
/>,
);
expect(html).toContain("Daily limit reached; resets at 00:00 UTC.");
expect(html).toContain(
"Daily limit reached. The budget refreshes at midnight UTC.",
);
// v1.4.33 IW7 — fallback copy rewritten ("AI provider" -> "Insights
// provider"). Both old and new strings should be absent here when
// the budget-exceeded error takes precedence.
+137 -1
View File
@@ -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<typeof vi.fn>;
$executeRaw: ReturnType<typeof vi.fn>;
};
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<typeof vi.fn> };
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);
});
});
+71 -2
View File
@@ -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<void> {
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
+9
View File
@@ -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,