Commit0b49751unknown_key
Wire AI budget hard enforcement — quota gate before AI feature calls
Wire AI budget hard enforcement — quota gate before AI feature calls
- billing.ts: Add `AiQuotaExceededError`, a 60s TTL per-user in-memory
Map cache (`checkAiQuotaCached`), and `assertAiQuota(userId)` which
throws on exhaustion and warns at 90% of limit.
- ai-review.ts: Call `assertAiQuota(pr.authorId)` inside
`triggerAiReview()` before any API call. On quota exhaustion, post a
PR comment telling the author to upgrade, then bail.
- ai-review-trio.ts: Call `assertAiQuota(prAuthorId)` inside
`runTrioReview()` before the 3-way fan-out. On exhaustion, post a
single quota-exceeded trio summary comment and return a fail-closed
result.
- ai-ci-healer.ts: Resolve the repo owner in `healOneRun()` and call
`assertAiQuota`. On exhaustion, skip silently (no audit marker written
so the healer will retry after the billing cycle resets).
- spec-to-pr.ts: Gate both `createSpecPR()` (uses `args.userId`) and
`runSpecToPr()` (uses `repoRow.ownerId`). Both return
`{ok: false, error: "..."}` on exhaustion — the UI surfaces the
error inline.
https://claude.ai/code/session_01DzJMTFASjMHt2f5ze4cNLR5 files changed+266−50b49751fe32281febcf56723eea89f83609b54db
5 changed files+266−5
Modifiedsrc/lib/ai-ci-healer.ts+39−5View fileUnifiedSplit
@@ -48,6 +48,7 @@ import {
4848 generatePatchForGateTestFinding,
4949 type GateTestFinding,
5050} from "./ai-patch-generator";
51import { assertAiQuota, AiQuotaExceededError } from "./billing";
5152
5253// ---------------------------------------------------------------------------
5354// Tunables
@@ -450,13 +451,10 @@ export async function healOneRun(
450451 return { outcome: "skipped", reason: "already processed" };
451452 }
452453
453 const analysis = await analyzeFailedWorkflowRun(runId, {
454 client: opts.client,
455 });
456
457 // Look up the run for audit metadata (repo + sha).
454 // Look up the run for audit metadata (repo + sha) and owner for quota check.
458455 let repositoryId: string | null = null;
459456 let commitSha: string | null = null;
457 let repoOwnerId: string | null = null;
460458 try {
461459 const [row] = await db
462460 .select({
@@ -474,6 +472,42 @@ export async function healOneRun(
474472 console.warn("[ai-ci-healer] post-analyze run lookup failed:", err);
475473 }
476474
475 // Resolve repo owner for the quota check.
476 if (repositoryId) {
477 try {
478 const [repoRow] = await db
479 .select({ ownerId: repositories.ownerId })
480 .from(repositories)
481 .where(eq(repositories.id, repositoryId))
482 .limit(1);
483 if (repoRow) repoOwnerId = repoRow.ownerId;
484 } catch {
485 /* tolerate */
486 }
487 }
488
489 // Hard quota gate — skip silently when the repo owner's AI budget is
490 // exhausted. We don't fail the CI run; the autopilot marker is not written
491 // so the healer will retry on the next tick once budget resets.
492 if (repoOwnerId) {
493 try {
494 await assertAiQuota(repoOwnerId);
495 } catch (err) {
496 if (err instanceof AiQuotaExceededError) {
497 console.log(
498 `[ai-ci-healer] skipping run=${runId}: AI quota exceeded for owner=${repoOwnerId}`
499 );
500 return { outcome: "skipped", reason: "AI quota exceeded" };
501 }
502 // Unexpected error — log and proceed (fail open).
503 console.warn("[ai-ci-healer] assertAiQuota failed unexpectedly:", err);
504 }
505 }
506
507 const analysis = await analyzeFailedWorkflowRun(runId, {
508 client: opts.client,
509 });
510
477511 if (!analysis) {
478512 await audit({
479513 userId: null,
Modifiedsrc/lib/ai-review-trio.ts+57−0View fileUnifiedSplit
@@ -32,6 +32,7 @@ import { pullRequests, prComments } from "../db/schema";
3232import { getAnthropic, MODEL_SONNET, parseJsonResponse } from "./ai-client";
3333import { audit } from "./notify";
3434import { recordAiCost, extractUsage } from "./ai-cost-tracker";
35import { assertAiQuota, AiQuotaExceededError } from "./billing";
3536
3637// ---------------------------------------------------------------------------
3738// Public types
@@ -253,6 +254,62 @@ export async function runTrioReview(
253254 ? opts.diff.slice(0, DIFF_BYTE_CAP)
254255 : opts.diff;
255256
257 // 0. Hard quota gate — bail before any API calls if the PR author is over
258 // budget. Resolve the author from the DB (needed for the comment insert
259 // in persistTrioComments anyway).
260 let prAuthorId: string | null = null;
261 try {
262 const [pr] = await db
263 .select({ authorId: pullRequests.authorId })
264 .from(pullRequests)
265 .where(eq(pullRequests.id, opts.pullRequestId))
266 .limit(1);
267 if (pr) prAuthorId = pr.authorId;
268 } catch {
269 /* tolerate — quota check will fail open */
270 }
271 if (prAuthorId) {
272 try {
273 await assertAiQuota(prAuthorId);
274 } catch (err) {
275 if (err instanceof AiQuotaExceededError) {
276 // Post a single summary comment so the PR author sees the skip reason.
277 try {
278 await db.insert(prComments).values({
279 pullRequestId: opts.pullRequestId,
280 authorId: prAuthorId,
281 isAiReview: true,
282 body: [
283 TRIO_SUMMARY_MARKER,
284 "## AI Trio Review skipped",
285 "",
286 "Your monthly AI token budget has been reached. Upgrade at [/settings/billing](/settings/billing) to re-enable AI code review.",
287 ].join("\n"),
288 });
289 } catch {
290 /* best-effort */
291 }
292 // Return a neutral fail-closed result so the caller doesn't crash.
293 const skippedVerdict = (persona: TrioPersona): TrioVerdict => ({
294 persona,
295 verdict: "fail",
296 findings: [],
297 rawText: "",
298 latencyMs: 0,
299 failed: true,
300 });
301 return {
302 securityVerdict: skippedVerdict("security"),
303 correctnessVerdict: skippedVerdict("correctness"),
304 styleVerdict: skippedVerdict("style"),
305 disagreements: [],
306 };
307 }
308 // Unexpected error — log and proceed (fail open).
309 console.warn("[ai-review-trio] assertAiQuota failed unexpectedly:", err);
310 }
311 }
312
256313 // 1. Fan out the three persona calls.
257314 const personas: TrioPersona[] = ["security", "correctness", "style"];
258315 const [securityVerdict, correctnessVerdict, styleVerdict] = await Promise.all(
Modifiedsrc/lib/ai-review.ts+24−0View fileUnifiedSplit
@@ -17,6 +17,7 @@ import {
1717 alreadyTrioReviewed,
1818 runTrioReview,
1919} from "./ai-review-trio";
20import { assertAiQuota, AiQuotaExceededError } from "./billing";
2021
2122interface ReviewComment {
2223 filePath: string;
@@ -304,6 +305,29 @@ export async function triggerAiReview(
304305 .limit(1);
305306 if (!pr) return;
306307
308 // Hard quota gate — post a comment and bail if the user's AI budget is
309 // exhausted. This runs after loading the PR so we have authorId for the
310 // comment insert.
311 try {
312 await assertAiQuota(pr.authorId);
313 } catch (err) {
314 if (err instanceof AiQuotaExceededError) {
315 await db
316 .insert(prComments)
317 .values({
318 pullRequestId: prId,
319 authorId: pr.authorId,
320 isAiReview: true,
321 body: `${AI_REVIEW_MARKER}\n## AI review skipped\n\nYour monthly AI token budget has been reached. Upgrade at [/settings/billing](/settings/billing) to re-enable AI code review.`,
322 })
323 .catch(() => {});
324 return;
325 }
326 // Any other error from assertAiQuota is unexpected — log and proceed
327 // (fail open so a billing glitch never silently kills reviews).
328 console.warn("[ai-review] assertAiQuota failed unexpectedly:", err);
329 }
330
307331 let diffText = await diffBetweenBranches(
308332 ownerName,
309333 repoName,
Modifiedsrc/lib/billing.ts+113−0View fileUnifiedSplit
@@ -316,3 +316,116 @@ export function formatPrice(cents: number): string {
316316 if (cents === 0) return "Free";
317317 return `$${(cents / 100).toFixed(2)}/mo`;
318318}
319
320// ---------------------------------------------------------------------------
321// AI quota hard-gate helpers
322// ---------------------------------------------------------------------------
323
324/**
325 * Thrown by `assertAiQuota` when the user's AI token budget is exhausted.
326 * Callers should catch this specifically and degrade gracefully — post a
327 * comment, return a user-facing error, or skip silently — rather than letting
328 * it bubble up as an unhandled exception.
329 */
330export class AiQuotaExceededError extends Error {
331 readonly userId: string;
332 readonly planSlug: string;
333 constructor(userId: string, planSlug: string) {
334 super(
335 `AI quota exceeded for user ${userId} (plan: ${planSlug}). Upgrade at /settings/billing.`
336 );
337 this.name = "AiQuotaExceededError";
338 this.userId = userId;
339 this.planSlug = planSlug;
340 }
341}
342
343interface QuotaCacheEntry {
344 /** True = allowed, false = blocked. */
345 allowed: boolean;
346 /** Usage as a fraction of the limit (0–1). Used to emit the 90% warning. */
347 fraction: number;
348 planSlug: string;
349 expiresAt: number;
350}
351
352/** Per-user TTL cache. A simple Map is fine — the process is single-replica. */
353const _quotaCache = new Map<string, QuotaCacheEntry>();
354
355/** Cache TTL in milliseconds. */
356const QUOTA_CACHE_TTL_MS = 60_000;
357
358/**
359 * Check whether the user has AI token budget remaining, using an in-process
360 * 60-second cache to avoid a DB round-trip on every AI call.
361 *
362 * - Returns `{ allowed: true }` when usage is under 100% of the plan limit.
363 * - Returns `{ allowed: false }` when at or above the limit.
364 * - Returns `{ allowed: true, nearLimit: true }` when usage is >= 90% but
365 * under 100%, so callers can log a warning.
366 *
367 * Fails **open** on any DB / billing error (same policy as `checkQuota`).
368 */
369export async function checkAiQuotaCached(
370 userId: string
371): Promise<{ allowed: boolean; nearLimit: boolean; planSlug: string }> {
372 if (!userId) return { allowed: true, nearLimit: false, planSlug: "free" };
373
374 const now = Date.now();
375 const cached = _quotaCache.get(userId);
376 if (cached && cached.expiresAt > now) {
377 return {
378 allowed: cached.allowed,
379 nearLimit: cached.fraction >= 0.9 && cached.allowed,
380 planSlug: cached.planSlug,
381 };
382 }
383
384 try {
385 const quota = await getUserQuota(userId);
386 const limit = quota.plan.aiTokensMonthly;
387 const used = quota.usage.aiTokensUsedThisMonth;
388 const allowed = limit <= 0 || used < limit;
389 const fraction = limit > 0 ? used / limit : 0;
390 const entry: QuotaCacheEntry = {
391 allowed,
392 fraction,
393 planSlug: quota.planSlug,
394 expiresAt: now + QUOTA_CACHE_TTL_MS,
395 };
396 _quotaCache.set(userId, entry);
397 return { allowed, nearLimit: fraction >= 0.9 && allowed, planSlug: quota.planSlug };
398 } catch {
399 // Fail open — billing is never a hard dependency for the primary path.
400 return { allowed: true, nearLimit: false, planSlug: "free" };
401 }
402}
403
404/**
405 * Invalidate the cached quota entry for a user. Call after `bumpUsage` when
406 * you want the next AI call to re-check immediately rather than waiting up to
407 * 60 seconds.
408 */
409export function invalidateQuotaCache(userId: string): void {
410 _quotaCache.delete(userId);
411}
412
413/**
414 * Assert that the user has AI token quota remaining. Throws
415 * `AiQuotaExceededError` if blocked; logs a warning (but proceeds) when
416 * usage is 90-99% of the limit.
417 *
418 * Designed to be called at the top of every AI feature entry point.
419 * Never throws on billing errors (fails open).
420 */
421export async function assertAiQuota(userId: string): Promise<void> {
422 const { allowed, nearLimit, planSlug } = await checkAiQuotaCached(userId);
423 if (nearLimit) {
424 console.warn(
425 `[billing] AI quota warning: user ${userId} (plan: ${planSlug}) is near their monthly AI token limit.`
426 );
427 }
428 if (!allowed) {
429 throw new AiQuotaExceededError(userId, planSlug);
430 }
431}
Modifiedsrc/lib/spec-to-pr.ts+33−0View fileUnifiedSplit
@@ -33,6 +33,7 @@ import {
3333 createOrUpdateFileOnBranch,
3434 getTreeRecursive,
3535} from "../git/repository";
36import { assertAiQuota, AiQuotaExceededError } from "./billing";
3637
3738export type SpecPRArgs = {
3839 repoId: string;
@@ -71,6 +72,22 @@ export async function createSpecPR(args: SpecPRArgs): Promise<SpecPRResult> {
7172 return { ok: false, error: "ANTHROPIC_API_KEY required for spec-to-PR" };
7273 }
7374
75 // 1b. AI quota hard gate. Returns a user-visible error when the budget is
76 // exhausted so the UI can surface an upgrade prompt.
77 try {
78 await assertAiQuota(args.userId);
79 } catch (err) {
80 if (err instanceof AiQuotaExceededError) {
81 return {
82 ok: false,
83 error:
84 "Monthly AI token budget reached. Upgrade at /settings/billing to continue using spec-to-PR.",
85 };
86 }
87 // Unexpected billing error — fail open so a DB glitch doesn't block users.
88 console.warn("[spec-to-pr] assertAiQuota failed unexpectedly:", err);
89 }
90
7491 const spec = typeof args.spec === "string" ? args.spec.trim() : "";
7592 if (!spec) return { ok: false, error: "spec is empty" };
7693
@@ -451,6 +468,22 @@ export async function runSpecToPr(
451468 return { ok: false, error: "repo not found" };
452469 }
453470
471 // 2b. AI quota hard gate — check against the repo owner's budget.
472 // runSpecToPr is called from the autopilot so we skip silently (return
473 // ok:false) rather than throwing; the autopilot loop logs the error.
474 try {
475 await assertAiQuota(repoRow.ownerId);
476 } catch (err) {
477 if (err instanceof AiQuotaExceededError) {
478 return {
479 ok: false,
480 error:
481 "Monthly AI token budget reached. Upgrade at /settings/billing to continue using spec-to-PR.",
482 };
483 }
484 console.warn("[spec-to-pr] assertAiQuota failed unexpectedly:", err);
485 }
486
454487 const ownerName = repoRow.ownerName;
455488 const repoName = repoRow.name;
456489 const defaultBranch = repoRow.defaultBranch || "main";
457490