diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 779f9c801c..2c2d251fe2 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -84,98 +84,100 @@ 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 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 @@ -220,16 +222,51 @@ 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 (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) + } + ( + 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, + 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..7b75cada51 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, 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,70 @@ 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()); + 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 + .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; + + if let Some(source_mode) = effective_source_mode { + ctx.source_mode = source_mode; + } +} + /// 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..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; @@ -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,21 @@ impl TokenAccountOverride { account, env_override, cookie_header, + 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, } } } @@ -796,6 +890,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 +900,132 @@ 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 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 = 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..b4f922b065 100644 --- a/rust/src/providers/opencodego/mod.rs +++ b/rust/src/providers/opencodego/mod.rs @@ -477,6 +477,12 @@ impl Provider for OpenCodeGoProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { tracing::debug!("Fetching OpenCode Go usage"); + 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; + } + match ctx.source_mode { SourceMode::Auto => { // Local-first unless workspace/token scope asks for web first @@ -513,12 +519,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 } } @@ -537,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, 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!(