diff --git a/src/db/repositories.ts b/src/db/repositories.ts index cd04903282..903f90e169 100644 --- a/src/db/repositories.ts +++ b/src/db/repositories.ts @@ -2169,7 +2169,12 @@ 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 { +/** + * Count a repo's maintainer-billed (BYOK) AI calls since `sinceIso`, across ALL AI features (review + + * slop + any future BYOK path). One shared per-repo/day budget governs every BYOK feature, so a repo + * cannot multiply its frontier-model spend by enabling more capabilities. + */ +export async function countByokAiEventsForRepoSince(env: Env, repoFullName: string, sinceIso: string): Promise { const db = getDb(env.DB); const [row] = await db .select({ total: sql`count(*)` }) @@ -2177,7 +2182,6 @@ export async function countByokAiReviewEventsForRepoSince(env: Env, repoFullName .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}`, diff --git a/src/queue/processors.ts b/src/queue/processors.ts index 1e14ef753a..d037f2ed17 100644 --- a/src/queue/processors.ts +++ b/src/queue/processors.ts @@ -936,6 +936,7 @@ export async function runAiReviewForAdvisory( export async function runAiSlopForAdvisory( env: Env, args: { + settings: RepositorySettings; advisory: Awaited>; repoFullName: string; pr: { number: number; title: string; body?: string | null | undefined }; @@ -946,6 +947,14 @@ export async function runAiSlopForAdvisory( ): Promise { if (!args.advisory.headSha) return; try { + // BYOK (opt-in): reuse the repo's encrypted key + aiReviewByok flag — one BYOK key serves both AI + // features. A declared provider must match the stored key's provider, else skip BYOK (Workers-AI + // fallback). The slop advisory stays advisory-only regardless of which model writes it. + const storedKey = args.settings.aiReviewByok ? await getDecryptedRepositoryAiKey(env, args.repoFullName) : null; + const providerKey = + storedKey && (!args.settings.aiReviewProvider || args.settings.aiReviewProvider === storedKey.provider) + ? { provider: storedKey.provider, key: storedKey.key, model: args.settings.aiReviewModel ?? storedKey.model } + : null; const result = await runGittensoryAiSlopAdvisory(env, { repoFullName: args.repoFullName, prNumber: args.pr.number, @@ -954,6 +963,7 @@ export async function runAiSlopForAdvisory( diff: buildAiReviewDiff(args.files), actor: args.author, deterministicBand: args.deterministicBand, + providerKey, }); if (result.status === "ok" && result.finding) args.advisory.findings.push(result.finding); } catch (error) { @@ -1159,7 +1169,7 @@ async function maybePublishPrPublicSurface( // AI-assisted slop advisory (#533, opt-in). Reuses the already-fetched files; appends at most one // advisory-only finding. Deliberately does NOT update slopRisk — only the deterministic core blocks. if (settings.slopAiAdvisory) { - await runAiSlopForAdvisory(env, { advisory, repoFullName, pr, author, files: slopFiles, deterministicBand: slop.band }); + await runAiSlopForAdvisory(env, { settings, advisory, repoFullName, pr, author, files: slopFiles, deterministicBand: slop.band }); } } diff --git a/src/services/ai-review.ts b/src/services/ai-review.ts index 2e7b15c8ab..0c4ce3a88c 100644 --- a/src/services/ai-review.ts +++ b/src/services/ai-review.ts @@ -16,7 +16,7 @@ // 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 { countByokAiEventsForRepoSince, recordAiUsageEvent, sumAiEstimatedNeuronsSince } from "../db/repositories"; import { sanitizePublicComment } from "../queue-intelligence"; /** @@ -230,16 +230,24 @@ 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; +/** Default per-repository/day cap for maintainer-paid BYOK calls (shared across all BYOK AI features). */ +export 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"; +/** Why a BYOK call produced no usable output — surfaced in the audit event for observability (never a key). */ +export type ProviderFailure = "timeout" | "http_error" | "exception"; type ProviderReviewOutcome = { review: ModelReview | null; failure?: ProviderFailure }; -/** Run the maintainer's BYOK frontier model for the advisory write-up. Never throws; the review is null on - * any error and `failure` names the reason (timeout/http_error/exception) for the audit trail. */ -async function runProviderReview(providerKey: AiReviewProviderKey, system: string, user: string, maxTokens: number): Promise { +/** + * POST to the maintainer's BYOK provider and return the raw response text (or null + a failure reason). + * Never throws. Shared by every BYOK AI path (review, slop, …) so the endpoint/timeout/error handling + * lives in one place; callers parse the returned text into their own shape. + */ +export async function callAiProvider( + providerKey: AiReviewProviderKey, + system: string, + user: string, + maxTokens: number, +): Promise<{ text: string | null; failure?: ProviderFailure }> { const model = providerKey.model || PROVIDER_DEFAULT_MODEL[providerKey.provider]; try { let response: Response; @@ -265,15 +273,22 @@ async function runProviderReview(providerKey: AiReviewProviderKey, system: strin signal: AbortSignal.timeout(AI_PROVIDER_TIMEOUT_MS), }); } - if (!response.ok) return { review: null, failure: "http_error" }; - return { review: parseModelReview(coerceAiText(await response.json())) }; + if (!response.ok) return { text: null, failure: "http_error" }; + return { text: coerceAiText(await response.json()) }; } catch (error) { // AbortSignal.timeout rejects with a TimeoutError; everything else is a network/parse exception. const failure: ProviderFailure = (error as { name?: string } | null)?.name === "TimeoutError" ? "timeout" : "exception"; - return { review: null, failure }; + return { text: null, failure }; } } +/** Run the maintainer's BYOK frontier model for the advisory write-up. Never throws; the review is null on + * any error and `failure` names the reason (timeout/http_error/exception) for the audit trail. */ +async function runProviderReview(providerKey: AiReviewProviderKey, system: string, user: string, maxTokens: number): Promise { + const { text, failure } = await callAiProvider(providerKey, system, user, maxTokens); + return { review: text ? parseModelReview(text) : null, ...(failure ? { failure } : {}) }; +} + /** Compose a public-safe markdown advisory blurb from one or two model reviews. Null if nothing safe. */ export function composeAdvisoryNotes(reviews: ModelReview[]): string | null { const assessments = reviews.map((r) => r.assessment).filter(Boolean); @@ -336,7 +351,7 @@ export async function runGittensoryAiReview(env: Env, input: GittensoryAiReviewI 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()); + const byokUsed = await countByokAiEventsForRepoSince(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 }; diff --git a/src/services/ai-slop.ts b/src/services/ai-slop.ts index 1f69cf9f9f..2efdac9904 100644 --- a/src/services/ai-slop.ts +++ b/src/services/ai-slop.ts @@ -17,10 +17,13 @@ // forced through `toPublicSafe`; anything tripping the public/private boundary is dropped, not published. import type { SignalFinding } from "../signals/engine"; import type { SlopBand } from "../signals/slop"; -import { recordAiUsageEvent, sumAiEstimatedNeuronsSince } from "../db/repositories"; +import { countByokAiEventsForRepoSince, recordAiUsageEvent, sumAiEstimatedNeuronsSince } from "../db/repositories"; import { + type AiReviewProviderKey, BEST_REVIEW_MODELS, + DEFAULT_BYOK_DAILY_REPO_LIMIT, RELIABLE_FALLBACK_MODELS, + callAiProvider, clampNumber, coerceAiText, estimateNeurons, @@ -59,6 +62,10 @@ export type AiSlopInput = { /** The deterministic band already computed for this PR — passed as context so the model can corroborate * or temper it. Never used to override the model's own judgement. */ deterministicBand?: SlopBand | undefined; + /** Optional BYOK: when present, the maintainer's frontier model writes the advisory (billed to their + * account, counted against the shared per-repo/day BYOK cap) instead of free Workers AI. Advisory-only + * either way — BYOK never changes whether this can block (it can't). */ + providerKey?: AiReviewProviderKey | null | undefined; }; export type AiSlopResult = @@ -173,7 +180,10 @@ export async function runGittensoryAiSlopAdvisory(env: Env, input: AiSlopInput): const maxTokens = clampNumber(Number(env.AI_MAX_OUTPUT_TOKENS || 256), 256, 1024); const user = buildUserPrompt(input); - const estimatedNeurons = estimateNeurons(SLOP_SYSTEM_PROMPT.length + user.length, maxTokens, 1); + // BYOK bills the maintainer's own account, so it does NOT draw on the free neuron budget — it has a + // separate per-repo/day cap shared with the AI review path. Free Workers-AI = one metered call. + const freeCalls = input.providerKey ? 0 : 1; + const estimatedNeurons = freeCalls === 0 ? 0 : estimateNeurons(SLOP_SYSTEM_PROMPT.length + user.length, maxTokens, 1); const budget = clampNumber(Number(env.AI_DAILY_NEURON_BUDGET || 10000), 0, 1_000_000); const used = await sumAiEstimatedNeuronsSince(env, utcDayStartIso()); const remainingBudget = Math.max(0, budget - used); @@ -181,12 +191,28 @@ export async function runGittensoryAiSlopAdvisory(env: Env, input: AiSlopInput): await record(env, input, "quota_exceeded", 0, `estimated ${estimatedNeurons} neurons exceeds remaining ${remainingBudget}`); 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 countByokAiEventsForRepoSince(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 }; + } + } - const opinion = await runWorkersSlopOpinion(env, SLOP_SYSTEM_PROMPT, user, maxTokens); + // BYOK frontier model if configured, else the free Workers-AI primary (with fallback). Both fail-safe to null. + let opinion: SlopOpinion | null; + if (input.providerKey) { + const { text } = await callAiProvider(input.providerKey, SLOP_SYSTEM_PROMPT, user, maxTokens); + opinion = text ? parseSlopOpinion(text) : null; + } else { + opinion = await runWorkersSlopOpinion(env, SLOP_SYSTEM_PROMPT, user, maxTokens); + } const finding = opinion ? slopFindingFromOpinion(opinion) : null; await record(env, input, "ok", estimatedNeurons, finding ? `advisory finding (${opinion?.band})` : opinion ? `clean/no-op (${opinion.band})` : "no usable output", { band: opinion?.band ?? null, surfaced: Boolean(finding), + byok: Boolean(input.providerKey), }); return { status: "ok", finding, band: opinion?.band ?? null, estimatedNeurons }; } @@ -196,7 +222,8 @@ async function record(env: Env, input: AiSlopInput, status: string, estimatedNeu feature: "ai_slop_pr", actor: input.actor ?? null, route: "github_app.ai_slop", - model: BEST_REVIEW_MODELS.join("+"), + // `byok:` so countByokAiEventsForRepoSince (model LIKE 'byok:%') counts it toward the cap. + model: input.providerKey ? `byok:${input.providerKey.provider}` : BEST_REVIEW_MODELS.join("+"), status, estimatedNeurons, detail, diff --git a/test/unit/ai-slop.test.ts b/test/unit/ai-slop.test.ts index 784e0c9dc5..7df1b63413 100644 --- a/test/unit/ai-slop.test.ts +++ b/test/unit/ai-slop.test.ts @@ -7,7 +7,8 @@ import { } from "../../src/services/ai-slop"; import { evaluateGateCheck } from "../../src/rules/advisory"; import { runAiSlopForAdvisory } from "../../src/queue/processors"; -import type { Advisory, PullRequestFileRecord } from "../../src/types"; +import { recordAiUsageEvent, upsertRepositoryAiKey } from "../../src/db/repositories"; +import type { Advisory, PullRequestFileRecord, RepositorySettings } from "../../src/types"; import { createTestEnv } from "../helpers/d1"; const { parseSlopOpinion, slopFindingFromOpinion, buildUserPrompt } = __aiSlopInternals; @@ -187,6 +188,20 @@ describe("runGittensoryAiSlopAdvisory gating + fail-safe", () => { if (result.status !== "ok") throw new Error("unreachable"); expect(result.band).toBe("low"); }); + + it("enforces the shared BYOK daily repo cap before any provider call (BYOK does not draw on the free budget)", async () => { + 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: "1", AI_BYOK_DAILY_REPO_LIMIT: "1" }); + // Seed one prior BYOK event for this repo so the cap (1) is already reached. + await recordAiUsageEvent(env, { feature: "ai_review_pr", actor: null, route: "x", model: "byok:anthropic", status: "ok", estimatedNeurons: 1, detail: "seed", metadata: { repoFullName: baseInput.repoFullName } }); + const fetchMock = vi.fn(async (_url: RequestInfo | URL, _init?: RequestInit) => new Response("{}", { status: 200 })); + vi.stubGlobal("fetch", fetchMock); + // Free budget is exhausted (1 neuron) but BYOK skips it; the BYOK cap is what stops the call. + const result = await runGittensoryAiSlopAdvisory(env, { ...baseInput, providerKey: { provider: "anthropic", key: "sk-ant-x" } }); + expect(result.status).toBe("quota_exceeded"); + expect(fetchMock).not.toHaveBeenCalled(); + expect(run).not.toHaveBeenCalled(); + }); }); describe("the AI slop advisory can never become a gate blocker", () => { @@ -248,10 +263,12 @@ describe("runAiSlopForAdvisory (processor wiring)", () => { { repoFullName: "acme/widgets", pullNumber: 3, path: "src/a.ts", status: "modified", additions: 80, deletions: 2, changes: 82, payload: { patch: "@@\n+// set x\n+const x = 1;" } }, ]; const pr = { number: 3, title: "Tidy", body: "cleanup" }; + const noByok = { aiReviewByok: false } as RepositorySettings; it("appends a single ai_slop_advisory finding when the model flags slop", async () => { const adv = advisory(); await runAiSlopForAdvisory(enabledEnv(async () => ({ response: slopJson({ band: "high" }) })), { + settings: noByok, advisory: adv, repoFullName: "acme/widgets", pr, @@ -266,7 +283,7 @@ describe("runAiSlopForAdvisory (processor wiring)", () => { const noSha = advisory(); delete (noSha as Partial).headSha; const run = vi.fn(); - await runAiSlopForAdvisory(enabledEnv(run), { advisory: noSha, repoFullName: "acme/widgets", pr, author: "alice", files, deterministicBand: "low" }); + await runAiSlopForAdvisory(enabledEnv(run), { settings: noByok, advisory: noSha, repoFullName: "acme/widgets", pr, author: "alice", files, deterministicBand: "low" }); expect(noSha.findings).toEqual([]); expect(run).not.toHaveBeenCalled(); }); @@ -274,6 +291,7 @@ describe("runAiSlopForAdvisory (processor wiring)", () => { it("adds nothing when the model judges the change clean", async () => { const adv = advisory(); await runAiSlopForAdvisory(enabledEnv(async () => ({ response: slopJson({ band: "clean", rationale: "genuine", signals: [] }) })), { + settings: noByok, advisory: adv, repoFullName: "acme/widgets", pr, @@ -287,7 +305,35 @@ describe("runAiSlopForAdvisory (processor wiring)", () => { it("is fail-safe: a thrown error (broken DB) yields no finding and never throws", async () => { const adv = advisory(); const env = { ...enabledEnv(async () => ({ response: slopJson() })), DB: undefined } as unknown as Env; - await expect(runAiSlopForAdvisory(env, { advisory: adv, repoFullName: "acme/widgets", pr, author: "alice", files, deterministicBand: "high" })).resolves.toBeUndefined(); + await expect(runAiSlopForAdvisory(env, { settings: noByok, advisory: adv, repoFullName: "acme/widgets", pr, author: "alice", files, deterministicBand: "high" })).resolves.toBeUndefined(); expect(adv.findings).toEqual([]); }); + + it("uses the maintainer's BYOK frontier model (not Workers AI) when aiReviewByok is on and a key is configured", async () => { + const run = vi.fn(async () => ({ response: slopJson({ band: "clean" }) })); // Workers AI must NOT be used + const env = createTestEnv({ + AI: { run } as unknown as Ai, + AI_SUMMARIES_ENABLED: "true", + AI_PUBLIC_COMMENTS_ENABLED: "true", + AI_DAILY_NEURON_BUDGET: "100000", + TOKEN_ENCRYPTION_SECRET: "ai-slop-byok-test-encryption-secret-32b", + }); + await upsertRepositoryAiKey(env, { repoFullName: "acme/widgets", provider: "anthropic", key: "sk-ant-byok-slop-9999", model: null }); + const fetchMock = vi.fn(async (_url: RequestInfo | URL, _init?: RequestInit) => new Response(JSON.stringify({ content: [{ type: "text", text: slopJson({ band: "high" }) }] }), { status: 200 })); + vi.stubGlobal("fetch", fetchMock); + const adv = advisory(); + await runAiSlopForAdvisory(env, { + settings: { aiReviewByok: true } as RepositorySettings, + advisory: adv, + repoFullName: "acme/widgets", + pr, + author: "alice", + files, + deterministicBand: "elevated", + }); + // The advisory came from the BYOK provider (high band → finding), and Workers AI was never called. + expect(adv.findings.map((f) => f.code)).toEqual([AI_SLOP_FINDING_CODE]); + expect(fetchMock.mock.calls[0]?.[0]).toBe("https://api.anthropic.com/v1/messages"); + expect(run).not.toHaveBeenCalled(); + }); });