tom/HealthLog
1
0 Fork 0
mirror of https://github.com/MBombeck/HealthLog.git synced 2026-11-09 13:01:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ai): the daily ledger records which tokens the operator paid for

Browse Source 
`coach_usage` counted one number per (user, day). A turn served by the
user's own ChatGPT plan after the operator's shared key answered 500 was
therefore charged to the operator-cost cap, and the operator's chat was
locked out every morning by background work he pays for himself.

The ledger gains `operator_tokens`: the share of `total_tokens` served by
an operator-funded provider. `reserveBudget` books a reservation to the
owner of the chain's primary and `reconcileSpend` settles it against the
provider that actually served, both inside the same statements as the
total so concurrent requests can never see one counter without the other.
Every call site names its cost owner and its serving hop; the argument is
required, so a new surface cannot inherit the old blind accounting.
This commit is contained in:
Marc-André Bombeck committed 2026-09-11 09:42:14 +02:00
1 parent 274923a060
commit badef28e96
53 files changed
+813 -117
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

prisma/migrations/0341_coach_usage_operator_tokens/migration.sql
+15
View File
@@ -0,0 +1,15 @@
-- v1.38.19 (Wave E) — the daily coach ledger gains a cost-owner dimension.
--
-- `coach_usage` counted one number per (user, day), so tokens the USER's own
-- plan paid for (a `codex` fallback hop after the operator's shared key
-- answered 500) were charged against the operator-cost cap and locked the
-- interactive chat out every morning. `operator_tokens` is the share of
-- `total_tokens` served by an operator-funded provider (`admin-openai` /
-- `admin-codex`); the operator cap is enforced against it, the user-plan cap
-- against `total_tokens`.
--
-- Existing rows start at 0: their owner split is unknown and 0 is the side
-- that cannot lock a user out of a chat the operator never paid for. The
-- counter is exact from the first turn after deploy.
ALTER TABLE "coach_usage"
ADD COLUMN IF NOT EXISTS "operator_tokens" INTEGER NOT NULL DEFAULT 0;
prisma/schema.prisma
+13 -7
View File
@@ -6267,13 +6267,19 @@ model HealthProfileFactRevision {
// secondary signal used by the frontend to render a "≈ N replies left"
// hint without exposing raw token figures.
model CoachUsage {
id String @id @default(cuid())
userId String @map("user_id")
dateKey String @map("date_key")
totalTokens Int @default(0) @map("total_tokens")
messageCount Int @default(0) @map("message_count")
updatedAt DateTime @updatedAt @map("updated_at")
createdAt DateTime @default(now()) @map("created_at")
id String @id @default(cuid())
userId String @map("user_id")
dateKey String @map("date_key")
totalTokens Int @default(0) @map("total_tokens")
/// v1.38.19 (Wave E) — the share of `totalTokens` served by a provider the
/// OPERATOR funds (`admin-openai` / `admin-codex`). The operator-cost cap is
/// enforced against this counter and the user-plan cap against
/// `totalTokens`, so background work served by the user's own plan can no
/// longer lock the interactive chat out of the operator's ceiling.
operatorTokens Int @default(0) @map("operator_tokens")
messageCount Int @default(0) @map("message_count")
updatedAt DateTime @updatedAt @map("updated_at")
createdAt DateTime @default(now()) @map("created_at")
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
src/app/api/ai/test/__tests__/route.test.ts
+5
View File
@@ -30,9 +30,12 @@ vi.mock("@/lib/ai/coach/budget", () => ({
allowed: true,
reserved: 32,
totalAfter: 32,
owner: "operator" as const,
operatorAfter: 32,
})),
reconcileSpend: vi.fn(async () => undefined),
resolveDailyCap: () => 200_000,
resolveCostOwner: () => "operator" as const,
}));
vi.mock("@/lib/logging/context", () => ({
@@ -376,6 +379,8 @@ describe("POST /api/ai/test — daily ceiling + ledger", () => {
32,
expect.any(Number),
"2026-01-01",
0,
{ servedBy: expect.any(String), reservedOwner: "operator" },
);
});
});
src/app/api/ai/test/route.ts
+15 -1
View File
@@ -19,6 +19,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { aiTestOverrideSchema } from "@/lib/validations/ai-provider";
@@ -113,6 +114,11 @@ export const POST = apiHandler(async (request: NextRequest) => {
providerType: provider.type === "admin-key" ? "admin-openai" : "local",
},
]),
resolveCostOwner([
{
providerType: provider.type === "admin-key" ? "admin-openai" : "local",
},
]),
);
if (!reservation.allowed) {
annotate({ action: { name: "ai.test.budget_exceeded" } });
@@ -137,6 +143,11 @@ export const POST = apiHandler(async (request: NextRequest) => {
reservation.reserved,
result.tokensUsed ?? reservation.reserved,
dateKey,
0,
{
servedBy: provider.type === "admin-key" ? "admin-openai" : "local",
reservedOwner: reservation.owner,
},
);
return apiSuccess({
@@ -155,7 +166,10 @@ export const POST = apiHandler(async (request: NextRequest) => {
// A failed probe produced no reported usage — refund the reservation so a
// provider that is simply misconfigured doesn't burn the caller's ledger
// while they fix it.
await reconcileSpend(user.id, reservation.reserved, 0, dateKey);
await reconcileSpend(user.id, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
annotate({
meta: {
ai_test_error: err.message.slice(0, 500),
src/app/api/documents/inbound/[id]/chat/__tests__/route.test.ts
+1
View File
@@ -51,6 +51,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 600 }),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 100000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/provider-runner", () => ({
runStreamingRawCompletionWithFallback: vi.fn(),
src/app/api/documents/inbound/[id]/extract/__tests__/route.test.ts
+1
View File
@@ -69,6 +69,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/documents/extract", () => ({
runInboundExtraction: vi.fn(),
src/app/api/documents/inbound/[id]/extract/route.ts
+22 -3
View File
@@ -41,6 +41,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { auditLog } from "@/lib/auth/audit";
@@ -254,6 +255,7 @@ async function handleTextExtract(
AI_BUDGETS.ocrExtractText.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -273,6 +275,8 @@ async function handleTextExtract(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
const updated = await stageExtraction(document.id, userId, result);
@@ -289,7 +293,10 @@ async function handleTextExtract(
return apiSuccess(serialiseDocumentDetail(updated, updated.facts));
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (err instanceof InboundExtractError) {
return apiError("Couldn't read the document. Try a clearer copy.", 422, {
errorCode: "documents.inbound.extractFailed",
@@ -351,6 +358,7 @@ async function handleStoredExtract(
AI_BUDGETS.ocrExtractText.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -370,6 +378,8 @@ async function handleStoredExtract(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
const updated = await stageExtraction(document.id, userId, result);
@@ -386,7 +396,10 @@ async function handleStoredExtract(
return apiSuccess(serialiseDocumentDetail(updated, updated.facts));
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (err instanceof InboundExtractError) {
return apiError("Couldn't read the document. Try a clearer copy.", 422, {
errorCode: "documents.inbound.extractFailed",
@@ -465,6 +478,7 @@ async function handleVisionExtract(
AI_BUDGETS.ocrExtract.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -493,6 +507,8 @@ async function handleVisionExtract(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
const updated = await stageExtraction(document.id, userId, result);
@@ -509,7 +525,10 @@ async function handleVisionExtract(
return apiSuccess(serialiseDocumentDetail(updated, updated.facts));
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (err instanceof InboundExtractError) {
return apiError("Couldn't read the document. Try a clearer copy.", 422, {
errorCode: "documents.inbound.extractFailed",
src/app/api/documents/inbound/[id]/index/__tests__/route.test.ts
+1
View File
@@ -52,6 +52,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/modules/gate", () => ({
requireModuleEnabled: vi.fn().mockResolvedValue({ enabled: true }),
src/app/api/documents/inbound/[id]/index/route.ts
+8 -1
View File
@@ -37,6 +37,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { auditLog } from "@/lib/auth/audit";
@@ -249,6 +250,7 @@ async function handleVisionIndex(
AI_BUDGETS.documentTranscribe.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -269,6 +271,8 @@ async function handleVisionIndex(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
// Refs #776 — the empty-transcription guard, same contract as the auto
// path (`tryProviderIndex`): a provider answer with no text must never
@@ -295,7 +299,10 @@ async function handleVisionIndex(
});
return finishIndex(request, userId, document.id, "vision", tokenCount);
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
await recordIndexAttempt(userId, document.id, {
indexed: false,
reason: "provider-error",
src/app/api/documents/inbound/[id]/suggest/__tests__/route.test.ts
+1
View File
@@ -40,6 +40,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/modules/gate", () => ({
requireModuleEnabled: vi.fn().mockResolvedValue({ enabled: true }),
src/app/api/documents/inbound/[id]/suggest/route.ts
+15 -2
View File
@@ -29,6 +29,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { auditLog } from "@/lib/auth/audit";
@@ -166,6 +167,7 @@ async function handleTextSuggest(
AI_BUDGETS.documentAssist.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -185,10 +187,15 @@ async function handleTextSuggest(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
return finishSuggest(request, userId, document.id, "text", suggestion);
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (err instanceof DocumentAssistError) {
return apiError("Couldn't read the document. Try a clearer copy.", 422, {
errorCode: "documents.inbound.extractFailed",
@@ -255,6 +262,7 @@ async function handleVisionSuggest(
AI_BUDGETS.documentAssist.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -275,10 +283,15 @@ async function handleVisionSuggest(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
return finishSuggest(request, userId, document.id, "vision", suggestion);
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (err instanceof DocumentAssistError) {
return apiError("Couldn't read the document. Try a clearer copy.", 422, {
errorCode: "documents.inbound.extractFailed",
src/app/api/documents/inbound/[id]/summary/__tests__/route.test.ts
+1
View File
@@ -53,6 +53,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 1 }),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/modules/gate", () => ({
requireModuleEnabled: vi.fn().mockResolvedValue({ enabled: true }),
src/app/api/documents/inbound/[id]/summary/route.ts
+15 -2
View File
@@ -36,6 +36,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { auditLog } from "@/lib/auth/audit";
@@ -345,6 +346,7 @@ async function handleTextSummary(
budgetFor(mode),
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -369,10 +371,15 @@ async function handleTextSummary(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
return finishSummary(request, userId, document.id, "text", mode, result);
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return summaryError(err, "text");
}
}
@@ -430,6 +437,7 @@ async function handleVisionSummary(
budgetFor(mode),
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
await refundDocumentAiSlot(userId);
@@ -455,10 +463,15 @@ async function handleVisionSummary(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
return finishSummary(request, userId, document.id, "vision", mode, result);
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return summaryError(err, "vision");
}
}
src/app/api/insights/chat/__tests__/route-cancelled-turn.test.ts
+21 -3
View File
@@ -93,7 +93,11 @@ vi.mock("@/lib/ai/coach/facts", () => ({
}));
const { reserveBudget, reconcileSpend } = vi.hoisted(() => ({
reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 3000 })),
reserveBudget: vi.fn(async () => ({
allowed: true,
reserved: 3000,
owner: "user" as const,
})),
reconcileSpend: vi.fn(async () => undefined),
}));
vi.mock("@/lib/ai/coach/budget", () => ({
@@ -101,6 +105,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget,
reconcileSpend,
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
vi.mock("@/lib/ai/coach/refusal", () => ({
@@ -205,7 +210,11 @@ function abortError(): Error {
describe("coach chat — client abort mid-generation (#781)", () => {
beforeEach(() => {
vi.clearAllMocks();
reserveBudget.mockResolvedValue({ allowed: true, reserved: 3000 });
reserveBudget.mockResolvedValue({
allowed: true,
reserved: 3000,
owner: "user" as const,
});
// No-tools chain so the turn runs the streaming runner (which rejects
// with the abort the torn-down connection produces).
resolveProviderChain.mockResolvedValue([
@@ -236,7 +245,16 @@ describe("coach chat — client abort mid-generation (#781)", () => {
await post(abortedChatReq({ message: "How is my BP?" }));
await sse.done;
expect(reconcileSpend).toHaveBeenCalledWith("u1", 3000, 0, "2026-08-13");
expect(reconcileSpend).toHaveBeenCalledWith(
"u1",
3000,
0,
"2026-08-13",
0,
// No hop served the cancelled turn, so nothing is attributed to an
// owner — the reservation is reversed on both counters.
{ servedBy: null, reservedOwner: "user" },
);
const actions = annotate.mock.calls
.map((c) => (c[0] as { action?: { name?: string } }).action?.name)
.filter(Boolean);
src/app/api/insights/chat/__tests__/route-fence-guard.test.ts
+1
View File
@@ -86,6 +86,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 1500 })),
reconcileSpend: vi.fn(async () => undefined),
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
vi.mock("@/lib/ai/coach/refusal", () => ({
detectRefusal: vi.fn(() => ({ refuse: false })),
src/app/api/insights/chat/__tests__/route-guard-activation.test.ts
+1
View File
@@ -134,6 +134,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget,
reconcileSpend,
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
const { detectRefusal } = vi.hoisted(() => ({
src/app/api/insights/chat/__tests__/route-guard-ledger.test.ts
+1
View File
@@ -127,6 +127,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget,
reconcileSpend,
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
const { detectRefusal } = vi.hoisted(() => ({
src/app/api/insights/chat/__tests__/route-guard.test.ts
+1
View File
@@ -109,6 +109,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget,
reconcileSpend,
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
const { detectRefusal } = vi.hoisted(() => ({
src/app/api/insights/chat/__tests__/route-list-query.test.ts
+1
View File
@@ -101,6 +101,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(),
reconcileSpend: vi.fn(),
resolveDailyCap: vi.fn(() => 200_000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn() }));
vi.mock("@/lib/ai/coach/system-prompt", () => ({
src/app/api/insights/chat/__tests__/route-module-gate.test.ts
+1
View File
@@ -100,6 +100,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
})),
reconcileSpend: vi.fn(),
resolveDailyCap: vi.fn(() => 200_000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/coach/refusal", () => ({ detectRefusal: vi.fn() }));
vi.mock("@/lib/ai/coach/system-prompt", () => ({
src/app/api/insights/chat/__tests__/route-snapshot-once.test.ts
+1
View File
@@ -103,6 +103,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 1500 })),
reconcileSpend: vi.fn(async () => undefined),
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
vi.mock("@/lib/ai/coach/refusal", () => ({
detectRefusal: vi.fn(() => ({ refuse: false })),
src/app/api/insights/chat/__tests__/route-tool-mode.test.ts
+12 -2
View File
@@ -123,8 +123,9 @@ const { reserveBudget, reconcileSpend } = vi.hoisted(() => ({
allowed: boolean;
reserved: number;
totalAfter?: number;
owner?: "operator" | "user";
}>
>(async () => ({ allowed: true, reserved: 3000 })),
>(async () => ({ allowed: true, reserved: 3000, owner: "user" as const })),
reconcileSpend: vi.fn(async () => undefined),
}));
vi.mock("@/lib/ai/coach/budget", () => ({
@@ -132,6 +133,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget,
reconcileSpend,
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
const { detectRefusal } = vi.hoisted(() => ({
@@ -266,7 +268,11 @@ function chatReq(body: Record<string, unknown>): Request {
describe("coach chat — tool-mode routing (F1)", () => {
beforeEach(() => {
vi.clearAllMocks();
reserveBudget.mockResolvedValue({ allowed: true, reserved: 3000 });
reserveBudget.mockResolvedValue({
allowed: true,
reserved: 3000,
owner: "user" as const,
});
detectRefusal.mockReturnValue({ refuse: false });
checkRateLimit.mockResolvedValue({ allowed: true });
assertConsentForChain.mockResolvedValue(undefined);
@@ -318,6 +324,7 @@ describe("coach chat — tool-mode routing (F1)", () => {
4500,
"2026-06-21",
2_000_000,
"user",
);
});
@@ -347,6 +354,8 @@ describe("coach chat — tool-mode routing (F1)", () => {
3000,
0,
expect.anything(),
0,
{ servedBy: null, reservedOwner: "user" },
);
});
@@ -558,6 +567,7 @@ describe("coach chat — tool-mode routing (F1)", () => {
1500,
"2026-06-21",
2_000_000,
"user",
);
// The legacy path ships the snapshot figures in the user turn.
const params = (
src/app/api/insights/chat/__tests__/route-workout-scope.test.ts
+1
View File
@@ -177,6 +177,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(async () => ({ allowed: true, reserved: 1500 })),
reconcileSpend: vi.fn(async () => undefined),
resolveDailyCap: vi.fn(() => 2_000_000),
resolveCostOwner: vi.fn(() => "user" as const),
}));
vi.mock("@/lib/ai/coach/refusal", () => ({
detectRefusal: vi.fn(() => ({ refuse: false })),
src/app/api/insights/chat/route.ts
+15 -5
View File
@@ -52,6 +52,7 @@ import {
runStreamingRawCompletionWithFallback,
} from "@/lib/ai/provider-runner";
import { resolveProviderChain, resolveProvider } from "@/lib/ai/provider";
import type { ProviderChainType } from "@/lib/ai/provider-chain";
import { assertConsentForChain } from "@/lib/ai/consent-guard";
import { singleUserTurn, type CompletionResult } from "@/lib/ai/types";
import { PROMPT_VERSION } from "@/lib/ai/prompts/insight-generator";
@@ -74,6 +75,7 @@ import {
buildDateKey,
reserveBudget,
reconcileSpend,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { detectRefusal } from "@/lib/ai/coach/refusal";
@@ -603,6 +605,7 @@ async function handleChatRequest(request: NextRequest): Promise<Response> {
: AI_BUDGETS.coach.maxTokens,
reqDateKey,
dailyCap,
resolveCostOwner(chain),
);
if (!reservation.allowed) {
annotate({
@@ -637,7 +640,10 @@ async function handleChatRequest(request: NextRequest): Promise<Response> {
async function produceReply(): Promise<ReplyOutcome> {
let result: CompletionResult;
let workingProviderType: string;
// v1.38.19 (Wave E) — typed as the chain's provider union, not a bare
// string: the budget reconcile attributes the turn's tokens to the cost
// owner of the hop named here.
let workingProviderType: ProviderChainType;
let toolTrace: CoachToolTrace[] = [];
// v1.21.0 (P6) — the present tool-result payloads this turn, for the post-hoc
// prose number-verifier. Empty on the no-tools path.
@@ -702,7 +708,9 @@ async function handleChatRequest(request: NextRequest): Promise<Response> {
timeoutMs: aiResponseTimeoutMs,
});
result = loop.result;
workingProviderType = loop.workingProviderType;
// The loop reports the hop it landed on as a bare string; it is
// assigned from `workingProvider.providerType` one frame up.
workingProviderType = loop.workingProviderType as ProviderChainType;
toolTrace = loop.toolTrace;
// v1.32.1 — the numeric verifier ACTIVATES only when this turn actually
// delivered figures the model was told to ground against: a pinned
@@ -784,9 +792,10 @@ async function handleChatRequest(request: NextRequest): Promise<Response> {
} 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(
() => {},
);
await reconcileSpend(userId, reservation.reserved, 0, reqDateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
}).catch(() => {});
// #781 — the client walked away mid-generation. The request's abort
// signal is threaded into every provider call, so the teardown surfaces
// here as an abort-shaped failure with `request.signal` already flipped.
@@ -887,6 +896,7 @@ async function handleChatRequest(request: NextRequest): Promise<Response> {
totalTokensSpent,
reqDateKey,
cachedTokensSpent,
{ servedBy: workingProviderType, reservedOwner: reservation.owner },
).catch(() => {
// Ledger reconcile is best-effort; a failure leaves the conservative
// reservation in place (never an undercount) and never breaks the turn.
src/app/api/labs/ocr/extract/__tests__/route.test.ts
+11
View File
@@ -45,6 +45,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 200_000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/labs/ocr-extract", async () => {
const actual = await vi.importActual<typeof import("@/lib/labs/ocr-extract")>(
@@ -95,6 +96,8 @@ beforeEach(() => {
allowed: true,
reserved: AI_BUDGETS.ocrExtractText.maxTokens ?? 0,
totalAfter: AI_BUDGETS.ocrExtractText.maxTokens ?? 0,
owner: "operator",
operatorAfter: AI_BUDGETS.ocrExtractText.maxTokens ?? 0,
} as never);
});
@@ -116,6 +119,8 @@ describe("POST /api/labs/ocr/extract — text mode budget", () => {
"2026-06-26",
// F1 — the provider-aware daily cap (mocked) is threaded as the 4th arg.
200_000,
// Wave E — and the cost owner that cap is enforced against as the 5th.
"operator",
);
});
@@ -134,6 +139,8 @@ describe("POST /api/labs/ocr/extract — text mode budget", () => {
AI_BUDGETS.ocrExtractText.maxTokens,
0,
"2026-06-26",
0,
{ servedBy: null, reservedOwner: "operator" },
);
});
@@ -182,6 +189,8 @@ describe("POST /api/labs/ocr/extract — vision PDF rasterization", () => {
allowed: true,
reserved: AI_BUDGETS.ocrExtract.maxTokens ?? 0,
totalAfter: AI_BUDGETS.ocrExtract.maxTokens ?? 0,
owner: "operator",
operatorAfter: AI_BUDGETS.ocrExtract.maxTokens ?? 0,
} as never);
});
@@ -222,6 +231,8 @@ describe("POST /api/labs/ocr/extract — vision PDF rasterization", () => {
AI_BUDGETS.ocrExtract.maxTokens,
0,
"2026-06-26",
0,
{ servedBy: null, reservedOwner: "operator" },
);
});
});
src/app/api/labs/ocr/extract/route.ts
+53 -12
View File
@@ -41,6 +41,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { prisma } from "@/lib/db";
@@ -183,6 +184,7 @@ async function handleTextExtract(
AI_BUDGETS.ocrExtractText.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
annotate({
@@ -208,12 +210,17 @@ async function handleTextExtract(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
return apiSuccess(result);
} catch (err) {
// A failed structuring call produced no usable rows; mirror the vision
// path and refund the reservation in full rather than charging it.
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (err instanceof OcrExtractError) {
return apiError("Couldn't read the report. Try a clearer photo.", 422, {
errorCode: "labs.ocr.extractFailed",
@@ -281,6 +288,7 @@ async function handleVisionExtract(
AI_BUDGETS.ocrExtract.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
annotate({
@@ -303,7 +311,10 @@ async function handleVisionExtract(
action: { name: "labs.ocr.fileRejected" },
meta: { reason: "content_length_exceeded" },
});
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return apiError("File is too large (max 12 MB).", 413, {
errorCode: "labs.ocr.fileTooLarge",
});
@@ -316,7 +327,10 @@ async function handleVisionExtract(
headers: { "content-type": request.headers.get("content-type") ?? "" },
}).formData();
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (err instanceof BodyTooLargeError) {
annotate({
action: { name: "labs.ocr.fileRejected" },
@@ -331,7 +345,10 @@ async function handleVisionExtract(
const file = formData.get("file");
if (!(file instanceof File)) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return apiError("Field 'file' must be a file", 422);
}
@@ -339,7 +356,10 @@ async function handleVisionExtract(
try {
buffer = Buffer.from(await file.arrayBuffer());
} catch {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return apiError("Failed to read uploaded file", 400);
}
@@ -350,7 +370,10 @@ async function handleVisionExtract(
action: { name: "labs.ocr.fileRejected" },
meta: { reason: "unsupported_mime" },
});
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return apiError("Upload a JPEG, PNG, WebP, or PDF.", 415, {
errorCode: "labs.ocr.fileType",
});
@@ -388,7 +411,10 @@ async function handleVisionExtract(
action: { name: "labs.ocr.fileRejected" },
meta: { reason: "pdf_rasterize_failed" },
});
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return apiError(
"Couldn't read this PDF; upload a photo instead.",
422,
@@ -416,12 +442,26 @@ async function handleVisionExtract(
// The orchestration does not surface token counts; reconcile against the
// reserved estimate as the spend ceiling (the provider already billed it).
actualTokens = reservation.reserved;
await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey);
await reconcileSpend(
userId,
reservation.reserved,
actualTokens,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
return apiSuccess(result);
} catch (err) {
// Provider/extraction failure — the call may still have burned tokens, so
// reconcile against the reserved estimate rather than refunding in full.
await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey);
await reconcileSpend(
userId,
reservation.reserved,
actualTokens,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
if (err instanceof OcrExtractError) {
return apiError("Couldn't read the report. Try a clearer photo.", 422, {
errorCode: "labs.ocr.extractFailed",
@@ -439,9 +479,10 @@ async function handleVisionExtract(
}
} catch (err) {
// A guard threw after the reservation (e.g. consent races) — refund fully.
await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch(
() => {},
);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
}).catch(() => {});
throw err;
}
}
src/app/api/medications/extract/__tests__/route.test.ts
+18 -1
View File
@@ -44,6 +44,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
// returns a distinctive figure so an accidental default-cap regression is
// visible in the assertion rather than silently passing.
resolveDailyCap: vi.fn(() => 1_234_567),
resolveCostOwner: vi.fn(() => "user" as const),
}));
vi.mock("@/lib/ai/provider", () => ({
@@ -118,6 +119,8 @@ beforeEach(() => {
allowed: true,
reserved: 700,
totalAfter: 700,
owner: "user",
operatorAfter: 0,
});
vi.mocked(resolveDailyCap).mockReturnValue(1_234_567);
vi.mocked(resolveProviderChain).mockResolvedValue([
@@ -274,6 +277,7 @@ describe("POST /api/medications/extract — daily budget", () => {
expect.any(Number),
"2026-05-28",
1_234_567,
"user",
);
});
@@ -283,6 +287,8 @@ describe("POST /api/medications/extract — daily budget", () => {
allowed: false,
reserved: 700,
totalAfter: 9_999_999,
owner: "user",
operatorAfter: 0,
});
const res = await POST(postReq({ text: "Mounjaro 5mg weekly" }) as never);
@@ -307,6 +313,8 @@ describe("POST /api/medications/extract — daily budget", () => {
700,
120,
"2026-05-28",
0,
{ servedBy: "openai", reservedOwner: "user" },
);
});
@@ -323,6 +331,8 @@ describe("POST /api/medications/extract — daily budget", () => {
700,
700,
"2026-05-28",
0,
{ servedBy: "openai", reservedOwner: "user" },
);
});
@@ -335,7 +345,14 @@ describe("POST /api/medications/extract — daily budget", () => {
const res = await POST(postReq({ text: "5mg weekly" }) as never);
expect(res.status).toBe(503);
expect(reconcileSpend).toHaveBeenCalledWith("user-1", 700, 0, "2026-05-28");
expect(reconcileSpend).toHaveBeenCalledWith(
"user-1",
700,
0,
"2026-05-28",
0,
{ servedBy: null, reservedOwner: "user" },
);
});
});
src/app/api/medications/extract/route.ts
+10 -4
View File
@@ -70,6 +70,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import {
@@ -210,6 +211,7 @@ async function handleExtract(request: NextRequest): Promise<Response> {
estimatedTokens,
dateKey,
resolveDailyCap(chain),
resolveCostOwner(chain),
);
if (!reservation.allowed) {
annotate({
@@ -241,9 +243,10 @@ async function handleExtract(request: NextRequest): Promise<Response> {
} catch (err) {
// Nothing was generated — refund the whole reservation rather than bill an
// invented figure. Best-effort: a failed refund must not mask the 503.
await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch(
() => {},
);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
}).catch(() => {});
if (err instanceof AllProvidersFailedError) {
annotate({
action: { name: "medications.extract.provider-failed" },
@@ -260,7 +263,10 @@ async function handleExtract(request: NextRequest): Promise<Response> {
// a free retry loop. Falls back to the reservation when the provider reports
// no count, so an unreported generation is never billed as zero.
const actualTokens = completion.result.tokensUsed ?? reservation.reserved;
await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey);
await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey, 0, {
servedBy: completion.workingProvider.providerType,
reservedOwner: reservation.owner,
});
const reply = completion.result.content?.trim() ?? "";
if (!reply) {
src/lib/ai/__tests__/provider-runner.test.ts
+5 -2
View File
@@ -13,13 +13,16 @@ import {
// v1.37.19 (A7-2) — the runner consults the day's recorded spend before an
// operator-funded hop. These pure chain tests do not stand up a DB, so the
// ledger read is swapped for a controllable in-memory figure.
const budgetState = vi.hoisted(() => ({ spent: 0 }));
const budgetState = vi.hoisted(() => ({ spent: 0, operatorSpent: 0 }));
vi.mock("../coach/budget", async () => {
const actual =
await vi.importActual<typeof import("../coach/budget")>("../coach/budget");
return {
...actual,
readDailySpend: vi.fn(async () => budgetState.spent),
readDailySpend: vi.fn(async () => ({
total: budgetState.spent,
operator: budgetState.operatorSpent,
})),
};
});
src/lib/ai/coach/__tests__/budget.test.ts
+36 -6
View File
@@ -4,6 +4,7 @@ import {
buildDateKey,
OPERATOR_COST_CAP,
USER_PLAN_CAP,
resolveCostOwner,
resolveDailyCap,
} from "../budget";
@@ -60,19 +61,37 @@ describe("reserveBudget — ledger clamps", () => {
it("clamps a non-finite reservation to 0", async () => {
const { reserveBudget } = await import("../budget");
const res = await reserveBudget("u", Number.NaN, "2026-05-10");
const res = await reserveBudget(
"u",
Number.NaN,
"2026-05-10",
OPERATOR_COST_CAP,
"operator",
);
expect(res.reserved).toBe(0);
});
it("clamps a negative reservation to 0", async () => {
const { reserveBudget } = await import("../budget");
const res = await reserveBudget("u", -42, "2026-05-10");
const res = await reserveBudget(
"u",
-42,
"2026-05-10",
OPERATOR_COST_CAP,
"operator",
);
expect(res.reserved).toBe(0);
});
it("floors a fractional reservation", async () => {
const { reserveBudget } = await import("../budget");
const res = await reserveBudget("u", 12.7, "2026-05-10");
const res = await reserveBudget(
"u",
12.7,
"2026-05-10",
OPERATOR_COST_CAP,
"operator",
);
expect(res.reserved).toBe(12);
});
});
@@ -143,6 +162,7 @@ describe("reserveBudget cap (F1 — user-plan path not locked out)", () => {
1_200,
"2026-05-10",
resolveDailyCap([{ providerType: "codex" }]),
resolveCostOwner([{ providerType: "codex" }]),
);
expect(res.allowed).toBe(true);
// The reservation upsert ran; no refund executeRaw fired.
@@ -161,6 +181,7 @@ describe("reserveBudget cap (F1 — user-plan path not locked out)", () => {
1_200,
"2026-05-10",
resolveDailyCap([{ providerType: "admin-openai" }]),
resolveCostOwner([{ providerType: "admin-openai" }]),
);
expect(res.allowed).toBe(false);
// Refund of the reservation fired on refusal.
@@ -181,7 +202,10 @@ describe("reconcileSpend cached-token subtraction (F3)", () => {
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);
await reconcileSpend("u", 1_200, 20_000, "2026-05-10", 13_000, {
servedBy: "admin-openai",
reservedOwner: "operator",
});
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);
@@ -190,7 +214,10 @@ describe("reconcileSpend cached-token subtraction (F3)", () => {
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);
await reconcileSpend("u", 1_200, 5_000, "2026-05-10", 9_999, {
servedBy: "admin-openai",
reservedOwner: "operator",
});
const interpolations = prismaMock.$executeRaw.mock.calls[0].slice(1);
// net actual clamped to 0 → delta = 0 - 1200 = -1200.
expect(interpolations).toContain(-1_200);
@@ -198,7 +225,10 @@ describe("reconcileSpend cached-token subtraction (F3)", () => {
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");
await reconcileSpend("u", 1_000, 4_000, "2026-05-10", 0, {
servedBy: "admin-openai",
reservedOwner: "operator",
});
const interpolations = prismaMock.$executeRaw.mock.calls[0].slice(1);
expect(interpolations).toContain(3_000);
});
src/lib/ai/coach/__tests__/self-context-questions-consent.test.ts
+11
View File
@@ -41,9 +41,12 @@ const budgetMocks = vi.hoisted(() => ({
allowed: true,
reserved: 300,
totalAfter: 300,
owner: "operator" as const,
operatorAfter: 300,
})),
reconcileSpend: vi.fn(async () => {}),
resolveDailyCap: vi.fn(() => 200_000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/coach/budget", () => budgetMocks);
@@ -87,6 +90,8 @@ beforeEach(() => {
allowed: true,
reserved: 300,
totalAfter: 300,
owner: "operator" as const,
operatorAfter: 300,
});
buildCoachSnapshot.mockResolvedValue({ snapshotJson: '{"weight":[]}' });
});
@@ -239,6 +244,7 @@ describe("deriveClarifyingQuestions — atomic budget", () => {
300,
"2026-07-18",
200_000,
"operator",
);
expect(budgetMocks.reconcileSpend).toHaveBeenCalledWith(
"user-1",
@@ -246,6 +252,7 @@ describe("deriveClarifyingQuestions — atomic budget", () => {
120,
"2026-07-18",
0,
{ servedBy: "openai", reservedOwner: "operator" },
);
});
@@ -254,6 +261,8 @@ describe("deriveClarifyingQuestions — atomic budget", () => {
allowed: false,
reserved: 300,
totalAfter: 200_000,
owner: "operator" as const,
operatorAfter: 200_000,
});
const out = await deriveClarifyingQuestions("user-1", ctx, "en");
@@ -273,6 +282,8 @@ describe("deriveClarifyingQuestions — atomic budget", () => {
300,
0,
"2026-07-18",
0,
{ servedBy: null, reservedOwner: "operator" },
);
});
});
src/lib/ai/coach/budget.ts
+115 -20
View File
@@ -63,8 +63,7 @@ export const USER_PLAN_CAP = 2_000_000;
export function resolveDailyCap(
chain: ReadonlyArray<{ providerType: ProviderChainType }>,
): number {
const primary = chain[0]?.providerType;
return primary === undefined || isOperatorFundedProvider(primary)
return resolveCostOwner(chain) === "operator"
? OPERATOR_COST_CAP
: USER_PLAN_CAP;
}
@@ -80,6 +79,28 @@ export function isOperatorFundedProvider(type: ProviderChainType): boolean {
return type === "admin-openai" || type === "admin-codex";
}
/**
* v1.38.19 (Wave E) — who pays for a turn. `"operator"` means the tokens land
* on the operator's invoice (`admin-openai` / `admin-codex`), `"user"` means
* the user's own plan, key or hardware carries them.
*/
export type BudgetCostOwner = "operator" | "user";
/**
* v1.38.19 (Wave E) — the cost owner a chain RESERVES under: the primary is
* the provider that will be tried first, so it is the expected payer. An empty
* chain defaults to the operator — the conservative side. `resolveDailyCap`
* and `reserveBudget` must agree on this classification, so both read it here.
*/
export function resolveCostOwner(
chain: ReadonlyArray<{ providerType: ProviderChainType }>,
): BudgetCostOwner {
const primary = chain[0]?.providerType;
return primary === undefined || isOperatorFundedProvider(primary)
? "operator"
: "user";
}
/**
* v1.37.19 (A7-2) — the day's recorded spend for one user (0 when no row).
*
@@ -90,15 +111,30 @@ export function isOperatorFundedProvider(type: ProviderChainType): boolean {
* hop must not spend past `OPERATOR_COST_CAP` just because the reservation
* was checked against the wrong owner's ceiling.
*/
export interface DailySpend {
/** Every token recorded for the day, whoever paid for it. */
total: number;
/**
* v1.38.19 (Wave E) — the share of `total` served by an operator-funded
* provider. The hop guard compares THIS against `OPERATOR_COST_CAP`: a day
* full of turns the user's own plan paid for must not close the operator's
* fallback hop.
*/
operator: number;
}
export async function readDailySpend(
userId: string,
dateKey: string = buildDateKey(),
): Promise<number> {
): Promise<DailySpend> {
const row = await prisma.coachUsage.findUnique({
where: { userId_dateKey: { userId, dateKey } },
select: { totalTokens: true },
select: { totalTokens: true, operatorTokens: true },
});
return row?.totalTokens ?? 0;
return {
total: row?.totalTokens ?? 0,
operator: row?.operatorTokens ?? 0,
};
}
/**
@@ -132,53 +168,79 @@ export function buildDateKey(at: Date = new Date()): string {
* reserved-but-failed turn still counts as an attempt.
*/
export interface ReserveBudgetResult {
/** True when the reservation kept the day's total within the cap. */
/** True when the reservation kept the day's spend 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;
/**
* v1.38.19 (Wave E) — the owner this reservation was charged to. The
* reconcile needs it: only a reservation booked to the operator has an
* amount to move back OUT of `operator_tokens` when a user-funded hop ends
* up serving the turn.
*/
owner: BudgetCostOwner;
/** The day's operator-funded spend AFTER this reservation. */
operatorAfter: number;
}
export async function reserveBudget(
userId: string,
estimatedTokens: number,
dateKey: string = buildDateKey(),
cap: number = OPERATOR_COST_CAP,
dateKey: string,
cap: number,
owner: BudgetCostOwner,
): Promise<ReserveBudgetResult> {
const reserved =
Number.isFinite(estimatedTokens) && estimatedTokens > 0
? Math.floor(estimatedTokens)
: 0;
// v1.38.19 (Wave E) — a reservation is booked to `operator_tokens` only when
// the chain's primary is operator-funded. The counter moves inside the SAME
// statement as the total, so two concurrent requests can never observe one
// counter without the other.
const operatorReserved = owner === "operator" ? reserved : 0;
// Single atomic upsert-increment returning the new total. Two concurrent
// Single atomic upsert-increment returning the new totals. 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())
const rows = await prisma.$queryRaw<
{ total_tokens: number; operator_tokens: number }[]
>`
INSERT INTO coach_usage (id, user_id, date_key, total_tokens, operator_tokens, message_count, created_at, updated_at)
VALUES (gen_random_uuid()::text, ${userId}, ${dateKey}, ${reserved}, ${operatorReserved}, 1, NOW(), NOW())
ON CONFLICT (user_id, date_key) DO UPDATE SET
total_tokens = coach_usage.total_tokens + ${reserved},
operator_tokens = coach_usage.operator_tokens + ${operatorReserved},
message_count = coach_usage.message_count + 1,
updated_at = NOW()
RETURNING total_tokens
RETURNING total_tokens, operator_tokens
`;
const totalAfter = Number(rows[0]?.total_tokens ?? reserved);
const operatorAfter = Number(rows[0]?.operator_tokens ?? operatorReserved);
// 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;
const priorOperator = operatorAfter - operatorReserved;
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 };
await refundReservation(userId, reserved, operatorReserved, dateKey);
return {
allowed: false,
reserved,
totalAfter: priorTotal,
owner,
operatorAfter: priorOperator,
};
}
return { allowed: true, reserved, totalAfter };
return { allowed: true, reserved, totalAfter, owner, operatorAfter };
}
/**
@@ -194,12 +256,24 @@ export async function reserveBudget(
* 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 interface ReconcileSpendOptions {
/**
* The provider that actually SERVED the turn (`runRawWithFallback`'s
* `workingProvider.providerType`), or `null` when no hop served — a failed,
* timed-out or cancelled call whose reservation is being refunded.
*/
servedBy: ProviderChainType | null;
/** The owner the reservation was booked to (`reserveBudget(...).owner`). */
reservedOwner: BudgetCostOwner;
}
export async function reconcileSpend(
userId: string,
reserved: number,
actualTokens: number,
dateKey: string = buildDateKey(),
cachedTokens = 0,
dateKey: string,
cachedTokens: number,
opts: ReconcileSpendOptions,
): Promise<void> {
const grossActual =
Number.isFinite(actualTokens) && actualTokens > 0
@@ -213,12 +287,31 @@ export async function reconcileSpend(
// (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;
// v1.38.19 (Wave E) — settle the operator's counter against the hop that
// SERVED, not the one the chain expected. Book the actual tokens when an
// operator-funded provider answered, and take back whatever this
// reservation had put there (nothing, when the reservation was the user's
// own plan — subtracting then would eat another request's legitimate
// operator balance). The four cases settle exactly:
// operator reservation + operator hop → `actual`
// operator reservation + user hop → 0 (the reservation moves out)
// user reservation + user hop → untouched
// user reservation + operator hop → `actual` (a fallback the
// operator really paid for)
const servedByOperator =
opts.servedBy !== null && isOperatorFundedProvider(opts.servedBy);
const operatorDelta =
(servedByOperator ? actual : 0) -
(opts.reservedOwner === "operator" ? reserved : 0);
if (delta === 0 && operatorDelta === 0) return;
// Clamp at zero so a smaller-than-reserved actual can't drive the row
// negative under a racing reconcile.
// negative under a racing reconcile. Both counters move in ONE statement:
// a concurrent reader can never see the total settled and the owner split
// still stale.
await prisma.$executeRaw`
UPDATE coach_usage
SET total_tokens = GREATEST(0, total_tokens + ${delta}),
operator_tokens = GREATEST(0, operator_tokens + ${operatorDelta}),
updated_at = NOW()
WHERE user_id = ${userId} AND date_key = ${dateKey}
`;
@@ -228,11 +321,13 @@ export async function reconcileSpend(
async function refundReservation(
userId: string,
reserved: number,
operatorReserved: number,
dateKey: string,
): Promise<void> {
await prisma.$executeRaw`
UPDATE coach_usage
SET total_tokens = GREATEST(0, total_tokens - ${reserved}),
operator_tokens = GREATEST(0, operator_tokens - ${operatorReserved}),
message_count = GREATEST(0, message_count - 1),
updated_at = NOW()
WHERE user_id = ${userId} AND date_key = ${dateKey}
src/lib/ai/coach/self-context-questions.ts
+7 -3
View File
@@ -45,6 +45,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import {
@@ -235,6 +236,7 @@ export async function deriveClarifyingQuestions(
QUESTIONS_MAX_TOKENS,
dateKey,
resolveDailyCap(chain),
resolveCostOwner(chain),
);
if (!reservation.allowed) return fallback();
@@ -266,9 +268,10 @@ export async function deriveClarifyingQuestions(
} catch (err) {
// Reconcile the reservation down to zero actual spend before the
// failure lands on the deterministic fallback.
await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch(
() => {},
);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
}).catch(() => {});
throw err;
}
@@ -278,6 +281,7 @@ export async function deriveClarifyingQuestions(
result.tokensUsed ?? 0,
dateKey,
result.cachedInputTokens ?? 0,
{ servedBy: chain[0].providerType, reservedOwner: reservation.owner },
).catch(() => {});
const questions = parseQuestionsReply(result.content);
src/lib/ai/provider-runner.ts
+6 -3
View File
@@ -431,8 +431,11 @@ async function runRawChain(
// Fail OPEN on a ledger read error: the reservation layer still
// meters every call against the same ledger, so a transient read
// failure here must not take the whole fallback chain down with it.
const spent = await readDailySpend(userId).catch(() => 0);
if (spent >= OPERATOR_COST_CAP) {
const spent = await readDailySpend(userId).catch(() => ({
total: 0,
operator: 0,
}));
if (spent.total >= OPERATOR_COST_CAP) {
const hop: FallbackHop = {
providerType: candidate.providerType,
attempt: i + 1,
@@ -445,7 +448,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,
operator_cap_spent: spent.total,
},
});
continue;
src/lib/documents/__tests__/fenced-chat-pipeline.test.ts
+1
View File
@@ -29,6 +29,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn().mockResolvedValue({ allowed: true, reserved: 600 }),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 100000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/provider-runner", () => ({
runStreamingRawCompletionWithFallback: vi.fn(),
src/lib/documents/__tests__/index-document.test.ts
+18 -2
View File
@@ -45,6 +45,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/ai-budgets", () => ({
AI_BUDGETS: { documentTranscribe: { temperature: 0, maxTokens: 4000 } },
@@ -106,6 +107,7 @@ beforeEach(() => {
vi.mocked(reserveBudget).mockResolvedValue({
allowed: true,
reserved: 1,
owner: "operator",
} as never);
vi.mocked(localExtractText).mockResolvedValue({
ok: true,
@@ -132,7 +134,14 @@ describe("indexDocumentContent — provider-first path", () => {
expect.objectContaining({ source: "vision", providerType: "anthropic" }),
);
// Budget reserved then reconciled at the full reserved amount (charged).
expect(reconcileSpend).toHaveBeenCalledWith("user-1", 1, 1, "2026-07-07");
expect(reconcileSpend).toHaveBeenCalledWith(
"user-1",
1,
1,
"2026-07-07",
0,
{ servedBy: expect.any(String), reservedOwner: "operator" },
);
// The local path is never touched when the provider succeeds.
expect(localExtractText).not.toHaveBeenCalled();
});
@@ -222,7 +231,14 @@ describe("indexDocumentContent — local fallback path", () => {
const outcome = await indexDocumentContent("user-1", "doc-1");
expect(outcome).toMatchObject({ indexed: true, source: "local-pdf" });
// Reservation refunded to zero spend.
expect(reconcileSpend).toHaveBeenCalledWith("user-1", 1, 0, "2026-07-07");
expect(reconcileSpend).toHaveBeenCalledWith(
"user-1",
1,
0,
"2026-07-07",
0,
{ servedBy: null, reservedOwner: "operator" },
);
expect(localExtractText).toHaveBeenCalledTimes(1);
});
src/lib/documents/__tests__/index-outcome-record.test.ts
+1
View File
@@ -56,6 +56,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/ai-budgets", () => ({
AI_BUDGETS: { documentTranscribe: { temperature: 0, maxTokens: 4000 } },
src/lib/documents/auto-stage-labs.ts
+10 -1
View File
@@ -283,6 +283,7 @@ export async function maybeAutoStageLabFacts(
AI_BUDGETS.ocrExtractText.maxTokens,
dateKey,
provider.dailyCap,
provider.costOwner,
);
if (!reservation.allowed) return { staged: false, reason: "budget" };
@@ -298,9 +299,17 @@ export async function maybeAutoStageLabFacts(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{
servedBy: provider.pick.entry.providerType,
reservedOwner: reservation.owner,
},
);
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
if (!(err instanceof InboundExtractError)) {
annotate({
action: { name: "documents.autoStage.failed" },
src/lib/documents/fenced-chat.ts
+7 -3
View File
@@ -43,6 +43,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { detectRefusal, type CoachRefusalReason } from "@/lib/ai/coach/refusal";
@@ -307,6 +308,7 @@ export async function streamFencedReply(
AI_BUDGETS.documentChat.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
annotate({ action: { name: "documents.chat.budget.exceeded" } });
@@ -383,9 +385,10 @@ Reply now as the assistant, grounded ONLY in the documents above, in ${
});
result = fallback.result;
} catch (err) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey).catch(
() => {},
);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
}).catch(() => {});
if (err instanceof AllProvidersFailedError) {
const allRateLimited =
err.attempts.length > 0 &&
@@ -425,6 +428,7 @@ Reply now as the assistant, grounded ONLY in the documents above, in ${
totalTokens,
dateKey,
cachedTokens,
{ servedBy: pick!.entry.providerType, reservedOwner: reservation.owner },
).catch(() => {});
let replyText = (result.content ?? "").trim();
src/lib/documents/index-document.ts
+22 -2
View File
@@ -36,7 +36,9 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
type BudgetCostOwner,
} from "@/lib/ai/coach/budget";
import {
loadOwnedDocument,
@@ -128,6 +130,12 @@ export interface ResolvedIndexProvider {
pick: Awaited<ReturnType<typeof resolveDocumentVisionProvider>>["pick"];
consentOk: boolean;
dailyCap: number;
/**
* v1.38.19 (Wave E) — whose budget `dailyCap` rations. The reservation books
* `operator_tokens` only for an operator-funded pick, so the cap and the
* counter it is compared against are resolved from one place.
*/
costOwner: BudgetCostOwner;
}
/**
@@ -160,7 +168,13 @@ export async function resolveIndexProvider(
pick && consentOk
? resolveDailyCap([{ providerType: pick.entry.providerType }])
: 0;
return { chain, pick, consentOk, dailyCap };
// v1.38.19 (Wave E) — the cap and its cost owner are resolved from the SAME
// pick, so the reservation books the counter the cap is enforced against.
const costOwner: BudgetCostOwner =
pick && consentOk
? resolveCostOwner([{ providerType: pick.entry.providerType }])
: "operator";
return { chain, pick, consentOk, dailyCap, costOwner };
}
/**
@@ -221,6 +235,7 @@ async function tryProviderIndex(
AI_BUDGETS.documentTranscribe.maxTokens,
dateKey,
provider.dailyCap,
provider.costOwner,
);
// Budget exhausted → fall through to the free local path rather than stall;
// a text-layer PDF stays searchable even once the AI allowance is spent.
@@ -238,6 +253,8 @@ async function tryProviderIndex(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
// Refs #776 — the empty-transcription guard: a provider answer with no
// text must NEVER become a "successful" empty index (it would mark the
@@ -261,7 +278,10 @@ async function tryProviderIndex(
} catch {
// Refund the reservation and let the caller fall through to local — a
// transient provider miss must never leave a text-layer PDF unsearchable.
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return { outcome: null, note: "provider-error" };
}
}
src/lib/insights/briefing-provider.ts
+13 -6
View File
@@ -39,6 +39,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { singleUserTurn } from "@/lib/ai/types";
@@ -122,6 +123,7 @@ export async function runBriefingCompletion(
estimatedTokens,
dateKey,
resolveDailyCap(args.chain),
resolveCostOwner(args.chain),
);
if (!reservation.allowed) {
annotate({
@@ -153,12 +155,13 @@ export async function runBriefingCompletion(
// invented figure. A partially-burned upstream call is possible here, but
// we have no reported count to charge, and over-charging a failed
// generation would ration the retry the user is entitled to.
await reconcileSpend(args.userId, reservation.reserved, 0, dateKey).catch(
() => {
// Best-effort: a failed refund leaves the conservative reservation in
// place (never an undercount) and must not mask the provider error.
},
);
await reconcileSpend(args.userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
}).catch(() => {
// Best-effort: a failed refund leaves the conservative reservation in
// place (never an undercount) and must not mask the provider error.
});
throw err;
}
@@ -169,6 +172,10 @@ export async function runBriefingCompletion(
actualTokens,
dateKey,
outcome.result.cachedInputTokens ?? 0,
{
servedBy: outcome.workingProvider.providerType,
reservedOwner: reservation.owner,
},
).catch(() => {
// Ledger reconcile is best-effort; a failure leaves the reservation in
// place and never breaks a generation that already succeeded.
src/lib/insights/status-provider.ts
+14 -3
View File
@@ -13,6 +13,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { AI_BUDGETS, REFERENCE_AI_SEED } from "@/lib/ai/ai-budgets";
@@ -223,6 +224,7 @@ export async function runStatusCompletion(
estimatedTokens,
dateKey,
resolveDailyCap(chain),
resolveCostOwner(chain),
);
if (!reservation.allowed) {
// Over the day's ceiling. Reported as `error` — a TRANSIENT miss the caller
@@ -264,11 +266,17 @@ export async function runStatusCompletion(
// A timed-out generation may still have burned upstream tokens, but we have
// no reported count to charge — refund the reservation rather than bill an
// invented figure.
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return { kind: "timeout" };
}
if (raced.errored || raced.value === null) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
return { kind: "error" };
}
@@ -279,7 +287,10 @@ export async function runStatusCompletion(
// into a free retry loop. Falls back to the reservation when the provider
// reports no count, so an unreported generation is never billed as zero.
const actualTokens = result.tokensUsed ?? reservation.reserved;
await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey);
await reconcileSpend(userId, reservation.reserved, actualTokens, dateKey, 0, {
servedBy: workingProvider.providerType,
reservedOwner: reservation.owner,
});
const content = result.content;
if (typeof content !== "string" || content.trim().length === 0) {
src/lib/jobs/__tests__/async-state-contract.test.ts
+1
View File
@@ -76,6 +76,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/ai-budgets", () => ({
AI_BUDGETS: {
src/lib/jobs/__tests__/coach-nudge-ai-consent.test.ts
+1
View File
@@ -35,6 +35,7 @@ const budgetMocks = vi.hoisted(() => ({
})),
reconcileSpend: vi.fn(async () => {}),
resolveDailyCap: vi.fn(() => 200_000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/coach/budget", () => budgetMocks);
src/lib/jobs/__tests__/document-content-index-backfill.test.ts
+1
View File
@@ -49,6 +49,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/jobs/boss-instance", () => ({ getGlobalBoss: vi.fn() }));
src/lib/jobs/__tests__/document-summary.test.ts
+18 -2
View File
@@ -50,6 +50,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: vi.fn(),
reconcileSpend: vi.fn().mockResolvedValue(undefined),
resolveDailyCap: vi.fn(() => 1000),
resolveCostOwner: vi.fn(() => "operator" as const),
}));
vi.mock("@/lib/ai/ai-budgets", () => ({
AI_BUDGETS: { documentSummary: { temperature: 0.3, maxTokens: 600 } },
@@ -124,6 +125,7 @@ beforeEach(() => {
vi.mocked(reserveBudget).mockResolvedValue({
allowed: true,
reserved: 5,
owner: "operator",
} as never);
});
@@ -151,7 +153,14 @@ describe("runDocumentSummaryJob — gating", () => {
"A lab report from a clinic listing routine blood values.",
);
// Budget reserved then reconciled at the full reserved amount (charged).
expect(reconcileSpend).toHaveBeenCalledWith("user-1", 5, 5, "2026-07-17");
expect(reconcileSpend).toHaveBeenCalledWith(
"user-1",
5,
5,
"2026-07-17",
0,
{ servedBy: expect.any(String), reservedOwner: "operator" },
);
// Write is scoped to the still-null column so a re-run cannot clobber it.
expect(prisma.inboundDocument.updateMany).toHaveBeenCalledWith(
expect.objectContaining({
@@ -216,7 +225,14 @@ describe("runDocumentSummaryJob — gating", () => {
await runDocumentSummaryJob({ userId: "user-1", documentId: "doc-1" });
// Reservation refunded to zero spend; no summary written, state recorded.
expect(reconcileSpend).toHaveBeenCalledWith("user-1", 5, 0, "2026-07-17");
expect(reconcileSpend).toHaveBeenCalledWith(
"user-1",
5,
0,
"2026-07-17",
0,
{ servedBy: null, reservedOwner: "operator" },
);
expect(prisma.inboundDocument.updateMany).toHaveBeenCalledWith(
expect.objectContaining({ data: { summaryState: "UNAVAILABLE" } }),
);
src/lib/jobs/__tests__/reaction-line-degradation.test.ts
+6
View File
@@ -88,6 +88,7 @@ vi.mock("@/lib/ai/coach/budget", () => ({
reserveBudget: (...a: unknown[]) => reserveBudget(...a),
reconcileSpend: (...a: unknown[]) => reconcileSpend(...a),
resolveDailyCap: () => 200_000,
resolveCostOwner: () => "operator" as const,
}));
const loadDailyDigest = vi.fn();
@@ -467,6 +468,7 @@ describe("reaction line — degradation", () => {
1_400,
"2026-07-16",
0,
{ servedBy: expect.any(String), reservedOwner: "operator" },
);
});
@@ -528,6 +530,8 @@ describe("reaction line — degradation", () => {
1_400,
1_400,
"2026-07-16",
0,
{ servedBy: expect.any(String), reservedOwner: "operator" },
);
expect(updateMany).toHaveBeenCalledWith(
expect.objectContaining({
@@ -623,6 +627,7 @@ describe("reaction line — degradation", () => {
900,
"2026-07-16",
0,
{ servedBy: expect.any(String), reservedOwner: "operator" },
);
expect(update).not.toHaveBeenCalled();
expectSurfaceStillWorks();
@@ -683,6 +688,7 @@ describe("reaction line — degradation", () => {
1_100,
"2026-07-16",
100,
{ servedBy: expect.any(String), reservedOwner: "operator" },
);
});
});
src/lib/jobs/coach-nudge-ai.ts
+7 -6
View File
@@ -43,6 +43,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { AI_BUDGETS } from "@/lib/ai/ai-budgets";
@@ -199,6 +200,7 @@ export const composeNudgeWithAI: ComposeNudgeWithAI = async (params) => {
maxTokens,
dateKey,
dailyCap,
resolveCostOwner(chain),
);
if (!reservation.allowed) {
annotate({ action: { name: "coach.nudge.ai.budget_exceeded" } });
@@ -224,12 +226,10 @@ export const composeNudgeWithAI: ComposeNudgeWithAI = async (params) => {
} catch {
// Timeout / network / provider error → refund what wasn't spent and
// fall back to the template.
await reconcileSpend(
params.userId,
reservation.reserved,
0,
dateKey,
).catch(() => {});
await reconcileSpend(params.userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
}).catch(() => {});
annotate({ action: { name: "coach.nudge.ai.fallback" } });
return null;
}
@@ -240,6 +240,7 @@ export const composeNudgeWithAI: ComposeNudgeWithAI = async (params) => {
result.tokensUsed ?? 0,
dateKey,
result.cachedInputTokens ?? 0,
{ servedBy: chain[0].providerType, reservedOwner: reservation.owner },
).catch(() => {});
const body = sanitiseAiBody(result.content, params.locale);
src/lib/jobs/document-content-index-backfill.ts
+19 -3
View File
@@ -38,6 +38,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import { prisma } from "@/lib/db";
@@ -231,6 +232,7 @@ export async function runContentIndexBackfillForUser(
AI_BUDGETS.documentTranscribe.maxTokens,
dateKey,
dailyCap,
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
if (!reservation.allowed) {
reason = "budget-reached";
@@ -239,7 +241,10 @@ export async function runContentIndexBackfillForUser(
const document = await loadOwnedDocument(userId, id);
if (!document) {
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
skipped += 1;
continue;
}
@@ -249,7 +254,10 @@ export async function runContentIndexBackfillForUser(
// or a PDF the rasteriser could not render) — refund and move on so
// the walk keeps converging. Counted as skipped, not failed: no
// provider call happened.
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
skipped += 1;
continue;
}
@@ -266,6 +274,11 @@ export async function runContentIndexBackfillForUser(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{
servedBy: pick.entry.providerType,
reservedOwner: reservation.owner,
},
);
await upsertContentIndex({
userId,
@@ -278,7 +291,10 @@ export async function runContentIndexBackfillForUser(
} catch {
// A provider miss on one document must not abort the batch — refund and
// leave it un-indexed (a later run retries it).
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
failed += 1;
}
}
src/lib/jobs/document-summary.ts
+8 -1
View File
@@ -30,6 +30,7 @@ import {
buildDateKey,
reconcileSpend,
reserveBudget,
resolveCostOwner,
resolveDailyCap,
} from "@/lib/ai/coach/budget";
import {
@@ -191,6 +192,7 @@ export async function runDocumentSummaryJob(
AI_BUDGETS.documentSummary.maxTokens,
dateKey,
resolveDailyCap([{ providerType: pick.entry.providerType }]),
resolveCostOwner([{ providerType: pick.entry.providerType }]),
);
// Budget exhausted → skip (no local fallback for a summary).
if (!reservation.allowed) {
@@ -224,6 +226,8 @@ export async function runDocumentSummaryJob(
reservation.reserved,
reservation.reserved,
dateKey,
0,
{ servedBy: pick.entry.providerType, reservedOwner: reservation.owner },
);
// WITHHOLD policy, v1.30.31. The model's prose still never lands — a
// summary that tripped the outbound screen must never be shown as if it
@@ -273,7 +277,10 @@ export async function runDocumentSummaryJob(
// Refund the reservation on a provider miss; the document keeps no summary
// (the on-demand route remains the manual fallback). Never rethrow — a
// transient provider error must not retry-loop or fail the queue.
await reconcileSpend(userId, reservation.reserved, 0, dateKey);
await reconcileSpend(userId, reservation.reserved, 0, dateKey, 0, {
servedBy: null,
reservedOwner: reservation.owner,
});
await markSummaryState(userId, documentId, "UNAVAILABLE");
annotate({
action: { name: "documents.summary.autoFailed" },
src/lib/jobs/reaction-line.ts
+26 -4
View File
@@ -41,7 +41,9 @@ import { isModuleEnabled } from "@/lib/modules/gate";
import {
buildDateKey,
reconcileSpend,
resolveCostOwner,
resolveDailyCap,
type BudgetCostOwner,
} from "@/lib/ai/coach/budget";
import { AI_BUDGETS } from "@/lib/ai/ai-budgets";
import { screenCoachReply } from "@/lib/ai/coach/outbound-guard";
@@ -314,6 +316,13 @@ type ReactionReservation = {
allowed: boolean;
reserved: number;
dateKey: string;
/**
* v1.38.19 (Wave E) — the cost owner the reservation was booked to. A
* RESUMED reservation (the claim survived a crash, its tokens are on the
* row) recomputes it from the same chain the reservation was made under, so
* the reconcile reverses exactly what the reservation added.
*/
owner: BudgetCostOwner;
};
class ReactionClaimLostError extends Error {
@@ -334,6 +343,7 @@ async function reserveClaimBudget(
revision: Date,
dateKey: string,
cap: number,
owner: BudgetCostOwner,
): Promise<ReactionReservation> {
return prisma.$transaction(async (tx) => {
const stillOwned = await tx.arrivalReaction.updateMany({
@@ -352,11 +362,17 @@ async function reserveClaimBudget(
if (stillOwned.count !== 1) throw new ReactionClaimLostError();
const reserved = ARRIVAL_REACTION_RESERVE_TOKENS;
// v1.38.19 (Wave E) — this reservation is the claim-linked twin of
// `reserveBudget`, so it books the owner split the same way: an
// operator-funded chain increments `operator_tokens` inside the SAME
// statement as the total, and the refusal path reverses both.
const operatorReserved = owner === "operator" ? reserved : 0;
const rows = await tx.$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())
INSERT INTO coach_usage (id, user_id, date_key, total_tokens, operator_tokens, message_count, created_at, updated_at)
VALUES (gen_random_uuid()::text, ${userId}, ${dateKey}, ${reserved}, ${operatorReserved}, 1, NOW(), NOW())
ON CONFLICT (user_id, date_key) DO UPDATE SET
total_tokens = coach_usage.total_tokens + ${reserved},
operator_tokens = coach_usage.operator_tokens + ${operatorReserved},
message_count = coach_usage.message_count + 1,
updated_at = NOW()
RETURNING total_tokens
@@ -366,6 +382,7 @@ async function reserveClaimBudget(
await tx.$executeRaw`
UPDATE coach_usage
SET total_tokens = GREATEST(0, total_tokens - ${reserved}),
operator_tokens = GREATEST(0, operator_tokens - ${operatorReserved}),
message_count = GREATEST(0, message_count - 1),
updated_at = NOW()
WHERE user_id = ${userId} AND date_key = ${dateKey}
@@ -374,7 +391,7 @@ async function reserveClaimBudget(
where: { id: rowId, userId, occurredAt: revision, generationClaimId },
data: { generationClaimId: null, generationClaimedAt: null },
});
return { allowed: false, reserved, dateKey };
return { allowed: false, reserved, dateKey, owner };
}
const linked = await tx.arrivalReaction.updateMany({
@@ -395,7 +412,7 @@ async function reserveClaimBudget(
});
if (linked.count !== 1) throw new ReactionClaimLostError();
return { allowed: true, reserved, dateKey };
return { allowed: true, reserved, dateKey, owner };
});
}
@@ -526,6 +543,7 @@ export async function runReactionLine(
allowed: true,
reserved: row.generationReservedTokens!,
dateKey: row.generationBudgetDateKey!,
owner: resolveCostOwner(chain),
};
} else {
try {
@@ -536,6 +554,7 @@ export async function runReactionLine(
revision,
buildDateKey(),
resolveDailyCap(chain),
resolveCostOwner(chain),
);
} catch (err) {
if (err instanceof ReactionClaimLostError) {
@@ -625,6 +644,8 @@ export async function runReactionLine(
reservation.reserved,
reservation.reserved,
reservation.dateKey,
0,
{ servedBy: chain[0].providerType, reservedOwner: reservation.owner },
).catch(() => {});
await finishTerminalAttempt().catch(() => {});
workerLog("error", "[reaction-line] generation failed", err);
@@ -638,6 +659,7 @@ export async function runReactionLine(
result.tokensUsed ?? reservation.reserved,
reservation.dateKey,
result.cachedInputTokens ?? 0,
{ servedBy: chain[0].providerType, reservedOwner: reservation.owner },
);
} catch (err) {
// Do not publish a line whose spend was not durably reconciled. The
tests/integration/coach-budget-owner.test.ts
+186
View File
@@ -0,0 +1,186 @@
/**
* v1.38.19 (Wave E) — the daily coach ledger attributes tokens to the COST
* OWNER of the hop that actually served them.
*
* Production evidence (operator instance, 2026-09-11): the chat refused with
* `coach.budget.exceeded` at `totalAfter: 513539` although almost every token
* of that day had been served by `codex` — the operator's own ChatGPT plan —
* after the operator-funded `admin-openai` primary answered HTTP 500. The
* ledger had no provider dimension, so tokens the USER's plan paid for were
* counted against the operator-key cap and the interactive chat was locked out
* by background jobs every morning.
*
* These contracts need real Postgres: the attribution lives inside the same
* atomic statements as the total (`INSERT … ON CONFLICT DO UPDATE … RETURNING`
* and the reconcile `UPDATE`), and a mocked client cannot prove that the two
* counters move together under the row lock.
*/
import { beforeEach, describe, expect, it } from "vitest";
import { getPrismaClient, truncateAllTables } from "./setup";
async function seedUser(): Promise<string> {
const user = await getPrismaClient().user.create({
data: {
username: "budget-owner-user",
email: "budget-owner@example.test",
role: "USER",
},
});
return user.id;
}
async function readRow(userId: string, dateKey: string) {
return getPrismaClient().coachUsage.findUnique({
where: { userId_dateKey: { userId, dateKey } },
});
}
beforeEach(async () => {
await truncateAllTables(getPrismaClient());
});
describe("coach_usage.operator_tokens (real Postgres)", () => {
it("bills an operator-served turn to both counters", async () => {
const { reserveBudget, reconcileSpend, OPERATOR_COST_CAP } =
await import("@/lib/ai/coach/budget");
const userId = await seedUser();
const dateKey = "2026-09-11";
const res = await reserveBudget(
userId,
1_000,
dateKey,
OPERATOR_COST_CAP,
"operator",
);
expect(res.allowed).toBe(true);
await reconcileSpend(userId, res.reserved, 3_000, dateKey, 0, {
servedBy: "admin-openai",
reservedOwner: "operator",
});
const row = await readRow(userId, dateKey);
expect(row?.totalTokens).toBe(3_000);
expect(row?.operatorTokens).toBe(3_000);
});
it("moves the reservation out of operator_tokens when the user's own plan served", async () => {
const { reserveBudget, reconcileSpend, OPERATOR_COST_CAP } =
await import("@/lib/ai/coach/budget");
const userId = await seedUser();
const dateKey = "2026-09-11";
// The chain's primary is operator-funded, so the reservation is charged to
// the operator — but the `admin-openai` hop failed and `codex` (the user's
// own ChatGPT plan) answered. The operator pays nothing for that turn.
const res = await reserveBudget(
userId,
1_000,
dateKey,
OPERATOR_COST_CAP,
"operator",
);
await reconcileSpend(userId, res.reserved, 3_000, dateKey, 0, {
servedBy: "codex",
reservedOwner: "operator",
});
const row = await readRow(userId, dateKey);
expect(row?.totalTokens).toBe(3_000);
expect(row?.operatorTokens).toBe(0);
});
it("a user-plan turn never touches an existing operator balance", async () => {
const { reserveBudget, reconcileSpend, USER_PLAN_CAP } =
await import("@/lib/ai/coach/budget");
const userId = await seedUser();
const dateKey = "2026-09-11";
// An operator-funded job already spent 150k earlier that day.
await getPrismaClient().coachUsage.create({
data: {
userId,
dateKey,
totalTokens: 150_000,
operatorTokens: 150_000,
messageCount: 12,
},
});
const res = await reserveBudget(
userId,
1_000,
dateKey,
USER_PLAN_CAP,
"user",
);
await reconcileSpend(userId, res.reserved, 4_000, dateKey, 0, {
servedBy: "codex",
reservedOwner: "user",
});
const row = await readRow(userId, dateKey);
expect(row?.totalTokens).toBe(154_000);
// The operator's balance is untouched: a user-plan turn must neither add
// to it nor deplete it.
expect(row?.operatorTokens).toBe(150_000);
});
it("bills an operator-funded FALLBACK hop under a user-plan reservation", async () => {
const { reserveBudget, reconcileSpend, USER_PLAN_CAP } =
await import("@/lib/ai/coach/budget");
const userId = await seedUser();
const dateKey = "2026-09-11";
// Primary is the user's own key (user-plan reservation), but the chain
// fell back onto the operator's shared key — those tokens land on the
// operator's invoice and must be counted there.
const res = await reserveBudget(
userId,
1_000,
dateKey,
USER_PLAN_CAP,
"user",
);
await reconcileSpend(userId, res.reserved, 2_500, dateKey, 0, {
servedBy: "admin-openai",
reservedOwner: "user",
});
const row = await readRow(userId, dateKey);
expect(row?.totalTokens).toBe(2_500);
expect(row?.operatorTokens).toBe(2_500);
});
it("refunds both counters when a refused reservation is rolled back", 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: OPERATOR_COST_CAP,
operatorTokens: OPERATOR_COST_CAP,
messageCount: 3,
},
});
const res = await reserveBudget(
userId,
1_000,
dateKey,
OPERATOR_COST_CAP,
"operator",
);
expect(res.allowed).toBe(false);
const row = await readRow(userId, dateKey);
expect(row?.totalTokens).toBe(OPERATOR_COST_CAP);
expect(row?.operatorTokens).toBe(OPERATOR_COST_CAP);
expect(row?.messageCount).toBe(3);
});
});
tests/integration/coach-budget-reservation.test.ts
+13 -7
View File
@@ -42,8 +42,8 @@ describe("reserveBudget (real Postgres)", () => {
});
const [a, b] = await Promise.all([
reserveBudget(userId, 600, dateKey, cap),
reserveBudget(userId, 600, dateKey, cap),
reserveBudget(userId, 600, dateKey, cap, "operator"),
reserveBudget(userId, 600, dateKey, cap, "operator"),
]);
// Exactly one is admitted: the first to commit sees prior 900 (< cap) and
@@ -65,7 +65,7 @@ describe("reserveBudget (real Postgres)", () => {
const userId = await seedUser();
const dateKey = "2026-06-20";
const res = await reserveBudget(userId, 600, dateKey, 25_000);
const res = await reserveBudget(userId, 600, dateKey, 25_000, "operator");
expect(res.allowed).toBe(true);
const row = await getPrismaClient().coachUsage.findUnique({
@@ -81,8 +81,11 @@ describe("reserveBudget (real Postgres)", () => {
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 res = await reserveBudget(userId, 600, dateKey, 25_000, "operator");
await reconcileSpend(userId, res.reserved, 120, dateKey, 0, {
servedBy: "admin-openai",
reservedOwner: res.owner,
});
const row = await getPrismaClient().coachUsage.findUnique({
where: { userId_dateKey: { userId, dateKey } },
@@ -99,8 +102,11 @@ describe("reserveBudget (real Postgres)", () => {
// 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 res = await reserveBudget(userId, 600, dateKey, 25_000, "operator");
await reconcileSpend(userId, res.reserved, 450, dateKey, 0, {
servedBy: "admin-openai",
reservedOwner: res.owner,
});
const row = await getPrismaClient().coachUsage.findUnique({
where: { userId_dateKey: { userId, dateKey } },