From 205b5ab71fe560cdc8eed046ae133d40343df156 Mon Sep 17 00:00:00 2001 From: kaitranntt Date: Sun, 28 Dec 2025 19:24:12 -0500 Subject: [PATCH] feat(cliproxy): add account quota display for Antigravity provider - Add quota-fetcher.ts for Google Cloud Code API integration - Add REST endpoint GET /api/cliproxy/quota/:provider/:accountId - Add useAccountQuota hook with React Query caching - Display quota bar in AccountItem with per-model tooltip - Create reusable Progress component - Consolidate quota types in api-client.ts --- src/cliproxy/index.ts | 4 + src/cliproxy/quota-fetcher.ts | 313 ++++++++++++++++++ .../routes/cliproxy-stats-routes.ts | 48 +++ .../cliproxy/provider-editor/account-item.tsx | 209 +++++++++--- .../provider-editor/accounts-section.tsx | 4 + .../cliproxy/provider-editor/index.tsx | 1 + .../provider-editor/model-config-tab.tsx | 4 + .../cliproxy/provider-editor/types.ts | 2 + ui/src/components/ui/progress.tsx | 40 +++ ui/src/hooks/use-cliproxy-stats.ts | 31 ++ ui/src/lib/api-client.ts | 32 ++ 11 files changed, 637 insertions(+), 51 deletions(-) create mode 100644 src/cliproxy/quota-fetcher.ts create mode 100644 ui/src/components/ui/progress.tsx diff --git a/src/cliproxy/index.ts b/src/cliproxy/index.ts index e42705be..5e327180 100644 --- a/src/cliproxy/index.ts +++ b/src/cliproxy/index.ts @@ -111,6 +111,10 @@ export { export type { CliproxyStats } from './stats-fetcher'; export { fetchCliproxyStats, isCliproxyRunning } from './stats-fetcher'; +// Quota fetcher +export type { ModelQuota, QuotaResult } from './quota-fetcher'; +export { fetchAccountQuota } from './quota-fetcher'; + // OpenAI compatibility layer export type { OpenAICompatProvider, OpenAICompatModel } from './openai-compat-manager'; export { diff --git a/src/cliproxy/quota-fetcher.ts b/src/cliproxy/quota-fetcher.ts new file mode 100644 index 00000000..f4d7c05b --- /dev/null +++ b/src/cliproxy/quota-fetcher.ts @@ -0,0 +1,313 @@ +/** + * Quota Fetcher for Antigravity Accounts + * + * Fetches quota information from Google Cloud Code internal API. + * Used for displaying remaining quota percentages and reset times. + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import { getAuthDir } from './config-generator'; +import { CLIProxyProvider } from './types'; + +/** Individual model quota info */ +export interface ModelQuota { + /** Model name, e.g., "gemini-3-pro-high" */ + name: string; + /** Display name from API, e.g., "Gemini 3 Pro" */ + displayName?: string; + /** Remaining quota as percentage (0-100) */ + percentage: number; + /** ISO timestamp when quota resets, null if unknown */ + resetTime: string | null; +} + +/** Quota fetch result */ +export interface QuotaResult { + /** Whether fetch succeeded */ + success: boolean; + /** Quota for each available model */ + models: ModelQuota[]; + /** Timestamp of fetch */ + lastUpdated: number; + /** True if account lacks quota access (403) */ + isForbidden?: boolean; + /** Error message if fetch failed */ + error?: string; +} + +/** Google Cloud Code API endpoints */ +const ANTIGRAVITY_API_BASE = 'https://cloudcode-pa.googleapis.com'; +const ANTIGRAVITY_API_VERSION = 'v1internal'; + +/** API client headers */ +const ANTIGRAVITY_HEADERS = { + 'Content-Type': 'application/json', + 'User-Agent': 'antigravity/1.11.5 linux/amd64', + 'X-Goog-Api-Client': 'gl-node/20.9.0', +}; + +/** Auth file structure */ +interface AntigravityAuthFile { + access_token: string; + refresh_token?: string; + email?: string; + expired?: string; + expires_in?: number; + timestamp?: number; + type?: string; +} + +/** loadCodeAssist response */ +interface LoadCodeAssistResponse { + cloudaicompanionProject?: string | { id?: string }; +} + +/** fetchAvailableModels response model */ +interface AvailableModel { + name?: string; + displayName?: string; + quotaInfo?: { + remainingFraction?: number; + remaining_fraction?: number; + remaining?: number; + resetTime?: string; + reset_time?: string; + }; + quota_info?: { + remainingFraction?: number; + remaining_fraction?: number; + remaining?: number; + resetTime?: string; + reset_time?: string; + }; +} + +/** fetchAvailableModels response */ +interface FetchAvailableModelsResponse { + models?: Record; +} + +/** + * Read access token from auth file + */ +function readAccessToken(provider: CLIProxyProvider, accountId: string): string | null { + const authDir = getAuthDir(); + + // Account ID format: email with @ and . replaced by _ + // Try to find matching token file + const files = fs.readdirSync(authDir); + const prefix = provider === 'agy' ? 'antigravity-' : `${provider}-`; + + for (const file of files) { + if (file.startsWith(prefix) && file.endsWith('.json')) { + // Check if this file matches the account ID + const baseName = file.replace(prefix, '').replace('.json', ''); + if (baseName === accountId || file === accountId || file === `${accountId}.json`) { + const filePath = path.join(authDir, file); + try { + const content = fs.readFileSync(filePath, 'utf-8'); + const data = JSON.parse(content) as AntigravityAuthFile; + return data.access_token || null; + } catch { + return null; + } + } + } + } + + return null; +} + +/** + * Get project ID via loadCodeAssist endpoint + */ +async function getProjectId(accessToken: string): Promise { + const url = `${ANTIGRAVITY_API_BASE}/${ANTIGRAVITY_API_VERSION}:loadCodeAssist`; + + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), 5000); + + try { + const response = await fetch(url, { + method: 'POST', + signal: controller.signal, + headers: { + ...ANTIGRAVITY_HEADERS, + Authorization: `Bearer ${accessToken}`, + }, + body: JSON.stringify({ + metadata: { + ideType: 'IDE_UNSPECIFIED', + platform: 'PLATFORM_UNSPECIFIED', + pluginType: 'GEMINI', + }, + }), + }); + + clearTimeout(timeoutId); + + if (!response.ok) { + return null; + } + + const data = (await response.json()) as LoadCodeAssistResponse; + + // Extract project ID from response + let projectId: string | undefined; + if (typeof data.cloudaicompanionProject === 'string') { + projectId = data.cloudaicompanionProject; + } else if (typeof data.cloudaicompanionProject === 'object') { + projectId = data.cloudaicompanionProject?.id; + } + + return projectId?.trim() || null; + } catch { + clearTimeout(timeoutId); + return null; + } +} + +/** + * Fetch available models with quota info + */ +async function fetchAvailableModels(accessToken: string, projectId: string): Promise { + const url = `${ANTIGRAVITY_API_BASE}/${ANTIGRAVITY_API_VERSION}:fetchAvailableModels`; + + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), 5000); + + try { + const response = await fetch(url, { + method: 'POST', + signal: controller.signal, + headers: { + ...ANTIGRAVITY_HEADERS, + Authorization: `Bearer ${accessToken}`, + }, + body: JSON.stringify({ + project: projectId, + }), + }); + + clearTimeout(timeoutId); + + if (response.status === 403) { + return { + success: false, + models: [], + lastUpdated: Date.now(), + isForbidden: true, + error: 'Quota access forbidden for this account', + }; + } + + if (response.status === 401) { + return { + success: false, + models: [], + lastUpdated: Date.now(), + error: 'Access token expired or invalid', + }; + } + + if (!response.ok) { + return { + success: false, + models: [], + lastUpdated: Date.now(), + error: `API error: ${response.status}`, + }; + } + + const data = (await response.json()) as FetchAvailableModelsResponse; + const models: ModelQuota[] = []; + + if (data.models && typeof data.models === 'object') { + for (const [modelId, modelData] of Object.entries(data.models)) { + const quotaInfo = modelData.quotaInfo || modelData.quota_info; + if (!quotaInfo) continue; + + // Extract remaining fraction (0-1 range) + const remaining = + quotaInfo.remainingFraction ?? quotaInfo.remaining_fraction ?? quotaInfo.remaining; + + if (typeof remaining !== 'number') continue; + + // Convert to percentage (0-100) + const percentage = Math.round(remaining * 100); + + // Extract reset time + const resetTime = quotaInfo.resetTime || quotaInfo.reset_time || null; + + models.push({ + name: modelId, + displayName: modelData.displayName, + percentage, + resetTime, + }); + } + } + + return { + success: true, + models, + lastUpdated: Date.now(), + }; + } catch (err) { + clearTimeout(timeoutId); + return { + success: false, + models: [], + lastUpdated: Date.now(), + error: err instanceof Error ? err.message : 'Unknown error', + }; + } +} + +/** + * Fetch quota for an Antigravity account + * + * @param provider - Provider name (only 'agy' supported) + * @param accountId - Account identifier (email with _ replacing @ and .) + * @returns Quota result with models and percentages + */ +export async function fetchAccountQuota( + provider: CLIProxyProvider, + accountId: string +): Promise { + // Only Antigravity supports quota fetching + if (provider !== 'agy') { + return { + success: false, + models: [], + lastUpdated: Date.now(), + error: `Quota not supported for provider: ${provider}`, + }; + } + + // Read access token from auth file + const accessToken = readAccessToken(provider, accountId); + if (!accessToken) { + return { + success: false, + models: [], + lastUpdated: Date.now(), + error: 'Access token not found for account', + }; + } + + // Get project ID first + const projectId = await getProjectId(accessToken); + if (!projectId) { + return { + success: false, + models: [], + lastUpdated: Date.now(), + error: 'Failed to retrieve project ID', + }; + } + + // Fetch models with quota + return fetchAvailableModels(accessToken, projectId); +} diff --git a/src/web-server/routes/cliproxy-stats-routes.ts b/src/web-server/routes/cliproxy-stats-routes.ts index a4ac196c..729ebd4e 100644 --- a/src/web-server/routes/cliproxy-stats-routes.ts +++ b/src/web-server/routes/cliproxy-stats-routes.ts @@ -12,6 +12,8 @@ import { fetchCliproxyErrorLogs, fetchCliproxyErrorLogContent, } from '../../cliproxy/stats-fetcher'; +import { fetchAccountQuota } from '../../cliproxy/quota-fetcher'; +import type { CLIProxyProvider } from '../../cliproxy/types'; import { getCliproxyWritablePath, getConfigPath, @@ -430,4 +432,50 @@ router.put('/models/:provider', async (req: Request, res: Response): Promise => { + const { provider, accountId } = req.params; + + // Validate provider + const validProviders: CLIProxyProvider[] = [ + 'agy', + 'gemini', + 'codex', + 'qwen', + 'iflow', + 'kiro', + 'ghcp', + ]; + if (!validProviders.includes(provider as CLIProxyProvider)) { + res.status(400).json({ + error: 'Invalid provider', + message: `Provider must be one of: ${validProviders.join(', ')}`, + }); + return; + } + + // Validate accountId - prevent path traversal + if ( + !accountId || + accountId.includes('..') || + accountId.includes('/') || + accountId.includes('\\') + ) { + res.status(400).json({ error: 'Invalid account ID' }); + return; + } + + try { + const result = await fetchAccountQuota(provider as CLIProxyProvider, accountId); + res.json(result); + } catch (error) { + res.status(500).json({ error: (error as Error).message }); + } +}); + export default router; diff --git a/ui/src/components/cliproxy/provider-editor/account-item.tsx b/ui/src/components/cliproxy/provider-editor/account-item.tsx index 589353ee..e2d008de 100644 --- a/ui/src/components/cliproxy/provider-editor/account-item.tsx +++ b/ui/src/components/cliproxy/provider-editor/account-item.tsx @@ -1,88 +1,195 @@ /** * Account Item Component - * Displays a single OAuth account with actions + * Displays a single OAuth account with actions and quota bar */ import { Button } from '@/components/ui/button'; import { Badge } from '@/components/ui/badge'; +import { Progress } from '@/components/ui/progress'; import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger, } from '@/components/ui/dropdown-menu'; -import { User, Star, MoreHorizontal, Clock, Trash2 } from 'lucide-react'; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '@/components/ui/tooltip'; +import { User, Star, MoreHorizontal, Clock, Trash2, Loader2 } from 'lucide-react'; import { cn } from '@/lib/utils'; import { PRIVACY_BLUR_CLASS } from '@/contexts/privacy-context'; +import { useAccountQuota } from '@/hooks/use-cliproxy-stats'; import type { AccountItemProps } from './types'; +/** + * Format reset time as relative time (e.g., "in 2 hours") + */ +function formatResetTime(resetTime: string | null): string | null { + if (!resetTime) return null; + try { + const reset = new Date(resetTime); + const now = new Date(); + const diff = reset.getTime() - now.getTime(); + if (diff <= 0) return 'soon'; + + const hours = Math.floor(diff / (1000 * 60 * 60)); + const minutes = Math.floor((diff % (1000 * 60 * 60)) / (1000 * 60)); + + if (hours > 0) return `in ${hours}h ${minutes}m`; + return `in ${minutes}m`; + } catch { + return null; + } +} + +/** + * Get color class based on quota percentage + */ +function getQuotaColor(percentage: number): string { + if (percentage <= 20) return 'bg-destructive'; + if (percentage <= 50) return 'bg-yellow-500'; + return 'bg-green-500'; +} + export function AccountItem({ account, onSetDefault, onRemove, isRemoving, privacyMode, + showQuota, }: AccountItemProps) { + // Fetch quota for 'agy' provider accounts + const { data: quota, isLoading: quotaLoading } = useAccountQuota( + account.provider, + account.id, + showQuota && account.provider === 'agy' + ); + + // Calculate average quota across all models + const avgQuota = + quota?.success && quota.models.length > 0 + ? Math.round(quota.models.reduce((sum, m) => sum + m.percentage, 0) / quota.models.length) + : null; + + // Get earliest reset time + const nextReset = + quota?.success && quota.models.length > 0 + ? quota.models.reduce( + (earliest, m) => { + if (!m.resetTime) return earliest; + if (!earliest) return m.resetTime; + return new Date(m.resetTime) < new Date(earliest) ? m.resetTime : earliest; + }, + null as string | null + ) + : null; + return (
-
-
- -
-
-
- - {account.email || account.id} - - {account.isDefault && ( - - - Default - +
+
+
+ +
+
+
+ + {account.email || account.id} + + {account.isDefault && ( + + + Default + + )} +
+ {account.lastUsedAt && ( +
+ + Last used: {new Date(account.lastUsedAt).toLocaleDateString()} +
)}
- {account.lastUsedAt && ( -
- - Last used: {new Date(account.lastUsedAt).toLocaleDateString()} -
- )}
+ + + + + + + {!account.isDefault && ( + + + Set as default + + )} + + + {isRemoving ? 'Removing...' : 'Remove account'} + + +
- - - - - - {!account.isDefault && ( - - - Set as default - - )} - - - {isRemoving ? 'Removing...' : 'Remove account'} - - - + {/* Quota bar - only for 'agy' provider */} + {showQuota && account.provider === 'agy' && ( +
+ {quotaLoading ? ( +
+ + Loading quota... +
+ ) : avgQuota !== null ? ( + + + +
+ + {avgQuota}% +
+ + +
+

Model Quotas:

+ {quota?.models.map((m) => ( +
+ {m.displayName || m.name} + {m.percentage}% +
+ ))} + {nextReset && ( +

+ Resets {formatResetTime(nextReset)} +

+ )} +
+ + + + ) : quota?.error ? ( +
{quota.error}
+ ) : null} +
+ )}
); } diff --git a/ui/src/components/cliproxy/provider-editor/accounts-section.tsx b/ui/src/components/cliproxy/provider-editor/accounts-section.tsx index 7d398ac3..939777a8 100644 --- a/ui/src/components/cliproxy/provider-editor/accounts-section.tsx +++ b/ui/src/components/cliproxy/provider-editor/accounts-section.tsx @@ -16,6 +16,8 @@ interface AccountsSectionProps { onRemoveAccount: (accountId: string) => void; isRemovingAccount?: boolean; privacyMode?: boolean; + /** Show quota bars for accounts (only applicable for 'agy' provider) */ + showQuota?: boolean; } export function AccountsSection({ @@ -25,6 +27,7 @@ export function AccountsSection({ onRemoveAccount, isRemovingAccount, privacyMode, + showQuota, }: AccountsSectionProps) { return (
@@ -54,6 +57,7 @@ export function AccountsSection({ onRemove={() => onRemoveAccount(account.id)} isRemoving={isRemovingAccount} privacyMode={privacyMode} + showQuota={showQuota} /> ))}
diff --git a/ui/src/components/cliproxy/provider-editor/index.tsx b/ui/src/components/cliproxy/provider-editor/index.tsx index ab278d5e..d1c13f43 100644 --- a/ui/src/components/cliproxy/provider-editor/index.tsx +++ b/ui/src/components/cliproxy/provider-editor/index.tsx @@ -177,6 +177,7 @@ export function ProviderEditor({ onRemoveAccount={onRemoveAccount} isRemovingAccount={isRemovingAccount} privacyMode={privacyMode} + provider={provider} /> void; isRemovingAccount?: boolean; privacyMode?: boolean; + /** Provider name for quota display */ + provider?: string; } export function ModelConfigTab({ @@ -56,6 +58,7 @@ export function ModelConfigTab({ onRemoveAccount, isRemovingAccount, privacyMode, + provider, }: ModelConfigTabProps) { return ( @@ -82,6 +85,7 @@ export function ModelConfigTab({ onRemoveAccount={onRemoveAccount} isRemovingAccount={isRemovingAccount} privacyMode={privacyMode} + showQuota={provider === 'agy'} />
diff --git a/ui/src/components/cliproxy/provider-editor/types.ts b/ui/src/components/cliproxy/provider-editor/types.ts index 6bea9288..600016f4 100644 --- a/ui/src/components/cliproxy/provider-editor/types.ts +++ b/ui/src/components/cliproxy/provider-editor/types.ts @@ -35,6 +35,8 @@ export interface AccountItemProps { onRemove: () => void; isRemoving?: boolean; privacyMode?: boolean; + /** Show quota bar (only for 'agy' provider) */ + showQuota?: boolean; } export interface ModelMappingValues { diff --git a/ui/src/components/ui/progress.tsx b/ui/src/components/ui/progress.tsx new file mode 100644 index 00000000..403613bc --- /dev/null +++ b/ui/src/components/ui/progress.tsx @@ -0,0 +1,40 @@ +/** + * Progress Component + * Simple progress bar with customizable indicator color + */ + +import * as React from 'react'; +import { cn } from '@/lib/utils'; + +interface ProgressProps extends React.HTMLAttributes { + value?: number; + max?: number; + indicatorClassName?: string; +} + +const Progress = React.forwardRef( + ({ className, value = 0, max = 100, indicatorClassName, ...props }, ref) => { + const percentage = Math.min(Math.max((value / max) * 100, 0), 100); + + return ( +
+
+
+ ); + } +); + +Progress.displayName = 'Progress'; + +export { Progress }; diff --git a/ui/src/hooks/use-cliproxy-stats.ts b/ui/src/hooks/use-cliproxy-stats.ts index 2e825e01..9f93ec2c 100644 --- a/ui/src/hooks/use-cliproxy-stats.ts +++ b/ui/src/hooks/use-cliproxy-stats.ts @@ -3,6 +3,7 @@ */ import { useQuery } from '@tanstack/react-query'; +import type { ModelQuota, QuotaResult } from '@/lib/api-client'; /** Per-account usage statistics */ export interface AccountUsageStats { @@ -189,3 +190,33 @@ export function useCliproxyErrorLogContent(name: string | null) { staleTime: 60000, // Cache log content for 1 minute }); } + +// Re-export for consumers +export type { ModelQuota, QuotaResult }; + +/** + * Fetch account quota from API + */ +async function fetchAccountQuota(provider: string, accountId: string): Promise { + const response = await fetch(`/api/cliproxy/quota/${provider}/${encodeURIComponent(accountId)}`); + if (!response.ok) { + const error = await response.json(); + throw new Error(error.message || 'Failed to fetch quota'); + } + return response.json(); +} + +/** + * Hook to get account quota + * Only enabled for 'agy' provider (Antigravity) as it's the only one supporting quota + */ +export function useAccountQuota(provider: string, accountId: string, enabled = true) { + return useQuery({ + queryKey: ['account-quota', provider, accountId], + queryFn: () => fetchAccountQuota(provider, accountId), + enabled: enabled && provider === 'agy' && !!accountId, + staleTime: 60000, // Cache for 1 minute + refetchInterval: 300000, // Refresh every 5 minutes + retry: 1, + }); +} diff --git a/ui/src/lib/api-client.ts b/ui/src/lib/api-client.ts index 87013507..8dbb2d31 100644 --- a/ui/src/lib/api-client.ts +++ b/ui/src/lib/api-client.ts @@ -107,6 +107,32 @@ export interface CliproxyModelsResponse { totalCount: number; } +/** Individual model quota info from Google Cloud Code API */ +export interface ModelQuota { + /** Model name, e.g., "gemini-3-pro-high" */ + name: string; + /** Display name from API, e.g., "Gemini 3 Pro" */ + displayName?: string; + /** Remaining quota as percentage (0-100) */ + percentage: number; + /** ISO timestamp when quota resets, null if unknown */ + resetTime: string | null; +} + +/** Quota fetch result */ +export interface QuotaResult { + /** Whether fetch succeeded */ + success: boolean; + /** Quota for each available model */ + models: ModelQuota[]; + /** Timestamp of fetch */ + lastUpdated: number; + /** True if account lacks quota access (403) */ + isForbidden?: boolean; + /** Error message if fetch failed */ + error?: string; +} + /** Provider accounts summary */ export type ProviderAccountsMap = Record; @@ -404,4 +430,10 @@ export const api = { body: JSON.stringify(params), }), }, + /** Account quota API */ + quota: { + /** Fetch quota for a specific account */ + get: (provider: string, accountId: string) => + request(`/cliproxy/quota/${provider}/${encodeURIComponent(accountId)}`), + }, };