From dbab66afc9dec2d69ab00df70d5856b5c6a47bfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=AC=E9=A3=8E=E5=90=B9=E9=9B=A8?= <2686367411@qq.com> Date: Wed, 11 Nov 2026 22:56:12 +0800 Subject: [PATCH] feat(quota): unified zero-risk quota monitoring with smart subscription toggle --- src-tauri/src/app_state.rs | 7 + src-tauri/src/bin/codeg_server.rs | 1 + src-tauri/src/commands/mod.rs | 1 + src-tauri/src/commands/quota.rs | 109 +++ src-tauri/src/lib.rs | 5 + src-tauri/src/models/mod.rs | 2 + src-tauri/src/models/quota.rs | 214 ++++++ src-tauri/src/quota/antigravity.rs | 369 ++++++++++ src-tauri/src/quota/codex.rs | 647 ++++++++++++++++++ src-tauri/src/quota/mod.rs | 441 ++++++++++++ src-tauri/src/web/handlers/mod.rs | 1 + src-tauri/src/web/handlers/quota.rs | 35 + src-tauri/src/web/mod.rs | 4 + src-tauri/src/web/router.rs | 5 + .../chat/composer-quota-badge.test.tsx | 257 +++++++ src/components/chat/composer-quota-badge.tsx | 558 +++++++++++++++ src/components/chat/message-input.test.tsx | 3 + src/components/chat/message-input.tsx | 2 + .../settings/acp-agent-settings.tsx | 203 ++++++ src/i18n/messages/ar.json | 39 +- src/i18n/messages/de.json | 23 + src/i18n/messages/en.json | 23 + src/i18n/messages/es.json | 23 + src/i18n/messages/fr.json | 23 + src/i18n/messages/ja.json | 23 + src/i18n/messages/ko.json | 23 + src/i18n/messages/pt.json | 23 + src/i18n/messages/zh-CN.json | 23 + src/i18n/messages/zh-TW.json | 23 + src/lib/api.ts | 23 + src/lib/transport/types.ts | 11 +- src/lib/transport/web-transport.test.ts | 58 ++ src/lib/transport/web-transport.ts | 13 +- src/lib/types.ts | 24 + 34 files changed, 3226 insertions(+), 13 deletions(-) create mode 100644 src-tauri/src/commands/quota.rs create mode 100644 src-tauri/src/models/quota.rs create mode 100644 src-tauri/src/quota/antigravity.rs create mode 100644 src-tauri/src/quota/codex.rs create mode 100644 src-tauri/src/quota/mod.rs create mode 100644 src-tauri/src/web/handlers/quota.rs create mode 100644 src/components/chat/composer-quota-badge.test.tsx create mode 100644 src/components/chat/composer-quota-badge.tsx diff --git a/src-tauri/src/app_state.rs b/src-tauri/src/app_state.rs index 77ab999fe21..b1ee91c2bd8 100644 --- a/src-tauri/src/app_state.rs +++ b/src-tauri/src/app_state.rs @@ -96,6 +96,8 @@ pub struct AppState { /// The upgrade UI subscribes to it and re-syncs from a snapshot on mount, /// so download progress survives settings-page navigation and reloads. pub update_state: crate::update::AppUpdateStateHandle, + /// Quota tracking and fetch manager for supported agents. + pub quota_manager: Arc, /// Headless `codeg-cli` runs. Built in both runtimes at startup; its /// routes are served under `/api` and on the loopback listener. pub cli_runs: Arc, @@ -124,6 +126,10 @@ pub fn default_chat_channel_manager() -> ChatChannelManager { ChatChannelManager::new() } +pub fn default_quota_manager() -> Arc { + Arc::new(crate::quota::QuotaManager::new()) +} + /// Build the delegation broker + token registry + per-process UDS socket /// path. Shared between codeg-server bootstrap and the Tauri `setup` block /// so both modes apply identical depth limit + timeout defaults. @@ -300,6 +306,7 @@ impl AppState { computer_service: std::sync::OnceLock::new(), system_op_lock: default_system_op_lock(), update_state: default_update_state(), + quota_manager: Arc::new(crate::quota::QuotaManager::new()), cli_runs, cli_access, } diff --git a/src-tauri/src/bin/codeg_server.rs b/src-tauri/src/bin/codeg_server.rs index 7179953157f..4c71f3e5c0d 100644 --- a/src-tauri/src/bin/codeg_server.rs +++ b/src-tauri/src/bin/codeg_server.rs @@ -399,6 +399,7 @@ async fn async_main( computer_service: std::sync::OnceLock::new(), system_op_lock: codeg_lib::app_state::default_system_op_lock(), update_state: codeg_lib::app_state::default_update_state(), + quota_manager: codeg_lib::app_state::default_quota_manager(), cli_runs, cli_access, }); diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index c3bdf783993..8d7d7f544f5 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -69,6 +69,7 @@ pub mod pet; pub mod project_boot; pub mod question; pub mod quick_messages; +pub mod quota; #[cfg(feature = "tauri-runtime")] pub mod remote_proxy; #[cfg(feature = "tauri-runtime")] diff --git a/src-tauri/src/commands/quota.rs b/src-tauri/src/commands/quota.rs new file mode 100644 index 00000000000..1350d0c43f9 --- /dev/null +++ b/src-tauri/src/commands/quota.rs @@ -0,0 +1,109 @@ +#[cfg(feature = "tauri-runtime")] +use std::sync::Arc; + +use crate::app_error::AppCommandError; +use crate::db::service::agent_setting_service; +use crate::db::AppDatabase; +use crate::models::{AgentQuotaInfo, AgentType}; +use crate::parsers::AgentEnv; +use crate::quota::QuotaManager; + +async fn resolve_agent_quota_context( + db: Option<&AppDatabase>, + agent_type_str: &str, +) -> (AgentEnv, Option) { + let Some(db) = db else { + return (AgentEnv::inherited(), None); + }; + let Ok(parsed_type) = serde_json::from_value::(serde_json::Value::String( + agent_type_str.to_string(), + )) else { + return (AgentEnv::inherited(), None); + }; + + let setting = agent_setting_service::get_by_agent_type(&db.conn, parsed_type) + .await + .ok() + .flatten(); + let model_provider_id = setting.as_ref().and_then(|s| s.model_provider_id); + let local_config_json = crate::commands::acp::load_agent_local_config_json(parsed_type); + let runtime_env = crate::commands::acp::build_runtime_env_from_setting( + parsed_type, + setting.as_ref(), + local_config_json.as_deref(), + ); + + (AgentEnv::from_overrides(runtime_env), model_provider_id) +} + +pub async fn get_agent_quota_core( + quota_manager: &QuotaManager, + agent_type: &str, +) -> Result { + get_agent_quota_with_db_core(quota_manager, None, agent_type).await +} + +pub async fn get_agent_quota_with_db_core( + quota_manager: &QuotaManager, + db: Option<&AppDatabase>, + agent_type: &str, +) -> Result { + let (env, model_provider_id) = resolve_agent_quota_context(db, agent_type).await; + quota_manager + .get_quota_for_env(agent_type, &env, model_provider_id) + .await + .map_err(|e| AppCommandError::external_command("Failed to get agent quota", e)) +} + +pub async fn refresh_agent_quota_core( + quota_manager: &QuotaManager, + agent_type: &str, +) -> Result { + refresh_agent_quota_with_db_core(quota_manager, None, agent_type).await +} + +pub async fn refresh_agent_quota_with_db_core( + quota_manager: &QuotaManager, + db: Option<&AppDatabase>, + agent_type: &str, +) -> Result { + let (env, model_provider_id) = resolve_agent_quota_context(db, agent_type).await; + quota_manager + .fetch_quota_for_env(agent_type, &env, model_provider_id) + .await + .map_err(|e| AppCommandError::external_command("Failed to refresh agent quota", e)) +} + +#[cfg_attr(feature = "tauri-runtime", tauri::command)] +pub async fn get_agent_quota( + #[cfg(feature = "tauri-runtime")] quota_manager: tauri::State<'_, Arc>, + #[cfg(feature = "tauri-runtime")] db: tauri::State<'_, AppDatabase>, + agent_type: String, +) -> Result { + #[cfg(feature = "tauri-runtime")] + { + get_agent_quota_with_db_core("a_manager, Some(&db), &agent_type).await + } + #[cfg(not(feature = "tauri-runtime"))] + { + let _ = agent_type; + Err(AppCommandError::configuration_invalid("tauri-only command")) + } +} + +#[cfg_attr(feature = "tauri-runtime", tauri::command)] +pub async fn refresh_agent_quota( + #[cfg(feature = "tauri-runtime")] quota_manager: tauri::State<'_, Arc>, + #[cfg(feature = "tauri-runtime")] db: tauri::State<'_, AppDatabase>, + agent_type: String, +) -> Result { + #[cfg(feature = "tauri-runtime")] + { + refresh_agent_quota_with_db_core("a_manager, Some(&db), &agent_type).await + } + #[cfg(not(feature = "tauri-runtime"))] + { + let _ = agent_type; + Err(AppCommandError::configuration_invalid("tauri-only command")) + } +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index e5900e5aef0..a91a124d531 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -59,6 +59,7 @@ pub mod pets; #[cfg(feature = "tauri-runtime")] pub mod preferences; pub mod process; +pub mod quota; #[cfg(all(feature = "tauri-runtime", target_os = "macos"))] mod relaunch; pub mod runtime_lock; @@ -118,6 +119,7 @@ mod tauri_app { folders, logging as logging_commands, mcp as mcp_commands, model_provider as model_provider_commands, notification, pet as pet_commands, project_boot, question as question_commands, quick_messages as quick_messages_commands, + quota as quota_commands, remote_proxy as remote_proxy_commands, remote_workspace as remote_workspace_commands, science as science_commands, session_info as session_info_commands, @@ -755,6 +757,7 @@ mod tauri_app { // embedded web server's AppState so HTTP and webview clients see the // same download progress; lets the upgrade UI survive navigation. .manage(crate::update::new_update_state_handle()) + .manage(std::sync::Arc::new(crate::quota::QuotaManager::new())) .setup(|app| { let app_data_dir = app.path().app_data_dir()?; @@ -2300,6 +2303,8 @@ mod tauri_app { token_usage_commands::token_usage_facets, token_usage_commands::token_usage_status, token_usage_commands::token_usage_sync, + quota_commands::get_agent_quota, + quota_commands::refresh_agent_quota, work_task_commands::work_task_list, work_task_commands::work_task_get, work_task_commands::work_task_events, diff --git a/src-tauri/src/models/mod.rs b/src-tauri/src/models/mod.rs index 43960b7cd48..2e1e1bf8117 100644 --- a/src-tauri/src/models/mod.rs +++ b/src-tauri/src/models/mod.rs @@ -10,6 +10,7 @@ pub mod message; pub mod model_provider; pub mod pet; pub mod quick_message; +pub mod quota; pub mod remote_workspace_connection; pub mod system; pub mod token_usage; @@ -43,6 +44,7 @@ pub use message::{ MessageTurn, TurnRole, TurnUsage, UnifiedMessage, }; pub use quick_message::QuickMessageInfo; +pub use quota::{AgentQuotaInfo, QuotaWindow, SpendLimit}; pub use remote_workspace_connection::{ RemoteWorkspaceConnectionInfo, RemoteWorkspaceHeader, ToHeaderMap, }; diff --git a/src-tauri/src/models/quota.rs b/src-tauri/src/models/quota.rs new file mode 100644 index 00000000000..bc6e16e25c5 --- /dev/null +++ b/src-tauri/src/models/quota.rs @@ -0,0 +1,214 @@ +//! Agent quota and rate limit data models. + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +/// One time-window quota metric (e.g. 5-hour rolling window or weekly quota). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct QuotaWindow { + pub label: String, + pub used_percent: f64, + pub remaining_percent: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub resets_at: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reset_in_seconds: Option, +} + +impl QuotaWindow { + pub fn new( + label: impl Into, + used_percent: f64, + remaining_percent: f64, + resets_at: Option>, + reset_in_seconds: Option, + ) -> Self { + Self { + label: label.into(), + used_percent, + remaining_percent, + resets_at, + reset_in_seconds, + } + } + + /// Calculate quota percentages and reset countdown from usage and limit counters. + pub fn from_usage( + label: impl Into, + used: f64, + limit: f64, + resets_at: Option>, + now: Option>, + ) -> Self { + let (used_percent, remaining_percent) = if limit > 0.0 { + let used_pct = (used / limit * 100.0).clamp(0.0, 100.0); + (used_pct, 100.0 - used_pct) + } else { + (0.0, 100.0) + }; + + let reset_in_seconds = resets_at.map(|reset_time| { + let current = now.unwrap_or_else(Utc::now); + (reset_time - current).num_seconds().max(0) + }); + + Self { + label: label.into(), + used_percent, + remaining_percent, + resets_at, + reset_in_seconds, + } + } +} + +/// Spending limit information in USD. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SpendLimit { + pub used_usd: f64, + pub limit_usd: f64, +} + +impl SpendLimit { + pub fn new(used_usd: f64, limit_usd: f64) -> Self { + Self { + used_usd, + limit_usd, + } + } +} + +/// Consolidated quota and rate limit status for an agent. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentQuotaInfo { + pub agent_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub plan_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub short_window: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub weekly_window: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub spend_limit: Option, + pub last_updated: DateTime, +} + +impl AgentQuotaInfo { + pub fn new(agent_type: impl Into, last_updated: DateTime) -> Self { + Self { + agent_type: agent_type.into(), + plan_name: None, + short_window: None, + weekly_window: None, + spend_limit: None, + last_updated, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + + #[test] + fn test_quota_window_serialization_camel_case() { + let resets_at = Utc.with_ymd_and_hms(2026, 9, 26, 20, 0, 0).unwrap(); + let window = QuotaWindow { + label: "5-Hour Window".to_string(), + used_percent: 35.5, + remaining_percent: 64.5, + resets_at: Some(resets_at), + reset_in_seconds: Some(13156), + }; + + let json_str = serde_json::to_string(&window).expect("serialize QuotaWindow"); + assert!(json_str.contains("\"label\":\"5-Hour Window\"")); + assert!(json_str.contains("\"usedPercent\":35.5")); + assert!(json_str.contains("\"remainingPercent\":64.5")); + assert!(json_str.contains("\"resetsAt\":\"2026-09-26T20:00:00Z\"")); + assert!(json_str.contains("\"resetInSeconds\":13156")); + + let deserialized: QuotaWindow = + serde_json::from_str(&json_str).expect("deserialize QuotaWindow"); + assert_eq!(deserialized, window); + } + + #[test] + fn test_spend_limit_serialization_camel_case() { + let spend = SpendLimit { + used_usd: 12.5, + limit_usd: 50.0, + }; + + let json_str = serde_json::to_string(&spend).expect("serialize SpendLimit"); + assert!(json_str.contains("\"usedUsd\":12.5")); + assert!(json_str.contains("\"limitUsd\":50.0")); + + let deserialized: SpendLimit = + serde_json::from_str(&json_str).expect("deserialize SpendLimit"); + assert_eq!(deserialized, spend); + } + + #[test] + fn test_agent_quota_info_full_serialization() { + let now = Utc.with_ymd_and_hms(2026, 9, 26, 16, 20, 0).unwrap(); + let short_reset = Utc.with_ymd_and_hms(2026, 9, 26, 21, 0, 0).unwrap(); + let weekly_reset = Utc.with_ymd_and_hms(2026, 10, 1, 0, 0, 0).unwrap(); + + let info = AgentQuotaInfo { + agent_type: "codex".to_string(), + plan_name: Some("ChatGPT Plus".to_string()), + short_window: Some(QuotaWindow::new( + "5-Hour Window", + 20.0, + 80.0, + Some(short_reset), + Some(16800), + )), + weekly_window: Some(QuotaWindow::new( + "Weekly Window", + 55.0, + 45.0, + Some(weekly_reset), + Some(373200), + )), + spend_limit: Some(SpendLimit::new(5.0, 20.0)), + last_updated: now, + }; + + let json_val = serde_json::to_value(&info).expect("to_value"); + assert_eq!(json_val["agentType"], "codex"); + assert_eq!(json_val["planName"], "ChatGPT Plus"); + assert_eq!(json_val["shortWindow"]["label"], "5-Hour Window"); + assert_eq!(json_val["shortWindow"]["usedPercent"], 20.0); + assert_eq!(json_val["shortWindow"]["remainingPercent"], 80.0); + assert_eq!( + json_val["shortWindow"]["resetsAt"], + "2026-09-26T21:00:00Z" + ); + assert_eq!(json_val["shortWindow"]["resetInSeconds"], 16800); + assert_eq!(json_val["weeklyWindow"]["label"], "Weekly Window"); + assert_eq!(json_val["spendLimit"]["usedUsd"], 5.0); + assert_eq!(json_val["spendLimit"]["limitUsd"], 20.0); + assert_eq!(json_val["lastUpdated"], "2026-09-26T16:20:00Z"); + + let deserialized: AgentQuotaInfo = + serde_json::from_value(json_val).expect("from_value AgentQuotaInfo"); + assert_eq!(deserialized, info); + } + + #[test] + fn test_quota_window_from_usage_helper() { + let now = Utc.with_ymd_and_hms(2026, 9, 26, 10, 0, 0).unwrap(); + let reset = Utc.with_ymd_and_hms(2026, 9, 26, 15, 0, 0).unwrap(); + + let window = QuotaWindow::from_usage("5-Hour Window", 25.0, 100.0, Some(reset), Some(now)); + assert_eq!(window.used_percent, 25.0); + assert_eq!(window.remaining_percent, 75.0); + assert_eq!(window.reset_in_seconds, Some(5 * 3600)); + } +} diff --git a/src-tauri/src/quota/antigravity.rs b/src-tauri/src/quota/antigravity.rs new file mode 100644 index 00000000000..8c1821c0e2a --- /dev/null +++ b/src-tauri/src/quota/antigravity.rs @@ -0,0 +1,369 @@ +//! Antigravity quota parsing and safe zero-network probe helper module. +//! +//! Zero-Risk Architectural Decision: +//! Active probing via simulated HTTP requests to Google's internal CCPA endpoints +//! (`daily-cloudcode-pa.googleapis.com`) with spoofed User-Agent headers is +//! strictly prohibited to eliminate account ban / risk-control triggers. +//! +//! Antigravity quota display operates in a zero-network passive mode: +//! 1. Parses session turn usage / status events sent natively by the running +//! `agy_acp_server` process over the ACP channel. +//! 2. If no native turn quota data is available, returns a safe fallback without +//! initiating any unsolicited background network requests. + +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; +use serde::Deserialize; + +use crate::models::{AgentQuotaInfo, QuotaWindow, SpendLimit}; +use crate::parsers::AgentEnv; + +/// Helper to parse an Antigravity quota window with percentage and reset time. +pub fn parse_antigravity_window( + label: impl Into, + used_percent: f64, + resets_at: Option>, + now: Option>, +) -> QuotaWindow { + let used_clamped = used_percent.clamp(0.0, 100.0); + let remaining = (100.0 - used_clamped).max(0.0); + let reset_in_seconds = resets_at.map(|reset_time| { + let current = now.unwrap_or_else(Utc::now); + (reset_time - current).num_seconds().max(0) + }); + + QuotaWindow::new( + label, + used_clamped, + remaining, + resets_at, + reset_in_seconds, + ) +} + +/// Helper to parse an Antigravity quota window from raw usage and limit counters. +pub fn parse_antigravity_usage_window( + label: impl Into, + used: f64, + limit: f64, + resets_at: Option>, + now: Option>, +) -> QuotaWindow { + QuotaWindow::from_usage(label, used, limit, resets_at, now) +} + +#[derive(Debug, Deserialize)] +struct RawWindowPayload { + #[serde(default)] + label: Option, + #[serde(default)] + used: Option, + #[serde(default)] + limit: Option, + #[serde(default)] + used_percent: Option, + #[serde(default)] + resets_at: Option>, +} + +#[derive(Debug, Deserialize)] +struct RawSpendPayload { + used_usd: f64, + limit_usd: f64, +} + +#[derive(Debug, Deserialize)] +struct RawAntigravityQuotaPayload { + #[serde(default)] + plan: Option, + #[serde(default)] + short_window: Option, + #[serde(default)] + weekly_window: Option, + #[serde(default)] + spend_limit: Option, +} + +/// Parse a raw JSON payload from Antigravity/Gemini into `AgentQuotaInfo`. +pub fn parse_antigravity_quota_json( + raw_json: &str, + now: Option>, +) -> Result { + let payload: RawAntigravityQuotaPayload = serde_json::from_str(raw_json) + .map_err(|e| format!("failed to parse antigravity quota JSON: {e}"))?; + + let short_window = payload.short_window.map(|w| { + let label = w.label.unwrap_or_else(|| "Short-Term Quota".to_string()); + if let (Some(used), Some(limit)) = (w.used, w.limit) { + parse_antigravity_usage_window(label, used, limit, w.resets_at, now) + } else { + let used_pct = w.used_percent.unwrap_or(0.0); + parse_antigravity_window(label, used_pct, w.resets_at, now) + } + }); + + let weekly_window = payload.weekly_window.map(|w| { + let label = w.label.unwrap_or_else(|| "Weekly Quota".to_string()); + if let (Some(used), Some(limit)) = (w.used, w.limit) { + parse_antigravity_usage_window(label, used, limit, w.resets_at, now) + } else { + let used_pct = w.used_percent.unwrap_or(0.0); + parse_antigravity_window(label, used_pct, w.resets_at, now) + } + }); + + let spend_limit = payload + .spend_limit + .map(|s| SpendLimit::new(s.used_usd, s.limit_usd)); + + let last_updated = now.unwrap_or_else(Utc::now); + + Ok(AgentQuotaInfo { + agent_type: "antigravity".to_string(), + plan_name: payload.plan, + short_window, + weekly_window, + spend_limit, + last_updated, + }) +} + +// --------------------------------------------------------------------------- +// Zero-Network Local Authentication & Multi-Instance Resolution +// --------------------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +struct GoogleOAuthToken { + #[allow(dead_code)] + access_token: Option, + #[allow(dead_code)] + expires_at: Option, +} + +/// Resolves the Antigravity `acp_token.json` or `oauth_token.json` path on disk. +/// Respects `GEMINI_HOME` so custom multi-instance agents (e.g. `custom:antigravity-b`) +/// find their own isolated tokens. +pub fn resolve_antigravity_token_path_for_env(env: &AgentEnv) -> Option { + if let Some(custom_gemini_home) = env.var("GEMINI_HOME") { + let trimmed = custom_gemini_home.trim(); + if !trimmed.is_empty() { + let base = Path::new(trimmed); + let candidate1 = base.join("antigravity-acp").join("acp_token.json"); + if candidate1.exists() { + return Some(candidate1); + } + let candidate2 = base + .join("antigravity-engine") + .join("oauth_token.json"); + if candidate2.exists() { + return Some(candidate2); + } + let candidate3 = base.join("acp_token.json"); + if candidate3.exists() { + return Some(candidate3); + } + return Some(candidate1); + } + } + + let home_dir = dirs::home_dir()?; + let candidate1 = home_dir + .join(".gemini") + .join("antigravity-acp") + .join("acp_token.json"); + if candidate1.exists() { + return Some(candidate1); + } + + let candidate2 = home_dir + .join(".gemini") + .join("antigravity-engine") + .join("oauth_token.json"); + if candidate2.exists() { + return Some(candidate2); + } + + Some(candidate1) +} + +/// Checks whether an Antigravity agent is configured for official subscription mode +/// (OAuth personal or business) as opposed to a custom Gemini API Key or 3rd-party Model Provider. +pub fn is_antigravity_official_subscription(env: &AgentEnv, model_provider_id: Option) -> bool { + if model_provider_id.is_some() { + return false; + } + + if let Some(method) = env.var("AGY_AUTH_METHOD") { + let trimmed = method.trim(); + if trimmed == "gemini-api-key" || trimmed == "agent-platform" { + return false; + } + if trimmed == "oauth-personal" || trimmed == "oauth-business" { + return true; + } + } + + if env + .var("AGY_GEMINI_API_KEY") + .is_some_and(|v| !v.trim().is_empty()) + { + return false; + } + + if env.var("GEMINI_API_KEY").is_some_and(|v| !v.trim().is_empty()) { + return false; + } + + true +} + +/// Fetches quota status for Google Antigravity in safe zero-network mode. +/// +/// To completely eliminate any account ban / risk-control hazard, Codeg never +/// initiates spoofed HTTP requests to Google's internal CCPA endpoints. +/// Quota information is derived strictly from local credentials or native ACP session turns. +pub async fn fetch_antigravity_quota() -> Result { + fetch_antigravity_quota_for_env("antigravity", &AgentEnv::inherited()).await +} + +pub async fn fetch_antigravity_quota_for_env( + agent_type_label: &str, + env: &AgentEnv, +) -> Result { + let now = Utc::now(); + + let Some(token_path) = resolve_antigravity_token_path_for_env(env) else { + return Ok(AgentQuotaInfo { + agent_type: agent_type_label.to_string(), + plan_name: Some("Not Logged In".to_string()), + short_window: None, + weekly_window: None, + spend_limit: None, + last_updated: now, + }); + }; + + // If local token file exists and parses, mark the account as logged in via official OAuth. + // Zero external HTTP requests are performed. + let is_logged_in = if token_path.exists() { + std::fs::read_to_string(&token_path) + .ok() + .and_then(|c| serde_json::from_str::(&c).ok()) + .is_some() + } else { + false + }; + + let plan_name = if is_logged_in { + Some("Google AI (OAuth)".to_string()) + } else { + Some("Not Logged In".to_string()) + }; + + Ok(AgentQuotaInfo { + agent_type: agent_type_label.to_string(), + plan_name, + short_window: None, + weekly_window: None, + spend_limit: None, + last_updated: now, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_antigravity_window() { + let now = Utc::now(); + let resets_at = now + chrono::Duration::hours(2); + let win = parse_antigravity_window("Gemini 5-Hour", 35.0, Some(resets_at), Some(now)); + + assert_eq!(win.label, "Gemini 5-Hour"); + assert_eq!(win.used_percent, 35.0); + assert_eq!(win.remaining_percent, 65.0); + assert!(win.reset_in_seconds.unwrap() > 7100); + } + + #[test] + fn test_parse_antigravity_usage_window() { + let now = Utc::now(); + let resets_at = now + chrono::Duration::days(1); + let win = parse_antigravity_usage_window("Tokens", 250.0, 1000.0, Some(resets_at), Some(now)); + + assert_eq!(win.used_percent, 25.0); + assert_eq!(win.remaining_percent, 75.0); + } + + #[test] + fn test_parse_antigravity_quota_json() { + let json = r#"{ + "plan": "Google AI Ultra", + "short_window": { + "label": "5-Hour Window", + "used_percent": 15.5 + }, + "weekly_window": { + "label": "Weekly Window", + "used": 40.0, + "limit": 100.0 + }, + "spend_limit": { + "used_usd": 12.50, + "limit_usd": 50.00 + } + }"#; + + let quota = parse_antigravity_quota_json(json, None).expect("parse antigravity json"); + assert_eq!(quota.agent_type, "antigravity"); + assert_eq!(quota.plan_name, Some("Google AI Ultra".to_string())); + + let short = quota.short_window.unwrap(); + assert_eq!(short.used_percent, 15.5); + assert_eq!(short.remaining_percent, 84.5); + + let weekly = quota.weekly_window.unwrap(); + assert_eq!(weekly.used_percent, 40.0); + assert_eq!(weekly.remaining_percent, 60.0); + + let spend = quota.spend_limit.unwrap(); + assert_eq!(spend.used_usd, 12.50); + assert_eq!(spend.limit_usd, 50.00); + } + + #[test] + fn test_is_antigravity_official_subscription() { + let empty_env = AgentEnv::inherited(); + assert!(is_antigravity_official_subscription(&empty_env, None)); + assert!(!is_antigravity_official_subscription(&empty_env, Some(1))); + + let oauth_env = + AgentEnv::from_env_json(Some(r#"{"AGY_AUTH_METHOD":"oauth-personal"}"#)); + assert!(is_antigravity_official_subscription(&oauth_env, None)); + + let apikey_env = + AgentEnv::from_env_json(Some(r#"{"AGY_AUTH_METHOD":"gemini-api-key"}"#)); + assert!(!is_antigravity_official_subscription(&apikey_env, None)); + } + + #[test] + fn test_resolve_antigravity_token_path_respects_gemini_home() { + let custom_dir = std::env::temp_dir().join("test_gemini_home_quota"); + let token_dir = custom_dir.join("antigravity-acp"); + std::fs::create_dir_all(&token_dir).ok(); + let token_file = token_dir.join("acp_token.json"); + std::fs::write(&token_file, "{}").ok(); + + let env = AgentEnv::from_env_json(Some(&format!( + r#"{{"GEMINI_HOME":{}}}"#, + serde_json::to_string(&custom_dir.to_string_lossy()).unwrap() + ))); + + let resolved = resolve_antigravity_token_path_for_env(&env); + assert_eq!(resolved, Some(token_file)); + + std::fs::remove_dir_all(&custom_dir).ok(); + } +} diff --git a/src-tauri/src/quota/codex.rs b/src-tauri/src/quota/codex.rs new file mode 100644 index 00000000000..58a616d73c2 --- /dev/null +++ b/src-tauri/src/quota/codex.rs @@ -0,0 +1,647 @@ +//! Codex quota parsing and probe helper module. +//! +//! Primary path (Zero Network / Zero Risk): +//! Codex CLI writes a `token_count` event into its local session rollout JSONL +//! (`/sessions/YYYY/MM/DD/rollout-*.jsonl`) after every turn, which +//! carries `payload.rate_limits` (`primary`, `secondary`, `plan_type`, +//! `individual_limit`, etc.) directly from the model response headers. +//! +//! Fallback path (On-Demand Official Wham Endpoint): +//! Used only when no local rollout with `rate_limits` exists yet. + +use std::fs::File; +use std::io::{BufRead, BufReader, Read, Seek, SeekFrom}; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, TimeZone, Utc}; +use serde::Deserialize; +use walkdir::WalkDir; + +use crate::models::{AgentQuotaInfo, QuotaWindow, SpendLimit}; +use crate::parsers::AgentEnv; + +/// Helper to parse a quota window with percentage and reset time. +pub fn parse_codex_window( + label: impl Into, + used_percent: f64, + resets_at: Option>, + now: Option>, +) -> QuotaWindow { + let used_clamped = used_percent.clamp(0.0, 100.0); + let remaining = (100.0 - used_clamped).max(0.0); + let reset_in_seconds = resets_at.map(|reset_time| { + let current = now.unwrap_or_else(Utc::now); + (reset_time - current).num_seconds().max(0) + }); + + QuotaWindow::new( + label, + used_clamped, + remaining, + resets_at, + reset_in_seconds, + ) +} + +/// Helper to parse a quota window from raw usage and limit counters. +pub fn parse_codex_usage_window( + label: impl Into, + used: f64, + limit: f64, + resets_at: Option>, + now: Option>, +) -> QuotaWindow { + QuotaWindow::from_usage(label, used, limit, resets_at, now) +} + +#[derive(Debug, Deserialize)] +struct RawWindowPayload { + #[serde(default)] + label: Option, + #[serde(default)] + used: Option, + #[serde(default)] + limit: Option, + #[serde(default)] + used_percent: Option, + #[serde(default)] + resets_at: Option>, +} + +#[derive(Debug, Deserialize)] +struct RawSpendPayload { + used_usd: f64, + limit_usd: f64, +} + +#[derive(Debug, Deserialize)] +struct RawCodexQuotaPayload { + #[serde(default)] + plan: Option, + #[serde(default)] + short_window: Option, + #[serde(default)] + weekly_window: Option, + #[serde(default)] + spend_limit: Option, +} + +/// Parse a raw JSON payload from Codex CLI or API into `AgentQuotaInfo`. +pub fn parse_codex_quota_json( + raw_json: &str, + now: Option>, +) -> Result { + let payload: RawCodexQuotaPayload = serde_json::from_str(raw_json) + .map_err(|e| format!("failed to parse codex quota JSON: {e}"))?; + + let short_window = payload.short_window.map(|w| { + let label = w.label.unwrap_or_else(|| "5-Hour Window".to_string()); + if let (Some(used), Some(limit)) = (w.used, w.limit) { + parse_codex_usage_window(label, used, limit, w.resets_at, now) + } else { + let used_pct = w.used_percent.unwrap_or(0.0); + parse_codex_window(label, used_pct, w.resets_at, now) + } + }); + + let weekly_window = payload.weekly_window.map(|w| { + let label = w.label.unwrap_or_else(|| "Weekly Limit".to_string()); + if let (Some(used), Some(limit)) = (w.used, w.limit) { + parse_codex_usage_window(label, used, limit, w.resets_at, now) + } else { + let used_pct = w.used_percent.unwrap_or(0.0); + parse_codex_window(label, used_pct, w.resets_at, now) + } + }); + + let spend_limit = payload + .spend_limit + .map(|s| SpendLimit::new(s.used_usd, s.limit_usd)); + + let last_updated = now.unwrap_or_else(Utc::now); + + Ok(AgentQuotaInfo { + agent_type: "codex".to_string(), + plan_name: payload.plan, + short_window, + weekly_window, + spend_limit, + last_updated, + }) +} + +// --------------------------------------------------------------------------- +// Zero-Network Rollout JSONL `token_count.rate_limits` Reader +// --------------------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +struct RolloutRateLimitWindow { + used_percent: Option, + #[allow(dead_code)] + window_minutes: Option, + resets_at: Option, +} + +#[derive(Debug, Deserialize)] +struct RolloutRateLimits { + primary: Option, + secondary: Option, + plan_type: Option, +} + +fn format_plan_label(raw_plan: &str) -> String { + let trimmed = raw_plan.trim(); + let mut chars = trimmed.chars(); + match chars.next() { + None => String::new(), + Some(f) => f.to_uppercase().collect::() + chars.as_str() + " Plan", + } +} + +/// Parse a `rate_limits` JSON value from a Codex rollout `token_count` event. +pub fn parse_codex_rollout_rate_limits( + rate_limits_val: &serde_json::Value, + event_ts: Option>, + now: Option>, +) -> Option { + let rl: RolloutRateLimits = serde_json::from_value(rate_limits_val.clone()).ok()?; + if rl.primary.is_none() && rl.secondary.is_none() && rl.plan_type.is_none() { + return None; + } + + let current_now = now.unwrap_or_else(Utc::now); + + let short_window = rl.primary.map(|pw| { + let used_pct = pw.used_percent.unwrap_or(0.0); + let resets_at = pw.resets_at.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); + parse_codex_window("5-Hour Window", used_pct, resets_at, Some(current_now)) + }); + + let weekly_window = rl.secondary.map(|sw| { + let used_pct = sw.used_percent.unwrap_or(0.0); + let resets_at = sw.resets_at.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); + parse_codex_window("Weekly Limit", used_pct, resets_at, Some(current_now)) + }); + + let plan_name = rl + .plan_type + .as_deref() + .filter(|s| !s.trim().is_empty()) + .map(format_plan_label); + + Some(AgentQuotaInfo { + agent_type: "codex".to_string(), + plan_name, + short_window, + weekly_window, + spend_limit: None, + last_updated: event_ts.unwrap_or(current_now), + }) +} + +/// Extract the latest `token_count.rate_limits` from a single rollout JSONL file. +/// Reads from the tail chunk first (last 256 KB) for speed on multi-megabyte rollouts. +pub fn read_latest_rate_limits_from_rollout_file( + path: &Path, + now: Option>, +) -> Option { + let mut file = File::open(path).ok()?; + let len = file.metadata().ok()?.len(); + const TAIL_BYTES: u64 = 256 * 1024; + + let lines: Vec = if len > TAIL_BYTES { + file.seek(SeekFrom::Start(len - TAIL_BYTES)).ok()?; + let mut reader = BufReader::new(file); + let mut discard = String::new(); + let _ = reader.read_line(&mut discard); + reader.lines().map_while(Result::ok).collect() + } else { + BufReader::new(file).lines().map_while(Result::ok).collect() + }; + + for line in lines.iter().rev() { + if !line.contains("\"rate_limits\"") { + continue; + } + let Ok(val) = serde_json::from_str::(line) else { + continue; + }; + if val.get("type").and_then(|v| v.as_str()) != Some("event_msg") { + continue; + } + let Some(payload) = val.get("payload") else { + continue; + }; + if payload.get("type").and_then(|v| v.as_str()) != Some("token_count") { + continue; + } + let Some(rate_limits) = payload.get("rate_limits") else { + continue; + }; + if rate_limits.is_null() { + continue; + } + let event_ts = val + .get("timestamp") + .and_then(|v| v.as_str()) + .and_then(|s| DateTime::parse_from_rfc3339(s).ok()) + .map(|dt| dt.with_timezone(&Utc)); + if let Some(info) = parse_codex_rollout_rate_limits(rate_limits, event_ts, now) { + return Some(info); + } + } + + if len > TAIL_BYTES { + let mut full_file = File::open(path).ok()?; + let mut text = String::new(); + full_file.read_to_string(&mut text).ok()?; + for line in text.lines().rev() { + if !line.contains("\"rate_limits\"") { + continue; + } + let Ok(val) = serde_json::from_str::(line) else { + continue; + }; + let Some(payload) = val.get("payload") else { + continue; + }; + if payload.get("type").and_then(|v| v.as_str()) != Some("token_count") { + continue; + } + let Some(rate_limits) = payload.get("rate_limits") else { + continue; + }; + if rate_limits.is_null() { + continue; + } + let event_ts = val + .get("timestamp") + .and_then(|v| v.as_str()) + .and_then(|s| DateTime::parse_from_rfc3339(s).ok()) + .map(|dt| dt.with_timezone(&Utc)); + if let Some(info) = parse_codex_rollout_rate_limits(rate_limits, event_ts, now) { + return Some(info); + } + } + } + + None +} + +/// Scans the most recently modified rollout files under `/sessions` +/// to extract the newest `rate_limits` snapshot with zero network I/O. +pub fn read_latest_codex_quota_from_sessions_dir( + sessions_dir: &Path, + now: Option>, +) -> Option { + if !sessions_dir.exists() { + return None; + } + + let mut rollout_files: Vec<(std::time::SystemTime, PathBuf)> = WalkDir::new(sessions_dir) + .into_iter() + .filter_map(Result::ok) + .filter(|entry| { + let path = entry.path(); + path.extension().and_then(|e| e.to_str()) == Some("jsonl") + && path + .file_name() + .unwrap_or_default() + .to_string_lossy() + .starts_with("rollout-") + }) + .filter_map(|entry| { + let mtime = entry.metadata().ok()?.modified().ok()?; + Some((mtime, entry.into_path())) + }) + .collect(); + + rollout_files.sort_by_key(|a| std::cmp::Reverse(a.0)); + + for (_, path) in rollout_files.into_iter().take(8) { + if let Some(info) = read_latest_rate_limits_from_rollout_file(&path, now) { + return Some(info); + } + } + + None +} + +// --------------------------------------------------------------------------- +// Production OpenAI / ChatGPT Wham API Structures (Cold-Start Fallback) +// --------------------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +struct WhamWindow { + used_percent: Option, + #[allow(dead_code)] + limit_window_seconds: Option, + reset_after_seconds: Option, + reset_at: Option, +} + +#[derive(Debug, Deserialize)] +struct WhamRateLimit { + #[allow(dead_code)] + allowed: Option, + primary_window: Option, + secondary_window: Option, +} + +#[derive(Debug, Deserialize)] +struct WhamSpendControl { + #[allow(dead_code)] + reached: Option, + #[allow(dead_code)] + individual_limit: Option, +} + +#[derive(Debug, Deserialize)] +struct WhamUsageResponse { + plan_type: Option, + rate_limit: Option, + #[allow(dead_code)] + spend_control: Option, +} + +#[derive(Debug, Deserialize)] +struct CodexAuthTokens { + access_token: Option, + account_id: Option, +} + +#[derive(Debug, Deserialize)] +struct CodexAuthJson { + #[serde(default)] + auth_mode: Option, + #[serde(rename = "OPENAI_API_KEY", default)] + openai_api_key: Option, + tokens: Option, +} + +/// Resolves the `/auth.json` path for the given `AgentEnv`. +pub(crate) fn resolve_codex_auth_path_for_env(env: &AgentEnv) -> PathBuf { + crate::parsers::codex::resolve_codex_home_dir_for_env(env).join("auth.json") +} + +/// Checks whether Codex is currently configured for official ChatGPT subscription +/// login rather than a 3rd-party API key or custom model provider. +pub fn is_codex_official_subscription( + env: &AgentEnv, + model_provider_id: Option, +) -> bool { + if model_provider_id.is_some() { + return false; + } + if env + .var("OPENAI_API_KEY") + .is_some_and(|v| !v.trim().is_empty()) + || env + .var("OPENAI_BASE_URL") + .is_some_and(|v| !v.trim().is_empty()) + { + return false; + } + + let auth_path = resolve_codex_auth_path_for_env(env); + let Ok(content) = std::fs::read_to_string(&auth_path) else { + return true; + }; + let Ok(parsed) = serde_json::from_str::(&content) else { + return true; + }; + + if parsed.auth_mode.as_deref() == Some("chatgpt") { + return true; + } + if let Some(ref key) = parsed.openai_api_key { + if !key.trim().is_empty() { + return false; + } + } + true +} + +/// Reads local `auth.json` and extracts credentials. +fn read_codex_credentials_for_env(env: &AgentEnv) -> Option<(String, Option)> { + let auth_path = resolve_codex_auth_path_for_env(env); + if !auth_path.exists() { + return None; + } + let content = std::fs::read_to_string(&auth_path).ok()?; + let parsed: CodexAuthJson = serde_json::from_str(&content).ok()?; + let tokens = parsed.tokens?; + let access_token = tokens.access_token.filter(|s| !s.trim().is_empty())?; + let account_id = tokens.account_id.filter(|s| !s.trim().is_empty()); + Some((access_token, account_id)) +} + +/// Fetches quota status for Codex. +/// +/// 1. First checks local `/sessions` rollout JSONL files for the +/// latest `token_count.rate_limits` record (100% zero network, zero risk). +/// 2. Only if no local rollout record exists yet, falls back to a single +/// on-demand request to `wham/usage`. +pub async fn fetch_codex_quota() -> Result { + fetch_codex_quota_for_env(&AgentEnv::inherited()).await +} + +pub async fn fetch_codex_quota_for_env(env: &AgentEnv) -> Result { + let now = Utc::now(); + let codex_home = crate::parsers::codex::resolve_codex_home_dir_for_env(env); + let sessions_dir = codex_home.join("sessions"); + + if let Some(local_quota) = read_latest_codex_quota_from_sessions_dir(&sessions_dir, Some(now)) { + return Ok(local_quota); + } + + let (access_token, account_id) = match read_codex_credentials_for_env(env) { + Some(creds) => creds, + None => { + return Ok(AgentQuotaInfo { + agent_type: "codex".to_string(), + plan_name: Some("Not Logged In".to_string()), + short_window: None, + weekly_window: None, + spend_limit: None, + last_updated: now, + }); + } + }; + + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(10)) + .build() + .map_err(|e| format!("failed to initialize HTTP client: {e}"))?; + + let mut req = client + .get("https://chatgpt.com/backend-api/wham/usage") + .header("Authorization", format!("Bearer {access_token}")) + .header("User-Agent", "codex-cli/0.156.0"); + + if let Some(ref acc_id) = account_id { + req = req.header("ChatGPT-Account-Id", acc_id); + } + + let response = req + .send() + .await + .map_err(|e| format!("Codex usage API request failed: {e}"))?; + + let status = response.status(); + if !status.is_success() { + return Err(format!("Codex usage API returned HTTP {status}")); + } + + let wham: WhamUsageResponse = response + .json() + .await + .map_err(|e| format!("failed to deserialize Codex usage payload: {e}"))?; + + let plan_name = wham.plan_type.as_deref().map(format_plan_label); + + let mut short_window = None; + let mut weekly_window = None; + + if let Some(rate_limit) = wham.rate_limit { + if let Some(pw) = rate_limit.primary_window { + let used_pct = pw.used_percent.unwrap_or(0.0); + let resets_at = pw.reset_at.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); + let mut win = parse_codex_window("5-Hour Window", used_pct, resets_at, Some(now)); + if win.reset_in_seconds.is_none() { + win.reset_in_seconds = pw.reset_after_seconds; + } + short_window = Some(win); + } + + if let Some(sw) = rate_limit.secondary_window { + let used_pct = sw.used_percent.unwrap_or(0.0); + let resets_at = sw.reset_at.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); + let mut win = parse_codex_window("Weekly Limit", used_pct, resets_at, Some(now)); + if win.reset_in_seconds.is_none() { + win.reset_in_seconds = sw.reset_after_seconds; + } + weekly_window = Some(win); + } + } + + Ok(AgentQuotaInfo { + agent_type: "codex".to_string(), + plan_name, + short_window, + weekly_window, + spend_limit: None, + last_updated: now, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_codex_window_percentage() { + let now = Utc.with_ymd_and_hms(2026, 9, 26, 12, 0, 0).unwrap(); + let reset = Utc.with_ymd_and_hms(2026, 9, 26, 14, 0, 0).unwrap(); + + let window = parse_codex_window("5-Hour Window", 42.0, Some(reset), Some(now)); + assert_eq!(window.label, "5-Hour Window"); + assert_eq!(window.used_percent, 42.0); + assert_eq!(window.remaining_percent, 58.0); + assert_eq!(window.resets_at, Some(reset)); + assert_eq!(window.reset_in_seconds, Some(7200)); + } + + #[test] + fn test_parse_codex_usage_window_counters() { + let now = Utc.with_ymd_and_hms(2026, 9, 26, 12, 0, 0).unwrap(); + let reset = Utc.with_ymd_and_hms(2026, 9, 27, 12, 0, 0).unwrap(); + + let window = parse_codex_usage_window("Weekly Limit", 30.0, 150.0, Some(reset), Some(now)); + assert_eq!(window.used_percent, 20.0); + assert_eq!(window.remaining_percent, 80.0); + assert_eq!(window.reset_in_seconds, Some(86400)); + } + + #[test] + fn test_parse_codex_quota_json() { + let now = Utc.with_ymd_and_hms(2026, 9, 26, 12, 0, 0).unwrap(); + let raw = r#"{ + "plan": "ChatGPT Pro", + "short_window": { + "label": "5-Hour Window", + "used": 15.0, + "limit": 100.0, + "resets_at": "2026-09-26T15:00:00Z" + }, + "weekly_window": { + "label": "Weekly Limit", + "used_percent": 60.0, + "resets_at": "2026-09-30T00:00:00Z" + }, + "spend_limit": { + "used_usd": 18.5, + "limit_usd": 100.0 + } + }"#; + + let info = parse_codex_quota_json(raw, Some(now)).expect("parse codex json"); + assert_eq!(info.agent_type, "codex"); + assert_eq!(info.plan_name.as_deref(), Some("ChatGPT Pro")); + + let short = info.short_window.expect("short_window"); + assert_eq!(short.used_percent, 15.0); + assert_eq!(short.remaining_percent, 85.0); + assert_eq!(short.reset_in_seconds, Some(3 * 3600)); + + let weekly = info.weekly_window.expect("weekly_window"); + assert_eq!(weekly.used_percent, 60.0); + assert_eq!(weekly.remaining_percent, 40.0); + + let spend = info.spend_limit.expect("spend_limit"); + assert_eq!(spend.used_usd, 18.5); + assert_eq!(spend.limit_usd, 100.0); + } + + #[test] + fn test_parse_codex_rollout_rate_limits() { + let now = Utc.with_ymd_and_hms(2026, 11, 10, 12, 0, 0).unwrap(); + let sample = serde_json::json!({ + "limit_id": "codex", + "primary": { + "used_percent": 2.0, + "window_minutes": 300, + "resets_at": 1794319200 + }, + "secondary": { + "used_percent": 98.0, + "window_minutes": 10080, + "resets_at": 1794388800 + }, + "plan_type": "team" + }); + + let info = parse_codex_rollout_rate_limits(&sample, Some(now), Some(now)) + .expect("should parse rollout rate limits"); + assert_eq!(info.agent_type, "codex"); + assert_eq!(info.plan_name.as_deref(), Some("Team Plan")); + assert_eq!(info.short_window.unwrap().remaining_percent, 98.0); + assert_eq!(info.weekly_window.unwrap().remaining_percent, 2.0); + } + + #[test] + fn test_read_latest_codex_quota_from_sessions_dir() { + let dir = tempfile::tempdir().unwrap(); + let day_dir = dir.path().join("2026").join("11").join("10"); + std::fs::create_dir_all(&day_dir).unwrap(); + let rollout_path = day_dir.join("rollout-2026-11-10T10-00-00-01a19a0a.jsonl"); + let line = r#"{"timestamp":"2026-11-10T10:05:00Z","type":"event_msg","payload":{"type":"token_count","info":{},"rate_limits":{"primary":{"used_percent":12.0,"window_minutes":300,"resets_at":1794319200},"secondary":{"used_percent":45.0,"window_minutes":10080,"resets_at":1794388800},"plan_type":"plus"}}}"#; + std::fs::write(&rollout_path, format!("{line}\n")).unwrap(); + + let info = read_latest_codex_quota_from_sessions_dir(dir.path(), None) + .expect("should find rollout quota"); + assert_eq!(info.plan_name.as_deref(), Some("Plus Plan")); + assert_eq!(info.short_window.unwrap().remaining_percent, 88.0); + assert_eq!(info.weekly_window.unwrap().remaining_percent, 55.0); + } +} diff --git a/src-tauri/src/quota/mod.rs b/src-tauri/src/quota/mod.rs new file mode 100644 index 00000000000..85d1a2aab4a --- /dev/null +++ b/src-tauri/src/quota/mod.rs @@ -0,0 +1,441 @@ +//! Quota and rate-limit manager for AI coding agents. +//! +//! Includes: +//! - `QuotaProbePolicy`: Smart per-agent `CODEG_QUOTA_PROBE_MODE` (`auto` | `on` | `off`) +//! that automatically enables quota probing when an agent is in official +//! account/subscription login mode and disables it when switched to a 3rd-party +//! API key, custom Base URL, or bound Model Provider. +//! - `QuotaManager`: Unified in-memory cache and dispatcher for Codex (zero-network +//! rollout log extraction + cold-start fallback) and Antigravity (including custom +//! `agy_acp_server` instances via `GEMINI_HOME`). + +pub mod antigravity; +pub mod codex; + +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use tokio::sync::RwLock; + +use crate::acp::file_system_runtime::env_value; +use crate::models::{AgentQuotaInfo, AgentType}; +use crate::parsers::AgentEnv; + +/// Per-agent `env_json` key controlling whether subscription quota probing is enabled. +/// Values: `"auto"` (default: enabled iff official account login is active), `"on"`, `"off"`. +pub const QUOTA_PROBE_MODE_ENV: &str = "CODEG_QUOTA_PROBE_MODE"; + +/// Default time-to-live for cached quota information (180 seconds = 3 minutes). +pub const DEFAULT_QUOTA_TTL: Duration = Duration::from_secs(180); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum QuotaProbeMode { + /// Follow authentication mode: enabled iff official account/subscription login is active. + Auto, + /// Explicitly enabled by the user. + On, + /// Explicitly disabled by the user even in official account login mode. + Off, +} + +impl QuotaProbeMode { + pub fn from_env_map(runtime_env: &BTreeMap) -> Self { + match env_value(runtime_env, QUOTA_PROBE_MODE_ENV).as_deref() { + None | Some("") | Some("auto") => Self::Auto, + Some("on") | Some("true") | Some("1") => Self::On, + Some("off") | Some("false") | Some("0") => Self::Off, + Some(_) => Self::Auto, + } + } + + pub fn from_agent_env(env: &AgentEnv) -> Self { + match env.var(QUOTA_PROBE_MODE_ENV).as_deref().map(str::trim) { + None | Some("") | Some("auto") => Self::Auto, + Some("on") | Some("true") | Some("1") => Self::On, + Some("off") | Some("false") | Some("0") => Self::Off, + Some(_) => Self::Auto, + } + } +} + +/// Checks whether a custom agent is an Antigravity (`agy_acp_server`) instance. +pub fn is_custom_antigravity_agent(agent_type: &str, env: &AgentEnv) -> bool { + let Some(custom_id) = agent_type.strip_prefix("custom:") else { + return false; + }; + if custom_id.to_ascii_lowercase().contains("antigravity") + || custom_id.to_ascii_lowercase().starts_with("agy") + { + return true; + } + if env + .var("GEMINI_HOME") + .is_some_and(|v| !v.trim().is_empty()) + || env + .var("AGY_AUTH_METHOD") + .is_some_and(|v| !v.trim().is_empty()) + { + return true; + } + if let Ok(parsed_type) = serde_json::from_value::(serde_json::Value::String( + agent_type.to_string(), + )) { + let meta = crate::acp::registry::get_agent_meta(parsed_type); + { + let desc_lower = meta.description.to_ascii_lowercase(); + let name_lower = meta.name.to_ascii_lowercase(); + if name_lower.contains("antigravity") || desc_lower.contains("antigravity") { + return true; + } + } + } + false +} + +/// Determines whether an agent is in official account/subscription login mode +/// (as opposed to 3rd-party API Key, custom Base URL, or bound Model Provider). +pub fn is_agent_official_subscription_mode( + agent_type: &str, + env: &AgentEnv, + model_provider_id: Option, +) -> bool { + if model_provider_id.is_some() { + return false; + } + match agent_type { + "codex" => codex::is_codex_official_subscription(env, model_provider_id), + "antigravity" => antigravity::is_antigravity_official_subscription(env, model_provider_id), + other if is_custom_antigravity_agent(other, env) => { + antigravity::is_antigravity_official_subscription(env, model_provider_id) + } + "claude_code" => { + env.var("ANTHROPIC_BASE_URL") + .is_none_or(|v| v.trim().is_empty()) + && env + .var("ANTHROPIC_API_KEY") + .is_none_or(|v| v.trim().is_empty()) + && env + .var("ANTHROPIC_AUTH_TOKEN") + .is_none_or(|v| v.trim().is_empty()) + } + "gemini" => { + env.var("GEMINI_API_KEY") + .is_none_or(|v| v.trim().is_empty()) + && env + .var("GOOGLE_API_KEY") + .is_none_or(|v| v.trim().is_empty()) + && env + .var("GOOGLE_GEMINI_BASE_URL") + .is_none_or(|v| v.trim().is_empty()) + && env + .var("GOOGLE_GENAI_USE_VERTEXAI") + .is_none_or(|v| v.trim() != "true") + } + "cursor" => { + env.var("CURSOR_AUTH_MODE") + .is_none_or(|v| v.trim() != "custom") + && env + .var("CURSOR_API_KEY") + .is_none_or(|v| v.trim().is_empty()) + } + "grok" => { + env.var("GROK_AUTH_MODE") + .is_none_or(|v| v.trim() == "subscription" || v.trim().is_empty()) + && env.var("XAI_API_KEY").is_none_or(|v| v.trim().is_empty()) + } + "kimi_code" => { + env.var("KIMI_API_KEY") + .is_none_or(|v| v.trim().is_empty()) + && env + .var("KIMI_BASE_URL") + .is_none_or(|v| v.trim().is_empty()) + } + "kiro" => env.var("KIRO_API_KEY").is_none_or(|v| v.trim().is_empty()), + "qoder" => env + .var("QODER_PERSONAL_ACCESS_TOKEN") + .is_none_or(|v| v.trim().is_empty()), + "trae" => env + .var("TRAECLI_PERSONAL_ACCESS_TOKEN") + .is_none_or(|v| v.trim().is_empty()), + "command_code" => env + .var("COMMAND_CODE_API_KEY") + .is_none_or(|v| v.trim().is_empty()), + "cline" => { + env.var("CLINE_API_KEY") + .is_none_or(|v| v.trim().is_empty()) + && env + .var("CLINE_BASE_URL") + .is_none_or(|v| v.trim().is_empty()) + } + "zcode" => true, + _ => false, + } +} + +/// Evaluates whether quota probing is effectively enabled for `agent_type` given +/// its environment and `model_provider_id`. +pub fn is_quota_probe_enabled( + agent_type: &str, + env: &AgentEnv, + model_provider_id: Option, +) -> bool { + if !is_agent_official_subscription_mode(agent_type, env, model_provider_id) { + return false; + } + match QuotaProbeMode::from_agent_env(env) { + QuotaProbeMode::Auto | QuotaProbeMode::On => true, + QuotaProbeMode::Off => false, + } +} + +/// Centralized manager for querying and caching agent quotas. +#[derive(Debug, Clone)] +pub struct QuotaManager { + cache: Arc>>, + ttl: Duration, +} + +impl Default for QuotaManager { + fn default() -> Self { + Self::new() + } +} + +impl QuotaManager { + /// Create a new QuotaManager with the default TTL (180s). + pub fn new() -> Self { + Self::with_ttl(DEFAULT_QUOTA_TTL) + } + + /// Create a new QuotaManager with a custom TTL. + pub fn with_ttl(ttl: Duration) -> Self { + Self { + cache: Arc::new(RwLock::new(HashMap::new())), + ttl, + } + } + + /// Get configured TTL. + pub fn ttl(&self) -> Duration { + self.ttl + } + + /// Retrieve cached quota info for `agent_type` if it exists and has not expired. + pub async fn get_cached(&self, agent_type: &str) -> Option { + let cache = self.cache.read().await; + if let Some((info, cached_at)) = cache.get(agent_type) { + if cached_at.elapsed() < self.ttl { + return Some(info.clone()); + } + } + None + } + + /// Insert or update the cached quota for `agent_type`. + pub async fn cache_quota(&self, agent_type: &str, info: AgentQuotaInfo) { + let mut cache = self.cache.write().await; + cache.insert(agent_type.to_string(), (info, Instant::now())); + } + + /// Invalidate the cache entry for a given agent type. + pub async fn invalidate(&self, agent_type: &str) { + let mut cache = self.cache.write().await; + cache.remove(agent_type); + } + + /// Clear all cached quota entries. + pub async fn clear(&self) { + let mut cache = self.cache.write().await; + cache.clear(); + } + + /// Get quota for an agent using the inherited environment. + pub async fn get_quota(&self, agent_type: &str) -> Result { + self.get_quota_for_env(agent_type, &AgentEnv::inherited(), None) + .await + } + + /// Get quota for an agent with its configured `AgentEnv` and `model_provider_id`. + /// Returns cached info if valid; otherwise fetches fresh info. + pub async fn get_quota_for_env( + &self, + agent_type: &str, + env: &AgentEnv, + model_provider_id: Option, + ) -> Result { + if !is_quota_probe_enabled(agent_type, env, model_provider_id) { + return Err(format!( + "Quota probing is disabled for agent '{agent_type}' (non-subscription mode or turned off)" + )); + } + if let Some(cached) = self.get_cached(agent_type).await { + return Ok(cached); + } + self.fetch_quota_for_env(agent_type, env, model_provider_id) + .await + } + + /// Fetch fresh quota from the respective agent backend using the inherited environment. + pub async fn fetch_quota(&self, agent_type: &str) -> Result { + self.fetch_quota_for_env(agent_type, &AgentEnv::inherited(), None) + .await + } + + /// Fetch fresh quota from the respective agent backend and cache it. + pub async fn fetch_quota_for_env( + &self, + agent_type: &str, + env: &AgentEnv, + model_provider_id: Option, + ) -> Result { + if !is_quota_probe_enabled(agent_type, env, model_provider_id) { + return Err(format!( + "Quota probing is disabled for agent '{agent_type}' (non-subscription mode or turned off)" + )); + } + + let info = match agent_type { + "codex" => codex::fetch_codex_quota_for_env(env).await?, + "antigravity" => antigravity::fetch_antigravity_quota_for_env("antigravity", env).await?, + other if is_custom_antigravity_agent(other, env) => { + antigravity::fetch_antigravity_quota_for_env(other, env).await? + } + other => { + return Err(format!("Quota fetching not implemented for agent '{other}'")); + } + }; + + self.cache_quota(agent_type, info.clone()).await; + Ok(info) + } + + /// Fetch Codex quota skeleton. + pub async fn fetch_codex_quota(&self) -> Result { + codex::fetch_codex_quota().await + } + + /// Fetch Antigravity quota skeleton. + pub async fn fetch_antigravity_quota(&self) -> Result { + antigravity::fetch_antigravity_quota().await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + use tokio::time::sleep; + + #[tokio::test] + async fn test_cache_and_get() { + let manager = QuotaManager::new(); + let now = Utc::now(); + let info = AgentQuotaInfo::new("codex", now); + + assert_eq!(manager.get_cached("codex").await, None); + + manager.cache_quota("codex", info.clone()).await; + let cached = manager.get_cached("codex").await; + assert_eq!(cached, Some(info)); + } + + #[tokio::test] + async fn test_cache_ttl_expiration() { + let ttl = Duration::from_millis(50); + let manager = QuotaManager::with_ttl(ttl); + let now = Utc::now(); + let info = AgentQuotaInfo::new("antigravity", now); + + manager.cache_quota("antigravity", info.clone()).await; + assert_eq!(manager.get_cached("antigravity").await, Some(info.clone())); + + sleep(Duration::from_millis(70)).await; + assert_eq!(manager.get_cached("antigravity").await, None); + } + + #[tokio::test] + async fn test_invalidate_and_clear() { + let manager = QuotaManager::new(); + let now = Utc::now(); + let codex_info = AgentQuotaInfo::new("codex", now); + let antigravity_info = AgentQuotaInfo::new("antigravity", now); + + manager.cache_quota("codex", codex_info.clone()).await; + manager + .cache_quota("antigravity", antigravity_info.clone()) + .await; + + manager.invalidate("codex").await; + assert_eq!(manager.get_cached("codex").await, None); + assert_eq!( + manager.get_cached("antigravity").await, + Some(antigravity_info) + ); + + manager.clear().await; + assert_eq!(manager.get_cached("antigravity").await, None); + } + + #[tokio::test] + async fn test_get_quota_caches_result() { + let manager = QuotaManager::new(); + assert_eq!(manager.get_cached("codex").await, None); + + let quota = manager.get_quota("codex").await.expect("get_quota"); + assert_eq!(quota.agent_type, "codex"); + + let cached = manager.get_cached("codex").await; + assert_eq!(cached, Some(quota)); + } + + #[tokio::test] + async fn test_get_quota_unsupported_agent() { + let manager = QuotaManager::new(); + let result = manager.get_quota("unknown_agent").await; + assert!(result.is_err()); + } + + #[test] + fn test_quota_probe_policy_across_agents() { + // 1. Antigravity in oauth-personal => auto enabled + let agy_oauth = AgentEnv::from_env_json(Some(r#"{"AGY_AUTH_METHOD":"oauth-personal"}"#)); + assert!(is_quota_probe_enabled("antigravity", &agy_oauth, None)); + + // 2. Antigravity in oauth-personal but user manually turned OFF + let agy_off = AgentEnv::from_env_json(Some( + r#"{"AGY_AUTH_METHOD":"oauth-personal","CODEG_QUOTA_PROBE_MODE":"off"}"#, + )); + assert!(!is_quota_probe_enabled("antigravity", &agy_off, None)); + + // 3. Antigravity in gemini-api-key => automatically disabled even if CODEG_QUOTA_PROBE_MODE=on + let agy_apikey = AgentEnv::from_env_json(Some( + r#"{"AGY_AUTH_METHOD":"gemini-api-key","CODEG_QUOTA_PROBE_MODE":"on"}"#, + )); + assert!(!is_quota_probe_enabled("antigravity", &agy_apikey, None)); + + // 4. Custom Antigravity instance with GEMINI_HOME => auto enabled + let custom_agy = + AgentEnv::from_env_json(Some(r#"{"GEMINI_HOME":"C:\\Users\\Admin\\.gemini-b"}"#)); + assert!(is_custom_antigravity_agent("custom:antigravity-b", &custom_agy)); + assert!(is_quota_probe_enabled("custom:antigravity-b", &custom_agy, None)); + + // 5. Claude Code with 3rd-party ANTHROPIC_BASE_URL => automatically disabled + let claude_proxy = AgentEnv::from_env_json(Some( + r#"{"ANTHROPIC_BASE_URL":"http://192.168.30.98:13000","ANTHROPIC_AUTH_TOKEN":"sk-xxx"}"#, + )); + assert!(!is_quota_probe_enabled("claude_code", &claude_proxy, None)); + + // 6. Pure API-key agents => always disabled + assert!(!is_quota_probe_enabled( + "open_code", + &AgentEnv::inherited(), + None + )); + assert!(!is_quota_probe_enabled( + "deepseek", + &AgentEnv::inherited(), + None + )); + } +} diff --git a/src-tauri/src/web/handlers/mod.rs b/src-tauri/src/web/handlers/mod.rs index b98d8706969..326a33db14b 100644 --- a/src-tauri/src/web/handlers/mod.rs +++ b/src-tauri/src/web/handlers/mod.rs @@ -40,6 +40,7 @@ pub mod pet; pub mod project_boot; pub mod question; pub mod quick_messages; +pub mod quota; pub mod science; pub mod session_info; pub mod system_settings; diff --git a/src-tauri/src/web/handlers/quota.rs b/src-tauri/src/web/handlers/quota.rs new file mode 100644 index 00000000000..7dd23aa9165 --- /dev/null +++ b/src-tauri/src/web/handlers/quota.rs @@ -0,0 +1,35 @@ +use std::sync::Arc; + +use axum::extract::{Extension, Path}; +use axum::Json; + +use crate::app_error::AppCommandError; +use crate::app_state::AppState; +use crate::commands::quota as quota_commands; +use crate::models::AgentQuotaInfo; + +pub async fn get_agent_quota( + Extension(state): Extension>, + Path(agent_type): Path, +) -> Result, AppCommandError> { + let result = quota_commands::get_agent_quota_with_db_core( + &state.quota_manager, + Some(&state.db), + &agent_type, + ) + .await?; + Ok(Json(result)) +} + +pub async fn refresh_agent_quota( + Extension(state): Extension>, + Path(agent_type): Path, +) -> Result, AppCommandError> { + let result = quota_commands::refresh_agent_quota_with_db_core( + &state.quota_manager, + Some(&state.db), + &agent_type, + ) + .await?; + Ok(Json(result)) +} diff --git a/src-tauri/src/web/mod.rs b/src-tauri/src/web/mod.rs index 54a8e1831b5..b121169dfb0 100644 --- a/src-tauri/src/web/mod.rs +++ b/src-tauri/src/web/mod.rs @@ -932,6 +932,10 @@ pub(crate) async fn do_start_web_server_tauri( .state::() .inner() .clone(), + quota_manager: app + .try_state::>() + .map(|s| s.inner().clone()) + .unwrap_or_else(|| Arc::new(crate::quota::QuotaManager::new())), }); // See do_start_web_server_with_state for rationale on the reset. diff --git a/src-tauri/src/web/router.rs b/src-tauri/src/web/router.rs index e2008192d32..0df6ea89b68 100644 --- a/src-tauri/src/web/router.rs +++ b/src-tauri/src/web/router.rs @@ -52,6 +52,11 @@ pub fn build_router_with_proxy_auth( let api = Router::new() .route("/health", post(health_check)) + // ?????? Agent Quota ?????? + .route( + "/quota/{agent_type}", + get(handlers::quota::get_agent_quota).post(handlers::quota::refresh_agent_quota), + ) // Debug endpoint: operator-facing snapshot of `EventBusMetrics` // (emit volume, lag/eviction counts, attach decision counts). // Sits behind the same auth middleware as every other route. diff --git a/src/components/chat/composer-quota-badge.test.tsx b/src/components/chat/composer-quota-badge.test.tsx new file mode 100644 index 00000000000..a08b35b6109 --- /dev/null +++ b/src/components/chat/composer-quota-badge.test.tsx @@ -0,0 +1,257 @@ +/** + * @vitest-environment jsdom + */ +import { cleanup, fireEvent, render, waitFor } from "@testing-library/react" +import { NextIntlClientProvider } from "next-intl" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" +import type { AgentQuotaInfo } from "@/lib/types" +import enMessages from "@/i18n/messages/en.json" + +const mockGetAgentQuota = vi.fn() +const mockRefreshAgentQuota = vi.fn() + +let mockConnState: { status: string; agentType: string } | undefined = undefined +let mockTabAgentType: string | null = null +const mockConnSubscribers = new Set<() => void>() + +const mockConnectionStore = { + subscribeKey: (_tabId: string, cb: () => void) => { + mockConnSubscribers.add(cb) + return () => { + mockConnSubscribers.delete(cb) + } + }, + getConnection: () => mockConnState, + getConnectPending: () => undefined, + getConnectError: () => undefined, +} + +vi.mock("@/lib/api", () => ({ + getAgentQuota: (...args: unknown[]) => mockGetAgentQuota(...args), + refreshAgentQuota: (...args: unknown[]) => mockRefreshAgentQuota(...args), +})) + +vi.mock("@/contexts/acp-connections-context", () => ({ + useConnectionStore: () => mockConnectionStore, +})) + +vi.mock("@/contexts/tab-context", () => ({ + useTabStore: ( + selector: (s: { + tabs: Array<{ id: string; agentType: string | null }> + }) => unknown + ) => + selector({ + tabs: [{ id: "tab-1", agentType: mockTabAgentType }], + }), +})) + +import { + ComposerQuotaBadge, + isQuotaSupportedAgent, + pickBindingWindow, + resetQuotaCache, +} from "./composer-quota-badge" +import { + isDraftOfficialSubscriptionMode, + isQuotaProbeCapableAgent, + quotaProbeEnabledForDraft, + setQuotaProbeMode, + type QuotaDraftSlice, +} from "@/components/settings/acp-agent-settings" + +function renderWithIntl(ui: React.ReactNode) { + return render( + + {ui} + + ) +} + +function makeSampleQuota(agentType = "codex"): AgentQuotaInfo { + return { + agentType, + planName: "Plus", + shortWindow: { + label: "5h Window", + usedPercent: 24, + remainingPercent: 76, + resetsAt: null, + resetInSeconds: 5400, + }, + weeklyWindow: { + label: "7d Window", + usedPercent: 42, + remainingPercent: 58, + resetsAt: null, + resetInSeconds: 172800, + }, + spendLimit: null, + modelQuotas: [], + lastUpdated: new Date().toISOString(), + } +} + +function baseDraft(overrides: Partial = {}): QuotaDraftSlice { + return { + envText: "", + modelProviderId: null, + apiKey: "", + apiBaseUrl: "", + claudeAuthMode: "official_subscription", + codexAuthMode: "chatgpt_subscription", + geminiAuthMode: "login_google", + grokAuthMode: "subscription", + clineProvider: "cline", + ...overrides, + } +} + +describe("ComposerQuotaBadge", () => { + beforeEach(() => { + resetQuotaCache() + mockGetAgentQuota.mockReset() + mockRefreshAgentQuota.mockReset() + mockConnSubscribers.clear() + mockConnState = undefined + mockTabAgentType = null + }) + + afterEach(() => { + cleanup() + }) + + it("recognizes codex, antigravity, and custom antigravity agents", () => { + expect(isQuotaSupportedAgent("codex")).toBe(true) + expect(isQuotaSupportedAgent("antigravity")).toBe(true) + expect(isQuotaSupportedAgent("custom:antigravity-b")).toBe(true) + expect(isQuotaSupportedAgent("custom:agy-work")).toBe(true) + expect(isQuotaSupportedAgent("open_code")).toBe(false) + expect(isQuotaSupportedAgent(null)).toBe(false) + }) + + it("does not render for unsupported agents", () => { + mockConnState = { status: "connected", agentType: "open_code" } + mockTabAgentType = "open_code" + const { container } = renderWithIntl() + expect(container.firstChild).toBeNull() + expect(mockGetAgentQuota).not.toHaveBeenCalled() + }) + + it("renders binding remaining percentage and popover windows for connected Codex session", async () => { + mockConnState = { status: "connected", agentType: "codex" } + mockTabAgentType = "codex" + mockGetAgentQuota.mockResolvedValueOnce(makeSampleQuota("codex")) + + const { getByRole, getByText } = renderWithIntl( + + ) + + await waitFor(() => { + expect(mockGetAgentQuota).toHaveBeenCalledWith("codex") + }) + + const trigger = await waitFor(() => + getByRole("button", { name: "View agent quota" }) + ) + await waitFor(() => { + // Weekly window has 58% remaining (< 76%), so it is the binding window + expect(trigger.textContent).toContain("7d: 58%") + }) + + fireEvent.click(trigger) + await waitFor(() => { + expect(getByText("5h Window")).toBeTruthy() + expect(getByText("7d Window")).toBeTruthy() + }) + }) + + it("hides badge when backend reports quota probing is disabled", async () => { + mockConnState = { status: "connected", agentType: "codex" } + mockTabAgentType = "codex" + mockGetAgentQuota.mockRejectedValueOnce( + new Error("Quota probing is disabled for codex") + ) + + const { queryByRole } = renderWithIntl() + + await waitFor(() => { + expect(mockGetAgentQuota).toHaveBeenCalledWith("codex") + }) + await waitFor(() => { + expect(queryByRole("button", { name: "View agent quota" })).toBeNull() + }) + }) + + it("refreshes quota when turn completes (prompting -> connected)", async () => { + mockConnState = { status: "prompting", agentType: "custom:antigravity-b" } + mockTabAgentType = "custom:antigravity-b" + mockGetAgentQuota.mockResolvedValue(makeSampleQuota("custom:antigravity-b")) + + const { rerender } = renderWithIntl() + + await waitFor(() => { + expect(mockGetAgentQuota).toHaveBeenCalledTimes(1) + }) + + mockConnState = { status: "connected", agentType: "custom:antigravity-b" } + mockConnSubscribers.forEach((cb) => cb()) + rerender( + + + + ) + + await waitFor(() => { + expect(mockGetAgentQuota).toHaveBeenCalledTimes(2) + }) + }) + + it("picks the binding quota window with lower remaining percentage", () => { + const sample = makeSampleQuota("codex") + const binding = pickBindingWindow(sample) + expect(binding?.label).toBe("7d Window") + expect(binding?.remainingPercent).toBe(58) + }) +}) + +describe("Smart Quota Toggle Switch Helpers", () => { + it("shows toggle only for subscription-capable agents", () => { + expect(isQuotaProbeCapableAgent({ agent_type: "codex" })).toBe(true) + expect(isQuotaProbeCapableAgent({ agent_type: "antigravity" })).toBe(true) + expect( + isQuotaProbeCapableAgent({ agent_type: "custom:antigravity-b" }) + ).toBe(true) + expect(isQuotaProbeCapableAgent({ agent_type: "claude_code" })).toBe(true) + expect(isQuotaProbeCapableAgent({ agent_type: "gemini" })).toBe(true) + expect(isQuotaProbeCapableAgent({ agent_type: "open_code" })).toBe(false) + expect(isQuotaProbeCapableAgent({ agent_type: "deepseek" })).toBe(false) + expect(isQuotaProbeCapableAgent({ agent_type: "hermes" })).toBe(false) + }) + + it("auto-enables on official subscription login and auto-disables on 3rd-party API / Model Provider", () => { + const officialDraft = baseDraft() + expect(isDraftOfficialSubscriptionMode("codex", officialDraft)).toBe(true) + expect(quotaProbeEnabledForDraft("codex", officialDraft)).toBe(true) + + const manuallyOffDraft = baseDraft({ + envText: setQuotaProbeMode("", false), + }) + expect(manuallyOffDraft.envText).toContain("CODEG_QUOTA_PROBE_MODE=off") + expect(quotaProbeEnabledForDraft("codex", manuallyOffDraft)).toBe(false) + + const apiKeyCodexDraft = baseDraft({ codexAuthMode: "api_key" }) + expect(isDraftOfficialSubscriptionMode("codex", apiKeyCodexDraft)).toBe( + false + ) + expect(quotaProbeEnabledForDraft("codex", apiKeyCodexDraft)).toBe(false) + + const providerBoundDraft = baseDraft({ modelProviderId: 42 }) + expect( + isDraftOfficialSubscriptionMode("antigravity", providerBoundDraft) + ).toBe(false) + expect(quotaProbeEnabledForDraft("antigravity", providerBoundDraft)).toBe( + false + ) + }) +}) diff --git a/src/components/chat/composer-quota-badge.tsx b/src/components/chat/composer-quota-badge.tsx new file mode 100644 index 00000000000..96d694038be --- /dev/null +++ b/src/components/chat/composer-quota-badge.tsx @@ -0,0 +1,558 @@ +"use client" + +import { + useCallback, + useEffect, + useRef, + useState, + useSyncExternalStore, +} from "react" +import { RefreshCw, Zap } from "lucide-react" +import { useTranslations } from "next-intl" +import { useConnectionStore } from "@/contexts/acp-connections-context" +import { useTabStore } from "@/contexts/tab-context" +import { AgentIcon } from "@/components/agent-icon" +import { Button } from "@/components/ui/button" +import { + Popover, + PopoverContent, + PopoverTrigger, +} from "@/components/ui/popover" +import { getAgentQuota, refreshAgentQuota } from "@/lib/api" +import { getAgentLabel } from "@/lib/custom-agents" +import type { AgentQuotaInfo, QuotaWindow } from "@/lib/types" +import { cn } from "@/lib/utils" + +const SUPPORTED_AGENTS = new Set(["codex", "antigravity"]) + +export function isQuotaSupportedAgent( + agentType: string | null | undefined +): boolean { + if (!agentType) return false + if (SUPPORTED_AGENTS.has(agentType)) return true + if (agentType.startsWith("custom:")) { + const customId = agentType.slice("custom:".length).toLowerCase() + return customId.includes("antigravity") || customId.startsWith("agy") + } + return false +} + +const quotaCache = new Map< + string, + { info: AgentQuotaInfo | null; disabled: boolean; fetchedAt: number } +>() +const quotaListeners = new Map void>>() + +export function resetQuotaCache() { + quotaCache.clear() + quotaListeners.clear() +} + +function notifyQuotaChanged(agentType: string) { + quotaListeners.get(agentType)?.forEach((cb) => cb()) +} + +function subscribeQuota(agentType: string, callback: () => void) { + if (!quotaListeners.has(agentType)) { + quotaListeners.set(agentType, new Set()) + } + quotaListeners.get(agentType)!.add(callback) + return () => { + quotaListeners.get(agentType)?.delete(callback) + } +} + +type QuotaTranslator = ReturnType< + typeof useTranslations<"Folder.statusBar.quota"> +> + +function formatCountdown( + resetInSeconds: number | null | undefined, + resetsAt: string | null | undefined, + t: QuotaTranslator +): string | null { + let seconds = resetInSeconds + if (seconds == null && resetsAt) { + const diff = Math.floor((new Date(resetsAt).getTime() - Date.now()) / 1000) + seconds = Math.max(0, diff) + } + if (seconds == null) return null + + if (seconds <= 0) { + return t("resetsSoon") + } + + const daysUnit = t("days") + const hoursUnit = t("hours") + const minutesUnit = t("minutes") + const secondsUnit = t("seconds") + + if (seconds < 60) { + return t("resetsIn", { time: `${seconds}${secondsUnit}` }) + } + + if (seconds < 3600) { + const mins = Math.floor(seconds / 60) + return t("resetsIn", { time: `${mins}${minutesUnit}` }) + } + + if (seconds < 86400) { + const hours = Math.floor(seconds / 3600) + const mins = Math.floor((seconds % 3600) / 60) + const timeStr = + mins > 0 + ? `${hours}${hoursUnit} ${mins}${minutesUnit}` + : `${hours}${hoursUnit}` + return t("resetsIn", { time: timeStr }) + } + + const days = Math.floor(seconds / 86400) + const hours = Math.floor((seconds % 86400) / 3600) + const timeStr = + hours > 0 ? `${days}${daysUnit} ${hours}${hoursUnit}` : `${days}${daysUnit}` + return t("resetsIn", { time: timeStr }) +} + +/** The window with the least remaining — the one that will run out first. */ +export function pickBindingWindow( + quota: AgentQuotaInfo | null | undefined +): QuotaWindow | null { + const short = quota?.shortWindow ?? null + const weekly = quota?.weeklyWindow ?? null + if (short && weekly) { + return weekly.remainingPercent < short.remainingPercent ? weekly : short + } + return short ?? weekly +} + +function getShortWindowLabel(rawLabel: string): string { + const lower = rawLabel.toLowerCase() + if ( + lower.includes("5-hour") || + lower.includes("5h") || + lower.includes("5小时") + ) + return "5h" + if (lower.includes("week") || lower.includes("7d") || lower.includes("周")) + return "7d" + if (lower.includes("daily") || lower.includes("day") || lower.includes("日")) + return "1d" + if (rawLabel.length > 6) return rawLabel.slice(0, 5) + return rawLabel +} + +/** + * Quota badge and popover shown in the composer status row. + * Displays rate limit / usage window for supported agents (Codex, Antigravity, + * and custom Antigravity instances) when official subscription mode is active. + */ +export function ComposerQuotaBadge({ tabId }: { tabId: string | null }) { + const t = useTranslations("Folder.statusBar.quota") + const store = useConnectionStore() + + const tabAgentType = useTabStore((s) => { + const tab = s.tabs.find((x) => x.id === tabId) + return tab?.agentType ?? null + }) + + const subscribeConn = useCallback( + (cb: () => void) => { + if (!tabId) return () => {} + return store.subscribeKey(tabId, cb) + }, + [store, tabId] + ) + + const getConnSnapshot = useCallback( + () => (tabId ? store.getConnection(tabId) : undefined), + [store, tabId] + ) + const conn = useSyncExternalStore( + subscribeConn, + getConnSnapshot, + getConnSnapshot + ) + + const getPendingSnapshot = useCallback( + () => (tabId ? store.getConnectPending(tabId) : undefined), + [store, tabId] + ) + const connectPending = useSyncExternalStore( + subscribeConn, + getPendingSnapshot, + getPendingSnapshot + ) + + const getConnectErrorSnapshot = useCallback( + () => (tabId ? store.getConnectError(tabId) : undefined), + [store, tabId] + ) + const connectError = useSyncExternalStore( + subscribeConn, + getConnectErrorSnapshot, + getConnectErrorSnapshot + ) + + const failedConnect = !conn && !connectPending ? connectError : undefined + + const agentType = + conn?.agentType ?? + connectPending?.agentType ?? + failedConnect?.agentType ?? + tabAgentType ?? + null + + const isSupported = isQuotaSupportedAgent(agentType) + + const subscribeAgentQuota = useCallback( + (cb: () => void) => { + if (!agentType) return () => {} + return subscribeQuota(agentType, cb) + }, + [agentType] + ) + + const getQuotaSnapshot = useCallback(() => { + return agentType ? (quotaCache.get(agentType)?.info ?? null) : null + }, [agentType]) + + const getDisabledSnapshot = useCallback(() => { + return agentType ? (quotaCache.get(agentType)?.disabled ?? false) : false + }, [agentType]) + + const quota = useSyncExternalStore( + subscribeAgentQuota, + getQuotaSnapshot, + getQuotaSnapshot + ) + + const quotaDisabled = useSyncExternalStore( + subscribeAgentQuota, + getDisabledSnapshot, + getDisabledSnapshot + ) + + const [loading, setLoading] = useState(false) + const [refreshing, setRefreshing] = useState(false) + const [, setTick] = useState(0) + + useEffect(() => { + const timer = setInterval(() => { + setTick((tick) => (tick + 1) % 100000) + }, 10_000) + return () => clearInterval(timer) + }, []) + + // Fetch once on mount / agent change if cache is stale (3 minutes TTL, zero background polling) + useEffect(() => { + if (!agentType || !isQuotaSupportedAgent(agentType)) return + + let mounted = true + const cached = quotaCache.get(agentType) + const isStale = !cached || Date.now() - cached.fetchedAt > 180_000 + + if (isStale) { + setLoading(true) + getAgentQuota(agentType) + .then((info) => { + if (!mounted) return + quotaCache.set(agentType, { + info, + disabled: false, + fetchedAt: Date.now(), + }) + notifyQuotaChanged(agentType) + }) + .catch((err) => { + if (!mounted) return + const msg = String(err ?? "") + if (msg.includes("Quota probing is disabled")) { + quotaCache.set(agentType, { + info: null, + disabled: true, + fetchedAt: Date.now(), + }) + notifyQuotaChanged(agentType) + } + }) + .finally(() => { + if (mounted) setLoading(false) + }) + } + + return () => { + mounted = false + } + }, [agentType]) + + // Event-driven refresh when a prompt turn finishes (status transitions from prompting -> connected) + const prevStatusRef = useRef(conn?.status) + useEffect(() => { + const prevStatus = prevStatusRef.current + const nextStatus = conn?.status + prevStatusRef.current = nextStatus + + if ( + prevStatus === "prompting" && + nextStatus === "connected" && + agentType && + isQuotaSupportedAgent(agentType) && + !quotaDisabled + ) { + getAgentQuota(agentType) + .then((info) => { + quotaCache.set(agentType, { + info, + disabled: false, + fetchedAt: Date.now(), + }) + notifyQuotaChanged(agentType) + }) + .catch(() => {}) + } + }, [conn?.status, agentType, quotaDisabled]) + + if (!isSupported || !agentType || quotaDisabled) { + return null + } + + const handleRefresh = async (e: React.MouseEvent) => { + e.stopPropagation() + if (!agentType || refreshing) return + setRefreshing(true) + try { + const info = await refreshAgentQuota(agentType) + quotaCache.set(agentType, { + info, + disabled: false, + fetchedAt: Date.now(), + }) + notifyQuotaChanged(agentType) + } catch (err) { + console.error("Failed to refresh agent quota:", err) + } finally { + setRefreshing(false) + } + } + + const agentLabel = getAgentLabel(agentType) + + // Show whichever window is actually binding — the one with less left. + const bindingWindow = pickBindingWindow(quota) + + let shortText = "" + if (bindingWindow) { + shortText = `${getShortWindowLabel(bindingWindow.label)}: ${Math.round(bindingWindow.remainingPercent)}%` + } else if (quota?.planName) { + shortText = quota.planName + } else { + shortText = t("title") + } + + const isLowQuota = + (quota?.shortWindow != null && quota.shortWindow.remainingPercent <= 20) || + (quota?.weeklyWindow != null && quota.weeklyWindow.remainingPercent <= 20) + const isCriticalQuota = + (quota?.shortWindow != null && quota.shortWindow.remainingPercent <= 5) || + (quota?.weeklyWindow != null && quota.weeklyWindow.remainingPercent <= 5) + + const triggerTitle = quota + ? `${agentLabel} ${t("title")}${quota.planName ? ` (${quota.planName})` : ""}: ${shortText}` + : `${agentLabel} ${t("title")}` + + return ( + + + + + + {/* Header: Agent Icon + Agent Label + Plan Badge + Refresh Button */} +
+
+ + {agentLabel} + {quota?.planName ? ( + + {quota.planName} + + ) : null} +
+ +
+ + {/* Short Window Section */} + {quota?.shortWindow ? ( +
+
+ + {quota.shortWindow.label || t("shortWindow")} + + + {t("remaining")}{" "} + {Math.round(quota.shortWindow.remainingPercent)}% + +
+
+
+
+ {formatCountdown( + quota.shortWindow.resetInSeconds, + quota.shortWindow.resetsAt, + t + ) ? ( +
+ {formatCountdown( + quota.shortWindow.resetInSeconds, + quota.shortWindow.resetsAt, + t + )} +
+ ) : null} +
+ ) : null} + + {/* Weekly Window Section */} + {quota?.weeklyWindow ? ( +
+
+ + {quota.weeklyWindow.label || t("weeklyWindow")} + + + {t("remaining")}{" "} + {Math.round(quota.weeklyWindow.remainingPercent)}% + +
+
+
+
+ {formatCountdown( + quota.weeklyWindow.resetInSeconds, + quota.weeklyWindow.resetsAt, + t + ) ? ( +
+ {formatCountdown( + quota.weeklyWindow.resetInSeconds, + quota.weeklyWindow.resetsAt, + t + )} +
+ ) : null} +
+ ) : null} + + {/* Spend Limit Section */} + {quota?.spendLimit ? ( +
+ {t("spendLimit")} + + ${quota.spendLimit.usedUsd.toFixed(2)} / $ + {quota.spendLimit.limitUsd.toFixed(2)} + +
+ ) : null} + + {/* Empty state if quota is null */} + {!quota && !loading ? ( +
+ {t("noQuota")} +
+ ) : null} + + {/* Footer: Last Updated */} + {quota?.lastUpdated ? ( +
+ {t("lastUpdated", { + time: new Date(quota.lastUpdated).toLocaleTimeString([], { + hour: "2-digit", + minute: "2-digit", + second: "2-digit", + }), + })} +
+ ) : null} + + + ) +} diff --git a/src/components/chat/message-input.test.tsx b/src/components/chat/message-input.test.tsx index 3cef470850b..af39c20d289 100644 --- a/src/components/chat/message-input.test.tsx +++ b/src/components/chat/message-input.test.tsx @@ -102,6 +102,9 @@ vi.mock("./composer-context-usage", () => ({ vi.mock("./composer-connection-status", () => ({ ComposerConnectionStatus: () => null, })) +vi.mock("./composer-quota-badge", () => ({ + ComposerQuotaBadge: () => null, +})) // The platform opener is the DESKTOP arm of the shared opener; this suite runs // in web mode, where a system-browser target lands on `window.open` instead. const platform = vi.hoisted(() => ({ openUrl: vi.fn(async () => {}) })) diff --git a/src/components/chat/message-input.tsx b/src/components/chat/message-input.tsx index 519b3ab5939..060f3be2376 100644 --- a/src/components/chat/message-input.tsx +++ b/src/components/chat/message-input.tsx @@ -74,6 +74,7 @@ import { type ConversationFolderPickerOverride, } from "@/components/chat/conversation-context-bar" import { ComposerContextUsage } from "@/components/chat/composer-context-usage" +import { ComposerQuotaBadge } from "@/components/chat/composer-quota-badge" import { ComposerConnectionStatus } from "@/components/chat/composer-connection-status" import { SessionSettingsPicker } from "@/components/chat/session-settings-picker" import { @@ -2325,6 +2326,7 @@ export function MessageInput({ send button's right edge in the action bar above — no centring slot, which would inset the narrow icon and break the alignment. */}
+
diff --git a/src/components/settings/acp-agent-settings.tsx b/src/components/settings/acp-agent-settings.tsx index 177895c7177..28c3295b1a8 100644 --- a/src/components/settings/acp-agent-settings.tsx +++ b/src/components/settings/acp-agent-settings.tsx @@ -692,6 +692,168 @@ export function setHostToolsAgentMode( }) } +const QUOTA_PROBE_MODE_ENV = "CODEG_QUOTA_PROBE_MODE" +const QUOTA_PROBE_MODE_AUTO = "auto" +const QUOTA_PROBE_MODE_OFF = "off" + +const QUOTA_PROBE_CAPABLE_AGENT_TYPES = new Set([ + "codex", + "antigravity", + "claude_code", + "gemini", + "cursor", + "kimi_code", + "minimax_code", + "grok", + "cline", + "kiro", + "qoder", + "trae", + "command_code", + "zcode", +]) + +/** + * Whether this agent supports official account / subscription login and can + * therefore display the smart quota monitoring switch. Pure BYO API key or + * multi-provider router agents (`open_code`, `open_claw`, `hermes`, + * `code_buddy`, `pi`, `deepseek`) do not show the switch. + */ +export function isQuotaProbeCapableAgent( + agent: Pick, + envText = "" +): boolean { + const agentType = agent.agent_type + if (QUOTA_PROBE_CAPABLE_AGENT_TYPES.has(agentType)) return true + if ( + agentType.startsWith("custom:antigravity") || + agentType.startsWith("custom:agy") + ) { + return true + } + if (agentType.startsWith("custom:")) { + const env = parseEnvText(envText) + return Boolean( + env.GEMINI_HOME?.trim() || + env.AGY_AUTH_METHOD?.trim() || + env.AGY_GEMINI_API_KEY?.trim() + ) + } + return false +} + +export type QuotaDraftSlice = Pick< + AgentDraft, + | "envText" + | "modelProviderId" + | "apiKey" + | "apiBaseUrl" + | "claudeAuthMode" + | "codexAuthMode" + | "geminiAuthMode" + | "grokAuthMode" + | "clineProvider" +> + +/** + * Whether the current settings draft is configured for official account / + * subscription login (as opposed to a 3rd-party API key, custom base URL, or + * bound Model Provider). + */ +export function isDraftOfficialSubscriptionMode( + agentType: AgentType, + draft: QuotaDraftSlice +): boolean { + if (draft.modelProviderId !== null) return false + const env = parseEnvText(draft.envText) + + if (agentType === "codex") { + return draft.codexAuthMode === "chatgpt_subscription" + } + if ( + agentType === "antigravity" || + agentType.startsWith("custom:antigravity") || + agentType.startsWith("custom:agy") + ) { + const method = (env.AGY_AUTH_METHOD ?? "").trim() + const apiKey = (env.AGY_GEMINI_API_KEY ?? "").trim() + if (method === "gemini-api-key" || method === "agent-platform") { + return false + } + if (!method && apiKey) return false + return true + } + if (agentType === "claude_code") { + return draft.claudeAuthMode === "official_subscription" + } + if (agentType === "gemini") { + return draft.geminiAuthMode === "login_google" + } + if (agentType === "grok") { + return draft.grokAuthMode === "subscription" + } + if (agentType === "cline") { + return isClineSignInProvider(draft.clineProvider) + } + if (agentType === "cursor") { + const cursorMode = (env.CURSOR_AUTH_MODE ?? "").trim() + if (cursorMode === "api_key") return false + if (!cursorMode && (env.CURSOR_API_KEY ?? "").trim()) return false + return true + } + if ( + agentType === "kimi_code" || + agentType === "minimax_code" || + agentType === "kiro" || + agentType === "qoder" || + agentType === "trae" || + agentType === "command_code" || + agentType === "zcode" + ) { + return !draft.apiKey.trim() && !draft.apiBaseUrl.trim() + } + if (agentType.startsWith("custom:")) { + const method = (env.AGY_AUTH_METHOD ?? "").trim() + if (method === "gemini-api-key" || method === "agent-platform") { + return false + } + if ((env.AGY_GEMINI_API_KEY ?? "").trim()) return false + return Boolean(env.GEMINI_HOME?.trim() || method) + } + return false +} + +/** + * Resolve whether quota monitoring is enabled for the current draft: + * - In 3rd-party API Key / Custom Base URL / Model Provider mode, returns false + * automatically unless explicitly forced to `on` in `CODEG_QUOTA_PROBE_MODE`. + * - In Official Account Subscription mode, defaults to true unless explicitly + * switched `off` by the user. + */ +export function quotaProbeEnabledForDraft( + agentType: AgentType, + draft: QuotaDraftSlice +): boolean { + const mode = (parseEnvText(draft.envText)[QUOTA_PROBE_MODE_ENV] ?? "") + .trim() + .toLowerCase() + if (mode === QUOTA_PROBE_MODE_OFF || mode === "false" || mode === "0") { + return false + } + if (mode === "on" || mode === "true" || mode === "1") { + return true + } + return isDraftOfficialSubscriptionMode(agentType, draft) +} + +export function setQuotaProbeMode(envText: string, enabled: boolean): string { + return patchEnvText(envText, { + [QUOTA_PROBE_MODE_ENV]: enabled + ? QUOTA_PROBE_MODE_AUTO + : QUOTA_PROBE_MODE_OFF, + }) +} + /** * How Claude Code shows its thinking. Claude Code asks the API for * `display: "updates"` by default, so a thinking block carries only the notes @@ -8925,6 +9087,47 @@ export function AcpAgentSettings() { aria-label={t("hostTools.label")} />
+ {isQuotaProbeCapableAgent( + selectedAgent, + selectedDraft.envText + ) && + (() => { + const officialSub = isDraftOfficialSubscriptionMode( + selectedAgent.agent_type, + selectedDraft + ) + return ( +
+
+ +

+ {officialSub + ? t("quotaProbe.description") + : t("quotaProbe.autoDisabledHint")} +

+
+ { + updateSelectedDraft((current) => ({ + ...current, + envText: setQuotaProbeMode( + current.envText, + checked + ), + })) + }} + disabled={!officialSub || selectedGrokSaving} + aria-label={t("quotaProbe.label")} + /> +
+ ) + })()} {selectedAgent.agent_type === "claude_code" && (
diff --git a/src/i18n/messages/ar.json b/src/i18n/messages/ar.json index cbfc27b92f7..515d7fccd2a 100644 --- a/src/i18n/messages/ar.json +++ b/src/i18n/messages/ar.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "يشغّل Codeg الحزمة {adapterPackage} كجسر ACP. ثبّت ZCode الرسمي وسجّل الدخول أولاً؛ يستخدم بيئته وإعداداته الحالية.", "runtimeHint": "ثبّت ZCode الرسمي وسجّل الدخول أولاً. يثبّت Codeg جسر ACP فقط. إذا فشل الاكتشاف، حدّد مسار CLI عبر ZCODE_BIN ومسار Node.js عبر ZCODE_NODE. تُدار النماذج وبيانات الاعتماد في ZCode." + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "تعذّرت قراءة الحالة: {message}", "openSettings": "فتح الإعدادات الكاملة", "refresh": "تحديث" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { @@ -3825,9 +3848,9 @@ "resumeReasonLabel": "سبب الاستئناف", "detailTitle": "محادثة الوكيل الفرعي", "detailDescription": "عرض للقراءة فقط لمحادثة الوكيل الفرعي المُفوَّض.", - "waitForResult": "في انتظار نتيجة المهمة \u2066{task}\u2069", + "waitForResult": "في انتظار نتيجة المهمة ⁦{task}⁩", "waitForResultNoTask": "في انتظار نتيجة المهمة", - "cancelTask": "جارٍ إلغاء المهمة \u2066{task}\u2069", + "cancelTask": "جارٍ إلغاء المهمة ⁦{task}⁩", "cancelTaskNoTask": "جارٍ إلغاء المهمة", "resultPageOf": "{current} / {total}", "prevResult": "النتيجة السابقة", @@ -3986,7 +4009,7 @@ "monitor": "مراقبة", "networkAccess": "الوصول إلى الشبكة", "codexScript": "نص برمجي", - "background": "في الخلفية · \u2066#{id}\u2069", + "background": "في الخلفية · ⁦#{id}⁩", "backgroundShort": "في الخلفية", "sessionWait": "انتظار", "sessionTerminate": "إيقاف", @@ -4281,9 +4304,9 @@ "tokens": "{count} رمز", "stopDeclined": "تعذّر إيقاف هذه المهمة، ربما تكون قد انتهت بالفعل.", "stopFailed": "فشل إيقاف المهمة: {error}", - "agentWorking": "\u2066{agent}\u2069 يعمل", + "agentWorking": "⁦{agent}⁩ يعمل", "scheduledWakeup": "تنبيه مجدول", - "scheduledWakeupAt": "تنبيه مجدول في \u2066{time}\u2069", + "scheduledWakeupAt": "تنبيه مجدول في ⁦{time}⁩", "scheduledTask": "مهمة مجدولة" }, "remoteControl": { @@ -6625,8 +6648,8 @@ "fetch": "جلب النماذج", "needKey": "أدخل مفتاح API أولًا.", "needBaseUrl": "أدخل عنوان API بصيغة http(s) أولًا.", - "summaryModels": "{count, plural, zero {لا نماذج} one {نموذج واحد} two {نموذجان} few {# نماذج} many {# نموذجًا} other {# نموذج}} · \u2066{host}\u2069", - "summaryPicked": "تم اختيار {picked} من {total} · \u2066{host}\u2069", + "summaryModels": "{count, plural, zero {لا نماذج} one {نموذج واحد} two {نموذجان} few {# نماذج} many {# نموذجًا} other {# نموذج}} · ⁦{host}⁩", + "summaryPicked": "تم اختيار {picked} من {total} · ⁦{host}⁩", "summaryRemoved": "أُزيل {count}", "fromPublished": "القائمة المطبّقة", "fromCache": "قائمة جُلبت سابقًا", @@ -6707,7 +6730,7 @@ "ProviderModels": { "fetch": "جلب النماذج", "fetching": "جارٍ جلب النماذج…", - "fromHost": "{count, plural, zero {لا نماذج} one {نموذج واحد} two {نموذجان} few {# نماذج} many {# نموذجًا} other {# نموذج}} · \u2066{host}\u2069", + "fromHost": "{count, plural, zero {لا نماذج} one {نموذج واحد} two {نموذجان} few {# نماذج} many {# نموذجًا} other {# نموذج}} · ⁦{host}⁩", "retry": "إعادة المحاولة", "needKey": "أدخل مفتاح API أولًا: تُجلب النماذج به.", "needBaseUrl": "أدخل عنوان API أولًا.", diff --git a/src/i18n/messages/de.json b/src/i18n/messages/de.json index e6c5b50fe62..6b6f685147c 100644 --- a/src/i18n/messages/de.json +++ b/src/i18n/messages/de.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg startet {adapterPackage} als ACP-Brücke. Installiere zuerst ZCode und melde dich an; die vorhandene Laufzeit und Konfiguration werden verwendet.", "runtimeHint": "Installiere zuerst ZCode und melde dich an. Codeg installiert nur die ACP-Brücke. Bei Erkennungsproblemen gib mit ZCODE_BIN den CLI-Einstieg und mit ZCODE_NODE Node.js an. Modelle und Zugangsdaten werden in ZCode verwaltet." + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "Status konnte nicht gelesen werden: {message}", "openSettings": "Alle Einstellungen öffnen", "refresh": "Aktualisieren" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 9d9dd7d1169..96f4d381466 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg runs {adapterPackage} as an ACP bridge. Install and sign in to official ZCode first; the bridge uses its runtime and existing configuration.", "runtimeHint": "Install and sign in to official ZCode first. Codeg installs only the ACP bridge. Use ZCODE_BIN for the CLI entry and ZCODE_NODE for Node.js if detection fails. Manage models and credentials in ZCode." + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "Could not read the status: {message}", "openSettings": "Open full settings", "refresh": "Refresh" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/es.json b/src/i18n/messages/es.json index 16b26cf6b29..eced44b138a 100644 --- a/src/i18n/messages/es.json +++ b/src/i18n/messages/es.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg ejecuta {adapterPackage} como puente ACP. Instala ZCode e inicia sesión primero; utiliza su entorno y configuración existentes.", "runtimeHint": "Instala ZCode e inicia sesión primero. Codeg solo instala el puente ACP. Si falla la detección, indica la entrada CLI con ZCODE_BIN y Node.js con ZCODE_NODE. Gestiona modelos y credenciales en ZCode." + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "No se pudo leer el estado: {message}", "openSettings": "Abrir ajustes completos", "refresh": "Actualizar" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/fr.json b/src/i18n/messages/fr.json index 07c58bd4432..b28db00ee99 100644 --- a/src/i18n/messages/fr.json +++ b/src/i18n/messages/fr.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg exécute {adapterPackage} comme passerelle ACP. Installez ZCode et connectez-vous d’abord ; son environnement et sa configuration existants sont utilisés.", "runtimeHint": "Installez ZCode et connectez-vous d’abord. Codeg installe uniquement la passerelle ACP. Si la détection échoue, indiquez l’entrée CLI avec ZCODE_BIN et Node.js avec ZCODE_NODE. Gérez modèles et identifiants dans ZCode." + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "Impossible de lire l'état : {message}", "openSettings": "Ouvrir les paramètres", "refresh": "Actualiser" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/ja.json b/src/i18n/messages/ja.json index 6fc728084d7..8561357ab26 100644 --- a/src/i18n/messages/ja.json +++ b/src/i18n/messages/ja.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg は {adapterPackage} を ACP ブリッジとして起動します。先に公式 ZCode をインストールしてログインしてください。既存の実行環境と設定を使用します。", "runtimeHint": "先に公式 ZCode をインストールしてログインしてください。Codeg は ACP ブリッジのみをインストールします。自動検出に失敗した場合は ZCODE_BIN で CLI、ZCODE_NODE で Node.js を指定できます。モデルと認証情報は ZCode で管理します。" + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "状態を取得できませんでした: {message}", "openSettings": "詳細設定を開く", "refresh": "更新" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/ko.json b/src/i18n/messages/ko.json index 720b2bdc646..bd901026a24 100644 --- a/src/i18n/messages/ko.json +++ b/src/i18n/messages/ko.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg는 {adapterPackage}를 ACP 브리지로 실행합니다. 먼저 공식 ZCode를 설치하고 로그인하세요. 기존 런타임과 설정을 사용합니다.", "runtimeHint": "먼저 공식 ZCode를 설치하고 로그인하세요. Codeg는 ACP 브리지만 설치합니다. 자동 감지가 실패하면 ZCODE_BIN으로 CLI 경로를, ZCODE_NODE로 Node.js를 지정하세요. 모델과 인증 정보는 ZCode에서 관리합니다." + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "상태를 읽을 수 없습니다: {message}", "openSettings": "전체 설정 열기", "refresh": "새로 고침" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/pt.json b/src/i18n/messages/pt.json index 4020ec2ed28..91d1a6e7fcb 100644 --- a/src/i18n/messages/pt.json +++ b/src/i18n/messages/pt.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "O Codeg executa {adapterPackage} como ponte ACP. Instale o ZCode e entre na sua conta primeiro; são usados o ambiente e a configuração existentes.", "runtimeHint": "Instale o ZCode e entre na sua conta primeiro. O Codeg instala apenas a ponte ACP. Se a detecção falhar, indique a entrada CLI com ZCODE_BIN e o Node.js com ZCODE_NODE. Gerencie modelos e credenciais no ZCode." + }, + "quotaProbe": { + "label": "Official account quota monitoring", + "description": "Display subscription quota in the composer bar via 100% zero-network local session logs and native ACP events (no spoofed HTTP requests, zero risk of account restriction). Automatically enabled for official account login.", + "autoDisabledHint": "Automatically disabled because this agent is currently configured with a 3rd-party API key, custom endpoint, or Model Provider." } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "Não foi possível ler o estado: {message}", "openSettings": "Abrir configurações", "refresh": "Atualizar" + }, + "quota": { + "title": "Quota", + "triggerAria": "View agent quota", + "refresh": "Refresh quota", + "refreshing": "Refreshing…", + "shortWindow": "5-Hour Window", + "weeklyWindow": "Weekly Window", + "remaining": "Remaining", + "resetsSoon": "Resets soon", + "resetsIn": "Resets in {time}", + "days": "d", + "hours": "h", + "minutes": "m", + "seconds": "s", + "spendLimit": "Spend Limit", + "noQuota": "Quota unavailable", + "lastUpdated": "Updated {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/zh-CN.json b/src/i18n/messages/zh-CN.json index aeb32c44119..521be3bb92a 100644 --- a/src/i18n/messages/zh-CN.json +++ b/src/i18n/messages/zh-CN.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg 通过 {adapterPackage} 桥接 ACP。请先安装并登录官方 ZCode；桥接器使用其运行时和现有配置。", "runtimeHint": "请先安装并登录官方 ZCode，Codeg 仅安装 ACP 桥接器。自动检测失败时，可用 ZCODE_BIN 指定 CLI 入口、ZCODE_NODE 指定 Node.js。模型和凭据在 ZCode 中管理。" + }, + "quotaProbe": { + "label": "官方账号配额监控", + "description": "在对话输入框底部展示订阅配额（100% 零网络读取本地 CLI 会话日志或原生 ACP 事件流，彻底杜绝伪造请求与账号风控风险）。检测到官方账号登录模式时默认开启。", + "autoDisabledHint": "当前处于第三方 API Key、自定义端点或模型提供商（Model Provider）模式，已自动关闭配额监控以避免无效请求。" } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "无法读取状态：{message}", "openSettings": "打开详细设置", "refresh": "刷新" + }, + "quota": { + "title": "额度", + "triggerAria": "查看智能体额度", + "refresh": "刷新额度", + "refreshing": "刷新中…", + "shortWindow": "5小时窗口", + "weeklyWindow": "周窗口", + "remaining": "剩余", + "resetsSoon": "即将重置", + "resetsIn": "{time}后重置", + "days": "天", + "hours": "小时", + "minutes": "分", + "seconds": "秒", + "spendLimit": "消费限额", + "noQuota": "暂无额度数据", + "lastUpdated": "更新于 {time}" } }, "auxPanel": { diff --git a/src/i18n/messages/zh-TW.json b/src/i18n/messages/zh-TW.json index 034d3eeb509..4bba0f8b88e 100644 --- a/src/i18n/messages/zh-TW.json +++ b/src/i18n/messages/zh-TW.json @@ -1893,6 +1893,11 @@ "zcode": { "adapterInfo": "Codeg 透過 {adapterPackage} 橋接 ACP。請先安裝並登入官方 ZCode；橋接器使用其執行環境與現有設定。", "runtimeHint": "請先安裝並登入官方 ZCode，Codeg 僅安裝 ACP 橋接器。自動偵測失敗時，可用 ZCODE_BIN 指定 CLI 入口、ZCODE_NODE 指定 Node.js。模型與憑證在 ZCode 中管理。" + }, + "quotaProbe": { + "label": "官方帳號配額監控", + "description": "在對話輸入框底部展示訂閱配額（100% 零網路讀取本機 CLI 工作階段日誌或原生 ACP 事件流，徹底杜絕偽造請求與帳號風控風險）。偵測到官方帳號登入模式時預設開啟。", + "autoDisabledHint": "目前處於第三方 API Key、自訂端點或模型提供商（Model Provider）模式，已自動關閉配額監控以避免無效請求。" } }, "AgentVocabulary": { @@ -2526,6 +2531,24 @@ "loadFailed": "無法讀取狀態：{message}", "openSettings": "開啟詳細設定", "refresh": "重新整理" + }, + "quota": { + "title": "額度", + "triggerAria": "檢視智慧體額度", + "refresh": "重新整理額度", + "refreshing": "重新整理中…", + "shortWindow": "5小時視窗", + "weeklyWindow": "週視窗", + "remaining": "剩餘", + "resetsSoon": "即將重設", + "resetsIn": "{time}後重設", + "days": "天", + "hours": "小時", + "minutes": "分", + "seconds": "秒", + "spendLimit": "消費限額", + "noQuota": "暫無額度資料", + "lastUpdated": "更新於 {time}" } }, "auxPanel": { diff --git a/src/lib/api.ts b/src/lib/api.ts index c7250ff793d..95a97cb89ac 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -22,6 +22,7 @@ import { TurnBusyError, isTurnInProgressRejection } from "./turn-busy" import type { FolderThemeColor } from "./theme-presets" import type { FollowUpIntent } from "./task-follow-up" import type { + AgentQuotaInfo, LeakedTempReclaim, LeakedTempScan, AgentType, @@ -6648,6 +6649,28 @@ export async function forgeSettingsSet( return getTransport().call("forge_settings_set", { folderId, settings }) } +// ?????? Agent Quota ?????????????????????????????????????????????????????????????????????????????????????????????????????????????????????????? + +export async function getAgentQuota( + agentType: string +): Promise { + return getTransport().call( + "get_agent_quota", + { agentType }, + { webPath: `/api/quota/${agentType}`, webMethod: "GET" } + ) +} + +export async function refreshAgentQuota( + agentType: string +): Promise { + return getTransport().call( + "refresh_agent_quota", + { agentType }, + { webPath: `/api/quota/${agentType}`, webMethod: "POST" } + ) +} + /** Where auto-processing stands for one folder: the last scan's report, the * next scheduled scan, the last webhook delivery. */ export async function forgeAutoStatus( diff --git a/src/lib/transport/types.ts b/src/lib/transport/types.ts index 134d9e07820..736e514f8a9 100644 --- a/src/lib/transport/types.ts +++ b/src/lib/transport/types.ts @@ -32,7 +32,7 @@ export type AttachDetachReason = * caller invokes `detach()` or the server emits `onDetached`. * * `eventSeq` / `highWaterSeq` are the high-water mark after the initial - * frame; subsequent `onEvent` envelopes have `envelope.seq > highWaterSeq`. + * frame; subsequent `onEvent` envelopes have `envelope.seq > highWaterSeq`Platform. */ export interface AttachHandlers { onSnapshot(snapshot: LiveSessionSnapshot, eventSeq: number): void @@ -92,6 +92,15 @@ export interface CallOptions { * backend can return its own structured error. */ timeoutMs?: number + /** + * Optional custom URL path for WebTransport (e.g. `/api/quota/${agentType}`). + * If not provided, WebTransport defaults to `/api/${command}`. + */ + webPath?: string + /** + * Optional HTTP method for WebTransport (defaults to "POST"). + */ + webMethod?: "GET" | "POST" | "PUT" | "DELETE" } export interface Transport { diff --git a/src/lib/transport/web-transport.test.ts b/src/lib/transport/web-transport.test.ts index 47e25e47a80..398add63ee2 100644 --- a/src/lib/transport/web-transport.test.ts +++ b/src/lib/transport/web-transport.test.ts @@ -87,6 +87,64 @@ function connectReady() { } const ok200 = () => ({ status: 200, ok: true, json: async () => ({}) }) + +describe("WebTransport quota endpoints", () => { + it.each([true, false])( + "GET quota endpoint omits body and Content-Type (hasToken=%s)", + async (hasToken) => { + if (!hasToken) localStorage.removeItem("codeg_token") + fetchMock.mockResolvedValue({ ...ok200(), text: async () => "{}" }) + const t = new WebTransport("http://localhost") + + await t.call( + "get_agent_quota", + {}, + { + webPath: "/api/quota/codex", + webMethod: "GET", + } + ) + + const [url, init] = fetchMock.mock.calls[0] + expect(url).toBe("http://localhost/api/quota/codex") + expect(init.method).toBe("GET") + expect(init.body).toBeUndefined() + expect(init.headers).not.toHaveProperty("Content-Type") + expect(init.headers["X-Requested-With"]).toBe("XMLHttpRequest") + if (hasToken) expect(init.headers.Authorization).toBe("Bearer tok") + else expect(init.headers).not.toHaveProperty("Authorization") + t.destroy() + } + ) + + it("POST refresh quota endpoint sends empty JSON body", async () => { + fetchMock.mockResolvedValue({ ...ok200(), text: async () => "{}" }) + const t = new WebTransport("http://localhost") + + await t.call( + "refresh_agent_quota", + {}, + { + webPath: "/api/quota/codex", + webMethod: "POST", + } + ) + + expect(fetchMock).toHaveBeenCalledWith( + "http://localhost/api/quota/codex", + expect.objectContaining({ + method: "POST", + body: "{}", + headers: { + "Content-Type": "application/json", + Authorization: "Bearer tok", + "X-Requested-With": "XMLHttpRequest", + }, + }) + ) + t.destroy() + }) +}) // Codeg's own 401 carries its bearer challenge (`web/auth.rs`); a 401 from an // identity-aware proxy in front of it carries the proxy's challenge instead. const resp401 = () => ({ diff --git a/src/lib/transport/web-transport.ts b/src/lib/transport/web-transport.ts index 9468c21e1da..07e236f1335 100644 --- a/src/lib/transport/web-transport.ts +++ b/src/lib/transport/web-transport.ts @@ -177,15 +177,20 @@ export class WebTransport implements Transport { () => controller.abort(), effectiveTimeoutMs ) + const method = options?.webMethod ?? "POST" + const isGet = method === "GET" + const url = options?.webPath + ? `${this.baseUrl}${options.webPath.startsWith("/") ? "" : "/"}${options.webPath}` + : `${this.baseUrl}/api/${command}` let res: Response try { - res = await fetch(`${this.baseUrl}/api/${command}`, { - method: "POST", + res = await fetch(url, { + method, headers: { - "Content-Type": "application/json", + ...(isGet ? {} : { "Content-Type": "application/json" }), ...codegAuthHeaders(token), }, - body: JSON.stringify(args ?? {}), + body: isGet ? undefined : JSON.stringify(args ?? {}), signal: controller.signal, }) } catch (err) { diff --git a/src/lib/types.ts b/src/lib/types.ts index 6748fd130dc..b18574b5831 100644 --- a/src/lib/types.ts +++ b/src/lib/types.ts @@ -6227,3 +6227,27 @@ export interface DeepSeekModelCatalog { * fixed, sessions run on the agent's built-in catalog instead. */ invalid: string | null } + +// ?????? Agent Quota ?????????????????????????????????????????????????????????????????????????????????????????????????????????????????????????? + +export interface QuotaWindow { + label: string + usedPercent: number + remainingPercent: number + resetsAt?: string | null + resetInSeconds?: number | null +} + +export interface SpendLimit { + usedUsd: number + limitUsd: number +} + +export interface AgentQuotaInfo { + agentType: string + planName?: string | null + shortWindow?: QuotaWindow | null + weeklyWindow?: QuotaWindow | null + spendLimit?: SpendLimit | null + lastUpdated: string +}