use serde::{Deserialize, Serialize}; use serde_json::json; use crate::models::QuotaData; const QUOTA_API_URL: &str = "https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels"; const USER_AGENT: &str = "antigravity/1.11.3 Darwin/arm64"; const CLOUD_CODE_BASE_URL: &str = "https://cloudcode-pa.googleapis.com"; #[derive(Debug, Serialize, Deserialize)] struct QuotaResponse { models: std::collections::HashMap, } #[derive(Debug, Serialize, Deserialize)] struct ModelInfo { #[serde(rename = "quotaInfo")] quota_info: Option, } #[derive(Debug, Serialize, Deserialize)] struct QuotaInfo { #[serde(rename = "remainingFraction")] remaining_fraction: Option, #[serde(rename = "resetTime")] reset_time: Option, } #[derive(Debug, Deserialize)] struct LoadProjectResponse { #[serde(rename = "cloudaicompanionProject")] project_id: Option, #[serde(rename = "currentTier")] current_tier: Option, #[serde(rename = "paidTier")] paid_tier: Option, } #[derive(Debug, Deserialize)] struct Tier { id: Option, #[serde(rename = "quotaTier")] #[allow(dead_code)] quota_tier: Option, #[allow(dead_code)] name: Option, #[allow(dead_code)] slug: Option, } /// Create configured HTTP client fn create_client() -> reqwest::Client { crate::utils::http::create_client(15) } /// Get project ID and subscription type async fn fetch_project_id(access_token: &str, email: &str) -> (Option, Option) { let client = create_client(); let meta = json!({"metadata": {"ideType": "ANTIGRAVITY"}}); let res = client .post(format!("{}/v1internal:loadCodeAssist", CLOUD_CODE_BASE_URL)) .header(reqwest::header::AUTHORIZATION, format!("Bearer {}", access_token)) .header(reqwest::header::CONTENT_TYPE, "application/json") .header(reqwest::header::USER_AGENT, "antigravity/windows/amd64") .json(&meta) .send() .await; match res { Ok(res) => { if res.status().is_success() { if let Ok(data) = res.json::().await { let project_id = data.project_id.clone(); // Core logic: prefer paid_tier ID, reflects true account rights better than current_tier let subscription_tier = data.paid_tier .and_then(|t| t.id) .or_else(|| data.current_tier.and_then(|t| t.id)); if let Some(ref tier) = subscription_tier { crate::modules::logger::log_info(&format!( "[{}] Subscription identified: {}", email, tier )); } return (project_id, subscription_tier); } } else { crate::modules::logger::log_warn(&format!( "[{}] loadCodeAssist failed: Status: {}", email, res.status() )); } } Err(e) => { crate::modules::logger::log_error(&format!("[{}] loadCodeAssist network error: {}", email, e)); } } (None, None) } /// Unified entry point for querying account quota pub async fn fetch_quota(access_token: &str, email: &str) -> crate::error::AppResult<(QuotaData, Option)> { fetch_quota_inner(access_token, email).await } /// Quota query logic pub async fn fetch_quota_inner(access_token: &str, email: &str) -> crate::error::AppResult<(QuotaData, Option)> { use crate::error::AppError; // 1. Get project ID and subscription type let (project_id, subscription_tier) = fetch_project_id(access_token, email).await; let final_project_id = project_id.as_deref().unwrap_or("bamboo-precept-lgxtn"); let client = create_client(); let payload = json!({ "project": final_project_id }); let url = QUOTA_API_URL; let max_retries = 3; let mut last_error: Option = None; for attempt in 1..=max_retries { match client .post(url) .bearer_auth(access_token) .header("User-Agent", USER_AGENT) .json(&json!(payload)) .send() .await { Ok(response) => { // Convert HTTP error status to AppError if response.error_for_status_ref().is_err() { let status = response.status(); // Special handling for 403 Forbidden - return directly, no retry if status == reqwest::StatusCode::FORBIDDEN { crate::modules::logger::log_warn("Account forbidden (403), marking as forbidden status"); let mut q = QuotaData::new(); q.is_forbidden = true; q.subscription_tier = subscription_tier.clone(); return Ok((q, project_id.clone())); } // Other errors continue retry logic if attempt < max_retries { let text = response.text().await.unwrap_or_default(); crate::modules::logger::log_warn(&format!("API error: {} - {} (attempt {}/{})", status, text, attempt, max_retries)); last_error = Some(AppError::Unknown(format!("HTTP {} - {}", status, text))); tokio::time::sleep(std::time::Duration::from_secs(1)).await; continue; } else { let text = response.text().await.unwrap_or_default(); return Err(AppError::Unknown(format!("API error: {} - {}", status, text))); } } let quota_response: QuotaResponse = response .json() .await .map_err(AppError::Network)?; let mut quota_data = QuotaData::new(); tracing::debug!("Quota API returned {} models", quota_response.models.len()); for (name, info) in quota_response.models { if let Some(quota_info) = info.quota_info { let percentage = quota_info.remaining_fraction .map(|f| (f * 100.0) as i32) .unwrap_or(0); let reset_time = quota_info.reset_time.unwrap_or_default(); // Only save models we care about if name.contains("gemini") || name.contains("claude") { quota_data.add_model(name, percentage, reset_time); } } } // Set subscription type quota_data.subscription_tier = subscription_tier.clone(); return Ok((quota_data, project_id.clone())); }, Err(e) => { crate::modules::logger::log_warn(&format!("Request failed: {} (attempt {}/{})", e, attempt, max_retries)); last_error = Some(AppError::Network(e)); if attempt < max_retries { tokio::time::sleep(std::time::Duration::from_secs(1)).await; } } } } Err(last_error.unwrap_or_else(|| AppError::Unknown("Quota query failed".to_string()))) } /// Batch query all account quotas (backup function) #[allow(dead_code)] pub async fn fetch_all_quotas(accounts: Vec<(String, String)>) -> Vec<(String, crate::error::AppResult)> { let mut results = Vec::new(); for (account_id, access_token) in accounts { let result = fetch_quota(&access_token, &account_id).await.map(|(q, _)| q); results.push((account_id, result)); } results }