From e7b4e14a66ea5b7b5dfa72cac7c753dd85ba1515 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:21:15 +0700 Subject: [PATCH 1/9] Add per-provider token account routing --- .../src-tauri/src/commands/providers.rs | 34 ++++- .../src-tauri/src/commands/tests.rs | 133 ++++++++++++++++ rust/src/cli/diagnose.rs | 2 + rust/src/cli/guard.rs | 2 + rust/src/cli/hooks.rs | 2 + rust/src/cli/serve/dashboard/source.rs | 4 + rust/src/cli/serve/data.rs | 2 + rust/src/cli/usage.rs | 2 + rust/src/cli/usage/fetch_helpers.rs | 76 +++++++++- rust/src/cli/usage_tests.rs | 82 +++++++++- rust/src/core/provider.rs | 9 ++ rust/src/core/token_accounts.rs | 142 +++++++++++++++++- rust/src/providers/doubao/mod.rs | 37 +++++ rust/src/providers/kimi/mod.rs | 23 ++- rust/src/providers/kimi/web.rs | 27 ++++ rust/src/providers/opencodego/mod.rs | 21 ++- rust/src/providers/opencodego/usage_api.rs | 23 +++ 17 files changed, 606 insertions(+), 15 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 779f9c801c..cd958bdb97 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -84,6 +84,7 @@ pub(crate) fn build_fetch_context( .and_then(|override_data| override_data.env_override.as_ref()); let active_token_api_key = active_token_env.and_then(|env| env.values().next().cloned()); let usage_source = SourceMode::parse(settings.usage_source(id)).unwrap_or_default(); + let token_account_kind = token_override.as_ref().map(|account| account.kind); // Selected token-account key overrides a stored provider apiKey (upstream #2271 / #1183). let api_key = active_token_api_key.or(stored_api_key); let has_kimi_code_api_key = @@ -220,16 +221,47 @@ pub(crate) fn build_fetch_context( // token account or manual cookie source scopes the session to web creds. let auto_prefer_web = token_override.is_some() || cookie_source == "manual"; + // These upstream account types are explicit identity selections. Keep the + // provider's saved region/source settings intact, but project the selected + // credential into the route required by that account. + let (source_mode, cookie_header, api_key) = match (id, token_account_kind) { + (ProviderId::Kimi, Some(_)) => (SourceMode::Web, active_token_cookie.clone(), None), + (ProviderId::Doubao, Some(_)) => (SourceMode::OAuth, None, active_token_api_key.clone()), + (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::ApiKey)) + if usage_source == SourceMode::Auto => + { + (SourceMode::Auto, None, active_token_api_key.clone()) + } + (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::ApiKey)) => { + (usage_source, cookie_header, api_key) + } + (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::Cookie)) + if usage_source == SourceMode::Auto => + { + (SourceMode::Web, active_token_cookie.clone(), api_key) + } + _ => (source_mode, cookie_header, api_key), + }; + let token_account_isolated = token_override.is_some() + && matches!( + id, + ProviderId::Kimi | ProviderId::Doubao | ProviderId::OpenCodeGo + ); + FetchContext { source_mode, manual_cookie_header: cookie_header, manual_cookie_missing: fails_closed_without_cookie, api_key, + token_account_kind, + token_account_isolated, workspace_id: (!workspace_id.is_empty()).then_some(workspace_id), seat_credit_entitlement: settings.seat_credit_entitlement(id), api_region: (!api_region.is_empty()).then_some(api_region), gateway_url, - auto_prefer_web, + auto_prefer_web: auto_prefer_web + && !(id == ProviderId::OpenCodeGo + && token_account_kind == Some(codexbar::core::TokenAccountKind::ApiKey)), ..FetchContext::default() } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 94ac348f11..8b6b43be7f 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -483,6 +483,139 @@ fn fetch_context_opencode_empty_manual_remaps_to_web() { assert_eq!(ctx.source_mode, SourceMode::Web); } +#[test] +fn kimi_selected_account_forces_web_and_keeps_saved_region() { + let mut settings = Settings::default(); + settings.set_usage_source(ProviderId::Kimi, "oauth"); + settings.set_api_region(ProviderId::Kimi, "international"); + let mut accounts = HashMap::new(); + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Work", "selected-kimi-session")); + accounts.insert(ProviderId::Kimi, data); + + let ctx = super::build_fetch_context( + ProviderId::Kimi, + &settings, + &ManualCookies::default(), + &ApiKeys::default(), + &accounts, + ); + + assert_eq!(ctx.source_mode, SourceMode::Web); + assert_eq!( + ctx.manual_cookie_header.as_deref(), + Some("kimi-auth=selected-kimi-session") + ); + assert_eq!(ctx.api_key, None); + assert_eq!(ctx.api_region.as_deref(), Some("international")); + assert!(ctx.token_account_isolated); + assert_eq!(settings.usage_source(ProviderId::Kimi), "oauth"); + assert_eq!(settings.api_region(ProviderId::Kimi), "international"); +} + +#[test] +fn doubao_selected_account_forces_ark_api_and_ignores_saved_source() { + let mut settings = Settings::default(); + settings.set_usage_source(ProviderId::Doubao, "cli"); + let mut accounts = HashMap::new(); + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Work", "selected-ark-key")); + accounts.insert(ProviderId::Doubao, data); + + let ctx = super::build_fetch_context( + ProviderId::Doubao, + &settings, + &ManualCookies::default(), + &ApiKeys::default(), + &accounts, + ); + + assert_eq!(ctx.source_mode, SourceMode::OAuth); + assert_eq!(ctx.api_key.as_deref(), Some("selected-ark-key")); + assert!(ctx.token_account_isolated); +} + +#[test] +fn opencodego_selected_api_account_overrides_global_key_without_changing_explicit_source() { + let mut settings = Settings::default(); + settings.set_usage_source(ProviderId::OpenCodeGo, "auto"); + let mut keys = ApiKeys::default(); + keys.set("opencodego", "global-key", None); + let mut accounts = HashMap::new(); + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Work", "selected-account-key")); + accounts.insert(ProviderId::OpenCodeGo, data); + + let ctx = super::build_fetch_context( + ProviderId::OpenCodeGo, + &settings, + &ManualCookies::default(), + &keys, + &accounts, + ); + + assert_eq!(ctx.source_mode, SourceMode::Auto); + assert_eq!(ctx.api_key.as_deref(), Some("selected-account-key")); + assert!(!ctx.auto_prefer_web); + assert!(ctx.token_account_isolated); + + for cookie_source in ["off", "manual"] { + settings.set_cookie_source(ProviderId::OpenCodeGo, cookie_source); + settings.set_usage_source(ProviderId::OpenCodeGo, "auto"); + let auto_ctx = super::build_fetch_context( + ProviderId::OpenCodeGo, + &settings, + &ManualCookies::default(), + &keys, + &accounts, + ); + assert_eq!(auto_ctx.source_mode, SourceMode::Auto); + assert_eq!(auto_ctx.api_key.as_deref(), Some("selected-account-key")); + assert!(auto_ctx.manual_cookie_header.is_none()); + } + + for (saved_source, expected_source) in [("web", SourceMode::Web), ("cli", SourceMode::Cli)] { + settings.set_cookie_source(ProviderId::OpenCodeGo, "off"); + settings.set_usage_source(ProviderId::OpenCodeGo, saved_source); + let explicit_ctx = super::build_fetch_context( + ProviderId::OpenCodeGo, + &settings, + &ManualCookies::default(), + &keys, + &accounts, + ); + assert_eq!(explicit_ctx.source_mode, expected_source); + } +} + +#[test] +fn opencodego_selected_cookie_account_uses_web_route() { + let settings = Settings::default(); + let mut accounts = HashMap::new(); + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Web", "Cookie: session=selected-session")); + accounts.insert(ProviderId::OpenCodeGo, data); + + let ctx = super::build_fetch_context( + ProviderId::OpenCodeGo, + &settings, + &ManualCookies::default(), + &ApiKeys::default(), + &accounts, + ); + + assert_eq!(ctx.source_mode, SourceMode::Web); + assert_eq!( + ctx.manual_cookie_header.as_deref(), + Some("Cookie: session=selected-session") + ); + assert_eq!( + ctx.token_account_kind, + Some(codexbar::core::TokenAccountKind::Cookie) + ); + assert!(ctx.token_account_isolated); +} + #[test] fn fetch_context_replicate_empty_manual_fails_closed_without_browser_import() { let settings = Settings::default(); diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 64d0b54d59..1d09877414 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -177,6 +177,8 @@ async fn collect_provider_diagnostic( .map(ToOwned::to_owned), manual_cookie_missing: false, api_key: api_keys.get(provider_id.cli_name()).map(ToOwned::to_owned), + token_account_kind: None, + token_account_isolated: false, workspace_id: settings .provider_config(provider_id) .and_then(|config| config.workspace_id.clone()), diff --git a/rust/src/cli/guard.rs b/rust/src/cli/guard.rs index b9ecf5dce8..b56945db1b 100644 --- a/rust/src/cli/guard.rs +++ b/rust/src/cli/guard.rs @@ -318,6 +318,8 @@ async fn fetch_guard_outcome( manual_cookie_header: None, manual_cookie_missing: false, api_key: None, + token_account_kind: None, + token_account_isolated: false, workspace_id: None, seat_credit_entitlement: None, api_region: None, diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index 2f26498cf1..c7b211f7c0 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -296,6 +296,8 @@ async fn hooks_watch_observation( manual_cookie_header: None, manual_cookie_missing: false, api_key: None, + token_account_kind: None, + token_account_isolated: false, workspace_id: (!workspace.is_empty()).then(|| workspace.to_string()), seat_credit_entitlement: settings.seat_credit_entitlement(provider_id), api_region: (!region.is_empty()).then(|| region.to_string()), diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index 8d28758df4..af7941d386 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -159,6 +159,8 @@ async fn fetch_provider_envelope( manual_cookie_header: None, manual_cookie_missing: false, api_key: None, + token_account_kind: None, + token_account_isolated: false, workspace_id: None, seat_credit_entitlement: None, api_region: None, @@ -292,6 +294,8 @@ async fn collect_claude_accounts(claude_enabled: bool) -> Option) -> String { manual_cookie_header: None, manual_cookie_missing: false, api_key: None, + token_account_kind: None, + token_account_isolated: false, workspace_id: None, seat_credit_entitlement: None, api_region: None, diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 2b1dfff0af..1a588c1c75 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -249,6 +249,8 @@ fn build_usage_fetch_context(args: &UsageArgs, source_mode: SourceMode) -> Fetch manual_cookie_header: None, manual_cookie_missing: false, api_key: None, + token_account_kind: None, + token_account_isolated: false, workspace_id: None, seat_credit_entitlement: None, api_region: None, diff --git a/rust/src/cli/usage/fetch_helpers.rs b/rust/src/cli/usage/fetch_helpers.rs index d093a2bbe7..e9e063f296 100644 --- a/rust/src/cli/usage/fetch_helpers.rs +++ b/rust/src/cli/usage/fetch_helpers.rs @@ -5,7 +5,8 @@ use super::render::{ render_brief_text, render_json_result, render_text_error, render_text_with_status, }; use crate::core::{ - ProviderFetchResult, ProviderId, TokenAccountStore, TokenAccountSupport, instantiate_provider, + ProviderFetchResult, ProviderId, SourceMode, TokenAccountKind, TokenAccountOverride, + TokenAccountStore, TokenAccountSupport, instantiate_provider, }; use crate::settings::ApiKeys; use crate::status::{ProviderStatus as StatusInfo, fetch_provider_status}; @@ -45,7 +46,8 @@ pub async fn fetch_provider_result( .fetch_status .then(|| fetch_provider_status(provider_id.cli_name())); let mut ctx = command.ctx.clone(); - if ctx.api_key.is_none() { + let account_projected = project_cli_account(provider_id, command.account.as_deref(), &mut ctx)?; + if !account_projected && ctx.api_key.is_none() { ctx.api_key = resolve_cli_api_key(provider_id, command.account.as_deref())?; } let result = provider.fetch_usage(&ctx).await?; @@ -57,6 +59,76 @@ pub async fn fetch_provider_result( Ok((result, status)) } +/// Apply the selected labeled account to the same route used by the desktop +/// shell. The three v0.65 account-source ports require provider-specific route +/// selection in addition to the shared credential injection. +fn project_cli_account( + provider: ProviderId, + account_ref: Option<&str>, + ctx: &mut crate::core::FetchContext, +) -> anyhow::Result { + if !matches!( + provider, + ProviderId::Kimi | ProviderId::Doubao | ProviderId::OpenCodeGo + ) { + return Ok(false); + } + let store = TokenAccountStore::new(); + let data = match store.load_provider(provider) { + Ok(data) => data, + Err(error) if account_ref.is_some() => { + return Err(error.into()); + } + Err(_) => return Ok(false), + }; + if data.accounts.is_empty() { + if account_ref.is_some() { + anyhow::bail!( + "No token accounts configured for {}", + provider.display_name() + ); + } + return Ok(false); + } + let account = if let Some(account_ref) = account_ref { + find_token_account(&data, account_ref)? + } else { + data.active_account().ok_or_else(|| { + anyhow::anyhow!("No active token account for {}", provider.display_name()) + })? + } + .clone(); + project_token_account(provider, &account, ctx); + Ok(true) +} + +pub(super) fn project_token_account( + provider: ProviderId, + account: &crate::core::TokenAccount, + ctx: &mut crate::core::FetchContext, +) { + let projected = TokenAccountOverride::from_account(provider, account.clone()); + ctx.token_account_kind = Some(projected.kind); + ctx.token_account_isolated = true; + ctx.api_key = projected + .env_override + .as_ref() + .and_then(|env| env.values().next().cloned()); + ctx.manual_cookie_header = projected.cookie_header; + ctx.auto_prefer_web = projected.kind == TokenAccountKind::Cookie; + + match (provider, projected.kind) { + (ProviderId::Kimi, _) => ctx.source_mode = SourceMode::Web, + (ProviderId::Doubao, _) => ctx.source_mode = SourceMode::OAuth, + (ProviderId::OpenCodeGo, TokenAccountKind::Cookie) + if ctx.source_mode == SourceMode::Auto => + { + ctx.source_mode = SourceMode::Web; + } + _ => {} + } +} + /// Resolve an API key from token accounts (active or `--account`) then stored keys. /// /// Token-account env injection takes precedence over `api_keys.json` so multi-key diff --git a/rust/src/cli/usage_tests.rs b/rust/src/cli/usage_tests.rs index d332131b81..7339b66a43 100644 --- a/rust/src/cli/usage_tests.rs +++ b/rust/src/cli/usage_tests.rs @@ -2,8 +2,9 @@ use super::*; use crate::core::{ - CostSnapshot, ProviderAccountData, ProviderDisplayDetail, ProviderInventoryItem, RateWindow, - TokenAccount, TokenAccountSupport, UsageSnapshot, + CostSnapshot, FetchContext, ProviderAccountData, ProviderDisplayDetail, ProviderId, + ProviderInventoryItem, RateWindow, SourceMode, TokenAccount, TokenAccountKind, + TokenAccountSupport, UsageSnapshot, }; use crate::providers::claude::claude_swap::ClaudeSwapAccount; use crate::status::{ProviderStatus as StatusInfo, StatusLevel}; @@ -175,6 +176,83 @@ fn openrouter_account_ref_resolves_labeled_key() { assert_eq!(by_index.token, "sk-or-v1-work"); } +#[test] +fn kimi_account_projection_forces_isolated_web_and_preserves_region() { + let account = TokenAccount::new("work", "selected-kimi-auth"); + let mut ctx = FetchContext { + source_mode: SourceMode::OAuth, + api_region: Some("international".into()), + api_key: Some("ambient-api-key".into()), + ..FetchContext::default() + }; + + super::fetch_helpers::project_token_account(ProviderId::Kimi, &account, &mut ctx); + + assert_eq!(ctx.source_mode, SourceMode::Web); + assert_eq!( + ctx.manual_cookie_header.as_deref(), + Some("kimi-auth=selected-kimi-auth") + ); + assert_eq!(ctx.api_key, None); + assert_eq!(ctx.api_region.as_deref(), Some("international")); + assert!(ctx.token_account_isolated); +} + +#[test] +fn doubao_account_projection_uses_only_the_selected_ark_key() { + let account = TokenAccount::new("work", "selected-ark-key"); + let mut ctx = FetchContext { + source_mode: SourceMode::Cli, + api_key: Some("ambient-key".into()), + ..FetchContext::default() + }; + + super::fetch_helpers::project_token_account(ProviderId::Doubao, &account, &mut ctx); + + assert_eq!(ctx.source_mode, SourceMode::OAuth); + assert_eq!(ctx.api_key.as_deref(), Some("selected-ark-key")); + assert_eq!(ctx.token_account_kind, Some(TokenAccountKind::ApiKey)); + assert!(ctx.token_account_isolated); +} + +#[test] +fn opencodego_account_projection_distinguishes_api_and_cookie_routes() { + let mut api_ctx = FetchContext::default(); + super::fetch_helpers::project_token_account( + ProviderId::OpenCodeGo, + &TokenAccount::new("api", "selected-opencode-key"), + &mut api_ctx, + ); + assert_eq!(api_ctx.source_mode, SourceMode::Auto); + assert_eq!(api_ctx.api_key.as_deref(), Some("selected-opencode-key")); + assert_eq!(api_ctx.token_account_kind, Some(TokenAccountKind::ApiKey)); + assert!(!api_ctx.auto_prefer_web); + + let mut cookie_ctx = FetchContext::default(); + super::fetch_helpers::project_token_account( + ProviderId::OpenCodeGo, + &TokenAccount::new("web", "Cookie: session=selected-session"), + &mut cookie_ctx, + ); + assert_eq!(cookie_ctx.source_mode, SourceMode::Web); + assert_eq!( + cookie_ctx.manual_cookie_header.as_deref(), + Some("Cookie: session=selected-session") + ); + assert_eq!( + cookie_ctx.token_account_kind, + Some(TokenAccountKind::Cookie) + ); + + cookie_ctx.source_mode = SourceMode::Cli; + super::fetch_helpers::project_token_account( + ProviderId::OpenCodeGo, + &TokenAccount::new("api", "another-key"), + &mut cookie_ctx, + ); + assert_eq!(cookie_ctx.source_mode, SourceMode::Cli); +} + #[test] fn text_rendering_shows_sub_one_percent_usage() { let result = fetch_result(UsageSnapshot::new(RateWindow::new(0.4))); diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 868b66fde6..2dc59a177f 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -766,6 +766,13 @@ pub struct FetchContext { /// API key for providers that require authentication pub api_key: Option, + /// Type of the explicitly selected labeled token account, if any. + pub token_account_kind: Option, + + /// A selected account is an identity boundary: providers must not retry + /// another ambient credential or account after its credential fails. + pub token_account_isolated: bool, + /// Optional provider workspace/project scope from persisted settings. pub workspace_id: Option, @@ -801,6 +808,8 @@ impl Default for FetchContext { manual_cookie_header: None, manual_cookie_missing: false, api_key: None, + token_account_kind: None, + token_account_isolated: false, workspace_id: None, seat_credit_entitlement: None, api_region: None, diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index 2b86ad8181..abd9cd2b86 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -20,6 +20,15 @@ pub enum TokenInjection { CookieHeader, /// Inject as environment variable Environment { key: String }, + /// Accept either an API key or a Cookie header, as with OpenCode Go. + EnvironmentOrCookie { key: String }, +} + +/// Credential route selected by a labeled account. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TokenAccountKind { + Cookie, + ApiKey, } /// Support definition for a provider's token accounts @@ -330,6 +339,34 @@ impl TokenAccountSupport { requires_manual_cookie_source: false, cookie_name: None, }), + ProviderId::Kimi => Some(TokenAccountSupport { + title: "Web sessions", + subtitle: "Store labeled Kimi kimi-auth web sessions.", + placeholder: "kimi-auth value or Cookie: kimi-auth=...", + injection: TokenInjection::CookieHeader, + requires_manual_cookie_source: true, + cookie_name: Some("kimi-auth"), + }), + ProviderId::Doubao => Some(TokenAccountSupport { + title: "Ark API keys", + subtitle: "Store labeled Volcengine Ark API keys.", + placeholder: "Ark API key", + injection: TokenInjection::Environment { + key: "ARK_API_KEY".to_string(), + }, + requires_manual_cookie_source: false, + cookie_name: None, + }), + ProviderId::OpenCodeGo => Some(TokenAccountSupport { + title: "API keys or sessions", + subtitle: "Store labeled OpenCode Go API keys or Cookie headers.", + placeholder: "API key or Cookie: ...", + injection: TokenInjection::EnvironmentOrCookie { + key: "OPENCODE_API_KEY".to_string(), + }, + requires_manual_cookie_source: false, + cookie_name: None, + }), // These providers don't support token accounts ProviderId::Codex | ProviderId::Pi @@ -337,7 +374,6 @@ impl TokenAccountSupport { | ProviderId::Antigravity | ProviderId::Kiro | ProviderId::VertexAI - | ProviderId::Kimi | ProviderId::KimiK2 | ProviderId::JetBrains | ProviderId::Warp @@ -346,14 +382,12 @@ impl TokenAccountSupport { | ProviderId::Infini | ProviderId::Perplexity | ProviderId::Abacus - | ProviderId::OpenCodeGo | ProviderId::Kilo | ProviderId::Bedrock | ProviderId::Codebuff | ProviderId::CodeRabbit | ProviderId::DeepSeek | ProviderId::Windsurf - | ProviderId::Doubao | ProviderId::StepFun | ProviderId::Venice | ProviderId::OpenAIApi @@ -404,6 +438,12 @@ impl TokenAccountSupport { map.insert(key.clone(), token.to_string()); Some(map) } + TokenInjection::EnvironmentOrCookie { key } => { + let api_key = Self::normalized_opencodego_api_key(token)?; + let mut map = HashMap::new(); + map.insert(key.clone(), api_key); + Some(map) + } TokenInjection::CookieHeader => { // Check for Claude OAuth token if provider == ProviderId::Claude @@ -419,6 +459,43 @@ impl TokenAccountSupport { } } + fn normalized_opencodego_api_key(token: &str) -> Option { + let token = token.trim(); + let token = if token.len() >= 2 + && ((token.starts_with('"') && token.ends_with('"')) + || (token.starts_with('\'') && token.ends_with('\''))) + { + token[1..token.len() - 1].trim() + } else { + token + }; + if token.is_empty() + || token + .chars() + .any(|ch| ch.is_whitespace() || matches!(ch, '=' | ':')) + { + return None; + } + Some(token.to_string()) + } + + pub fn account_kind(provider: ProviderId, token: &str) -> TokenAccountKind { + if provider == ProviderId::OpenCodeGo { + if Self::normalized_opencodego_api_key(token).is_some() { + TokenAccountKind::ApiKey + } else { + TokenAccountKind::Cookie + } + } else if matches!( + Self::for_provider(provider).map(|support| support.injection), + Some(TokenInjection::Environment { .. }) + ) { + TokenAccountKind::ApiKey + } else { + TokenAccountKind::Cookie + } + } + /// Normalize a cookie header for a provider pub fn normalized_cookie_header(provider: ProviderId, token: &str) -> String { let trimmed = token.trim(); @@ -758,11 +835,13 @@ pub struct TokenAccountOverride { pub env_override: Option>, /// Cookie header to use pub cookie_header: Option, + pub kind: TokenAccountKind, } impl TokenAccountOverride { /// Create an override from an account pub fn from_account(provider: ProviderId, account: TokenAccount) -> Self { + let kind = TokenAccountSupport::account_kind(provider, &account.token); let env_override = TokenAccountSupport::env_override(provider, &account.token); let cookie_header = if env_override.is_none() { Some(TokenAccountSupport::normalized_cookie_header( @@ -778,6 +857,7 @@ impl TokenAccountOverride { account, env_override, cookie_header, + kind, } } } @@ -796,6 +876,9 @@ mod tests { assert!(TokenAccountSupport::is_supported(ProviderId::Copilot)); assert!(TokenAccountSupport::is_supported(ProviderId::OpenRouter)); assert!(TokenAccountSupport::is_supported(ProviderId::Grok)); + assert!(TokenAccountSupport::is_supported(ProviderId::Kimi)); + assert!(TokenAccountSupport::is_supported(ProviderId::Doubao)); + assert!(TokenAccountSupport::is_supported(ProviderId::OpenCodeGo)); assert!(!TokenAccountSupport::is_supported(ProviderId::Codex)); assert!(!TokenAccountSupport::is_supported(ProviderId::Gemini)); assert!(!TokenAccountSupport::is_supported(ProviderId::Hyper)); @@ -803,6 +886,59 @@ mod tests { assert!(!TokenAccountSupport::is_supported(ProviderId::Bifrost)); } + #[test] + fn upstream_account_sources_normalize_and_classify_selected_credentials() { + assert_eq!( + TokenAccountSupport::normalized_cookie_header(ProviderId::Kimi, "selected-session"), + "kimi-auth=selected-session" + ); + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::Kimi, "selected-session"), + TokenAccountKind::Cookie + ); + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::Doubao, "ark-key"), + TokenAccountKind::ApiKey + ); + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::OpenCodeGo, "opencode-key"), + TokenAccountKind::ApiKey + ); + assert_eq!( + TokenAccountSupport::account_kind( + ProviderId::OpenCodeGo, + "Cookie: session=opencode-session" + ), + TokenAccountKind::Cookie + ); + assert_eq!( + TokenAccountSupport::env_override(ProviderId::OpenCodeGo, "opencode-key") + .and_then(|env| env.get("OPENCODE_API_KEY").cloned()) + .as_deref(), + Some("opencode-key") + ); + assert!( + TokenAccountSupport::env_override( + ProviderId::OpenCodeGo, + "Cookie: session=opencode-session" + ) + .is_none() + ); + for malformed in ["", " ", "Cookie: broken", "auth=fixture", "two words"] { + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::OpenCodeGo, malformed), + TokenAccountKind::Cookie + ); + assert!(TokenAccountSupport::env_override(ProviderId::OpenCodeGo, malformed).is_none()); + } + assert_eq!( + TokenAccountSupport::env_override(ProviderId::OpenCodeGo, " 'go_key' ") + .and_then(|env| env.get("OPENCODE_API_KEY").cloned()) + .as_deref(), + Some("go_key") + ); + } + #[test] fn grok_token_accounts_route_bearer_and_cookie_credentials() { let bearer = diff --git a/rust/src/providers/doubao/mod.rs b/rust/src/providers/doubao/mod.rs index b11e6e34ad..e094e62188 100644 --- a/rust/src/providers/doubao/mod.rs +++ b/rust/src/providers/doubao/mod.rs @@ -888,6 +888,15 @@ impl Provider for DoubaoProvider { } async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + if ctx.token_account_isolated + && ctx.token_account_kind == Some(crate::core::TokenAccountKind::ApiKey) + { + let api_key = selected_ark_api_key(ctx)?; + return Ok(ProviderFetchResult::new( + self.fetch_api(&api_key).await?, + "api", + )); + } match ctx.source_mode { SourceMode::Auto | SourceMode::OAuth => { if let Some(credentials) = Self::coding_plan_credentials(ctx.api_key.as_deref()) { @@ -945,6 +954,15 @@ impl Provider for DoubaoProvider { } } +fn selected_ark_api_key(ctx: &FetchContext) -> Result { + ctx.api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .ok_or(ProviderError::AuthRequired) +} + fn resolve_api_key( explicit: Option<&str>, credential_target: &str, @@ -977,6 +995,25 @@ fn resolve_api_key( #[cfg(test)] mod tests { use super::*; + + #[test] + fn selected_ark_account_requires_its_projected_key() { + let isolated = FetchContext { + token_account_isolated: true, + token_account_kind: Some(crate::core::TokenAccountKind::ApiKey), + ..FetchContext::default() + }; + assert!(matches!( + selected_ark_api_key(&isolated), + Err(ProviderError::AuthRequired) + )); + + let selected = FetchContext { + api_key: Some(" selected-key ".into()), + ..isolated + }; + assert_eq!(selected_ark_api_key(&selected).unwrap(), "selected-key"); + } use reqwest::header::{HeaderMap, HeaderValue}; #[test] diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index e326afa3d6..70a704bf2f 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -263,7 +263,14 @@ impl KimiProvider { } fn auth_token_from_cookie_header(cookie_header: &str) -> Result { - for cookie in cookie_header.split(';') { + let header = cookie_header.trim(); + let header = header + .get(..7) + .filter(|prefix| prefix.eq_ignore_ascii_case("cookie:")) + .map(|_| &header[7..]) + .unwrap_or(header) + .trim(); + for cookie in header.split(';') { let cookie = cookie.trim(); if cookie.starts_with("kimi-auth=") || cookie.starts_with("authorization=") @@ -375,7 +382,12 @@ impl Provider for KimiProvider { } } - let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref(), region).await?; + let usage = web::fetch_via_web( + ctx.manual_cookie_header.as_deref(), + region, + ctx.token_account_isolated, + ) + .await?; Ok(ProviderFetchResult::new(usage, "web")) } SourceMode::OAuth => { @@ -384,7 +396,12 @@ impl Provider for KimiProvider { Ok(ProviderFetchResult::new(usage, "code-api")) } SourceMode::Web => { - let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref(), region).await?; + let usage = web::fetch_via_web( + ctx.manual_cookie_header.as_deref(), + region, + ctx.token_account_isolated, + ) + .await?; Ok(ProviderFetchResult::new(usage, "web")) } SourceMode::Cli => Err(ProviderError::UnsupportedSource(SourceMode::Cli)), diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index c5ed5b8305..b18644640a 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -132,7 +132,12 @@ fn browser_auth_token(region: KimiRegion) -> Option { pub(crate) async fn fetch_via_web( cookie_header: Option<&str>, region: KimiRegion, + account_isolated: bool, ) -> Result { + if account_isolated { + let token = selected_account_auth_token(cookie_header)?; + return fetch_via_web_token(&client()?, &token, region).await; + } let source = cookie_source(); if let Some(token) = cookie_header.and_then(|header| KimiProvider::auth_token_from_cookie_header(header).ok()) @@ -176,6 +181,12 @@ pub(crate) async fn fetch_via_web( Err(ProviderError::AuthRequired) } +fn selected_account_auth_token(cookie_header: Option<&str>) -> Result { + cookie_header + .and_then(|header| KimiProvider::auth_token_from_cookie_header(header).ok()) + .ok_or(ProviderError::AuthRequired) +} + fn client() -> Result { crate::core::credentialed_http_client_builder() .timeout(std::time::Duration::from_secs(30)) @@ -327,6 +338,22 @@ pub(super) async fn fetch_subscription_for_enrichment_result( mod tests { use super::*; + #[test] + fn selected_session_rejects_missing_or_invalid_cookie_without_fallback() { + assert!(matches!( + selected_account_auth_token(None), + Err(ProviderError::AuthRequired) + )); + assert!(matches!( + selected_account_auth_token(Some("locale=en-US")), + Err(ProviderError::AuthRequired) + )); + assert_eq!( + selected_account_auth_token(Some("Cookie: kimi-auth=selected")).unwrap(), + "selected" + ); + } + fn static_desktop(_: KimiRegion) -> Option { Some("desktop-token".to_string()) } diff --git a/rust/src/providers/opencodego/mod.rs b/rust/src/providers/opencodego/mod.rs index d676e3dd94..29a39bfea6 100644 --- a/rust/src/providers/opencodego/mod.rs +++ b/rust/src/providers/opencodego/mod.rs @@ -477,6 +477,15 @@ impl Provider for OpenCodeGoProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { tracing::debug!("Fetching OpenCode Go usage"); + if ctx.token_account_isolated + && ctx.token_account_kind == Some(crate::core::TokenAccountKind::ApiKey) + && ctx.source_mode == SourceMode::Auto + { + let api_key = + usage_api::selected_account_api_key(ctx).ok_or(ProviderError::AuthRequired)?; + return usage_api::fetch(&self.client, ctx, &api_key, "api").await; + } + match ctx.source_mode { SourceMode::Auto => { // Local-first unless workspace/token scope asks for web first @@ -513,12 +522,16 @@ impl Provider for OpenCodeGoProvider { SourceMode::Web => self.fetch_web(ctx).await, SourceMode::Cli => self.fetch_local_with_balance(ctx).await, SourceMode::OAuth => { - let api_key = usage_api::resolve_api_key(ctx).ok_or_else(|| { - ProviderError::NotInstalled( + let api_key = if ctx.token_account_isolated { + usage_api::selected_account_api_key(ctx).ok_or(ProviderError::AuthRequired)? + } else { + usage_api::resolve_api_key(ctx).ok_or_else(|| { + ProviderError::NotInstalled( "Missing OpenCode Go API key. Add one in Settings or set OPENCODE_API_KEY." .to_string(), - ) - })?; + ) + })? + }; usage_api::fetch(&self.client, ctx, &api_key, "api").await } } diff --git a/rust/src/providers/opencodego/usage_api.rs b/rust/src/providers/opencodego/usage_api.rs index 00c5b59a2b..13427919d9 100644 --- a/rust/src/providers/opencodego/usage_api.rs +++ b/rust/src/providers/opencodego/usage_api.rs @@ -24,6 +24,10 @@ pub(super) fn resolve_api_key(ctx: &FetchContext) -> Option { }) } +pub(super) fn selected_account_api_key(ctx: &FetchContext) -> Option { + normalized_api_key(ctx.api_key.as_deref()) +} + pub(super) async fn fetch( client: &Client, ctx: &FetchContext, @@ -147,6 +151,25 @@ fn api_window( mod tests { use super::*; + #[test] + fn selected_account_key_does_not_consult_global_environment() { + let ctx = FetchContext { + token_account_isolated: true, + token_account_kind: Some(crate::core::TokenAccountKind::ApiKey), + api_key: Some(" selected-account-key ".into()), + ..FetchContext::default() + }; + assert_eq!( + selected_account_api_key(&ctx).as_deref(), + Some("selected-account-key") + ); + let missing = FetchContext { + api_key: None, + ..ctx + }; + assert_eq!(selected_account_api_key(&missing), None); + } + #[test] fn api_key_normalization_matches_upstream_settings_reader() { assert_eq!( From 6878cc01873b173d5bcd876e630062129161e17b Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:48:06 +0700 Subject: [PATCH 2/9] Fail closed for mismatched OpenCode Go account sources --- rust/src/providers/opencodego/mod.rs | 24 ++++++++++++++++++++---- rust/src/providers/opencodego/tests.rs | 26 ++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/rust/src/providers/opencodego/mod.rs b/rust/src/providers/opencodego/mod.rs index 29a39bfea6..b4f922b065 100644 --- a/rust/src/providers/opencodego/mod.rs +++ b/rust/src/providers/opencodego/mod.rs @@ -477,10 +477,7 @@ impl Provider for OpenCodeGoProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { tracing::debug!("Fetching OpenCode Go usage"); - if ctx.token_account_isolated - && ctx.token_account_kind == Some(crate::core::TokenAccountKind::ApiKey) - && ctx.source_mode == SourceMode::Auto - { + if selected_api_account_requires_api_route(ctx)? { let api_key = usage_api::selected_account_api_key(ctx).ok_or(ProviderError::AuthRequired)?; return usage_api::fetch(&self.client, ctx, &api_key, "api").await; @@ -550,6 +547,25 @@ impl Provider for OpenCodeGoProvider { } } +/// Keep a selected API-key account on its own identity. Explicit web and local +/// sources cannot represent that account, so reject them instead of fetching +/// browser or device-wide data under the selected account's label. +fn selected_api_account_requires_api_route(ctx: &FetchContext) -> Result { + if !ctx.token_account_isolated + || ctx.token_account_kind != Some(crate::core::TokenAccountKind::ApiKey) + { + return Ok(false); + } + + match ctx.source_mode { + SourceMode::Auto | SourceMode::OAuth => Ok(true), + SourceMode::Web | SourceMode::Cli => Err(ProviderError::Other(format!( + "Selected OpenCode Go API-key account is incompatible with explicit {:?} source", + ctx.source_mode + ))), + } +} + impl OpenCodeGoProvider { /// Auto prefers web when a workspace override or active token-account scope /// is present (upstream `requiresScopedWebStrategy`). diff --git a/rust/src/providers/opencodego/tests.rs b/rust/src/providers/opencodego/tests.rs index 880fd87e2b..05b2e14a0d 100644 --- a/rust/src/providers/opencodego/tests.rs +++ b/rust/src/providers/opencodego/tests.rs @@ -4,6 +4,32 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, }; +#[test] +fn selected_api_account_uses_api_in_auto_and_rejects_explicit_other_sources() { + let auto = FetchContext { + source_mode: SourceMode::Auto, + token_account_kind: Some(crate::core::TokenAccountKind::ApiKey), + token_account_isolated: true, + api_key: Some("selected-key".into()), + ..FetchContext::default() + }; + assert!(selected_api_account_requires_api_route(&auto).unwrap()); + assert_eq!( + usage_api::selected_account_api_key(&auto).as_deref(), + Some("selected-key") + ); + + for source in [SourceMode::Web, SourceMode::Cli] { + let context = FetchContext { + source_mode: source, + ..auto.clone() + }; + let error = selected_api_account_requires_api_route(&context).unwrap_err(); + assert!(matches!(error, ProviderError::Other(_))); + assert!(error.to_string().contains(&format!("{source:?}"))); + } +} + #[derive(Clone, Debug, Eq, PartialEq)] struct FakeLegacySession { workspace_id: String, From 220c80409e80af7a00e9e701011abc4640a3c535 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Fri, 25 Sep 2026 11:03:14 +0700 Subject: [PATCH 3/9] Preserve token account credentials in fetch routing --- .../src-tauri/src/commands/providers.rs | 165 +++++++++--------- 1 file changed, 83 insertions(+), 82 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index cd958bdb97..f0cf85e8c2 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -86,97 +86,98 @@ pub(crate) fn build_fetch_context( let usage_source = SourceMode::parse(settings.usage_source(id)).unwrap_or_default(); let token_account_kind = token_override.as_ref().map(|account| account.kind); // Selected token-account key overrides a stored provider apiKey (upstream #2271 / #1183). - let api_key = active_token_api_key.or(stored_api_key); + let api_key = active_token_api_key.clone().or(stored_api_key); let has_kimi_code_api_key = id == ProviderId::Kimi && api_key.as_deref().is_some_and(|key| !key.trim().is_empty()); let has_opencodego_api_key = id == ProviderId::OpenCodeGo && api_key.as_deref().is_some_and(|key| !key.trim().is_empty()); - let (mut source_mode, mut cookie_header, fails_closed_without_cookie) = - if id.cookie_domain().is_none() { - let source_mode = if active_token_env.is_some() { - SourceMode::OAuth - } else { - usage_source - }; - (source_mode, None, false) + let (mut source_mode, mut cookie_header, fails_closed_without_cookie) = if id + .cookie_domain() + .is_none() + { + let source_mode = if active_token_env.is_some() { + SourceMode::OAuth } else { - match cookie_source { - // #433: an explicitly selected, non-empty Claude manual cookie is - // authoritative. Do not let an active OAuth token account silently - // replace it; this keeps tray refresh behavior aligned with diagnose, - // whose Claude Auto path tries the supplied Web cookie before OAuth. - "manual" - if provider.manual_cookie_precedes_token_account() - && stored_cookie - .as_deref() - .is_some_and(|cookie| !cookie.trim().is_empty()) => - { - (SourceMode::Web, stored_cookie.clone(), false) - } - _ if active_token_env.is_some() => (SourceMode::OAuth, None, false), - "off" if provider_uses_oauth_without_cookies(id, usage_source) => { - (SourceMode::OAuth, None, false) - } - "off" - if (has_kimi_code_api_key || has_opencodego_api_key) - && usage_source == SourceMode::Auto => + usage_source + }; + (source_mode, None, false) + } else { + match cookie_source { + // #433: an explicitly selected, non-empty Claude manual cookie is + // authoritative. Do not let an active OAuth token account silently + // replace it; this keeps tray refresh behavior aligned with diagnose, + // whose Claude Auto path tries the supplied Web cookie before OAuth. + "manual" + if provider.manual_cookie_precedes_token_account() + && stored_cookie + .as_deref() + .is_some_and(|cookie| !cookie.trim().is_empty()) => + { + (SourceMode::Web, stored_cookie.clone(), false) + } + _ if active_token_env.is_some() => (SourceMode::OAuth, None, false), + "off" if provider_uses_oauth_without_cookies(id, usage_source) => { + (SourceMode::OAuth, None, false) + } + "off" + if (has_kimi_code_api_key || has_opencodego_api_key) + && usage_source == SourceMode::Auto => + { + (SourceMode::Auto, None, false) + } + // Droid/Factory: cookie-off must never scrape browser cookies. Map to + // Cli (API-only in the provider) so Auto does not fall through to web. + "off" if id == ProviderId::Factory => (SourceMode::Cli, None, false), + "off" => (SourceMode::Cli, None, false), + "manual" => { + let cookie_header = active_token_cookie.clone().or(stored_cookie); + let fails_closed_without_cookie = cookie_header.is_none() + && provider.manual_empty_cookie_policy() + == ManualEmptyCookiePolicy::FailClosedWeb; + let source_mode = if (has_kimi_code_api_key || has_opencodego_api_key) + && usage_source == SourceMode::Auto { - (SourceMode::Auto, None, false) - } - // Droid/Factory: cookie-off must never scrape browser cookies. Map to - // Cli (API-only in the provider) so Auto does not fall through to web. - "off" if id == ProviderId::Factory => (SourceMode::Cli, None, false), - "off" => (SourceMode::Cli, None, false), - "manual" => { - let cookie_header = active_token_cookie.or(stored_cookie); - let fails_closed_without_cookie = cookie_header.is_none() - && provider.manual_empty_cookie_policy() - == ManualEmptyCookiePolicy::FailClosedWeb; - let source_mode = if (has_kimi_code_api_key || has_opencodego_api_key) - && usage_source == SourceMode::Auto - { - SourceMode::Auto - } else if let Some(mode) = grok_source_mode_for_manual_cookie(id, usage_source) - { - // Grok Switch writes ~/.grok/auth.json. Leftover grok.com - // cookies must not force Web, or Weekly/notifications keep - // showing the previous browser account. - mode - } else if cookie_header.is_some() { - SourceMode::Web - } else if fails_closed_without_cookie { - // The provider owns this policy; Web with no header means - // it fails closed instead of importing a browser account - // the user did not select. - SourceMode::Web - } else if provider_uses_oauth_without_cookies(id, usage_source) { - SourceMode::OAuth + SourceMode::Auto + } else if let Some(mode) = grok_source_mode_for_manual_cookie(id, usage_source) { + // Grok Switch writes ~/.grok/auth.json. Leftover grok.com + // cookies must not force Web, or Weekly/notifications keep + // showing the previous browser account. + mode + } else if cookie_header.is_some() { + SourceMode::Web + } else if fails_closed_without_cookie { + // The provider owns this policy; Web with no header means + // it fails closed instead of importing a browser account + // the user did not select. + SourceMode::Web + } else if provider_uses_oauth_without_cookies(id, usage_source) { + SourceMode::OAuth + } else { + SourceMode::Cli + }; + (source_mode, cookie_header, fails_closed_without_cookie) + } + // `browser` is accepted as a legacy alias from older settings. + "auto" | "browser" | "web" => { + // Claude resolves its cached cookie and browser fallback inside + // the provider; other providers retain the shell fallback. + let cookie_header = active_token_cookie.clone().or(stored_cookie).or_else(|| { + if defer_provider_browser_cookie_lookup { + None } else { - SourceMode::Cli - }; - (source_mode, cookie_header, fails_closed_without_cookie) - } - // `browser` is accepted as a legacy alias from older settings. - "auto" | "browser" | "web" => { - // Claude resolves its cached cookie and browser fallback inside - // the provider; other providers retain the shell fallback. - let cookie_header = active_token_cookie.or(stored_cookie).or_else(|| { - if defer_provider_browser_cookie_lookup { - None - } else { - provider_cookie_domain(id, settings).and_then(|domain| { - codexbar::browser::cookies::get_cookie_header(domain) - .ok() - .filter(|h| !h.is_empty()) - }) - } - }); - (usage_source, cookie_header, false) - } - _ => (usage_source, stored_cookie, false), + provider_cookie_domain(id, settings).and_then(|domain| { + codexbar::browser::cookies::get_cookie_header(domain) + .ok() + .filter(|h| !h.is_empty()) + }) + } + }); + (usage_source, cookie_header, false) } - }; + _ => (usage_source, stored_cookie, false), + } + }; // Cookie-web providers (Cursor, OpenCode, …) reject SourceMode::Cli. The shell // historically mapped "manual + no cookie" to Cli, which surfaces as From eddd57f5cc2119be51f93a973fefcaebbbd6e987 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Fri, 25 Sep 2026 11:56:12 +0700 Subject: [PATCH 4/9] Share token account source routing policy --- .../src-tauri/src/commands/providers.rs | 36 ++++---- rust/src/cli/usage/fetch_helpers.rs | 15 +--- rust/src/core/token_accounts.rs | 89 ++++++++++++++++++- 3 files changed, 112 insertions(+), 28 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index f0cf85e8c2..2c2d251fe2 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -225,24 +225,28 @@ pub(crate) fn build_fetch_context( // These upstream account types are explicit identity selections. Keep the // provider's saved region/source settings intact, but project the selected // credential into the route required by that account. - let (source_mode, cookie_header, api_key) = match (id, token_account_kind) { - (ProviderId::Kimi, Some(_)) => (SourceMode::Web, active_token_cookie.clone(), None), - (ProviderId::Doubao, Some(_)) => (SourceMode::OAuth, None, active_token_api_key.clone()), - (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::ApiKey)) - if usage_source == SourceMode::Auto => - { - (SourceMode::Auto, None, active_token_api_key.clone()) - } - (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::ApiKey)) => { - (usage_source, cookie_header, api_key) - } - (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::Cookie)) - if usage_source == SourceMode::Auto => - { - (SourceMode::Web, active_token_cookie.clone(), api_key) + let (cookie_header, api_key) = match (id, token_account_kind, usage_source) { + (ProviderId::Kimi, Some(_), _) => (active_token_cookie.clone(), None), + (ProviderId::Doubao, Some(_), _) => (None, active_token_api_key.clone()), + ( + ProviderId::OpenCodeGo, + Some(codexbar::core::TokenAccountKind::ApiKey), + SourceMode::Auto, + ) => (None, active_token_api_key.clone()), + (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::ApiKey), _) => { + (cookie_header, api_key) } - _ => (source_mode, cookie_header, api_key), + ( + ProviderId::OpenCodeGo, + Some(codexbar::core::TokenAccountKind::Cookie), + SourceMode::Auto, + ) => (active_token_cookie.clone(), api_key), + _ => (cookie_header, api_key), }; + let source_mode = token_override + .as_ref() + .and_then(|account| account.effective_source_mode(usage_source)) + .unwrap_or(source_mode); let token_account_isolated = token_override.is_some() && matches!( id, diff --git a/rust/src/cli/usage/fetch_helpers.rs b/rust/src/cli/usage/fetch_helpers.rs index e9e063f296..f8c519d86b 100644 --- a/rust/src/cli/usage/fetch_helpers.rs +++ b/rust/src/cli/usage/fetch_helpers.rs @@ -5,8 +5,8 @@ use super::render::{ render_brief_text, render_json_result, render_text_error, render_text_with_status, }; use crate::core::{ - ProviderFetchResult, ProviderId, SourceMode, TokenAccountKind, TokenAccountOverride, - TokenAccountStore, TokenAccountSupport, instantiate_provider, + ProviderFetchResult, ProviderId, TokenAccountKind, TokenAccountOverride, TokenAccountStore, + TokenAccountSupport, instantiate_provider, }; use crate::settings::ApiKeys; use crate::status::{ProviderStatus as StatusInfo, fetch_provider_status}; @@ -117,15 +117,8 @@ pub(super) fn project_token_account( ctx.manual_cookie_header = projected.cookie_header; ctx.auto_prefer_web = projected.kind == TokenAccountKind::Cookie; - match (provider, projected.kind) { - (ProviderId::Kimi, _) => ctx.source_mode = SourceMode::Web, - (ProviderId::Doubao, _) => ctx.source_mode = SourceMode::OAuth, - (ProviderId::OpenCodeGo, TokenAccountKind::Cookie) - if ctx.source_mode == SourceMode::Auto => - { - ctx.source_mode = SourceMode::Web; - } - _ => {} + if let Some(source_mode) = projected.effective_source_mode(ctx.source_mode) { + ctx.source_mode = source_mode; } } diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index abd9cd2b86..9c6566e465 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -3,7 +3,7 @@ //! Store and manage multiple accounts/tokens per provider. //! Supports parallel fetching and account switching. -use crate::core::ProviderId; +use crate::core::{ProviderId, SourceMode}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -860,6 +860,20 @@ impl TokenAccountOverride { kind, } } + + /// Normalize source selection for account types whose credential requires + /// a specific route. `None` leaves unrelated providers' source policy alone. + pub fn effective_source_mode(&self, requested: SourceMode) -> Option { + match (self.provider, self.kind, requested) { + (ProviderId::Kimi, _, _) => Some(SourceMode::Web), + (ProviderId::Doubao, _, _) => Some(SourceMode::OAuth), + (ProviderId::OpenCodeGo, TokenAccountKind::Cookie, SourceMode::Auto) => { + Some(SourceMode::Web) + } + (ProviderId::OpenCodeGo, _, _) => Some(requested), + _ => None, + } + } } /// Maximum number of accounts to fetch per provider @@ -939,6 +953,79 @@ mod tests { ); } + #[test] + fn selected_account_effective_source_normalization() { + let cases = [ + ( + ProviderId::Kimi, + "kimi-session", + SourceMode::Auto, + Some(SourceMode::Web), + ), + ( + ProviderId::Kimi, + "kimi-session", + SourceMode::OAuth, + Some(SourceMode::Web), + ), + ( + ProviderId::Kimi, + "kimi-session", + SourceMode::Cli, + Some(SourceMode::Web), + ), + ( + ProviderId::Doubao, + "ark-key", + SourceMode::Cli, + Some(SourceMode::OAuth), + ), + ( + ProviderId::Doubao, + "ark-key", + SourceMode::Web, + Some(SourceMode::OAuth), + ), + ( + ProviderId::OpenCodeGo, + "Cookie: session=web", + SourceMode::Auto, + Some(SourceMode::Web), + ), + ( + ProviderId::OpenCodeGo, + "Cookie: session=web", + SourceMode::Cli, + Some(SourceMode::Cli), + ), + ( + ProviderId::OpenCodeGo, + "api-key", + SourceMode::Auto, + Some(SourceMode::Auto), + ), + ( + ProviderId::OpenCodeGo, + "api-key", + SourceMode::Web, + Some(SourceMode::Web), + ), + ( + ProviderId::OpenCodeGo, + "api-key", + SourceMode::Cli, + Some(SourceMode::Cli), + ), + (ProviderId::OpenRouter, "api-key", SourceMode::Auto, None), + ]; + + for (provider, token, requested, expected) in cases { + let account = + TokenAccountOverride::from_account(provider, TokenAccount::new("selected", token)); + assert_eq!(account.effective_source_mode(requested), expected); + } + } + #[test] fn grok_token_accounts_route_bearer_and_cookie_credentials() { let bearer = From e0a4bdc2f12c0fbf1b8d8c6083f80bff7385ae77 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Fri, 25 Sep 2026 12:13:47 +0700 Subject: [PATCH 5/9] Compute account source before moving credentials --- rust/src/cli/usage/fetch_helpers.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rust/src/cli/usage/fetch_helpers.rs b/rust/src/cli/usage/fetch_helpers.rs index f8c519d86b..7b75cada51 100644 --- a/rust/src/cli/usage/fetch_helpers.rs +++ b/rust/src/cli/usage/fetch_helpers.rs @@ -108,6 +108,7 @@ pub(super) fn project_token_account( ctx: &mut crate::core::FetchContext, ) { let projected = TokenAccountOverride::from_account(provider, account.clone()); + let effective_source_mode = projected.effective_source_mode(ctx.source_mode); ctx.token_account_kind = Some(projected.kind); ctx.token_account_isolated = true; ctx.api_key = projected @@ -117,7 +118,7 @@ pub(super) fn project_token_account( ctx.manual_cookie_header = projected.cookie_header; ctx.auto_prefer_web = projected.kind == TokenAccountKind::Cookie; - if let Some(source_mode) = projected.effective_source_mode(ctx.source_mode) { + if let Some(source_mode) = effective_source_mode { ctx.source_mode = source_mode; } } From 3fb049ab0f402220309f62eddc8dfe681f0960ff Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:25:43 +0700 Subject: [PATCH 6/9] Port upstream 0.65.0: Bifrost identity, unknown rate limits, model labels --- rust/src/providers/bifrost/mod.rs | 517 ++++++++------------- rust/src/providers/bifrost/model_labels.rs | 160 +++++++ rust/src/providers/bifrost/tests.rs | 390 ++++++++++++++++ 3 files changed, 731 insertions(+), 336 deletions(-) create mode 100644 rust/src/providers/bifrost/model_labels.rs create mode 100644 rust/src/providers/bifrost/tests.rs diff --git a/rust/src/providers/bifrost/mod.rs b/rust/src/providers/bifrost/mod.rs index d872103c12..cddce9474f 100644 --- a/rust/src/providers/bifrost/mod.rs +++ b/rust/src/providers/bifrost/mod.rs @@ -10,7 +10,7 @@ use serde_json::Value; use std::{net::IpAddr, time::Duration}; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, + CostSnapshot, FetchContext, NamedRateWindow, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; @@ -18,8 +18,11 @@ const CREDENTIAL_TARGET: &str = "codexbar-bifrost"; const API_KEY_ENV: &str = "BIFROST_API_KEY"; const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); const MAX_RESPONSE_BYTES: usize = 2 * 1024 * 1024; +const MAX_MODEL_ROWS: usize = 5; const QUOTA_PATH: &str = "/api/governance/virtual-keys/quota"; +mod model_labels; + pub struct BifrostProvider { metadata: ProviderMetadata, client: Option, @@ -38,7 +41,6 @@ struct Budget { scope: Scope, source: Option, used: f64, - usage_known: bool, limit: f64, reset: ResetTiming, models: Vec, @@ -285,8 +287,8 @@ fn parse_usage(value: &Value, now: DateTime) -> Result) -> Result) -> Result) -> Result, + key_name: Option, budgets: Vec, limits: Vec<(Scope, Value, usize)>, } @@ -406,7 +409,7 @@ fn add_scopes( } fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { - let mut budget_windows = usage + let budget_windows = usage .budgets .iter() .filter(|budget| budget.limit > 0.0) @@ -426,8 +429,7 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { budget.reset.window_minutes, budget.reset.resets_at, Some(budget_description(budget)), - ) - .with_usage_known(budget.usage_known); + ); (budget, title, window) }) .collect::>(); @@ -441,90 +443,28 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { .unwrap_or_else(|| RateWindow::informational("No Bifrost budget quota reported")); let secondary = root_windows.get(1).map(|(_, _, window)| (*window).clone()); let mut snapshot = UsageSnapshot::new(primary).with_login_method("API"); + if let Some(key_name) = &usage.key_name { + snapshot = snapshot.with_email(key_name.clone()); + } + if let Some(source) = first_root_budget(&usage).and_then(|first| first.source.as_deref()) { + snapshot = snapshot.with_organization(source); + } if let Some(secondary) = secondary { snapshot = snapshot.with_secondary(secondary); } let mut extra_windows = root_windows .iter() .skip(2) - .map(|(_, title, window)| { - ( - format!( - "bifrost-budget-{}", - window.reset_description.as_deref().unwrap_or("extra") - ), - title.clone(), - window.clone(), - ) - }) + .map(|(budget, title, window)| budget_named_window(budget, title, window)) .collect::>(); extra_windows.extend( budget_windows - .drain(..) + .iter() .filter(|(budget, _, _)| !budget.scope.id.is_empty()) - .map(|(budget, title, window)| { - ( - format!("bifrost-{}budget-{}", budget.scope.id, budget.id), - title, - window, - ) - }), + .map(|(budget, title, window)| budget_named_window(budget, title, window)), ); - - for (scope, limit, index) in &usage.limits { - let source = optional_text(limit, "source_name").ok().flatten(); - for (key, title) in [("token", "Tokens"), ("request", "Requests")] { - let max = optional_number(limit, &format!("{key}_max_limit")) - .ok() - .flatten(); - let reset_raw = optional_text(limit, &format!("{key}_reset_duration")) - .ok() - .flatten(); - let Some(max_or_reset) = max - .filter(|max| *max > 0.0) - .or_else(|| reset_raw.as_ref().map(|_| 0.0)) - else { - continue; - }; - let used = optional_number(limit, &format!("{key}_current_usage")) - .ok() - .flatten() - .unwrap_or(0.0); - let reset = reset_timing( - reset_raw.as_deref(), - optional_text(limit, &format!("{key}_last_reset")) - .ok() - .flatten() - .as_deref(), - Utc::now(), - ); - let title = bounded( - &[scope.title.as_deref(), source.as_deref(), Some(title)] - .into_iter() - .flatten() - .collect::>() - .join(" "), - ); - let window = RateWindow::with_details( - if max_or_reset > 0.0 { - percent(used, max_or_reset) - } else { - 0.0 - }, - reset.window_minutes, - reset.resets_at, - reset.label.map(str::to_owned), - ); - extra_windows.push(( - format!("bifrost-{}{}s-{index}", scope.id, key), - title, - window, - )); - } - } - for (id, title, window) in extra_windows { - snapshot = snapshot.with_extra_rate_window(id, title, window); - } + extra_windows.extend(rate_limit_windows(&usage.limits)); + snapshot.extra_rate_windows.extend(extra_windows); if usage.active == Some(false) { snapshot = snapshot.with_extra_rate_window( "bifrost-key-inactive", @@ -534,11 +474,7 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { } let mut result = ProviderFetchResult::new(snapshot, "api"); - if let Some(first) = usage - .budgets - .iter() - .find(|budget| budget.scope.id.is_empty()) - { + if let Some(first) = first_root_budget(&usage) { let mut cost = CostSnapshot::new( first.used, "USD", @@ -585,59 +521,175 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { result = result.with_display_detail(detail); } } - if let Some(first) = usage + if let Some(first) = first_root_budget(&usage) { + for detail in model_details(&first.models) { + result = result.with_display_detail(detail); + } + } + result +} + +fn first_root_budget(usage: &ParsedUsage) -> Option<&Budget> { + usage .budgets .iter() .find(|budget| budget.scope.id.is_empty()) - { - let mut models = first - .models - .iter() - .filter_map(|model| { - let name = optional_text(model, "model") +} + +fn budget_named_window(budget: &Budget, title: &str, window: &RateWindow) -> NamedRateWindow { + NamedRateWindow::new( + format!("bifrost-{}budget-{}", budget.scope.id, budget.id), + title, + window.clone(), + ) +} + +/// Token and request windows for every rate limit. A dimension with a reset +/// duration but no positive max is reported as an unknown-usage window; a +/// dimension with neither is unconfigured and skipped. +fn rate_limit_windows(limits: &[(Scope, Value, usize)]) -> Vec { + let mut windows = Vec::new(); + for (scope, limit, index) in limits { + let source = optional_text(limit, "source_name").ok().flatten(); + for (key, title) in [("token", "Tokens"), ("request", "Requests")] { + let max = optional_number(limit, &format!("{key}_max_limit")) + .ok() + .flatten() + .filter(|max| *max > 0.0); + let reset_raw = optional_text(limit, &format!("{key}_reset_duration")) + .ok() + .flatten(); + if max.is_none() && reset_raw.is_none() { + continue; + } + let used = optional_number(limit, &format!("{key}_current_usage")) + .ok() + .flatten() + .unwrap_or(0.0); + let reset = reset_timing( + reset_raw.as_deref(), + optional_text(limit, &format!("{key}_last_reset")) .ok() .flatten() - .unwrap_or_else(|| "Model".into()); - let provider = optional_text(model, "provider").ok().flatten(); - let cost = optional_number(model, "total_cost").ok().flatten(); - let tokens = optional_number(model, "total_tokens").ok().flatten(); - (cost.unwrap_or(0.0) != 0.0 || tokens.unwrap_or(0.0) != 0.0) - .then_some((name, provider, cost, tokens)) - }) - .collect::>(); - models.sort_by(|a, b| { - b.2.unwrap_or(0.0) - .total_cmp(&a.2.unwrap_or(0.0)) - .then_with(|| b.3.unwrap_or(0.0).total_cmp(&a.3.unwrap_or(0.0))) - .then_with(|| a.0.cmp(&b.0)) - }); - for (index, (name, provider, cost, tokens)) in models.iter().take(5).enumerate() { - let label = bounded( - &[provider.as_deref(), Some(name.as_str())] + .as_deref(), + Utc::now(), + ); + let title = bounded( + &[scope.title.as_deref(), source.as_deref(), Some(title)] .into_iter() .flatten() .collect::>() - .join(" · "), + .join(" "), + ); + let known = max.is_some(); + let window = RateWindow::with_details( + max.map_or(0.0, |max| percent(used, max)), + reset.window_minutes, + reset.resets_at, + reset.label.map(str::to_owned), + ) + .with_usage_known(known); + windows.push( + NamedRateWindow::new( + format!("bifrost-{}{}s-{index}", scope.id, key), + title, + window, + ) + .with_usage_known(known), ); - let value = cost.map(usd).unwrap_or_else(|| "—".into()); - let mut detail = - ProviderDisplayDetail::new(format!("bifrost-model-{index}"), label, value); - if let Some(tokens) = tokens { - detail = detail.and_then(|row| { - row.with_secondary_value(format!("{} tokens", format_count(*tokens))) - }); - } - result = result.with_display_detail(detail); - } - if models.len() > 5 { - result = result.with_display_detail(ProviderDisplayDetail::new( - "bifrost-model-other", - "Other models", - (models.len() - 5).to_string(), - )); } } - result + windows +} + +struct ModelRow { + raw: String, + provider: Option, + cost: Option, + tokens: Option, +} + +/// Spend rows for the first root budget, most expensive first, capped at +/// `MAX_MODEL_ROWS` plus an "Other models" count. +fn model_details(models: &[Value]) -> Vec> { + let mut rows = models + .iter() + .filter_map(|model| { + let cost = optional_number(model, "total_cost").ok().flatten(); + // Upstream truncates the count and drops values outside i64. + let tokens = optional_number(model, "total_tokens") + .ok() + .flatten() + .filter(|tokens| tokens.abs() < 9_223_372_036_854_775_808.0) + .map(f64::trunc); + (cost.unwrap_or(0.0) != 0.0 || tokens.unwrap_or(0.0) != 0.0).then(|| ModelRow { + raw: optional_text(model, "model") + .ok() + .flatten() + .unwrap_or_else(|| "Model".into()), + provider: optional_text(model, "provider").ok().flatten(), + cost, + tokens, + }) + }) + .collect::>(); + rows.sort_by(|a, b| { + b.cost + .unwrap_or(0.0) + .total_cmp(&a.cost.unwrap_or(0.0)) + .then_with(|| b.tokens.unwrap_or(0.0).total_cmp(&a.tokens.unwrap_or(0.0))) + .then_with(|| a.raw.cmp(&b.raw)) + }); + let total = rows.len(); + let visible = &rows[..total.min(MAX_MODEL_ROWS)]; + let names = visible + .iter() + .map(|row| model_labels::model_name(&row.raw)) + .collect::>(); + // The provider prefix only helps when the visible rows mix providers. + let mut providers = visible.iter().filter_map(|row| row.provider.as_deref()); + let mixed = providers + .next() + .is_some_and(|first| providers.any(|provider| provider != first)); + + let mut details = visible + .iter() + .zip(&names) + .enumerate() + .map(|(index, (row, name))| { + // Two rows sharing a short name fall back to the raw model id. + let duplicate = names.iter().filter(|other| *other == name).count() > 1; + let label = bounded( + &[ + mixed.then_some(row.provider.as_deref()).flatten(), + Some(if duplicate { &row.raw } else { name }), + ] + .into_iter() + .flatten() + .collect::>() + .join(" · "), + ); + let value = row.cost.map(usd).unwrap_or_else(|| "—".into()); + let detail = ProviderDisplayDetail::new(format!("bifrost-model-{index}"), label, value); + match row.tokens { + Some(tokens) => detail.and_then(|detail| { + detail.with_secondary_value(format!( + "{} tokens", + model_labels::token_count(tokens) + )) + }), + None => detail, + } + }) + .collect::>(); + if total > MAX_MODEL_ROWS { + details.push(ProviderDisplayDetail::new( + "bifrost-model-other", + "Other models", + (total - MAX_MODEL_ROWS).to_string(), + )); + } + details } fn budget_description(budget: &Budget) -> String { @@ -820,13 +872,6 @@ fn usd(value: f64) -> String { fn bounded(value: &str) -> String { value.chars().take(120).collect() } -fn format_count(value: f64) -> String { - if value.fract() == 0.0 && value.abs() <= i64::MAX as f64 { - format!("{value:.0}") - } else { - format!("{value:.2}") - } -} fn parse_error(reason: &str) -> ProviderError { ProviderError::Parse(format!( "Bifrost returned an unrecognized quota response ({reason})." @@ -834,204 +879,4 @@ fn parse_error(reason: &str) -> ProviderError { } #[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, - time::timeout, - }; - - fn now() -> DateTime { - DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z") - .unwrap() - .with_timezone(&Utc) - } - - #[test] - fn component_rate_limits_take_precedence_over_aggregate_compatibility_field() { - let value = json!({ - "rate_limit": { "token_max_limit": 900, "token_current_usage": 450, "token_reset_duration": "1h" }, - "rate_limits": [ - { "source_name": "provider-a", "token_max_limit": 100, "token_current_usage": 25, "token_reset_duration": "1h" }, - { "source_name": "provider-b", "request_max_limit": 20, "request_current_usage": 5, "request_reset_duration": "1d" } - ] - }); - let parsed = parse_usage(&value, now()).unwrap(); - assert_eq!(parsed.limits.len(), 2); - assert_eq!(parsed.limits[0].1["token_max_limit"], 100); - assert_eq!(parsed.limits[1].1["request_max_limit"], 20); - assert!( - !parsed - .limits - .iter() - .any(|(_, limit, _)| limit["token_max_limit"] == 900) - ); - } - - #[test] - fn aggregate_rate_limit_is_fallback_when_components_are_absent() { - let parsed = parse_usage(&json!({ - "rate_limit": { "token_max_limit": 50, "token_current_usage": 10, "token_reset_duration": "1h" } - }), now()).unwrap(); - assert_eq!(parsed.limits.len(), 1); - assert_eq!(parsed.limits[0].1["token_max_limit"], 50); - } - - #[test] - fn missing_current_usage_at_a_positive_rate_limit_is_known_zero() { - for (usage_field, id) in [ - ("token_current_usage", "bifrost-tokens-0"), - ("request_current_usage", "bifrost-requests-0"), - ] { - let mut limit = json!({ - "token_max_limit": 100, - "token_current_usage": 25, - "request_max_limit": 20, - "request_current_usage": 5 - }); - limit.as_object_mut().unwrap().remove(usage_field); - let result = - result_from_usage(parse_usage(&json!({ "rate_limit": limit }), now()).unwrap()); - let named = result - .usage - .extra_rate_windows - .iter() - .find(|window| window.id == id) - .unwrap_or_else(|| panic!("missing named rate window {id}")); - - assert_eq!(named.window.used_percent, 0.0, "{usage_field}"); - assert!(named.usage_known, "{usage_field}"); - assert!(named.window.usage_known(), "{usage_field}"); - } - } - - #[test] - fn validates_gateway_before_request_url_is_built() { - assert!( - quota_url_for_test("https://bifrost.example.com/base/") - .unwrap() - .as_str() - .starts_with("https://bifrost.example.com/base/api/governance/virtual-keys/quota") - ); - assert!(quota_url_for_test("http://10.1.2.3:8080").is_ok()); - assert!(quota_url_for_test("http://bifrost.example.com").is_err()); - assert!(quota_url_for_test("https://user:secret@bifrost.example.com").is_err()); - assert!(quota_url_for_test("ftp://10.1.2.3").is_err()); - } - - #[tokio::test] - async fn rejects_public_http_before_resolving_any_credential() { - let provider = BifrostProvider::new(); - let ctx = FetchContext { - gateway_url: Some("http://public.example.com".into()), - ..FetchContext::default() - }; - let error = provider.fetch_api(&ctx).await.unwrap_err(); - assert!(error.to_string().contains("must use HTTPS")); - } - - #[tokio::test] - async fn does_not_forward_virtual_key_through_gateway_redirects() { - let redirect_target = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let redirect_target_addr = redirect_target.local_addr().unwrap(); - let gateway = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let gateway_addr = gateway.local_addr().unwrap(); - let gateway_task = tokio::spawn(async move { - let (mut stream, _) = gateway.accept().await.unwrap(); - let mut request = vec![0; 4096]; - let read = stream.read(&mut request).await.unwrap(); - let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase(); - assert!(request.contains("x-bf-vk: test-virtual-key")); - let response = format!( - "HTTP/1.1 302 Found\r\nLocation: http://{redirect_target_addr}/redirected\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - }); - let redirect_target_task = tokio::spawn(async move { - timeout(Duration::from_millis(250), redirect_target.accept()).await - }); - - let provider = BifrostProvider::new(); - let ctx = FetchContext { - gateway_url: Some(format!("http://{gateway_addr}")), - api_key: Some("test-virtual-key".into()), - ..FetchContext::default() - }; - let error = provider.fetch_api(&ctx).await.unwrap_err(); - assert!(error.to_string().contains("HTTP 302")); - gateway_task.await.unwrap(); - assert!(redirect_target_task.await.unwrap().is_err()); - } - - #[test] - fn parses_fixed_reset_and_calendar_label_without_inventing_calendar_time() { - let fixed = reset_timing(Some("1h"), Some("2025-12-31T23:30:00Z"), now()); - assert_eq!(fixed.window_minutes, Some(60)); - assert_eq!( - fixed.resets_at.unwrap().to_rfc3339(), - "2026-01-01T00:30:00+00:00" - ); - let calendar = reset_timing(Some("1M"), Some("2025-12-01T00:00:00Z"), now()); - assert_eq!(calendar.label, Some("Monthly")); - assert_eq!(calendar.resets_at, None); - assert_eq!(calendar.window_minutes, None); - assert_eq!(parse_duration("1h30m"), Some(5_400.0)); - } - - #[test] - fn budgets_and_overrides_are_mapped_to_percent_and_spend() { - let parsed = parse_usage( - &json!({ - "virtual_key_name": "Build key", - "budgets": [{ "id": "b1", "max_limit": 10, "current_usage": 5, - "override_amount": 5, "override_mode": "forever", "source_name": "Team" }] - }), - now(), - ) - .unwrap(); - let result = result_from_usage(parsed); - assert!((result.usage.primary.used_percent - (100.0 / 3.0)).abs() < 0.001); - assert_eq!(result.cost.unwrap().limit, Some(15.0)); - } - - #[test] - fn shortest_root_budget_is_primary_and_cost_without_summing_budgets() { - let parsed = parse_usage( - &json!({ - "budgets": [ - { "id": "monthly", "max_limit": 1_000, "current_usage": 200, - "reset_duration": "1M" }, - { "id": "daily", "max_limit": 100, "current_usage": 10, - "reset_duration": "1d" } - ] - }), - now(), - ) - .unwrap(); - let result = result_from_usage(parsed); - - assert_eq!(result.usage.primary.used_percent, 10.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("Daily · $10.00 / $100.00") - ); - let secondary = result.usage.secondary.as_ref().unwrap(); - assert_eq!(secondary.used_percent, 20.0); - assert_eq!( - secondary.reset_description.as_deref(), - Some("Monthly · $200.00 / $1000.00") - ); - - let cost = result.cost.unwrap(); - assert_eq!(cost.used, 10.0); - assert_eq!(cost.limit, Some(100.0)); - assert_eq!(cost.period, "Daily"); - } - - #[test] - fn malformed_optional_scope_collection_fails_closed() { - assert!(parse_usage(&json!({ "provider_configs": {} }), now()).is_err()); - } -} +mod tests; diff --git a/rust/src/providers/bifrost/model_labels.rs b/rust/src/providers/bifrost/model_labels.rs new file mode 100644 index 0000000000..432361d0a4 --- /dev/null +++ b/rust/src/providers/bifrost/model_labels.rs @@ -0,0 +1,160 @@ +//! Display formatting for Bifrost per-model spend rows, matching the tagged +//! upstream `modelName()` and `tokenCount()` helpers. + +const REGIONS: [&str; 5] = ["us-gov", "us", "eu", "apac", "global"]; +const VENDORS: [&str; 13] = [ + "ai21", + "amazon", + "anthropic", + "cohere", + "deepseek", + "luma", + "meta", + "mistral", + "openai", + "qwen", + "stability", + "twelvelabs", + "writer", +]; + +/// Strip a leading region and vendor token, a Bedrock `-vN:N` revision, and a +/// trailing date so a versioned gateway model id reads as its short name. +pub(super) fn model_name(raw: &str) -> String { + let name = strip_dotted_prefix(raw, ®IONS); + let name = strip_dotted_prefix(name, &VENDORS); + let name = strip_revision(name); + let name = strip_date(name); + let name = name.trim_end_matches([' ', '\t', '-']); + if name.is_empty() { raw } else { name }.to_owned() +} + +fn strip_dotted_prefix<'a>(value: &'a str, tokens: &[&str]) -> &'a str { + for token in tokens { + if let Some(head) = value.get(..token.len()) + && head.eq_ignore_ascii_case(token) + && let Some(rest) = value[token.len()..].strip_prefix('.') + { + return rest; + } + } + value +} + +/// Remove a trailing `-v:`. +fn strip_revision(value: &str) -> &str { + let Some((head, tail)) = value.rsplit_once("-v") else { + return value; + }; + match tail.split_once(':') { + Some((major, minor)) if all_digits(major) && all_digits(minor) => head, + _ => value, + } +} + +/// Remove a trailing `-`/whitespace separator plus `YYYYMMDD` or `YYYY-MM-DD`. +fn strip_date(value: &str) -> &str { + let bytes = value.as_bytes(); + for len in [10usize, 8] { + let Some(start) = bytes.len().checked_sub(len + 1) else { + continue; + }; + let (separator, date) = (bytes[start], &bytes[start + 1..]); + let dated = if len == 10 { + date[4] == b'-' + && date[7] == b'-' + && [0, 1, 2, 3, 5, 6, 8, 9] + .iter() + .all(|&i| date[i].is_ascii_digit()) + } else { + date.iter().all(u8::is_ascii_digit) + }; + if dated && (separator == b'-' || separator.is_ascii_whitespace()) { + return &value[..start]; + } + } + value +} + +fn all_digits(value: &str) -> bool { + !value.is_empty() && value.bytes().all(|b| b.is_ascii_digit()) +} + +/// Compact token count: `K`/`M`/`B` above 1,000 / 999,500 / 999,500,000 with +/// one decimal below 10 and no trailing `.0`. +pub(super) fn token_count(value: f64) -> String { + let magnitude = value.abs(); + let sign = if value < 0.0 { "-" } else { "" }; + for (threshold, divisor, unit) in [ + (999_500_000.0, 1e9, "B"), + (999_500.0, 1e6, "M"), + (1_000.0, 1e3, "K"), + ] { + if magnitude >= threshold { + let scaled = magnitude / divisor; + let fixed = to_fixed(scaled, usize::from(scaled < 10.0)); + return format!("{sign}{}{unit}", fixed.strip_suffix(".0").unwrap_or(&fixed)); + } + } + format!("{value:.0}") +} + +/// `Number.prototype.toFixed`: exact ties round up, where Rust's formatter +/// rounds them to even. +fn to_fixed(value: f64, digits: usize) -> String { + let exact = format!("{value:.40}"); + let fraction = exact.split_once('.').map_or("", |(_, fraction)| fraction); + let tie = fraction + .get(digits..) + .is_some_and(|tail| tail.starts_with('5') && tail[1..].bytes().all(|b| b == b'0')); + let value = if tie { value.next_up() } else { value }; + format!("{value:.digits$}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_names_drop_region_vendor_revision_and_date() { + for (raw, expected) in [ + ( + "us.anthropic.claude-sonnet-4-20250514-v1:0", + "claude-sonnet-4", + ), + ("EU.Amazon.nova-pro-v1:0", "nova-pro"), + ("us-gov.anthropic.claude-3-haiku", "claude-3-haiku"), + ("gpt-4o-2024-08-06", "gpt-4o"), + ("claude 20250514", "claude"), + ("gpt-4o", "gpt-4o"), + ("global.qwen.qwen3-32b", "qwen3-32b"), + ("usa.model", "usa.model"), + ("anthropic.", "anthropic."), + ("-20250514", "-20250514"), + ("model-v1:", "model-v1:"), + ("20250514", "20250514"), + ] { + assert_eq!(model_name(raw), expected, "{raw}"); + } + } + + #[test] + fn token_counts_use_upstream_thresholds_and_half_up_ties() { + for (value, expected) in [ + (0.0, "0"), + (999.0, "999"), + (1_000.0, "1K"), + (1_250.0, "1.3K"), + (1_500.0, "1.5K"), + (12_500.0, "13K"), + (999_499.0, "999K"), + (999_500.0, "1M"), + (2_340_000.0, "2.3M"), + (999_499_999.0, "999M"), + (999_500_000.0, "1B"), + (-1_250.0, "-1.3K"), + ] { + assert_eq!(token_count(value), expected, "{value}"); + } + } +} diff --git a/rust/src/providers/bifrost/tests.rs b/rust/src/providers/bifrost/tests.rs new file mode 100644 index 0000000000..331c0dedad --- /dev/null +++ b/rust/src/providers/bifrost/tests.rs @@ -0,0 +1,390 @@ +use super::*; +use serde_json::json; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + time::timeout, +}; + +fn now() -> DateTime { + DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z") + .unwrap() + .with_timezone(&Utc) +} + +#[test] +fn component_rate_limits_take_precedence_over_aggregate_compatibility_field() { + let value = json!({ + "rate_limit": { "token_max_limit": 900, "token_current_usage": 450, "token_reset_duration": "1h" }, + "rate_limits": [ + { "source_name": "provider-a", "token_max_limit": 100, "token_current_usage": 25, "token_reset_duration": "1h" }, + { "source_name": "provider-b", "request_max_limit": 20, "request_current_usage": 5, "request_reset_duration": "1d" } + ] + }); + let parsed = parse_usage(&value, now()).unwrap(); + assert_eq!(parsed.limits.len(), 2); + assert_eq!(parsed.limits[0].1["token_max_limit"], 100); + assert_eq!(parsed.limits[1].1["request_max_limit"], 20); + assert!( + !parsed + .limits + .iter() + .any(|(_, limit, _)| limit["token_max_limit"] == 900) + ); +} + +#[test] +fn aggregate_rate_limit_is_fallback_when_components_are_absent() { + let parsed = parse_usage(&json!({ + "rate_limit": { "token_max_limit": 50, "token_current_usage": 10, "token_reset_duration": "1h" } + }), now()).unwrap(); + assert_eq!(parsed.limits.len(), 1); + assert_eq!(parsed.limits[0].1["token_max_limit"], 50); +} + +#[test] +fn missing_current_usage_at_a_positive_rate_limit_is_known_zero() { + for (usage_field, id) in [ + ("token_current_usage", "bifrost-tokens-0"), + ("request_current_usage", "bifrost-requests-0"), + ] { + let mut limit = json!({ + "token_max_limit": 100, + "token_current_usage": 25, + "request_max_limit": 20, + "request_current_usage": 5 + }); + limit.as_object_mut().unwrap().remove(usage_field); + let result = + result_from_usage(parse_usage(&json!({ "rate_limit": limit }), now()).unwrap()); + let named = result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == id) + .unwrap_or_else(|| panic!("missing named rate window {id}")); + + assert_eq!(named.window.used_percent, 0.0, "{usage_field}"); + assert!(named.usage_known, "{usage_field}"); + assert!(named.window.usage_known(), "{usage_field}"); + } +} + +#[test] +fn validates_gateway_before_request_url_is_built() { + assert!( + quota_url_for_test("https://bifrost.example.com/base/") + .unwrap() + .as_str() + .starts_with("https://bifrost.example.com/base/api/governance/virtual-keys/quota") + ); + assert!(quota_url_for_test("http://10.1.2.3:8080").is_ok()); + assert!(quota_url_for_test("http://bifrost.example.com").is_err()); + assert!(quota_url_for_test("https://user:secret@bifrost.example.com").is_err()); + assert!(quota_url_for_test("ftp://10.1.2.3").is_err()); +} + +#[tokio::test] +async fn rejects_public_http_before_resolving_any_credential() { + let provider = BifrostProvider::new(); + let ctx = FetchContext { + gateway_url: Some("http://public.example.com".into()), + ..FetchContext::default() + }; + let error = provider.fetch_api(&ctx).await.unwrap_err(); + assert!(error.to_string().contains("must use HTTPS")); +} + +#[tokio::test] +async fn does_not_forward_virtual_key_through_gateway_redirects() { + let redirect_target = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let redirect_target_addr = redirect_target.local_addr().unwrap(); + let gateway = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let gateway_addr = gateway.local_addr().unwrap(); + let gateway_task = tokio::spawn(async move { + let (mut stream, _) = gateway.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let read = stream.read(&mut request).await.unwrap(); + let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase(); + assert!(request.contains("x-bf-vk: test-virtual-key")); + let response = format!( + "HTTP/1.1 302 Found\r\nLocation: http://{redirect_target_addr}/redirected\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ); + stream.write_all(response.as_bytes()).await.unwrap(); + }); + let redirect_target_task = + tokio::spawn( + async move { timeout(Duration::from_millis(250), redirect_target.accept()).await }, + ); + + let provider = BifrostProvider::new(); + let ctx = FetchContext { + gateway_url: Some(format!("http://{gateway_addr}")), + api_key: Some("test-virtual-key".into()), + ..FetchContext::default() + }; + let error = provider.fetch_api(&ctx).await.unwrap_err(); + assert!(error.to_string().contains("HTTP 302")); + gateway_task.await.unwrap(); + assert!(redirect_target_task.await.unwrap().is_err()); +} + +#[test] +fn parses_fixed_reset_and_calendar_label_without_inventing_calendar_time() { + let fixed = reset_timing(Some("1h"), Some("2025-12-31T23:30:00Z"), now()); + assert_eq!(fixed.window_minutes, Some(60)); + assert_eq!( + fixed.resets_at.unwrap().to_rfc3339(), + "2026-01-01T00:30:00+00:00" + ); + let calendar = reset_timing(Some("1M"), Some("2025-12-01T00:00:00Z"), now()); + assert_eq!(calendar.label, Some("Monthly")); + assert_eq!(calendar.resets_at, None); + assert_eq!(calendar.window_minutes, None); + assert_eq!(parse_duration("1h30m"), Some(5_400.0)); +} + +#[test] +fn budgets_and_overrides_are_mapped_to_percent_and_spend() { + let parsed = parse_usage( + &json!({ + "virtual_key_name": "Build key", + "budgets": [{ "id": "b1", "max_limit": 10, "current_usage": 5, + "override_amount": 5, "override_mode": "forever", "source_name": "Team" }] + }), + now(), + ) + .unwrap(); + let result = result_from_usage(parsed); + assert!((result.usage.primary.used_percent - (100.0 / 3.0)).abs() < 0.001); + assert_eq!(result.cost.unwrap().limit, Some(15.0)); +} + +#[test] +fn shortest_root_budget_is_primary_and_cost_without_summing_budgets() { + let parsed = parse_usage( + &json!({ + "budgets": [ + { "id": "monthly", "max_limit": 1_000, "current_usage": 200, + "reset_duration": "1M" }, + { "id": "daily", "max_limit": 100, "current_usage": 10, + "reset_duration": "1d" } + ] + }), + now(), + ) + .unwrap(); + let result = result_from_usage(parsed); + + assert_eq!(result.usage.primary.used_percent, 10.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Daily · $10.00 / $100.00") + ); + let secondary = result.usage.secondary.as_ref().unwrap(); + assert_eq!(secondary.used_percent, 20.0); + assert_eq!( + secondary.reset_description.as_deref(), + Some("Monthly · $200.00 / $1000.00") + ); + + let cost = result.cost.unwrap(); + assert_eq!(cost.used, 10.0); + assert_eq!(cost.limit, Some(100.0)); + assert_eq!(cost.period, "Daily"); +} + +#[test] +fn malformed_optional_scope_collection_fails_closed() { + assert!(parse_usage(&json!({ "provider_configs": {} }), now()).is_err()); +} + +fn result_for(value: Value) -> ProviderFetchResult { + result_from_usage(parse_usage(&value, now()).unwrap()) +} + +fn detail_rows(result: &ProviderFetchResult) -> Vec<(String, String, Option)> { + result + .display_details() + .iter() + .filter(|row| row.id().starts_with("bifrost-model-")) + .map(|row| { + ( + row.title().to_owned(), + row.value().to_owned(), + row.secondary_value().map(str::to_owned), + ) + }) + .collect() +} + +#[test] +fn identity_uses_key_name_and_first_root_budget_source() { + let result = result_for(json!({ + "virtual_key_name": " Build key ", + "budgets": [ + { "id": "monthly", "max_limit": 1_000, "current_usage": 1, "reset_duration": "1M", + "source_name": "Monthly team" }, + { "id": "daily", "max_limit": 100, "current_usage": 1, "reset_duration": "1d", + "source_name": "Daily team" } + ], + "provider_configs": [ + { "provider": "openai", "budgets": [ + { "id": "p", "max_limit": 5, "current_usage": 1, "reset_duration": "1h", + "source_name": "Provider team" } + ] } + ] + })); + assert_eq!(result.usage.account_email.as_deref(), Some("Build key")); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Daily team") + ); + assert_eq!(result.usage.login_method.as_deref(), Some("API")); +} + +#[test] +fn identity_is_omitted_when_key_name_and_source_are_empty() { + let result = result_for(json!({ + "virtual_key_name": " ", + "budgets": [{ "id": "b", "max_limit": 5, "current_usage": 1, "source_name": " " }] + })); + assert_eq!(result.usage.account_email, None); + assert_eq!(result.usage.account_organization, None); + assert!( + parse_usage(&json!({ "virtual_key_name": 7 }), now()).is_err(), + "a non-string key name is a malformed response" + ); +} + +#[test] +fn reset_only_rate_limit_dimension_is_an_unknown_named_window() { + let result = result_for(json!({ + "rate_limit": { + "token_max_limit": 0, "token_current_usage": 40, "token_reset_duration": "1h", + "request_max_limit": 0, "request_current_usage": 3, + "request_last_reset": "2025-12-31T23:30:00Z" + } + })); + let windows = &result.usage.extra_rate_windows; + assert_eq!( + windows.len(), + 1, + "unconfigured request dimension is skipped" + ); + let tokens = &windows[0]; + assert_eq!(tokens.id, "bifrost-tokens-0"); + assert_eq!(tokens.window.used_percent, 0.0); + assert!(!tokens.usage_known); + assert!(!tokens.window.usage_known()); + assert_eq!(tokens.window.window_minutes, Some(60)); +} + +#[test] +fn root_budget_extras_use_scope_and_budget_id() { + let result = result_for(json!({ + "budgets": [ + { "id": "a", "max_limit": 10, "current_usage": 1, "reset_duration": "1h" }, + { "id": "b", "max_limit": 10, "current_usage": 1, "reset_duration": "1d" }, + { "id": "c", "max_limit": 10, "current_usage": 1, "reset_duration": "1w" } + ], + "model_configs": [ + { "model_name": "m", "budgets": [{ "id": "d", "max_limit": 10, "current_usage": 1 }] } + ] + })); + let ids = result + .usage + .extra_rate_windows + .iter() + .map(|window| window.id.as_str()) + .collect::>(); + assert_eq!(ids, ["bifrost-budget-c", "bifrost-model-0-budget-d"]); +} + +#[test] +fn missing_current_usage_is_a_known_zero_budget_window() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10 }] + })); + assert_eq!(result.usage.primary.used_percent, 0.0); + assert!(result.usage.primary.usage_known()); +} + +#[test] +fn model_rows_normalize_names_and_format_token_counts() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": [ + { "model": "us.anthropic.claude-sonnet-4-20250514-v1:0", "provider": "bedrock", + "total_cost": 2.5, "total_tokens": 1_250_000 }, + { "model": "gpt-4o-2024-08-06", "provider": "bedrock", + "total_cost": 0.5, "total_tokens": 999 }, + { "model": "idle", "total_cost": 0, "total_tokens": 0 } + ] }] + })); + assert_eq!( + detail_rows(&result), + [ + ( + "claude-sonnet-4".to_owned(), + "$2.50".to_owned(), + Some("1.3M tokens".to_owned()) + ), + ( + "gpt-4o".to_owned(), + "$0.50".to_owned(), + Some("999 tokens".to_owned()) + ), + ] + ); +} + +#[test] +fn model_rows_prefix_provider_only_when_providers_are_mixed() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": [ + { "model": "eu.mistral.large-v1:0", "provider": "bedrock", "total_cost": 2 }, + { "model": "gpt-4o", "provider": "openai", "total_cost": 1 } + ] }] + })); + let titles = detail_rows(&result) + .into_iter() + .map(|row| row.0) + .collect::>(); + assert_eq!(titles, ["bedrock · large", "openai · gpt-4o"]); +} + +#[test] +fn model_rows_use_raw_name_when_short_names_collide() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": [ + { "model": "us.anthropic.claude-3-haiku-20240307-v1:0", "total_cost": 3 }, + { "model": "claude-3-haiku-20240307", "total_cost": 2 }, + { "model": "gpt-4o", "total_cost": 1 } + ] }] + })); + let titles = detail_rows(&result) + .into_iter() + .map(|row| row.0) + .collect::>(); + assert_eq!( + titles, + [ + "us.anthropic.claude-3-haiku-20240307-v1:0", + "claude-3-haiku-20240307", + "gpt-4o" + ] + ); +} + +#[test] +fn model_rows_cap_at_five_with_an_other_models_count() { + let models = (0..7) + .map(|index| json!({ "model": format!("m{index}"), "total_cost": 10 - index })) + .collect::>(); + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": models }] + })); + let rows = detail_rows(&result); + assert_eq!(rows.len(), 6); + assert_eq!(rows[5], ("Other models".to_owned(), "2".to_owned(), None)); +} From fce4584a6d0968c77ffabe15ddf4a68c59121094 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:13:06 +0700 Subject: [PATCH 7/9] Address thermo review --- rust/src/providers/bifrost/model_labels.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/src/providers/bifrost/model_labels.rs b/rust/src/providers/bifrost/model_labels.rs index 432361d0a4..8e8859a87a 100644 --- a/rust/src/providers/bifrost/model_labels.rs +++ b/rust/src/providers/bifrost/model_labels.rs @@ -25,7 +25,6 @@ pub(super) fn model_name(raw: &str) -> String { let name = strip_dotted_prefix(name, &VENDORS); let name = strip_revision(name); let name = strip_date(name); - let name = name.trim_end_matches([' ', '\t', '-']); if name.is_empty() { raw } else { name }.to_owned() } @@ -129,6 +128,7 @@ mod tests { ("gpt-4o", "gpt-4o"), ("global.qwen.qwen3-32b", "qwen3-32b"), ("usa.model", "usa.model"), + ("model-", "model-"), ("anthropic.", "anthropic."), ("-20250514", "-20250514"), ("model-v1:", "model-v1:"), From 250f1d8c3313911352fa68f09e554bc95f4151c1 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Fri, 2 Oct 2026 04:55:19 +0700 Subject: [PATCH 8/9] Fix the no-cookie-domain branch after the #619 merge --- .../src-tauri/src/commands/providers.rs | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index dd2b778721..2adfdffaf0 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -95,8 +95,8 @@ pub(crate) fn build_fetch_context( .cookie_domain() .is_none() { - let source_mode = if active_token_env.is_some() { - SourceMode::OAuth + let (source_mode, cookie_header, missing_cookie) = if active_token_env.is_some() { + (SourceMode::OAuth, None, false) } else { match cookie_source { // #433: an explicitly selected, non-empty Claude manual cookie is @@ -118,7 +118,7 @@ pub(crate) fn build_fetch_context( // API-key fallback. "off" | "manual" if provider.cookie_source_scopes_session_only() => { let cookie_header = if cookie_source == "manual" { - active_token_cookie.or(stored_cookie) + active_token_cookie.clone().or(stored_cookie) } else { None }; @@ -144,7 +144,7 @@ pub(crate) fn build_fetch_context( "off" if id == ProviderId::Factory => (SourceMode::Cli, None, false), "off" => (SourceMode::Cli, None, false), "manual" => { - let cookie_header = active_token_cookie.or(stored_cookie); + let cookie_header = active_token_cookie.clone().or(stored_cookie); let fails_closed_without_cookie = cookie_header.is_none() && provider.manual_empty_cookie_policy() == ManualEmptyCookiePolicy::FailClosedWeb; @@ -176,23 +176,24 @@ pub(crate) fn build_fetch_context( "auto" | "browser" | "web" => { // Claude resolves its cached cookie and browser fallback inside // the provider; other providers retain the shell fallback. - let cookie_header = active_token_cookie.or(stored_cookie).or_else(|| { - if defer_provider_browser_cookie_lookup { - None - } else { - provider_cookie_domain(id, settings).and_then(|domain| { - codexbar::browser::cookies::get_cookie_header(domain) - .ok() - .filter(|h| !h.is_empty()) - }) - } - }); + let cookie_header = + active_token_cookie.clone().or(stored_cookie).or_else(|| { + if defer_provider_browser_cookie_lookup { + None + } else { + provider_cookie_domain(id, settings).and_then(|domain| { + codexbar::browser::cookies::get_cookie_header(domain) + .ok() + .filter(|h| !h.is_empty()) + }) + } + }); (usage_source, cookie_header, false) } _ => (usage_source, stored_cookie, false), } }; - (source_mode, None, false) + (source_mode, cookie_header, missing_cookie) } else { match cookie_source { // #433: an explicitly selected, non-empty Claude manual cookie is From 4663218def72a806c3bd35053d8e2270426611e2 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Fri, 2 Oct 2026 07:28:46 +0700 Subject: [PATCH 9/9] Fix the #619 merge resolution in build_fetch_context --- .../src-tauri/src/commands/providers.rs | 28 +++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 2adfdffaf0..2a487504e4 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -209,6 +209,24 @@ pub(crate) fn build_fetch_context( (SourceMode::Web, stored_cookie.clone(), false) } _ if active_token_env.is_some() => (SourceMode::OAuth, None, false), + // Charm Hyper: the cookie source only picks the session, and + // the usage source keeps routing. Off and an empty Manual + // source never import a browser session, while Auto keeps its + // API-key fallback. + "off" | "manual" if provider.cookie_source_scopes_session_only() => { + let cookie_header = if cookie_source == "manual" { + active_token_cookie.clone().or(stored_cookie) + } else { + None + }; + let source_mode = if provider.available_sources().contains(&usage_source) { + usage_source + } else { + SourceMode::Auto + }; + let cookie_missing = cookie_header.is_none(); + (source_mode, cookie_header, cookie_missing) + } "off" if provider_uses_oauth_without_cookies(id, usage_source) => { (SourceMode::OAuth, None, false) } @@ -274,8 +292,14 @@ pub(crate) fn build_fetch_context( // Cookie-web providers (Cursor, OpenCode, …) reject SourceMode::Cli. The shell // historically mapped "manual + no cookie" to Cli, which surfaces as // "Source mode 'Cli' not supported". Remap to Web and try browser cookies - // unless the user explicitly disabled cookies ("off"). - if source_mode == SourceMode::Cli && cookie_source != "off" && !provider.supports_cli() { + // unless the user explicitly disabled cookies ("off"). Providers whose + // cookie source only scopes the session (Charm Hyper) own this contract in + // the provider, so the shell must not remap their source mode. + if source_mode == SourceMode::Cli + && cookie_source != "off" + && !provider.supports_cli() + && !provider.cookie_source_scopes_session_only() + { if cookie_header .as_deref() .map(str::trim)