diff --git a/src-tauri/src/app_config.rs b/src-tauri/src/app_config.rs index 0262bf2c4..4623d4e39 100644 --- a/src-tauri/src/app_config.rs +++ b/src-tauri/src/app_config.rs @@ -316,7 +316,7 @@ use crate::prompt_files::prompt_file_path; use crate::provider::ProviderManager; /// 应用类型 -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum AppType { Claude, diff --git a/src-tauri/src/commands/provider.rs b/src-tauri/src/commands/provider.rs index 565cc8bac..30db8f7ff 100644 --- a/src-tauri/src/commands/provider.rs +++ b/src-tauri/src/commands/provider.rs @@ -1,5 +1,5 @@ use indexmap::IndexMap; -use tauri::State; +use tauri::{Emitter, State}; use crate::app_config::AppType; use crate::commands::copilot::CopilotAuthState; @@ -153,19 +153,55 @@ pub fn import_default_config(state: State<'_, AppState>, app: String) -> Result< #[allow(non_snake_case)] #[tauri::command] pub async fn queryProviderUsage( + app_handle: tauri::AppHandle, state: State<'_, AppState>, copilot_state: State<'_, CopilotAuthState>, #[allow(non_snake_case)] providerId: String, // 使用 camelCase 匹配前端 app: String, ) -> Result { let app_type = AppType::from_str(&app).map_err(|e| e.to_string())?; + // inner 可能以两种形式失败： + // 1) 返回 Ok(UsageResult { success: false, .. }) —— 业务失败（401、脚本报错等） + // 2) 返回 Err(String) —— RPC/DB/Copilot fetch_usage 等 transport 层失败 + // 两种都要把"失败"写进 UsageCache 并刷新托盘，让 format_script_summary 的 + // success 守卫生效、suffix 自然消失，避免旧 success 快照长期滞留。 + // 同时保持原始 Err 返回给前端 React Query 的 onError 回调，不吞错误。 + let inner = + query_provider_usage_inner(&state, &copilot_state, app_type.clone(), &providerId).await; + let snapshot = match &inner { + Ok(r) => r.clone(), + Err(err_msg) => crate::provider::UsageResult { + success: false, + data: None, + error: Some(err_msg.clone()), + }, + }; + let payload = serde_json::json!({ + "kind": "script", + "appType": app_type.as_str(), + "providerId": &providerId, + "data": &snapshot, + }); + if let Err(e) = app_handle.emit("usage-cache-updated", payload) { + log::error!("emit usage-cache-updated (script) 失败: {e}"); + } + state.usage_cache.put_script(app_type, providerId, snapshot); + crate::tray::schedule_tray_refresh(&app_handle); + inner +} +async fn query_provider_usage_inner( + state: &AppState, + copilot_state: &CopilotAuthState, + app_type: AppType, + provider_id: &str, +) -> Result { // 从数据库读取供应商信息，检查特殊模板类型 let providers = state .db .get_all_providers(app_type.as_str()) .map_err(|e| format!("Failed to get providers: {e}"))?; - let provider = providers.get(&providerId); + let provider = providers.get(provider_id); let usage_script = provider .and_then(|p| p.meta.as_ref()) .and_then(|m| m.usage_script.as_ref()); @@ -294,7 +330,7 @@ pub async fn queryProviderUsage( } // ── 通用 JS 脚本路径 ── - ProviderService::query_usage(state.inner(), app_type, &providerId) + ProviderService::query_usage(state, app_type, provider_id) .await .map_err(|e| e.to_string()) } @@ -406,7 +442,7 @@ pub fn update_providers_sort_order( use crate::provider::UniversalProvider; use std::collections::HashMap; -use tauri::{AppHandle, Emitter}; +use tauri::AppHandle; #[derive(Clone, serde::Serialize)] pub struct UniversalProviderSyncedEvent { diff --git a/src-tauri/src/commands/subscription.rs b/src-tauri/src/commands/subscription.rs index 00b456eb7..053045b71 100644 --- a/src-tauri/src/commands/subscription.rs +++ b/src-tauri/src/commands/subscription.rs @@ -1,10 +1,41 @@ -use crate::services::subscription::SubscriptionQuota; +use std::str::FromStr; +use tauri::{Emitter, State}; + +use crate::app_config::AppType; +use crate::services::subscription::{CredentialStatus, SubscriptionQuota}; +use crate::store::AppState; /// 查询官方订阅额度 /// /// 读取 CLI 工具已有的 OAuth 凭据并调用官方 API 获取使用额度。 -/// 不需要 AppState（不访问数据库），直接读文件 + 发 HTTP。 +/// 结果（无论业务失败还是 transport 层 Err）都会写入 `UsageCache`、通知托盘 +/// 刷新，并 emit `usage-cache-updated`，让前端 React Query 与托盘共享同一份 +/// 最新数据。失败快照写入后 `format_subscription_summary` 会通过 `success=false` +/// 守卫返回 `None`，托盘 suffix 自然消失，避免长期滞留旧配额数字。 +/// Err 原样向前端返回，React Query 的 onError 不会被吞掉。 #[tauri::command] -pub async fn get_subscription_quota(tool: String) -> Result { - crate::services::subscription::get_subscription_quota(&tool).await +pub async fn get_subscription_quota( + app: tauri::AppHandle, + state: State<'_, AppState>, + tool: String, +) -> Result { + let inner = crate::services::subscription::get_subscription_quota(&tool).await; + let snapshot = match &inner { + Ok(q) => q.clone(), + // transport 层 Err —— 凭据状态不明，用 Valid 表达"凭据没问题，是通信/parse 出错"。 + Err(err_msg) => SubscriptionQuota::error(&tool, CredentialStatus::Valid, err_msg.clone()), + }; + if let Ok(app_type) = AppType::from_str(&tool) { + let payload = serde_json::json!({ + "kind": "subscription", + "appType": app_type.as_str(), + "data": &snapshot, + }); + if let Err(e) = app.emit("usage-cache-updated", payload) { + log::error!("emit usage-cache-updated (subscription) 失败: {e}"); + } + state.usage_cache.put_subscription(app_type, snapshot); + crate::tray::schedule_tray_refresh(&app); + } + inner } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index ac2b227e9..0966cd49d 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -748,9 +748,16 @@ pub fn run() { // 构建托盘 let mut tray_builder = TrayIconBuilder::with_id(tray::TRAY_ID) - .on_tray_icon_event(|_tray, event| match event { - // 左键点击已通过 show_menu_on_left_click(true) 打开菜单，这里不再额外处理 - TrayIconEvent::Click { .. } => {} + .on_tray_icon_event(|tray, event| match event { + // 鼠标悬停/点击到托盘图标时，后台异步刷新用量缓存， + // 让用户下一次（或快速打开菜单的那一刻）看到较新的数字。 + // refresh_all_usage_in_tray 内部有 10 秒防抖。 + TrayIconEvent::Enter { .. } | TrayIconEvent::Click { .. } => { + let app = tray.app_handle().clone(); + tauri::async_runtime::spawn(async move { + crate::tray::refresh_all_usage_in_tray(&app).await; + }); + } _ => log::debug!("unhandled event {event:?}"), }) .menu(&menu) diff --git a/src-tauri/src/provider.rs b/src-tauri/src/provider.rs index 6e076a310..60c591fe7 100644 --- a/src-tauri/src/provider.rs +++ b/src-tauri/src/provider.rs @@ -65,6 +65,25 @@ impl Provider { in_failover_queue: false, } } + + pub fn is_codex_oauth(&self) -> bool { + self.meta.as_ref().and_then(|m| m.provider_type.as_deref()) == Some("codex_oauth") + } + + pub fn codex_fast_mode_enabled(&self) -> bool { + self.meta + .as_ref() + .map(|m| m.codex_fast_mode_enabled()) + .unwrap_or(false) + } + + pub fn has_usage_script_enabled(&self) -> bool { + self.meta + .as_ref() + .and_then(|m| m.usage_script.as_ref()) + .map(|s| s.enabled) + .unwrap_or(false) + } } /// 供应商管理器 diff --git a/src-tauri/src/proxy/providers/claude.rs b/src-tauri/src/proxy/providers/claude.rs index 585a83990..dd2e76f94 100644 --- a/src-tauri/src/proxy/providers/claude.rs +++ b/src-tauri/src/proxy/providers/claude.rs @@ -89,11 +89,7 @@ pub fn transform_claude_request_for_api_format( session_id: Option<&str>, shadow_store: Option<&super::gemini_shadow::GeminiShadowStore>, ) -> Result { - let is_codex_oauth = provider - .meta - .as_ref() - .and_then(|m| m.provider_type.as_deref()) - == Some("codex_oauth"); + let is_codex_oauth = provider.is_codex_oauth(); // Copilot 场景：优先从 metadata.user_id 提取 session ID 作为 cache key // 格式: "uuid_sessionId" → 提取 "_" 后面的部分作为 session 标识 @@ -151,11 +147,7 @@ pub fn transform_claude_request_for_api_format( "openai_responses" => { // Codex OAuth (ChatGPT Plus/Pro 反代) 需要在请求体里强制 store: false // + include: ["reasoning.encrypted_content"]，由 transform 层统一处理。 - let codex_fast_mode = provider - .meta - .as_ref() - .map(|m| m.codex_fast_mode_enabled()) - .unwrap_or(false); + let codex_fast_mode = provider.codex_fast_mode_enabled(); super::transform_responses::anthropic_to_responses( body, Some(cache_key), diff --git a/src-tauri/src/proxy/providers/transform_responses.rs b/src-tauri/src/proxy/providers/transform_responses.rs index 9ed4209c9..ddff5136e 100644 --- a/src-tauri/src/proxy/providers/transform_responses.rs +++ b/src-tauri/src/proxy/providers/transform_responses.rs @@ -528,7 +528,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["model"], "gpt-4o"); assert_eq!(result["max_output_tokens"], 1024); assert_eq!(result["input"][0]["role"], "user"); @@ -547,7 +547,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["instructions"], "You are a helpful assistant."); // system should not appear in input assert_eq!(result["input"].as_array().unwrap().len(), 1); @@ -565,7 +565,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["instructions"], "Part 1\n\nPart 2"); } @@ -582,7 +582,7 @@ mod tests { }] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["tools"][0]["type"], "function"); assert_eq!(result["tools"][0]["name"], "get_weather"); assert!(result["tools"][0].get("parameters").is_some()); @@ -599,7 +599,7 @@ mod tests { "tool_choice": {"type": "any"} }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["tool_choice"], "required"); } @@ -612,7 +612,7 @@ mod tests { "tool_choice": {"type": "tool", "name": "get_weather"} }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["tool_choice"]["type"], "function"); assert_eq!(result["tool_choice"]["name"], "get_weather"); } @@ -631,7 +631,7 @@ mod tests { }] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); let input_arr = result["input"].as_array().unwrap(); // Should produce: assistant message (text) + function_call item @@ -661,7 +661,7 @@ mod tests { }] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); let input_arr = result["input"].as_array().unwrap(); // Should produce: function_call_output item (lifted) @@ -685,7 +685,7 @@ mod tests { }] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); let input_arr = result["input"].as_array().unwrap(); // thinking should be discarded, only text remains @@ -708,7 +708,7 @@ mod tests { }] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); let content = result["input"][0]["content"].as_array().unwrap(); assert_eq!(content[0]["type"], "input_text"); @@ -866,7 +866,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["model"], "o3-mini"); } @@ -878,7 +878,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, Some("my-provider-id"), false, true).unwrap(); + let result = anthropic_to_responses(input, Some("my-provider-id"), false, false).unwrap(); assert_eq!(result["prompt_cache_key"], "my-provider-id"); } @@ -896,7 +896,7 @@ mod tests { }] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert!(result["tools"][0].get("cache_control").is_none()); } @@ -913,7 +913,7 @@ mod tests { }] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert!(result["input"][0]["content"][0] .get("cache_control") .is_none()); @@ -975,7 +975,7 @@ mod tests { "max_tokens": 4096, "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["max_output_tokens"], 4096); assert!(result.get("max_completion_tokens").is_none()); } @@ -989,7 +989,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["reasoning"]["effort"], "xhigh"); } @@ -1003,7 +1003,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["reasoning"]["effort"], "low"); } @@ -1016,7 +1016,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["reasoning"]["effort"], "low"); } @@ -1029,7 +1029,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["reasoning"]["effort"], "medium"); } @@ -1042,7 +1042,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["reasoning"]["effort"], "high"); } @@ -1055,7 +1055,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["reasoning"]["effort"], "xhigh"); } @@ -1068,7 +1068,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert!(result.get("reasoning").is_none()); } @@ -1102,7 +1102,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert!(result.get("store").is_none()); assert!(result.get("service_tier").is_none()); @@ -1184,7 +1184,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["max_output_tokens"], json!(1024)); } @@ -1316,7 +1316,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert_eq!(result["temperature"], json!(0.7)); assert_eq!(result["top_p"], json!(0.9)); @@ -1332,7 +1332,7 @@ mod tests { "messages": [{"role": "user", "content": "Hello"}] }); - let result = anthropic_to_responses(input, None, false, true).unwrap(); + let result = anthropic_to_responses(input, None, false, false).unwrap(); assert!( result.get("parallel_tool_calls").is_none(), diff --git a/src-tauri/src/services/mod.rs b/src-tauri/src/services/mod.rs index 5a6bb798f..d9fb5077c 100644 --- a/src-tauri/src/services/mod.rs +++ b/src-tauri/src/services/mod.rs @@ -16,6 +16,7 @@ pub mod skill; pub mod speedtest; pub mod stream_check; pub mod subscription; +pub mod usage_cache; pub mod usage_stats; pub mod webdav; pub mod webdav_auto_sync; @@ -30,6 +31,7 @@ pub use proxy::ProxyService; #[allow(unused_imports)] pub use skill::{DiscoverableSkill, Skill, SkillRepo, SkillService}; pub use speedtest::{EndpointLatency, SpeedtestService}; +pub use usage_cache::UsageCache; #[allow(unused_imports)] pub use usage_stats::{ DailyStats, LogFilters, ModelStats, PaginatedLogs, ProviderLimitStatus, ProviderStats, diff --git a/src-tauri/src/services/skill.rs b/src-tauri/src/services/skill.rs index d395fe58f..06d159677 100644 --- a/src-tauri/src/services/skill.rs +++ b/src-tauri/src/services/skill.rs @@ -941,10 +941,6 @@ impl SkillService { None => continue, }; - if !remote_skill_dir.is_dir() { - continue; - } - let remote_hash = match Self::compute_dir_hash(&remote_skill_dir) { Ok(h) => h, Err(e) => { diff --git a/src-tauri/src/services/stream_check.rs b/src-tauri/src/services/stream_check.rs index c4f7fea76..4befbefe1 100644 --- a/src-tauri/src/services/stream_check.rs +++ b/src-tauri/src/services/stream_check.rs @@ -355,16 +355,8 @@ impl StreamCheckService { }); // Codex OAuth (ChatGPT Plus/Pro 反代) 需要 store:false + include 标记， // 否则 Stream Check 会和生产路径一样被服务端 400 拒绝。 - let is_codex_oauth = provider - .meta - .as_ref() - .and_then(|m| m.provider_type.as_deref()) - == Some("codex_oauth"); - let codex_fast_mode = provider - .meta - .as_ref() - .map(|m| m.codex_fast_mode_enabled()) - .unwrap_or(false); + let is_codex_oauth = provider.is_codex_oauth(); + let codex_fast_mode = provider.codex_fast_mode_enabled(); let body = if is_openai_responses { anthropic_to_responses( diff --git a/src-tauri/src/services/subscription.rs b/src-tauri/src/services/subscription.rs index ee4a0b115..e4530a08c 100644 --- a/src-tauri/src/services/subscription.rs +++ b/src-tauri/src/services/subscription.rs @@ -291,12 +291,22 @@ struct ApiExtraUsage { currency: Option, } -/// 已知的 Claude 用量窗口名称 +/// 已知的 Claude 用量窗口名称。`QuotaTier::name` 会是其中之一。 +pub const TIER_FIVE_HOUR: &str = "five_hour"; +pub const TIER_SEVEN_DAY: &str = "seven_day"; +pub const TIER_SEVEN_DAY_OPUS: &str = "seven_day_opus"; +pub const TIER_SEVEN_DAY_SONNET: &str = "seven_day_sonnet"; + +/// Gemini 用量分组名称（按模型而非时间窗口）。`classify_gemini_model` 输出。 +pub const TIER_GEMINI_PRO: &str = "gemini_pro"; +pub const TIER_GEMINI_FLASH: &str = "gemini_flash"; +pub const TIER_GEMINI_FLASH_LITE: &str = "gemini_flash_lite"; + const KNOWN_TIERS: &[&str] = &[ - "five_hour", - "seven_day", - "seven_day_opus", - "seven_day_sonnet", + TIER_FIVE_HOUR, + TIER_SEVEN_DAY, + TIER_SEVEN_DAY_OPUS, + TIER_SEVEN_DAY_SONNET, ]; /// 查询 Claude 官方订阅额度 @@ -993,11 +1003,11 @@ fn extract_project_id(value: &serde_json::Value) -> Option { /// 将 Gemini 模型 ID 分类为 Pro / Flash / Flash Lite fn classify_gemini_model(model_id: &str) -> &str { if model_id.contains("flash-lite") { - "gemini_flash_lite" + TIER_GEMINI_FLASH_LITE } else if model_id.contains("flash") { - "gemini_flash" + TIER_GEMINI_FLASH } else if model_id.contains("pro") { - "gemini_pro" + TIER_GEMINI_PRO } else { model_id } @@ -1152,9 +1162,9 @@ async fn query_gemini_quota(access_token: &str) -> SubscriptionQuota { // 转换为 tiers（remainingFraction → utilization: 已用百分比） let sort_order = |name: &str| -> usize { match name { - "gemini_pro" => 0, - "gemini_flash" => 1, - "gemini_flash_lite" => 2, + TIER_GEMINI_PRO => 0, + TIER_GEMINI_FLASH => 1, + TIER_GEMINI_FLASH_LITE => 2, _ => 3, } }; diff --git a/src-tauri/src/services/usage_cache.rs b/src-tauri/src/services/usage_cache.rs new file mode 100644 index 000000000..638040fca --- /dev/null +++ b/src-tauri/src/services/usage_cache.rs @@ -0,0 +1,143 @@ +//! 托盘展示用的用量缓存（进程内、写穿式）。 +//! +//! 各 usage 查询命令成功时写入；系统托盘构建菜单时读取。不持久化， +//! 进程重启即空，由下一次自动查询或托盘悬停触发的刷新重新填充。 + +use std::collections::HashMap; +use std::sync::RwLock; + +use crate::app_config::AppType; +use crate::provider::UsageResult; +use crate::services::subscription::SubscriptionQuota; + +#[derive(Default)] +pub struct UsageCache { + subscription: RwLock>, + script: RwLock>, +} + +impl UsageCache { + pub fn new() -> Self { + Self::default() + } + + pub fn put_subscription(&self, app_type: AppType, quota: SubscriptionQuota) { + if let Ok(mut w) = self.subscription.write() { + w.insert(app_type, quota); + } + } + + pub fn put_script(&self, app_type: AppType, provider_id: String, result: UsageResult) { + if let Ok(mut w) = self.script.write() { + w.insert((app_type, provider_id), result); + } + } + + /// 以借用形式暴露订阅快照，避免托盘每次重建时深拷贝整个 `SubscriptionQuota`。 + pub fn with_subscription( + &self, + app_type: &AppType, + f: impl FnOnce(&SubscriptionQuota) -> R, + ) -> Option { + self.subscription + .read() + .ok() + .and_then(|r| r.get(app_type).map(f)) + } + + /// 以借用形式暴露脚本型用量结果，同上。 + pub fn with_script( + &self, + app_type: &AppType, + provider_id: &str, + f: impl FnOnce(&UsageResult) -> R, + ) -> Option { + self.script + .read() + .ok() + .and_then(|r| r.get(&(app_type.clone(), provider_id.to_string())).map(f)) + } + + pub fn invalidate_script(&self, app_type: &AppType, provider_id: &str) { + // 热路径会对每个禁用脚本的 provider 在托盘重建时调用一次：先走读锁 + // `contains_key` 快速放行"本来就不在缓存里"的常见情况，避免无谓的写锁升级。 + let key = (app_type.clone(), provider_id.to_string()); + if !self.script.read().is_ok_and(|r| r.contains_key(&key)) { + return; + } + if let Ok(mut w) = self.script.write() { + w.remove(&key); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::services::subscription::CredentialStatus; + + fn fake_quota() -> SubscriptionQuota { + SubscriptionQuota { + tool: "claude".to_string(), + credential_status: CredentialStatus::Valid, + credential_message: None, + success: true, + tiers: vec![], + extra_usage: None, + error: None, + queried_at: Some(0), + } + } + + fn fake_result() -> UsageResult { + UsageResult { + success: true, + data: None, + error: None, + } + } + + #[test] + fn subscription_round_trip() { + let cache = UsageCache::new(); + assert!(cache + .with_subscription(&AppType::Claude, |q| q.success) + .is_none()); + cache.put_subscription(AppType::Claude, fake_quota()); + let got = cache + .with_subscription(&AppType::Claude, |q| q.success) + .unwrap(); + assert!(got); + assert!(cache + .with_subscription(&AppType::Codex, |q| q.success) + .is_none()); + } + + #[test] + fn script_round_trip_and_invalidate() { + let cache = UsageCache::new(); + assert!(cache + .with_script(&AppType::Codex, "pid", |r| r.success) + .is_none()); + cache.put_script(AppType::Codex, "pid".to_string(), fake_result()); + assert!(cache + .with_script(&AppType::Codex, "pid", |r| r.success) + .is_some()); + cache.invalidate_script(&AppType::Codex, "pid"); + assert!(cache + .with_script(&AppType::Codex, "pid", |r| r.success) + .is_none()); + } + + #[test] + fn script_keys_isolated_by_app_type() { + let cache = UsageCache::new(); + cache.put_script(AppType::Claude, "same".to_string(), fake_result()); + assert!(cache + .with_script(&AppType::Claude, "same", |r| r.success) + .is_some()); + assert!(cache + .with_script(&AppType::Codex, "same", |r| r.success) + .is_none()); + } +} diff --git a/src-tauri/src/store.rs b/src-tauri/src/store.rs index 9c626a5ec..f34e1de9a 100644 --- a/src-tauri/src/store.rs +++ b/src-tauri/src/store.rs @@ -1,11 +1,12 @@ use crate::database::Database; -use crate::services::ProxyService; +use crate::services::{ProxyService, UsageCache}; use std::sync::Arc; /// 全局应用状态 pub struct AppState { pub db: Arc, pub proxy_service: ProxyService, + pub usage_cache: Arc, } impl AppState { @@ -13,6 +14,10 @@ impl AppState { pub fn new(db: Arc) -> Self { let proxy_service = ProxyService::new(db.clone()); - Self { db, proxy_service } + Self { + db, + proxy_service, + usage_cache: Arc::new(UsageCache::new()), + } } } diff --git a/src-tauri/src/tray.rs b/src-tauri/src/tray.rs index a1d852bc9..73088f208 100644 --- a/src-tauri/src/tray.rs +++ b/src-tauri/src/tray.rs @@ -2,13 +2,20 @@ //! //! 负责系统托盘图标和菜单的创建、更新和事件处理。 -use tauri::menu::{CheckMenuItem, Menu, MenuBuilder, MenuItem, SubmenuBuilder}; +use once_cell::sync::Lazy; +use tauri::menu::{CheckMenuItem, Menu, MenuBuilder, MenuItem, Submenu, SubmenuBuilder}; use tauri::{Emitter, Manager}; use crate::app_config::AppType; use crate::error::AppError; use crate::store::AppState; +/// 每个 app 分区的子菜单句柄，用于 usage 更新时就地改 label 而非整菜单重建。 +/// `create_tray_menu` 每次重建都会整表覆盖写入，保证句柄始终指向当前活跃菜单。 +static TRAY_SECTION_SUBMENUS: Lazy< + std::sync::Mutex>>, +> = Lazy::new(|| std::sync::Mutex::new(std::collections::HashMap::new())); + /// 托盘菜单文本（国际化） #[derive(Clone, Copy)] pub struct TrayTexts { @@ -84,6 +91,140 @@ pub const TRAY_SECTIONS: [TrayAppSection; 3] = [ }, ]; +/// 配色阈值（与前端 `utilizationColor` 语义一致）。 +const UTIL_WARN_PCT: f64 = 70.0; +const UTIL_DANGER_PCT: f64 = 90.0; + +fn emoji_for_utilization(pct: f64) -> &'static str { + if pct >= UTIL_DANGER_PCT { + "\u{1F534}" // 🔴 + } else if pct >= UTIL_WARN_PCT { + "\u{1F7E0}" // 🟠 + } else { + "\u{1F7E2}" // 🟢 + } +} + +fn format_subscription_summary( + quota: &crate::services::subscription::SubscriptionQuota, +) -> Option { + use crate::services::subscription::{ + TIER_FIVE_HOUR, TIER_GEMINI_FLASH, TIER_GEMINI_FLASH_LITE, TIER_GEMINI_PRO, TIER_SEVEN_DAY, + }; + if !quota.success { + return None; + } + + // 按 tool 选取主卡槽 tier 并映射到短 label： + // Claude / Codex 沿用时间窗口（h=5 小时，w=7 天）； + // Gemini 用模型维度（p=pro，f=flash，l=flash-lite）——Gemini 后端 tier + // 命名是 gemini_pro / gemini_flash / gemini_flash_lite，与时间窗口不同命名空间。 + // flash_lite 必须纳入：否则 lite 利用率最高时色标偏低，与前端 footer 行为不一致。 + let parts: Vec<(&'static str, f64)> = match quota.tool.as_str() { + "gemini" => { + let mut v = Vec::new(); + if let Some(t) = quota.tiers.iter().find(|t| t.name == TIER_GEMINI_PRO) { + v.push(("p", t.utilization)); + } + if let Some(t) = quota.tiers.iter().find(|t| t.name == TIER_GEMINI_FLASH) { + v.push(("f", t.utilization)); + } + if let Some(t) = quota + .tiers + .iter() + .find(|t| t.name == TIER_GEMINI_FLASH_LITE) + { + v.push(("l", t.utilization)); + } + v + } + _ => { + let mut v = Vec::new(); + if let Some(t) = quota.tiers.iter().find(|t| t.name == TIER_FIVE_HOUR) { + v.push(("h", t.utilization)); + } + if let Some(t) = quota.tiers.iter().find(|t| t.name == TIER_SEVEN_DAY) { + v.push(("w", t.utilization)); + } + v + } + }; + + if parts.is_empty() { + return None; + } + + // 色标取所有已选 tier 里最高的利用率——用户更关心"离上限多近"。 + let worst = parts + .iter() + .map(|(_, u)| *u) + .fold(f64::NEG_INFINITY, f64::max); + if !worst.is_finite() { + return None; + } + + let emoji = emoji_for_utilization(worst); + let body = parts + .iter() + .map(|(label, u)| format!("{label}{}%", u.round() as i64)) + .collect::>() + .join(" "); + Some(format!("{emoji} {body}")) +} + +fn format_script_summary(result: &crate::provider::UsageResult) -> Option { + if !result.success { + return None; + } + let data = result.data.as_ref()?.first()?; + let pct = match (data.used, data.total) { + (Some(used), Some(total)) if total > 0.0 => used / total * 100.0, + _ => return None, + }; + let emoji = emoji_for_utilization(pct); + let plan = data.plan_name.as_deref().unwrap_or(""); + let rounded = pct.round() as i64; + if plan.is_empty() { + Some(format!("{} {}%", emoji, rounded)) + } else { + Some(format!("{} {} {}%", emoji, plan, rounded)) + } +} + +fn format_usage_suffix( + app_state: &AppState, + app_type: &AppType, + provider: &crate::provider::Provider, + provider_id: &str, +) -> Option { + // 当前脚本是否启用：禁用/删除时不再沿用旧 UsageCache 结果， + // 并顺手 invalidate，防止后续重建继续命中过期数据。 + if provider.has_usage_script_enabled() { + // 脚本缓存优先（覆盖 Copilot/coding_plan/balance/自定义脚本），借用访问避免克隆整条 UsageResult。 + if let Some(Some(s)) = + app_state + .usage_cache + .with_script(app_type, provider_id, format_script_summary) + { + return Some(format!(" · {s}")); + } + } else { + app_state + .usage_cache + .invalidate_script(app_type, provider_id); + } + + if provider.category.as_deref() == Some("official") { + if let Some(Some(s)) = app_state + .usage_cache + .with_subscription(app_type, format_subscription_summary) + { + return Some(format!(" · {s}")); + } + } + None +} + /// 对供应商列表排序：sort_index → created_at → name fn sort_providers( providers: &indexmap::IndexMap, @@ -291,6 +432,8 @@ pub fn create_tray_menu( let visible_apps = app_settings.visible_apps.unwrap_or_default(); let mut menu_builder = MenuBuilder::new(app); + let mut section_handles: std::collections::HashMap> = + std::collections::HashMap::new(); // 顶部：打开主界面 let show_main_item = @@ -323,10 +466,13 @@ pub fn create_tray_menu( })?; menu_builder = menu_builder.item(&empty_item); } else { - // 有供应商：构建子菜单 - let current_name = providers.get(¤t_id).map(|p| p.name.as_str()); - let submenu_label = match current_name { - Some(name) => format!("{} · {}", section.header_label, name), + let current_provider = providers.get(¤t_id); + let submenu_label = match current_provider { + Some(p) => { + let suffix = format_usage_suffix(app_state, §ion.app_type, p, ¤t_id) + .unwrap_or_default(); + format!("{} · {}{}", section.header_label, p.name, suffix) + } None => section.header_label.to_string(), }; let submenu_id = format!("submenu_{}", app_type_str); @@ -369,6 +515,7 @@ pub fn create_tray_menu( let submenu = submenu_builder.build().map_err(|e| { AppError::Message(format!("构建{}子菜单失败: {e}", section.log_name)) })?; + section_handles.insert(section.app_type.clone(), submenu.clone()); menu_builder = menu_builder.item(&submenu); } @@ -393,9 +540,51 @@ pub fn create_tray_menu( menu_builder = menu_builder.item(&quit_item); - menu_builder + let menu = menu_builder .build() - .map_err(|e| AppError::Message(format!("构建菜单失败: {e}"))) + .map_err(|e| AppError::Message(format!("构建菜单失败: {e}")))?; + + *TRAY_SECTION_SUBMENUS + .lock() + .unwrap_or_else(|p| p.into_inner()) = section_handles; + + Ok(menu) +} + +/// 就地更新各 app 分区子菜单的标题（usage 后缀变化时走这条）， +/// 避免 `set_menu` 导致用户打开中的菜单被关闭。 +/// 句柄由上一次 `create_tray_menu` 填充；为空（从未构建过菜单）时无事发生。 +fn update_tray_usage_labels(app: &tauri::AppHandle) { + let Some(app_state) = app.try_state::() else { + return; + }; + let handles = match TRAY_SECTION_SUBMENUS.lock() { + Ok(g) => g, + Err(poisoned) => poisoned.into_inner(), + }; + + for section in TRAY_SECTIONS.iter() { + let Some(submenu) = handles.get(§ion.app_type) else { + continue; + }; + let Ok(providers) = app_state.db.get_all_providers(section.app_type.as_str()) else { + continue; + }; + let Ok(Some(current_id)) = + crate::settings::get_effective_current_provider(&app_state.db, §ion.app_type) + else { + continue; + }; + let Some(provider) = providers.get(¤t_id) else { + continue; + }; + let suffix = format_usage_suffix(&app_state, §ion.app_type, provider, ¤t_id) + .unwrap_or_default(); + let new_label = format!("{} · {}{}", section.header_label, provider.name, suffix); + if let Err(e) = submenu.set_text(&new_label) { + log::debug!("[Tray] 更新{}子菜单标题失败: {e}", section.log_name); + } + } } pub fn refresh_tray_menu(app: &tauri::AppHandle) { @@ -412,17 +601,6 @@ pub fn refresh_tray_menu(app: &tauri::AppHandle) { } } -#[cfg(test)] -mod tests { - use super::TRAY_ID; - - #[test] - fn tray_id_is_unique_to_app() { - assert_eq!(TRAY_ID, "cc-switch"); - assert_ne!(TRAY_ID, "main"); - } -} - #[cfg(target_os = "macos")] pub fn apply_tray_policy(app: &tauri::AppHandle, dock_visible: bool) { use tauri::ActivationPolicy; @@ -491,3 +669,269 @@ pub fn handle_tray_menu_event(app: &tauri::AppHandle, event_id: &str) { } } } + +static LAST_TRAY_USAGE_REFRESH: std::sync::Mutex> = + std::sync::Mutex::new(None); +const MIN_TRAY_USAGE_REFRESH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(10); + +/// 合并多次快速触发的"usage 标题软更新"：批量刷新期间多个 usage 命令 +/// 同时成功时，只会产生一次就地 `set_text` 批量调用。走软更新而不是 +/// `refresh_tray_menu` 整建，避免用户打开中的菜单被 macOS 系统关闭。 +static TRAY_REBUILD_SCHEDULED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + +pub fn schedule_tray_refresh(app: &tauri::AppHandle) { + use std::sync::atomic::Ordering; + if TRAY_REBUILD_SCHEDULED.swap(true, Ordering::AcqRel) { + return; + } + let app = app.clone(); + tauri::async_runtime::spawn_blocking(move || { + // 50ms 合窗：让同一轮 React Query / 托盘批量刷新触发的多个写入 + // 共享一次标题更新。 + std::thread::sleep(std::time::Duration::from_millis(50)); + TRAY_REBUILD_SCHEDULED.store(false, Ordering::Release); + update_tray_usage_labels(&app); + }); +} + +/// 并行刷新每个可见 app "当前 provider" 的用量；成功 / 失败结果都通过各 +/// command 的 write-through 逻辑写入 `UsageCache`，单次重建菜单由 +/// `schedule_tray_refresh` 做合并。内部 10 秒节流防止鼠标悬停反复进出时 +/// 雪崩请求；互斥锁被毒化时以上次状态为准继续推进，不会永久阻塞。 +/// +/// 刷新面与 `format_usage_suffix` 的展示面严格对齐 —— 每次悬停最多发 +/// `TRAY_SECTIONS.len()` 次外部请求，script 优先（覆盖 coding_plan / balance / +/// Copilot / 自定义脚本），否则当前 provider 必须是 `official` 才查订阅。 +pub(crate) async fn refresh_all_usage_in_tray(app: &tauri::AppHandle) { + use crate::commands::CopilotAuthState; + use futures::future::join_all; + + { + let mut guard = LAST_TRAY_USAGE_REFRESH + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let now = std::time::Instant::now(); + if let Some(last) = *guard { + if now.duration_since(last) < MIN_TRAY_USAGE_REFRESH_INTERVAL { + return; + } + } + *guard = Some(now); + } + + let Some(app_state) = app.try_state::() else { + return; + }; + + // 与 `create_tray_menu` 保持一致：用户隐藏的 app 不参与外部 API 查询， + // 避免在未使用的 app 上浪费请求、撞 rate limit 或反复触发鉴权失败日志。 + let visible_apps = crate::settings::get_settings() + .visible_apps + .unwrap_or_default(); + + let mut subscription_futures = Vec::new(); + let mut script_futures = Vec::new(); + + for section in TRAY_SECTIONS.iter() { + if !visible_apps.is_visible(§ion.app_type) { + continue; + } + + let app_type_str = section.app_type.as_str(); + let log_name = section.log_name; + + // 解析 effective current provider；未设置 / 出错都静默跳过， + // 与 create_tray_menu 的行为保持一致。 + let current_id = + match crate::settings::get_effective_current_provider(&app_state.db, §ion.app_type) + { + Ok(Some(id)) => id, + Ok(None) => continue, + Err(e) => { + log::warn!("[Tray] 读取{log_name}当前供应商失败: {e}"); + continue; + } + }; + // 只需当前 provider —— by-id 查询避免把整个 app 的 provider 列表加载 + // 进内存（每次悬停 × 3 sections 的热路径）。 + let current = match app_state.db.get_provider_by_id(¤t_id, app_type_str) { + Ok(Some(p)) => p, + Ok(None) => continue, + Err(e) => { + log::warn!("[Tray] 读取{log_name}当前供应商失败: {e}"); + continue; + } + }; + + // 与 format_usage_suffix 同一优先级：脚本启用 → 查脚本； + // 否则当前 provider 是 official → 查订阅；其它情况不发请求。 + if current.has_usage_script_enabled() { + let app_clone = app.clone(); + let state = app.state::(); + let copilot_state = app.state::(); + let provider_id = current_id.clone(); + let app_str = app_type_str.to_string(); + script_futures.push(async move { + if let Err(e) = crate::commands::queryProviderUsage( + app_clone, + state, + copilot_state, + provider_id.clone(), + app_str, + ) + .await + { + log::debug!("[Tray] 刷新{log_name}供应商 {provider_id} 用量失败: {e}"); + } + }); + } else if current.category.as_deref() == Some("official") { + let app_clone = app.clone(); + let state = app.state::(); + let tool = app_type_str.to_string(); + subscription_futures.push(async move { + if let Err(e) = + crate::commands::get_subscription_quota(app_clone, state, tool).await + { + log::debug!("[Tray] 刷新{log_name}订阅用量失败（可能未登录）: {e}"); + } + }); + } + } + + // 两组并行启动，整体等待 —— 订阅/脚本互不依赖，没必要串行。 + futures::future::join(join_all(subscription_futures), join_all(script_futures)).await; +} + +#[cfg(test)] +mod tests { + use super::{format_subscription_summary, TRAY_ID}; + use crate::services::subscription::{CredentialStatus, QuotaTier, SubscriptionQuota}; + + #[test] + fn tray_id_is_unique_to_app() { + assert_eq!(TRAY_ID, "cc-switch"); + assert_ne!(TRAY_ID, "main"); + } + + fn make_quota(tool: &str, success: bool, tiers: Vec) -> SubscriptionQuota { + SubscriptionQuota { + tool: tool.to_string(), + credential_status: CredentialStatus::Valid, + credential_message: None, + success, + tiers, + extra_usage: None, + error: None, + queried_at: Some(0), + } + } + + fn tier(name: &str, utilization: f64) -> QuotaTier { + QuotaTier { + name: name.to_string(), + utilization, + resets_at: None, + } + } + + #[test] + fn claude_summary_uses_h_and_w_labels() { + let quota = make_quota( + "claude", + true, + vec![tier("five_hour", 9.0), tier("seven_day", 27.0)], + ); + let s = format_subscription_summary("a).expect("should format"); + assert!(s.contains("h9%"), "expected h9% in {s}"); + assert!(s.contains("w27%"), "expected w27% in {s}"); + } + + #[test] + fn gemini_summary_uses_p_and_f_labels() { + let quota = make_quota( + "gemini", + true, + vec![tier("gemini_pro", 15.0), tier("gemini_flash", 42.0)], + ); + let s = format_subscription_summary("a).expect("should format"); + assert!(s.contains("p15%"), "expected p15% in {s}"); + assert!(s.contains("f42%"), "expected f42% in {s}"); + } + + #[test] + fn gemini_summary_includes_all_three_tiers() { + let quota = make_quota( + "gemini", + true, + vec![ + tier("gemini_pro", 5.0), + tier("gemini_flash", 42.0), + tier("gemini_flash_lite", 80.0), + ], + ); + let s = format_subscription_summary("a).expect("should format"); + assert!(s.contains("p5%"), "expected p5% in {s}"); + assert!(s.contains("f42%"), "expected f42% in {s}"); + assert!(s.contains("l80%"), "expected l80% in {s}"); + } + + #[test] + fn gemini_summary_lite_only_still_renders() { + // flash_lite 如果是 API 返回的唯一 tier，仍应显示（避免前端 footer 能看到、 + // 托盘空白的不对称）。 + let quota = make_quota("gemini", true, vec![tier("gemini_flash_lite", 80.0)]); + let s = format_subscription_summary("a).expect("should format"); + assert!(s.contains("l80%"), "expected l80% in {s}"); + } + + #[test] + fn gemini_summary_emoji_reflects_highest_tier_including_lite() { + // lite 是利用率最高的那条 → emoji 必须是红色，不能被 pro/flash 掩盖。 + let quota = make_quota( + "gemini", + true, + vec![ + tier("gemini_pro", 10.0), + tier("gemini_flash", 20.0), + tier("gemini_flash_lite", 95.0), + ], + ); + let s = format_subscription_summary("a).unwrap(); + assert!( + s.starts_with("\u{1F534}"), + "expected red emoji (lite worst) in {s}" + ); + } + + #[test] + fn worst_emoji_reflects_highest_utilization() { + // 🔴 = \u{1F534}; 任一 tier ≥ 90% 时预期显示红色。 + let quota = make_quota( + "claude", + true, + vec![tier("five_hour", 10.0), tier("seven_day", 95.0)], + ); + let s = format_subscription_summary("a).unwrap(); + assert!(s.starts_with("\u{1F534}"), "expected red emoji in {s}"); + } + + #[test] + fn failure_quota_returns_none() { + let quota = make_quota("claude", false, vec![tier("five_hour", 50.0)]); + assert!(format_subscription_summary("a).is_none()); + } + + #[test] + fn unknown_tiers_return_none() { + let quota = make_quota("claude", true, vec![tier("one_hour", 80.0)]); + assert!(format_subscription_summary("a).is_none()); + } + + #[test] + fn gemini_without_any_known_tiers_returns_none() { + // 完全没有 pro/flash/flash_lite 三种 tier 的退化响应 → None。 + let quota = make_quota("gemini", true, vec![tier("some_future_tier", 80.0)]); + assert!(format_subscription_summary("a).is_none()); + } +} diff --git a/src-tauri/tests/deeplink_import.rs b/src-tauri/tests/deeplink_import.rs index 5a0238c61..6cee273e8 100644 --- a/src-tauri/tests/deeplink_import.rs +++ b/src-tauri/tests/deeplink_import.rs @@ -1,8 +1,6 @@ use std::sync::Arc; -use cc_switch_lib::{ - import_provider_from_deeplink, parse_deeplink_url, AppState, Database, ProxyService, -}; +use cc_switch_lib::{import_provider_from_deeplink, parse_deeplink_url, AppState, Database}; #[path = "support.rs"] mod support; @@ -18,11 +16,7 @@ fn deeplink_import_claude_provider_persists_to_db() { let request = parse_deeplink_url(url).expect("parse deeplink url"); let db = Arc::new(Database::memory().expect("create memory db")); - let proxy_service = ProxyService::new(db.clone()); - let state = AppState { - db: db.clone(), - proxy_service, - }; + let state = AppState::new(db.clone()); let provider_id = import_provider_from_deeplink(&state, request.clone()) .expect("import provider from deeplink"); @@ -58,11 +52,7 @@ fn deeplink_import_codex_provider_builds_auth_and_config() { let request = parse_deeplink_url(url).expect("parse deeplink url"); let db = Arc::new(Database::memory().expect("create memory db")); - let proxy_service = ProxyService::new(db.clone()); - let state = AppState { - db: db.clone(), - proxy_service, - }; + let state = AppState::new(db.clone()); let provider_id = import_provider_from_deeplink(&state, request.clone()) .expect("import provider from deeplink"); diff --git a/src-tauri/tests/support.rs b/src-tauri/tests/support.rs index 83aa34d13..86f8f76e0 100644 --- a/src-tauri/tests/support.rs +++ b/src-tauri/tests/support.rs @@ -1,9 +1,7 @@ use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex, OnceLock}; -use cc_switch_lib::{ - update_settings, AppSettings, AppState, Database, MultiAppConfig, ProxyService, -}; +use cc_switch_lib::{update_settings, AppSettings, AppState, Database, MultiAppConfig}; /// 为测试设置隔离的 HOME 目录，避免污染真实用户数据。 pub fn ensure_test_home() -> &'static Path { @@ -62,8 +60,7 @@ pub fn test_mutex() -> &'static Mutex<()> { #[allow(dead_code)] pub fn create_test_state() -> Result> { let db = Arc::new(Database::init()?); - let proxy_service = ProxyService::new(db.clone()); - Ok(AppState { db, proxy_service }) + Ok(AppState::new(db)) } /// 创建测试用的 AppState，并从 MultiAppConfig 迁移数据 @@ -73,6 +70,5 @@ pub fn create_test_state_with_config( ) -> Result> { let db = Arc::new(Database::init()?); db.migrate_from_json(config)?; - let proxy_service = ProxyService::new(db.clone()); - Ok(AppState { db, proxy_service }) + Ok(AppState::new(db)) } diff --git a/src/App.tsx b/src/App.tsx index f621bc651..d3295100e 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -49,6 +49,7 @@ import { import { hermesApi } from "@/lib/api/hermes"; import { useProxyStatus } from "@/hooks/useProxyStatus"; import { useAutoCompact } from "@/hooks/useAutoCompact"; +import { useUsageCacheBridge } from "@/hooks/useUsageCacheBridge"; import { useLastValidValue } from "@/hooks/useLastValidValue"; import { extractErrorMessage } from "@/utils/errorUtils"; import { isTextEditableTarget } from "@/utils/domUtils"; @@ -236,6 +237,8 @@ function App() { const toolbarRef = useRef(null); const isToolbarCompact = useAutoCompact(toolbarRef); + useUsageCacheBridge(); + const promptPanelRef = useRef(null); const mcpPanelRef = useRef(null); const skillsPageRef = useRef(null); diff --git a/src/hooks/useUsageCacheBridge.ts b/src/hooks/useUsageCacheBridge.ts new file mode 100644 index 000000000..ed0deb13d --- /dev/null +++ b/src/hooks/useUsageCacheBridge.ts @@ -0,0 +1,66 @@ +import { useEffect } from "react"; +import { listen, type UnlistenFn } from "@tauri-apps/api/event"; +import { useQueryClient } from "@tanstack/react-query"; +import type { AppId } from "@/lib/api/types"; +import type { UsageResult } from "@/types"; +import type { SubscriptionQuota } from "@/types/subscription"; +import { usageKeys } from "@/lib/query/usage"; +import { subscriptionKeys } from "@/lib/query/subscription"; + +type UsageCacheUpdatedPayload = + | { + kind: "script"; + appType: AppId; + providerId: string; + data: UsageResult; + } + | { + kind: "subscription"; + appType: AppId; + data: SubscriptionQuota; + }; + +/** + * 后端 `UsageCache` 写入后会 emit `usage-cache-updated`，本 hook 把 payload 同步到 + * React Query 缓存，让托盘触发的刷新（不经前端）也能立刻反映到主界面，避免 + * React Query 与 Rust 侧两份缓存各自为战。 + */ +export function useUsageCacheBridge() { + const queryClient = useQueryClient(); + + useEffect(() => { + let unlisten: UnlistenFn | undefined; + let disposed = false; + + (async () => { + const off = await listen( + "usage-cache-updated", + (event) => { + const payload = event.payload; + if (payload.kind === "script") { + queryClient.setQueryData( + usageKeys.script(payload.providerId, payload.appType), + payload.data, + ); + } else if (payload.kind === "subscription") { + queryClient.setQueryData( + subscriptionKeys.quota(payload.appType), + payload.data, + ); + } + }, + ); + + if (disposed) { + off(); + } else { + unlisten = off; + } + })(); + + return () => { + disposed = true; + unlisten?.(); + }; + }, [queryClient]); +} diff --git a/src/lib/query/queries.ts b/src/lib/query/queries.ts index 43034d000..ef0711c36 100644 --- a/src/lib/query/queries.ts +++ b/src/lib/query/queries.ts @@ -17,6 +17,7 @@ import type { SessionMeta, SessionMessage, } from "@/types"; +import { usageKeys } from "@/lib/query/usage"; const sortProviders = ( providers: Record, @@ -113,7 +114,7 @@ export const useUsageQuery = ( : 5 * 60 * 1000; // 默认 5 分钟 const query = useQuery({ - queryKey: ["usage", providerId, appId], + queryKey: usageKeys.script(providerId, appId), queryFn: async () => usageApi.query(providerId, appId), enabled: enabled && !!providerId, refetchInterval: diff --git a/src/lib/query/subscription.ts b/src/lib/query/subscription.ts index f79dd1d6b..4262b55aa 100644 --- a/src/lib/query/subscription.ts +++ b/src/lib/query/subscription.ts @@ -7,13 +7,18 @@ import { PROVIDER_TYPES } from "@/config/constants"; const REFETCH_INTERVAL = 5 * 60 * 1000; // 5 minutes +export const subscriptionKeys = { + all: ["subscription"] as const, + quota: (appId: AppId) => [...subscriptionKeys.all, "quota", appId] as const, +}; + export function useSubscriptionQuota( appId: AppId, enabled: boolean, autoQuery = false, ) { return useQuery({ - queryKey: ["subscription", "quota", appId], + queryKey: subscriptionKeys.quota(appId), queryFn: () => subscriptionApi.getQuota(appId), enabled: enabled && ["claude", "codex", "gemini"].includes(appId), refetchInterval: autoQuery ? REFETCH_INTERVAL : false, diff --git a/src/lib/query/usage.ts b/src/lib/query/usage.ts index ed565508d..22cea395a 100644 --- a/src/lib/query/usage.ts +++ b/src/lib/query/usage.ts @@ -106,6 +106,8 @@ export const usageKeys = { pricing: () => [...usageKeys.all, "pricing"] as const, limits: (providerId: string, appType: string) => [...usageKeys.all, "limits", providerId, appType] as const, + script: (providerId: string, appType: string) => + [...usageKeys.all, providerId, appType] as const, }; // Hooks