/** * Usage Fetcher - Get usage data from provider APIs */ import { getGitHubCopilotInternalUserHeaders, getKiroServiceHeaders, } from "@omniroute/open-sse/config/providerHeaderProfiles.ts"; import { applyAntigravityClientProfileHeaders } from "@omniroute/open-sse/services/antigravityClientProfile.ts"; import { getAntigravityHeaders } 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"; import { generateAntigravityRequestId, getAntigravitySessionId, } from "@omniroute/open-sse/services/antigravityIdentity.ts"; /** * Get usage data for a provider connection * @param {Object} connection - Provider connection with accessToken * @returns {Object} Usage data with quotas */ export async function getUsageForProvider(connection) { const { provider, accessToken, providerSpecificData } = connection; switch (provider) { case "github": return await getGitHubUsage(accessToken, providerSpecificData); case "antigravity": case "agy": return await getAntigravityUsage( accessToken, providerSpecificData, connection.projectId, connection.id ); case "claude": return await getClaudeUsage(accessToken); case "codex": return await getCodexUsage(accessToken, providerSpecificData); case "qwen": return await getQwenUsage(accessToken, providerSpecificData); case "qoder": return await getQoderUsage(accessToken); case "kiro": case "amazon-q": return await getKiroUsage(accessToken); default: return { message: `Usage API not implemented for ${provider}` }; } } /** * GitHub Copilot Usage */ async function getGitHubUsage(accessToken, providerSpecificData) { try { // Use copilotToken for copilot_internal API, not GitHub OAuth accessToken const copilotToken = providerSpecificData?.copilotToken; if (!copilotToken) { throw new Error("Copilot token not found. Please refresh token first."); } const response = await fetch("https://api.github.com/copilot_internal/user", { headers: getGitHubCopilotInternalUserHeaders(`Bearer ${copilotToken}`), }); if (!response.ok) { const error = await response.text(); throw new Error(`GitHub API error: ${error}`); } const data = await response.json(); // Handle different response formats (paid vs free) if (data.quota_snapshots) { // Paid plan format const snapshots = data.quota_snapshots; return { plan: data.copilot_plan, resetDate: data.quota_reset_date, quotas: { chat: formatGitHubQuotaSnapshot(snapshots.chat), completions: formatGitHubQuotaSnapshot(snapshots.completions), premium_interactions: formatGitHubQuotaSnapshot(snapshots.premium_interactions), }, }; } else if (data.monthly_quotas || data.limited_user_quotas) { // Free/limited plan format const monthlyQuotas = data.monthly_quotas || {}; const usedQuotas = data.limited_user_quotas || {}; return { plan: data.copilot_plan || data.access_type_sku, resetDate: data.limited_user_reset_date, quotas: { chat: { used: usedQuotas.chat || 0, total: monthlyQuotas.chat || 0, unlimited: false, }, completions: { used: usedQuotas.completions || 0, total: monthlyQuotas.completions || 0, unlimited: false, }, }, }; } return { message: "GitHub Copilot connected. Unable to parse quota data." }; } catch (error) { throw new Error(`Failed to fetch GitHub usage: ${error.message}`); } } function formatGitHubQuotaSnapshot(quota) { if (!quota) return { used: 0, total: 0, unlimited: true }; return { used: quota.entitlement - quota.remaining, total: quota.entitlement, remaining: quota.remaining, unlimited: quota.unlimited || false, }; } /** * 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, providerSpecificData: Record = {} ): 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: generateAntigravityRequestId(), enabledCreditTypes: ["GOOGLE_ONE_AI"], request: { model: "gemini-2-flash", contents: [{ role: "user", parts: [{ text: "hi" }] }], generationConfig: { maxOutputTokens: 1 }, sessionId: getAntigravitySessionId({ connectionId: accountId, projectId }), }, }; const headers: Record = { "Content-Type": "application/json", Authorization: `Bearer ${accessToken}`, Accept: "text/event-stream", }; applyAntigravityClientProfileHeaders( headers, { connectionId: accountId, projectId, providerSpecificData }, body ); 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 = {}, projectId?: string | null, connectionId?: string | null ) { try { // 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) 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, providerSpecificData ); } // fetchAvailableModels — resolves project from token, no projectId needed let res: Response | null = null; let lastError: Error | null = null; for (const endpoint of getAntigravityFetchAvailableModelsUrls()) { try { res = await fetch(endpoint, { method: "POST", headers: getAntigravityHeaders("fetchAvailableModels", accessToken), body: JSON.stringify({}), signal: AbortSignal.timeout(15_000), }); if (res.ok || res.status === 401 || res.status === 403) { break; } } catch (error) { lastError = error as Error; } } if (!res) { throw lastError || new Error("Antigravity API unavailable"); } if (!res.ok) { return { plan: "Antigravity", message: "Antigravity connected. Unable to fetch model quotas.", ...(creditBalance !== null && { quotas: { credits: { used: 0, total: 0, remaining: creditBalance, unlimited: false, resetAt: null, }, }, }), }; } const data = await res.json(); const models: Record = data?.models ?? {}; // Walk quota-based models (those with remainingFraction in quotaInfo) let quotaModelsTotal = 0; let quotaModelsAvailable = 0; const modelQuotas: Record< string, { remaining: number; resetAt: string | null; limited: boolean } > = {}; for (const [modelId, rawInfo] of Object.entries(models)) { const info = rawInfo as Record; if (info.isInternal) continue; const quotaInfo = (info.quotaInfo as Record) ?? {}; // Process models with any quota metadata (exhausted models may have only resetTime, active models have remainingFraction, credit-based models have empty quotaInfo) if (Object.keys(quotaInfo).length > 0) { // Default to 0 when remainingFraction is undefined/null/invalid, clamp to valid range [0, 1] const fraction = typeof quotaInfo.remainingFraction === "number" ? Math.max(0, Math.min(1, quotaInfo.remainingFraction)) : 0; const resetTime = typeof quotaInfo.resetTime === "string" ? quotaInfo.resetTime : null; modelQuotas[modelId] = { remaining: Math.round(fraction * 100), resetAt: resetTime, limited: fraction <= 0, }; quotaModelsTotal++; if (fraction > 0) quotaModelsAvailable++; } // Credit-based models have empty quotaInfo — their availability is // tracked via the GOOGLE_ONE_AI credit balance cached from SSE responses. } const allLimited = quotaModelsTotal > 0 && quotaModelsAvailable === 0; return { plan: "Antigravity", quotas: { models: { used: quotaModelsTotal - quotaModelsAvailable, total: quotaModelsTotal, remaining: quotaModelsAvailable, limited: allLimited, unlimited: false, resetAt: null, }, ...(creditBalance !== null && { credits: { used: 0, total: 0, remaining: creditBalance, unlimited: false, resetAt: null, }, }), }, modelQuotas, limitReached: allLimited, }; } catch (error) { return { message: `Unable to fetch Antigravity usage: ${(error as Error).message}` }; } } /** * Claude Usage (legacy fallback) * Real Claude OAuth quota windows are fetched in @omniroute/open-sse/services/usage.ts. */ async function getClaudeUsage(accessToken?: string) { try { return { message: "Claude connected. Detailed quota windows are handled by the open-sse usage service.", }; } catch (error) { return { message: "Unable to fetch Claude usage." }; } } /** * Codex (OpenAI) Usage * Note: Actual quota tracking is handled by open-sse/services/usage.ts * This fallback returns a message directing users to the dashboard. */ async function getCodexUsage(accessToken, providerSpecificData: Record = {}) { try { // Check if workspace is bound const workspaceId = providerSpecificData?.workspaceId; if (workspaceId) { return { message: `Codex connected (workspace: ${workspaceId.slice(0, 8)}...). Check dashboard for quota.`, }; } return { message: "Codex connected. Check OpenAI dashboard for usage." }; } catch (error) { return { message: "Unable to fetch Codex usage." }; } } /** * Qwen Usage */ async function getQwenUsage(accessToken, providerSpecificData) { try { const resourceUrl = providerSpecificData?.resourceUrl; if (!resourceUrl) { return { message: "Qwen connected. No resource URL available." }; } // Qwen may have usage endpoint at resource URL return { message: "Qwen connected. Usage tracked per request." }; } catch (error) { return { message: "Unable to fetch Qwen usage." }; } } /** * Qoder Usage */ async function getQoderUsage(accessToken) { try { // Qoder may have usage endpoint return { message: "Qoder connected. Usage tracked per request." }; } catch (error) { return { message: "Unable to fetch Qoder usage." }; } } /** * Kiro Credits * Fetches credit balance from Kiro's AWS CodeWhisperer backend. * The endpoint mirrors what Kiro IDE uses internally for the credit badge. */ async function getKiroUsage(accessToken: string) { try { const response = await fetch("https://codewhisperer.us-east-1.amazonaws.com/getUserCredits", { method: "GET", headers: { Authorization: `Bearer ${accessToken}`, ...getKiroServiceHeaders("application/json"), }, }); if (!response.ok) { const errText = await response.text(); // 401/403 = expired token, show user-friendly message if (response.status === 401 || response.status === 403) { return { message: "Kiro token expired. Please reconnect in Dashboard → Providers → Kiro." }; } throw new Error(`Kiro credits API error (${response.status}): ${errText}`); } const data = await response.json(); // Response shape: { remainingCredits, totalCredits, resetDate, subscriptionType } const remaining = data.remainingCredits ?? data.remaining_credits ?? null; const total = data.totalCredits ?? data.total_credits ?? null; const resetDate = data.resetDate ?? data.reset_date ?? null; const plan = data.subscriptionType ?? data.subscription_type ?? "unknown"; if (remaining === null) { return { message: "Kiro connected. Credit data unavailable — check Kiro IDE for balance." }; } return { plan, credits: { remaining, total: total ?? remaining, used: total != null ? total - remaining : 0, unlimited: total === null || total === 0, resetDate, }, }; } catch (error: any) { return { message: `Unable to fetch Kiro credits: ${error.message}` }; } }