mirror of
https://github.com/MBombeck/HealthLog.git
synced 2026-11-15 06:52:04 +00:00
feat(ai): the daily ledger records which tokens the operator paid for
`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:
1 parent
274923a060
commit
badef28e96
53 files changed
+813
-117
No files matched your search
@@ -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" },
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -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),
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 }),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 }),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 }),
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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 })),
|
||||
|
||||
@@ -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(() => ({
|
||||
|
||||
@@ -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(() => ({
|
||||
|
||||
@@ -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(() => ({
|
||||
|
||||
@@ -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", () => ({
|
||||
|
||||
@@ -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", () => ({
|
||||
|
||||
@@ -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 })),
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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 })),
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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" },
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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" },
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in new issue
Block a user