mirror of
https://github.com/MBombeck/HealthLog.git
synced 2026-11-08 12:32:43 +00:00
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.
This commit is contained in:
1 parent
badef28e96
commit
2c46712e13
16 files changed
+285
-26
No files matched your search
@@ -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);
|
||||
|
||||
@@ -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 }] });
|
||||
|
||||
@@ -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<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();
|
||||
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<typeof vi.fn> };
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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" });
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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");
|
||||
|
||||
Reference in new issue
Block a user