From 05e95410d74223eb9a2e7d0db21bf98f119e7a4b Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Sun, 1 Feb 2026 19:37:18 +0200 Subject: [PATCH] feat(desktop/vscode): add quota providers API and UI integration (#266) * feat(openchamber): persist usage auto-refresh settings Enable a switch to toggle automatic usage refresh Provide input to configure refresh interval in milliseconds Persist changes to desktop settings and server API when changed * feat: add UsageCard component Add a new UsageCard component to display a usage window with title, optional subtitle, and a progress bar Show current usage percentage and a reset time label for the window Render a formatted window label and a compact subtitle for concise UI * feat(usage): add UsagePage UI for quota usage Add a dedicated UsagePage with provider-based usage details Show last updated time and auto-refresh status Handle empty, not configured, and error states with informative banners * feat(usage): add UsageProgressBar component Introduce a new UsageProgressBar component to visualize quota usage Display a gradient fill that changes with critical, warn, or normal tones Expose accessible progress attributes for screen readers * feat(usage): add UsageSidebar component Display quotas for all providers in a scrollable sidebar Refresh quotas with a button and loading indicator Colorize provider rows based on usage status and runtime context * feat: add usage section to Settings Add a new Usage item to the Settings sidebar for desktop and mobile Render UsagePage when the Usage tab is selected in Settings Wire up new UsageSidebar and UsagePage components under usage * feat: add Usage section to sidebar Add new Usage section in the sidebar for API quota monitoring. Display a bar chart icon and description for the Usage item. Monitor and display API quota usage across providers. * feat(desktop): add usage auto-refresh settings Enable automatic refresh for usage data with new settings Store refresh interval in milliseconds for usage updates * feat(persistence): support usageAutoRefresh and usageRefreshIntervalMs Persist usageAutoRefresh in desktop settings Persist usageRefreshIntervalMs in desktop settings Validate types for new fields during sanitizeWebSettings * feat(quota): export providers and utilities Expose QUOTA_PROVIDERS and QUOTA_PROVIDER_MAP for consumers Export QuotaProviderMeta type for user code Make formatting and usage resolution utilities available from quota module * feat(quota): define base quota provider interface Define QuotaProvider interface with id, name, isConfigured, and fetchQuota Expose ProviderResult type in fetchQuota contract * feat: add quota providers index and map Expose QUOTA_PROVIDERS with OpenAI, Google and z.ai Provide QUOTA_PROVIDER_MAP for quick provider lookup by id * feat: add quota utils for percent formatting and tone Add clampPercent to sanitize and clamp numbers to 0-100 Add formatPercent to render '-' for null and 'x%' for values Add resolveUsageTone to categorize percent as safe, warn, or critical * feat: add useQuotaStore for quota data Load usage settings from desktop, VSCode, or API at startup Fetch quotas for all providers in parallel and update loading state Expose lastUpdated timestamp and error state for UI feedback * feat: export quota types from quota module Expose quota-related types in UI type definitions Allow downstream code to import QuotaProviderId and related types Aggregate quota exports under the quota module in index * feat: add quota types for usage providers Add QuotaProviderId and UsageWindow shapes to model quota data Add ProviderUsage, ProviderResult, and related usage mapping for providers * feat: add quota provider endpoints API List available quota providers via GET /api/quota/providers Retrieve quota details for a specific provider with GET /api/quota/:providerId Log errors and return 500 with error message on quota fetch failures * feat: add quota providers discovery and formatting Detect configured quota providers from auth and account files Normalize auth entries to tokens or objects for API usage Expose formatted reset times and remaining window metrics * feat: persist usage settings in UsageSidebar and remove from defaults Load usage settings on mount for the sidebar Persist changes to auto-refresh and refresh interval to server Remove usage settings state and effects from DefaultsSettings * fix(usage): guard auto-select in UsagePage when results empty Guard auto-select in UsagePage when results are empty Prevent unexpected provider selection on initial render * fix(quota): set error state during quota updates Reset error to null when a new quota result is added Set error to the error message on fetch failure or fallback Keep error state alongside results in all update paths * feat(desktop/vscode): add quota providers API and UI integration * fix(quotaProviders): correct quota payload parsing --------- Co-authored-by: Nelson Pires --- .../src-tauri/src/commands/settings.rs | 13 + packages/desktop/src-tauri/src/main.rs | 41 +- .../desktop/src-tauri/src/quota_providers.rs | 797 ++++++++++++++++++ .../components/sections/usage/UsageCard.tsx | 2 +- .../components/sections/usage/UsagePage.tsx | 8 +- .../sections/usage/UsageProgressBar.tsx | 14 +- .../sections/usage/UsageSidebar.tsx | 37 +- packages/vscode/src/bridge.ts | 35 + packages/vscode/src/quotaProviders.ts | 612 ++++++++++++++ packages/vscode/webview/main.tsx | 22 + 10 files changed, 1555 insertions(+), 26 deletions(-) create mode 100644 packages/desktop/src-tauri/src/quota_providers.rs create mode 100644 packages/vscode/src/quotaProviders.ts diff --git a/packages/desktop/src-tauri/src/commands/settings.rs b/packages/desktop/src-tauri/src/commands/settings.rs index 62b7c031..75343ccb 100644 --- a/packages/desktop/src-tauri/src/commands/settings.rs +++ b/packages/desktop/src-tauri/src/commands/settings.rs @@ -300,6 +300,9 @@ fn sanitize_settings_update(payload: &Value) -> Value { if let Some(Value::Bool(b)) = obj.get("notifyOnSubtasks") { result_obj.insert("notifyOnSubtasks".to_string(), json!(b)); } + if let Some(Value::Bool(b)) = obj.get("usageAutoRefresh") { + result_obj.insert("usageAutoRefresh".to_string(), json!(b)); + } if let Some(Value::String(s)) = obj.get("notificationMode") { let trimmed = s.trim(); if trimmed == "always" || trimmed == "hidden-only" { @@ -377,6 +380,16 @@ fn sanitize_settings_update(payload: &Value) -> Value { result_obj.insert("inputBarOffset".to_string(), json!(clamped)); } } + if let Some(Value::Number(n)) = obj.get("usageRefreshIntervalMs") { + let parsed = n + .as_u64() + .or_else(|| n.as_i64().and_then(|v| if v >= 0 { Some(v as u64) } else { None })) + .or_else(|| n.as_f64().map(|v| v.round().max(0.0) as u64)); + if let Some(value) = parsed { + let clamped = value.max(30000).min(300000); + result_obj.insert("usageRefreshIntervalMs".to_string(), json!(clamped)); + } + } // Memory limit fields if let Some(Value::Number(n)) = obj.get("memoryLimitHistorical") { diff --git a/packages/desktop/src-tauri/src/main.rs b/packages/desktop/src-tauri/src/main.rs index 7d24d5f7..6dcf8a5c 100644 --- a/packages/desktop/src-tauri/src/main.rs +++ b/packages/desktop/src-tauri/src/main.rs @@ -7,6 +7,7 @@ mod opencode_auth; mod opencode_config; mod opencode_manager; mod path_utils; +mod quota_providers; mod session_activity; mod skills_catalog; mod window_state; @@ -22,7 +23,7 @@ use anyhow::{anyhow, Result}; use assistant_notifications::spawn_assistant_notifications; use axum::{ body::{to_bytes, Body}, - extract::{Request, State}, + extract::{Path, Request, State}, http::{Method, StatusCode}, response::{IntoResponse, Response}, routing::{any, get, post}, @@ -275,6 +276,11 @@ struct ServerInfoPayload { has_last_directory: bool, } +#[derive(Serialize)] +struct QuotaProvidersResponse { + providers: Vec, +} + #[tauri::command] async fn desktop_server_info( state: tauri::State<'_, DesktopRuntime>, @@ -1098,6 +1104,8 @@ async fn run_http_server( "/api/openchamber/models-metadata", get(models_metadata_handler), ) + .route("/api/quota/providers", get(quota_providers_handler)) + .route("/api/quota/{providerId}", get(quota_provider_handler)) .route("/api/opencode/directory", post(change_directory_handler)) .route("/api", any(proxy_to_opencode)) .route("/api/{*rest}", any(proxy_to_opencode)) @@ -1179,6 +1187,37 @@ async fn models_metadata_handler( Ok(Json(payload)) } +async fn quota_providers_handler(State(_state): State) -> Response { + match quota_providers::list_configured_quota_providers().await { + Ok(providers) => json_response( + StatusCode::OK, + QuotaProvidersResponse { providers }, + ), + Err(err) => { + error!("[desktop:quota] Failed to list quota providers: {}", err); + config_error_response(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) + } + } +} + +async fn quota_provider_handler( + State(state): State, + Path(provider_id): Path, +) -> Response { + let trimmed = provider_id.trim(); + if trimmed.is_empty() { + return config_error_response(StatusCode::BAD_REQUEST, "Provider ID is required"); + } + + match quota_providers::fetch_quota_for_provider(&state.client, trimmed).await { + Ok(result) => json_response(StatusCode::OK, result), + Err(err) => { + error!("[desktop:quota] Failed to fetch quota: {}", err); + config_error_response(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) + } + } +} + #[derive(Deserialize)] struct DirectoryChangeRequest { path: String, diff --git a/packages/desktop/src-tauri/src/quota_providers.rs b/packages/desktop/src-tauri/src/quota_providers.rs new file mode 100644 index 00000000..025844eb --- /dev/null +++ b/packages/desktop/src-tauri/src/quota_providers.rs @@ -0,0 +1,797 @@ +use anyhow::{anyhow, Result}; +use chrono::{DateTime, Local, TimeZone}; +use log::warn; +use reqwest::Client; +use serde::Serialize; +use serde_json::Value; +use std::{ + collections::{HashMap, HashSet}, + path::PathBuf, + time::Duration, +}; + +use crate::opencode_auth; + +const OPENCODE_CONFIG_DIR: &str = ".config/opencode"; +const OPENCODE_DATA_DIR: &str = ".local/share/opencode"; + +const GOOGLE_CLIENT_ID: &str = + "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"; +const GOOGLE_CLIENT_SECRET: &str = "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"; +const DEFAULT_PROJECT_ID: &str = "rising-fact-p41fc"; +const GOOGLE_WINDOW_SECONDS: i64 = 5 * 60 * 60; + +const GOOGLE_ENDPOINTS: [&str; 3] = [ + "https://daily-cloudcode-pa.sandbox.googleapis.com", + "https://autopush-cloudcode-pa.sandbox.googleapis.com", + "https://cloudcode-pa.googleapis.com", +]; + +const GOOGLE_USER_AGENT: &str = "antigravity/1.11.5 windows/amd64"; +const GOOGLE_API_CLIENT: &str = "google-cloud-sdk vscode_cloudshelleditor/0.1"; +const GOOGLE_CLIENT_METADATA: &str = + "{\"ideType\":\"IDE_UNSPECIFIED\",\"platform\":\"PLATFORM_UNSPECIFIED\",\"pluginType\":\"GEMINI\"}"; + +#[derive(Clone, Debug, Default)] +struct AuthEntry { + token: Option, + access: Option, + refresh: Option, + expires: Option, + key: Option, +} + +#[derive(Clone, Debug, Default)] +struct GoogleAuth { + access_token: Option, + refresh_token: Option, + expires: Option, + project_id: Option, +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ProviderResult { + provider_id: String, + provider_name: String, + ok: bool, + configured: bool, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, + usage: Option, + fetched_at: i64, +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct ProviderUsage { + windows: HashMap, + #[serde(skip_serializing_if = "Option::is_none")] + models: Option>, +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct UsageWindow { + used_percent: Option, + remaining_percent: Option, + window_seconds: Option, + reset_after_seconds: Option, + reset_at: Option, + reset_at_formatted: Option, + reset_after_formatted: Option, +} + +fn get_home_dir() -> PathBuf { + dirs::home_dir().unwrap_or_else(|| PathBuf::from("/")) +} + +fn opencode_config_dir() -> PathBuf { + get_home_dir().join(OPENCODE_CONFIG_DIR) +} + +fn opencode_data_dir() -> PathBuf { + get_home_dir().join(OPENCODE_DATA_DIR) +} + +fn antigravity_accounts_paths() -> [PathBuf; 2] { + [ + opencode_config_dir().join("antigravity-accounts.json"), + opencode_data_dir().join("antigravity-accounts.json"), + ] +} + +async fn read_json_file(path: &PathBuf) -> Option { + if !path.exists() { + return None; + } + let raw = tokio::fs::read_to_string(path).await.ok()?; + let trimmed = raw.trim(); + if trimmed.is_empty() { + return None; + } + serde_json::from_str(trimmed).map_err(|err| { + warn!("Failed to read JSON file {}: {}", path.display(), err); + err + }).ok() +} + +fn get_auth_entry<'a>(auth: &'a serde_json::Map, aliases: &[&str]) -> Option<&'a Value> { + for alias in aliases { + if let Some(value) = auth.get(*alias) { + return Some(value); + } + } + None +} + +fn normalize_auth_entry(value: Option<&Value>) -> Option { + let value = value?; + match value { + Value::String(token) => Some(AuthEntry { + token: Some(token.clone()), + ..AuthEntry::default() + }), + Value::Object(map) => { + let token = map.get("token").and_then(|v| v.as_str()).map(|s| s.to_string()); + let access = map.get("access").and_then(|v| v.as_str()).map(|s| s.to_string()); + let refresh = map.get("refresh").and_then(|v| v.as_str()).map(|s| s.to_string()); + let key = map.get("key").and_then(|v| v.as_str()).map(|s| s.to_string()); + let expires = map + .get("expires") + .and_then(|v| v.as_i64()) + .or_else(|| map.get("expires").and_then(|v| v.as_f64()).map(|v| v.round() as i64)); + + Some(AuthEntry { + token, + access, + refresh, + expires, + key, + }) + } + _ => None, + } +} + +fn format_reset_at(timestamp_ms: i64) -> Option { + let dt = Local.timestamp_millis_opt(timestamp_ms).single()?; + Some(dt.format("%-I:%M %p").to_string()) +} + +fn format_duration(seconds: i64) -> Option { + if seconds < 0 { + return None; + } + let clamped = seconds.max(0) as i64; + let hours = clamped / 3600; + let minutes = (clamped % 3600) / 60; + if hours == 0 && minutes == 0 { + return Some("0m".to_string()); + } + if hours == 0 { + return Some(format!("{}m", minutes)); + } + if minutes == 0 { + return Some(format!("{}h", hours)); + } + Some(format!("{}h {}m", hours, minutes)) +} + +fn calculate_reset_after_seconds(reset_at: Option) -> Option { + let reset_at = reset_at?; + let now_ms = chrono::Utc::now().timestamp_millis(); + let delta = (reset_at - now_ms) / 1000; + Some(delta.max(0)) +} + +fn to_usage_window(used_percent: Option, window_seconds: Option, reset_at: Option) -> UsageWindow { + let remaining_percent = used_percent.map(|value| (100.0 - value).max(0.0)); + let reset_after_seconds = calculate_reset_after_seconds(reset_at); + let reset_at_formatted = reset_at.and_then(format_reset_at); + let reset_after_formatted = reset_after_seconds.and_then(format_duration); + + UsageWindow { + used_percent, + remaining_percent, + window_seconds, + reset_after_seconds, + reset_at, + reset_at_formatted, + reset_after_formatted, + } +} + +fn build_result( + provider_id: &str, + provider_name: &str, + ok: bool, + configured: bool, + usage: Option, + error: Option, +) -> ProviderResult { + ProviderResult { + provider_id: provider_id.to_string(), + provider_name: provider_name.to_string(), + ok, + configured, + error, + usage, + fetched_at: chrono::Utc::now().timestamp_millis(), + } +} + +async fn load_auth_map() -> Result> { + let auth = opencode_auth::read_auth().await?; + auth.as_object() + .cloned() + .ok_or_else(|| anyhow!("Auth file is not a valid JSON object")) +} + +async fn has_antigravity_accounts() -> bool { + for path in antigravity_accounts_paths() { + if let Some(data) = read_json_file(&path).await { + if data + .get("accounts") + .and_then(|value| value.as_array()) + .is_some_and(|accounts| !accounts.is_empty()) + { + return true; + } + } + } + false +} + +pub async fn list_configured_quota_providers() -> Result> { + let auth = load_auth_map().await?; + let mut configured: HashSet = HashSet::new(); + + let openai_auth = normalize_auth_entry(get_auth_entry(&auth, &["openai", "codex", "chatgpt"])); + if let Some(entry) = openai_auth { + if entry.access.is_some() || entry.token.is_some() { + configured.insert("openai".to_string()); + } + } + + let google_auth = normalize_auth_entry(get_auth_entry(&auth, &["google", "antigravity"])); + if let Some(entry) = google_auth { + if entry.access.is_some() || entry.token.is_some() || entry.refresh.is_some() { + configured.insert("google".to_string()); + } + } + + let zai_auth = + normalize_auth_entry(get_auth_entry(&auth, &["zai-coding-plan", "zai", "z.ai"])); + if let Some(entry) = zai_auth { + if entry.key.is_some() || entry.token.is_some() { + configured.insert("zai-coding-plan".to_string()); + } + } + + if has_antigravity_accounts().await { + configured.insert("google".to_string()); + } + + Ok(configured.into_iter().collect()) +} + +fn parse_number(value: Option<&Value>) -> Option { + let value = value?; + value.as_f64().or_else(|| value.as_i64().map(|v| v as f64)) +} + +async fn fetch_openai_quota(client: &Client) -> Result { + let auth = load_auth_map().await?; + let entry = normalize_auth_entry(get_auth_entry(&auth, &["openai", "codex", "chatgpt"])); + let access_token = entry + .as_ref() + .and_then(|entry| entry.access.clone().or(entry.token.clone())); + + let Some(access_token) = access_token else { + return Ok(build_result( + "openai", + "OpenAI", + false, + false, + None, + Some("Not configured".to_string()), + )); + }; + + let response = client + .get("https://chatgpt.com/backend-api/wham/usage") + .bearer_auth(access_token) + .header("Content-Type", "application/json") + .send() + .await; + + let response = match response { + Ok(resp) => resp, + Err(err) => { + return Ok(build_result( + "openai", + "OpenAI", + false, + true, + None, + Some(err.to_string()), + )) + } + }; + + if !response.status().is_success() { + return Ok(build_result( + "openai", + "OpenAI", + false, + true, + None, + Some(format!("API error: {}", response.status().as_u16())), + )); + } + + let payload: Value = match response.json().await { + Ok(value) => value, + Err(err) => { + return Ok(build_result( + "openai", + "OpenAI", + false, + true, + None, + Some(err.to_string()), + )) + } + }; + + let primary = payload + .get("rate_limit") + .and_then(|value| value.get("primary_window")); + let secondary = payload + .get("rate_limit") + .and_then(|value| value.get("secondary_window")); + + let mut windows: HashMap = HashMap::new(); + + if let Some(primary) = primary { + let used_percent = parse_number(primary.get("used_percent")); + let window_seconds = primary + .get("limit_window_seconds") + .and_then(|value| value.as_i64()); + let reset_at = primary + .get("reset_at") + .and_then(|value| value.as_i64()) + .map(|value| value * 1000); + windows.insert( + "5h".to_string(), + to_usage_window(used_percent, window_seconds, reset_at), + ); + } + + if let Some(secondary) = secondary { + let used_percent = parse_number(secondary.get("used_percent")); + let window_seconds = secondary + .get("limit_window_seconds") + .and_then(|value| value.as_i64()); + let reset_at = secondary + .get("reset_at") + .and_then(|value| value.as_i64()) + .map(|value| value * 1000); + windows.insert( + "weekly".to_string(), + to_usage_window(used_percent, window_seconds, reset_at), + ); + } + + Ok(build_result( + "openai", + "OpenAI", + true, + true, + Some(ProviderUsage { + windows, + models: None, + }), + None, + )) +} + +async fn resolve_google_auth() -> Result> { + let auth = load_auth_map().await?; + let entry = normalize_auth_entry(get_auth_entry(&auth, &["google", "antigravity"])); + + if let Some(entry) = entry { + let mut refresh = entry.refresh.clone(); + let mut project_id = None; + if let Some(value) = entry.refresh.clone() { + if let Some((first, second)) = value.split_once('|') { + refresh = Some(first.to_string()); + project_id = Some(second.to_string()); + } + } + return Ok(Some(GoogleAuth { + access_token: entry.access.or(entry.token), + refresh_token: refresh, + expires: entry.expires, + project_id, + })); + } + + for path in antigravity_accounts_paths() { + let data = match read_json_file(&path).await { + Some(data) => data, + None => continue, + }; + let accounts = data.get("accounts").and_then(|value| value.as_array()); + if let Some(accounts) = accounts { + if accounts.is_empty() { + continue; + } + let index = data + .get("activeIndex") + .and_then(|value| value.as_i64()) + .unwrap_or(0) + .max(0) as usize; + let account = accounts.get(index).or_else(|| accounts.first()); + if let Some(account) = account { + let refresh_token = account + .get("refreshToken") + .and_then(|value| value.as_str()) + .map(|value| value.to_string()); + if refresh_token.is_none() { + continue; + } + let project_id = account + .get("projectId") + .and_then(|value| value.as_str()) + .or_else(|| { + account + .get("managedProjectId") + .and_then(|value| value.as_str()) + }) + .map(|value| value.to_string()); + + return Ok(Some(GoogleAuth { + access_token: None, + refresh_token, + expires: None, + project_id, + })); + } + } + } + + Ok(None) +} + +async fn refresh_google_access_token(client: &Client, refresh_token: &str) -> Result> { + let body = format!( + "client_id={}&client_secret={}&refresh_token={}&grant_type=refresh_token", + urlencoding::encode(GOOGLE_CLIENT_ID), + urlencoding::encode(GOOGLE_CLIENT_SECRET), + urlencoding::encode(refresh_token) + ); + + let response = client + .post("https://oauth2.googleapis.com/token") + .header("Content-Type", "application/x-www-form-urlencoded") + .body(body) + .send() + .await; + + let response = match response { + Ok(resp) => resp, + Err(err) => { + warn!("Failed to refresh Google token: {}", err); + return Ok(None); + } + }; + + if !response.status().is_success() { + return Ok(None); + } + + let payload: Value = response.json().await.unwrap_or(Value::Null); + Ok(payload + .get("access_token") + .and_then(|value| value.as_str()) + .map(|value| value.to_string())) +} + +async fn fetch_google_models(client: &Client, access_token: &str, project_id: Option<&str>) -> Option { + let body = if let Some(project_id) = project_id { + serde_json::json!({ "project": project_id }) + } else { + serde_json::json!({}) + }; + + for endpoint in GOOGLE_ENDPOINTS { + let response = client + .post(format!("{}/v1internal:fetchAvailableModels", endpoint)) + .header("Authorization", format!("Bearer {}", access_token)) + .header("Content-Type", "application/json") + .header("User-Agent", GOOGLE_USER_AGENT) + .header("X-Goog-Api-Client", GOOGLE_API_CLIENT) + .header("Client-Metadata", GOOGLE_CLIENT_METADATA) + .json(&body) + .timeout(Duration::from_secs(15)) + .send() + .await; + + let response = match response { + Ok(resp) => resp, + Err(_) => continue, + }; + + if response.status().is_success() { + if let Ok(payload) = response.json::().await { + return Some(payload); + } + } + } + + None +} + +fn parse_reset_time(value: Option<&Value>) -> Option { + let value = value?; + if let Some(num) = value.as_i64() { + if num > 0 { + return Some(num); + } + } + if let Some(text) = value.as_str() { + if let Ok(parsed) = DateTime::parse_from_rfc3339(text) { + return Some(parsed.timestamp_millis()); + } + } + None +} + +async fn fetch_google_quota(client: &Client) -> Result { + let auth = resolve_google_auth().await?; + let Some(auth) = auth else { + return Ok(build_result( + "google", + "Google", + false, + false, + None, + Some("Not configured".to_string()), + )); + }; + + let now = chrono::Utc::now().timestamp_millis(); + let mut access_token = auth.access_token; + if access_token.is_none() + || auth + .expires + .is_some_and(|expires| expires <= now) + { + let Some(refresh_token) = auth.refresh_token.as_ref() else { + return Ok(build_result( + "google", + "Google", + false, + true, + None, + Some("Missing refresh token".to_string()), + )); + }; + access_token = refresh_google_access_token(client, refresh_token).await?; + } + + let Some(access_token) = access_token else { + return Ok(build_result( + "google", + "Google", + false, + true, + None, + Some("Failed to refresh OAuth token".to_string()), + )); + }; + + let project_id = auth.project_id.unwrap_or_else(|| DEFAULT_PROJECT_ID.to_string()); + let payload = fetch_google_models(client, &access_token, Some(project_id.as_str())).await; + let Some(payload) = payload else { + return Ok(build_result( + "google", + "Google", + false, + true, + None, + Some("Failed to fetch models".to_string()), + )); + }; + + let mut models: HashMap = HashMap::new(); + if let Some(model_map) = payload.get("models").and_then(|value| value.as_object()) { + for (model_name, model_data) in model_map { + let remaining_fraction = parse_number(model_data.get("quotaInfo").and_then(|v| v.get("remainingFraction"))); + let remaining_percent = remaining_fraction.map(|value| (value * 100.0).round()); + let used_percent = remaining_percent.map(|value| (100.0 - value).max(0.0)); + let reset_at = parse_reset_time(model_data.get("quotaInfo").and_then(|v| v.get("resetTime"))); + + let mut windows = HashMap::new(); + windows.insert( + "5h".to_string(), + to_usage_window(used_percent, Some(GOOGLE_WINDOW_SECONDS), reset_at), + ); + models.insert( + model_name.to_string(), + ProviderUsage { + windows, + models: None, + }, + ); + } + } + + Ok(build_result( + "google", + "Google", + true, + true, + Some(ProviderUsage { + windows: HashMap::new(), + models: if models.is_empty() { None } else { Some(models) }, + }), + None, + )) +} + +fn normalize_timestamp(value: Option<&Value>) -> Option { + let value = value?; + if let Some(num) = value.as_i64() { + if num < 1_000_000_000_000 { + return Some(num * 1000); + } + return Some(num); + } + None +} + +fn resolve_window_seconds(limit: &Value) -> Option { + let number = limit.get("number").and_then(|value| value.as_i64())?; + let unit = limit.get("unit").and_then(|value| value.as_i64())?; + let unit_seconds = match unit { + 3 => Some(3600), + _ => None, + }?; + Some(unit_seconds * number) +} + +fn resolve_window_label(window_seconds: Option) -> String { + let Some(window_seconds) = window_seconds else { + return "tokens".to_string(); + }; + if window_seconds % 86400 == 0 { + let days = window_seconds / 86400; + if days == 7 { + return "weekly".to_string(); + } + return format!("{}d", days); + } + if window_seconds % 3600 == 0 { + return format!("{}h", window_seconds / 3600); + } + format!("{}s", window_seconds) +} + +async fn fetch_zai_quota(client: &Client) -> Result { + let auth = load_auth_map().await?; + let entry = normalize_auth_entry(get_auth_entry(&auth, &["zai-coding-plan", "zai", "z.ai"])); + let api_key = entry + .as_ref() + .and_then(|entry| entry.key.clone().or(entry.token.clone())); + + let Some(api_key) = api_key else { + return Ok(build_result( + "zai-coding-plan", + "z.ai", + false, + false, + None, + Some("Not configured".to_string()), + )); + }; + + let response = client + .get("https://api.z.ai/api/monitor/usage/quota/limit") + .bearer_auth(api_key) + .header("Content-Type", "application/json") + .send() + .await; + + let response = match response { + Ok(resp) => resp, + Err(err) => { + return Ok(build_result( + "zai-coding-plan", + "z.ai", + false, + true, + None, + Some(err.to_string()), + )) + } + }; + + if !response.status().is_success() { + return Ok(build_result( + "zai-coding-plan", + "z.ai", + false, + true, + None, + Some(format!("API error: {}", response.status().as_u16())), + )); + } + + let payload: Value = match response.json().await { + Ok(value) => value, + Err(err) => { + return Ok(build_result( + "zai-coding-plan", + "z.ai", + false, + true, + None, + Some(err.to_string()), + )) + } + }; + + let limits = payload + .get("data") + .and_then(|value| value.get("limits")) + .and_then(|value| value.as_array()) + .cloned() + .unwrap_or_default(); + let tokens_limit = limits + .iter() + .find(|limit| limit.get("type").and_then(|value| value.as_str()) == Some("TOKENS_LIMIT")); + + let mut windows = HashMap::new(); + if let Some(limit) = tokens_limit { + let window_seconds = resolve_window_seconds(limit); + let window_label = resolve_window_label(window_seconds); + let reset_at = normalize_timestamp(limit.get("nextResetTime")); + let used_percent = parse_number(limit.get("percentage")); + + windows.insert( + window_label, + to_usage_window(used_percent, window_seconds, reset_at), + ); + } + + Ok(build_result( + "zai-coding-plan", + "z.ai", + true, + true, + Some(ProviderUsage { + windows, + models: None, + }), + None, + )) +} + +pub async fn fetch_quota_for_provider(client: &Client, provider_id: &str) -> Result { + match provider_id { + "openai" => fetch_openai_quota(client).await, + "google" => fetch_google_quota(client).await, + "zai-coding-plan" => fetch_zai_quota(client).await, + _ => Ok(build_result( + provider_id, + provider_id, + false, + false, + None, + Some("Unsupported provider".to_string()), + )), + } +} diff --git a/packages/ui/src/components/sections/usage/UsageCard.tsx b/packages/ui/src/components/sections/usage/UsageCard.tsx index f295c25b..0bd1d924 100644 --- a/packages/ui/src/components/sections/usage/UsageCard.tsx +++ b/packages/ui/src/components/sections/usage/UsageCard.tsx @@ -15,7 +15,7 @@ export const UsageCard: React.FC = ({ title, window, subtitle }) const windowLabel = formatWindowLabel(title); return ( -
+
{windowLabel}
diff --git a/packages/ui/src/components/sections/usage/UsagePage.tsx b/packages/ui/src/components/sections/usage/UsagePage.tsx index de9f97a6..a4409051 100644 --- a/packages/ui/src/components/sections/usage/UsagePage.tsx +++ b/packages/ui/src/components/sections/usage/UsagePage.tsx @@ -69,20 +69,20 @@ export const UsagePage: React.FC = () => {
{!selectedResult && ( -
+

No usage data available yet.

)} {error && ( -
+

Failed to refresh usage data.

{error}

)} {selectedResult && !selectedResult.configured && ( -
+

Provider is not configured yet.

Add credentials in the Providers tab to enable usage tracking. @@ -114,7 +114,7 @@ export const UsagePage: React.FC = () => { {selectedResult?.configured && usage && Object.keys(usage.windows ?? {}).length === 0 && Object.keys(usage.models ?? {}).length === 0 && ( -

+

No quota windows reported for this provider.

)} diff --git a/packages/ui/src/components/sections/usage/UsageProgressBar.tsx b/packages/ui/src/components/sections/usage/UsageProgressBar.tsx index e8d58840..b0d02228 100644 --- a/packages/ui/src/components/sections/usage/UsageProgressBar.tsx +++ b/packages/ui/src/components/sections/usage/UsageProgressBar.tsx @@ -11,17 +11,17 @@ export const UsageProgressBar: React.FC = ({ percent, cla const clamped = clampPercent(percent) ?? 0; const tone = resolveUsageTone(percent); - const fillClass = tone === 'critical' - ? 'from-rose-500 to-rose-400' + const fillStyle = tone === 'critical' + ? { backgroundColor: 'var(--status-error)' } : tone === 'warn' - ? 'from-amber-500 to-amber-400' - : 'from-emerald-500 to-emerald-400'; + ? { backgroundColor: 'var(--status-warning)' } + : { backgroundColor: 'var(--status-success)' }; return ( -
+
= ({ onItemSelect }) => {
Total {QUOTA_PROVIDERS.length}
- + + + + + + + + Auto-refresh usage data at set interval + +