From bfa5e424096f9eb106ca4772ff3b78294107edb0 Mon Sep 17 00:00:00 2001 From: ghost <49853598+JSONbored@users.noreply.github.com> Date: Sat, 13 Jun 2026 23:24:10 -0700 Subject: [PATCH] fix(ai-review): cap BYOK daily review usage --- src/db/repositories.ts | 18 ++++++++++++++++++ src/env.d.ts | 2 ++ src/services/ai-review.ts | 19 ++++++++++++++++--- test/unit/ai-review.test.ts | 19 +++++++++++++++++-- wrangler.jsonc | 1 + 5 files changed, 54 insertions(+), 5 deletions(-) diff --git a/src/db/repositories.ts b/src/db/repositories.ts index 2fb6c8141f..3019388345 100644 --- a/src/db/repositories.ts +++ b/src/db/repositories.ts @@ -1954,6 +1954,24 @@ export async function sumAiEstimatedNeuronsSince(env: Env, sinceIso: string): Pr return Number(row?.total ?? 0); } +export async function countByokAiReviewEventsForRepoSince(env: Env, repoFullName: string, sinceIso: string): Promise { + const db = getDb(env.DB); + const [row] = await db + .select({ total: sql`count(*)` }) + .from(aiUsageEvents) + .where( + and( + gte(aiUsageEvents.createdAt, sinceIso), + eq(aiUsageEvents.feature, "ai_review_pr"), + eq(aiUsageEvents.status, "ok"), + sql`${aiUsageEvents.model} like 'byok:%'`, + sql`json_extract(${aiUsageEvents.metadataJson}, '$.repoFullName') = ${repoFullName}`, + ), + ); + /* v8 ignore next -- SQL aggregate count always returns one row; fallback protects D1 driver anomalies. */ + return Number(row?.total ?? 0); +} + export async function upsertContributorScoringProfile(env: Env, profile: ContributorScoringProfileRecord): Promise { const db = getDb(env.DB); await db diff --git a/src/env.d.ts b/src/env.d.ts index e54c1462a5..c89d83a88c 100644 --- a/src/env.d.ts +++ b/src/env.d.ts @@ -10,6 +10,8 @@ declare global { AI_PUBLIC_COMMENTS_ENABLED?: string; WORKERS_AI_SUMMARY_MODEL?: string; AI_DAILY_NEURON_BUDGET?: string; + /** Per-repository/day cap for maintainer-paid BYOK AI review provider calls. */ + AI_BYOK_DAILY_REPO_LIMIT?: string; AI_MAX_OUTPUT_TOKENS?: string; ADMIN_GITHUB_LOGINS?: string; GITHUB_WEBHOOK_SECRET: string; diff --git a/src/services/ai-review.ts b/src/services/ai-review.ts index c6eba0de28..04e6c10877 100644 --- a/src/services/ai-review.ts +++ b/src/services/ai-review.ts @@ -13,9 +13,10 @@ // confirmed Gittensor contributors (the gate enforces that downstream). // // Every public string (notes + defect title/detail) is forced through `sanitizePublicComment`; anything -// that trips the public/private boundary is dropped, not published. Every model call is metered against -// the shared daily neuron budget and audited via `recordAiUsageEvent`. -import { recordAiUsageEvent, sumAiEstimatedNeuronsSince } from "../db/repositories"; +// that trips the public/private boundary is dropped, not published. Free Workers-AI calls are metered against +// the shared daily neuron budget; maintainer-paid BYOK calls have a separate repo/day cap. All calls +// are audited via `recordAiUsageEvent`. +import { countByokAiReviewEventsForRepoSince, recordAiUsageEvent, sumAiEstimatedNeuronsSince } from "../db/repositories"; import { sanitizePublicComment } from "../queue-intelligence"; /** @@ -217,6 +218,9 @@ const PROVIDER_DEFAULT_MODEL: Record = * the existing fail-safe null path. Mirrors the github/gittensor fetch-timeout convention. */ const AI_PROVIDER_TIMEOUT_MS = 20_000; +/** Default per-repository/day cap for maintainer-paid BYOK advisory calls. */ +const DEFAULT_BYOK_DAILY_REPO_LIMIT = 25; + /** Why a BYOK advisory call produced no review — surfaced in the audit event for observability (never a key). */ type ProviderFailure = "timeout" | "http_error" | "exception"; type ProviderReviewOutcome = { review: ModelReview | null; failure?: ProviderFailure }; @@ -318,6 +322,15 @@ export async function runGittensoryAiReview(env: Env, input: GittensoryAiReviewI return { status: "quota_exceeded", estimatedNeurons, remainingBudget }; } + if (input.providerKey) { + const byokDailyLimit = clampNumber(Number(env.AI_BYOK_DAILY_REPO_LIMIT || DEFAULT_BYOK_DAILY_REPO_LIMIT), 0, 10_000); + const byokUsed = await countByokAiReviewEventsForRepoSince(env, input.repoFullName, utcDayStartIso()); + if (byokUsed >= byokDailyLimit) { + await record(env, input, "quota_exceeded", 0, `BYOK daily repo limit ${byokDailyLimit} reached`); + return { status: "quota_exceeded", estimatedNeurons, remainingBudget }; + } + } + // Advisory write-up: BYOK frontier model if configured, else the free Workers-AI primary (with fallback). let byokFailure: ProviderFailure | undefined; let advisoryReview: ModelReview | null; diff --git a/test/unit/ai-review.test.ts b/test/unit/ai-review.test.ts index f3cf6347cb..c04e45a1ee 100644 --- a/test/unit/ai-review.test.ts +++ b/test/unit/ai-review.test.ts @@ -64,8 +64,9 @@ describe("runGittensoryAiReview gating", () => { const fetchMock = vi.fn(async () => new Response(JSON.stringify({ content: [{ type: "text", text: reviewJson({ assessment: "BYOK advisory." }) }] }), { status: 200 })); vi.stubGlobal("fetch", fetchMock); const run = vi.fn(); - // Free budget is exhausted (1 neuron), but a BYOK advisory bills the maintainer's account, so it still runs. - const env = createTestEnv({ AI: { run } as unknown as Ai, AI_SUMMARIES_ENABLED: "true", AI_PUBLIC_COMMENTS_ENABLED: "true", AI_DAILY_NEURON_BUDGET: "1" }); + // Free budget is exhausted (1 neuron), but a BYOK advisory bills the maintainer's account, so it still runs + // while the separate BYOK repo/day quota has capacity. + const env = createTestEnv({ AI: { run } as unknown as Ai, AI_SUMMARIES_ENABLED: "true", AI_PUBLIC_COMMENTS_ENABLED: "true", AI_DAILY_NEURON_BUDGET: "1", AI_BYOK_DAILY_REPO_LIMIT: "1" }); const result = await runGittensoryAiReview(env, { ...baseInput, providerKey: { provider: "anthropic", key: "sk-ant-secret" } }); expect(result.status).toBe("ok"); expect(result.status === "ok" && result.advisoryNotes).toContain("BYOK advisory."); @@ -73,6 +74,20 @@ describe("runGittensoryAiReview gating", () => { expect(fetchMock).toHaveBeenCalled(); expect(run).not.toHaveBeenCalled(); }); + + it("enforces a separate per-repo daily quota before BYOK provider calls", async () => { + const fetchMock = vi.fn(async () => new Response(JSON.stringify({ content: [{ type: "text", text: reviewJson({ assessment: "BYOK advisory." }) }] }), { status: 200 })); + vi.stubGlobal("fetch", fetchMock); + const run = vi.fn(); + const env = createTestEnv({ AI: { run } as unknown as Ai, AI_SUMMARIES_ENABLED: "true", AI_PUBLIC_COMMENTS_ENABLED: "true", AI_DAILY_NEURON_BUDGET: "0", AI_BYOK_DAILY_REPO_LIMIT: "1" }); + const providerKey = { provider: "anthropic" as const, key: "sk-ant-secret" }; + + await expect(runGittensoryAiReview(env, { ...baseInput, providerKey })).resolves.toMatchObject({ status: "ok" }); + await expect(runGittensoryAiReview(env, { ...baseInput, prNumber: 8, providerKey })).resolves.toMatchObject({ status: "quota_exceeded" }); + + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(run).not.toHaveBeenCalled(); + }); }); describe("runGittensoryAiReview advisory mode", () => { diff --git a/wrangler.jsonc b/wrangler.jsonc index e5b295f628..73cd23bcda 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -38,6 +38,7 @@ "AI_PUBLIC_COMMENTS_ENABLED": "false", "WORKERS_AI_SUMMARY_MODEL": "@cf/meta/llama-3.1-8b-instruct-fp8-fast", "AI_DAILY_NEURON_BUDGET": "10000", + "AI_BYOK_DAILY_REPO_LIMIT": "25", "AI_MAX_OUTPUT_TOKENS": "256", "ADMIN_GITHUB_LOGINS": "JSONbored", },