diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 098ea52710..8846ec7a20 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -29,7 +29,8 @@ 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 api_key = stored_api_key.or(active_token_api_key); + // 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 = id == ProviderId::Kimi && api_key.as_deref().is_some_and(|key| !key.trim().is_empty()); diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index b1e2877f4c..da25cdd81f 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -579,6 +579,51 @@ fn fetch_context_token_account_takes_precedence_over_manual_cookie() { ); } +#[test] +fn fetch_context_openrouter_token_account_overrides_stored_api_key() { + let settings = Settings::default(); + let cookies = ManualCookies::default(); + let mut api_keys = ApiKeys::default(); + api_keys.set("openrouter", "sk-or-v1-stored-decoy", None); + let mut token_accounts = HashMap::new(); + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Personal", "sk-or-v1-personal")); + data.add_account(TokenAccount::new("Work", "sk-or-v1-work")); + data.set_active(1); + token_accounts.insert(ProviderId::OpenRouter, data); + + let ctx = super::build_fetch_context( + ProviderId::OpenRouter, + &settings, + &cookies, + &api_keys, + &token_accounts, + ); + + assert_eq!(ctx.source_mode, SourceMode::OAuth); + assert!(ctx.manual_cookie_header.is_none()); + assert_eq!(ctx.api_key.as_deref(), Some("sk-or-v1-work")); +} + +#[test] +fn fetch_context_openrouter_falls_back_to_stored_api_key_without_token_accounts() { + let settings = Settings::default(); + let cookies = ManualCookies::default(); + let mut api_keys = ApiKeys::default(); + api_keys.set("openrouter", "sk-or-v1-stored", None); + let token_accounts = HashMap::new(); + + let ctx = super::build_fetch_context( + ProviderId::OpenRouter, + &settings, + &cookies, + &api_keys, + &token_accounts, + ); + + assert_eq!(ctx.api_key.as_deref(), Some("sk-or-v1-stored")); +} + #[test] fn provider_region_set_rejects_non_regional_provider() { let mut s = Settings::default(); diff --git a/rust/src/cli/mod.rs b/rust/src/cli/mod.rs index 84288a9d1f..a3a00f4a06 100755 --- a/rust/src/cli/mod.rs +++ b/rust/src/cli/mod.rs @@ -88,6 +88,10 @@ pub struct Cli { #[arg(long = "all-accounts")] pub all_accounts: bool, + /// Token-account label or 1-based index (requires a single provider) + #[arg(long = "account")] + pub account: Option, + /// Skip credits line in output #[arg(long = "no-credits")] pub no_credits: bool, @@ -161,6 +165,7 @@ impl Cli { pretty: self.pretty, status: self.status, all_accounts: self.all_accounts, + account: self.account.clone(), source: self.source.clone(), web_timeout: self.web_timeout, web_debug_dump_html: self.web_debug_dump_html, diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index ecd2eff79a..1d787f61f6 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -4,9 +4,10 @@ use clap::Args; use serde::Serialize; use crate::core::{ - CostSnapshot, FetchContext, ProviderFetchResult, ProviderId, RateWindow, SourceMode, UsagePace, - UsageSnapshot, instantiate_provider, + CostSnapshot, FetchContext, ProviderFetchResult, ProviderId, RateWindow, SourceMode, + TokenAccountStore, TokenAccountSupport, UsagePace, UsageSnapshot, instantiate_provider, }; +use crate::settings::ApiKeys; use crate::status::{ProviderStatus as StatusInfo, StatusLevel, fetch_provider_status}; pub const PROVIDER_ARG_HELP: &str = "Provider to query (for example: codex, claude, gemini, antigravity/agy, nanogpt, deepseek, codebuff, windsurf, all, both)"; @@ -45,6 +46,10 @@ pub struct UsageArgs { #[arg(long = "all-accounts")] pub all_accounts: bool, + /// Token-account label or 1-based index (requires a single provider) + #[arg(long = "account")] + pub account: Option, + /// Data source: auto, oauth, web, cli #[arg(long, default_value = "auto", value_parser = ["auto", "web", "cli", "oauth"])] pub source: String, @@ -155,6 +160,8 @@ struct UsageCommand { brief: bool, fetch_status: bool, pretty: bool, + /// Optional token-account label/index for a single-provider fetch. + account: Option, ctx: FetchContext, } @@ -163,6 +170,9 @@ impl UsageCommand { let format = effective_format(&args); let source_mode = SourceMode::parse(&args.source).unwrap_or(SourceMode::Auto); let providers = ProviderSelection::from_arg(args.provider.as_deref())?.as_list(); + if args.account.is_some() && providers.len() != 1 { + anyhow::bail!("--account requires a single --provider (not all/both)"); + } Ok(Self { format, @@ -171,6 +181,7 @@ impl UsageCommand { brief: args.brief, fetch_status: args.status, pretty: args.pretty, + account: args.account.clone(), ctx: build_usage_fetch_context(&args, source_mode), }) } @@ -272,7 +283,11 @@ async fn fetch_provider_result( let status_future = command .fetch_status .then(|| fetch_provider_status(provider_id.cli_name())); - let result = provider.fetch_usage(&command.ctx).await?; + let mut ctx = command.ctx.clone(); + if 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?; let status = if let Some(fut) = status_future { fut.await } else { @@ -281,6 +296,63 @@ async fn fetch_provider_result( Ok((result, status)) } +/// 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 +/// providers (OpenRouter, z.ai, ...) honor the selected labeled account. +fn resolve_cli_api_key( + provider_id: ProviderId, + account_ref: Option<&str>, +) -> anyhow::Result> { + if TokenAccountSupport::is_supported(provider_id) + && let Ok(data) = TokenAccountStore::new().load_provider(provider_id) + && !data.accounts.is_empty() + { + 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_id.display_name() + ) + })? + }; + if let Some(env) = TokenAccountSupport::env_override(provider_id, &account.token) + && let Some(key) = env.into_values().next() + { + return Ok(Some(key)); + } + } + + Ok(ApiKeys::load() + .get(provider_id.cli_name()) + .map(|s| s.to_string())) +} + +fn find_token_account<'a>( + data: &'a crate::core::ProviderAccountData, + account_ref: &str, +) -> anyhow::Result<&'a crate::core::TokenAccount> { + if let Ok(idx) = account_ref.parse::() + && idx > 0 + && idx <= data.accounts.len() + { + return Ok(&data.accounts[idx - 1]); + } + if let Some(account) = data + .accounts + .iter() + .find(|a| a.label.eq_ignore_ascii_case(account_ref)) + { + return Ok(account); + } + anyhow::bail!( + "Account '{}' not found. Use 'codexbar account list ' to see accounts.", + account_ref + ) +} + fn render_text_error(provider_id: ProviderId, error_msg: &str, use_color: bool) -> String { let header = if use_color { format!("\x1b[1m{}\x1b[0m", provider_id.display_name()) @@ -573,11 +645,30 @@ fn render_progress_bar(percent: f64, width: usize, use_color: bool) -> String { #[cfg(test)] mod tests { use super::*; + use crate::core::{ProviderAccountData, TokenAccount, TokenAccountSupport}; fn fetch_result(usage: UsageSnapshot) -> ProviderFetchResult { ProviderFetchResult::new(usage, "test") } + #[test] + fn openrouter_account_ref_resolves_labeled_key() { + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Personal", "sk-or-v1-personal")); + data.add_account(TokenAccount::new("Work", "sk-or-v1-work")); + data.set_active(0); + + let work = find_token_account(&data, "Work").unwrap(); + let env = TokenAccountSupport::env_override(ProviderId::OpenRouter, &work.token).unwrap(); + assert_eq!( + env.get("OPENROUTER_API_KEY").map(String::as_str), + Some("sk-or-v1-work") + ); + + let by_index = find_token_account(&data, "2").unwrap(); + assert_eq!(by_index.token, "sk-or-v1-work"); + } + #[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/token_accounts.rs b/rust/src/core/token_accounts.rs index fa9a4ac74a..3b53b8b46e 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -209,6 +209,17 @@ impl TokenAccountSupport { requires_manual_cookie_source: false, cookie_name: None, }), + // Upstream 0.45 #2271: labeled OpenRouter API keys via token accounts. + ProviderId::OpenRouter => Some(TokenAccountSupport { + title: "API keys", + subtitle: "Store multiple OpenRouter API keys.", + placeholder: "sk-or-v1-...", + injection: TokenInjection::Environment { + key: "OPENROUTER_API_KEY".to_string(), + }, + requires_manual_cookie_source: false, + cookie_name: None, + }), ProviderId::Copilot => Some(TokenAccountSupport { title: "GitHub accounts", subtitle: "Store GitHub OAuth tokens for Copilot plan usage.", @@ -230,7 +241,6 @@ impl TokenAccountSupport { | ProviderId::JetBrains | ProviderId::Warp | ProviderId::AzureOpenAI - | ProviderId::OpenRouter | ProviderId::NanoGPT | ProviderId::Infini | ProviderId::Perplexity @@ -646,10 +656,48 @@ mod tests { assert!(TokenAccountSupport::is_supported(ProviderId::Claude)); assert!(TokenAccountSupport::is_supported(ProviderId::Cursor)); assert!(TokenAccountSupport::is_supported(ProviderId::Copilot)); + assert!(TokenAccountSupport::is_supported(ProviderId::OpenRouter)); assert!(!TokenAccountSupport::is_supported(ProviderId::Codex)); assert!(!TokenAccountSupport::is_supported(ProviderId::Gemini)); } + #[test] + fn openrouter_token_accounts_inject_api_key_env() { + let support = TokenAccountSupport::for_provider(ProviderId::OpenRouter).unwrap(); + assert_eq!(support.title, "API keys"); + assert_eq!(support.placeholder, "sk-or-v1-..."); + assert!(!support.requires_manual_cookie_source); + match &support.injection { + TokenInjection::Environment { key } => assert_eq!(key, "OPENROUTER_API_KEY"), + other => panic!("expected environment injection, got {other:?}"), + } + + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Personal", "sk-or-v1-personal")); + data.add_account(TokenAccount::new("Work", "sk-or-v1-work")); + data.set_active(1); + + let active = data.active_account().unwrap(); + assert_eq!(active.label, "Work"); + let env = TokenAccountSupport::env_override(ProviderId::OpenRouter, &active.token).unwrap(); + assert_eq!( + env.get("OPENROUTER_API_KEY").map(String::as_str), + Some("sk-or-v1-work") + ); + + let override_data = + TokenAccountOverride::from_account(ProviderId::OpenRouter, active.clone()); + assert_eq!( + override_data + .env_override + .as_ref() + .and_then(|m| m.get("OPENROUTER_API_KEY")) + .map(String::as_str), + Some("sk-or-v1-work") + ); + assert!(override_data.cookie_header.is_none()); + } + #[test] fn test_claude_oauth_detection() { assert!(TokenAccountSupport::is_claude_oauth_token(