feat: 新增 ZenMux Token Plan 供应商，支持手动凭证与 USD 额度富展示 (#2709)

* feat(Token plan): 增加 ZenMux 支持 * chore: format code with prettier * chore: format code with cargo fmt --------- Co-authored-by: 明桓 <jihaodong.jhd@oceanbase.com> Co-authored-by: Jason <farion1231@gmail.com>
2026-08-01 12:22:09 +08:00 · 2026-06-03 14:45:49 +08:00
parent 43ae1e5f2c
commit c1dff06625
11 changed files with 438 additions and 17 deletions
@@ -417,6 +417,42 @@ fn resolve_native_credentials(app_type: &AppType, provider: Option<&Provider>) -
        .unwrap_or_default()
 }

+fn resolve_coding_plan_credentials(
+    app_type: &AppType,
+    provider: Option<&Provider>,
+    usage_script: Option<&crate::provider::UsageScript>,
+) -> (String, String) {
+    let is_zenmux = usage_script
+        .and_then(|s| s.coding_plan_provider.as_deref())
+        .map(|provider| provider.eq_ignore_ascii_case("zenmux"))
+        .unwrap_or(false);
+
+    if !is_zenmux {
+        return resolve_native_credentials(app_type, provider);
+    }
+
+    let script_base_url = usage_script
+        .and_then(|s| s.base_url.as_deref())
+        .unwrap_or("")
+        .trim_end_matches('/')
+        .to_string();
+    let script_api_key = usage_script
+        .and_then(|s| s.api_key.as_deref())
+        .unwrap_or("")
+        .to_string();
+
+    if !script_base_url.is_empty() && !script_api_key.is_empty() {
+        return (script_base_url, script_api_key);
+    }
+
+    let native = resolve_native_credentials(app_type, provider);
+    if !native.0.is_empty() && !native.1.is_empty() {
+        native
+    } else {
+        (script_base_url, script_api_key)
+    }
+}
+
 async fn query_provider_usage_inner(
    state: &AppState,
    copilot_state: &CopilotAuthState,
@@ -474,8 +510,8 @@ async fn query_provider_usage_inner(

    // ── Coding Plan 专用路径 ──
    if template_type == TEMPLATE_TYPE_TOKEN_PLAN {
-        // 从供应商配置中提取 API Key 和 Base URL（按 app 区分存储格式）
-        let (base_url, api_key) = resolve_native_credentials(&app_type, provider);
+        let (base_url, api_key) =
+            resolve_coding_plan_credentials(&app_type, provider, usage_script);

        let quota = crate::services::coding_plan::get_coding_plan_quota(&base_url, &api_key)
            .await
@@ -490,6 +526,19 @@ async fn query_provider_usage_inner(
            });
        }

+        // ZenMux 的 tier 携带 USD 额度信息，需要编码为 JSON extra
+        let has_usd = quota
+            .tiers
+            .first()
+            .map(|t| t.used_value_usd.is_some())
+            .unwrap_or(false);
+        let plan_label = quota
+            .credential_message
+            .as_deref()
+            .and_then(|msg| msg.split(' ').next())
+            .map(|tier| format!("ZenMux·{}", tier.to_uppercase()));
+        let mut first_tier = true;
+
        let data: Vec<crate::provider::UsageData> = quota
            .tiers
            .iter()
@@ -497,6 +546,26 @@ async fn query_provider_usage_inner(
                let total = 100.0;
                let used = tier.utilization;
                let remaining = total - used;
+                let extra = if has_usd {
+                    let mut extra_json = serde_json::json!({
+                        "resetsAt": tier.resets_at,
+                    });
+                    if let Some(v) = tier.used_value_usd {
+                        extra_json["usedValueUsd"] = serde_json::json!(v);
+                    }
+                    if let Some(v) = tier.max_value_usd {
+                        extra_json["maxValueUsd"] = serde_json::json!(v);
+                    }
+                    if first_tier {
+                        if let Some(ref label) = plan_label {
+                            extra_json["planLabel"] = serde_json::json!(label);
+                        }
+                        first_tier = false;
+                    }
+                    Some(extra_json.to_string())
+                } else {
+                    tier.resets_at.clone()
+                };
                crate::provider::UsageData {
                    plan_name: Some(tier.name.clone()),
                    remaining: Some(remaining),
@@ -505,7 +574,7 @@ async fn query_provider_usage_inner(
                    unit: Some("%".to_string()),
                    is_valid: Some(true),
                    invalid_message: None,
-                    extra: tier.resets_at.clone(),
+                    extra,
                }
            })
            .collect();
@@ -950,11 +1019,31 @@ mod import_claude_desktop_tests {

 #[cfg(test)]
 mod native_query_credentials_tests {
-    use super::resolve_native_credentials;
+    use super::{resolve_coding_plan_credentials, resolve_native_credentials};
    use crate::app_config::AppType;
-    use crate::provider::Provider;
+    use crate::provider::{Provider, UsageScript};
    use serde_json::json;

+    fn usage_script(
+        coding_plan_provider: Option<&str>,
+        base_url: Option<&str>,
+        api_key: Option<&str>,
+    ) -> UsageScript {
+        UsageScript {
+            enabled: true,
+            language: "javascript".to_string(),
+            code: String::new(),
+            timeout: Some(10),
+            api_key: api_key.map(str::to_string),
+            base_url: base_url.map(str::to_string),
+            access_token: None,
+            user_id: None,
+            template_type: Some("token_plan".to_string()),
+            auto_query_interval: None,
+            coding_plan_provider: coding_plan_provider.map(str::to_string),
+        }
+    }
+
    #[test]
    fn delegates_to_provider_for_codex() {
        let provider = Provider::with_id(
@@ -979,4 +1068,52 @@ mod native_query_credentials_tests {
        assert!(base_url.is_empty());
        assert!(api_key.is_empty());
    }
+
+    #[test]
+    fn zenmux_coding_plan_uses_script_credentials_first() {
+        let provider = Provider::with_id(
+            "test".to_string(),
+            "Test".to_string(),
+            json!({
+                "env": {
+                    "ANTHROPIC_BASE_URL": "https://provider.zenmux.example/v1",
+                    "ANTHROPIC_AUTH_TOKEN": "sk-provider"
+                }
+            }),
+            None,
+        );
+        let script = usage_script(
+            Some("zenmux"),
+            Some("https://script.zenmux.example/api/usage/"),
+            Some("sk-script"),
+        );
+
+        let (base_url, api_key) =
+            resolve_coding_plan_credentials(&AppType::Claude, Some(&provider), Some(&script));
+
+        assert_eq!(base_url, "https://script.zenmux.example/api/usage");
+        assert_eq!(api_key, "sk-script");
+    }
+
+    #[test]
+    fn zenmux_coding_plan_falls_back_to_provider_credentials() {
+        let provider = Provider::with_id(
+            "test".to_string(),
+            "Test".to_string(),
+            json!({
+                "env": {
+                    "ANTHROPIC_BASE_URL": "https://provider.zenmux.example/v1",
+                    "ANTHROPIC_AUTH_TOKEN": "sk-provider"
+                }
+            }),
+            None,
+        );
+        let script = usage_script(Some("zenmux"), Some("https://script.zenmux.example"), None);
+
+        let (base_url, api_key) =
+            resolve_coding_plan_credentials(&AppType::Claude, Some(&provider), Some(&script));
+
+        assert_eq!(base_url, "https://provider.zenmux.example/v1");
+        assert_eq!(api_key, "sk-provider");
+    }
 }