diff --git a/open-sse/executors/antigravity.ts b/open-sse/executors/antigravity.ts index f8d5223854..83eeabe720 100644 --- a/open-sse/executors/antigravity.ts +++ b/open-sse/executors/antigravity.ts @@ -7,8 +7,11 @@ import { classify429, decide429, type Decision } from "../services/antigravity42 import { injectCreditsField, shouldRetryWithCredits, + shouldUseCreditsFirst, + getCreditsMode, handleCreditsFailure, } from "../services/antigravityCredits.ts"; +import { persistCreditBalance, getAllPersistedCreditBalances } from "@/lib/db/creditBalance"; import { obfuscateSensitiveWords } from "../services/antigravityObfuscation.ts"; const MAX_RETRY_AFTER_MS = 60_000; @@ -28,11 +31,30 @@ const creditsExhaustedUntil = new Map(); * Per-account GOOGLE_ONE_AI remaining credit balance cache. * Populated from the final SSE chunk's `remainingCredits` field after every * successful credit-injected request. Keyed by accountId. + * On first access, hydrated from the DB-persisted balances so values survive restarts. */ const creditBalanceCache = new Map(); +let creditCacheHydrated = false; + +function hydrateCreditCacheFromDb(): void { + if (creditCacheHydrated) return; + creditCacheHydrated = true; + try { + const persisted = getAllPersistedCreditBalances(); + for (const [accountId, balance] of persisted) { + // Only fill in accounts not already populated by a live SSE response + if (!creditBalanceCache.has(accountId)) { + creditBalanceCache.set(accountId, balance); + } + } + } catch { + // DB not ready yet (build phase, etc.) — ignore silently + } +} /** Read the last-known GOOGLE_ONE_AI credit balance for a given account. */ export function getAntigravityRemainingCredits(accountId: string): number | null { + hydrateCreditCacheFromDb(); const balance = creditBalanceCache.get(accountId); return balance !== undefined ? balance : null; } @@ -40,6 +62,12 @@ export function getAntigravityRemainingCredits(accountId: string): number | null /** Update the balance cache — called when we parse `remainingCredits` from an SSE stream. */ export function updateAntigravityRemainingCredits(accountId: string, balance: number): void { creditBalanceCache.set(accountId, balance); + // Persist to DB so the value survives server restarts + try { + persistCreditBalance(accountId, balance); + } catch { + // Non-critical — in-memory cache is the primary source + } } function isCreditsExhausted(accountId: string): boolean { @@ -426,17 +454,31 @@ export class AntigravityExecutor extends BaseExecutor { // non-streaming Response so chatCore's non-streaming path stays unchanged. const upstreamStream = true; - // Account ID for credits-exhausted tracking. - // Key must match getAntigravityUsage() in fetcher.ts (providerSpecificData?.email || sub). - // credentials.email and credentials.sub are populated from the same OAuth token store, - // so the cache keys written here and read in the fetcher will always match. - const accountId: string = credentials?.email || credentials?.sub || "unknown"; + // Account ID for credits tracking. + // Use connectionId as the stable cache key — it's available in both the executor + // (via credentials.connectionId) and the usage fetcher (via connection.id). + // The email-based key was unreliable because email isn't always on the credentials object. + const accountId: string = credentials?.connectionId || "unknown"; + + // Resolve credits mode once per execute() call. "always" injects + // enabledCreditTypes: ["GOOGLE_ONE_AI"] on the first request so the + // preflight normal call is skipped entirely. + const creditsMode = getCreditsMode(); + const useCreditsFirst = shouldUseCreditsFirst(credentials?.accessToken || "", creditsMode); for (let urlIndex = 0; urlIndex < fallbackCount; urlIndex++) { const url = this.buildUrl(model, upstreamStream, urlIndex); const headers = this.buildHeaders(credentials, upstreamStream); mergeUpstreamExtraHeaders(headers, upstreamExtraHeaders); - const transformedBody = await this.transformRequest(model, body, upstreamStream, credentials); + let transformedBody = await this.transformRequest(model, body, upstreamStream, credentials); + + // Credits-first: inject GOOGLE_ONE_AI upfront so we never try the normal + // quota path. If credits are exhausted / disabled shouldUseCreditsFirst() + // returns false and we fall back to the legacy retry-on-429 flow. + if (useCreditsFirst) { + transformedBody = injectCreditsField(transformedBody); + log?.debug?.("AG_CREDITS", "Credits-first enabled (ANTIGRAVITY_CREDITS=always)"); + } // Initialize retry counter for this URL if (!retryAttemptsByUrl[urlIndex]) { @@ -471,17 +513,29 @@ export class AntigravityExecutor extends BaseExecutor { // 1. Try to parse explicit retry time from message const parsedRetryMs = this.parseRetryFromErrorMessage(errorMessage); - // 2. Classify 429 - const category = classify429(errorMessage); + // 2. Classify 429 (pass header-parsed retry hint as fallback + // signal — multi-hour Retry-After upgrades rate_limited to + // quota_exhausted so the GOOGLE_ONE_AI credits retry fires). + const effectiveRetryHintMs = retryMs ?? parsedRetryMs ?? null; + const category = classify429(errorMessage, effectiveRetryHintMs); // 3. For quota_exhausted, attempt Google One AI credits retry FIRST! + // Skip if credits were already injected on the first call + // (creditsMode === "always") — no point re-running with the + // same body. Record the failure so the 5h breaker kicks in. + const creditsAlreadyInjected = + (transformedBody as { enabledCreditTypes?: unknown }).enabledCreditTypes != null; + + if (category === "quota_exhausted" && creditsAlreadyInjected) { + handleCreditsFailure(credentials?.accessToken || ""); + log?.warn?.("AG_CREDITS", "Credits-first request 429'd — credits likely exhausted"); + markCreditsExhausted(accountId); + } + if ( category === "quota_exhausted" && - shouldRetryWithCredits( - credentials?.accessToken || "", - process.env.ANTIGRAVITY_CREDITS === "1" || - process.env.ANTIGRAVITY_CREDITS === "true" - ) + !creditsAlreadyInjected && + shouldRetryWithCredits(credentials?.accessToken || "", creditsMode !== "off") ) { log?.info?.("AG_CREDITS", "Retrying with Google One AI credits"); const creditsBody = injectCreditsField(transformedBody); @@ -627,7 +681,7 @@ export class AntigravityExecutor extends BaseExecutor { // For non-streaming clients, collect the SSE stream and return a synthetic // non-streaming Response so chatCore doesn't need to handle SSE conversion. if (!stream) { - return this.collectStreamToResponse( + const collected = await this.collectStreamToResponse( response, model, url, @@ -636,6 +690,82 @@ export class AntigravityExecutor extends BaseExecutor { log, signal ); + // When credits were injected (credits-first or credits-retry), the + // synthetic body contains _remainingCredits — mirror it into the + // balance cache so the dashboard stays fresh. + try { + const syntheticJson = await collected.response.clone().json(); + const rc = syntheticJson?._remainingCredits; + if (Array.isArray(rc)) { + const googleCredit = rc.find( + (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI" + ); + if (googleCredit) { + const balance = parseInt(googleCredit.creditAmount, 10); + if (!isNaN(balance)) updateAntigravityRemainingCredits(accountId, balance); + } + } + } catch { + /* balance cache is best-effort */ + } + return collected; + } + + // Streaming path: wrap the response body in a pass-through TransformStream + // that extracts remainingCredits from the final SSE chunk(s) without + // consuming the stream. The client receives the unmodified SSE data. + if (response.body) { + let sseBuffer = ""; + const passThrough = new TransformStream({ + transform(chunk, controller) { + controller.enqueue(chunk); + // Accumulate text to scan for remainingCredits + try { + const text = new TextDecoder().decode(chunk, { stream: true }); + sseBuffer += text; + } catch { + /* decoding best-effort */ + } + }, + flush() { + // Parse the accumulated SSE data for remainingCredits + try { + const lines = sseBuffer.split("\n"); + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith("data:")) continue; + const payload = trimmed.slice(5).trim(); + if (!payload || payload === "[DONE]") continue; + try { + const parsed = JSON.parse(payload); + if (Array.isArray(parsed?.remainingCredits)) { + const googleCredit = parsed.remainingCredits.find( + (c) => c?.creditType === "GOOGLE_ONE_AI" + ); + if (googleCredit) { + const balance = parseInt(googleCredit.creditAmount, 10); + if (!isNaN(balance)) { + updateAntigravityRemainingCredits(accountId, balance); + } + } + } + } catch { + /* skip malformed lines */ + } + } + } catch { + /* credits extraction is best-effort */ + } + sseBuffer = ""; + }, + }); + const tappedBody = response.body.pipeThrough(passThrough); + const tappedResponse = new Response(tappedBody, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }); + return { response: tappedResponse, url, headers, transformedBody }; } return { response, url, headers, transformedBody }; diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index af1816ed8c..d3b8d5123c 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -3,15 +3,24 @@ */ import { PROVIDERS } from "../config/constants.ts"; -import { getAntigravityFetchAvailableModelsUrls } from "../config/antigravityUpstream.ts"; +import { + getAntigravityFetchAvailableModelsUrls, + ANTIGRAVITY_BASE_URLS, +} from "../config/antigravityUpstream.ts"; import { getGlmQuotaUrl } from "../config/glmProvider.ts"; import { safePercentage } from "@/shared/utils/formatting"; import { fetchBailianQuota, type BailianTripleWindowQuota } from "./bailianQuotaFetcher.ts"; import { antigravityUserAgent, + googApiClientHeader, getAntigravityHeaders, getAntigravityLoadCodeAssistMetadata, } from "./antigravityHeaders.ts"; +import { + getAntigravityRemainingCredits, + updateAntigravityRemainingCredits, +} from "../executors/antigravity.ts"; +import { getCreditsMode } from "./antigravityCredits.ts"; // GitHub API config const GITHUB_CONFIG = { @@ -208,7 +217,7 @@ async function getBailianCodingPlanUsage( * @returns {Promise} Usage data with quotas */ export async function getUsageForProvider(connection) { - const { id, provider, accessToken, apiKey, providerSpecificData, projectId } = connection; + const { id, provider, accessToken, apiKey, providerSpecificData, projectId, email } = connection; switch (provider) { case "github": @@ -216,7 +225,7 @@ export async function getUsageForProvider(connection) { case "gemini-cli": return await getGeminiUsage(accessToken, providerSpecificData, projectId); case "antigravity": - return await getAntigravityUsage(accessToken, undefined); + return await getAntigravityUsage(accessToken, providerSpecificData, projectId, id); case "claude": return await getClaudeUsage(accessToken); case "codex": @@ -881,16 +890,129 @@ function getAntigravityPlanLabel(subscriptionInfo) { return "Free"; } +/** + * Proactive credit balance probe for Antigravity. + * + * Fires a minimal streamGenerateContent request with GOOGLE_ONE_AI credits enabled + * and maxOutputTokens=1 to extract the `remainingCredits` field from the SSE stream. + * This uses ~1 credit but lets us show the balance on the dashboard without waiting + * for a real user request. + * + * Returns the credit balance, or null if the probe failed. + */ +async function probeAntigravityCreditBalance( + accessToken: string, + accountId: string, + projectId?: string | null +): Promise { + try { + if (!projectId) return null; + + // Try all base URLs (some accounts only work with specific endpoints) + for (const baseUrl of ANTIGRAVITY_BASE_URLS) { + const url = `${baseUrl}/v1internal:streamGenerateContent?alt=sse`; + + const sessionId = `-${Math.floor(Math.random() * 9_000_000_000_000_000_000)}`; + const body = { + project: projectId, + model: "gemini-2-flash", + userAgent: "antigravity", + requestType: "agent", + requestId: `credits-probe-${Date.now()}`, + enabledCreditTypes: ["GOOGLE_ONE_AI"], + request: { + model: "gemini-2-flash", + contents: [{ role: "user", parts: [{ text: "hi" }] }], + generationConfig: { maxOutputTokens: 1 }, + sessionId, + }, + }; + + const headers = { + "Content-Type": "application/json", + Authorization: `Bearer ${accessToken}`, + "User-Agent": antigravityUserAgent(), + "X-Goog-Api-Client": googApiClientHeader(), + Accept: "text/event-stream", + }; + + try { + const res = await fetch(url, { + method: "POST", + headers, + body: JSON.stringify(body), + signal: AbortSignal.timeout(10_000), + }); + + if (!res.ok) continue; + + // Read the full SSE response and scan for remainingCredits + const rawSSE = await res.text(); + const lines = rawSSE.split("\n"); + + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith("data:")) continue; + const payload = trimmed.slice(5).trim(); + if (payload === "[DONE]") break; + try { + const parsed = JSON.parse(payload); + if (Array.isArray(parsed?.remainingCredits)) { + const googleCredit = parsed.remainingCredits.find( + (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI" + ); + if (googleCredit) { + const balance = parseInt(googleCredit.creditAmount, 10); + if (!isNaN(balance)) { + updateAntigravityRemainingCredits(accountId, balance); + return balance; + } + } + } + } catch { + // Skip malformed SSE lines + } + } + } catch { + // Individual endpoint failure; try next + } + } + + return null; + } catch { + // Probe is best-effort — don't let it break the usage fetch + return null; + } +} + /** * Antigravity Usage - Fetch quota from Google Cloud Code API * Uses fetchAvailableModels API which returns ALL models (including Claude) * with per-model quotaInfo (remainingFraction, resetTime). * retrieveUserQuota only returns Gemini models — not suitable for Antigravity. */ -async function getAntigravityUsage(accessToken, providerSpecificData) { +async function getAntigravityUsage( + accessToken, + providerSpecificData, + connectionProjectId?, + connectionId? +) { try { const subscriptionInfo = await getAntigravitySubscriptionInfoCached(accessToken); - const projectId = subscriptionInfo?.cloudaicompanionProject || null; + const projectId = connectionProjectId || subscriptionInfo?.cloudaicompanionProject || null; + + // Derive accountId for credit balance cache. + // Must match executor key: credentials.connectionId + const accountId: string = connectionId || "unknown"; + + // Read cached credit balance (hydrated from DB on first access) + let creditBalance = getAntigravityRemainingCredits(accountId); + + // If no cached balance and credits mode is enabled, fire a minimal probe + const creditsMode = getCreditsMode(); + if (creditBalance === null && creditsMode !== "off") { + creditBalance = await probeAntigravityCreditBalance(accessToken, accountId, projectId); + } // Fetch model list with quota info from fetchAvailableModels let response: Response | null = null; @@ -987,7 +1109,18 @@ async function getAntigravityUsage(accessToken, providerSpecificData) { return { plan: getAntigravityPlanLabel(subscriptionInfo), - quotas, + quotas: { + ...quotas, + ...(creditBalance !== null && { + credits: { + used: 0, + total: 0, + remaining: creditBalance, + unlimited: false, + resetAt: null, + }, + }), + }, subscriptionInfo, }; } catch (error) { diff --git a/src/lib/db/creditBalance.ts b/src/lib/db/creditBalance.ts new file mode 100644 index 0000000000..b495633877 --- /dev/null +++ b/src/lib/db/creditBalance.ts @@ -0,0 +1,92 @@ +/** + * Antigravity GOOGLE_ONE_AI credit balance persistence. + * + * Stores last-known credit balances in the `key_value` table so they survive + * server restarts. The in-memory `creditBalanceCache` in the executor is the + * primary source at runtime; this module provides the fallback read/write layer. + */ + +import { getDbInstance, isBuildPhase, isCloud } from "./core"; + +interface StatementLike { + get: (...params: unknown[]) => TRow | undefined; + run: (...params: unknown[]) => { changes?: number }; + all: (...params: unknown[]) => TRow[]; +} + +interface DbLike { + prepare: (sql: string) => StatementLike; +} + +interface KeyValueRow { + key: string; + value: string; +} + +interface CreditBalanceEntry { + balance: number; + updatedAt: string; +} + +const NAMESPACE = "antigravityCreditBalance"; + +function parseJson(raw: string): unknown { + try { + return JSON.parse(raw); + } catch { + return null; + } +} + +/** + * Read the persisted credit balance for an accountId. + * Returns the balance number, or null if not found. + */ +export function getPersistedCreditBalance(accountId: string): number | null { + if (isBuildPhase || isCloud) return null; + const db = getDbInstance() as unknown as DbLike; + const row = db + .prepare("SELECT value FROM key_value WHERE namespace = ? AND key = ?") + .get(NAMESPACE, accountId) as KeyValueRow | undefined; + if (!row?.value) return null; + const parsed = parseJson(row.value) as CreditBalanceEntry | null; + if (!parsed || typeof parsed.balance !== "number") return null; + return parsed.balance; +} + +/** + * Read all persisted credit balances. + * Returns a Map of accountId → balance. + */ +export function getAllPersistedCreditBalances(): Map { + const result = new Map(); + if (isBuildPhase || isCloud) return result; + const db = getDbInstance() as unknown as DbLike; + const rows = db + .prepare("SELECT key, value FROM key_value WHERE namespace = ?") + .all(NAMESPACE) as KeyValueRow[]; + for (const row of rows) { + const parsed = parseJson(row.value) as CreditBalanceEntry | null; + if (parsed && typeof parsed.balance === "number") { + result.set(row.key, parsed.balance); + } + } + return result; +} + +/** + * Persist a credit balance for an accountId. + */ +export function persistCreditBalance(accountId: string, balance: number): void { + if (isBuildPhase || isCloud) return; + const db = getDbInstance() as unknown as DbLike; + const entry: CreditBalanceEntry = { + balance, + updatedAt: new Date().toISOString(), + }; + db.prepare("INSERT OR REPLACE INTO key_value (namespace, key, value) VALUES (?, ?, ?)").run( + NAMESPACE, + accountId, + JSON.stringify(entry) + ); +} diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index 8917f8ba8b..eebbc80604 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -267,3 +267,9 @@ export { } from "./db/providerLimits"; export type { ProviderLimitsCacheEntry } from "./db/providerLimits"; + +export { + getPersistedCreditBalance, + getAllPersistedCreditBalances, + persistCreditBalance, +} from "./db/creditBalance"; diff --git a/src/lib/usage/fetcher.ts b/src/lib/usage/fetcher.ts index 16abdceac8..11591a080c 100644 --- a/src/lib/usage/fetcher.ts +++ b/src/lib/usage/fetcher.ts @@ -3,9 +3,20 @@ */ import { GITHUB_CONFIG, GEMINI_CONFIG } from "@/lib/oauth/constants/oauth"; -import { getAntigravityHeaders } from "@omniroute/open-sse/services/antigravityHeaders.ts"; -import { getAntigravityFetchAvailableModelsUrls } from "@omniroute/open-sse/config/antigravityUpstream.ts"; -import { getAntigravityRemainingCredits } from "@omniroute/open-sse/executors/antigravity.ts"; +import { + getAntigravityHeaders, + antigravityUserAgent, + googApiClientHeader, +} from "@omniroute/open-sse/services/antigravityHeaders.ts"; +import { + getAntigravityFetchAvailableModelsUrls, + ANTIGRAVITY_BASE_URLS, +} from "@omniroute/open-sse/config/antigravityUpstream.ts"; +import { + getAntigravityRemainingCredits, + updateAntigravityRemainingCredits, +} from "@omniroute/open-sse/executors/antigravity.ts"; +import { getCreditsMode } from "@omniroute/open-sse/services/antigravityCredits.ts"; /** * Get usage data for a provider connection @@ -21,7 +32,12 @@ export async function getUsageForProvider(connection) { case "gemini-cli": return await getGeminiUsage(accessToken); case "antigravity": - return await getAntigravityUsage(accessToken, providerSpecificData); + return await getAntigravityUsage( + accessToken, + providerSpecificData, + connection.projectId, + connection.id + ); case "claude": return await getClaudeUsage(accessToken); case "codex": @@ -147,23 +163,120 @@ async function getGeminiUsage(accessToken) { } } +/** + * Proactive credit balance probe for Antigravity. + * + * Fires a minimal streamGenerateContent request with GOOGLE_ONE_AI credits enabled + * and maxOutputTokens=1 to extract the `remainingCredits` field from the SSE stream. + * This uses ~1 credit but lets us show the balance on the dashboard without waiting + * for a real user request. + * + * Returns the credit balance, or null if the probe failed. + */ +async function probeAntigravityCreditBalance( + accessToken: string, + accountId: string, + projectId?: string | null +): Promise { + try { + if (!projectId) return null; // Can't call streamGenerateContent without a projectId + + const baseUrl = ANTIGRAVITY_BASE_URLS[0]; + const url = `${baseUrl}/v1internal:streamGenerateContent?alt=sse`; + + const body = { + project: projectId, + model: "gemini-2-flash", + userAgent: "antigravity", + requestType: "agent", + requestId: `credits-probe-${Date.now()}`, + enabledCreditTypes: ["GOOGLE_ONE_AI"], + request: { + model: "gemini-2-flash", + contents: [{ role: "user", parts: [{ text: "hi" }] }], + generationConfig: { maxOutputTokens: 1 }, + }, + }; + + const headers = { + "Content-Type": "application/json", + Authorization: `Bearer ${accessToken}`, + "User-Agent": antigravityUserAgent(), + "X-Goog-Api-Client": googApiClientHeader(), + Accept: "text/event-stream", + }; + + const res = await fetch(url, { + method: "POST", + headers, + body: JSON.stringify(body), + signal: AbortSignal.timeout(10_000), + }); + + if (!res.ok) return null; + + // Read the full SSE response and scan for remainingCredits + const rawSSE = await res.text(); + const lines = rawSSE.split("\n"); + + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith("data:")) continue; + const payload = trimmed.slice(5).trim(); + if (payload === "[DONE]") break; + try { + const parsed = JSON.parse(payload); + if (Array.isArray(parsed?.remainingCredits)) { + const googleCredit = parsed.remainingCredits.find( + (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI" + ); + if (googleCredit) { + const balance = parseInt(googleCredit.creditAmount, 10); + if (!isNaN(balance)) { + // Cache the balance for future reads (also persists to DB) + updateAntigravityRemainingCredits(accountId, balance); + return balance; + } + } + } + } catch { + // Skip malformed SSE lines + } + } + + return null; + } catch { + // Probe is best-effort — don't let it break the usage fetch + return null; + } +} + /** * Antigravity Usage * Calls fetchAvailableModels to get per-model quota fractions. * Credit balance (GOOGLE_ONE_AI) is read from the executor's in-memory cache, * which is populated automatically after each successful credit-injected SSE call. + * If the cache is empty and credits mode is enabled, fires a minimal probe request + * to fetch the balance proactively. */ async function getAntigravityUsage( accessToken: string, - providerSpecificData: Record = {} + providerSpecificData: Record = {}, + projectId?: string | null, + connectionId?: string | null ) { try { - // Derive accountId (same key used in AntigravityExecutor.execute) - const accountId: string = - (providerSpecificData?.email as string) || (providerSpecificData?.sub as string) || "unknown"; + // Use connectionId as the cache key — matches executor's credentials.connectionId + const accountId: string = connectionId || "unknown"; // Read cached credit balance from executor module (populated from SSE remainingCredits) - const creditBalance = getAntigravityRemainingCredits(accountId); + let creditBalance = getAntigravityRemainingCredits(accountId); + + // If no cached balance and credits mode is enabled, fire a minimal probe + const creditsMode = getCreditsMode(); + if (creditBalance === null && creditsMode !== "off") { + creditBalance = await probeAntigravityCreditBalance(accessToken, accountId, projectId); + } // fetchAvailableModels — resolves project from token, no projectId needed let res: Response | null = null;