diff --git a/src/db/repositories.ts b/src/db/repositories.ts index 871997fa81..8da21ddf19 100644 --- a/src/db/repositories.ts +++ b/src/db/repositories.ts @@ -3456,10 +3456,24 @@ export async function sumAiEstimatedNeuronsSince(env: Env, sinceIso: string): Pr return Number(row?.total ?? 0); } +/** Spend-attempt statuses `countByokAiEventsForRepoSince`/`sumByokAiUsageForRepoSince` count: a real request + * reached the provider, whether or not it returned something usable ("ok") or genuinely failed ("error" -- + * timeout/http_error/exception, see e.g. queue/processors.ts's recordVisualVisionUsage). Deliberately an + * ALLOWLIST, not an exclusion of "quota_exceeded": `ai_usage_events` is also reused for BYOK key-lifecycle + * audit rows (recordAiKeyChange's "set"/"replace"/"delete", `model` also `byok:`-prefixed so they + * match this query's model filter too) -- an exclusion-based filter would silently start counting those + * (or any future non-spend status added to this shared table) as spend. */ +const BYOK_SPEND_ATTEMPT_STATUSES = ["ok", "error"] as const; + /** * 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. + * + * Counts every ATTEMPTED call, not just ones tagged "ok" -- a caller that records a distinct "error" status + * for a genuine provider failure still made a real request against the maintainer's key, so it must still + * count; excluding attempted-but-failed calls would turn a flaky or misconfigured provider into a way to + * bypass this cap entirely via forced failures. See BYOK_SPEND_ATTEMPT_STATUSES for why this is an allowlist. */ export async function countByokAiEventsForRepoSince(env: Env, repoFullName: string, sinceIso: string): Promise { const db = getDb(env.DB); @@ -3469,7 +3483,7 @@ export async function countByokAiEventsForRepoSince(env: Env, repoFullName: stri .where( and( gte(aiUsageEvents.createdAt, sinceIso), - eq(aiUsageEvents.status, "ok"), + inArray(aiUsageEvents.status, BYOK_SPEND_ATTEMPT_STATUSES), sql`${aiUsageEvents.model} like 'byok:%'`, sql`json_extract(${aiUsageEvents.metadataJson}, '$.repoFullName') = ${repoFullName}`, ), @@ -3493,6 +3507,10 @@ export async function sumByokAiUsageForRepoSince( sinceIso: string, ): Promise<{ calls: number; inputTokens: number; outputTokens: number; totalTokens: number; costUsd: number }> { const db = getDb(env.DB); + // Mirrors countByokAiEventsForRepoSince's own WHERE clause (see BYOK_SPEND_ATTEMPT_STATUSES's doc comment) + // -- an attempted-but-failed call still counts as a "call" for reporting purposes, same as it counts toward + // the daily cap. A failed attempt's usage columns are all 0/null (no billable usage was ever returned), so + // it contributes to `calls` but not to the token/cost sums. const [row] = await db .select({ calls: sql`count(*)`, @@ -3505,7 +3523,7 @@ export async function sumByokAiUsageForRepoSince( .where( and( gte(aiUsageEvents.createdAt, sinceIso), - eq(aiUsageEvents.status, "ok"), + inArray(aiUsageEvents.status, BYOK_SPEND_ATTEMPT_STATUSES), 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 464b580757..e6e62bb058 100644 --- a/src/queue/processors.ts +++ b/src/queue/processors.ts @@ -12,6 +12,7 @@ import { getRepoAuthorPullRequestHistory, getRepository, getDecryptedRepositoryAiKey, + countByokAiEventsForRepoSince, getRepositorySettings, listCheckSummaries, listAllIssues, @@ -79,6 +80,7 @@ import { terminalizeActiveReviewTracking, bumpPullRequestDraftConversionCount, recordProductUsageEvent, + recordAiUsageEvent, persistSignalSnapshot, recordWebhookEvent, replaceCollisionEdges, @@ -448,9 +450,13 @@ import { isRepoDocRefreshDue } from "../review/repo-doc-refresh-schedule"; import type { LocalBranchAnalysisInput } from "../signals/local-branch"; import { callAiProvider, + clampNumber, + DEFAULT_BYOK_DAILY_REPO_LIMIT, hasPublicReviewAssessment, isEnabled, runGittensoryAiReview, + utcDayStartIso, + type AiReviewActualUsage, type InlineFinding, } from "../services/ai-review"; import { @@ -8558,6 +8564,31 @@ export async function runVisualVisionForAdvisory( // false case: if neither is set here, the gate itself would already have returned run:false above. /* v8 ignore next 2 -- see comment above */ if (!visionProviderKey && !selfHostVisionAvailable) return; + // BYOK (a maintainer's own anthropic/openai key) takes priority when both are configured -- matches every + // other dual-path AI call site's convention (BYOK bills the maintainer's own account, so it's preferred + // over the shared/free local resource when the operator has explicitly set one up). Only the BYOK branch + // is metered/capped below -- self-host vision consumes the operator's own resources, already gated + // separately by selfHostVisionAllowed above, and was never part of the BYOK daily-spend surface. The cap + // check runs BEFORE the shot-fetching loop so a repo that's already over budget never even pays for the + // screenshot fetches, not just the provider call. + if (visionProviderKey) { + const byokDailyLimit = clampNumber( + Number(env.AI_BYOK_DAILY_REPO_LIMIT || DEFAULT_BYOK_DAILY_REPO_LIMIT), + 0, + 10_000, + ); + const byokUsed = await countByokAiEventsForRepoSince(env, args.repoFullName, utcDayStartIso()); + if (byokUsed >= byokDailyLimit) { + await recordVisualVisionUsage( + env, + args, + visionProviderKey, + "quota_exceeded", + "BYOK daily repo limit reached", + ); + return; + } + } const images: AiContentBlock[] = []; for (const route of visionGate.routes) { // Show the model the viewport that actually crossed the pixel-diff threshold — a route can qualify via @@ -8576,19 +8607,46 @@ export async function runVisualVisionForAdvisory( if (afterBlock) images.push(afterBlock); } if (images.length === 0) return; - // BYOK (a maintainer's own anthropic/openai key) takes priority when both are configured -- matches - // every other dual-path AI call site's convention (BYOK bills the maintainer's own account, so it's - // preferred over the shared/free local resource when the operator has explicitly set one up). let visionText: string | null; + let visionUsage: AiReviewActualUsage | undefined; if (visionProviderKey) { const visionResponse = await callAiProvider(visionProviderKey, VISUAL_VISION_SYSTEM_PROMPT, buildVisualVisionUserPrompt(visionGate.routes), 600, images); visionText = visionResponse.text; + visionUsage = visionResponse.usage; + if (!visionText) { + // "error" (not "ok") when the provider call itself failed (timeout/http_error/exception) -- matches + // runAgentSummary's convention (services/ai-summaries.ts) of a distinct status for a genuine call + // failure vs. a call that completed but returned nothing usable. countByokAiEventsForRepoSince + // deliberately still counts "error" rows toward the daily cap (it only excludes "quota_exceeded", + // not "ok" specifically) -- a repo hitting a flaky/misconfigured provider must not get a free, + // uncapped retry budget just because every attempt happens to fail. + await recordVisualVisionUsage( + env, + args, + visionProviderKey, + visionResponse.failure ? "error" : "ok", + visionResponse.failure ? `provider failure: ${String(visionResponse.failure)}` : "no usable output", + visionResponse.usage, + ); + return; + } } else { visionText = await runSelfHostVisualVision(env, VISUAL_VISION_SYSTEM_PROMPT, buildVisualVisionUserPrompt(visionGate.routes), images); } if (!visionText) return; const visionFindings = parseVisualVisionResponse(visionText); - args.advisory.findings.push(...buildVisualRegressionFindings(visionFindings)); + const findings = buildVisualRegressionFindings(visionFindings); + args.advisory.findings.push(...findings); + if (visionProviderKey) { + await recordVisualVisionUsage( + env, + args, + visionProviderKey, + "ok", + findings.length > 0 ? `advisory findings (${findings.length})` : "no usable output", + visionUsage, + ); + } } catch (error) { console.log( JSON.stringify({ @@ -8601,6 +8659,32 @@ export async function runVisualVisionForAdvisory( } } +async function recordVisualVisionUsage( + env: Env, + args: { repoFullName: string; pr: { number: number }; author: string | null }, + providerKey: { provider: string }, + status: string, + detail: string, + usage?: AiReviewActualUsage | undefined, +): Promise { + await recordAiUsageEvent(env, { + feature: "visual_vision", + actor: args.author ?? null, + route: "github_app.visual_vision", + model: `byok:${providerKey.provider}`, + status, + estimatedNeurons: 0, + provider: usage?.provider, + effort: usage?.effort, + inputTokens: usage?.inputTokens, + outputTokens: usage?.outputTokens, + totalTokens: usage?.totalTokens, + costUsd: usage?.costUsd, + detail, + metadata: { repoFullName: args.repoFullName, pullNumber: args.pr.number }, + }); +} + /** * Resolve `manifest_missing_tests`' `passedValidationCount` signal (gate-review finding, #4719): a PR-body * validation-note match (`hasValidationNote`) is checked FIRST since it's free; only when that misses, AND diff --git a/test/unit/visual-vision-wiring.test.ts b/test/unit/visual-vision-wiring.test.ts index 9acb156981..dfbb9d3332 100644 --- a/test/unit/visual-vision-wiring.test.ts +++ b/test/unit/visual-vision-wiring.test.ts @@ -1,10 +1,11 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { runVisualVisionForAdvisory } from "../../src/queue/processors"; import * as repositories from "../../src/db/repositories"; -import { upsertRepositoryAiKey } from "../../src/db/repositories"; +import { countByokAiEventsForRepoSince, upsertRepositoryAiKey } from "../../src/db/repositories"; import * as submitterReputation from "../../src/review/submitter-reputation"; import type { CaptureRoute } from "../../src/review/visual/capture"; import type { AdvisoryFinding, RepositorySettings } from "../../src/types"; +import { utcDayStartIso } from "../../src/services/ai-review"; import { createTestEnv } from "../helpers/d1"; afterEach(() => { @@ -235,6 +236,83 @@ describe("runVisualVisionForAdvisory", () => { expect(fetchMock.mock.calls.map((c) => String(c[0]))).toEqual(["https://api.gittensor.io/miners"]); }); + it("enforces the shared BYOK daily cap before fetching screenshots or calling the vision provider", async () => { + const env = createTestEnv({ + TOKEN_ENCRYPTION_SECRET: "vision-test-encryption-secret-32b", + AI_BYOK_DAILY_REPO_LIMIT: "0", + }); + await upsertRepositoryAiKey(env, { repoFullName, provider: "anthropic", key: "sk-ant-vision-key", model: null }); + // #4513: the reputation/miner-identity check at the top of runVisualVisionForAdvisory runs regardless of + // the BYOK cap outcome -- only the shot fetches and the provider call are gated by the cap. + const fetchMock = stubMinerCheckOnly(); + vi.stubGlobal("fetch", fetchMock); + const adv = findingsHolder(); + await runVisualVisionForAdvisory(env, { + mode: "live", + repoFullName, + pr, + author: "alice", + confirmedContributor: true, + settings: byokSettings(), + advisory: adv, + routes: [ + route({ + path: "/app", + diffUrl: "https://x/gittensory/shot?key=diff", + beforeUrl: "https://x/gittensory/shot?key=before", + afterUrl: "https://x/gittensory/shot?key=after", + }), + ], + }); + expect(fetchMock.mock.calls.map((c) => String(c[0]))).toEqual(["https://api.gittensor.io/miners"]); + expect(adv.findings).toEqual([]); + expect(await countByokAiEventsForRepoSince(env, repoFullName, utcDayStartIso())).toBe(0); + }); + + // REGRESSION guard: `ai_usage_events` is shared with BYOK key-lifecycle audit rows (recordAiKeyChange's + // "set"/"replace"/"delete", src/db/repositories.ts) whose `model` is ALSO `byok:`-prefixed, so + // they match this query's model filter too -- only their `status` (never "ok" or "error") keeps them out. + // upsertRepositoryAiKey (used by nearly every test in this file to seed a BYOK key) always writes exactly + // one such "set" row, so this asserts it alone never counts toward the cap. + it("does not count a BYOK key-lifecycle audit event (upsertRepositoryAiKey's own 'set' row) toward the daily cap", async () => { + const env = byokEnv(); + await upsertRepositoryAiKey(env, { repoFullName, provider: "anthropic", key: "sk-ant-vision-key", model: null }); + const keyChangeEvents = await env.DB.prepare("select status, feature, model from ai_usage_events").all<{ status: string; feature: string; model: string }>(); + expect(keyChangeEvents.results).toEqual([{ status: "set", feature: "ai_key_change", model: "byok:anthropic" }]); + expect(await countByokAiEventsForRepoSince(env, repoFullName, utcDayStartIso())).toBe(0); + }); + + it("records successful visual BYOK calls so later passes count toward the shared daily cap", async () => { + const env = byokEnv(); + await upsertRepositoryAiKey(env, { + repoFullName, + provider: "anthropic", + key: "sk-ant-vision-key", + model: null, + }); + stubShotsAndProvider(findingsResponse([])); + const adv = findingsHolder(); + await runVisualVisionForAdvisory(env, { + mode: "live", + repoFullName, + pr, + author: "alice", + confirmedContributor: true, + settings: byokSettings(), + advisory: adv, + routes: [ + route({ + path: "/app", + diffUrl: "https://x/gittensory/shot?key=diff", + beforeUrl: "https://x/gittensory/shot?key=before", + afterUrl: "https://x/gittensory/shot?key=after", + }), + ], + }); + expect(adv.findings).toEqual([]); + expect(await countByokAiEventsForRepoSince(env, repoFullName, utcDayStartIso())).toBe(1); + }); + it("calls the BYOK vision provider with before+after images and publishes a returned finding (desktop route)", async () => { const env = byokEnv(); await upsertRepositoryAiKey(env, { repoFullName, provider: "anthropic", key: "sk-ant-vision-key", model: null }); @@ -396,7 +474,7 @@ describe("runVisualVisionForAdvisory", () => { expect(adv.findings).toEqual([]); }); - it("adds no finding when the provider call itself fails (non-2xx) -- callAiProvider's own fail-safe", async () => { + it("adds no finding when the provider call itself fails (non-2xx) -- callAiProvider's own fail-safe, but STILL records the attempt as a distinct 'error' status that counts toward the daily cap", async () => { const env = byokEnv(); await upsertRepositoryAiKey(env, { repoFullName, provider: "anthropic", key: "sk-ant-vision-key", model: null }); stubShotsAndProvider(null); @@ -412,6 +490,37 @@ describe("runVisualVisionForAdvisory", () => { routes: [route({ path: "/app", diffUrl: "https://x/gittensory/shot?key=diff", beforeUrl: "https://x/gittensory/shot?key=before", afterUrl: "https://x/gittensory/shot?key=after" })], }); expect(adv.findings).toEqual([]); + // A genuine provider failure is a distinct "error" status (not "ok") -- but it's still a real request + // against the maintainer's key, so it must still count toward the shared daily cap (see + // BYOK_SPEND_ATTEMPT_STATUSES's doc comment, src/db/repositories.ts): a repo hitting a flaky/misconfigured + // provider must not get unlimited free retries just because every attempt happens to fail. + const events = await env.DB.prepare("select status, detail from ai_usage_events where feature = 'visual_vision'").all<{ status: string; detail: string }>(); + expect(events.results).toEqual([{ status: "error", detail: "provider failure: http_error" }]); + expect(await countByokAiEventsForRepoSince(env, repoFullName, utcDayStartIso())).toBe(1); + }); + + it("adds no finding when the provider returns 200 with no usable text (distinct from an http_error failure)", async () => { + const env = byokEnv(); + await upsertRepositoryAiKey(env, { repoFullName, provider: "anthropic", key: "sk-ant-vision-key", model: null }); + // An empty string is a genuine 2xx response, unlike stubShotsAndProvider(null)'s 500 -- callAiProvider + // returns { text: "", failure: undefined } here (no "http_error"), exercising the "no usable output" + // fallback in recordVisualVisionUsage's detail message rather than the provider-failure one. Also uses a + // null author (ghost/deleted account, `args.author ?? undefined` short-circuits the reputation/miner + // check to neutral with no fetch) to exercise recordVisualVisionUsage's own `actor: args.author ?? null` + // fallback alongside it. + stubShotsAndProvider(""); + const adv = findingsHolder(); + await runVisualVisionForAdvisory(env, { + mode: "live", + repoFullName, + pr, + author: null, + confirmedContributor: true, + settings: byokSettings(), + advisory: adv, + routes: [route({ path: "/app", diffUrl: "https://x/gittensory/shot?key=diff", beforeUrl: "https://x/gittensory/shot?key=before", afterUrl: "https://x/gittensory/shot?key=after" })], + }); + expect(adv.findings).toEqual([]); }); it("swallows a thrown error from the BYOK key lookup and never lets it escape (visual_vision_error)", async () => {