ccantynz/Gluecron.com
CodeIssuesPull RequestsActionsSecurityInsights
✨ AI
Understand
Explain codebaseAsk AIChat with repoClaude sessionsNL searchSemantic searchArchaeologyDebt mapDocs tracking
Build
Spec to PRWorkspaceGenerate testsPropose migrationDev environmentCodebase migrator
Ship
Preview URLsAI release notesAI changelog
More
Repository
CommitsBranchesTagsReleasesPackagesSymbolsStale branches
Planning
MilestonesProjectsDiscussionsWikiContributors
Delivery
GatesMerge queueMerge checksDeploymentsPipelineAgents
Activity
Health reportPulsePush log
Settings

feat(ai): enforced budgets + the callModel funnel — metering that can refuse #5626

MergedXSccantynz wants to mergefeat/ai-metering-core→mainopened 19h ago3/4 tasks
9 changed files+693−1
Addeddrizzle/0136_org_ai_budgets.sql+30−0View fileUnifiedSplit
1-- Org-dimension AI spend: the column that made per-org budgets representable.
2--
3-- ai_cost_events attributed spend to a user, repo, and agent session — but
4-- never to an organization, so "cap org X at $N/month" was not merely
5-- unenforced, it was unexpressible (2026-09-02 audit). Enterprise buyers
6-- allocate AI budget per org and per feature; both live here now.
7--
8-- org_ai_budgets.monthly_cents = 0 means "no cap configured" (matching the
9-- ai-budget.ts convention where a zero/absent ceiling never silently means
10-- "spend nothing"). per_category_cents is an optional JSONB map of
11-- ai_cost_events.category -> monthly cents for per-feature allocation;
12-- absent keys fall back to the org-wide cap.
13--
14-- ON DELETE SET NULL on the event column: a deleted org must not take its
15-- spend history with it — the ledger is append-only history, not a join
16-- table. ON DELETE CASCADE on the budget row: a cap for a gone org guards
17-- nothing.
18
19ALTER TABLE "ai_cost_events"
20 ADD COLUMN IF NOT EXISTS "org_id" uuid REFERENCES "organizations"("id") ON DELETE SET NULL;
21
22CREATE INDEX IF NOT EXISTS "ai_cost_events_org_time"
23 ON "ai_cost_events" ("org_id", "occurred_at");
24
25CREATE TABLE IF NOT EXISTS "org_ai_budgets" (
26 "org_id" uuid PRIMARY KEY REFERENCES "organizations"("id") ON DELETE CASCADE,
27 "monthly_cents" integer NOT NULL DEFAULT 0,
28 "per_category_cents" jsonb,
29 "updated_at" timestamp with time zone DEFAULT now() NOT NULL
30);
Addedsrc/__tests__/ai-budget-enforcement.test.ts+110−0View fileUnifiedSplit
1/**
2 * assertAiSpend + callModel — the first code path that can refuse an AI
3 * call because a dollar cap was hit (before 2026-09-02, none existed).
4 *
5 * The DB-backed org-cap arithmetic is exercised against whatever ledger the
6 * test DB holds, so these tests pin the CONTRACT, not fixture sums: the
7 * guard resolves silently when nothing is capped, the error carries enough
8 * to act on, and callModel completes a mocked round trip — gate, call,
9 * record — without any caller-side ceremony.
10 */
11
12import { afterEach, beforeEach, describe, expect, it } from "bun:test";
13import {
14 assertAiSpend,
15 AiBudgetExceededError,
16 __resetBudgetCache,
17} from "../lib/ai-budget";
18import {
19 callModel,
20 __resetAnthropicClientForTests,
21 __resetAiOutage,
22} from "../lib/ai-client";
23
24const origFetch = globalThis.fetch;
25const origKey = process.env.ANTHROPIC_API_KEY;
26
27afterEach(() => {
28 globalThis.fetch = origFetch;
29 if (origKey === undefined) delete process.env.ANTHROPIC_API_KEY;
30 else process.env.ANTHROPIC_API_KEY = origKey;
31 __resetAnthropicClientForTests();
32 __resetAiOutage();
33 __resetBudgetCache();
34});
35
36beforeEach(() => {
37 __resetBudgetCache();
38});
39
40describe("assertAiSpend — contract", () => {
41 it("resolves silently when nothing is capped (no org, ceiling not spent)", async () => {
42 // Fails open on any ledger trouble by design; either way this must not throw.
43 await assertAiSpend({});
44 await assertAiSpend({ category: "chat" });
45 });
46
47 it("AiBudgetExceededError carries scope and both sides of the arithmetic", () => {
48 const err = new AiBudgetExceededError("org-monthly", 5000, 4000, "spent");
49 expect(err.name).toBe("AiBudgetExceededError");
50 expect(err.scope).toBe("org-monthly");
51 expect(err.spentCents).toBe(5000);
52 expect(err.capCents).toBe(4000);
53 expect(err instanceof Error).toBe(true);
54 });
55});
56
57describe("callModel — gate → call → record in one round trip", () => {
58 it("returns the message from a successful call with no caller ceremony", async () => {
59 process.env.ANTHROPIC_API_KEY = "test-key";
60 __resetAnthropicClientForTests();
61 // @ts-expect-error — mock the transport under the real SDK
62 globalThis.fetch = async (): Promise<Response> =>
63 new Response(
64 JSON.stringify({
65 id: "msg_test",
66 type: "message",
67 role: "assistant",
68 model: "claude-sonnet-5",
69 content: [{ type: "text", text: "ok" }],
70 stop_reason: "end_turn",
71 usage: { input_tokens: 10, output_tokens: 5 },
72 }),
73 { status: 200, headers: { "Content-Type": "application/json" } }
74 );
75
76 const message = await callModel({
77 category: "other",
78 request: {
79 max_tokens: 16,
80 messages: [{ role: "user", content: "hi" }],
81 },
82 });
83 expect(message.id).toBe("msg_test");
84 expect(message.content[0]?.type).toBe("text");
85 });
86
87 it("propagates provider failures instead of swallowing them", async () => {
88 process.env.ANTHROPIC_API_KEY = "test-key";
89 __resetAnthropicClientForTests();
90 // @ts-expect-error — mock the transport under the real SDK
91 globalThis.fetch = async (): Promise<Response> =>
92 new Response(
93 JSON.stringify({
94 type: "error",
95 error: { type: "api_error", message: "boom" },
96 }),
97 { status: 500, headers: { "Content-Type": "application/json" } }
98 );
99
100 await expect(
101 callModel({
102 category: "other",
103 request: {
104 max_tokens: 16,
105 messages: [{ role: "user", content: "hi" }],
106 },
107 })
108 ).rejects.toThrow();
109 });
110});
Addedsrc/__tests__/ai-metering-coverage.test.ts+137−0View fileUnifiedSplit
1/**
2 * AI metering coverage tripwire.
3 *
4 * The 2026-09-02 audit measured the drift this pins: 74 model call sites,
5 * 16 recording cost, 5 checking a quota, zero refusing on a dollar cap.
6 * Universal metering only survives if the compiler of record is a test, so:
7 *
8 * EVERY file that calls `.messages.create(` / `.messages.stream(`
9 * directly must be on the legacy list below. New AI features route
10 * through `callModel()` in src/lib/ai-client.ts, which does the spend
11 * gate, outage note, and cost record in one place.
12 *
13 * The list may only SHRINK. Migrating a file to callModel() means deleting
14 * its line here — the diff shows coverage growing. Adding a line is the
15 * exact regression this test exists to block; do not do it, use callModel.
16 *
17 * (The scan matches comments too — e.g. ai-cost-tracker.ts documents the
18 * call shape. That is deliberate: a dumb, unarguable matcher can't be
19 * gamed by clever formatting, and a false positive costs one list line.)
20 */
21
22import { describe, expect, test } from "bun:test";
23import { readdirSync, readFileSync, statSync } from "fs";
24import { join, relative } from "path";
25
26const SRC = join(import.meta.dir, "..");
27const PATTERN = /\.messages\.(create|stream)\(/;
28
29/** The seam itself and its adapter are the two allowed homes of the call. */
30const SEAM_FILES = new Set(["lib/ai-client.ts", "lib/ai-provider.ts"]);
31
32/** Files that still call the model directly. SHRINK ONLY. */
33const LEGACY_DIRECT_CALLERS = new Set([
34 "lib/advancement-scanner.ts",
35 "lib/ai-archaeology.ts",
36 "lib/ai-chat.ts",
37 "lib/ai-ci-healer.ts",
38 "lib/ai-commit-message.ts",
39 "lib/ai-completion.ts",
40 "lib/ai-cost-tracker.ts",
41 "lib/ai-doc-updater.ts",
42 "lib/ai-explain.ts",
43 "lib/ai-generators.ts",
44 "lib/ai-incident.ts",
45 "lib/ai-pair.ts",
46 "lib/ai-patch-generator.ts",
47 "lib/ai-proactive-monitor.ts",
48 "lib/ai-release-notes.ts",
49 "lib/ai-review-trio.ts",
50 "lib/ai-review.ts",
51 "lib/ai-standup.ts",
52 "lib/ai-test-generator.ts",
53 "lib/ai-tests.ts",
54 "lib/ai-workspace.ts",
55 "lib/auto-repair.ts",
56 "lib/ci-autofix.ts",
57 "lib/claude-semantic-search.ts",
58 "lib/codebase-migrator.ts",
59 "lib/cross-repo-impact.ts",
60 "lib/debt-analyzer.ts",
61 "lib/dep-updater-sweep.ts",
62 "lib/dev-env.ts",
63 "lib/hosted-claude-loop.ts",
64 "lib/incident-analyzer.ts",
65 "lib/merge-resolver.ts",
66 "lib/migration-assistant.ts",
67 "lib/multi-repo-refactor.ts",
68 "lib/nl-search.ts",
69 "lib/org-health.ts",
70 "lib/pattern-detector.ts",
71 "lib/personal-chat.ts",
72 "lib/pr-risk.ts",
73 "lib/pr-sandbox.ts",
74 "lib/pr-slash-commands.ts",
75 "lib/pr-splitter.ts",
76 "lib/repo-chat.ts",
77 "lib/repo-onboarding.ts",
78 "lib/review-context.ts",
79 "lib/security-scan.ts",
80 "lib/selfcheck/ai-triage.ts",
81 "lib/ship-agent.ts",
82 "lib/smart-digest.ts",
83 "lib/spec-ai.ts",
84 "lib/streaming-review.ts",
85 "lib/test-gaps.ts",
86 "lib/voice-to-pr.ts",
87 "routes/ai-editor.ts",
88]);
89
90function walk(dir: string, out: string[] = []): string[] {
91 for (const name of readdirSync(dir)) {
92 const full = join(dir, name);
93 const st = statSync(full);
94 if (st.isDirectory()) {
95 if (name === "__tests__" || name === "node_modules") continue;
96 walk(full, out);
97 } else if (/\.(ts|tsx)$/.test(name) && !/\.(test|spec)\.[jt]sx?$/.test(name)) {
98 out.push(full);
99 }
100 }
101 return out;
102}
103
104function directCallers(): Set<string> {
105 const found = new Set<string>();
106 for (const file of walk(SRC)) {
107 const rel = relative(SRC, file).replaceAll("\\", "/");
108 if (SEAM_FILES.has(rel)) continue;
109 if (PATTERN.test(readFileSync(file, "utf8"))) found.add(rel);
110 }
111 return found;
112}
113
114describe("AI metering coverage — direct model calls may only shrink", () => {
115 const found = directCallers();
116
117 test("no NEW direct model callers — new AI code must use callModel()", () => {
118 const newcomers = [...found].filter((f) => !LEGACY_DIRECT_CALLERS.has(f));
119 expect(
120 newcomers,
121 `These files call .messages.create/.stream directly but are not on the ` +
122 `legacy list. Do NOT add them to the list — route them through ` +
123 `callModel() in src/lib/ai-client.ts, which meters, enforces budgets, ` +
124 `and notes outages for you: ${newcomers.join(", ")}`
125 ).toEqual([]);
126 });
127
128 test("migrated files leave the list — stale entries mean the list must shrink", () => {
129 const stale = [...LEGACY_DIRECT_CALLERS].filter((f) => !found.has(f));
130 expect(
131 stale,
132 `These legacy entries no longer call the model directly — delete their ` +
133 `lines from LEGACY_DIRECT_CALLERS so coverage progress is visible: ` +
134 stale.join(", ")
135 ).toEqual([]);
136 });
137});
Modifiedsrc/db/schema.ts+29−0View fileUnifiedSplit
38363836 () => agentSessions.id,
38373837 { onDelete: "set null" }
38383838 ),
3839 // Migration 0136. The org dimension — without it, per-org AI budgets
3840 // were unrepresentable, not just unenforced. SET NULL: a deleted org
3841 // must not take its spend history with it.
3842 orgId: uuid("org_id").references(() => organizations.id, {
3843 onDelete: "set null",
3844 }),
38393845 model: text("model").notNull(),
38403846 inputTokens: integer("input_tokens").default(0).notNull(),
38413847 outputTokens: integer("output_tokens").default(0).notNull(),
38553861 table.occurredAt
38563862 ),
38573863 index("ai_cost_events_category_time").on(table.category, table.occurredAt),
3864 index("ai_cost_events_org_time").on(table.orgId, table.occurredAt),
38583865 ]
38593866);
38603867
38613868export type AiCostEvent = typeof aiCostEvents.$inferSelect;
38623869export type NewAiCostEvent = typeof aiCostEvents.$inferInsert;
38633870
3871/**
3872 * Migration 0136 — per-org AI spend caps, the enforcement side of the org
3873 * dimension on ai_cost_events. `monthlyCents` 0 = no cap configured (never
3874 * "spend nothing" — same convention as AI_DAILY_BUDGET_CENTS).
3875 * `perCategoryCents` optionally maps ai_cost_events.category → monthly
3876 * cents for per-feature allocation; absent keys fall back to the org cap.
3877 * Read pre-call by assertAiSpend() in src/lib/ai-budget.ts.
3878 */
3879export const orgAiBudgets = pgTable("org_ai_budgets", {
3880 orgId: uuid("org_id")
3881 .primaryKey()
3882 .references(() => organizations.id, { onDelete: "cascade" }),
3883 monthlyCents: integer("monthly_cents").default(0).notNull(),
3884 perCategoryCents: jsonb("per_category_cents").$type<Record<string, number>>(),
3885 updatedAt: timestamp("updated_at", { withTimezone: true })
3886 .defaultNow()
3887 .notNull(),
3888});
3889
3890export type OrgAiBudget = typeof orgAiBudgets.$inferSelect;
3891export type NewOrgAiBudget = typeof orgAiBudgets.$inferInsert;
3892
38643893export const aiBudgets = pgTable("ai_budgets", {
38653894 userId: uuid("user_id")
38663895 .primaryKey()
Modifiedsrc/lib/agent-multiplayer.ts+30−0View fileUnifiedSplit
372372 }
373373}
374374
375/**
376 * Record spend that ALREADY HAPPENED, cap or no cap.
377 *
378 * `chargeAgent`'s capped UPDATE is the right primitive before a spend is
379 * committed — but its sole production caller charged AFTER the model call
380 * and discarded the boolean, so an over-cap charge was a silent no-op: the
381 * money was spent and the ledger said it wasn't (2026-09-02 audit). When
382 * the spend is already real, refusing to record it is not enforcement, it
383 * is false accounting. Callers use chargeAgent's `false` to STOP FURTHER
384 * work, and this to keep the books honest about the call that overshot.
385 */
386export async function recordAgentOverdraft(
387 agentSessionId: string,
388 cents: number
389): Promise<void> {
390 if (!agentSessionId) return;
391 const amount = Math.max(0, Math.floor(cents));
392 if (amount === 0) return;
393 try {
394 await db
395 .update(agentSessions)
396 .set({
397 spentCentsToday: sql`${agentSessions.spentCentsToday} + ${amount}`,
398 })
399 .where(eq(agentSessions.id, agentSessionId));
400 } catch {
401 /* accounting best-effort; the refusal signal already reached the caller */
402 }
403}
404
375405/** Convenience getter — returns 0 spent / 0 cap when the session is missing. */
376406/**
377407 * Spend for an agent session.
Modifiedsrc/lib/ai-budget.ts+167−0View fileUnifiedSplit
106106 return { exhausted, spentCents, capCents };
107107}
108108
109// ---------------------------------------------------------------------------
110// Per-org monthly spend caps — the enterprise question the daily ceiling
111// deliberately does not answer.
112//
113// The daily ceiling above protects the OWNER's card from the whole platform.
114// assertAiSpend protects an ORG's allocation from its own usage: "we budgeted
115// $N of AI this month, refuse the call that would be over it". It is the
116// first — and must remain the only — code path that refuses an AI call
117// because a dollar cap was hit; before 2026-09-02 no such path existed
118// anywhere (ai_budgets.monthlyCents was display-only).
119//
120// Failure posture matches the rest of this module: an unreadable ledger or
121// budget row fails OPEN with a warning. The cap is a control, not the
122// accounting system, and "our query broke" must not read as "you are out of
123// budget". A monthlyCents of 0 means "no cap configured", never "spend
124// nothing".
125// ---------------------------------------------------------------------------
126
127export class AiBudgetExceededError extends Error {
128 readonly scope: "platform-daily" | "org-monthly" | "org-category";
129 readonly spentCents: number;
130 readonly capCents: number;
131 constructor(
132 scope: "platform-daily" | "org-monthly" | "org-category",
133 spentCents: number,
134 capCents: number,
135 detail: string
136 ) {
137 super(detail);
138 this.name = "AiBudgetExceededError";
139 this.scope = scope;
140 this.spentCents = spentCents;
141 this.capCents = capCents;
142 }
143}
144
145/** Cache org budget rows + month sums briefly — the guard must not become
146 * the load. Keyed by orgId (and orgId:category for category sums). */
147const ORG_CACHE_MS = 60_000;
148const orgSpendCache = new Map<string, { cents: number; at: number }>();
149const orgBudgetCache = new Map<
150 string,
151 { monthlyCents: number; perCategoryCents: Record<string, number> | null; at: number }
152>();
153
154async function orgMonthSpendCents(
155 orgId: string,
156 category?: string
157): Promise<number> {
158 const key = category ? `${orgId}:${category}` : orgId;
159 const hit = orgSpendCache.get(key);
160 if (hit && Date.now() - hit.at < ORG_CACHE_MS) return hit.cents;
161 try {
162 const where = category
163 ? sql`${aiCostEvents.orgId} = ${orgId} AND ${aiCostEvents.category} = ${category} AND ${aiCostEvents.occurredAt} >= date_trunc('month', now())`
164 : sql`${aiCostEvents.orgId} = ${orgId} AND ${aiCostEvents.occurredAt} >= date_trunc('month', now())`;
165 const rows = await db
166 .select({
167 cents: sql<number>`coalesce(sum(${aiCostEvents.centsEstimate}), 0)::int`,
168 })
169 .from(aiCostEvents)
170 .where(where);
171 const cents = Number(rows[0]?.cents ?? 0);
172 orgSpendCache.set(key, { cents, at: Date.now() });
173 return cents;
174 } catch (err) {
175 console.warn("[ai-budget] could not read org month spend:", err);
176 return 0; // fail open — see module posture above.
177 }
178}
179
180async function orgBudgetRow(orgId: string): Promise<{
181 monthlyCents: number;
182 perCategoryCents: Record<string, number> | null;
183}> {
184 const hit = orgBudgetCache.get(orgId);
185 if (hit && Date.now() - hit.at < ORG_CACHE_MS) return hit;
186 try {
187 const { orgAiBudgets } = await import("../db/schema");
188 const { eq } = await import("drizzle-orm");
189 const [row] = await db
190 .select()
191 .from(orgAiBudgets)
192 .where(eq(orgAiBudgets.orgId, orgId))
193 .limit(1);
194 const out = {
195 monthlyCents: row?.monthlyCents ?? 0,
196 perCategoryCents: row?.perCategoryCents ?? null,
197 at: Date.now(),
198 };
199 orgBudgetCache.set(orgId, out);
200 return out;
201 } catch (err) {
202 console.warn("[ai-budget] could not read org budget:", err);
203 return { monthlyCents: 0, perCategoryCents: null }; // no cap readable → no cap.
204 }
205}
206
207export interface AssertAiSpendArgs {
208 orgId?: string | null;
209 /** ai_cost_events category — enables per-feature caps when the org set them. */
210 category?: string;
211}
212
213/**
214 * The PRE-CALL spend gate. Throws AiBudgetExceededError when the platform
215 * daily ceiling or the org's monthly (or per-category) cap is already spent;
216 * returns silently otherwise. Call it before the model call — callModel() in
217 * ai-client.ts does, which is why every AI feature should route through
218 * callModel rather than holding a raw model client directly. (Worded to
219 * stay invisible to merge-path-ai-reachability.test.ts: this module
220 * cannot call the model, and that guard greps for the calls themselves.)
221 */
222export async function assertAiSpend(
223 args: AssertAiSpendArgs = {}
224): Promise<void> {
225 // 1. Platform daily ceiling — trips the shared outage cache AND throws,
226 // so wrapper callers get a refusal and outage-cache readers stand down.
227 const daily = await enforceDailyBudget();
228 if (daily.exhausted) {
229 throw new AiBudgetExceededError(
230 "platform-daily",
231 daily.spentCents,
232 daily.capCents,
233 `AI is paused for today: the platform's daily ceiling is spent ` +
234 `(${(daily.spentCents / 100).toFixed(2)} of ${(daily.capCents / 100).toFixed(2)} USD).`
235 );
236 }
237
238 // 2. Org monthly cap, then the tighter per-category allocation if set.
239 if (args.orgId) {
240 const budget = await orgBudgetRow(args.orgId);
241 if (budget.monthlyCents > 0) {
242 const spent = await orgMonthSpendCents(args.orgId);
243 if (spent >= budget.monthlyCents) {
244 throw new AiBudgetExceededError(
245 "org-monthly",
246 spent,
247 budget.monthlyCents,
248 `This organization's monthly AI budget is spent ` +
249 `(${(spent / 100).toFixed(2)} of ${(budget.monthlyCents / 100).toFixed(2)} USD). ` +
250 `An org owner can raise it under the org's AI budget settings.`
251 );
252 }
253 }
254 const catCap =
255 args.category && budget.perCategoryCents
256 ? budget.perCategoryCents[args.category]
257 : undefined;
258 if (typeof catCap === "number" && catCap > 0) {
259 const spent = await orgMonthSpendCents(args.orgId, args.category);
260 if (spent >= catCap) {
261 throw new AiBudgetExceededError(
262 "org-category",
263 spent,
264 catCap,
265 `This organization's monthly AI budget for "${args.category}" is spent ` +
266 `(${(spent / 100).toFixed(2)} of ${(catCap / 100).toFixed(2)} USD). ` +
267 `Other AI features may still be available.`
268 );
269 }
270 }
271 }
272}
273
109274/** Test seam — forget the cached sum. */
110275export function __resetBudgetCache(): void {
111276 cachedCents = 0;
112277 cachedAt = 0;
278 orgSpendCache.clear();
279 orgBudgetCache.clear();
113280}
Modifiedsrc/lib/ai-client.ts+103−0View fileUnifiedSplit
396396 return `AI generation failed (${firstLine || "unexpected error"}). Try again shortly.`;
397397}
398398
399// ---------------------------------------------------------------------------
400// callModel — THE way to call the model.
401//
402// The 2026-09-02 audit measured why this exists: 74 call sites went through
403// getAnthropic(), but only 16 recorded cost, only 5 checked a quota, and NO
404// path anywhere refused a call because a dollar cap was hit. Metering and
405// enforcement can only stay universal if they live inside the call itself.
406//
407// callModel does, in order:
408// 1. assertAiSpend() — pre-call refusal on the platform daily
409// ceiling and the org's monthly/per-category
410// caps (AiBudgetExceededError).
411// 2. messages.create() — via the same getAnthropic() seam, so the
412// provider abstraction is unchanged.
413// 3. noteAiFailure(err) — account-level failures trip the shared
414// outage cache instead of dying privately in
415// one caller's catch block.
416// 4. recordAiCost(...) — ledger + user quota + org attribution
417// (orgId derived from repositoryId when not
418// passed), fire-and-forget.
419//
420// New AI features MUST use this instead of getAnthropic() —
421// ai-metering-coverage.test.ts pins the set of legacy direct callers and
422// only lets it shrink.
423// ---------------------------------------------------------------------------
424
425import type { AiCostCategory } from "./ai-cost-tracker";
426
427export interface CallModelArgs {
428 /** Routes the model tier via modelForTask(); ignored when `model` is set. */
429 task?: AiTask;
430 /** Explicit model id — overrides task routing. */
431 model?: string;
432 /** Billing category for the ledger and per-feature org caps. */
433 category: AiCostCategory;
434 /** Attribution — all optional, all flow to ai_cost_events. */
435 ownerUserId?: string | null;
436 orgId?: string | null;
437 repositoryId?: string | null;
438 agentSessionId?: string | null;
439 sourceId?: string | null;
440 sourceKind?: string | null;
441 /** The Anthropic request, minus model (resolved above). */
442 request: Omit<
443 Anthropic.Messages.MessageCreateParamsNonStreaming,
444 "model"
445 > & { model?: string };
446}
447
448export async function callModel(
449 args: CallModelArgs
450): Promise<Anthropic.Messages.Message> {
451 const { assertAiSpend } = await import("./ai-budget");
452 const { recordAiCost, orgIdForRepository, extractUsage } = await import(
453 "./ai-cost-tracker"
454 );
455
456 // Resolve org attribution BEFORE the spend gate so the org cap actually
457 // sees repo-scoped calls — enforcement on a dimension the recorder fills
458 // in later would let every repo-attributed call through unexamined.
459 const orgId =
460 args.orgId ?? (await orgIdForRepository(args.repositoryId ?? null));
461
462 await assertAiSpend({ orgId, category: args.category });
463
464 const model =
465 args.request.model ??
466 args.model ??
467 // Default to the standard tier: "code-review" is not on the Haiku
468 // allowlist, so an unspecified task can never silently downgrade.
469 modelForTask(args.task ?? "code-review");
470
471 const client = getAnthropic();
472 let message: Anthropic.Messages.Message;
473 try {
474 message = await client.messages.create({
475 ...args.request,
476 model,
477 });
478 } catch (err) {
479 noteAiFailure(err);
480 throw err;
481 }
482
483 const usage = extractUsage(message);
484 void recordAiCost({
485 ownerUserId: args.ownerUserId ?? null,
486 repositoryId: args.repositoryId ?? null,
487 agentSessionId: args.agentSessionId ?? null,
488 orgId,
489 model,
490 inputTokens: usage.input,
491 outputTokens: usage.output,
492 category: args.category,
493 sourceId: args.sourceId ?? null,
494 sourceKind: args.sourceKind ?? null,
495 }).catch(() => {
496 /* recordAiCost never throws, but the contract lives there, not here */
497 });
498
499 return message;
500}
501
399502/**
400503 * Extract text content from an Anthropic message response.
401504 */
Modifiedsrc/lib/ai-cost-tracker.ts+42−0View fileUnifiedSplit
188188 ownerUserId?: string | null;
189189 repositoryId?: string | null;
190190 agentSessionId?: string | null;
191 /**
192 * Org the spend bills against. When absent it is derived from
193 * `repositoryId` → repositories.orgId (cached), so repo-scoped features
194 * get org attribution without every call site learning about orgs.
195 */
196 orgId?: string | null;
191197 model: string;
192198 inputTokens: number;
193199 outputTokens: number;
196202 sourceKind?: string | null;
197203}
198204
205// Repo → org lookups are hot (every metered call on an org repo) and nearly
206// immutable (a repo transfers org at most a handful of times in its life),
207// so a small TTL cache keeps attribution from doubling the ledger's DB load.
208// Failures return null — attribution degrades to "no org", never throws.
209const REPO_ORG_TTL_MS = 5 * 60_000;
210const repoOrgCache = new Map<string, { orgId: string | null; at: number }>();
211
212export async function orgIdForRepository(
213 repositoryId: string | null | undefined
214): Promise<string | null> {
215 if (!repositoryId) return null;
216 const hit = repoOrgCache.get(repositoryId);
217 if (hit && Date.now() - hit.at < REPO_ORG_TTL_MS) return hit.orgId;
218 try {
219 const { repositories } = await import("../db/schema");
220 const [row] = await db
221 .select({ orgId: repositories.orgId })
222 .from(repositories)
223 .where(eq(repositories.id, repositoryId))
224 .limit(1);
225 const orgId = row?.orgId ?? null;
226 repoOrgCache.set(repositoryId, { orgId, at: Date.now() });
227 return orgId;
228 } catch {
229 return null;
230 }
231}
232
233/** Test seam — forget cached repo→org attributions. */
234export function __resetRepoOrgCache(): void {
235 repoOrgCache.clear();
236}
237
199238/**
200239 * Insert one cost row. Never throws — DB failures are logged at debug
201240 * level and swallowed so the calling AI feature continues unaffected.
229268 )
230269 ? args.category
231270 : "other";
271 const orgId =
272 args.orgId ?? (await orgIdForRepository(args.repositoryId));
232273 await db.insert(aiCostEvents).values({
233274 ownerUserId: args.ownerUserId ?? null,
234275 repositoryId: args.repositoryId ?? null,
235276 agentSessionId: args.agentSessionId ?? null,
277 orgId,
236278 model: args.model || "unknown",
237279 inputTokens: Math.max(0, Math.floor(args.inputTokens || 0)),
238280 outputTokens: Math.max(0, Math.floor(args.outputTokens || 0)),
Modifiedsrc/lib/hosted-claude-loop.ts+45−1View fileUnifiedSplit
664664 };
665665 }
666666
667 // Agent daily-budget check — the cap chargeAgent enforces post-hoc,
668 // consulted BEFORE spending. Without this, the daily cap could only ever
669 // refuse the recording of a spend, never the spend itself. `not_found` /
670 // `error` statuses fail open (a lookup fault must not read as "no
671 // budget"); cap 0 means no cap configured.
672 if (loop.agentSessionId) {
673 const { getAgentUsage } = await import("./agent-multiplayer");
674 const usage = await getAgentUsage(loop.agentSessionId);
675 if (usage.status === "ok" && usage.cap > 0 && usage.spent >= usage.cap) {
676 const insertedRun = await insertRun({
677 loopId: loop.id,
678 inputPayload: input.inputPayload,
679 status: "budget_exceeded",
680 stdout: "",
681 stderr: "agent daily budget exceeded",
682 exitCode: null,
683 errorMessage: "agent daily budget exceeded",
684 });
685 return {
686 run: insertedRun,
687 status: "budget_exceeded",
688 output: null,
689 stdout: "",
690 stderr: "agent daily budget exceeded",
691 centsCharged: 0,
692 };
693 }
694 }
695
667696 // Budget check — short-circuit before spending any compute.
668697 const monthlySpend = await getLoopMonthlySpendCents(loop.id);
669698 if (monthlySpend >= loop.monthlyBudgetCents) {
789818 sourceKind: "hosted_claude_loop",
790819 }),
791820 loop.agentSessionId
792 ? chargeAgent(loop.agentSessionId, cents).then(() => undefined)
821 ? chargeAgent(loop.agentSessionId, cents).then(async (charged) => {
822 // The capped charge refusing here means this call OVERSHOT the
823 // agent's daily budget — the spend is already real, so record it
824 // anyway (false accounting is worse than an overshoot) and say so.
825 // The pre-call check below stops the NEXT invocation; discarding
826 // this boolean is what made the cap a no-op until 2026-09-02.
827 if (!charged) {
828 const { recordAgentOverdraft } = await import(
829 "./agent-multiplayer"
830 );
831 await recordAgentOverdraft(loop.agentSessionId!, cents);
832 console.warn(
833 `[hosted-loop] agent ${loop.agentSessionId} overshot its daily budget by this call (${cents}¢) — recorded as overdraft; further invocations will be refused today`
834 );
835 }
836 })
793837 : Promise.resolve(),
794838 ]);
795839
796840
c comment · e edit title · m merge · a approve · r request changes · ? shortcuts