From b6bb1e150cfd6dc090a7510bcd55c7268d903442 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 16 Jul 2026 12:03:53 -0700 Subject: [PATCH 01/24] feat: remove profiles from the authentication system first draft BREAKING-CHANGE: --profile and --prefer-profile removed Breaks any script using them --- README.md | 36 +- src/args.rs | 74 +- src/auth.rs | 1673 +++++++++++++++++++++++++++---------- src/config/mod.rs | 48 +- src/datasets/pipeline.rs | 3 +- src/datasets/snapshots.rs | 17 +- src/eval.rs | 30 +- src/functions/push.rs | 3 +- src/init.rs | 26 +- src/main.rs | 9 +- src/setup/mod.rs | 186 ++--- src/status.rs | 99 ++- src/switch.rs | 114 +-- src/traces.rs | 82 +- src/utils/mod.rs | 2 +- src/utils/profile.rs | 37 +- tests/cli.rs | 21 +- tests/functions.rs | 59 +- 18 files changed, 1545 insertions(+), 974 deletions(-) diff --git a/README.md b/README.md index cfa19e70..237e27d8 100644 --- a/README.md +++ b/README.md @@ -312,34 +312,34 @@ Local version and pagination-key conversion helpers: ## `bt auth` -- Authenticate interactively (prompts for auth method, profile name defaults to org name): +- Authenticate interactively (prompts for auth method and organization): - `bt auth login` - First prompt chooses: `OAuth (browser)` (default) or `API key`. - - If your API key can access multiple orgs, `bt` uses a searchable picker (alphabetized) and lets you choose a specific org or no default org (cross-org mode). - - After login, `bt` updates the active profile/org context immediately. If `--project` is set, it also switches that project; otherwise it clears any stale default project for the new login. + - OAuth can be saved for an org or in cross-org mode. API-key logins are saved per org; if a key can access multiple orgs, `bt` uses a searchable org picker. + - After login, `bt` updates the active org context immediately. If `--project` is set, it also switches that project; otherwise it clears any stale default project for the new login. - `bt` confirms the resolved API URL before saving. - Login with OAuth (browser-based, stores refresh token in secure credential store): - - `bt auth login --oauth --profile work` + - `bt auth login --oauth --org myorg` - You can pass `--no-browser` to print the URL without auto-opening. - On remote/SSH hosts, paste the final callback URL from your local browser if localhost callback cannot be delivered. -- List profiles: +- List saved auth logins: - `bt auth profiles` -- Log out (remove a saved profile): - - `bt auth logout` +- Log out: + - `bt auth logout --org myorg` + - `bt auth logout --org myorg --api-key-hint sk-****abcde` - `bt auth logout --force` (skip confirmation) -- Show current auth source/profile: - - `bt auth status` +- Show current auth context: + - `bt status` - Force-refresh OAuth access token for debugging: - - `bt auth refresh --profile work` + - `bt auth refresh --org myorg` Auth resolution order for commands is: -1. Explicit `--profile` -2. `--api-key` or `BRAINTRUST_API_KEY` (unless `--prefer-profile` is set) -3. `BRAINTRUST_PROFILE` -4. Org-based profile match (profile whose org matches `--org`/config org) -5. Single-profile auto-select (if only one profile exists) -6. Interactive profile picker (if multiple profiles exist and a TTY is available) +1. Explicit `--api-key sk-...` +2. `--prefer-api-key` / `BRAINTRUST_PREFER_API_KEY` (uses `BRAINTRUST_API_KEY` first, then a stored API key for the selected org, then falls back to OAuth) +3. Stored OAuth login for the selected org (or cross-org OAuth when selected) +4. Stored API key login for the selected org when no OAuth login is available +5. `BRAINTRUST_API_KEY` On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Service daemon. On macOS, it uses the `security` keychain utility. If a secure store is unavailable, `bt` falls back to a plaintext secrets file with `0600` permissions. @@ -357,8 +357,8 @@ Interactively switch org and project context: Show current org and project context: -- `bt status` — display current org, project, and config source -- `bt status --verbose` — show detailed config resolution +- `bt status` — display current org, project, selected auth method, and config source +- `bt status --verbose` — show detailed config and auth resolution - `bt status -j` — JSON output ## `bt setup` and `bt docs` diff --git a/src/args.rs b/src/args.rs index c15c6478..f0f20d64 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,4 +1,3 @@ -use std::ffi::OsString; use std::path::{Path, PathBuf}; use clap::Args; @@ -39,13 +38,9 @@ pub struct BaseArgs { #[arg(long, env = "BRAINTRUST_NO_INPUT", global = true, value_parser = clap::builder::BoolishValueParser::new(), default_value_t = false)] pub no_input: bool, - /// Use a saved login profile (or via BRAINTRUST_PROFILE) - #[arg(long, env = "BRAINTRUST_PROFILE", global = true)] + #[arg(skip)] pub profile: Option, - #[arg(skip = false)] - pub profile_explicit: bool, - /// Override active org (or via BRAINTRUST_ORG_NAME) #[arg(short = 'o', long = "org", env = "BRAINTRUST_ORG_NAME", global = true)] pub org_name: Option, @@ -67,9 +62,9 @@ pub struct BaseArgs { #[arg(skip)] pub api_key_source: Option, - /// Prefer profile credentials even if BRAINTRUST_API_KEY/--api-key is set. - #[arg(long, global = true)] - pub prefer_profile: bool, + /// Prefer API key credentials for the selected org when available. + #[arg(long = "prefer-api-key", env = "BRAINTRUST_PREFER_API_KEY", global = true, value_parser = clap::builder::BoolishValueParser::new(), default_value_t = false)] + pub prefer_api_key: bool, /// Override API URL (or via BRAINTRUST_API_URL) #[arg( @@ -126,64 +121,3 @@ impl BaseArgs { self.verbose && self.verbose_source.is_some() } } - -pub fn has_explicit_profile_arg(args: &[OsString]) -> bool { - let mut idx = 1usize; - while idx < args.len() { - let Some(arg) = args[idx].to_str() else { - idx += 1; - continue; - }; - - if arg == "--" { - break; - } - - if arg == "--profile" || arg.starts_with("--profile=") { - return true; - } - - idx += 1; - } - - false -} - -#[cfg(test)] -mod tests { - use super::has_explicit_profile_arg; - use std::ffi::OsString; - - #[test] - fn has_explicit_profile_arg_detects_split_flag() { - let args = vec![ - OsString::from("bt"), - OsString::from("status"), - OsString::from("--profile"), - OsString::from("work"), - ]; - assert!(has_explicit_profile_arg(&args)); - } - - #[test] - fn has_explicit_profile_arg_detects_equals_flag() { - let args = vec![ - OsString::from("bt"), - OsString::from("status"), - OsString::from("--profile=work"), - ]; - assert!(has_explicit_profile_arg(&args)); - } - - #[test] - fn has_explicit_profile_arg_ignores_passthrough_args() { - let args = vec![ - OsString::from("bt"), - OsString::from("eval"), - OsString::from("--"), - OsString::from("--profile"), - OsString::from("work"), - ]; - assert!(!has_explicit_profile_arg(&args)); - } -} diff --git a/src/auth.rs b/src/auth.rs index 44db3d77..32f1ae1d 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1,7 +1,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::error::Error as StdError; use std::fs; -use std::io::{IsTerminal, Write}; +use std::io::Write; use std::net::TcpListener; use std::path::{Path, PathBuf}; use std::process::Command; @@ -25,6 +25,7 @@ use oauth2::{ }; use serde::{Deserialize, Serialize}; use serde_json::json; +use sha2::{Digest, Sha256}; use tokio::sync::oneshot; use crate::{ @@ -62,6 +63,7 @@ pub struct ResolvedAuth { #[derive(Debug, Clone)] pub struct ProfileInfo { pub name: String, + pub auth_method: String, pub org_name: Option, pub user_name: Option, pub email: Option, @@ -71,7 +73,6 @@ pub struct ProfileInfo { #[derive(Debug, Clone)] pub(crate) struct StoredProfileInfo { pub name: String, - pub is_oauth: bool, pub org_name: Option, } @@ -84,7 +85,6 @@ pub struct AvailableOrg { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RecoverableAuthErrorKind { - OauthProfileSelection, OauthClientId, OauthRefreshToken, StoredCredential, @@ -115,8 +115,7 @@ pub fn is_missing_credential_error(err: &anyhow::Error) -> bool { .is_some_and(|err| { matches!( err.kind, - RecoverableAuthErrorKind::OauthProfileSelection - | RecoverableAuthErrorKind::OauthClientId + RecoverableAuthErrorKind::OauthClientId | RecoverableAuthErrorKind::OauthRefreshToken | RecoverableAuthErrorKind::StoredCredential ) @@ -129,13 +128,7 @@ pub fn list_profiles() -> Result> { Ok(store .profiles .iter() - .map(|(name, p)| ProfileInfo { - name: name.clone(), - org_name: p.org_name.clone(), - user_name: p.user_name.clone(), - email: p.email.clone(), - api_key_hint: p.api_key_hint.clone(), - }) + .map(|(name, p)| profile_info_from_store_entry(name, p)) .collect()) } @@ -146,12 +139,12 @@ pub(crate) fn list_stored_profiles() -> Result> { .iter() .map(|(name, profile)| StoredProfileInfo { name: name.clone(), - is_oauth: profile.auth_kind == AuthKind::Oauth, org_name: profile.org_name.clone(), }) .collect()) } +#[cfg(test)] pub fn resolve_org_to_profile(identifier: &str, profiles: &[ProfileInfo]) -> Result { if profiles.is_empty() { bail!("no auth profiles found. Run `bt auth login` to create one."); @@ -185,7 +178,7 @@ pub fn resolve_org_to_profile(identifier: &str, profiles: &[ProfileInfo]) -> Res _ => { if !ui::can_prompt() { bail!( - "multiple profiles for org '{identifier}': {}. Use --profile to disambiguate.", + "multiple auth logins for org '{identifier}': {}. Use --org to disambiguate.", matches .iter() .map(|p| p.name.as_str()) @@ -204,22 +197,38 @@ pub fn resolve_org_to_profile(identifier: &str, profiles: &[ProfileInfo]) -> Res } } +fn profile_info_identity_label(profile: &ProfileInfo) -> Option { + if let Some(email) = profile.email.as_deref() { + return match profile.user_name.as_deref() { + Some(name) => Some(format!("{name} ({email})")), + None => Some(email.to_string()), + }; + } + profile.api_key_hint.clone() +} + +pub(crate) fn profile_info_label(profile: &ProfileInfo) -> String { + let mut parts = vec![profile + .org_name + .clone() + .unwrap_or_else(|| "cross-org".to_string())]; + parts.push(profile.auth_method.clone()); + if let Some(identity) = profile_info_identity_label(profile) { + parts.push(identity); + } + parts.join(" — ") +} + pub fn select_profile_interactive(current: Option<&str>) -> Result> { let profiles = list_profiles()?; if profiles.is_empty() { - bail!("no auth profiles found. Run `bt auth login` to create one."); + bail!("no auth logins found. Run `bt auth login` to create one."); } if profiles.len() == 1 { return Ok(Some(profiles[0].name.clone())); } - let labels: Vec = profiles - .iter() - .map(|p| match &p.org_name { - Some(org) if org != &p.name => format!("{} (profile: {})", org, p.name), - _ => p.name.clone(), - }) - .collect(); + let labels: Vec = profiles.iter().map(profile_info_label).collect(); let default = current .and_then(|c| { @@ -228,7 +237,7 @@ pub fn select_profile_interactive(current: Option<&str>) -> Result, #[serde(default)] + org_id: Option, + #[serde(default)] org_name: Option, #[serde(default)] oauth_client_id: Option, @@ -317,10 +328,14 @@ struct AuthProfile { #[serde(default)] email: Option, #[serde(default)] + api_key_hash: Option, + #[serde(default)] api_key_hint: Option, + #[serde(default)] + legacy_secret_key: Option, } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] #[serde(rename_all = "snake_case")] enum AuthKind { #[default] @@ -328,6 +343,13 @@ enum AuthKind { Oauth, } +fn auth_kind_label(kind: AuthKind) -> &'static str { + match kind { + AuthKind::ApiKey => "api_key", + AuthKind::Oauth => "oauth", + } +} + #[derive(Debug, Clone, Deserialize)] struct ApiKeyLoginResponse { org_info: Vec, @@ -377,8 +399,8 @@ struct OAuthErrorResponse { Examples: bt auth login bt auth profiles - bt auth refresh - bt auth logout --profile work + bt auth refresh --org acme + bt auth logout --org acme ")] pub struct AuthArgs { #[command(subcommand)] @@ -389,20 +411,16 @@ pub struct AuthArgs { enum AuthCommand { /// Authenticate with Braintrust (OAuth or API key) Login(AuthLoginArgs), - /// Force-refresh OAuth access token for a profile + /// Force-refresh OAuth access token for the selected org Refresh, - /// List auth profiles and check connection status + /// List saved auth logins and check connection status Profiles(AuthProfilesArgs), - /// Log out by removing a saved profile + /// Log out by removing a saved auth login Logout(AuthLogoutArgs), } #[derive(Debug, Clone, Args)] -struct AuthProfilesArgs { - /// Only show the profile with this name - #[arg(long, value_name = "NAME")] - profile: Option, -} +struct AuthProfilesArgs {} #[derive(Debug, Clone, Args)] struct AuthLoginArgs { @@ -410,7 +428,7 @@ struct AuthLoginArgs { #[arg(long)] oauth: bool, - /// OAuth client id (defaults to bt_cli_) + /// OAuth client id (defaults to bt_cli) #[arg(long, value_name = "CLIENT_ID")] client_id: Option, @@ -421,9 +439,9 @@ struct AuthLoginArgs { #[derive(Debug, Clone, Args)] struct AuthLogoutArgs { - /// Profile name to log out of (interactive picker if omitted) - #[arg(long)] - profile: Option, + /// API key hint to log out of when multiple API keys exist for an org + #[arg(long = "api-key-hint", value_name = "HINT")] + api_key_hint: Option, /// Skip confirmation prompt #[arg(long, short = 'f')] @@ -524,11 +542,8 @@ pub async fn login(base: &BaseArgs) -> Result { } let login = match builder.build().await { Ok(client) => client.wait_for_login().await?, - Err(err) if auth.is_oauth => { - let org_name = auth - .org_name - .clone() - .ok_or_else(|| anyhow::anyhow!("oauth profile is missing org_name: {err}"))?; + Err(_err) if auth.is_oauth => { + let org_name = auth.org_name.clone().unwrap_or_default(); let login = LoginState::new(); login.set( api_key.clone(), @@ -797,42 +812,27 @@ fn has_cached_project_id(base: &BaseArgs) -> bool { .is_some_and(|project_id| !project_id.trim().is_empty()) } -fn maybe_warn_api_key_override(base: &BaseArgs) { - if base.json || !std::io::stderr().is_terminal() { - return; - } - if resolve_api_key_override(base).is_none() { - return; - } +fn maybe_warn_api_key_override(_base: &BaseArgs) {} - let ignored_profile = base - .profile - .as_ref() - .map(|value| value.trim()) - .filter(|value| !value.is_empty()); - - if let Some(profile_name) = ignored_profile { - eprintln!( - "Info: using --api-key/BRAINTRUST_API_KEY credentials; selected profile '{profile_name}' is ignored for this command. Use --prefer-profile or unset BRAINTRUST_API_KEY to use a profile with OAuth login.", - ); +fn resolve_cli_api_key_override(base: &BaseArgs) -> Option { + if matches!( + base.api_key_source, + Some(crate::args::ArgValueSource::EnvVariable) + ) { + return None; } + let value = base.api_key.as_deref()?.trim(); + if value.is_empty() { + return None; + } + Some(value.to_string()) } -fn has_explicit_profile_selection(base: &BaseArgs) -> bool { - base.profile_explicit - && base - .profile - .as_deref() - .is_some_and(|value| !value.trim().is_empty()) -} - -fn resolve_api_key_override(base: &BaseArgs) -> Option { - if (base.prefer_profile || has_explicit_profile_selection(base)) - && !matches!( - base.api_key_source, - Some(crate::args::ArgValueSource::CommandLine) - ) - { +fn resolve_env_api_key(base: &BaseArgs) -> Option { + if !matches!( + base.api_key_source, + Some(crate::args::ArgValueSource::EnvVariable) + ) { return None; } let value = base.api_key.as_deref()?.trim(); @@ -842,115 +842,278 @@ fn resolve_api_key_override(base: &BaseArgs) -> Option { Some(value.to_string()) } -fn config_auth_context(base: &BaseArgs) -> (Option, Option) { +#[cfg(test)] +fn resolve_api_key_override(base: &BaseArgs) -> Option { + resolve_cli_api_key_override(base) +} + +fn config_auth_context(base: &BaseArgs) -> Option { let cfg = crate::config::load().unwrap_or_default(); config_auth_context_from_config(base, &cfg) } -fn config_auth_context_from_config( - base: &BaseArgs, - cfg: &crate::config::Config, -) -> (Option, Option) { - let profile = if crate::config::trimmed_option(base.profile.as_deref()).is_none() { - crate::config::trimmed_option(cfg.profile.as_deref()).map(str::to_string) - } else { - None - }; - - let org = if crate::config::trimmed_option(base.org_name.as_deref()).is_none() { +fn config_auth_context_from_config(base: &BaseArgs, cfg: &crate::config::Config) -> Option { + if crate::config::trimmed_option(base.org_name.as_deref()).is_none() { crate::config::trimmed_option(cfg.org.as_deref()).map(str::to_string) } else { None - }; + } +} - (profile, org) +fn effective_org_name<'a>(base: &'a BaseArgs, cfg_org: &'a Option) -> Option<&'a str> { + crate::config::trimmed_option(base.org_name.as_deref()) + .or_else(|| crate::config::trimmed_option(cfg_org.as_deref())) } pub async fn resolve_auth(base: &BaseArgs) -> Result { let mut store = load_auth_store()?; - let mut auth_base = base.clone(); - let (cfg_profile, cfg_org) = config_auth_context(base); - if let Some(profile) = cfg_profile { - auth_base.profile = Some(profile); + let cfg_org = config_auth_context(base); + + if let Some(api_key) = resolve_cli_api_key_override(base) { + return Ok(ResolvedAuth { + api_key: Some(api_key), + api_url: base.api_url.clone(), + app_url: base.app_url.clone(), + org_name: effective_org_name(base, &cfg_org).map(str::to_string), + is_oauth: false, + }); + } + + if base.prefer_api_key { + if let Some(auth) = resolve_preferred_api_key_auth(base, &mut store, &cfg_org)? { + return Ok(auth); + } + if let Some(profile_name) = + select_oauth_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? + { + return resolve_oauth_profile_auth(base, &mut store, &cfg_org, &profile_name).await; + } + bail!("--prefer-api-key requires an API key or OAuth login for the selected org"); } if let Some(profile_name) = - maybe_select_profile_for_auth(&auth_base, &store, &cfg_org, ui::can_prompt())? + select_oauth_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? { - auth_base.profile = Some(profile_name); + return resolve_oauth_profile_auth(base, &mut store, &cfg_org, &profile_name).await; } - let mut auth = resolve_auth_from_store_with_secret_lookup( - &auth_base, - &store, - load_profile_secret, - &cfg_org, - )?; - if !auth.is_oauth { - return Ok(auth); + if let Some(profile_name) = + select_api_key_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? + { + return resolve_api_key_profile_auth(base, &mut store, &cfg_org, &profile_name); } - let effective_org = auth_base.org_name.as_deref().or(cfg_org.as_deref()); - let profile_name = auth_base - .profile + if let Some(api_key) = resolve_env_api_key(base) { + return Ok(ResolvedAuth { + api_key: Some(api_key), + api_url: base.api_url.clone(), + app_url: base.app_url.clone(), + org_name: effective_org_name(base, &cfg_org).map(str::to_string), + is_oauth: false, + }); + } + + if effective_org_name(base, &cfg_org).is_none() { + if let Some(err) = missing_org_for_stored_logins_error(&store) { + return Err(err); + } + } + + Ok(ResolvedAuth { + api_key: None, + api_url: base.api_url.clone(), + app_url: base.app_url.clone(), + org_name: effective_org_name(base, &cfg_org).map(str::to_string), + is_oauth: false, + }) +} + +fn resolve_preferred_api_key_auth( + base: &BaseArgs, + store: &mut AuthStore, + cfg_org: &Option, +) -> Result> { + let Some(org) = effective_org_name(base, cfg_org) else { + bail!( + "--prefer-api-key requires an org; pass --org or run `bt switch /`" + ); + }; + + if let Some(api_key) = resolve_env_api_key(base) { + return Ok(Some(ResolvedAuth { + api_key: Some(api_key), + api_url: base.api_url.clone(), + app_url: base.app_url.clone(), + org_name: Some(org.to_string()), + is_oauth: false, + })); + } + + if let Some(profile_name) = + select_api_key_profile_for_auth(base, store, cfg_org, ui::can_prompt())? + { + return resolve_api_key_profile_auth(base, store, cfg_org, &profile_name).map(Some); + } + + Ok(None) +} + +fn resolve_api_key_profile_auth( + base: &BaseArgs, + store: &mut AuthStore, + cfg_org: &Option, + profile_name: &str, +) -> Result { + let profile = store.profiles.get(profile_name).cloned().ok_or_else(|| { + anyhow::anyhow!("auth login '{profile_name}' not found; run `bt auth profiles`") + })?; + let api_key = load_profile_secret_with_legacy( + profile_name, + profile.legacy_secret_key.as_deref(), + )? + .ok_or_else(|| { + recoverable_auth_error( + RecoverableAuthErrorKind::StoredCredential, + format!( + "no keychain credential found for auth login '{}'; re-run `bt auth login --org --api-key `", + auth_slot_label(profile_name, &profile) + ), + ) + })?; + + let resolved = ResolvedAuth { + api_key: Some(api_key.clone()), + api_url: base.api_url.clone().or_else(|| profile.api_url.clone()), + app_url: base.app_url.clone().or_else(|| profile.app_url.clone()), + org_name: effective_org_name(base, cfg_org) + .map(str::to_string) + .or_else(|| profile.org_name.clone()), + is_oauth: false, + }; + + maybe_rekey_api_key_profile_after_secret_load(store, profile_name, &api_key)?; + Ok(resolved) +} + +fn maybe_rekey_api_key_profile_after_secret_load( + store: &mut AuthStore, + profile_name: &str, + api_key: &str, +) -> Result<()> { + let Some(profile) = store.profiles.get(profile_name).cloned() else { + return Ok(()); + }; + if profile.auth_kind != AuthKind::ApiKey { + return Ok(()); + } + let Some(org_id) = profile + .org_id .as_deref() - .filter(|value| !value.trim().is_empty()) - .or_else(|| effective_org.and_then(|org| resolve_profile_for_org(org, &store))) - .or_else(|| { - (store.profiles.len() == 1).then(|| store.profiles.keys().next().unwrap().as_str()) - }) - .ok_or_else(|| { - recoverable_auth_error( - RecoverableAuthErrorKind::OauthProfileSelection, - "oauth profile requested but none selected".to_string(), - ) - })? - .to_string(); + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + return Ok(()); + }; + + let hash = api_key_hash(api_key); + let new_key = api_key_slot_key(&hash, org_id); + let already_canonical = + profile_name == new_key && profile.api_key_hash.as_deref() == Some(&hash); + if already_canonical { + return Ok(()); + } + + let mut updated = profile; + updated.api_key_hash = Some(hash); + if profile_name != new_key && updated.legacy_secret_key.is_none() { + updated.legacy_secret_key = Some(profile_name.to_string()); + } + + if profile_name != new_key { + store.profiles.remove(profile_name); + } + store.profiles.insert(new_key, updated); + save_auth_store(store) +} + +async fn resolve_oauth_profile_auth( + base: &BaseArgs, + store: &mut AuthStore, + cfg_org: &Option, + profile_name: &str, +) -> Result { let profile = store .profiles - .get(profile_name.as_str()) - .ok_or_else(|| anyhow::anyhow!("profile '{profile_name}' not found"))?; + .get(profile_name) + .cloned() + .ok_or_else(|| anyhow::anyhow!("oauth login '{profile_name}' not found"))?; let client_id = profile.oauth_client_id.as_deref().ok_or_else(|| { recoverable_auth_error( RecoverableAuthErrorKind::OauthClientId, format!( - "oauth profile '{profile_name}' is missing client_id; re-run `bt auth login --oauth --profile {profile_name}`" + "oauth login for '{}' is missing client_id; re-run `bt auth login --oauth --org `", + auth_slot_label(profile_name, &profile) ), ) })?; - let cached_expires_at = profile.oauth_access_expires_at; - let api_url = auth + let api_url = base .api_url .clone() + .or_else(|| profile.api_url.clone()) .unwrap_or_else(|| DEFAULT_API_URL.to_string()); + let app_url = base.app_url.clone().or_else(|| profile.app_url.clone()); + let org_name = effective_org_name(base, cfg_org) + .map(str::to_string) + .or_else(|| profile.org_name.clone()); - if let Some(cached_access_token) = - load_valid_cached_oauth_access_token(&profile_name, cached_expires_at)? - { + let mut auth = ResolvedAuth { + api_key: None, + api_url: Some(api_url.clone()), + app_url, + org_name, + is_oauth: true, + }; + + if let Some(cached_access_token) = load_valid_cached_oauth_access_token( + profile_name, + &profile, + profile.oauth_access_expires_at, + )? { auth.api_key = Some(cached_access_token); return Ok(auth); } - let refresh_token = load_profile_oauth_refresh_token(&profile_name)?.ok_or_else(|| { - recoverable_auth_error( - RecoverableAuthErrorKind::OauthRefreshToken, - format!( - "oauth refresh token missing for profile '{profile_name}'; re-run `bt auth login --oauth --profile {profile_name}`" - ), - ) - })?; + let refresh_token = load_profile_oauth_refresh_token_for_profile(profile_name, &profile)? + .ok_or_else(|| { + recoverable_auth_error( + RecoverableAuthErrorKind::OauthRefreshToken, + format!( + "oauth refresh token missing for '{}'; re-run `bt auth login --oauth --org `", + auth_slot_label(profile_name, &profile) + ), + ) + })?; let refreshed = - refresh_oauth_access_token(&api_url, &refresh_token, client_id, &profile_name).await?; - save_profile_oauth_access_token(&profile_name, &refreshed.access_token)?; + refresh_oauth_access_token(&api_url, &refresh_token, client_id, profile_name).await?; + save_profile_oauth_access_token(profile_name, &refreshed.access_token)?; + let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { if next_refresh_token != &refresh_token { - save_profile_oauth_refresh_token(&profile_name, next_refresh_token)?; + save_profile_oauth_refresh_token(profile_name, next_refresh_token)?; + refresh_rotated = true; } } - if let Some(profile) = store.profiles.get_mut(&profile_name) { + if !refresh_rotated && profile.legacy_secret_key.is_some() { + save_profile_oauth_refresh_token(profile_name, &refresh_token)?; + } + if let Some(profile) = store.profiles.get_mut(profile_name) { profile.oauth_access_expires_at = determine_oauth_access_expiry_epoch(&refreshed); + if refresh_rotated || profile.legacy_secret_key.is_some() { + delete_legacy_profile_secrets(profile); + profile.legacy_secret_key = None; + } } - save_auth_store(&store)?; + save_auth_store(store)?; auth.api_key = Some(refreshed.access_token); Ok(auth) } @@ -986,22 +1149,48 @@ pub async fn resolved_runner_env(base: &BaseArgs) -> Result(org: &str, store: &'a AuthStore) -> Option<&'a str> { - if store.profiles.contains_key(org) { - return Some( - store - .profiles - .keys() - .find(|k| k.as_str() == org) - .map(|k| k.as_str()) - .unwrap(), - ); +fn profile_matches_org_identifier(profile: &AuthProfile, org: &str) -> bool { + profile.org_id.as_deref() == Some(org) || profile.org_name.as_deref() == Some(org) +} + +fn profile_org_label(profile: &AuthProfile) -> String { + profile + .org_name + .as_deref() + .filter(|org| !org.trim().is_empty()) + .map(str::to_string) + .unwrap_or_else(|| "cross-org".to_string()) +} + +fn profile_identity_label(profile: &AuthProfile) -> Option { + match profile.auth_kind { + AuthKind::Oauth => match (profile.user_name.as_deref(), profile.email.as_deref()) { + (Some(name), Some(email)) => Some(format!("{name} ({email})")), + (None, Some(email)) => Some(email.to_string()), + (Some(name), None) => Some(name.to_string()), + (None, None) => None, + }, + AuthKind::ApiKey => profile.api_key_hint.clone(), } +} +fn auth_slot_label(name: &str, profile: &AuthProfile) -> String { + let mut parts = vec![profile_org_label(profile)]; + parts.push(auth_kind_label(profile.auth_kind).to_string()); + if let Some(identity) = profile_identity_label(profile) { + parts.push(identity); + } else if !name.contains("::") { + parts.push(name.to_string()); + } + parts.join(" — ") +} + +#[cfg(test)] +fn resolve_profile_for_org<'a>(org: &str, store: &'a AuthStore) -> Option<&'a str> { let matches: Vec<&str> = store .profiles .iter() - .filter(|(_, p)| p.org_name.as_deref() == Some(org)) + .filter(|(name, p)| name.as_str() == org || profile_matches_org_identifier(p, org)) .map(|(name, _)| name.as_str()) .collect(); @@ -1012,26 +1201,172 @@ fn resolve_profile_for_org<'a>(org: &str, store: &'a AuthStore) -> Option<&'a st } } +#[cfg(test)] fn profile_names_for_org<'a>(org: &str, store: &'a AuthStore) -> Vec<&'a str> { store .profiles .iter() - .filter(|(_, profile)| profile.org_name.as_deref() == Some(org)) + .filter(|(_, profile)| profile_matches_org_identifier(profile, org)) .map(|(name, _)| name.as_str()) .collect() } -fn profile_label_from_store(name: &str, store: &AuthStore) -> String { - match store +fn is_cross_org_oauth_profile(profile: &AuthProfile) -> bool { + profile.auth_kind == AuthKind::Oauth + && profile + .org_id + .as_deref() + .is_none_or(|org_id| org_id.trim().is_empty()) + && profile + .org_name + .as_deref() + .is_none_or(|org_name| org_name.trim().is_empty()) +} + +fn auth_profile_names_by_kind<'a>( + store: &'a AuthStore, + org: Option<&str>, + kind: AuthKind, +) -> Vec<&'a str> { + store + .profiles + .iter() + .filter(|(_, profile)| profile.auth_kind == kind) + .filter(|(_, profile)| match org { + Some(org) => profile_matches_org_identifier(profile, org), + None if kind == AuthKind::Oauth => is_cross_org_oauth_profile(profile), + None => false, + }) + .map(|(name, _)| name.as_str()) + .collect() +} + +fn profile_info_from_store_entry(name: &str, profile: &AuthProfile) -> ProfileInfo { + ProfileInfo { + name: name.to_string(), + auth_method: auth_kind_label(profile.auth_kind).to_string(), + org_name: profile.org_name.clone(), + user_name: profile.user_name.clone(), + email: profile.email.clone(), + api_key_hint: profile.api_key_hint.clone(), + } +} + +fn profile_info_for_candidate(store: &AuthStore, name: &str) -> Option { + store .profiles .get(name) - .and_then(|profile| profile.org_name.as_deref()) - { - Some(org) if org != name => format!("{} (profile: {})", org, name), - _ => name.to_string(), + .map(|profile| profile_info_from_store_entry(name, profile)) +} + +fn ad_hoc_api_key_profile(org: Option<&str>, api_key: &str) -> ProfileInfo { + ProfileInfo { + name: String::new(), + auth_method: auth_kind_label(AuthKind::ApiKey).to_string(), + org_name: org.map(str::to_string), + user_name: None, + email: None, + api_key_hint: Some(obscure_api_key(api_key)), } } +pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Option { + if let Some(api_key) = resolve_cli_api_key_override(base) { + return Some(ad_hoc_api_key_profile(org, &api_key)); + } + + let store = load_auth_store().unwrap_or_default(); + + if base.prefer_api_key { + let org = org?; + if let Some(api_key) = resolve_env_api_key(base) { + return Some(ad_hoc_api_key_profile(Some(org), &api_key)); + } + + let api_key_candidates = auth_profile_names_by_kind(&store, Some(org), AuthKind::ApiKey); + match api_key_candidates.as_slice() { + [name] => return profile_info_for_candidate(&store, name), + [] => {} + _ => return None, + } + + let oauth_candidates = auth_profile_names_by_kind(&store, Some(org), AuthKind::Oauth); + return match oauth_candidates.as_slice() { + [name] => profile_info_for_candidate(&store, name), + [] => None, + _ => None, + }; + } + + let oauth_candidates = auth_profile_names_by_kind(&store, org, AuthKind::Oauth); + match oauth_candidates.as_slice() { + [name] => return profile_info_for_candidate(&store, name), + [] => {} + _ => return None, + } + + let api_key_candidates = auth_profile_names_by_kind(&store, org, AuthKind::ApiKey); + match api_key_candidates.as_slice() { + [name] => return profile_info_for_candidate(&store, name), + [] => {} + _ => return None, + } + + resolve_env_api_key(base).map(|api_key| ad_hoc_api_key_profile(org, &api_key)) +} + +fn missing_org_for_stored_logins_error(store: &AuthStore) -> Option { + let candidates = store + .profiles + .iter() + .filter(|(_, profile)| { + !is_cross_org_oauth_profile(profile) + && (profile + .org_id + .as_deref() + .is_some_and(|org_id| !org_id.trim().is_empty()) + || profile + .org_name + .as_deref() + .is_some_and(|org_name| !org_name.trim().is_empty())) + }) + .collect::>(); + if candidates.is_empty() { + return None; + } + + let labels = candidates + .iter() + .map(|(name, profile)| auth_slot_label(name, profile)) + .collect::>() + .join(", "); + let all_api_key = candidates + .iter() + .all(|(_, profile)| profile.auth_kind == AuthKind::ApiKey); + + Some(if candidates.len() == 1 { + anyhow::anyhow!( + "auth org selection required; pass --org to use saved auth login: {labels}" + ) + } else if all_api_key { + anyhow::anyhow!( + "multiple API key logins available: {labels}. Pass --org to disambiguate." + ) + } else { + anyhow::anyhow!( + "multiple auth logins available: {labels}. Pass --org to disambiguate." + ) + }) +} + +fn profile_label_from_store(name: &str, store: &AuthStore) -> String { + store + .profiles + .get(name) + .map(|profile| auth_slot_label(name, profile)) + .unwrap_or_else(|| name.to_string()) +} + fn select_profile_from_store( prompt: &str, names: &[&str], @@ -1049,15 +1384,81 @@ fn select_profile_from_store( || store .profiles .get(*name) - .and_then(|profile| profile.org_name.as_deref()) - == Some(current) + .is_some_and(|profile| profile_matches_org_identifier(profile, current)) }) }) .unwrap_or(0); - let idx = ui::fuzzy_select(prompt, &labels, default)?; + let label_refs: Vec<&str> = labels.iter().map(String::as_str).collect(); + let idx = ui::fuzzy_select(prompt, &label_refs, default)?; Ok(names[idx].to_string()) } +fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec { + names + .iter() + .map(|name| { + store + .profiles + .get(*name) + .and_then(profile_identity_label) + .unwrap_or_else(|| (*name).to_string()) + }) + .collect() +} + +fn select_oauth_profile_for_auth( + base: &BaseArgs, + store: &AuthStore, + cfg_org: &Option, + can_prompt: bool, +) -> Result> { + let org = effective_org_name(base, cfg_org); + let candidates = auth_profile_names_by_kind(store, org, AuthKind::Oauth); + select_auth_profile_candidate("OAuth login", org, &candidates, store, can_prompt) +} + +fn select_api_key_profile_for_auth( + base: &BaseArgs, + store: &AuthStore, + cfg_org: &Option, + can_prompt: bool, +) -> Result> { + let org = effective_org_name(base, cfg_org); + let candidates = auth_profile_names_by_kind(store, org, AuthKind::ApiKey); + select_auth_profile_candidate("API key", org, &candidates, store, can_prompt) +} + +fn select_auth_profile_candidate( + kind_label: &str, + org: Option<&str>, + candidates: &[&str], + store: &AuthStore, + can_prompt: bool, +) -> Result> { + match candidates.len() { + 0 => Ok(None), + 1 => Ok(Some(candidates[0].to_string())), + _ if can_prompt => { + let prompt = org + .map(|org| format!("Multiple {kind_label} logins for '{org}'. Select one")) + .unwrap_or_else(|| format!("Select {kind_label} login")); + select_profile_from_store(&prompt, candidates, org, store).map(Some) + } + _ => { + let identities = candidate_identities(candidates, store).join(", "); + if let Some(org) = org { + bail!( + "multiple {kind_label} logins for org '{org}': {identities}. Rerun interactively or remove one with `bt auth logout`." + ); + } + bail!( + "multiple cross-org {kind_label} logins available: {identities}. Rerun interactively or remove one with `bt auth logout`." + ); + } + } +} + +#[cfg(test)] fn maybe_select_profile_for_auth( base: &BaseArgs, store: &AuthStore, @@ -1090,7 +1491,7 @@ fn maybe_select_profile_for_auth( if !can_prompt { bail!( - "multiple profiles for org '{org}': {}. Use --profile to disambiguate.", + "multiple auth logins for org '{org}': {}. Use --org to disambiguate.", matching_profiles.join(", ") ); } @@ -1111,7 +1512,7 @@ fn maybe_select_profile_for_auth( let names: Vec<&str> = store.profiles.keys().map(|name| name.as_str()).collect(); if !can_prompt { bail!( - "multiple auth profiles available: {}. Pass --profile , set BRAINTRUST_PROFILE, or configure an org.", + "multiple auth logins available: {}. Pass --org to disambiguate.", names.join(", ") ); } @@ -1119,6 +1520,7 @@ fn maybe_select_profile_for_auth( select_profile_from_store("Select org", &names, None, store).map(Some) } +#[cfg(test)] fn resolve_auth_from_store_with_secret_lookup( base: &BaseArgs, store: &AuthStore, @@ -1159,7 +1561,7 @@ where if let Some(profile_name) = selected_profile_name { let profile = store.profiles.get(profile_name).ok_or_else(|| { anyhow::anyhow!( - "profile '{profile_name}' not found; run `bt auth profiles` or `bt auth login --profile {profile_name}`" + "auth login '{profile_name}' not found; run `bt auth profiles` or `bt auth login --org `" ) })?; let is_oauth = profile.auth_kind == AuthKind::Oauth; @@ -1170,7 +1572,7 @@ where recoverable_auth_error( RecoverableAuthErrorKind::StoredCredential, format!( - "no keychain credential found for profile '{profile_name}'; re-run `bt auth login --profile {profile_name}`" + "no keychain credential found for auth login '{profile_name}'; re-run `bt auth login --org `" ), ) })?) @@ -1225,7 +1627,6 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { .clone() .unwrap_or_else(|| DEFAULT_APP_URL.to_string()); let login_orgs = fetch_login_orgs(&api_key, &login_app_url).await?; - let store = load_auth_store()?; let requested_org_resolution = resolve_requested_org_for_api_key_login( &login_orgs, base.org_name.as_deref(), @@ -1235,8 +1636,6 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { if requested_org_resolution == RequestedOrgResolution::SwitchToOauth { return run_login_oauth(base, args).await; } - let default_org_name = - default_login_org_name(&store, base.profile.as_deref(), base.org_name.as_deref()); let selected_org = select_login_org( login_orgs.clone(), match requested_org_resolution { @@ -1246,51 +1645,45 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { } RequestedOrgResolution::SwitchToOauth => unreachable!("handled above"), }, - default_org_name.as_deref(), + None, interactive, base.verbose, - true, + false, explicitly_quiet(base), )?; + let selected_org = selected_org.ok_or_else(|| { + anyhow::anyhow!("API-key login requires an org; pass --org or rerun interactively") + })?; let selected_api_url = - resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)?; - let (profile_name, should_confirm_overwrite) = resolve_api_key_login_profile_name( - base.profile.as_deref(), - selected_org.as_ref().map(|org| org.name.as_str()), - &selected_api_url, - &store, - )?; - if should_confirm_overwrite { - confirm_profile_overwrite(&profile_name)?; - } + resolve_profile_api_url(base.api_url.clone(), Some(&selected_org), &login_orgs)?; - commit_api_key_profile( - &profile_name, + let _slot_key = commit_api_key_profile( &api_key, selected_api_url.clone(), base.app_url.clone(), - selected_org.as_ref().map(|org| org.name.clone()), + selected_org.id.clone(), + selected_org.name.clone(), )?; let context_update = persist_post_login_context( base, - &profile_name, &api_key, &selected_api_url, &login_app_url, - selected_org.as_ref(), + Some(&selected_org), ) .await .context("login succeeded, but failed to update active context")?; - let human = format_login_success(&selected_org, &profile_name, &selected_api_url); + let human = format_login_success(Some(&selected_org), &selected_api_url); emit_result( base.json, serde_json::json!({ - "name": profile_name, "auth": "api_key", - "org": selected_org.as_ref().map(|org| org.name.clone()), + "org": selected_org.name, + "org_id": selected_org.id, "api_url": selected_api_url, - "app_url": base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + "app_url": login_app_url, + "api_key_hint": obscure_api_key(&api_key), "status": "ok", }), || { @@ -1315,16 +1708,10 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { .app_url .clone() .unwrap_or_else(|| DEFAULT_APP_URL.to_string()); - let provisional_profile = base - .profile - .as_deref() - .map(str::trim) - .filter(|name| !name.is_empty()) - .unwrap_or("default"); let client_id = args .client_id .clone() - .unwrap_or_else(|| default_oauth_client_id(provisional_profile)); + .unwrap_or_else(default_oauth_client_id); let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256(); let state = generate_random_token(32)?; @@ -1375,13 +1762,10 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { ) .await?; let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; - let store = load_auth_store()?; - let default_org_name = - default_login_org_name(&store, base.profile.as_deref(), base.org_name.as_deref()); let selected_org = select_login_org( login_orgs.clone(), base.org_name.as_deref(), - default_org_name.as_deref(), + None, ui::can_prompt(), base.verbose, true, @@ -1389,30 +1773,16 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { )?; let selected_api_url = resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)?; - let jwt_id = decode_jwt_identity(&oauth_tokens.access_token); - let (profile_name, should_confirm_overwrite) = resolve_oauth_login_profile_name( - base.profile.as_deref(), - selected_org.as_ref().map(|org| org.name.as_str()), - &selected_api_url, - &app_url, - &jwt_id, - &store, - )?; - if should_confirm_overwrite { - confirm_profile_overwrite(&profile_name)?; - } - commit_oauth_profile( - &profile_name, + let _slot_key = commit_oauth_profile( &oauth_tokens, selected_api_url.clone(), app_url.clone(), client_id.clone(), - selected_org.as_ref().map(|org| org.name.clone()), + selected_org.as_ref(), )?; let context_update = persist_post_login_context( base, - &profile_name, &oauth_tokens.access_token, &selected_api_url, &app_url, @@ -1421,13 +1791,14 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { .await .context("login succeeded, but failed to update active context")?; - let human = format_login_success(&selected_org, &profile_name, &selected_api_url); + let human = format_login_success(selected_org.as_ref(), &selected_api_url); emit_result( base.json, serde_json::json!({ - "name": profile_name, "auth": "oauth", "org": selected_org.as_ref().map(|org| org.name.clone()), + "org_id": selected_org.as_ref().map(|org| org.id.clone()), + "cross_org": selected_org.is_none(), "api_url": selected_api_url, "app_url": app_url, "status": "ok", @@ -1446,84 +1817,116 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { } pub(crate) fn commit_api_key_profile( - profile_name: &str, api_key: &str, api_url: String, app_url: Option, - org_name: Option, -) -> Result<()> { - save_profile_secret(profile_name, api_key)?; - let _ = delete_profile_oauth_refresh_token(profile_name); - let _ = delete_profile_oauth_access_token(profile_name); + org_id: String, + org_name: String, +) -> Result { + let hash = api_key_hash(api_key); + let slot_key = api_key_slot_key(&hash, &org_id); + save_profile_secret(&slot_key, api_key)?; let mut store = load_auth_store()?; + if let Some(old_profile) = store.profiles.get(&slot_key) { + delete_legacy_profile_secrets(old_profile); + } store.profiles.insert( - profile_name.to_string(), + slot_key.clone(), AuthProfile { auth_kind: AuthKind::ApiKey, api_url: Some(api_url), app_url, - org_name, + org_id: Some(org_id), + org_name: Some(org_name), oauth_client_id: None, oauth_access_expires_at: None, user_name: None, email: None, + api_key_hash: Some(hash), api_key_hint: Some(obscure_api_key(api_key)), + legacy_secret_key: None, }, ); - save_auth_store(&store) + save_auth_store(&store)?; + Ok(slot_key) } fn commit_oauth_profile( - profile_name: &str, tokens: &OAuthTokenResponse, api_url: String, app_url: String, client_id: String, - org_name: Option, -) -> Result<()> { + selected_org: Option<&LoginOrgInfo>, +) -> Result { let refresh_token = tokens.refresh_token.as_ref().ok_or_else(|| { anyhow::anyhow!( - "oauth token response did not include a refresh_token; cannot create persistent oauth profile" + "oauth token response did not include a refresh_token; cannot create persistent oauth login" ) })?; - save_profile_oauth_refresh_token(profile_name, refresh_token)?; - save_profile_oauth_access_token(profile_name, &tokens.access_token)?; - let _ = delete_profile_secret(profile_name); let oauth_access_expires_at = determine_oauth_access_expiry_epoch(tokens); let jwt_id = decode_jwt_identity(&tokens.access_token); + let email = jwt_id + .email + .clone() + .filter(|email| !email.trim().is_empty()) + .ok_or_else(|| { + anyhow::anyhow!( + "oauth token did not include an email; cannot create persistent oauth login" + ) + })?; + let org_id = selected_org.map(|org| org.id.clone()).unwrap_or_default(); + let slot_key = oauth_slot_key(&org_id, &email); + + save_profile_oauth_refresh_token(&slot_key, refresh_token)?; + save_profile_oauth_access_token(&slot_key, &tokens.access_token)?; + let _ = delete_profile_secret(&slot_key); let mut store = load_auth_store()?; + if let Some(old_profile) = store.profiles.get(&slot_key) { + delete_legacy_profile_secrets(old_profile); + } store.profiles.insert( - profile_name.to_string(), + slot_key.clone(), AuthProfile { auth_kind: AuthKind::Oauth, api_url: Some(api_url), app_url: Some(app_url), - org_name, + org_id: Some(org_id), + org_name: selected_org.map(|org| org.name.clone()), oauth_client_id: Some(client_id), oauth_access_expires_at, user_name: jwt_id.name, email: jwt_id.email, + api_key_hash: None, api_key_hint: None, + legacy_secret_key: None, }, ); - save_auth_store(&store) + save_auth_store(&store)?; + Ok(slot_key) } async fn run_login_refresh(base: &BaseArgs) -> Result<()> { let mut store = load_auth_store()?; - let (profile_name, source) = resolve_selected_profile_name_for_debug(base, &store)?; + let cfg_org = config_auth_context(base); + let profile_name = select_oauth_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? + .ok_or_else(|| { + anyhow::anyhow!( + "no OAuth login selected; pass --org or run `bt auth profiles` to see available logins" + ) + })?; let profile = store .profiles .get(profile_name.as_str()) - .ok_or_else(|| profile_not_found_err(&profile_name, &store))?; - if profile.auth_kind != AuthKind::Oauth { - bail!( - "profile '{profile_name}' uses api key auth; `bt auth refresh` only applies to oauth profiles" - ); - } + .cloned() + .ok_or_else(|| { + anyhow::anyhow!( + "OAuth login '{}' not found; run `bt auth profiles` to see available logins", + profile_name + ) + })?; let api_url = profile .api_url @@ -1531,18 +1934,24 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { .unwrap_or_else(|| DEFAULT_API_URL.to_string()); let client_id = profile.oauth_client_id.clone().ok_or_else(|| { anyhow::anyhow!( - "oauth profile '{profile_name}' is missing client_id; re-run `bt auth login --oauth --profile {profile_name}`" + "OAuth login for '{}' is missing client_id; re-run `bt auth login --oauth --org `", + auth_slot_label(&profile_name, &profile) ) })?; let previous_expires_at = profile.oauth_access_expires_at; - let refresh_token = load_profile_oauth_refresh_token(profile_name.as_str())?.ok_or_else(|| { - anyhow::anyhow!( - "oauth refresh token missing for profile '{profile_name}'; re-run `bt auth login --oauth --profile {profile_name}`" + let refresh_token = + load_profile_oauth_refresh_token_for_profile(profile_name.as_str(), &profile)?.ok_or_else( + || { + anyhow::anyhow!( + "OAuth refresh token missing for '{}'; re-run `bt auth login --oauth --org `", + auth_slot_label(&profile_name, &profile) ) - })?; + }, + )?; eprintln!( - "Refreshing OAuth token for profile '{profile_name}' (source: {source}, api_url: {api_url})" + "Refreshing OAuth token for {} (api_url: {api_url})", + auth_slot_label(&profile_name, &profile) ); if let Some(expires_at) = previous_expires_at { let now = current_unix_timestamp(); @@ -1565,10 +1974,17 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { refresh_rotated = true; } } + if !refresh_rotated && profile.legacy_secret_key.is_some() { + save_profile_oauth_refresh_token(profile_name.as_str(), &refresh_token)?; + } let new_expires_at = determine_oauth_access_expiry_epoch(&refreshed); if let Some(profile) = store.profiles.get_mut(profile_name.as_str()) { profile.oauth_access_expires_at = new_expires_at; + if refresh_rotated || profile.legacy_secret_key.is_some() { + delete_legacy_profile_secrets(profile); + profile.legacy_secret_key = None; + } } save_auth_store(&store)?; @@ -1588,8 +2004,10 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { emit_result( base.json, serde_json::json!({ - "name": profile_name, "auth": "oauth", + "org": profile.org_name, + "org_id": profile.org_id.filter(|org_id| !org_id.trim().is_empty()), + "user_email": profile.email, "access_expires_at": new_expires_at, "refresh_token_rotated": refresh_rotated, "status": "ok", @@ -1598,6 +2016,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { ) } +#[cfg(test)] fn resolve_selected_profile_name_for_debug( base: &BaseArgs, store: &AuthStore, @@ -1605,7 +2024,7 @@ fn resolve_selected_profile_name_for_debug( if let Some(profile_name) = base.profile.as_deref() { let profile_name = profile_name.trim(); if !profile_name.is_empty() { - return Ok((profile_name.to_string(), "--profile/BRAINTRUST_PROFILE")); + return Ok((profile_name.to_string(), "legacy profile")); } } @@ -1626,9 +2045,10 @@ fn resolve_selected_profile_name_for_debug( } } - bail!("no profile selected; pass --profile , set BRAINTRUST_PROFILE, or configure an org") + bail!("no auth login selected; pass --org or rerun interactively") } +#[cfg(test)] fn resolve_profile_name( explicit_profile: Option<&str>, suggested_org_name: Option<&str>, @@ -1648,6 +2068,7 @@ fn resolve_profile_name( .to_string()) } +#[cfg(test)] fn default_login_org_name( store: &AuthStore, profile_name: Option<&str>, @@ -1673,6 +2094,7 @@ fn default_login_org_name( Some(stored_org_name.unwrap_or(profile_name).to_string()) } +#[cfg(test)] fn default_profile_name(suggested_org_name: Option<&str>) -> String { suggested_org_name .map(str::trim) @@ -1681,6 +2103,7 @@ fn default_profile_name(suggested_org_name: Option<&str>) -> String { .to_string() } +#[cfg(test)] fn next_available_profile_name(base_name: &str, store: &AuthStore) -> String { if !store.profiles.contains_key(base_name) { return base_name.to_string(); @@ -1692,6 +2115,7 @@ fn next_available_profile_name(base_name: &str, store: &AuthStore) -> String { .expect("profile name sequence is infinite") } +#[cfg(test)] fn resolve_api_key_login_profile_name( explicit_profile: Option<&str>, suggested_org_name: Option<&str>, @@ -1723,6 +2147,7 @@ fn resolve_api_key_login_profile_name( )) } +#[cfg(test)] fn resolve_oauth_login_profile_name( explicit_profile: Option<&str>, suggested_org_name: Option<&str>, @@ -1776,6 +2201,7 @@ fn resolve_oauth_login_profile_name( )) } +#[cfg(test)] fn profile_matches_api_key_login_target( profile: &AuthProfile, selected_api_url: &str, @@ -1786,6 +2212,7 @@ fn profile_matches_api_key_login_target( && profile.org_name.as_deref() == suggested_org_name } +#[cfg(test)] fn profile_matches_oauth_login_target( profile: &AuthProfile, selected_api_url: &str, @@ -1801,6 +2228,8 @@ fn profile_matches_oauth_login_target( && profile.email == jwt_id.email } +#[cfg(test)] +#[allow(dead_code)] fn confirm_profile_overwrite(profile_name: &str) -> Result<()> { let store = load_auth_store()?; if !store.profiles.contains_key(profile_name) { @@ -1821,17 +2250,10 @@ fn confirm_profile_overwrite(profile_name: &str) -> Result<()> { Ok(()) } -fn format_login_success( - selected_org: &Option, - profile_name: &str, - api_url: &str, -) -> String { - match selected_org.as_ref() { - Some(org) => format!( - "Logged in as {} (profile: {profile_name}, api: {api_url})", - org.name - ), - None => format!("Logged in (cross-org, profile: {profile_name}, api: {api_url})"), +fn format_login_success(selected_org: Option<&LoginOrgInfo>, api_url: &str) -> String { + match selected_org { + Some(org) => format!("Logged in as {} (api: {api_url})", org.name), + None => format!("Logged in (cross-org, api: {api_url})"), } } @@ -1893,7 +2315,6 @@ async fn resolve_post_login_project( async fn persist_post_login_context( base: &BaseArgs, - profile_name: &str, credential: &str, api_url: &str, app_url: &str, @@ -1910,7 +2331,7 @@ async fn persist_post_login_context( let mut cfg = config::load_file(&path); switch::apply_switch_config( &mut cfg, - Some(profile_name), + None, selected_org.map(|org| org.name.as_str()), project.as_ref(), ); @@ -1923,20 +2344,6 @@ async fn persist_post_login_context( }) } -/// Build an actionable "profile not found" error that lists the available -/// profiles, so non-interactive callers can see what they can pick from. -fn profile_not_found_err(name: &str, store: &AuthStore) -> anyhow::Error { - let available: Vec = store.profiles.keys().cloned().collect(); - let suffix = if available.is_empty() { - String::new() - } else { - format!(": {}", available.join(", ")) - }; - anyhow::anyhow!( - "profile '{name}' not found; run `bt auth profiles` to see available profiles{suffix}" - ) -} - /// Emit a machine-readable JSON payload on stdout when `--json` is set, /// otherwise run the human-readable printer. Keeps stdout pure JSON. fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) -> Result<()> { @@ -1948,36 +2355,21 @@ fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) -> Ok(()) } -async fn run_profiles(base: &BaseArgs, args: AuthProfilesArgs) -> Result<()> { +async fn run_profiles(base: &BaseArgs, _args: AuthProfilesArgs) -> Result<()> { let store = load_auth_store()?; - - // Filter to a single profile when --profile is given; error out if it doesn't match. - let filtered_store = match &args.profile { - Some(name) => { - let profile = store - .profiles - .get(name) - .ok_or_else(|| profile_not_found_err(name, &store))?; - let mut s = AuthStore::default(); - s.profiles.insert(name.clone(), profile.clone()); - s - } - None => store, - }; - - if filtered_store.profiles.is_empty() { + if store.profiles.is_empty() { return emit_result(base.json, serde_json::json!([]), || { - println!("No saved profiles. Run `bt auth login` to create one.") + println!("No saved auth logins. Run `bt auth login` to create one.") }); } - let verifications = verify_all_profiles_from_store(&filtered_store).await; + let verifications = verify_all_profiles_from_store(&store).await; let all_network_errors = verifications .iter() .all(|v| v.status == "error" && !v.error.as_deref().unwrap_or("").contains("invalid")); if all_network_errors { - eprintln!("Could not reach Braintrust API. Showing saved profiles:"); - print_saved_profiles(&filtered_store, base.json)?; + eprintln!("Could not reach Braintrust API. Showing saved auth logins:"); + print_saved_profiles(&store, base.json)?; return Ok(()); } @@ -2004,29 +2396,40 @@ async fn run_profiles(base: &BaseArgs, args: AuthProfilesArgs) -> Result<()> { Ok(()) } +fn auth_profile_json(profile: &AuthProfile, status: &str) -> serde_json::Value { + serde_json::json!({ + "auth": auth_kind_label(profile.auth_kind), + "org": profile.org_name, + "org_id": profile.org_id.as_deref().filter(|org_id| !org_id.trim().is_empty()), + "user_name": profile.user_name, + "user_email": profile.email, + "api_key_hint": profile.api_key_hint, + "status": status, + }) +} + fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result<()> { let profile_name = profile_name.trim(); if profile_name.is_empty() { - bail!("profile name cannot be empty"); + bail!("auth login key cannot be empty"); } let mut store = load_auth_store()?; - if !store.profiles.contains_key(profile_name) { - return Err(profile_not_found_err(profile_name, &store)); - } + let profile = store.profiles.get(profile_name).cloned().ok_or_else(|| { + anyhow::anyhow!("auth login not found; run `bt auth profiles` to see available logins") + })?; + let label = auth_slot_label(profile_name, &profile); if !force { if let Some(term) = ui::prompt_term() { let confirmed = Confirm::new() - .with_prompt(format!("Delete profile '{profile_name}'?")) + .with_prompt(format!("Delete {label}?")) .default(false) .interact_on(&term)?; if !confirmed { - return emit_result( - base_json, - serde_json::json!({ "name": profile_name, "status": "cancelled" }), - || eprintln!("Cancelled"), - ); + return emit_result(base_json, auth_profile_json(&profile, "cancelled"), || { + eprintln!("Cancelled") + }); } } } @@ -2034,49 +2437,72 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result< store.profiles.remove(profile_name); save_auth_store(&store)?; if let Err(err) = delete_profile_secret(profile_name) { - eprintln!("warning: failed to delete keychain credential for '{profile_name}': {err}"); + eprintln!("warning: failed to delete keychain credential for '{label}': {err}"); } if let Err(err) = delete_profile_oauth_refresh_token(profile_name) { - eprintln!("warning: failed to delete oauth refresh token for '{profile_name}': {err}"); + eprintln!("warning: failed to delete oauth refresh token for '{label}': {err}"); } if let Err(err) = delete_profile_oauth_access_token(profile_name) { - eprintln!("warning: failed to delete oauth access token for '{profile_name}': {err}"); + eprintln!("warning: failed to delete oauth access token for '{label}': {err}"); } + delete_legacy_profile_secrets(&profile); - emit_result( - base_json, - serde_json::json!({ "name": profile_name, "status": "deleted" }), - || { - ui::print_command_status( - ui::CommandStatus::Success, - &format!("Deleted profile '{profile_name}'"), - ) - }, - ) + emit_result(base_json, auth_profile_json(&profile, "deleted"), || { + ui::print_command_status(ui::CommandStatus::Success, &format!("Deleted {label}")); + }) } fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { let store = load_auth_store()?; if store.profiles.is_empty() { return emit_result(base.json, serde_json::json!({ "status": "empty" }), || { - println!("No saved profiles.") + println!("No saved auth logins.") }); } - let profile_name = if let Some(p) = args.profile.or(base.profile) { - let p = p.trim().to_string(); - if !store.profiles.contains_key(&p) { - return Err(profile_not_found_err(&p, &store)); + let cfg_org = config_auth_context(&base); + let org = effective_org_name(&base, &cfg_org); + let mut candidates: Vec<&str> = store + .profiles + .iter() + .filter(|(_, profile)| org.is_none_or(|org| profile_matches_org_identifier(profile, org))) + .filter(|(_, profile)| { + args.api_key_hint.as_deref().is_none_or(|hint| { + profile.auth_kind == AuthKind::ApiKey + && profile.api_key_hint.as_deref() == Some(hint.trim()) + }) + }) + .map(|(name, _)| name.as_str()) + .collect(); + + if args.api_key_hint.is_none() && org.is_some() { + let oauth_candidates: Vec<&str> = candidates + .iter() + .copied() + .filter(|name| { + store + .profiles + .get(*name) + .is_some_and(|profile| profile.auth_kind == AuthKind::Oauth) + }) + .collect(); + if !oauth_candidates.is_empty() { + candidates = oauth_candidates; + } + } + + let profile_name = match candidates.len() { + 0 => bail!("no matching auth login found; run `bt auth profiles` to see available logins"), + 1 => candidates[0].to_string(), + _ if ui::can_prompt() => { + select_profile_from_store("Select auth login to log out", &candidates, org, &store)? + } + _ => { + let labels = candidate_identities(&candidates, &store).join(", "); + bail!( + "multiple auth logins match: {labels}. Pass --org or --api-key-hint to disambiguate." + ); } - p - } else if store.profiles.len() == 1 { - store.profiles.keys().next().unwrap().clone() - } else if ui::can_prompt() { - let names: Vec<&str> = store.profiles.keys().map(|k| k.as_str()).collect(); - let idx = crate::ui::fuzzy_select("Select profile to log out", &names, 0)?; - names[idx].to_string() - } else { - bail!("multiple profiles exist. Use --profile to specify which one."); }; run_login_delete(&profile_name, args.force, base.json) @@ -2098,18 +2524,20 @@ enum CredentialLoad { fn load_credential_for_profile(name: &str, profile: &AuthProfile) -> CredentialLoad { match profile.auth_kind { - AuthKind::ApiKey => match load_profile_secret(name) { - Ok(Some(k)) => CredentialLoad::Found(k), - Ok(None) => CredentialLoad::Missing, - Err(e) => CredentialLoad::Error(e.to_string()), - }, + AuthKind::ApiKey => { + match load_profile_secret_with_legacy(name, profile.legacy_secret_key.as_deref()) { + Ok(Some(k)) => CredentialLoad::Found(k), + Ok(None) => CredentialLoad::Missing, + Err(e) => CredentialLoad::Error(e.to_string()), + } + } AuthKind::Oauth => { if let Some(ts) = profile.oauth_access_expires_at { if !oauth_access_token_is_fresh(ts) { return CredentialLoad::Expired; } } - match load_profile_oauth_access_token(name) { + match load_profile_oauth_access_token_for_profile(name, profile) { Ok(Some(k)) => CredentialLoad::Found(k), Ok(None) => CredentialLoad::Missing, Err(e) => CredentialLoad::Error(e.to_string()), @@ -2120,11 +2548,14 @@ fn load_credential_for_profile(name: &str, profile: &AuthProfile) -> CredentialL #[derive(Debug, Clone, Serialize)] pub struct ProfileVerification { + #[serde(skip_serializing)] pub name: String, pub auth: String, #[serde(skip_serializing_if = "Option::is_none")] pub org: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub org_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub user_name: Option, #[serde(skip_serializing_if = "Option::is_none")] pub user_email: Option, @@ -2139,6 +2570,7 @@ fn build_verification( name: &str, auth_kind: &str, org: Option, + org_id: Option, jwt_id: Option, api_key_hint: Option, status: ProfileStatus, @@ -2153,6 +2585,7 @@ fn build_verification( name: name.to_string(), auth: auth_kind.to_string(), org, + org_id, user_name: jwt_id.as_ref().and_then(|j| j.name.clone()), user_email: jwt_id.as_ref().and_then(|j| j.email.clone()), api_key_hint, @@ -2163,15 +2596,16 @@ fn build_verification( async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerification { let app_url = profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); - let auth_kind = match profile.auth_kind { - AuthKind::ApiKey => "api_key", - AuthKind::Oauth => "oauth", - }; + let auth_kind = auth_kind_label(profile.auth_kind); let mk = |status, jwt_id: Option, hint: Option| { build_verification( name, auth_kind, profile.org_name.clone(), + profile + .org_id + .clone() + .filter(|org_id| !org_id.trim().is_empty()), jwt_id, hint, status, @@ -2180,9 +2614,15 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifi let credential = match load_credential_for_profile(name, profile) { CredentialLoad::Found(k) => k, - CredentialLoad::Missing => return mk(ProfileStatus::Missing, None, None), - CredentialLoad::Expired => return mk(ProfileStatus::Expired, None, None), - CredentialLoad::Error(e) => return mk(ProfileStatus::Error(e), None, None), + CredentialLoad::Missing => { + return mk(ProfileStatus::Missing, None, profile.api_key_hint.clone()) + } + CredentialLoad::Expired => { + return mk(ProfileStatus::Expired, None, profile.api_key_hint.clone()) + } + CredentialLoad::Error(e) => { + return mk(ProfileStatus::Error(e), None, profile.api_key_hint.clone()) + } }; let (jwt_id, hint) = match profile.auth_kind { @@ -2203,7 +2643,7 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifi } else { ProfileStatus::Error(msg) }; - mk(status, None, None) + mk(status, None, hint) } } } @@ -2227,10 +2667,10 @@ async fn verify_all_profiles_from_store(store: &AuthStore) -> Vec String { - let mut parts = vec![v.name.clone(), v.auth.clone()]; - if let Some(ref org) = v.org { - parts.push(format!("org: {org}")); - } + let mut parts = vec![ + v.org.clone().unwrap_or_else(|| "cross-org".to_string()), + v.auth.clone(), + ]; match v.status.as_str() { "ok" => { let id = match (&v.user_name, &v.user_email) { @@ -2243,7 +2683,10 @@ fn format_verification_line(v: &ProfileVerification) -> String { } } "expired" => parts.push("token expired".into()), - "missing" => parts.push("credential missing".into()), + "missing" => match v.api_key_hint.as_deref() { + Some(hint) => parts.push(format!("{hint} credential missing")), + None => parts.push("credential missing".into()), + }, _ => { if let Some(ref e) = v.error { parts.push(e.clone()); @@ -2258,11 +2701,11 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { let output: Vec = store .profiles .iter() - .map(|(name, p)| { + .map(|(_name, p)| { serde_json::json!({ - "name": name, - "auth": match p.auth_kind { AuthKind::ApiKey => "api_key", AuthKind::Oauth => "oauth" }, + "auth": auth_kind_label(p.auth_kind), "org": p.org_name, + "org_id": p.org_id, "user_name": p.user_name, "user_email": p.email, "api_key_hint": p.api_key_hint, @@ -2273,25 +2716,7 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { println!("{}", serde_json::to_string(&output)?); } else { for (name, profile) in &store.profiles { - let kind = match profile.auth_kind { - AuthKind::ApiKey => "api_key", - AuthKind::Oauth => "oauth", - }; - let org = profile - .org_name - .as_deref() - .map(|o| format!(" org={o}")) - .unwrap_or_default(); - let id = match (profile.user_name.as_deref(), profile.email.as_deref()) { - (Some(n), Some(e)) => format!(" {n} ({e})"), - (None, Some(e)) => format!(" {e}"), - _ => profile - .api_key_hint - .as_deref() - .map(|h| format!(" {h}")) - .unwrap_or_default(), - }; - println!(" {name} {kind}{org}{id}"); + println!(" {}", auth_slot_label(name, profile)); } } Ok(()) @@ -2357,6 +2782,11 @@ fn select_login_org( } if !interactive { + if allow_cross_org { + bail!( + "organization selection required in non-interactive mode; pass --org or rerun interactively to choose cross-org mode" + ); + } return Ok(None); } @@ -2417,7 +2847,7 @@ fn find_login_org<'a>( fn find_login_org_index(orgs: &[LoginOrgInfo], requested_org_name: &str) -> Option { orgs.iter() - .position(|org| org.name == requested_org_name) + .position(|org| org.id == requested_org_name || org.name == requested_org_name) .or_else(|| { let lowered = requested_org_name.to_ascii_lowercase(); orgs.iter() @@ -2515,23 +2945,8 @@ fn resolve_profile_api_url( ) } -fn default_oauth_client_id(profile_name: &str) -> String { - let sanitized = profile_name - .chars() - .map(|ch| { - if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { - ch - } else { - '-' - } - }) - .collect::(); - let trimmed = sanitized.trim_matches('-'); - if trimmed.is_empty() { - "bt_cli_default".to_string() - } else { - format!("bt_cli_{trimmed}") - } +fn default_oauth_client_id() -> String { + "bt_cli".to_string() } fn generate_random_token(num_bytes: usize) -> Result { @@ -2917,20 +3332,19 @@ async fn exchange_oauth_authorization_code( fn map_refresh_oauth_error( api_url: &str, - profile_name: &str, + auth_login: &str, status: reqwest::StatusCode, body: &str, ) -> anyhow::Error { if let Ok(server_err) = serde_json::from_str::(body) { if matches!(server_err.error.as_deref(), Some("invalid_grant")) { - let mut message = - format!("oauth refresh token expired or was rejected for profile '{profile_name}'"); + let mut message = format!( + "oauth refresh token expired or was rejected for auth login '{auth_login}'" + ); if let Some(description) = server_err.error_description.as_deref() { message.push_str(&format!(" ({description})")); } - message.push_str(&format!( - "; re-run `bt auth login --oauth --profile {profile_name}`" - )); + message.push_str("; re-run `bt auth login --oauth --org `"); return recoverable_auth_error(RecoverableAuthErrorKind::OauthRefreshToken, message); } } @@ -2945,7 +3359,7 @@ async fn refresh_oauth_access_token( api_url: &str, refresh_token: &str, client_id: &str, - profile_name: &str, + auth_login: &str, ) -> Result { let http_client = build_http_client_from_builder( reqwest::Client::builder() @@ -2967,12 +3381,7 @@ async fn refresh_oauth_access_token( if !response.status().is_success() { let status = response.status(); let body = response.text().await.unwrap_or_default(); - return Err(map_refresh_oauth_error( - api_url, - profile_name, - status, - &body, - )); + return Err(map_refresh_oauth_error(api_url, auth_login, status, &body)); } response @@ -3093,6 +3502,47 @@ fn load_profile_secret(profile_name: &str) -> Result> { } } +fn load_profile_secret_with_legacy( + primary_key: &str, + legacy_key: Option<&str>, +) -> Result> { + if let Some(secret) = load_profile_secret(primary_key)? { + return Ok(Some(secret)); + } + + let Some(legacy_key) = legacy_key + .map(str::trim) + .filter(|key| !key.is_empty() && *key != primary_key) + else { + return Ok(None); + }; + + let Some(secret) = load_profile_secret(legacy_key)? else { + return Ok(None); + }; + let _ = relocate_plaintext_secret_if_present(primary_key, legacy_key, &secret); + Ok(Some(secret)) +} + +fn relocate_plaintext_secret_if_present( + primary_key: &str, + legacy_key: &str, + secret: &str, +) -> Result<()> { + let path = secret_store_path()?; + if !path.exists() { + return Ok(()); + } + let mut store = load_secret_store()?; + if store.secrets.remove(legacy_key).is_some() { + store + .secrets + .insert(primary_key.to_string(), secret.to_string()); + save_secret_store(&store)?; + } + Ok(()) +} + fn delete_profile_secret(profile_name: &str) -> Result<()> { let keychain_err = delete_profile_secret_keychain(profile_name).err(); let plaintext_err = delete_profile_secret_plaintext(profile_name).err(); @@ -3445,11 +3895,25 @@ fn save_profile_oauth_refresh_token(profile_name: &str, refresh_token: &str) -> save_profile_secret(&key, refresh_token) } +#[cfg(test)] +#[allow(dead_code)] fn load_profile_oauth_refresh_token(profile_name: &str) -> Result> { let key = oauth_refresh_secret_key(profile_name); load_profile_secret(&key) } +fn load_profile_oauth_refresh_token_for_profile( + profile_name: &str, + profile: &AuthProfile, +) -> Result> { + let primary = oauth_refresh_secret_key(profile_name); + let legacy = profile + .legacy_secret_key + .as_deref() + .map(oauth_refresh_secret_key); + load_profile_secret_with_legacy(&primary, legacy.as_deref()) +} + fn delete_profile_oauth_refresh_token(profile_name: &str) -> Result<()> { let key = oauth_refresh_secret_key(profile_name); delete_profile_secret(&key) @@ -3460,18 +3924,48 @@ fn save_profile_oauth_access_token(profile_name: &str, access_token: &str) -> Re save_profile_secret(&key, access_token) } +#[cfg(test)] +#[allow(dead_code)] fn load_profile_oauth_access_token(profile_name: &str) -> Result> { let key = oauth_access_secret_key(profile_name); load_profile_secret(&key) } +fn load_profile_oauth_access_token_for_profile( + profile_name: &str, + profile: &AuthProfile, +) -> Result> { + let primary = oauth_access_secret_key(profile_name); + let legacy = profile + .legacy_secret_key + .as_deref() + .map(oauth_access_secret_key); + load_profile_secret_with_legacy(&primary, legacy.as_deref()) +} + fn delete_profile_oauth_access_token(profile_name: &str) -> Result<()> { let key = oauth_access_secret_key(profile_name); delete_profile_secret(&key) } +fn delete_legacy_profile_secrets(profile: &AuthProfile) { + let Some(legacy_key) = profile.legacy_secret_key.as_deref() else { + return; + }; + match profile.auth_kind { + AuthKind::ApiKey => { + let _ = delete_profile_secret(legacy_key); + } + AuthKind::Oauth => { + let _ = delete_profile_oauth_refresh_token(legacy_key); + let _ = delete_profile_oauth_access_token(legacy_key); + } + } +} + fn load_valid_cached_oauth_access_token( profile_name: &str, + profile: &AuthProfile, expires_at: Option, ) -> Result> { let Some(expires_at) = expires_at else { @@ -3480,7 +3974,7 @@ fn load_valid_cached_oauth_access_token( if !oauth_access_token_is_fresh(expires_at) { return Ok(None); } - load_profile_oauth_access_token(profile_name) + load_profile_oauth_access_token_for_profile(profile_name, profile) } fn oauth_access_token_is_fresh(expires_at: u64) -> bool { @@ -3543,6 +4037,36 @@ pub fn obscure_api_key(key: &str) -> String { format!("{}****{}", &key[..prefix_end], &key[suffix_start..]) } +fn sha256_hex(value: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(value.as_bytes()); + let digest = hasher.finalize(); + let mut out = String::with_capacity(digest.len() * 2); + for byte in digest { + use std::fmt::Write as _; + let _ = write!(&mut out, "{byte:02x}"); + } + out +} + +fn api_key_hash(api_key: &str) -> String { + sha256_hex(api_key) +} + +fn oauth_slot_key(org_id: &str, email: &str) -> String { + format!("{org_id}::{email}") +} + +fn api_key_slot_key(api_key_hash: &str, org_id: &str) -> String { + format!("{api_key_hash}::{org_id}") +} + +#[cfg(test)] +#[allow(dead_code)] +fn split_slot_key(slot_key: &str) -> Option<(&str, &str)> { + slot_key.split_once("::") +} + fn current_unix_timestamp() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -3562,8 +4086,99 @@ fn load_auth_store_from_path(path: &Path) -> Result { let data = fs::read_to_string(path) .with_context(|| format!("failed to read auth config {}", path.display()))?; - serde_json::from_str(&data) - .with_context(|| format!("failed to parse auth config {}", path.display())) + let store: AuthStore = serde_json::from_str(&data) + .with_context(|| format!("failed to parse auth config {}", path.display()))?; + Ok(migrate_auth_store(store)) +} + +fn migrate_auth_store(store: AuthStore) -> AuthStore { + let mut migrated = AuthStore::default(); + for (old_key, mut profile) in store.profiles { + normalize_profile_cached_fields_from_key(&old_key, &mut profile); + let new_key = canonical_profile_key(&old_key, &profile); + if new_key != old_key && profile.legacy_secret_key.is_none() { + profile.legacy_secret_key = Some(old_key.clone()); + } + if migrated + .profiles + .get(&new_key) + .is_some_and(|existing| !should_replace_migrated_profile(existing, &profile)) + { + continue; + } + migrated.profiles.insert(new_key, profile); + } + migrated +} + +fn should_replace_migrated_profile(existing: &AuthProfile, candidate: &AuthProfile) -> bool { + match (existing.auth_kind, candidate.auth_kind) { + (AuthKind::Oauth, AuthKind::Oauth) => { + candidate.oauth_access_expires_at.unwrap_or_default() + > existing.oauth_access_expires_at.unwrap_or_default() + } + _ => false, + } +} + +fn looks_like_sha256_hex(value: &str) -> bool { + value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) +} + +fn normalize_profile_cached_fields_from_key(current_key: &str, profile: &mut AuthProfile) { + let Some((left, right)) = current_key.split_once("::") else { + return; + }; + + match profile.auth_kind { + AuthKind::Oauth => { + let key_matches_email = profile.email.as_deref() == Some(right); + if key_matches_email || (profile.email.is_none() && right.contains('@')) { + if profile.org_id.is_none() { + profile.org_id = Some(left.to_string()); + } + if profile.email.is_none() { + profile.email = Some(right.to_string()); + } + } + } + AuthKind::ApiKey => { + if looks_like_sha256_hex(left) && !right.trim().is_empty() { + if profile.api_key_hash.is_none() { + profile.api_key_hash = Some(left.to_string()); + } + if profile.org_id.is_none() { + profile.org_id = Some(right.to_string()); + } + } + } + } +} + +fn canonical_profile_key(current_key: &str, profile: &AuthProfile) -> String { + match profile.auth_kind { + AuthKind::Oauth => { + let Some(email) = profile.email.as_deref().filter(|value| !value.is_empty()) else { + return current_key.to_string(); + }; + let org_id = profile.org_id.as_deref().unwrap_or_default(); + if profile.org_id.is_some() || profile.org_name.is_none() { + oauth_slot_key(org_id, email) + } else { + current_key.to_string() + } + } + AuthKind::ApiKey => match ( + profile + .api_key_hash + .as_deref() + .filter(|value| !value.is_empty()), + profile.org_id.as_deref().filter(|value| !value.is_empty()), + ) { + (Some(hash), Some(org_id)) => api_key_slot_key(hash, org_id), + _ => current_key.to_string(), + }, + } } fn save_auth_store(store: &AuthStore) -> Result<()> { @@ -3669,12 +4284,11 @@ mod tests { no_color: false, no_input: false, profile: None, - profile_explicit: false, project: None, org_name: None, api_key: None, api_key_source: None, - prefer_profile: false, + prefer_api_key: false, api_url: None, app_url: None, ca_cert: None, @@ -3957,12 +4571,15 @@ mod tests { auth_kind: AuthKind::ApiKey, api_url: Some((*api_url).to_string()), app_url: Some((*app_url).to_string()), + org_id: None, org_name: Some((*org_name).to_string()), oauth_client_id: None, oauth_access_expires_at: None, user_name: None, email: None, + api_key_hash: None, api_key_hint: None, + legacy_secret_key: None, }, ); } @@ -4116,6 +4733,103 @@ mod tests { assert_eq!(DEFAULT_APP_URL, "https://www.braintrust.dev"); } + #[tokio::test] + async fn active_auth_info_without_org_uses_cross_org_oauth_only() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + store.profiles.insert( + api_key_slot_key(&api_key_hash("test-api-key"), "org_fake"), + AuthProfile { + auth_kind: AuthKind::ApiKey, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + api_key_hint: Some("sk-****abcde".to_string()), + ..Default::default() + }, + ); + store.profiles.insert( + oauth_slot_key("", "user@example.test"), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some(String::new()), + org_name: None, + user_name: Some("Test User".to_string()), + email: Some("user@example.test".to_string()), + ..Default::default() + }, + ); + save_auth_store(&store).expect("save auth store"); + + let info = active_auth_info(&make_base(), None).expect("active auth info"); + + assert_eq!(info.auth_method, "oauth"); + assert_eq!(info.email.as_deref(), Some("user@example.test")); + assert_eq!(info.org_name, None); + } + + #[tokio::test] + async fn active_auth_info_prefer_api_key_selects_stored_key_for_org() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + store.profiles.insert( + oauth_slot_key("org_fake", "user@example.test"), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + user_name: Some("Test User".to_string()), + email: Some("user@example.test".to_string()), + ..Default::default() + }, + ); + store.profiles.insert( + api_key_slot_key(&api_key_hash("test-api-key"), "org_fake"), + AuthProfile { + auth_kind: AuthKind::ApiKey, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + api_key_hint: Some("sk-****abcde".to_string()), + ..Default::default() + }, + ); + save_auth_store(&store).expect("save auth store"); + let mut base = make_base(); + base.prefer_api_key = true; + + let info = active_auth_info(&base, Some("test-org")).expect("active auth info"); + + assert_eq!(info.auth_method, "api_key"); + assert_eq!(info.api_key_hint.as_deref(), Some("sk-****abcde")); + } + + #[tokio::test] + async fn api_key_profile_rekeys_after_legacy_secret_load() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + store.profiles.insert( + "work".to_string(), + AuthProfile { + auth_kind: AuthKind::ApiKey, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + api_key_hash: None, + ..Default::default() + }, + ); + + maybe_rekey_api_key_profile_after_secret_load(&mut store, "work", "test-api-key") + .expect("rekey api key profile"); + + let key = api_key_slot_key(&api_key_hash("test-api-key"), "org_fake"); + let profile = store.profiles.get(&key).expect("rekeyed profile"); + assert!(!store.profiles.contains_key("work")); + assert_eq!( + profile.api_key_hash.as_deref(), + Some(api_key_hash("test-api-key").as_str()) + ); + assert_eq!(profile.legacy_secret_key.as_deref(), Some("work")); + } + #[test] fn save_and_load_auth_store_round_trip() { let unique = SystemTime::now() @@ -4148,6 +4862,80 @@ mod tests { let _ = fs::remove_dir_all(&dir); } + #[test] + fn migrate_auth_store_rekeys_oauth_slots_and_preserves_legacy_secret_key() { + let mut store = AuthStore::default(); + store.profiles.insert( + "work".to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + email: Some("user@example.test".to_string()), + oauth_client_id: Some("bt_cli_work".to_string()), + ..Default::default() + }, + ); + + let migrated = migrate_auth_store(store); + let key = oauth_slot_key("org_fake", "user@example.test"); + let profile = migrated.profiles.get(&key).expect("migrated profile"); + + assert_eq!(profile.legacy_secret_key.as_deref(), Some("work")); + assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_work")); + } + + #[test] + fn migrate_auth_store_rekeys_api_key_slots_by_hash_and_org() { + let hash = api_key_hash("test-api-key"); + let mut store = AuthStore::default(); + store.profiles.insert( + "work".to_string(), + AuthProfile { + auth_kind: AuthKind::ApiKey, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + api_key_hash: Some(hash.clone()), + api_key_hint: Some("test-****i-key".to_string()), + ..Default::default() + }, + ); + + let migrated = migrate_auth_store(store); + let key = api_key_slot_key(&hash, "org_fake"); + let profile = migrated.profiles.get(&key).expect("migrated profile"); + + assert_eq!(profile.legacy_secret_key.as_deref(), Some("work")); + assert_eq!(profile.api_key_hint.as_deref(), Some("test-****i-key")); + } + + #[test] + fn migrate_auth_store_dedupes_oauth_slots_by_latest_expiry() { + let mut store = AuthStore::default(); + for (name, expires_at, client_id) in [("old", 10, "bt_cli_old"), ("new", 20, "bt_cli_new")] + { + store.profiles.insert( + name.to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + email: Some("user@example.test".to_string()), + oauth_client_id: Some(client_id.to_string()), + oauth_access_expires_at: Some(expires_at), + ..Default::default() + }, + ); + } + + let migrated = migrate_auth_store(store); + let key = oauth_slot_key("org_fake", "user@example.test"); + assert_eq!(migrated.profiles.len(), 1); + let profile = migrated.profiles.get(&key).expect("migrated profile"); + assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_new")); + assert_eq!(profile.legacy_secret_key.as_deref(), Some("new")); + } + #[test] fn resolve_auth_uses_profile_when_no_api_key_override() { let mut base = make_base(); @@ -4181,25 +4969,23 @@ mod tests { } #[test] - fn config_auth_context_returns_profile_and_org_independently() { + fn config_auth_context_returns_config_org() { let base = make_base(); let cfg = auth_config(Some("default-profile"), Some("local-org")); - let (profile, org) = config_auth_context_from_config(&base, &cfg); + let org = config_auth_context_from_config(&base, &cfg); - assert_eq!(profile.as_deref(), Some("default-profile")); assert_eq!(org.as_deref(), Some("local-org")); } #[test] - fn config_auth_context_preserves_explicit_profile_and_config_org() { + fn config_auth_context_ignores_legacy_profile_and_preserves_config_org() { let mut base = make_base(); base.profile = Some("explicit-profile".to_string()); let cfg = auth_config(Some("config-profile"), Some("local-org")); - let (profile, org) = config_auth_context_from_config(&base, &cfg); + let org = config_auth_context_from_config(&base, &cfg); - assert_eq!(profile, None); assert_eq!(org.as_deref(), Some("local-org")); } @@ -4239,10 +5025,10 @@ mod tests { } #[test] - fn resolve_auth_prefer_profile_ignores_api_key_override() { + fn resolve_auth_prefer_api_key_uses_api_key_override() { let mut base = make_base(); base.api_key = Some("explicit-key".to_string()); - base.prefer_profile = true; + base.prefer_api_key = true; base.profile = Some("work".to_string()); let mut store = AuthStore::default(); @@ -4266,16 +5052,16 @@ mod tests { &None, ) .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("profile-key")); - assert_eq!(resolved.org_name.as_deref(), Some("Example Org")); + assert_eq!(resolved.api_key.as_deref(), Some("explicit-key")); + assert_eq!(resolved.org_name, None); } #[test] - fn resolve_auth_prefers_cli_api_key_even_with_prefer_profile() { + fn resolve_auth_prefers_cli_api_key_even_with_prefer_api_key() { let mut base = make_base(); base.api_key = Some("explicit-key".to_string()); base.api_key_source = Some(crate::args::ArgValueSource::CommandLine); - base.prefer_profile = true; + base.prefer_api_key = true; base.profile = Some("work".to_string()); let mut store = AuthStore::default(); @@ -4304,11 +5090,10 @@ mod tests { } #[test] - fn resolve_auth_explicit_profile_ignores_env_api_key_override() { + fn resolve_auth_legacy_explicit_profile_does_not_override_api_key() { let mut base = make_base(); base.api_key = Some("explicit-key".to_string()); base.profile = Some("work".to_string()); - base.profile_explicit = true; let mut store = AuthStore::default(); store.profiles.insert( @@ -4331,8 +5116,8 @@ mod tests { &None, ) .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("profile-key")); - assert_eq!(resolved.org_name.as_deref(), Some("Example Org")); + assert_eq!(resolved.api_key.as_deref(), Some("explicit-key")); + assert_eq!(resolved.org_name, None); } #[test] @@ -4375,7 +5160,7 @@ mod tests { let (profile_name, source) = resolve_selected_profile_name_for_debug(&base, &store).expect("resolve"); assert_eq!(profile_name, "work"); - assert_eq!(source, "--profile/BRAINTRUST_PROFILE"); + assert_eq!(source, "legacy profile"); } #[test] @@ -4487,10 +5272,10 @@ mod tests { let err = maybe_select_profile_for_auth(&base, &store, &None, false) .expect_err("selection should be required"); - assert!(err.to_string().contains("multiple auth profiles available")); + assert!(err.to_string().contains("multiple auth logins available")); assert!(err.to_string().contains("alpha")); assert!(err.to_string().contains("beta")); - assert!(err.to_string().contains("--profile ")); + assert!(err.to_string().contains("--org ")); } #[test] @@ -4517,7 +5302,9 @@ mod tests { let err = maybe_select_profile_for_auth(&base, &store, &None, false) .expect_err("org selection should be required"); - assert!(err.to_string().contains("multiple profiles for org 'acme'")); + assert!(err + .to_string() + .contains("multiple auth logins for org 'acme'")); assert!(err.to_string().contains("work-1")); assert!(err.to_string().contains("work-2")); } @@ -4946,7 +5733,6 @@ mod tests { let update = persist_post_login_context( &make_base(), - "work", "test-api-key", "https://api.example.test", "https://www.example.test", @@ -4957,7 +5743,7 @@ mod tests { let cfg = crate::config::load_global().expect("load global config"); assert_eq!(update.display, "acme"); - assert_eq!(cfg.profile.as_deref(), Some("work")); + assert_eq!(cfg.profile, None); assert_eq!(cfg.org.as_deref(), Some("acme")); assert_eq!(cfg.project, None); assert_eq!(cfg.project_id, None); @@ -5103,6 +5889,7 @@ mod tests { name: "work".into(), auth: "oauth".into(), org: Some("acme".into()), + org_id: None, user_name: Some("Alice".into()), user_email: Some("alice@example.com".into()), api_key_hint: None, @@ -5111,7 +5898,7 @@ mod tests { }; assert_eq!( format_verification_line(&v), - "work — oauth — org: acme — Alice (alice@example.com)" + "acme — oauth — Alice (alice@example.com)" ); } @@ -5121,6 +5908,7 @@ mod tests { name: "work".into(), auth: "api_key".into(), org: Some("acme".into()), + org_id: None, user_name: None, user_email: None, api_key_hint: Some("sk-****zhJwO".into()), @@ -5129,7 +5917,7 @@ mod tests { }; assert_eq!( format_verification_line(&v), - "work — api_key — org: acme — sk-****zhJwO" + "acme — api_key — sk-****zhJwO" ); } @@ -5139,13 +5927,17 @@ mod tests { name: "old".into(), auth: "oauth".into(), org: None, + org_id: None, user_name: None, user_email: None, api_key_hint: None, status: "expired".into(), error: None, }; - assert_eq!(format_verification_line(&v), "old — oauth — token expired"); + assert_eq!( + format_verification_line(&v), + "cross-org — oauth — token expired" + ); } #[test] @@ -5154,6 +5946,7 @@ mod tests { name: "bad".into(), auth: "api_key".into(), org: Some("corp".into()), + org_id: None, user_name: None, user_email: None, api_key_hint: None, @@ -5162,7 +5955,7 @@ mod tests { }; assert_eq!( format_verification_line(&v), - "bad — api_key — org: corp — invalid API key" + "corp — api_key — invalid API key" ); } diff --git a/src/config/mod.rs b/src/config/mod.rs index 779499d5..c6b51604 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -18,6 +18,7 @@ mod set; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] #[serde(default)] pub struct Config { + #[serde(default, skip_serializing_if = "Option::is_none")] pub profile: Option, pub org: Option, pub project: Option, @@ -26,12 +27,11 @@ pub struct Config { pub extra: serde_json::Map, } -pub const KNOWN_KEYS: &[&str] = &["profile", "org", "project", "project_id"]; +pub const KNOWN_KEYS: &[&str] = &["org", "project", "project_id"]; impl Config { pub fn get_field(&self, key: &str) -> Option<&str> { match key { - "profile" => self.profile.as_deref(), "org" => self.org.as_deref(), "project" => self.project.as_deref(), "project_id" => self.project_id.as_deref(), @@ -41,7 +41,6 @@ impl Config { pub fn set_field(&mut self, key: &str, value: String) -> bool { match key { - "profile" => self.profile = Some(value), "org" => self.org = Some(value), "project" => { self.project = Some(value); @@ -55,7 +54,6 @@ impl Config { pub fn unset_field(&mut self, key: &str) -> bool { match key { - "profile" => self.profile = None, "org" => self.org = None, "project" => { self.project = None; @@ -84,7 +82,7 @@ impl Config { self.project_id.clone() }; Config { - profile: other.profile.clone().or_else(|| self.profile.clone()), + profile: None, org: other.org.clone().or_else(|| self.org.clone()), project, project_id, @@ -119,7 +117,7 @@ pub fn load_file(path: &Path) -> Config { } }; - let config: Config = match serde_json::from_str(&file_contents) { + let mut config: Config = match serde_json::from_str(&file_contents) { Ok(c) => c, Err(e) => { print_command_status( @@ -130,6 +128,10 @@ pub fn load_file(path: &Path) -> Config { } }; + // Legacy config files may contain `profile`; profiles are no longer a + // user-facing selector, so ignore it rather than warning or preserving it. + config.profile = None; + for key in config.extra.keys() { print_command_status( CommandStatus::Error, @@ -180,21 +182,13 @@ pub(crate) fn project_from_config_for_context( .flatten() } -fn config_matches_context(base: &BaseArgs, cfg: &Config, resolved_org: Option<&str>) -> bool { - let selected_profile = trimmed_option(base.profile.as_deref()); - let cfg_profile = trimmed_option(cfg.profile.as_deref()); +fn config_matches_context(_base: &BaseArgs, cfg: &Config, resolved_org: Option<&str>) -> bool { let cfg_org = trimmed_option(cfg.org.as_deref()); let resolved_org = trimmed_option(resolved_org); - match selected_profile { - Some(profile) => { - cfg_profile == Some(profile) - || (cfg_profile.is_none() && cfg_org.is_some() && cfg_org == resolved_org) - } - None => cfg_org - .zip(resolved_org) - .is_none_or(|(cfg, resolved)| cfg == resolved), - } + cfg_org + .zip(resolved_org) + .is_none_or(|(cfg, resolved)| cfg == resolved) } pub(crate) fn trimmed_option(value: Option<&str>) -> Option<&str> { @@ -318,14 +312,14 @@ enum ConfigCommands { }, /// Get a config value Get { - /// Config key (profile, org, project, project_id) + /// Config key (org, project, project_id) key: String, #[command(flatten)] scope: ScopeArgs, }, /// Set a config value Set { - /// Config key (profile, org, project, project_id) + /// Config key (org, project, project_id) key: String, /// Value to set value: String, @@ -334,7 +328,7 @@ enum ConfigCommands { }, /// Remove a config value Unset { - /// Config key (profile, org, project, project_id) + /// Config key (org, project, project_id) key: String, #[command(flatten)] scope: ScopeArgs, @@ -443,12 +437,11 @@ mod tests { no_color: false, no_input: false, profile: profile.map(str::to_string), - profile_explicit: profile.is_some(), org_name: None, project: None, api_key: None, api_key_source: None, - prefer_profile: false, + prefer_api_key: false, api_url: None, app_url: None, ca_cert: None, @@ -466,13 +459,16 @@ mod tests { } #[test] - fn project_config_matches_explicit_profile_or_legacy_org() { + fn project_config_matches_org_and_ignores_legacy_profile() { let base = base_with_profile(Some("work")); let cases = [ (config(None, Some("acme"), Some("demo")), Some("demo")), (config(None, Some("other"), Some("demo")), None), - (config(None, None, Some("demo")), None), - (config(Some("other"), Some("acme"), Some("demo")), None), + (config(None, None, Some("demo")), Some("demo")), + ( + config(Some("other"), Some("acme"), Some("demo")), + Some("demo"), + ), ( config(Some("work"), Some("acme"), Some("demo")), Some("demo"), diff --git a/src/datasets/pipeline.rs b/src/datasets/pipeline.rs index f1c43e49..ae8bbd0f 100644 --- a/src/datasets/pipeline.rs +++ b/src/datasets/pipeline.rs @@ -2074,12 +2074,11 @@ mod tests { no_color: false, no_input: false, profile: None, - profile_explicit: false, org_name: None, project: None, api_key: None, api_key_source: None, - prefer_profile: false, + prefer_api_key: false, api_url: None, app_url: None, ca_cert: None, diff --git a/src/datasets/snapshots.rs b/src/datasets/snapshots.rs index 00f4a804..047f2d3e 100644 --- a/src/datasets/snapshots.rs +++ b/src/datasets/snapshots.rs @@ -8,11 +8,12 @@ use serde_json::json; use crate::{ args::BaseArgs, + auth, ui::{ apply_column_padding, header, is_interactive, print_command_status, print_with_pager, styled_table, truncate, with_spinner, with_spinner_visible, CommandStatus, }, - utils::{pluralize, profile_author_slug, resolve_profile_info, sanitize_name_segment}, + utils::{pluralize, profile_author_slug, sanitize_name_segment}, }; use super::{ @@ -1092,7 +1093,7 @@ fn resolve_default_snapshot_author(base: &BaseArgs, ctx: &ResolvedContext) -> Op return None; } - let profile = resolve_profile_info(base.profile.as_deref(), Some(ctx.client.org_name()))?; + let profile = auth::active_auth_info(base, Some(ctx.client.org_name()))?; profile_author_slug(&profile) } @@ -1102,11 +1103,13 @@ fn default_snapshot_name(author: &str, now: DateTime) -> String { } fn api_key_override_active(base: &BaseArgs) -> bool { - !base.prefer_profile - && base - .api_key - .as_deref() - .is_some_and(|value| !value.trim().is_empty()) + matches!( + base.api_key_source, + Some(crate::args::ArgValueSource::CommandLine) + ) && base + .api_key + .as_deref() + .is_some_and(|value| !value.trim().is_empty()) } #[cfg(test)] diff --git a/src/eval.rs b/src/eval.rs index ebc4f4dc..d279951a 100644 --- a/src/eval.rs +++ b/src/eval.rs @@ -746,7 +746,7 @@ async fn run_eval_attempt( options.jsonl, options.list, options.verbose, - base.profile.clone(), + base.org_name.clone(), ); let output = drive_eval_runner(spawned.process, console_policy, |event| ui.handle(event)).await?; @@ -2805,7 +2805,7 @@ struct EvalUi { deferred_errors: Vec, suppressed_stderr_lines: usize, finished: bool, - profile: Option, + org: Option, } struct EvalBarState { @@ -2815,7 +2815,7 @@ struct EvalBarState { } impl EvalUi { - fn new(jsonl: bool, list: bool, verbose: bool, profile: Option) -> Self { + fn new(jsonl: bool, list: bool, verbose: bool, org: Option) -> Self { let draw_target = if std::io::stderr().is_terminal() && animations_enabled() && !is_quiet() { ProgressDrawTarget::stderr_with_hz(10) @@ -2841,7 +2841,7 @@ impl EvalUi { deferred_errors: Vec::new(), suppressed_stderr_lines: 0, finished: false, - profile, + org, } } @@ -2869,7 +2869,7 @@ impl EvalUi { } } EvalEvent::Summary(summary) => { - let summary = enrich_experiment_summary(summary, self.profile.as_deref()); + let summary = enrich_experiment_summary(summary, self.org.as_deref()); if self.jsonl { if let Ok(line) = serde_json::to_string(&summary) { println!("{line}"); @@ -3350,10 +3350,10 @@ const COMPARE_MORE_HINT: &str = "append more experiment names at the end; max 7 fn enrich_experiment_summary( mut summary: ExperimentSummary, - profile: Option<&str>, + org: Option<&str>, ) -> ExperimentSummary { if summary.compare_command.is_none() { - summary.compare_command = build_experiment_compare_command(&summary, profile); + summary.compare_command = build_experiment_compare_command(&summary, org); } if summary.compare_command.is_some() && summary.compare_more.is_none() { summary.compare_more = Some(COMPARE_MORE_HINT.to_string()); @@ -3375,7 +3375,7 @@ fn format_experiment_compare_command(summary: &ExperimentSummary) -> Option, + org: Option<&str>, ) -> Option { let baseline = summary.comparison_experiment_name.as_deref()?.trim(); if baseline.is_empty() { @@ -3387,14 +3387,14 @@ fn build_experiment_compare_command( return None; } - let profile_args = profile + let org_args = org .map(str::trim) - .filter(|profile| !profile.is_empty()) - .map(|profile| format!(" --profile {}", shell_quote_arg(profile))) + .filter(|org| !org.is_empty()) + .map(|org| format!(" --org {}", shell_quote_arg(org))) .unwrap_or_default(); Some(format!( - "bt experiments compare{profile_args} -p {} {} {}", + "bt experiments compare{org_args} -p {} {} {}", shell_quote_arg(project), shell_quote_arg(baseline), shell_quote_arg(experiment), @@ -4544,7 +4544,7 @@ mod tests { } #[test] - fn enrich_experiment_summary_preserves_profile_in_compare_command() { + fn enrich_experiment_summary_preserves_org_in_compare_command() { let summary = ExperimentSummary { project_name: "test-project".to_string(), experiment_name: "challenger-a".to_string(), @@ -4564,11 +4564,11 @@ mod tests { compare_more: None, }; - let summary = enrich_experiment_summary(summary, Some("test-profile")); + let summary = enrich_experiment_summary(summary, Some("test-org")); assert_eq!( summary.compare_command.as_deref(), - Some("bt experiments compare --profile test-profile -p test-project baseline challenger-a") + Some("bt experiments compare --org test-org -p test-project baseline challenger-a") ); } diff --git a/src/functions/push.rs b/src/functions/push.rs index 4d986b2f..d3771ab2 100644 --- a/src/functions/push.rs +++ b/src/functions/push.rs @@ -4064,12 +4064,11 @@ mod tests { no_color: false, no_input: false, profile: None, - profile_explicit: false, org_name: None, project: None, api_key: None, api_key_source: None, - prefer_profile: false, + prefer_api_key: false, api_url: None, app_url: None, ca_cert: None, diff --git a/src/init.rs b/src/init.rs index e9976c1b..f9e1e03c 100644 --- a/src/init.rs +++ b/src/init.rs @@ -17,6 +17,26 @@ Examples: ")] pub struct InitArgs {} +fn select_saved_auth_org_for_init() -> Result> { + let mut orgs = auth::list_profiles()? + .into_iter() + .filter_map(|profile| profile.org_name) + .filter(|org| !org.trim().is_empty()) + .collect::>(); + orgs.sort(); + orgs.dedup(); + + match orgs.len() { + 0 => Ok(None), + 1 => Ok(orgs.into_iter().next()), + _ => { + let labels = orgs.iter().map(String::as_str).collect::>(); + let idx = crate::ui::fuzzy_select("Select organization", &labels, 0)?; + Ok(Some(orgs[idx].clone())) + } + } +} + pub async fn run(base: BaseArgs, _args: InitArgs) -> Result<()> { let config_path = config::local_save_path()?; if config_path.exists() { @@ -44,10 +64,8 @@ pub async fn run(base: BaseArgs, _args: InitArgs) -> Result<()> { bail!("--org and --project required in non-interactive mode"); } else { let mut login_base = base.clone(); - if login_base.org_name.is_none() && login_base.profile.is_none() { - if let Some(profile) = auth::select_profile_interactive(None)? { - login_base.profile = Some(profile); - } + if login_base.org_name.is_none() { + login_base.org_name = select_saved_auth_org_for_init()?; } let ctx = login(&login_base).await?; let client = ApiClient::new(&ctx)?; diff --git a/src/main.rs b/src/main.rs index 5f12e60f..9a7a5f12 100644 --- a/src/main.rs +++ b/src/main.rs @@ -35,7 +35,7 @@ mod ui; mod util_cmd; mod utils; -use crate::args::{has_explicit_profile_arg, ArgValueSource, BaseArgs, CLIArgs}; +use crate::args::{ArgValueSource, BaseArgs, CLIArgs}; const DEFAULT_CANARY_VERSION: &str = concat!(env!("CARGO_PKG_VERSION"), "-canary.dev"); pub(crate) const CLI_VERSION: &str = match option_env!("BT_VERSION_STRING") { @@ -84,9 +84,9 @@ Additional update Update bt in-place Flags - --profile Use a saved login profile [env: BRAINTRUST_PROFILE] -o, --org Override active org [env: BRAINTRUST_ORG_NAME] - -p, --project Override active project [env: BRAINTRUST_DEFAULT_PROJECT] + -p, --project Override active project [env: BRAINTRUST_DEFAULT_PROJECT] + --prefer-api-key Prefer API key credentials for the selected org [env: BRAINTRUST_PREFER_API_KEY=] --json Output as JSON -v, --verbose Increase output verbosity [env: BRAINTRUST_VERBOSE=] -q, --quiet Reduce interactive UI output [env: BRAINTRUST_QUIET=] @@ -296,7 +296,6 @@ fn try_main() -> Result<()> { let matches = Cli::command().get_matches_from(&argv); let mut cli = Cli::from_arg_matches(&matches).expect("clap matches should parse"); apply_base_arg_sources(&matches, cli.command.base_mut()); - cli.command.base_mut().profile_explicit = has_explicit_profile_arg(&argv); apply_base_output_defaults(&mut cli.command); configure_output(cli.command.base()); apply_runtime_env_overrides(cli.command.base()); @@ -500,7 +499,7 @@ fn looks_like_user_error(err: &anyhow::Error) -> bool { fn print_error(err: &anyhow::Error, code: ExitCode, missing_credential: bool) { eprintln!("error: {err}"); if code == ExitCode::Auth && !missing_credential { - eprintln!("Your credentials may be expired or invalid. For OAuth profiles, try `bt auth refresh --profile `; if refresh fails, re-run `bt auth login --oauth --profile `. Run `bt auth profiles` and `bt status` to inspect profile status."); + eprintln!("Your credentials may be expired or invalid. For OAuth login, try `bt auth refresh --org `; if refresh fails, re-run `bt auth login --oauth --org `. Run `bt auth profiles` and `bt status` to inspect auth status."); } if code == ExitCode::Error { eprintln!("If this seems like a bug, file an issue at https://github.com/braintrustdata/bt/issues/new and include `bt --version`, `bt status --json`, and the command you ran."); diff --git a/src/setup/mod.rs b/src/setup/mod.rs index 45ebfbcd..9d23a4f3 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -1467,21 +1467,18 @@ async fn run_setup_browser_auth( org_id: completed.org_id.clone(), description: None, }; - let stored_profiles = auth::list_profiles()?; - let profile_name = setup_browser_profile_name(profile_name, &org.name, &stored_profiles); - + let _ = profile_name; auth::commit_api_key_profile( - &profile_name, &completed.api_key, login.api_url.clone(), Some(login.app_url.clone()), - Some(org.name.clone()), + org.id.clone(), + org.name.clone(), ) - .context("failed to save Braintrust auth profile after browser setup")?; + .context("failed to save Braintrust auth login after browser setup")?; base.api_key = Some(completed.api_key); base.api_key_source = None; - base.profile = Some(profile_name); base.org_name = Some(org.name); base.project = Some(selected_project.name.clone()); @@ -1492,6 +1489,7 @@ async fn run_setup_browser_auth( }) } +#[cfg(test)] fn setup_browser_profile_name( selected_profile_name: Option<&str>, org_name: &str, @@ -1534,7 +1532,7 @@ fn resolve_profile_name_for_setup( return Ok(Some(profile_name.to_string())); } bail!( - "profile '{profile_name}' not found; run `bt auth profiles` to see available profiles" + "auth login '{profile_name}' not found; run `bt auth profiles` to see available logins" ); } @@ -1550,20 +1548,25 @@ fn resolve_profile_name_for_setup( let mut matches = profiles .iter() .filter(|profile| profile.org_name.as_deref() == Some(org_name)) - .map(|profile| profile.name.clone()) .collect::>(); - matches.sort(); + matches.sort_by(|left, right| left.name.cmp(&right.name)); return match matches.len() { 0 => Ok(None), - 1 => Ok(Some(matches.remove(0))), + 1 => Ok(Some(matches[0].name.clone())), _ if prompt_for_choice => auth::select_profile_interactive(Some(org_name))? .map(Some) - .ok_or_else(|| anyhow!("no profile selected")), - _ => bail!( - "multiple profiles for org '{org_name}': {}. Use --profile to disambiguate.", - matches.join(", ") - ), + .ok_or_else(|| anyhow!("no auth login selected")), + _ => { + let labels = matches + .iter() + .map(|profile| auth::profile_info_label(profile)) + .collect::>() + .join(", "); + bail!( + "multiple auth logins for org '{org_name}': {labels}. Rerun interactively or remove one with `bt auth logout`." + ) + } }; } @@ -1574,7 +1577,7 @@ fn resolve_profile_name_for_setup( if prompt_for_choice && !profiles.is_empty() { auth::select_profile_interactive(None)? .map(Some) - .ok_or_else(|| anyhow!("no profile selected")) + .ok_or_else(|| anyhow!("no auth login selected")) } else { Ok(None) } @@ -1612,11 +1615,13 @@ fn find_available_org<'a>( orgs: &'a [auth::AvailableOrg], org_name: &str, ) -> Option<&'a auth::AvailableOrg> { - orgs.iter().find(|org| org.name == org_name).or_else(|| { - let lowered = org_name.to_ascii_lowercase(); - orgs.iter() - .find(|org| org.name.to_ascii_lowercase() == lowered) - }) + orgs.iter() + .find(|org| org.id == org_name || org.name == org_name) + .or_else(|| { + let lowered = org_name.to_ascii_lowercase(); + orgs.iter() + .find(|org| org.name.to_ascii_lowercase() == lowered) + }) } fn build_api_key_login_context( @@ -1724,24 +1729,6 @@ fn select_api_key_org_for_setup( bail!("organization choice required in non-interactive mode; pass --org ") } -fn matching_profile_org_names( - profiles: &[auth::StoredProfileInfo], - only_oauth: Option, -) -> Vec { - let mut org_names = profiles - .iter() - .filter(|profile| { - only_oauth - .map(|expected| profile.is_oauth == expected) - .unwrap_or(true) - }) - .filter_map(|profile| profile.org_name.clone()) - .collect::>(); - org_names.sort(); - org_names.dedup(); - org_names -} - async fn ensure_profile_or_setup_browser_auth( base: &mut BaseArgs, prompt_for_profile_choice: bool, @@ -1760,10 +1747,17 @@ async fn ensure_profile_or_setup_browser_auth( if let Some(profile_name) = selected_profile { auth_base.profile = Some(profile_name.clone()); + if auth_base.org_name.is_none() { + auth_base.org_name = profiles + .iter() + .find(|profile| profile.name == profile_name) + .and_then(|profile| profile.org_name.clone()); + } match auth::login(&auth_base).await { Ok(ctx) => { base.profile = auth_base.profile.clone(); + base.org_name = auth_base.org_name.clone(); let is_oauth = auth::resolve_auth(&auth_base).await?.is_oauth; return Ok(SetupAuthLogin { login: ctx, @@ -1774,7 +1768,7 @@ async fn ensure_profile_or_setup_browser_auth( Err(err) if auth::is_missing_credential_error(&err) => { if base.verbose { eprintln!( - " Profile '{}' credentials inaccessible ({}). Re-authenticating in the browser...", + " Auth login '{}' credentials inaccessible ({}). Re-authenticating in the browser...", profile_name, err ); } @@ -1799,10 +1793,10 @@ async fn ensure_profile_or_setup_browser_auth( if !can_prompt { if profiles.is_empty() { bail!( - "setup needs interactive browser authentication; rerun without --no-input/--json or pass a valid API key/profile" + "setup needs interactive browser authentication; rerun without --no-input/--json or pass a valid API key or org login" ); } - bail!("profile selection required in non-interactive mode; pass --profile "); + bail!("auth org selection required in non-interactive mode; pass --org "); } if base.verbose { @@ -1903,12 +1897,12 @@ async fn ensure_setup_auth( let profile = stored_profiles .iter() .find(|profile| profile.name == profile_name) - .ok_or_else(|| anyhow!("profile '{profile_name}' not found"))?; + .ok_or_else(|| anyhow!("auth login '{profile_name}' not found"))?; let target_org = match org_name.as_deref() { Some(org_name) => { if profile.org_name.as_deref() != Some(org_name) { bail!( - "profile '{profile_name}' belongs to org '{}' but '{}' was requested", + "auth login '{profile_name}' belongs to org '{}' but '{}' was requested", profile.org_name.as_deref().unwrap_or("(none)"), org_name ); @@ -1916,7 +1910,7 @@ async fn ensure_setup_auth( org_name } None => profile.org_name.as_deref().ok_or_else(|| { - anyhow!("profile '{profile_name}' does not have a default org") + anyhow!("auth login '{profile_name}' does not have a default org") })?, }; @@ -1969,9 +1963,7 @@ async fn ensure_setup_auth( .iter() .filter(|profile| profile.org_name.as_deref() == Some(org_name)) .count(); - if base.prefer_profile && matching_profile_count == 0 { - bail!("no profile found for org '{org_name}'"); - } + let _ = matching_profile_count; if let Some(org) = find_available_org(&available_orgs, org_name) { let client = build_api_key_client(base, api_key, org).await?; @@ -2017,92 +2009,6 @@ async fn ensure_setup_auth( .await; } - if base.prefer_profile { - if stored_profiles.is_empty() { - let matched_orgs = match project_name.as_deref() { - Some(project_name) => { - orgs_with_project(base, api_key, &available_orgs, project_name).await? - } - None => available_orgs.clone(), - }; - if matched_orgs.is_empty() { - return ensure_profile_or_setup_browser_auth_context( - base, - prompt_for_profile_choice, - needs_api_key, - project_name.as_deref(), - project_was_explicit, - org_name.as_deref(), - ) - .await; - } - let org = select_api_key_org_for_setup( - base, - &matched_orgs, - project_name.as_deref(), - &[], - )?; - let client = build_api_key_client(base, api_key, &org).await?; - let api_url = base - .api_url - .clone() - .or_else(|| org.api_url.clone()) - .unwrap_or_else(|| DEFAULT_API_URL.to_string()); - auth::commit_api_key_profile( - &org.name, - api_key, - api_url, - base.app_url.clone(), - Some(org.name.clone()), - )?; - return build_setup_auth_context(base, client, false, needs_api_key, None).await; - } - - let preferred_org_names = matching_profile_org_names(&stored_profiles, None); - let matching_orgs = available_orgs - .iter() - .filter(|org| preferred_org_names.iter().any(|name| name == &org.name)) - .cloned() - .collect::>(); - if matching_orgs.is_empty() { - return ensure_profile_or_setup_browser_auth_context( - base, - prompt_for_profile_choice, - needs_api_key, - project_name.as_deref(), - project_was_explicit, - org_name.as_deref(), - ) - .await; - } - - let candidate_orgs = match project_name.as_deref() { - Some(project_name) => { - orgs_with_project(base, api_key, &matching_orgs, project_name).await? - } - None => matching_orgs, - }; - if candidate_orgs.is_empty() { - return ensure_profile_or_setup_browser_auth_context( - base, - prompt_for_profile_choice, - needs_api_key, - project_name.as_deref(), - project_was_explicit, - org_name.as_deref(), - ) - .await; - } - let org = select_api_key_org_for_setup( - base, - &candidate_orgs, - project_name.as_deref(), - &preferred_org_names, - )?; - let client = build_api_key_client(base, api_key, &org).await?; - return build_setup_auth_context(base, client, false, needs_api_key, None).await; - } - let candidate_orgs = match project_name.as_deref() { Some(project_name) => { orgs_with_project(base, api_key, &available_orgs, project_name).await? @@ -5435,12 +5341,11 @@ mod tests { no_color: false, no_input: false, profile: None, - profile_explicit: false, org_name: None, project: None, api_key: None, api_key_source: None, - prefer_profile: false, + prefer_api_key: false, api_url: None, app_url: None, ca_cert: None, @@ -5778,6 +5683,7 @@ mod tests { let profiles = vec![ auth::ProfileInfo { name: "zeta".to_string(), + auth_method: "api_key".to_string(), org_name: Some("Zeta Org".to_string()), user_name: None, email: None, @@ -5785,6 +5691,7 @@ mod tests { }, auth::ProfileInfo { name: "alpha".to_string(), + auth_method: "api_key".to_string(), org_name: Some("Alpha Org".to_string()), user_name: None, email: None, @@ -5813,6 +5720,7 @@ mod tests { base.profile = Some("missing".to_string()); let profiles = vec![auth::ProfileInfo { name: "work".to_string(), + auth_method: "api_key".to_string(), org_name: Some("Acme".to_string()), user_name: None, email: None, @@ -5821,7 +5729,7 @@ mod tests { let err = resolve_profile_name_for_setup(&base, &profiles, false).expect_err("missing profile"); - assert!(err.to_string().contains("profile 'missing' not found")); + assert!(err.to_string().contains("auth login 'missing' not found")); } #[test] @@ -5829,6 +5737,7 @@ mod tests { let profiles = vec![ auth::ProfileInfo { name: "Acme".to_string(), + auth_method: "api_key".to_string(), org_name: Some("Acme".to_string()), user_name: None, email: None, @@ -5836,6 +5745,7 @@ mod tests { }, auth::ProfileInfo { name: "Acme-2".to_string(), + auth_method: "api_key".to_string(), org_name: Some("Acme".to_string()), user_name: None, email: None, diff --git a/src/status.rs b/src/status.rs index 93b555ed..6893975e 100644 --- a/src/status.rs +++ b/src/status.rs @@ -4,7 +4,7 @@ use serde::Serialize; use crate::args::BaseArgs; use crate::auth; -use crate::{config, utils::resolve_profile_info}; +use crate::config; #[derive(Debug, Clone, Args)] #[command(after_help = "\ @@ -19,13 +19,14 @@ pub struct StatusArgs {} struct StatusOutput { org: Option, project: Option, - profile: Option, #[serde(skip_serializing_if = "Option::is_none")] user_name: Option, #[serde(skip_serializing_if = "Option::is_none")] user_email: Option, #[serde(skip_serializing_if = "Option::is_none")] api_key_hint: Option, + #[serde(skip_serializing_if = "Option::is_none")] + auth_method: Option, source: Option, } @@ -40,6 +41,13 @@ fn format_identity(p: &auth::ProfileInfo) -> Option { } } +fn format_auth(p: &auth::ProfileInfo) -> String { + match format_identity(p) { + Some(identity) => format!("{} — {identity}", p.auth_method), + None => p.auth_method.clone(), + } +} + pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { let global_path = config::global_path().ok(); let global_cfg = config::load_global().unwrap_or_default(); @@ -51,7 +59,7 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { let cli_org = cli_flag_value(&["--org", "-o"]); let cli_project = cli_flag_value(&["--project", "-p"]); - let (mut org, mut project, source) = resolve_config( + let (org, mut project, source) = resolve_config( cli_org, cli_project, &global_cfg, @@ -60,27 +68,7 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { &global_path, ); let merged_cfg = global_cfg.merge(&local_cfg); - let selected_profile = base - .profile - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .or_else(|| { - merged_cfg - .profile - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - }); - let profile_info = resolve_profile_info(selected_profile.as_deref(), org.as_deref()); - - if selected_profile.is_some() && org.as_deref().map(str::trim).is_none_or(str::is_empty) { - if let Some(profile_org) = profile_info.as_ref().and_then(|p| p.org_name.clone()) { - org = Some(profile_org); - } - } + let auth_info = auth::active_auth_info(&base, org.as_deref()); if base .project @@ -88,27 +76,17 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { .map(str::trim) .is_none_or(str::is_empty) { - let mut project_base = base.clone(); - if project_base - .profile - .as_deref() - .map(str::trim) - .is_none_or(str::is_empty) - { - project_base.profile = selected_profile.clone(); - } - project = - config::project_from_config_for_context(&project_base, &merged_cfg, org.as_deref()); + project = config::project_from_config_for_context(&base, &merged_cfg, org.as_deref()); } if base.json { let output = StatusOutput { org, project, - profile: profile_info.as_ref().map(|p| p.name.clone()), - user_name: profile_info.as_ref().and_then(|p| p.user_name.clone()), - user_email: profile_info.as_ref().and_then(|p| p.email.clone()), - api_key_hint: profile_info.as_ref().and_then(|p| p.api_key_hint.clone()), + user_name: auth_info.as_ref().and_then(|p| p.user_name.clone()), + user_email: auth_info.as_ref().and_then(|p| p.email.clone()), + api_key_hint: auth_info.as_ref().and_then(|p| p.api_key_hint.clone()), + auth_method: auth_info.as_ref().map(|p| p.auth_method.clone()), source, }; println!("{}", serde_json::to_string(&output)?); @@ -116,13 +94,19 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { } if base.verbose { - println!("org: {}", org.as_deref().unwrap_or("(unset)")); + let org_display = if org.is_none() + && auth_info + .as_ref() + .is_some_and(|p| p.auth_method == "oauth" && p.org_name.is_none()) + { + "cross-org" + } else { + org.as_deref().unwrap_or("(unset)") + }; + println!("org: {org_display}"); println!("project: {}", project.as_deref().unwrap_or("(unset)")); - if let Some(ref p) = profile_info { - println!("profile: {}", p.name); - if let Some(id) = format_identity(p) { - println!("user: {id}"); - } + if let Some(ref p) = auth_info { + println!("auth: {}", format_auth(p)); } if let Some(src) = source { println!("source: {src}"); @@ -134,14 +118,19 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { _ => unreachable!(), }; println!("{scope}"); - let profile_line = match &profile_info { - Some(p) => match format_identity(p) { - Some(id) => format!(" profile: {} — {id}", p.name), - None => format!(" profile: {}", p.name), - }, - None => " profile: (none)".to_string(), + let auth_line = match &auth_info { + Some(p) => format!(" auth: {}", format_auth(p)), + None => " auth: (none)".to_string(), }; - println!("{profile_line}"); + println!("{auth_line}"); + } else if auth_info + .as_ref() + .is_some_and(|p| p.auth_method == "oauth" && p.org_name.is_none()) + { + println!("cross-org"); + if let Some(p) = &auth_info { + println!(" auth: {}", format_auth(p)); + } } else { println!("No org/project configured. Run `bt switch` to set one."); } @@ -345,6 +334,12 @@ mod tests { ) -> auth::ProfileInfo { auth::ProfileInfo { name: name.into(), + auth_method: if api_key_hint.is_some() { + "api_key" + } else { + "oauth" + } + .into(), org_name: None, user_name: user_name.map(Into::into), email: email.map(Into::into), diff --git a/src/switch.rs b/src/switch.rs index 35495bb6..d079a212 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -54,59 +54,15 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { let current_cfg = config::load().unwrap_or_default(); let (resolved_org, resolved_project) = args.resolve_target(&base); let mut interactive = false; - let has_api_key_override = base - .api_key - .as_ref() - .is_some_and(|value| !value.trim().is_empty()); - - let profile_name = match &resolved_org { - Some(org_or_profile) => { - if base.profile.is_some() { - None - } else { - let profiles = auth::list_profiles()?; - Some(auth::resolve_org_to_profile(org_or_profile, &profiles)?) - } - } - None => resolve_profile_for_switch( - has_api_key_override, - resolved_project.is_none(), - is_interactive(), - || auth::select_profile_interactive(current_cfg.org.as_deref()), - &mut interactive, - )?, - }; - // When we resolved a profile from an org identifier, clear org_name — the raw identifier - // (e.g. "staging") may differ from the profile's actual org (e.g. "staging-org"). Letting - // org_name stay would override the profile's stored org_name in resolve_auth_from_store. - // - // When no org was specified (project-only switch), load the current config org so - // resolve_auth can find the right profile for authentication. - let login_base = match &profile_name { - Some(profile) if base.profile.is_none() => BaseArgs { - profile: Some(profile.clone()), - org_name: None, - ..base.clone() - }, - _ => { - let mut b = base.clone(); - if !has_api_key_override && b.org_name.is_none() && b.profile.is_none() { - b.org_name = current_cfg.org.clone(); - } - if !has_api_key_override && b.org_name.is_none() && b.profile.is_none() { - let profiles = auth::list_profiles()?; - if profiles.len() > 1 { - let names: Vec<&str> = profiles.iter().map(|p| p.name.as_str()).collect(); - bail!( - "multiple auth profiles found: {}. Use --profile to disambiguate.", - names.join(", ") - ); - } - } - b - } - }; + let mut login_base = base.clone(); + if login_base.org_name.is_none() { + let saved_org = select_saved_auth_org_for_switch()?; + login_base.org_name = resolved_org + .clone() + .or_else(|| current_cfg.org.clone()) + .or(saved_org); + } let ctx = login(&login_base).await?; let client = ApiClient::new(&ctx)?; @@ -147,15 +103,7 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { }; let mut cfg = config::load_file(&path); - let config_profile = - config::trimmed_option(profile_name.as_deref().or(base.profile.as_deref())) - .map(str::to_string); - apply_switch_config( - &mut cfg, - config_profile.as_deref(), - Some(&org_name), - Some(&project), - ); + apply_switch_config(&mut cfg, None, Some(&org_name), Some(&project)); config::save_file(&path, &cfg) .context(format!("Could not save config to {}", path.display()))?; @@ -164,7 +112,6 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { "org": org_name, "project": project.name, "project_id": project.id, - "profile": config_profile, "scope": scope, "path": path.display().to_string(), }); @@ -181,6 +128,30 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { Ok(()) } +fn select_saved_auth_org_for_switch() -> Result> { + if !is_interactive() { + return Ok(None); + } + + let mut orgs = auth::list_profiles()? + .into_iter() + .filter_map(|profile| profile.org_name) + .filter(|org| !org.trim().is_empty()) + .collect::>(); + orgs.sort(); + orgs.dedup(); + + match orgs.len() { + 0 => Ok(None), + 1 => Ok(orgs.into_iter().next()), + _ => { + let labels = orgs.iter().map(String::as_str).collect::>(); + let idx = crate::ui::fuzzy_select("Select organization", &labels, 0)?; + Ok(Some(orgs[idx].clone())) + } + } +} + pub(crate) fn select_scope() -> Result<(std::path::PathBuf, &'static str)> { let global = config::global_path()?; let local = config::local_path().unwrap(); @@ -258,13 +229,11 @@ pub(crate) async fn validate_or_create_project( pub(crate) fn apply_switch_config( cfg: &mut config::Config, - profile_name: Option<&str>, + _profile_name: Option<&str>, org_name: Option<&str>, project: Option<&api::Project>, ) { - if let Some(profile_name) = config::trimmed_option(profile_name) { - cfg.profile = Some(profile_name.to_string()); - } + cfg.profile = None; cfg.org = config::trimmed_option(org_name).map(str::to_string); match project { Some(project) => { @@ -278,6 +247,7 @@ pub(crate) fn apply_switch_config( } } +#[cfg(test)] fn resolve_profile_for_switch( has_api_key_override: bool, prompting_for_project_only: bool, @@ -326,12 +296,11 @@ mod tests { no_color: false, no_input: false, profile: None, - profile_explicit: false, org_name: org.map(String::from), project: project.map(String::from), api_key: None, api_key_source: None, - prefer_profile: false, + prefer_api_key: false, api_url: None, app_url: None, ca_cert: None, @@ -342,6 +311,7 @@ mod tests { fn profile_info(name: &str, org_name: Option<&str>) -> ProfileInfo { ProfileInfo { name: name.to_string(), + auth_method: "api_key".to_string(), org_name: org_name.map(String::from), user_name: None, email: None, @@ -548,14 +518,14 @@ mod tests { apply_switch_config(&mut cfg, Some("work"), Some("acme-org"), Some(&project)); - assert_eq!(cfg.profile.as_deref(), Some("work")); + assert_eq!(cfg.profile, None); assert_eq!(cfg.org.as_deref(), Some("acme-org")); assert_eq!(cfg.project.as_deref(), Some("my-project")); assert_eq!(cfg.project_id.as_deref(), Some("proj_123")); } #[test] - fn apply_switch_config_preserves_existing_profile_when_no_new_profile() { + fn apply_switch_config_clears_existing_profile() { let mut cfg = config::Config { profile: Some("work".to_string()), ..Default::default() @@ -569,7 +539,7 @@ mod tests { apply_switch_config(&mut cfg, None, Some("acme-org"), Some(&project)); - assert_eq!(cfg.profile.as_deref(), Some("work")); + assert_eq!(cfg.profile, None); assert_eq!(cfg.project.as_deref(), Some("next-project")); } @@ -585,7 +555,7 @@ mod tests { apply_switch_config(&mut cfg, Some("next"), None, None); - assert_eq!(cfg.profile.as_deref(), Some("next")); + assert_eq!(cfg.profile, None); assert_eq!(cfg.org, None); assert_eq!(cfg.project, None); assert_eq!(cfg.project_id, None); diff --git a/src/traces.rs b/src/traces.rs index 6ef7ff1c..cc59c4ca 100644 --- a/src/traces.rs +++ b/src/traces.rs @@ -33,7 +33,7 @@ use unicode_width::UnicodeWidthStr; use urlencoding::{decode, encode}; use crate::args::BaseArgs; -use crate::auth::{self, login}; +use crate::auth::login; use crate::experiments::api as experiments_api; use crate::http::ApiClient; use crate::ui::{fuzzy_select, is_interactive, with_spinner}; @@ -1060,7 +1060,7 @@ async fn run_logs_command(base: BaseArgs, client: ApiClient, args: LogsArgs) -> let has_more = next_cursor.is_some(); let rows = parse_summary_rows(response.data); let object_ref_arg = format_object_ref_arg(&object_ref); - let profile_flag = base.profile.as_deref(); + let org_flag = base.org_name.as_deref(); if base.json { let payload = json!({ @@ -1082,7 +1082,7 @@ async fn run_logs_command(base: BaseArgs, client: ApiClient, args: LogsArgs) -> has_more, &object_ref_arg, next_cursor.as_deref(), - profile_flag, + org_flag, args.limit, ), }); @@ -1092,7 +1092,7 @@ async fn run_logs_command(base: BaseArgs, client: ApiClient, args: LogsArgs) -> &rows, args.list_mode, &object_ref_arg, - profile_flag, + org_flag, args.limit, args.preview_length, next_cursor.as_deref(), @@ -1141,7 +1141,7 @@ async fn run_trace_command(base: BaseArgs, client: ApiClient, args: TraceArgs) - next_cursor_if_full_page(response.cursor.clone(), response.data.len(), selector.limit); let has_more = next_cursor.is_some(); let object_ref_arg = format_object_ref_arg(&target.object_ref); - let profile_flag = base.profile.as_deref(); + let org_flag = base.org_name.as_deref(); if base.json { let payload = json!({ @@ -1164,7 +1164,7 @@ async fn run_trace_command(base: BaseArgs, client: ApiClient, args: TraceArgs) - &object_ref_arg, &target.root_span_id, next_cursor.as_deref(), - profile_flag, + org_flag, selector.limit, ), }); @@ -1174,7 +1174,7 @@ async fn run_trace_command(base: BaseArgs, client: ApiClient, args: TraceArgs) - &target.root_span_id, &response.data, &object_ref_arg, - profile_flag, + org_flag, selector.limit, selector.preview_length, next_cursor.as_deref(), @@ -1231,11 +1231,11 @@ async fn run_thread_command(base: BaseArgs, client: ApiClient, args: ThreadArgs) }, "summary": summary, "messages": messages, - "hints": thread_hints(base.profile.as_deref()), + "hints": thread_hints(base.org_name.as_deref()), }); println!("{}", serde_json::to_string_pretty(&payload)?); } else { - print_thread_text(&target, &messages, &summary, base.profile.as_deref()); + print_thread_text(&target, &messages, &summary, base.org_name.as_deref()); } Ok(()) @@ -1287,14 +1287,14 @@ async fn run_waterfall_command( "has_more": has_more, }, "waterfall": waterfall, - "hints": waterfall_hints(base.profile.as_deref(), has_more), + "hints": waterfall_hints(base.org_name.as_deref(), has_more), }); println!("{}", serde_json::to_string_pretty(&payload)?); } else { waterfall::print_waterfall_text( &target, &waterfall, - base.profile.as_deref(), + base.org_name.as_deref(), selector.limit, has_more, ); @@ -5683,9 +5683,9 @@ fn format_object_ref_arg(object_ref: &ObjectRef) -> String { format!("{}:{}", object_ref.object_type, object_ref.object_name) } -fn profile_flag_suffix(profile: Option<&str>) -> String { - match profile.filter(|p| !p.trim().is_empty()) { - Some(profile) => format!(" --profile {}", btql_quote(profile)), +fn profile_flag_suffix(org: Option<&str>) -> String { + match org.filter(|p| !p.trim().is_empty()) { + Some(org) => format!(" --org {}", btql_quote(org)), None => String::new(), } } @@ -5715,21 +5715,7 @@ fn parse_startup_trace_url_from_view_args(args: &ViewArgs) -> Result) -> BaseArgs { - apply_url_hints_with_profile_resolver(base, parsed_url, |org| { - let profiles = auth::list_profiles().ok()?; - auth::resolve_org_to_profile(org, &profiles).ok() - }) -} - -fn apply_url_hints_with_profile_resolver( - mut base: BaseArgs, - parsed_url: Option<&ParsedTraceUrl>, - resolve_profile_for_org: F, -) -> BaseArgs -where - F: Fn(&str) -> Option, -{ +fn apply_url_hints_to_base(mut base: BaseArgs, parsed_url: Option<&ParsedTraceUrl>) -> BaseArgs { let Some(parsed) = parsed_url else { return base; }; @@ -5742,28 +5728,30 @@ where return base; }; - let has_profile_override = base - .profile - .as_deref() - .map(str::trim) - .is_some_and(|v| !v.is_empty()); let has_org_override = base .org_name .as_deref() .map(str::trim) .is_some_and(|v| !v.is_empty()); - if !has_org_override && !has_profile_override { + if !has_org_override { base.org_name = Some(url_org.to_string()); } - if !has_profile_override && !has_org_override { - if let Some(profile_name) = resolve_profile_for_org(url_org) { - base.profile = Some(profile_name); - } - } base } +#[cfg(test)] +fn apply_url_hints_with_profile_resolver( + base: BaseArgs, + parsed_url: Option<&ParsedTraceUrl>, + _resolve_profile_for_org: F, +) -> BaseArgs +where + F: Fn(&str) -> Option, +{ + apply_url_hints_to_base(base, parsed_url) +} + fn select_startup_url( long_url: Option<&str>, positional_url: Option<&str>, @@ -6729,12 +6717,11 @@ mod tests { no_color: false, no_input: false, profile: None, - profile_explicit: false, org_name: None, project: None, api_key: None, api_key_source: None, - prefer_profile: false, + prefer_api_key: false, api_url: None, app_url: None, ca_cert: None, @@ -6963,7 +6950,7 @@ mod tests { } #[test] - fn apply_url_hints_infers_profile_from_url_org() { + fn apply_url_hints_infers_org_from_url() { let base = base_args(); let parsed = parsed_url_with_org("Lovable"); @@ -6972,21 +6959,20 @@ mod tests { }); assert_eq!(updated.org_name.as_deref(), Some("Lovable")); - assert_eq!(updated.profile.as_deref(), Some("lovable-profile")); + assert_eq!(updated.profile, None); } #[test] - fn apply_url_hints_preserves_explicit_profile() { + fn apply_url_hints_preserves_explicit_org() { let mut base = base_args(); - base.profile = Some("explicit-profile".to_string()); + base.org_name = Some("explicit-org".to_string()); let parsed = parsed_url_with_org("Lovable"); let updated = apply_url_hints_with_profile_resolver(base, Some(&parsed), |_| { Some("other".to_string()) }); - assert_eq!(updated.profile.as_deref(), Some("explicit-profile")); - assert!(updated.org_name.is_none()); + assert_eq!(updated.org_name.as_deref(), Some("explicit-org")); } #[test] diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 05429ee0..52a4f874 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -14,4 +14,4 @@ pub use git::GitRepo; pub(crate) use ids::new_uuid_id; pub(crate) use json_object::lookup_object_path; pub use plurals::pluralize; -pub(crate) use profile::{profile_author_slug, resolve_profile_info, sanitize_name_segment}; +pub(crate) use profile::{profile_author_slug, sanitize_name_segment}; diff --git a/src/utils/profile.rs b/src/utils/profile.rs index b97812b2..620d681c 100644 --- a/src/utils/profile.rs +++ b/src/utils/profile.rs @@ -1,13 +1,6 @@ -use crate::auth::{self, ProfileInfo}; - -pub(crate) fn resolve_profile_info( - profile: Option<&str>, - org: Option<&str>, -) -> Option { - let profiles = auth::list_profiles().ok()?; - resolve_profile_info_from_profiles(profile, org, profiles) -} +use crate::auth::ProfileInfo; +#[cfg(test)] fn resolve_profile_info_from_profiles( profile: Option<&str>, org: Option<&str>, @@ -40,6 +33,17 @@ fn resolve_profile_info_from_profiles( .into_iter() .find(|profile| profile.name == profile_name); } + let oauth_matches: Vec<&ProfileInfo> = org_matches + .iter() + .copied() + .filter(|profile| profile.email.is_some() || profile.user_name.is_some()) + .collect(); + if oauth_matches.len() == 1 { + let profile_name = oauth_matches[0].name.clone(); + return profiles + .into_iter() + .find(|profile| profile.name == profile_name); + } return None; } @@ -54,7 +58,7 @@ pub(crate) fn profile_author_slug(profile: &ProfileInfo) -> Option { [ profile.user_name.as_deref(), profile.email.as_deref().and_then(email_local_part), - Some(profile.name.as_str()), + profile.org_name.as_deref(), ] .into_iter() .flatten() @@ -105,6 +109,12 @@ mod tests { ) -> ProfileInfo { ProfileInfo { name: name.to_string(), + auth_method: if email.is_some() || user_name.is_some() { + "oauth" + } else { + "api_key" + } + .to_string(), org_name: org_name.map(ToOwned::to_owned), user_name: user_name.map(ToOwned::to_owned), email: email.map(ToOwned::to_owned), @@ -155,12 +165,9 @@ mod tests { } #[test] - fn profile_author_slug_falls_back_to_profile_name() { + fn profile_author_slug_ignores_internal_profile_name() { let profile = profile_info("Work Profile", None, None, None); - assert_eq!( - profile_author_slug(&profile).as_deref(), - Some("work-profile") - ); + assert_eq!(profile_author_slug(&profile), None); } #[test] diff --git a/tests/cli.rs b/tests/cli.rs index acb09bfd..a120455c 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -10,7 +10,6 @@ fn bt_command() -> Command { fn clear_braintrust_auth_env(cmd: &mut Command) { for key in [ "BRAINTRUST_API_KEY", - "BRAINTRUST_PROFILE", "BRAINTRUST_ORG_NAME", "BRAINTRUST_DEFAULT_PROJECT", ] { @@ -113,13 +112,7 @@ fn top_level_help_shows_update_not_self() { fn topics_report_help_accepts_global_org_short_conflict_free() { bt_command() .args([ - "topics", - "report", - "--profile", - "test-profile", - "--id", - "fn_123", - "--help", + "topics", "report", "--org", "test-org", "--id", "fn_123", "--help", ]) .assert() .success() @@ -167,7 +160,7 @@ fn status_json_keeps_local_org_when_global_profile_has_different_org() { .success() .stdout(predicate::str::contains(r#""org":"local-org""#)) .stdout(predicate::str::contains(r#""project":"local-project""#)) - .stdout(predicate::str::contains(r#""profile":"default-profile""#)) + .stdout(predicate::str::contains(r#""profile""#).not()) .stdout(predicate::str::contains(r#""org":"profile-org""#).not()); } @@ -459,12 +452,12 @@ fn setup_mcp_only_requires_auth_in_non_interactive_mode() { .assert() .failure() .stderr(predicate::str::contains( - "profile selection required in non-interactive mode", + "auth org selection required in non-interactive mode", )); } #[test] -fn datasets_requires_profile_selection_when_multiple_profiles_exist() { +fn datasets_requires_org_selection_when_multiple_api_key_logins_exist() { let repo = make_git_repo(); let home = tempfile::tempdir().expect("home tempdir"); let config_home = tempfile::tempdir().expect("config tempdir"); @@ -481,8 +474,10 @@ fn datasets_requires_profile_selection_when_multiple_profiles_exist() { .args(["datasets", "--no-input"]) .assert() .failure() - .stderr(predicate::str::contains("multiple auth profiles available")) - .stderr(predicate::str::contains("--profile ")) + .stderr(predicate::str::contains( + "multiple API key logins available", + )) + .stderr(predicate::str::contains("--org ")) .stderr(predicate::str::contains("alpha")) .stderr(predicate::str::contains("beta")); } diff --git a/tests/functions.rs b/tests/functions.rs index 03f70300..f91c25d2 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -669,15 +669,17 @@ fn functions_push_rejects_invalid_language() { #[test] fn functions_push_requires_app_url_with_custom_api_url() { + let config_dir = tempdir().expect("config tempdir"); let output = Command::new(bt_binary_path()) .arg("functions") .arg("--json") .arg("push") + .env("XDG_CONFIG_HOME", config_dir.path()) + .env("APPDATA", config_dir.path()) .env("BRAINTRUST_API_KEY", "test-key") .env("BRAINTRUST_API_URL", "http://127.0.0.1:1") .env_remove("BRAINTRUST_APP_URL") .env_remove("BRAINTRUST_ORG_NAME") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run push with custom API URL and no app URL"); @@ -713,10 +715,8 @@ fn auth_profiles_ignores_api_key_env_override() { assert!(output.status.success()); let stdout = String::from_utf8_lossy(&output.stdout); - let stderr = String::from_utf8_lossy(&output.stderr); - assert!(stdout.contains("No saved profiles. Run `bt auth login` to create one.")); + assert!(stdout.contains("No saved auth logins. Run `bt auth login` to create one.")); assert!(!stdout.contains("Auth source: --api-key/BRAINTRUST_API_KEY override")); - assert!(!stderr.contains("pass --prefer-profile or unset BRAINTRUST_API_KEY")); } #[test] @@ -732,10 +732,8 @@ fn auth_profiles_ignores_api_key_from_dotenv() { assert!(output.status.success()); let stdout = String::from_utf8_lossy(&output.stdout); - let stderr = String::from_utf8_lossy(&output.stderr); - assert!(stdout.contains("No saved profiles. Run `bt auth login` to create one.")); + assert!(stdout.contains("No saved auth logins. Run `bt auth login` to create one.")); assert!(!stdout.contains("Auth source: --api-key/BRAINTRUST_API_KEY override")); - assert!(!stderr.contains("pass --prefer-profile or unset BRAINTRUST_API_KEY")); } #[test] @@ -753,28 +751,8 @@ fn auth_profiles_json_with_no_profiles_emits_empty_array() { assert_eq!(stdout, "[]"); } -#[test] -fn auth_profiles_profile_not_found_is_actionable() { - let cwd = tempdir().expect("create temp cwd"); - let config_dir = tempdir().expect("create temp config dir"); - - let output = auth_profiles_command(cwd.path(), config_dir.path()) - .arg("--profile") - .arg("test-profile") - .output() - .expect("run bt auth profiles --profile test-profile"); - - assert!(!output.status.success()); - let stderr = String::from_utf8_lossy(&output.stderr); - assert!(stderr.contains("profile 'test-profile' not found")); - assert!( - stderr.contains("run `bt auth profiles` to see available profiles"), - "expected actionable hint, got: {stderr}" - ); -} - -/// Seed a synthetic api_key `test-profile` so verification reports "missing" -/// without touching the network or keychain. +/// Seed a synthetic API-key auth login so refresh can fail before touching the +/// network or keychain. fn seed_api_key_profile(config_dir: &Path) { fs::create_dir_all(config_dir.join("bt")).expect("create bt config dir"); fs::write( @@ -794,7 +772,6 @@ fn auth_sub_command(cwd: &Path, config_dir: &Path, sub: &[&str]) -> Command { .env("XDG_CONFIG_HOME", config_dir) .env("APPDATA", config_dir) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .env_remove("BRAINTRUST_ORG_NAME") .env_remove("BRAINTRUST_API_URL") .env_remove("BRAINTRUST_APP_URL") @@ -817,9 +794,9 @@ fn auth_logout_json_with_no_profiles_emits_empty_status() { } #[test] -fn auth_refresh_errors_for_api_key_profile_even_with_json() { - // --json must not swallow a real error: an api_key profile cannot be - // refreshed, and the command should fail with an actionable message. +fn auth_refresh_errors_when_no_oauth_login_exists_even_with_json() { + // --json must not swallow a real error: refresh only applies to OAuth + // logins, and an org with only API-key auth should fail actionably. let cwd = tempdir().expect("create temp cwd"); let config_dir = tempdir().expect("create temp config dir"); @@ -828,15 +805,15 @@ fn auth_refresh_errors_for_api_key_profile_even_with_json() { let output = auth_sub_command( cwd.path(), config_dir.path(), - &["refresh", "--profile", "test-profile", "--json"], + &["refresh", "--org", "test-org", "--json"], ) .output() - .expect("run bt auth refresh --profile test-profile --json"); + .expect("run bt auth refresh --org test-org --json"); assert!(!output.status.success()); let stderr = String::from_utf8_lossy(&output.stderr); assert!( - stderr.contains("only applies to oauth profiles"), + stderr.contains("no OAuth login selected"), "expected oauth-only refresh hint, got: {stderr}" ); assert!(String::from_utf8_lossy(&output.stdout).trim().is_empty()); @@ -1877,7 +1854,6 @@ exit 24 .env("BRAINTRUST_API_URL", &server.base_url) .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run bt functions push"); @@ -2066,7 +2042,6 @@ exit 24 .env("BRAINTRUST_API_URL", &server.base_url) .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run bt functions push"); @@ -2200,7 +2175,6 @@ exit 24 .env("BRAINTRUST_API_URL", &server.base_url) .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run bt functions push"); @@ -2347,7 +2321,6 @@ exit 24 .env("BRAINTRUST_API_URL", &server.base_url) .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run bt functions push"); @@ -2410,7 +2383,6 @@ async fn functions_view_by_positional_id_does_not_require_project_context() { .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") .env("BRAINTRUST_NO_INPUT", "1") - .env_remove("BRAINTRUST_PROFILE") .env_remove("BRAINTRUST_DEFAULT_PROJECT") .env_remove("BT_FUNCTIONS_VIEW_ID") .env_remove("BT_FUNCTIONS_VIEW_VERSION") @@ -2485,7 +2457,6 @@ async fn functions_view_by_id_passes_version() { .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") .env("BRAINTRUST_NO_INPUT", "1") - .env_remove("BRAINTRUST_PROFILE") .env_remove("BRAINTRUST_DEFAULT_PROJECT") .env_remove("BT_FUNCTIONS_VIEW_ID") .env_remove("BT_FUNCTIONS_VIEW_VERSION") @@ -2569,7 +2540,6 @@ async fn functions_view_by_slug_passes_version() { .env("BRAINTRUST_DEFAULT_PROJECT", "mock-project") .env("BRAINTRUST_NO_COLOR", "1") .env("BRAINTRUST_NO_INPUT", "1") - .env_remove("BRAINTRUST_PROFILE") .env_remove("BT_FUNCTIONS_VIEW_ID") .env_remove("BT_FUNCTIONS_VIEW_VERSION") .output() @@ -2660,7 +2630,6 @@ async fn functions_pull_works_against_mock_api() { .env("BRAINTRUST_API_URL", &server.base_url) .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run bt functions pull"); @@ -2779,7 +2748,6 @@ async fn functions_pull_skips_untracked_existing_file_without_force() { .env("BRAINTRUST_API_URL", &server.base_url) .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run bt functions pull"); @@ -2859,7 +2827,6 @@ async fn functions_pull_selector_with_unsupported_only_rows_still_succeeds() { .env("BRAINTRUST_API_URL", &server.base_url) .env("BRAINTRUST_APP_URL", &server.base_url) .env("BRAINTRUST_NO_COLOR", "1") - .env_remove("BRAINTRUST_PROFILE") .output() .expect("run bt functions pull"); From 894226823d4c36e9bb01bd10a2289aa166c512a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 16 Jul 2026 13:19:33 -0700 Subject: [PATCH 02/24] follow up and remove all profile code --- src/args.rs | 9 +- src/auth.rs | 1368 ++------------------------------------ src/config/mod.rs | 33 +- src/datasets/pipeline.rs | 3 +- src/experiments/mod.rs | 7 +- src/functions/push.rs | 3 +- src/main.rs | 28 + src/setup/mod.rs | 357 +++------- src/status.rs | 164 +++-- src/switch.rs | 213 +----- src/traces.rs | 41 +- src/traces/waterfall.rs | 8 +- src/utils/profile.rs | 91 +-- tests/cli.rs | 9 +- 14 files changed, 353 insertions(+), 1981 deletions(-) diff --git a/src/args.rs b/src/args.rs index f0f20d64..930cfe9d 100644 --- a/src/args.rs +++ b/src/args.rs @@ -38,13 +38,13 @@ pub struct BaseArgs { #[arg(long, env = "BRAINTRUST_NO_INPUT", global = true, value_parser = clap::builder::BoolishValueParser::new(), default_value_t = false)] pub no_input: bool, - #[arg(skip)] - pub profile: Option, - /// Override active org (or via BRAINTRUST_ORG_NAME) #[arg(short = 'o', long = "org", env = "BRAINTRUST_ORG_NAME", global = true)] pub org_name: Option, + #[arg(skip)] + pub org_name_source: Option, + /// Override active project #[arg( short = 'p', @@ -55,6 +55,9 @@ pub struct BaseArgs { )] pub project: Option, + #[arg(skip)] + pub project_source: Option, + /// Override stored API key (or via BRAINTRUST_API_KEY) #[arg(long, env = "BRAINTRUST_API_KEY", global = true, hide = true)] pub api_key: Option, diff --git a/src/auth.rs b/src/auth.rs index 32f1ae1d..8ab0df8c 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -62,7 +62,6 @@ pub struct ResolvedAuth { #[derive(Debug, Clone)] pub struct ProfileInfo { - pub name: String, pub auth_method: String, pub org_name: Option, pub user_name: Option, @@ -70,12 +69,6 @@ pub struct ProfileInfo { pub api_key_hint: Option, } -#[derive(Debug, Clone)] -pub(crate) struct StoredProfileInfo { - pub name: String, - pub org_name: Option, -} - #[derive(Debug, Clone, PartialEq, Eq)] pub struct AvailableOrg { pub id: String, @@ -127,120 +120,11 @@ pub fn list_profiles() -> Result> { let store = load_auth_store()?; Ok(store .profiles - .iter() - .map(|(name, p)| profile_info_from_store_entry(name, p)) - .collect()) -} - -pub(crate) fn list_stored_profiles() -> Result> { - let store = load_auth_store()?; - Ok(store - .profiles - .iter() - .map(|(name, profile)| StoredProfileInfo { - name: name.clone(), - org_name: profile.org_name.clone(), - }) + .values() + .map(profile_info_from_store_entry) .collect()) } -#[cfg(test)] -pub fn resolve_org_to_profile(identifier: &str, profiles: &[ProfileInfo]) -> Result { - if profiles.is_empty() { - bail!("no auth profiles found. Run `bt auth login` to create one."); - } - - if let Some(p) = profiles.iter().find(|p| p.name == identifier) { - return Ok(p.name.clone()); - } - - let matches: Vec<&ProfileInfo> = profiles - .iter() - .filter(|p| p.org_name.as_deref() == Some(identifier)) - .collect(); - - match matches.len() { - 0 => { - let available: Vec = profiles - .iter() - .filter_map(|p| { - p.org_name - .as_ref() - .map(|org| format!(" {} (profile: {})", org, p.name)) - }) - .collect(); - bail!( - "no profile found for '{identifier}'.\nAvailable:\n{}", - available.join("\n") - ); - } - 1 => Ok(matches[0].name.clone()), - _ => { - if !ui::can_prompt() { - bail!( - "multiple auth logins for org '{identifier}': {}. Use --org to disambiguate.", - matches - .iter() - .map(|p| p.name.as_str()) - .collect::>() - .join(", ") - ); - } - let names: Vec<&str> = matches.iter().map(|p| p.name.as_str()).collect(); - let idx = crate::ui::fuzzy_select( - &format!("Multiple profiles for '{identifier}'. Select one"), - &names, - 0, - )?; - Ok(matches[idx].name.clone()) - } - } -} - -fn profile_info_identity_label(profile: &ProfileInfo) -> Option { - if let Some(email) = profile.email.as_deref() { - return match profile.user_name.as_deref() { - Some(name) => Some(format!("{name} ({email})")), - None => Some(email.to_string()), - }; - } - profile.api_key_hint.clone() -} - -pub(crate) fn profile_info_label(profile: &ProfileInfo) -> String { - let mut parts = vec![profile - .org_name - .clone() - .unwrap_or_else(|| "cross-org".to_string())]; - parts.push(profile.auth_method.clone()); - if let Some(identity) = profile_info_identity_label(profile) { - parts.push(identity); - } - parts.join(" — ") -} - -pub fn select_profile_interactive(current: Option<&str>) -> Result> { - let profiles = list_profiles()?; - if profiles.is_empty() { - bail!("no auth logins found. Run `bt auth login` to create one."); - } - if profiles.len() == 1 { - return Ok(Some(profiles[0].name.clone())); - } - - let labels: Vec = profiles.iter().map(profile_info_label).collect(); - - let default = current - .and_then(|c| { - profiles - .iter() - .position(|p| p.name == c || p.org_name.as_deref() == Some(c)) - }) - .unwrap_or(0); - let idx = crate::ui::fuzzy_select("Select auth login", &labels, default)?; - Ok(Some(profiles[idx].name.clone())) -} - pub async fn list_available_orgs(base: &BaseArgs) -> Result> { let resolved = resolve_auth(base).await?; let app_url = resolved @@ -842,11 +726,6 @@ fn resolve_env_api_key(base: &BaseArgs) -> Option { Some(value.to_string()) } -#[cfg(test)] -fn resolve_api_key_override(base: &BaseArgs) -> Option { - resolve_cli_api_key_override(base) -} - fn config_auth_context(base: &BaseArgs) -> Option { let cfg = crate::config::load().unwrap_or_default(); config_auth_context_from_config(base, &cfg) @@ -1185,32 +1064,6 @@ fn auth_slot_label(name: &str, profile: &AuthProfile) -> String { parts.join(" — ") } -#[cfg(test)] -fn resolve_profile_for_org<'a>(org: &str, store: &'a AuthStore) -> Option<&'a str> { - let matches: Vec<&str> = store - .profiles - .iter() - .filter(|(name, p)| name.as_str() == org || profile_matches_org_identifier(p, org)) - .map(|(name, _)| name.as_str()) - .collect(); - - match matches.len() { - 0 => None, - 1 => Some(matches[0]), - _ => None, - } -} - -#[cfg(test)] -fn profile_names_for_org<'a>(org: &str, store: &'a AuthStore) -> Vec<&'a str> { - store - .profiles - .iter() - .filter(|(_, profile)| profile_matches_org_identifier(profile, org)) - .map(|(name, _)| name.as_str()) - .collect() -} - fn is_cross_org_oauth_profile(profile: &AuthProfile) -> bool { profile.auth_kind == AuthKind::Oauth && profile @@ -1241,9 +1094,8 @@ fn auth_profile_names_by_kind<'a>( .collect() } -fn profile_info_from_store_entry(name: &str, profile: &AuthProfile) -> ProfileInfo { +fn profile_info_from_store_entry(profile: &AuthProfile) -> ProfileInfo { ProfileInfo { - name: name.to_string(), auth_method: auth_kind_label(profile.auth_kind).to_string(), org_name: profile.org_name.clone(), user_name: profile.user_name.clone(), @@ -1253,15 +1105,11 @@ fn profile_info_from_store_entry(name: &str, profile: &AuthProfile) -> ProfileIn } fn profile_info_for_candidate(store: &AuthStore, name: &str) -> Option { - store - .profiles - .get(name) - .map(|profile| profile_info_from_store_entry(name, profile)) + store.profiles.get(name).map(profile_info_from_store_entry) } fn ad_hoc_api_key_profile(org: Option<&str>, api_key: &str) -> ProfileInfo { ProfileInfo { - name: String::new(), auth_method: auth_kind_label(AuthKind::ApiKey).to_string(), org_name: org.map(str::to_string), user_name: None, @@ -1458,148 +1306,6 @@ fn select_auth_profile_candidate( } } -#[cfg(test)] -fn maybe_select_profile_for_auth( - base: &BaseArgs, - store: &AuthStore, - cfg_org: &Option, - can_prompt: bool, -) -> Result> { - if resolve_api_key_override(base).is_some() { - return Ok(None); - } - - let requested_profile = base - .profile - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()); - if requested_profile.is_some() { - return Ok(None); - } - - let effective_org = base.org_name.as_deref().or(cfg_org.as_deref()); - if let Some(org) = effective_org { - if resolve_profile_for_org(org, store).is_some() { - return Ok(None); - } - - let matching_profiles = profile_names_for_org(org, store); - if matching_profiles.is_empty() { - return Ok(None); - } - - if !can_prompt { - bail!( - "multiple auth logins for org '{org}': {}. Use --org to disambiguate.", - matching_profiles.join(", ") - ); - } - - return select_profile_from_store( - &format!("Multiple profiles for '{org}'. Select one"), - &matching_profiles, - Some(org), - store, - ) - .map(Some); - } - - if store.profiles.len() <= 1 { - return Ok(None); - } - - let names: Vec<&str> = store.profiles.keys().map(|name| name.as_str()).collect(); - if !can_prompt { - bail!( - "multiple auth logins available: {}. Pass --org to disambiguate.", - names.join(", ") - ); - } - - select_profile_from_store("Select org", &names, None, store).map(Some) -} - -#[cfg(test)] -fn resolve_auth_from_store_with_secret_lookup( - base: &BaseArgs, - store: &AuthStore, - load_secret: F, - cfg_org: &Option, -) -> Result -where - F: Fn(&str) -> Result>, -{ - if let Some(api_key) = resolve_api_key_override(base) { - return Ok(ResolvedAuth { - api_key: Some(api_key), - api_url: base.api_url.clone(), - app_url: base.app_url.clone(), - org_name: base.org_name.clone().or_else(|| cfg_org.clone()), - is_oauth: false, - }); - } - - let requested_profile = base - .profile - .as_ref() - .map(|s| s.trim()) - .filter(|s| !s.is_empty()); - - let effective_org = base.org_name.as_deref().or(cfg_org.as_deref()); - - let selected_profile_name = if let Some(profile) = requested_profile { - Some(profile) - } else if let Some(org) = effective_org { - resolve_profile_for_org(org, store) - } else if store.profiles.len() == 1 { - store.profiles.keys().next().map(|k| k.as_str()) - } else { - None - }; - - if let Some(profile_name) = selected_profile_name { - let profile = store.profiles.get(profile_name).ok_or_else(|| { - anyhow::anyhow!( - "auth login '{profile_name}' not found; run `bt auth profiles` or `bt auth login --org `" - ) - })?; - let is_oauth = profile.auth_kind == AuthKind::Oauth; - let api_key = if is_oauth { - None - } else { - Some(load_secret(profile_name)?.ok_or_else(|| { - recoverable_auth_error( - RecoverableAuthErrorKind::StoredCredential, - format!( - "no keychain credential found for auth login '{profile_name}'; re-run `bt auth login --org `" - ), - ) - })?) - }; - - return Ok(ResolvedAuth { - api_key, - api_url: base.api_url.clone().or_else(|| profile.api_url.clone()), - app_url: base.app_url.clone().or_else(|| profile.app_url.clone()), - org_name: base - .org_name - .clone() - .or_else(|| cfg_org.clone()) - .or_else(|| profile.org_name.clone()), - is_oauth, - }); - } - - Ok(ResolvedAuth { - api_key: None, - api_url: base.api_url.clone(), - app_url: base.app_url.clone(), - org_name: base.org_name.clone().or_else(|| cfg_org.clone()), - is_oauth: false, - }) -} - async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { if args.oauth { return run_login_oauth(base, args).await; @@ -2016,301 +1722,67 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { ) } -#[cfg(test)] -fn resolve_selected_profile_name_for_debug( - base: &BaseArgs, - store: &AuthStore, -) -> Result<(String, &'static str)> { - if let Some(profile_name) = base.profile.as_deref() { - let profile_name = profile_name.trim(); - if !profile_name.is_empty() { - return Ok((profile_name.to_string(), "legacy profile")); - } - } - - if let Some(org) = base.org_name.as_deref() { - if let Some(profile_name) = resolve_profile_for_org(org, store) { - return Ok((profile_name.to_string(), "org-based resolution")); - } - } - - if store.profiles.len() == 1 { - let name = store.profiles.keys().next().unwrap().clone(); - return Ok((name, "only profile")); - } - - if store.profiles.len() > 1 && ui::can_prompt() { - if let Some(name) = select_profile_interactive(None)? { - return Ok((name, "interactive selection")); - } - } - - bail!("no auth login selected; pass --org or rerun interactively") -} - -#[cfg(test)] -fn resolve_profile_name( - explicit_profile: Option<&str>, - suggested_org_name: Option<&str>, -) -> Result { - if let Some(profile) = explicit_profile { - let profile = profile.trim(); - if profile.is_empty() { - bail!("profile name cannot be empty"); - } - return Ok(profile.to_string()); +fn format_login_success(selected_org: Option<&LoginOrgInfo>, api_url: &str) -> String { + match selected_org { + Some(org) => format!("Logged in as {} (api: {api_url})", org.name), + None => format!("Logged in (cross-org, api: {api_url})"), } - - Ok(suggested_org_name - .map(str::trim) - .filter(|name| !name.is_empty()) - .unwrap_or("profile") - .to_string()) } -#[cfg(test)] -fn default_login_org_name( - store: &AuthStore, - profile_name: Option<&str>, - requested_org_name: Option<&str>, -) -> Option { - if requested_org_name - .map(str::trim) - .is_some_and(|name| !name.is_empty()) - { - return None; +fn build_login_context_for_selected_org( + credential: &str, + api_url: &str, + app_url: &str, + selected_org: Option<&LoginOrgInfo>, +) -> LoginContext { + let login = LoginState::new(); + let _ = login.set( + credential.to_string(), + selected_org.map(|org| org.id.clone()).unwrap_or_default(), + selected_org.map(|org| org.name.clone()).unwrap_or_default(), + api_url.to_string(), + app_url.to_string(), + ); + LoginContext { + login, + api_url: api_url.to_string(), + app_url: app_url.to_string(), } - - let profile_name = profile_name - .map(str::trim) - .filter(|name| !name.is_empty())?; - let stored_org_name = store - .profiles - .get(profile_name) - .and_then(|profile| profile.org_name.as_deref()) - .map(str::trim) - .filter(|org_name| !org_name.is_empty()); - - Some(stored_org_name.unwrap_or(profile_name).to_string()) -} - -#[cfg(test)] -fn default_profile_name(suggested_org_name: Option<&str>) -> String { - suggested_org_name - .map(str::trim) - .filter(|name| !name.is_empty()) - .unwrap_or("profile") - .to_string() } -#[cfg(test)] -fn next_available_profile_name(base_name: &str, store: &AuthStore) -> String { - if !store.profiles.contains_key(base_name) { - return base_name.to_string(); +fn format_post_login_context( + selected_org: Option<&LoginOrgInfo>, + project: Option<&api::Project>, +) -> String { + match (selected_org, project) { + (Some(org), Some(project)) => format!("{}/{}", org.name, project.name), + (Some(org), None) => org.name.clone(), + (None, _) => "cross-org mode".to_string(), } - - (2u32..) - .map(|idx| format!("{base_name}-{idx}")) - .find(|candidate| !store.profiles.contains_key(candidate)) - .expect("profile name sequence is infinite") } -#[cfg(test)] -fn resolve_api_key_login_profile_name( - explicit_profile: Option<&str>, - suggested_org_name: Option<&str>, - selected_api_url: &str, - store: &AuthStore, -) -> Result<(String, bool)> { - if let Some(profile_name) = explicit_profile { - let profile_name = resolve_profile_name(Some(profile_name), suggested_org_name)?; - let should_confirm_overwrite = store.profiles.get(&profile_name).is_some_and(|profile| { - !profile_matches_api_key_login_target(profile, selected_api_url, suggested_org_name) - }); - return Ok((profile_name.clone(), should_confirm_overwrite)); - } - - let default_name = default_profile_name(suggested_org_name); - let has_matching_api_key_profile = store.profiles.values().any(|profile| { - profile.auth_kind == AuthKind::ApiKey - && profile.api_url.as_deref() == Some(selected_api_url) - && profile.org_name.as_deref() == suggested_org_name - }); - - if has_matching_api_key_profile { - return Ok((next_available_profile_name(&default_name, store), false)); - } +async fn resolve_post_login_project( + base: &BaseArgs, + credential: &str, + api_url: &str, + app_url: &str, + selected_org: Option<&LoginOrgInfo>, +) -> Result> { + let Some(project_name) = config::trimmed_option(base.project.as_deref()) else { + return Ok(None); + }; - Ok(( - default_name.clone(), - store.profiles.contains_key(&default_name), - )) -} - -#[cfg(test)] -fn resolve_oauth_login_profile_name( - explicit_profile: Option<&str>, - suggested_org_name: Option<&str>, - selected_api_url: &str, - app_url: &str, - jwt_id: &JwtIdentity, - store: &AuthStore, -) -> Result<(String, bool)> { - if let Some(profile_name) = explicit_profile { - let profile_name = resolve_profile_name(Some(profile_name), suggested_org_name)?; - let should_confirm_overwrite = store.profiles.get(&profile_name).is_some_and(|profile| { - !profile_matches_oauth_login_target( - profile, - selected_api_url, - app_url, - suggested_org_name, - jwt_id, - ) - }); - return Ok((profile_name.clone(), should_confirm_overwrite)); - } - - let matched_profile = store - .profiles - .iter() - .filter(|(_, profile)| { - profile_matches_oauth_login_target( - profile, - selected_api_url, - app_url, - suggested_org_name, - jwt_id, - ) - }) - .max_by(|(left_name, left), (right_name, right)| { - left.oauth_access_expires_at - .unwrap_or_default() - .cmp(&right.oauth_access_expires_at.unwrap_or_default()) - .then_with(|| left_name.cmp(right_name)) - }) - .map(|(name, _)| name.clone()); - - if let Some(profile_name) = matched_profile { - return Ok((profile_name, false)); - } - - let default_name = default_profile_name(suggested_org_name); - Ok(( - default_name.clone(), - store.profiles.contains_key(&default_name), - )) -} - -#[cfg(test)] -fn profile_matches_api_key_login_target( - profile: &AuthProfile, - selected_api_url: &str, - suggested_org_name: Option<&str>, -) -> bool { - profile.auth_kind == AuthKind::ApiKey - && profile.api_url.as_deref() == Some(selected_api_url) - && profile.org_name.as_deref() == suggested_org_name -} - -#[cfg(test)] -fn profile_matches_oauth_login_target( - profile: &AuthProfile, - selected_api_url: &str, - app_url: &str, - suggested_org_name: Option<&str>, - jwt_id: &JwtIdentity, -) -> bool { - profile.auth_kind == AuthKind::Oauth - && profile.api_url.as_deref() == Some(selected_api_url) - && profile.app_url.as_deref() == Some(app_url) - && profile.org_name.as_deref() == suggested_org_name - && profile.user_name == jwt_id.name - && profile.email == jwt_id.email -} - -#[cfg(test)] -#[allow(dead_code)] -fn confirm_profile_overwrite(profile_name: &str) -> Result<()> { - let store = load_auth_store()?; - if !store.profiles.contains_key(profile_name) { - return Ok(()); - } - let Some(term) = ui::prompt_term() else { - return Ok(()); - }; - let confirmed = Confirm::new() - .with_prompt(format!( - "Profile '{profile_name}' already exists. Overwrite?" - )) - .default(false) - .interact_on(&term)?; - if !confirmed { - bail!("login cancelled"); - } - Ok(()) -} - -fn format_login_success(selected_org: Option<&LoginOrgInfo>, api_url: &str) -> String { - match selected_org { - Some(org) => format!("Logged in as {} (api: {api_url})", org.name), - None => format!("Logged in (cross-org, api: {api_url})"), - } -} - -fn build_login_context_for_selected_org( - credential: &str, - api_url: &str, - app_url: &str, - selected_org: Option<&LoginOrgInfo>, -) -> LoginContext { - let login = LoginState::new(); - let _ = login.set( - credential.to_string(), - selected_org.map(|org| org.id.clone()).unwrap_or_default(), - selected_org.map(|org| org.name.clone()).unwrap_or_default(), - api_url.to_string(), - app_url.to_string(), - ); - LoginContext { - login, - api_url: api_url.to_string(), - app_url: app_url.to_string(), - } -} - -fn format_post_login_context( - selected_org: Option<&LoginOrgInfo>, - project: Option<&api::Project>, -) -> String { - match (selected_org, project) { - (Some(org), Some(project)) => format!("{}/{}", org.name, project.name), - (Some(org), None) => org.name.clone(), - (None, _) => "cross-org mode".to_string(), - } -} - -async fn resolve_post_login_project( - base: &BaseArgs, - credential: &str, - api_url: &str, - app_url: &str, - selected_org: Option<&LoginOrgInfo>, -) -> Result> { - let Some(project_name) = config::trimmed_option(base.project.as_deref()) else { - return Ok(None); - }; - - let selected_org = selected_org.ok_or_else(|| { - anyhow::anyhow!( - "cannot set a default project in cross-org mode; rerun `bt auth login --org --project `" - ) - })?; - let ctx = - build_login_context_for_selected_org(credential, api_url, app_url, Some(selected_org)); - let client = ApiClient::new(&ctx)?; - switch::validate_or_create_project(&client, project_name) - .await - .map(Some) + let selected_org = selected_org.ok_or_else(|| { + anyhow::anyhow!( + "cannot set a default project in cross-org mode; rerun `bt auth login --org --project `" + ) + })?; + let ctx = + build_login_context_for_selected_org(credential, api_url, app_url, Some(selected_org)); + let client = ApiClient::new(&ctx)?; + switch::validate_or_create_project(&client, project_name) + .await + .map(Some) } async fn persist_post_login_context( @@ -2331,7 +1803,6 @@ async fn persist_post_login_context( let mut cfg = config::load_file(&path); switch::apply_switch_config( &mut cfg, - None, selected_org.map(|org| org.name.as_str()), project.as_ref(), ); @@ -2700,8 +2171,8 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { if json { let output: Vec = store .profiles - .iter() - .map(|(_name, p)| { + .values() + .map(|p| { serde_json::json!({ "auth": auth_kind_label(p.auth_kind), "org": p.org_name, @@ -4283,9 +3754,10 @@ mod tests { quiet_source: None, no_color: false, no_input: false, - profile: None, project: None, + project_source: None, org_name: None, + org_name_source: None, api_key: None, api_key_source: None, prefer_api_key: false, @@ -4296,9 +3768,8 @@ mod tests { } } - fn auth_config(profile: Option<&str>, org: Option<&str>) -> crate::config::Config { + fn auth_config(org: Option<&str>) -> crate::config::Config { crate::config::Config { - profile: profile.map(str::to_string), org: org.map(str::to_string), ..Default::default() } @@ -4936,233 +4407,16 @@ mod tests { assert_eq!(profile.legacy_secret_key.as_deref(), Some("new")); } - #[test] - fn resolve_auth_uses_profile_when_no_api_key_override() { - let mut base = make_base(); - base.profile = Some("work".to_string()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".to_string(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), - app_url: Some("https://www.example.com".to_string()), - org_name: Some("Example Org".to_string()), - oauth_client_id: None, - oauth_access_expires_at: None, - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".to_string())), - &None, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("profile-key")); - assert_eq!(resolved.api_url.as_deref(), Some("https://api.example.com")); - assert_eq!(resolved.org_name.as_deref(), Some("Example Org")); - assert!(!resolved.is_oauth); - } - #[test] fn config_auth_context_returns_config_org() { let base = make_base(); - let cfg = auth_config(Some("default-profile"), Some("local-org")); + let cfg = auth_config(Some("local-org")); let org = config_auth_context_from_config(&base, &cfg); assert_eq!(org.as_deref(), Some("local-org")); } - #[test] - fn config_auth_context_ignores_legacy_profile_and_preserves_config_org() { - let mut base = make_base(); - base.profile = Some("explicit-profile".to_string()); - let cfg = auth_config(Some("config-profile"), Some("local-org")); - - let org = config_auth_context_from_config(&base, &cfg); - - assert_eq!(org.as_deref(), Some("local-org")); - } - - #[test] - fn resolve_auth_prefers_explicit_api_key() { - let mut base = make_base(); - base.api_key = Some("explicit-key".to_string()); - base.api_url = Some("https://override.example.com".to_string()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".to_string(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), - app_url: None, - org_name: None, - oauth_client_id: None, - oauth_access_expires_at: None, - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".to_string())), - &None, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("explicit-key")); - assert_eq!( - resolved.api_url.as_deref(), - Some("https://override.example.com") - ); - assert!(!resolved.is_oauth); - } - - #[test] - fn resolve_auth_prefer_api_key_uses_api_key_override() { - let mut base = make_base(); - base.api_key = Some("explicit-key".to_string()); - base.prefer_api_key = true; - base.profile = Some("work".to_string()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".to_string(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), - app_url: None, - org_name: Some("Example Org".to_string()), - oauth_client_id: None, - oauth_access_expires_at: None, - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".to_string())), - &None, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("explicit-key")); - assert_eq!(resolved.org_name, None); - } - - #[test] - fn resolve_auth_prefers_cli_api_key_even_with_prefer_api_key() { - let mut base = make_base(); - base.api_key = Some("explicit-key".to_string()); - base.api_key_source = Some(crate::args::ArgValueSource::CommandLine); - base.prefer_api_key = true; - base.profile = Some("work".to_string()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".to_string(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), - app_url: None, - org_name: Some("Example Org".to_string()), - oauth_client_id: None, - oauth_access_expires_at: None, - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".to_string())), - &None, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("explicit-key")); - assert_eq!(resolved.org_name, None); - } - - #[test] - fn resolve_auth_legacy_explicit_profile_does_not_override_api_key() { - let mut base = make_base(); - base.api_key = Some("explicit-key".to_string()); - base.profile = Some("work".to_string()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".to_string(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), - app_url: None, - org_name: Some("Example Org".to_string()), - oauth_client_id: None, - oauth_access_expires_at: None, - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".to_string())), - &None, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("explicit-key")); - assert_eq!(resolved.org_name, None); - } - - #[test] - fn resolve_auth_marks_oauth_profiles() { - let mut base = make_base(); - base.profile = Some("work".to_string()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".to_string(), - AuthProfile { - auth_kind: AuthKind::Oauth, - api_url: Some("https://api.example.com".to_string()), - app_url: Some("https://www.example.com".to_string()), - org_name: Some("Example Org".to_string()), - oauth_client_id: Some("bt_cli_work".to_string()), - oauth_access_expires_at: None, - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("should-not-be-used".to_string())), - &None, - ) - .expect("resolve"); - - assert!(resolved.is_oauth); - assert_eq!(resolved.api_key, None); - assert_eq!(resolved.org_name.as_deref(), Some("Example Org")); - } - - #[test] - fn refresh_profile_selector_prefers_explicit_profile() { - let mut base = make_base(); - base.profile = Some(" work ".to_string()); - let store = AuthStore::default(); - let (profile_name, source) = - resolve_selected_profile_name_for_debug(&base, &store).expect("resolve"); - assert_eq!(profile_name, "work"); - assert_eq!(source, "legacy profile"); - } - #[test] fn parse_oauth_callback_input_accepts_json_payload() { let parsed = @@ -5191,388 +4445,6 @@ mod tests { ); } - #[test] - fn resolve_profile_for_org_exact_profile_name() { - let mut store = AuthStore::default(); - store.profiles.insert( - "acme".into(), - AuthProfile { - org_name: Some("acme-corp".into()), - ..Default::default() - }, - ); - assert_eq!(resolve_profile_for_org("acme", &store), Some("acme")); - } - - #[test] - fn resolve_profile_for_org_by_org_name() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - org_name: Some("acme-corp".into()), - ..Default::default() - }, - ); - assert_eq!(resolve_profile_for_org("acme-corp", &store), Some("work")); - } - - #[test] - fn resolve_profile_for_org_no_match() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - org_name: Some("acme-corp".into()), - ..Default::default() - }, - ); - assert_eq!(resolve_profile_for_org("unknown", &store), None); - } - - #[test] - fn resolve_profile_for_org_multiple_returns_none() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work-1".into(), - AuthProfile { - org_name: Some("acme".into()), - ..Default::default() - }, - ); - store.profiles.insert( - "work-2".into(), - AuthProfile { - org_name: Some("acme".into()), - ..Default::default() - }, - ); - assert_eq!(resolve_profile_for_org("acme", &store), None); - } - - #[test] - fn profile_selection_requires_choice_when_multiple_profiles_without_prompt() { - let base = make_base(); - let mut store = AuthStore::default(); - store.profiles.insert( - "alpha".into(), - AuthProfile { - org_name: Some("alpha-org".into()), - ..Default::default() - }, - ); - store.profiles.insert( - "beta".into(), - AuthProfile { - org_name: Some("beta-org".into()), - ..Default::default() - }, - ); - - let err = maybe_select_profile_for_auth(&base, &store, &None, false) - .expect_err("selection should be required"); - - assert!(err.to_string().contains("multiple auth logins available")); - assert!(err.to_string().contains("alpha")); - assert!(err.to_string().contains("beta")); - assert!(err.to_string().contains("--org ")); - } - - #[test] - fn profile_selection_requires_choice_for_ambiguous_org_without_prompt() { - let mut base = make_base(); - base.org_name = Some("acme".into()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work-1".into(), - AuthProfile { - org_name: Some("acme".into()), - ..Default::default() - }, - ); - store.profiles.insert( - "work-2".into(), - AuthProfile { - org_name: Some("acme".into()), - ..Default::default() - }, - ); - - let err = maybe_select_profile_for_auth(&base, &store, &None, false) - .expect_err("org selection should be required"); - - assert!(err - .to_string() - .contains("multiple auth logins for org 'acme'")); - assert!(err.to_string().contains("work-1")); - assert!(err.to_string().contains("work-2")); - } - - #[test] - fn profile_selection_skips_when_api_key_override_is_active() { - let mut base = make_base(); - base.api_key = Some("explicit-key".into()); - - let mut store = AuthStore::default(); - store - .profiles - .insert("alpha".into(), AuthProfile::default()); - store.profiles.insert("beta".into(), AuthProfile::default()); - - let selection = maybe_select_profile_for_auth(&base, &store, &None, false) - .expect("api key override should skip profile selection"); - - assert_eq!(selection, None); - } - - #[test] - fn resolve_auth_uses_org_to_find_profile() { - let mut base = make_base(); - base.org_name = Some("acme-corp".into()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - org_name: Some("acme-corp".into()), - api_url: Some("https://api.acme.com".into()), - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".into())), - &None, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("profile-key")); - assert_eq!(resolved.org_name.as_deref(), Some("acme-corp")); - } - - #[test] - fn resolve_auth_uses_config_org_to_find_profile() { - let base = make_base(); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - org_name: Some("acme-corp".into()), - api_url: Some("https://api.acme.com".into()), - ..Default::default() - }, - ); - let cfg_org = Some("acme-corp".to_string()); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".into())), - &cfg_org, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("profile-key")); - assert_eq!(resolved.org_name.as_deref(), Some("acme-corp")); - } - - #[test] - fn resolve_auth_config_org_overrides_profile_org() { - let mut base = make_base(); - base.profile = Some("default-profile".to_string()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "default-profile".into(), - AuthProfile { - org_name: Some("profile-org".into()), - ..Default::default() - }, - ); - let cfg_org = Some("local-org".to_string()); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("profile-key".into())), - &cfg_org, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("profile-key")); - assert_eq!(resolved.org_name.as_deref(), Some("local-org")); - } - - #[test] - fn resolve_auth_api_key_override_keeps_config_org() { - let mut base = make_base(); - base.api_key = Some("explicit-key".into()); - - let store = AuthStore::default(); - let cfg_org = Some("local-org".to_string()); - - let resolved = - resolve_auth_from_store_with_secret_lookup(&base, &store, |_| Ok(None), &cfg_org) - .expect("resolve"); - - assert_eq!(resolved.api_key.as_deref(), Some("explicit-key")); - assert_eq!(resolved.org_name.as_deref(), Some("local-org")); - } - - #[test] - fn resolve_auth_explicit_profile_overrides_org_resolution() { - let mut base = make_base(); - base.profile = Some("other".into()); - base.org_name = Some("acme-corp".into()); - - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - org_name: Some("acme-corp".into()), - ..Default::default() - }, - ); - store.profiles.insert( - "other".into(), - AuthProfile { - org_name: Some("other-org".into()), - api_url: Some("https://api.other.com".into()), - ..Default::default() - }, - ); - - let resolved = resolve_auth_from_store_with_secret_lookup( - &base, - &store, - |_| Ok(Some("other-key".into())), - &None, - ) - .expect("resolve"); - assert_eq!(resolved.api_key.as_deref(), Some("other-key")); - assert_eq!(resolved.org_name.as_deref(), Some("acme-corp")); - } - - #[test] - fn resolve_api_key_login_profile_name_creates_new_profile_for_matching_org() { - let mut store = AuthStore::default(); - store.profiles.insert( - "acme".into(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.acme.example".into()), - org_name: Some("acme".into()), - ..Default::default() - }, - ); - - let (profile_name, should_confirm) = resolve_api_key_login_profile_name( - None, - Some("acme"), - "https://api.acme.example", - &store, - ) - .expect("resolve"); - - assert_eq!(profile_name, "acme-2"); - assert!(!should_confirm); - } - - #[test] - fn resolve_api_key_login_profile_name_updates_explicit_matching_profile_without_confirm() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.test.example".into()), - org_name: Some("test-org".into()), - ..Default::default() - }, - ); - - let (profile_name, should_confirm) = resolve_api_key_login_profile_name( - Some("work"), - Some("test-org"), - "https://api.test.example", - &store, - ) - .expect("resolve"); - - assert_eq!(profile_name, "work"); - assert!(!should_confirm); - } - - #[test] - fn resolve_api_key_login_profile_name_confirms_explicit_different_target() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.test.example".into()), - org_name: Some("test-org".into()), - ..Default::default() - }, - ); - - let (profile_name, should_confirm) = resolve_api_key_login_profile_name( - Some("work"), - Some("other-org"), - "https://api.test.example", - &store, - ) - .expect("resolve"); - - assert_eq!(profile_name, "work"); - assert!(should_confirm); - } - - #[test] - fn default_login_org_name_uses_profile_org_when_org_not_requested() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - org_name: Some("acme".into()), - ..Default::default() - }, - ); - - assert_eq!( - default_login_org_name(&store, Some(" work "), None).as_deref(), - Some("acme") - ); - } - - #[test] - fn default_login_org_name_falls_back_to_profile_name() { - let store = AuthStore::default(); - - assert_eq!( - default_login_org_name(&store, Some(" acme "), None).as_deref(), - Some("acme") - ); - } - - #[test] - fn default_login_org_name_ignores_profile_when_org_requested() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - org_name: Some("acme".into()), - ..Default::default() - }, - ); - - assert_eq!( - default_login_org_name(&store, Some("work"), Some("other")), - None - ); - } - #[test] fn move_default_login_org_first_moves_matching_org() { let mut orgs = vec![ @@ -5595,122 +4467,6 @@ mod tests { assert_eq!(orgs[1].name, "beta"); } - #[test] - fn resolve_oauth_login_profile_name_reuses_most_recent_matching_profile() { - let mut store = AuthStore::default(); - store.profiles.insert( - "older".into(), - AuthProfile { - auth_kind: AuthKind::Oauth, - api_url: Some("https://api.acme.example".into()), - app_url: Some("https://www.acme.example".into()), - org_name: Some("acme".into()), - oauth_access_expires_at: Some(100), - user_name: Some("Alice".into()), - email: Some("alice@example.com".into()), - ..Default::default() - }, - ); - store.profiles.insert( - "newer".into(), - AuthProfile { - auth_kind: AuthKind::Oauth, - api_url: Some("https://api.acme.example".into()), - app_url: Some("https://www.acme.example".into()), - org_name: Some("acme".into()), - oauth_access_expires_at: Some(200), - user_name: Some("Alice".into()), - email: Some("alice@example.com".into()), - ..Default::default() - }, - ); - - let jwt_id = JwtIdentity { - name: Some("Alice".into()), - email: Some("alice@example.com".into()), - }; - let (profile_name, should_confirm) = resolve_oauth_login_profile_name( - None, - Some("acme"), - "https://api.acme.example", - "https://www.acme.example", - &jwt_id, - &store, - ) - .expect("resolve"); - - assert_eq!(profile_name, "newer"); - assert!(!should_confirm); - } - - #[test] - fn resolve_oauth_login_profile_name_updates_explicit_matching_profile_without_confirm() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - auth_kind: AuthKind::Oauth, - api_url: Some("https://api.test.example".into()), - app_url: Some("https://app.test.example".into()), - org_name: Some("test-org".into()), - user_name: Some("Test User".into()), - email: Some("user@test.example".into()), - ..Default::default() - }, - ); - let jwt_id = JwtIdentity { - name: Some("Test User".into()), - email: Some("user@test.example".into()), - }; - - let (profile_name, should_confirm) = resolve_oauth_login_profile_name( - Some("work"), - Some("test-org"), - "https://api.test.example", - "https://app.test.example", - &jwt_id, - &store, - ) - .expect("resolve"); - - assert_eq!(profile_name, "work"); - assert!(!should_confirm); - } - - #[test] - fn resolve_oauth_login_profile_name_confirms_explicit_different_target() { - let mut store = AuthStore::default(); - store.profiles.insert( - "work".into(), - AuthProfile { - auth_kind: AuthKind::Oauth, - api_url: Some("https://api.test.example".into()), - app_url: Some("https://app.test.example".into()), - org_name: Some("test-org".into()), - user_name: Some("Test User".into()), - email: Some("user@test.example".into()), - ..Default::default() - }, - ); - let jwt_id = JwtIdentity { - name: Some("Test User".into()), - email: Some("user@test.example".into()), - }; - - let (profile_name, should_confirm) = resolve_oauth_login_profile_name( - Some("work"), - Some("other-org"), - "https://api.test.example", - "https://app.test.example", - &jwt_id, - &store, - ) - .expect("resolve"); - - assert_eq!(profile_name, "work"); - assert!(should_confirm); - } - fn login_org(id: &str, name: &str) -> LoginOrgInfo { LoginOrgInfo { id: id.to_string(), @@ -5723,7 +4479,6 @@ mod tests { async fn persist_post_login_context_clears_stale_project_for_org_only_login() { let _env = TestEnv::new(None, None).await; crate::config::save_global(&crate::config::Config { - profile: Some("old-profile".to_string()), org: Some("old-org".to_string()), project: Some("stale-project".to_string()), project_id: Some("proj_stale".to_string()), @@ -5743,7 +4498,6 @@ mod tests { let cfg = crate::config::load_global().expect("load global config"); assert_eq!(update.display, "acme"); - assert_eq!(cfg.profile, None); assert_eq!(cfg.org.as_deref(), Some("acme")); assert_eq!(cfg.project, None); assert_eq!(cfg.project_id, None); diff --git a/src/config/mod.rs b/src/config/mod.rs index c6b51604..b6dc85fb 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -18,8 +18,6 @@ mod set; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] #[serde(default)] pub struct Config { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub profile: Option, pub org: Option, pub project: Option, pub project_id: Option, @@ -82,7 +80,6 @@ impl Config { self.project_id.clone() }; Config { - profile: None, org: other.org.clone().or_else(|| self.org.clone()), project, project_id, @@ -128,9 +125,7 @@ pub fn load_file(path: &Path) -> Config { } }; - // Legacy config files may contain `profile`; profiles are no longer a - // user-facing selector, so ignore it rather than warning or preserving it. - config.profile = None; + config.extra.remove("profile"); for key in config.extra.keys() { print_command_status( @@ -427,7 +422,7 @@ mod tests { assert_eq!(merged.project, Some("other-proj".into())); } - fn base_with_profile(profile: Option<&str>) -> BaseArgs { + fn base_args() -> BaseArgs { BaseArgs { json: false, verbose: false, @@ -436,9 +431,10 @@ mod tests { quiet_source: None, no_color: false, no_input: false, - profile: profile.map(str::to_string), org_name: None, + org_name_source: None, project: None, + project_source: None, api_key: None, api_key_source: None, prefer_api_key: false, @@ -449,9 +445,8 @@ mod tests { } } - fn config(profile: Option<&str>, org: Option<&str>, project: Option<&str>) -> Config { + fn config(org: Option<&str>, project: Option<&str>) -> Config { Config { - profile: profile.map(str::to_string), org: org.map(str::to_string), project: project.map(str::to_string), ..Default::default() @@ -459,20 +454,12 @@ mod tests { } #[test] - fn project_config_matches_org_and_ignores_legacy_profile() { - let base = base_with_profile(Some("work")); + fn project_config_matches_org() { + let base = base_args(); let cases = [ - (config(None, Some("acme"), Some("demo")), Some("demo")), - (config(None, Some("other"), Some("demo")), None), - (config(None, None, Some("demo")), Some("demo")), - ( - config(Some("other"), Some("acme"), Some("demo")), - Some("demo"), - ), - ( - config(Some("work"), Some("acme"), Some("demo")), - Some("demo"), - ), + (config(Some("acme"), Some("demo")), Some("demo")), + (config(Some("other"), Some("demo")), None), + (config(None, Some("demo")), Some("demo")), ]; for (cfg, expected) in cases { diff --git a/src/datasets/pipeline.rs b/src/datasets/pipeline.rs index ae8bbd0f..cd726927 100644 --- a/src/datasets/pipeline.rs +++ b/src/datasets/pipeline.rs @@ -2073,9 +2073,10 @@ mod tests { quiet_source: None, no_color: false, no_input: false, - profile: None, org_name: None, + org_name_source: None, project: None, + project_source: None, api_key: None, api_key_source: None, prefer_api_key: false, diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index a3eaabbb..341ec84d 100644 --- a/src/experiments/mod.rs +++ b/src/experiments/mod.rs @@ -219,17 +219,12 @@ fn apply_experiment_url_hints_to_base( } } - let has_profile_override = base - .profile - .as_deref() - .map(str::trim) - .is_some_and(|v| !v.is_empty()); let has_org_override = base .org_name .as_deref() .map(str::trim) .is_some_and(|v| !v.is_empty()); - if !has_profile_override && !has_org_override { + if !has_org_override { if let Some(org) = parsed_url .org .as_deref() diff --git a/src/functions/push.rs b/src/functions/push.rs index d3771ab2..b4d1f1fb 100644 --- a/src/functions/push.rs +++ b/src/functions/push.rs @@ -4063,9 +4063,10 @@ mod tests { quiet_source: None, no_color: false, no_input: false, - profile: None, org_name: None, + org_name_source: None, project: None, + project_source: None, api_key: None, api_key_source: None, prefer_api_key: false, diff --git a/src/main.rs b/src/main.rs index 9a7a5f12..009cc335 100644 --- a/src/main.rs +++ b/src/main.rs @@ -346,6 +346,8 @@ fn try_main() -> Result<()> { fn apply_base_arg_sources(matches: &ArgMatches, base: &mut BaseArgs) { base.verbose_source = find_value_source(matches, "verbose").and_then(map_value_source); base.quiet_source = find_value_source(matches, "quiet").and_then(map_value_source); + base.org_name_source = find_value_source(matches, "org_name").and_then(map_value_source); + base.project_source = find_value_source(matches, "project").and_then(map_value_source); base.api_key_source = find_value_source(matches, "api_key").and_then(map_value_source); } @@ -546,6 +548,32 @@ mod tests { ); } + #[test] + fn apply_base_arg_sources_tracks_cli_org_and_project() { + let matches = Cli::command() + .try_get_matches_from([ + "bt", + "status", + "--org", + "test-org", + "--project", + "test-project", + ]) + .expect("matches"); + let mut cli = Cli::from_arg_matches(&matches).expect("cli"); + + apply_base_arg_sources(&matches, cli.command.base_mut()); + + assert_eq!( + cli.command.base().org_name_source, + Some(ArgValueSource::CommandLine) + ); + assert_eq!( + cli.command.base().project_source, + Some(ArgValueSource::CommandLine) + ); + } + #[test] fn apply_base_arg_sources_leaves_api_key_source_empty_when_unset() { let _guard = env_test_lock().lock().expect("env test lock"); diff --git a/src/setup/mod.rs b/src/setup/mod.rs index 9d23a4f3..b98e79f4 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -1401,7 +1401,6 @@ async fn poll_setup_wizard_completion( async fn run_setup_browser_auth( base: &mut BaseArgs, - profile_name: Option<&str>, _project_name: Option<&str>, _project_was_explicit: bool, _requested_org: Option<&str>, @@ -1467,7 +1466,6 @@ async fn run_setup_browser_auth( org_id: completed.org_id.clone(), description: None, }; - let _ = profile_name; auth::commit_api_key_profile( &completed.api_key, login.api_url.clone(), @@ -1489,97 +1487,43 @@ async fn run_setup_browser_auth( }) } -#[cfg(test)] -fn setup_browser_profile_name( - selected_profile_name: Option<&str>, - org_name: &str, - profiles: &[auth::ProfileInfo], -) -> String { - if let Some(profile_name) = selected_profile_name +fn saved_auth_orgs(profiles: &[auth::ProfileInfo]) -> Vec { + let mut orgs = profiles + .iter() + .filter_map(|profile| profile.org_name.as_deref()) .map(str::trim) - .filter(|value| !value.is_empty()) - { - return profile_name.to_string(); - } - - let base_name = if org_name.trim().is_empty() { - "default" - } else { - org_name.trim() - }; - if !profiles.iter().any(|profile| profile.name == base_name) { - return base_name.to_string(); - } - - (2u32..) - .map(|idx| format!("{base_name}-{idx}")) - .find(|candidate| !profiles.iter().any(|profile| profile.name == *candidate)) - .expect("profile name sequence is infinite") + .filter(|org| !org.is_empty()) + .map(str::to_string) + .collect::>(); + orgs.sort(); + orgs.dedup(); + orgs } -fn resolve_profile_name_for_setup( +fn resolve_org_name_for_setup( base: &BaseArgs, profiles: &[auth::ProfileInfo], prompt_for_choice: bool, ) -> Result> { - if let Some(profile_name) = base - .profile + if let Some(org_name) = base + .org_name .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) { - if profiles.iter().any(|profile| profile.name == profile_name) { - return Ok(Some(profile_name.to_string())); - } - bail!( - "auth login '{profile_name}' not found; run `bt auth profiles` to see available logins" - ); + return Ok(Some(org_name.to_string())); } - if let Some(org_name) = base.org_name.as_deref() { - if let Some(profile_name) = profiles - .iter() - .find(|profile| profile.name == org_name) - .map(|profile| profile.name.clone()) - { - return Ok(Some(profile_name)); + let orgs = saved_auth_orgs(profiles); + match orgs.len() { + 0 => Ok(None), + 1 => Ok(orgs.into_iter().next()), + _ if prompt_for_choice => { + let labels = orgs.iter().map(String::as_str).collect::>(); + let idx = ui::fuzzy_select("Select organization", &labels, 0)?; + Ok(Some(orgs[idx].clone())) } - - let mut matches = profiles - .iter() - .filter(|profile| profile.org_name.as_deref() == Some(org_name)) - .collect::>(); - matches.sort_by(|left, right| left.name.cmp(&right.name)); - - return match matches.len() { - 0 => Ok(None), - 1 => Ok(Some(matches[0].name.clone())), - _ if prompt_for_choice => auth::select_profile_interactive(Some(org_name))? - .map(Some) - .ok_or_else(|| anyhow!("no auth login selected")), - _ => { - let labels = matches - .iter() - .map(|profile| auth::profile_info_label(profile)) - .collect::>() - .join(", "); - bail!( - "multiple auth logins for org '{org_name}': {labels}. Rerun interactively or remove one with `bt auth logout`." - ) - } - }; - } - - if profiles.len() == 1 { - return Ok(Some(profiles[0].name.clone())); - } - - if prompt_for_choice && !profiles.is_empty() { - auth::select_profile_interactive(None)? - .map(Some) - .ok_or_else(|| anyhow!("no auth login selected")) - } else { - Ok(None) + _ => Ok(None), } } @@ -1729,64 +1673,60 @@ fn select_api_key_org_for_setup( bail!("organization choice required in non-interactive mode; pass --org ") } -async fn ensure_profile_or_setup_browser_auth( +async fn ensure_org_or_setup_browser_auth( base: &mut BaseArgs, - prompt_for_profile_choice: bool, + prompt_for_org_choice: bool, project_name: Option<&str>, project_was_explicit: bool, requested_org: Option<&str>, ) -> Result { let profiles = auth::list_profiles()?; let can_prompt = setup_can_prompt(base); - let should_prompt_for_profile_choice = prompt_for_profile_choice && can_prompt; - let selected_profile = - resolve_profile_name_for_setup(base, &profiles, should_prompt_for_profile_choice)?; + let should_prompt_for_org_choice = prompt_for_org_choice && can_prompt; + let selected_org = resolve_org_name_for_setup(base, &profiles, should_prompt_for_org_choice)?; let mut auth_base = base.clone(); auth_base.api_key = None; auth_base.api_key_source = None; - if let Some(profile_name) = selected_profile { - auth_base.profile = Some(profile_name.clone()); - if auth_base.org_name.is_none() { - auth_base.org_name = profiles - .iter() - .find(|profile| profile.name == profile_name) - .and_then(|profile| profile.org_name.clone()); - } - - match auth::login(&auth_base).await { - Ok(ctx) => { - base.profile = auth_base.profile.clone(); - base.org_name = auth_base.org_name.clone(); - let is_oauth = auth::resolve_auth(&auth_base).await?.is_oauth; - return Ok(SetupAuthLogin { - login: ctx, - is_oauth, - selected_project: None, - }); - } - Err(err) if auth::is_missing_credential_error(&err) => { - if base.verbose { - eprintln!( - " Auth login '{}' credentials inaccessible ({}). Re-authenticating in the browser...", - profile_name, err - ); + if let Some(org_name) = selected_org.as_deref() { + auth_base.org_name = Some(org_name.to_string()); + let has_saved_auth_for_org = profiles + .iter() + .any(|profile| profile.org_name.as_deref() == Some(org_name)); + + if has_saved_auth_for_org { + match auth::login(&auth_base).await { + Ok(ctx) => { + base.org_name = auth_base.org_name.clone(); + let is_oauth = auth::resolve_auth(&auth_base).await?.is_oauth; + return Ok(SetupAuthLogin { + login: ctx, + is_oauth, + selected_project: None, + }); } - if !can_prompt { - bail!( - "setup needs interactive browser authentication; rerun without --no-input/--json or pass a working API key" - ); + Err(err) if auth::is_missing_credential_error(&err) => { + if base.verbose { + eprintln!( + " Auth login for org '{}' has inaccessible credentials ({}). Re-authenticating in the browser...", + org_name, err + ); + } + if !can_prompt { + bail!( + "setup needs interactive browser authentication; rerun without --no-input/--json or pass a working API key" + ); + } + return run_setup_browser_auth( + base, + project_name, + project_was_explicit, + requested_org, + ) + .await; } - return run_setup_browser_auth( - base, - Some(&profile_name), - project_name, - project_was_explicit, - requested_org, - ) - .await; + Err(err) => return Err(err), } - Err(err) => return Err(err), } } @@ -1802,27 +1742,20 @@ async fn ensure_profile_or_setup_browser_auth( if base.verbose { eprintln!("Starting browser setup.\n"); } - run_setup_browser_auth( - base, - None, - project_name, - project_was_explicit, - requested_org, - ) - .await + run_setup_browser_auth(base, project_name, project_was_explicit, requested_org).await } -async fn ensure_profile_or_setup_browser_auth_context( +async fn ensure_org_or_setup_browser_auth_context( base: &mut BaseArgs, - prompt_for_profile_choice: bool, + prompt_for_org_choice: bool, needs_api_key: bool, project_name: Option<&str>, project_was_explicit: bool, requested_org: Option<&str>, ) -> Result { - let login = ensure_profile_or_setup_browser_auth( + let login = ensure_org_or_setup_browser_auth( base, - prompt_for_profile_choice, + prompt_for_org_choice, project_name, project_was_explicit, requested_org, @@ -1841,7 +1774,7 @@ async fn ensure_profile_or_setup_browser_auth_context( async fn ensure_setup_auth( base: &mut BaseArgs, - prompt_for_profile_choice: bool, + prompt_for_org_choice: bool, needs_api_key: bool, ) -> Result { let project_was_explicit = base @@ -1857,7 +1790,6 @@ async fn ensure_setup_auth( .map(str::trim) .filter(|value| !value.is_empty()) .map(str::to_string); - let stored_profiles = auth::list_stored_profiles()?; let mut project_name = base .project .as_deref() @@ -1870,12 +1802,6 @@ async fn ensure_setup_auth( .map(str::trim) .filter(|value| !value.is_empty()) .map(str::to_string); - let profile_name = base - .profile - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string); if let Some(api_key) = explicit_api_key.as_deref() { let app_url = base @@ -1883,7 +1809,7 @@ async fn ensure_setup_auth( .clone() .unwrap_or_else(|| DEFAULT_APP_URL.to_string()); let available_orgs = list_available_orgs_for_setup(api_key, &app_url).await?; - if project_name.is_some() && org_name.is_none() && profile_name.is_none() { + if project_name.is_some() && org_name.is_none() { let name = project_name .as_deref() .expect("project name exists when probing all orgs"); @@ -1893,78 +1819,7 @@ async fn ensure_setup_auth( } } - if let Some(profile_name) = profile_name.as_deref() { - let profile = stored_profiles - .iter() - .find(|profile| profile.name == profile_name) - .ok_or_else(|| anyhow!("auth login '{profile_name}' not found"))?; - let target_org = match org_name.as_deref() { - Some(org_name) => { - if profile.org_name.as_deref() != Some(org_name) { - bail!( - "auth login '{profile_name}' belongs to org '{}' but '{}' was requested", - profile.org_name.as_deref().unwrap_or("(none)"), - org_name - ); - } - org_name - } - None => profile.org_name.as_deref().ok_or_else(|| { - anyhow!("auth login '{profile_name}' does not have a default org") - })?, - }; - - if let Some(org) = find_available_org(&available_orgs, target_org) { - let client = build_api_key_client(base, api_key, org).await?; - ensure_selected_setup_project( - base, - &client, - &mut project_name, - project_was_explicit, - &org.name, - ) - .await?; - return build_setup_auth_context(base, client, false, needs_api_key, None).await; - } - - let login = ensure_profile_or_setup_browser_auth( - base, - prompt_for_profile_choice, - project_name.as_deref(), - project_was_explicit, - org_name.as_deref(), - ) - .await?; - let selected_project = login.selected_project.clone(); - let client = ApiClient::new(&login.login)?; - let resolved_org = client.org_name().to_string(); - if selected_project.is_none() { - ensure_selected_setup_project( - base, - &client, - &mut project_name, - project_was_explicit, - &resolved_org, - ) - .await?; - } - return build_setup_auth_context( - base, - client, - login.is_oauth, - needs_api_key, - selected_project, - ) - .await; - } - if let Some(org_name) = org_name.as_deref() { - let matching_profile_count = stored_profiles - .iter() - .filter(|profile| profile.org_name.as_deref() == Some(org_name)) - .count(); - let _ = matching_profile_count; - if let Some(org) = find_available_org(&available_orgs, org_name) { let client = build_api_key_client(base, api_key, org).await?; ensure_selected_setup_project( @@ -1978,9 +1833,9 @@ async fn ensure_setup_auth( return build_setup_auth_context(base, client, false, needs_api_key, None).await; } - let login = ensure_profile_or_setup_browser_auth( + let login = ensure_org_or_setup_browser_auth( base, - prompt_for_profile_choice, + prompt_for_org_choice, project_name.as_deref(), project_was_explicit, Some(org_name), @@ -2016,9 +1871,9 @@ async fn ensure_setup_auth( None => available_orgs.clone(), }; if candidate_orgs.is_empty() { - return ensure_profile_or_setup_browser_auth_context( + return ensure_org_or_setup_browser_auth_context( base, - prompt_for_profile_choice, + prompt_for_org_choice, needs_api_key, project_name.as_deref(), project_was_explicit, @@ -2032,9 +1887,9 @@ async fn ensure_setup_auth( return build_setup_auth_context(base, client, false, needs_api_key, None).await; } - ensure_profile_or_setup_browser_auth_context( + ensure_org_or_setup_browser_auth_context( base, - prompt_for_profile_choice, + prompt_for_org_choice, needs_api_key, project_name.as_deref(), project_was_explicit, @@ -5340,9 +5195,10 @@ mod tests { quiet_source: None, no_color: false, no_input: false, - profile: None, org_name: None, + org_name_source: None, project: None, + project_source: None, api_key: None, api_key_source: None, prefer_api_key: false, @@ -5678,11 +5534,10 @@ mod tests { } #[test] - fn resolve_profile_name_for_setup_requires_prompt_when_multiple_profiles_exist() { + fn resolve_org_name_for_setup_requires_prompt_when_multiple_orgs_exist() { let base = make_base_args(); let profiles = vec![ auth::ProfileInfo { - name: "zeta".to_string(), auth_method: "api_key".to_string(), org_name: Some("Zeta Org".to_string()), user_name: None, @@ -5690,7 +5545,6 @@ mod tests { api_key_hint: None, }, auth::ProfileInfo { - name: "alpha".to_string(), auth_method: "api_key".to_string(), org_name: Some("Alpha Org".to_string()), user_name: None, @@ -5699,72 +5553,69 @@ mod tests { }, ]; - let resolved = - resolve_profile_name_for_setup(&base, &profiles, false).expect("resolve profile"); + let resolved = resolve_org_name_for_setup(&base, &profiles, false).expect("resolve org"); assert_eq!(resolved, None); } #[test] - fn resolve_profile_name_for_setup_allows_oauth_fallback_when_no_profiles_exist() { + fn resolve_org_name_for_setup_allows_browser_fallback_when_no_orgs_exist() { let base = make_base_args(); let profiles = Vec::new(); - let resolved = - resolve_profile_name_for_setup(&base, &profiles, true).expect("resolve profile"); + let resolved = resolve_org_name_for_setup(&base, &profiles, true).expect("resolve org"); assert_eq!(resolved, None); } #[test] - fn resolve_profile_name_for_setup_errors_for_unknown_explicit_profile() { + fn resolve_org_name_for_setup_uses_explicit_org_without_profile_lookup() { let mut base = make_base_args(); - base.profile = Some("missing".to_string()); + base.org_name = Some("Acme".to_string()); let profiles = vec![auth::ProfileInfo { - name: "work".to_string(), auth_method: "api_key".to_string(), - org_name: Some("Acme".to_string()), + org_name: Some("Other".to_string()), user_name: None, email: None, api_key_hint: None, }]; - let err = - resolve_profile_name_for_setup(&base, &profiles, false).expect_err("missing profile"); - assert!(err.to_string().contains("auth login 'missing' not found")); + let resolved = resolve_org_name_for_setup(&base, &profiles, false).expect("resolve org"); + assert_eq!(resolved.as_deref(), Some("Acme")); } #[test] - fn setup_browser_profile_name_reuses_selected_profile_or_suffixed_org_name() { + fn saved_auth_orgs_dedupes_and_sorts_profile_orgs() { let profiles = vec![ auth::ProfileInfo { - name: "Acme".to_string(), auth_method: "api_key".to_string(), - org_name: Some("Acme".to_string()), + org_name: Some("Zeta".to_string()), user_name: None, email: None, api_key_hint: None, }, auth::ProfileInfo { - name: "Acme-2".to_string(), + auth_method: "oauth".to_string(), + org_name: Some("Acme".to_string()), + user_name: Some("Test User".to_string()), + email: Some("user@example.test".to_string()), + api_key_hint: None, + }, + auth::ProfileInfo { auth_method: "api_key".to_string(), org_name: Some("Acme".to_string()), user_name: None, email: None, + api_key_hint: Some("sk-****abcde".to_string()), + }, + auth::ProfileInfo { + auth_method: "oauth".to_string(), + org_name: None, + user_name: Some("Test User".to_string()), + email: Some("user@example.test".to_string()), api_key_hint: None, }, ]; - assert_eq!( - setup_browser_profile_name(Some("work"), "Acme", &profiles), - "work" - ); - assert_eq!( - setup_browser_profile_name(None, "Acme", &profiles), - "Acme-3" - ); - assert_eq!( - setup_browser_profile_name(None, "New Org", &profiles), - "New Org" - ); + assert_eq!(saved_auth_orgs(&profiles), vec!["Acme", "Zeta"]); } #[test] diff --git a/src/status.rs b/src/status.rs index 6893975e..e6073031 100644 --- a/src/status.rs +++ b/src/status.rs @@ -57,11 +57,9 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { .map(|p| config::load_file(p)) .unwrap_or_default(); - let cli_org = cli_flag_value(&["--org", "-o"]); - let cli_project = cli_flag_value(&["--project", "-p"]); + let overrides = ConfigOverrides::from_base(&base); let (org, mut project, source) = resolve_config( - cli_org, - cli_project, + overrides, &global_cfg, &local_cfg, &local_path, @@ -138,21 +136,54 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { Ok(()) } -/// Precedence (clig.dev): CLI flag > env var > local config > global config. -pub(crate) fn resolve_config( +#[derive(Default)] +pub(crate) struct ConfigOverrides { cli_org: Option, + env_org: Option, cli_project: Option, + env_project: Option, +} + +impl ConfigOverrides { + fn from_base(base: &BaseArgs) -> Self { + use crate::args::ArgValueSource; + + let (cli_org, env_org) = match base.org_name_source { + Some(ArgValueSource::CommandLine) => (base.org_name.clone(), None), + Some(ArgValueSource::EnvVariable) => (None, base.org_name.clone()), + None => (None, None), + }; + let (cli_project, env_project) = match base.project_source { + Some(ArgValueSource::CommandLine) => (base.project.clone(), None), + Some(ArgValueSource::EnvVariable) => (None, base.project.clone()), + None => (None, None), + }; + + Self { + cli_org, + env_org, + cli_project, + env_project, + } + } +} + +/// Precedence (clig.dev): CLI flag > env var > local config > global config. +pub(crate) fn resolve_config( + overrides: ConfigOverrides, global: &config::Config, local: &config::Config, local_path: &Option, global_path: &Option, ) -> (Option, Option, Option) { - let env_org = std::env::var("BRAINTRUST_ORG_NAME") - .ok() - .filter(|s| !s.is_empty()); - let env_project = std::env::var("BRAINTRUST_DEFAULT_PROJECT") - .ok() - .filter(|s| !s.is_empty()); + let ConfigOverrides { + cli_org, + env_org, + cli_project, + env_project, + } = overrides; + let env_org = env_org.filter(|s| !s.is_empty()); + let env_project = env_project.filter(|s| !s.is_empty()); let org = cli_org .clone() @@ -181,32 +212,6 @@ pub(crate) fn resolve_config( (org, project, source) } -fn cli_flag_value(flags: &[&str]) -> Option { - let args: Vec = std::env::args().collect(); - for (i, arg) in args.iter().enumerate() { - if arg == "--" { - break; - } - for flag in flags { - if arg == *flag { - return args.get(i + 1).cloned(); - } - if let Some(val) = arg.strip_prefix(&format!("{flag}=")) { - return Some(val.to_string()); - } - // Handle -oVALUE style (short flags only) - if flag.len() == 2 && flag.starts_with('-') { - if let Some(val) = arg.strip_prefix(flag) { - if !val.is_empty() { - return Some(val.to_string()); - } - } - } - } - } - None -} - #[cfg(test)] mod tests { use super::*; @@ -232,8 +237,11 @@ mod tests { let global_path = Some(PathBuf::from("/home/.bt/config.json")); let (org, project, source) = resolve_config( - s("cli-org"), - s("cli-proj"), + ConfigOverrides { + cli_org: s("cli-org"), + cli_project: s("cli-proj"), + ..Default::default() + }, &global, &local, &local_path, @@ -245,6 +253,30 @@ mod tests { assert_eq!(source, s("cli")); } + #[test] + fn env_overrides_config_below_cli() { + let global = config(Some("global-org"), Some("global-proj")); + let local = config(Some("local-org"), Some("local-proj")); + let local_path = Some(PathBuf::from("/project/.bt/config.json")); + let global_path = Some(PathBuf::from("/home/.bt/config.json")); + + let (org, project, source) = resolve_config( + ConfigOverrides { + env_org: s("env-org"), + env_project: s("env-proj"), + ..Default::default() + }, + &global, + &local, + &local_path, + &global_path, + ); + + assert_eq!(org, s("env-org")); + assert_eq!(project, s("env-proj")); + assert_eq!(source, s("env")); + } + #[test] fn local_overrides_global() { let global = config(Some("global-org"), Some("global-proj")); @@ -252,8 +284,13 @@ mod tests { let local_path = Some(PathBuf::from("/project/.bt/config.json")); let global_path = Some(PathBuf::from("/home/.bt/config.json")); - let (org, project, source) = - resolve_config(None, None, &global, &local, &local_path, &global_path); + let (org, project, source) = resolve_config( + ConfigOverrides::default(), + &global, + &local, + &local_path, + &global_path, + ); assert_eq!(org, s("local-org")); assert_eq!(project, s("local-proj")); @@ -267,8 +304,13 @@ mod tests { let local_path = Some(PathBuf::from("/project/.bt/config.json")); let global_path = Some(PathBuf::from("/home/.bt/config.json")); - let (org, project, source) = - resolve_config(None, None, &global, &local, &local_path, &global_path); + let (org, project, source) = resolve_config( + ConfigOverrides::default(), + &global, + &local, + &local_path, + &global_path, + ); assert_eq!(org, s("global-org")); assert_eq!(project, s("global-proj")); @@ -282,8 +324,13 @@ mod tests { let local_path = Some(PathBuf::from("/project/.bt/config.json")); let global_path = Some(PathBuf::from("/home/.bt/config.json")); - let (org, project, source) = - resolve_config(None, None, &global, &local, &local_path, &global_path); + let (org, project, source) = resolve_config( + ConfigOverrides::default(), + &global, + &local, + &local_path, + &global_path, + ); assert_eq!(org, None); assert_eq!(project, None); @@ -298,8 +345,10 @@ mod tests { let global_path = Some(PathBuf::from("/home/.bt/config.json")); let (org, project, source) = resolve_config( - s("cli-org"), - None, + ConfigOverrides { + cli_org: s("cli-org"), + ..Default::default() + }, &global, &local, &local_path, @@ -318,8 +367,13 @@ mod tests { let local_path = Some(PathBuf::from("/project/.bt/config.json")); let global_path = Some(PathBuf::from("/home/.bt/config.json")); - let (org, project, source) = - resolve_config(None, None, &global, &local, &local_path, &global_path); + let (org, project, source) = resolve_config( + ConfigOverrides::default(), + &global, + &local, + &local_path, + &global_path, + ); assert_eq!(org, s("global-org")); assert_eq!(project, s("local-proj")); @@ -327,13 +381,11 @@ mod tests { } fn profile( - name: &str, user_name: Option<&str>, email: Option<&str>, api_key_hint: Option<&str>, ) -> auth::ProfileInfo { auth::ProfileInfo { - name: name.into(), auth_method: if api_key_hint.is_some() { "api_key" } else { @@ -349,7 +401,7 @@ mod tests { #[test] fn format_identity_name_and_email() { - let p = profile("work", Some("Alice"), Some("alice@example.com"), None); + let p = profile(Some("Alice"), Some("alice@example.com"), None); assert_eq!( format_identity(&p), Some("Alice (alice@example.com)".into()) @@ -358,19 +410,19 @@ mod tests { #[test] fn format_identity_email_only() { - let p = profile("work", None, Some("alice@example.com"), None); + let p = profile(None, Some("alice@example.com"), None); assert_eq!(format_identity(&p), Some("alice@example.com".into())); } #[test] fn format_identity_api_key_hint() { - let p = profile("work", None, None, Some("sk-****zhJwO")); + let p = profile(None, None, Some("sk-****zhJwO")); assert_eq!(format_identity(&p), Some("sk-****zhJwO".into())); } #[test] fn format_identity_none() { - let p = profile("work", None, None, None); + let p = profile(None, None, None); assert_eq!(format_identity(&p), None); } } diff --git a/src/switch.rs b/src/switch.rs index d079a212..021ccc16 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -103,7 +103,7 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { }; let mut cfg = config::load_file(&path); - apply_switch_config(&mut cfg, None, Some(&org_name), Some(&project)); + apply_switch_config(&mut cfg, Some(&org_name), Some(&project)); config::save_file(&path, &cfg) .context(format!("Could not save config to {}", path.display()))?; @@ -229,11 +229,9 @@ pub(crate) async fn validate_or_create_project( pub(crate) fn apply_switch_config( cfg: &mut config::Config, - _profile_name: Option<&str>, org_name: Option<&str>, project: Option<&api::Project>, ) { - cfg.profile = None; cfg.org = config::trimmed_option(org_name).map(str::to_string); match project { Some(project) => { @@ -247,37 +245,9 @@ pub(crate) fn apply_switch_config( } } -#[cfg(test)] -fn resolve_profile_for_switch( - has_api_key_override: bool, - prompting_for_project_only: bool, - is_interactive: bool, - select_profile_interactive: F, - interactive: &mut bool, -) -> Result> -where - F: FnOnce() -> Result>, -{ - if has_api_key_override { - if prompting_for_project_only && is_interactive { - *interactive = true; - } - return Ok(None); - } - - if prompting_for_project_only && is_interactive { - *interactive = true; - select_profile_interactive() - } else { - Ok(None) - } -} - #[cfg(test)] mod tests { use super::*; - use crate::auth::{resolve_org_to_profile, ProfileInfo}; - fn switch_args(target: Option<&str>) -> SwitchArgs { SwitchArgs { global: false, @@ -295,9 +265,10 @@ mod tests { quiet_source: None, no_color: false, no_input: false, - profile: None, org_name: org.map(String::from), + org_name_source: None, project: project.map(String::from), + project_source: None, api_key: None, api_key_source: None, prefer_api_key: false, @@ -308,17 +279,6 @@ mod tests { } } - fn profile_info(name: &str, org_name: Option<&str>) -> ProfileInfo { - ProfileInfo { - name: name.to_string(), - auth_method: "api_key".to_string(), - org_name: org_name.map(String::from), - user_name: None, - email: None, - api_key_hint: None, - } - } - // --- resolve_target tests (unchanged) --- #[test] @@ -413,99 +373,6 @@ mod tests { ); } - // --- resolve_org_to_profile tests --- - - #[test] - fn resolve_by_exact_profile_name() { - let profiles = vec![profile_info("acme", Some("acme-corp"))]; - assert_eq!(resolve_org_to_profile("acme", &profiles).unwrap(), "acme"); - } - - #[test] - fn resolve_by_org_name_when_profile_name_differs() { - let profiles = vec![profile_info("work", Some("acme-corp"))]; - assert_eq!( - resolve_org_to_profile("acme-corp", &profiles).unwrap(), - "work" - ); - } - - #[test] - fn resolve_no_match_errors() { - let profiles = vec![profile_info("work", Some("acme-corp"))]; - assert!(resolve_org_to_profile("unknown", &profiles).is_err()); - } - - #[test] - fn resolve_empty_profiles_errors() { - let profiles: Vec = vec![]; - let err = resolve_org_to_profile("anything", &profiles).unwrap_err(); - assert!(err.to_string().contains("no auth profiles found")); - } - - #[test] - fn resolve_prefers_profile_name_over_org_name() { - let profiles = vec![ - profile_info("acme", Some("other")), - profile_info("x", Some("acme")), - ]; - assert_eq!(resolve_org_to_profile("acme", &profiles).unwrap(), "acme"); - } - - #[test] - fn resolve_profile_without_org() { - let profiles = vec![profile_info("default", None)]; - assert_eq!( - resolve_org_to_profile("default", &profiles).unwrap(), - "default" - ); - } - - // --- login_base org_name clearing tests --- - - #[test] - fn login_base_clears_org_name_when_profile_resolved() { - let base = BaseArgs { - org_name: Some("staging".into()), - ..base_args(None, Some("foobar")) - }; - let profile_name = Some("staging".to_string()); - - let login_base = match &profile_name { - Some(profile) if base.profile.is_none() => BaseArgs { - profile: Some(profile.clone()), - org_name: None, - ..base.clone() - }, - _ => base.clone(), - }; - - assert_eq!(login_base.profile, Some("staging".into())); - assert_eq!(login_base.org_name, None); - } - - #[test] - fn login_base_preserves_org_when_explicit_profile_flag() { - let base = BaseArgs { - profile: Some("staging".into()), - org_name: Some("custom-org".into()), - ..base_args(None, Some("foobar")) - }; - let profile_name: Option = None; - - let login_base = match &profile_name { - Some(profile) if base.profile.is_none() => BaseArgs { - profile: Some(profile.clone()), - org_name: None, - ..base.clone() - }, - _ => base.clone(), - }; - - assert_eq!(login_base.profile, Some("staging".into())); - assert_eq!(login_base.org_name, Some("custom-org".into())); - } - #[test] fn apply_switch_config_sets_project_id_with_project_name_and_org() { let mut cfg = config::Config::default(); @@ -516,96 +383,26 @@ mod tests { description: None, }; - apply_switch_config(&mut cfg, Some("work"), Some("acme-org"), Some(&project)); + apply_switch_config(&mut cfg, Some("acme-org"), Some(&project)); - assert_eq!(cfg.profile, None); assert_eq!(cfg.org.as_deref(), Some("acme-org")); assert_eq!(cfg.project.as_deref(), Some("my-project")); assert_eq!(cfg.project_id.as_deref(), Some("proj_123")); } - #[test] - fn apply_switch_config_clears_existing_profile() { - let mut cfg = config::Config { - profile: Some("work".to_string()), - ..Default::default() - }; - let project = api::Project { - id: "proj_456".to_string(), - name: "next-project".to_string(), - org_id: "org_123".to_string(), - description: None, - }; - - apply_switch_config(&mut cfg, None, Some("acme-org"), Some(&project)); - - assert_eq!(cfg.profile, None); - assert_eq!(cfg.project.as_deref(), Some("next-project")); - } - #[test] fn apply_switch_config_clears_project_and_org_when_context_is_org_only() { let mut cfg = config::Config { - profile: Some("work".to_string()), org: Some("old-org".to_string()), project: Some("stale-project".to_string()), project_id: Some("proj_stale".to_string()), ..Default::default() }; - apply_switch_config(&mut cfg, Some("next"), None, None); + apply_switch_config(&mut cfg, None, None); - assert_eq!(cfg.profile, None); assert_eq!(cfg.org, None); assert_eq!(cfg.project, None); assert_eq!(cfg.project_id, None); } - - #[test] - fn resolve_profile_for_switch_skips_org_prompt_when_api_key_infers_profile() { - let mut interactive = false; - let profile = resolve_profile_for_switch( - true, - true, - true, - || panic!("org picker should not be called"), - &mut interactive, - ) - .expect("resolve"); - - assert_eq!(profile, None); - assert!(interactive); - } - - #[test] - fn resolve_profile_for_switch_prompts_when_no_inferred_profile() { - let mut interactive = false; - let profile = resolve_profile_for_switch( - false, - true, - true, - || Ok(Some("picked-profile".to_string())), - &mut interactive, - ) - .expect("resolve"); - - assert_eq!(profile.as_deref(), Some("picked-profile")); - assert!(interactive); - } - - #[test] - fn resolve_profile_for_switch_skips_org_prompt_when_api_key_override_has_no_profile_match() { - let mut interactive = false; - let profile = resolve_profile_for_switch( - true, - true, - true, - || panic!("org picker should not be called"), - &mut interactive, - ) - .expect("resolve"); - - assert_eq!(profile, None); - assert!(interactive); - } } diff --git a/src/traces.rs b/src/traces.rs index cc59c4ca..396b4de0 100644 --- a/src/traces.rs +++ b/src/traces.rs @@ -5277,7 +5277,7 @@ fn logs_hints( if let Some(cursor) = next_cursor { hints.push(format!( "Next page: bt view logs{} --object-ref {object_ref} --cursor {cursor}{limit_suffix}", - profile_flag_suffix(profile) + org_flag_suffix(profile) )); } else { hints.push(format!( @@ -5299,7 +5299,7 @@ fn trace_hints( let mut hints = vec![ format!( "Trace fetch returns truncated span rows; use `bt view span{} --object-ref {object_ref} --id ` for full single-span payloads.", - profile_flag_suffix(profile) + org_flag_suffix(profile) ), "Write output to a file for long traces.".to_string(), ]; @@ -5312,7 +5312,7 @@ fn trace_hints( if let Some(cursor) = next_cursor { hints.push(format!( "Next page: bt view trace{} --object-ref {object_ref} --trace-id {trace_id} --cursor {cursor}{limit_suffix}", - profile_flag_suffix(profile) + org_flag_suffix(profile) )); } else { hints.push(format!( @@ -5341,7 +5341,7 @@ fn thread_hints(profile: Option<&str>) -> Vec { vec![ format!( "Open an interactive thread view with `bt view thread{} --trace-id `.", - profile_flag_suffix(profile) + org_flag_suffix(profile) ), "Use `--non-interactive` for a compact text transcript.".to_string(), ] @@ -5351,7 +5351,7 @@ fn waterfall_hints(profile: Option<&str>, has_more: bool) -> Vec { let mut hints = vec![ format!( "Render an agent-readable trace report with `bt view waterfall{} --trace-id `.", - profile_flag_suffix(profile) + org_flag_suffix(profile) ), "Use `--json` for computed offsets, token counts, costs, cache metrics, and raw ids." .to_string(), @@ -5406,7 +5406,7 @@ fn print_logs_text( println!("next_cursor: {cursor}"); println!( "next: bt view logs{} --object-ref {object_ref} --cursor {cursor} --non-interactive{limit_suffix}", - profile_flag_suffix(profile) + org_flag_suffix(profile) ); } else { println!("No additional rows."); @@ -5552,7 +5552,7 @@ fn print_thread_text( println!( "\njson: bt view thread --json{} --trace-id {}", - profile_flag_suffix(profile), + org_flag_suffix(profile), target.root_span_id ); } @@ -5589,7 +5589,7 @@ fn print_trace_text( } println!( "\nTrace output is truncated. Re-run with --json for the full trace (`bt view trace --json{0} --object-ref {1} --trace-id {2}`), or fetch a single span with `bt view span{0} --object-ref {1} --id `.", - profile_flag_suffix(profile), + org_flag_suffix(profile), object_ref, trace_id ); @@ -5602,7 +5602,7 @@ fn print_trace_text( println!("next_cursor: {cursor}"); println!( "next: bt view trace{} --object-ref {} --trace-id {} --cursor {} --non-interactive{limit_suffix}", - profile_flag_suffix(profile), + org_flag_suffix(profile), object_ref, trace_id, cursor @@ -5683,7 +5683,7 @@ fn format_object_ref_arg(object_ref: &ObjectRef) -> String { format!("{}:{}", object_ref.object_type, object_ref.object_name) } -fn profile_flag_suffix(org: Option<&str>) -> String { +fn org_flag_suffix(org: Option<&str>) -> String { match org.filter(|p| !p.trim().is_empty()) { Some(org) => format!(" --org {}", btql_quote(org)), None => String::new(), @@ -5741,14 +5741,7 @@ fn apply_url_hints_to_base(mut base: BaseArgs, parsed_url: Option<&ParsedTraceUr } #[cfg(test)] -fn apply_url_hints_with_profile_resolver( - base: BaseArgs, - parsed_url: Option<&ParsedTraceUrl>, - _resolve_profile_for_org: F, -) -> BaseArgs -where - F: Fn(&str) -> Option, -{ +fn apply_url_hints_for_test(base: BaseArgs, parsed_url: Option<&ParsedTraceUrl>) -> BaseArgs { apply_url_hints_to_base(base, parsed_url) } @@ -6716,9 +6709,10 @@ mod tests { quiet_source: None, no_color: false, no_input: false, - profile: None, org_name: None, + org_name_source: None, project: None, + project_source: None, api_key: None, api_key_source: None, prefer_api_key: false, @@ -6954,12 +6948,9 @@ mod tests { let base = base_args(); let parsed = parsed_url_with_org("Lovable"); - let updated = apply_url_hints_with_profile_resolver(base, Some(&parsed), |org| { - (org == "Lovable").then(|| "lovable-profile".to_string()) - }); + let updated = apply_url_hints_for_test(base, Some(&parsed)); assert_eq!(updated.org_name.as_deref(), Some("Lovable")); - assert_eq!(updated.profile, None); } #[test] @@ -6968,9 +6959,7 @@ mod tests { base.org_name = Some("explicit-org".to_string()); let parsed = parsed_url_with_org("Lovable"); - let updated = apply_url_hints_with_profile_resolver(base, Some(&parsed), |_| { - Some("other".to_string()) - }); + let updated = apply_url_hints_for_test(base, Some(&parsed)); assert_eq!(updated.org_name.as_deref(), Some("explicit-org")); } diff --git a/src/traces/waterfall.rs b/src/traces/waterfall.rs index fbdf5fc7..84caefa2 100644 --- a/src/traces/waterfall.rs +++ b/src/traces/waterfall.rs @@ -6,8 +6,8 @@ use serde_json::{Map, Value}; use super::{ build_span_entries, extract_duration_seconds, extract_model_name, extract_parent_span_id, extract_span_end_seconds, extract_span_name_and_type, extract_start_time, - format_compact_duration, format_object_ref_arg, format_u64_with_commas, parse_f64ish, - parse_u64ish, profile_flag_suffix, project_label, span_has_error, value_as_object_owned, + format_compact_duration, format_object_ref_arg, format_u64_with_commas, org_flag_suffix, + parse_f64ish, parse_u64ish, project_label, span_has_error, value_as_object_owned, ResolvedTraceCommandTarget, }; @@ -197,13 +197,13 @@ pub(super) fn print_waterfall_text( } println!( "\nspan detail: bt view span{} --object-ref {} --id ", - profile_flag_suffix(profile), + org_flag_suffix(profile), format_object_ref_arg(&target.object_ref) ); println!("Use the inline `id=` value from a span row."); println!( "json: bt view waterfall --json{} --trace-id {}", - profile_flag_suffix(profile), + org_flag_suffix(profile), target.root_span_id ); } diff --git a/src/utils/profile.rs b/src/utils/profile.rs index 620d681c..613710a0 100644 --- a/src/utils/profile.rs +++ b/src/utils/profile.rs @@ -1,59 +1,5 @@ use crate::auth::ProfileInfo; -#[cfg(test)] -fn resolve_profile_info_from_profiles( - profile: Option<&str>, - org: Option<&str>, - profiles: Vec, -) -> Option { - if let Some(profile_name) = profile { - if let Some(profile) = profiles - .iter() - .find(|profile| profile.name == profile_name) - .cloned() - { - return Some(profile); - } - } - - if let Some(org_name) = org { - if profiles.iter().any(|profile| profile.name == org_name) { - return profiles - .into_iter() - .find(|profile| profile.name == org_name); - } - - let org_matches: Vec<&ProfileInfo> = profiles - .iter() - .filter(|profile| profile.org_name.as_deref() == Some(org_name)) - .collect(); - if org_matches.len() == 1 { - let profile_name = org_matches[0].name.clone(); - return profiles - .into_iter() - .find(|profile| profile.name == profile_name); - } - let oauth_matches: Vec<&ProfileInfo> = org_matches - .iter() - .copied() - .filter(|profile| profile.email.is_some() || profile.user_name.is_some()) - .collect(); - if oauth_matches.len() == 1 { - let profile_name = oauth_matches[0].name.clone(); - return profiles - .into_iter() - .find(|profile| profile.name == profile_name); - } - return None; - } - - if profiles.len() == 1 { - return profiles.into_iter().next(); - } - - None -} - pub(crate) fn profile_author_slug(profile: &ProfileInfo) -> Option { [ profile.user_name.as_deref(), @@ -102,13 +48,11 @@ mod tests { use super::*; fn profile_info( - name: &str, org_name: Option<&str>, user_name: Option<&str>, email: Option<&str>, ) -> ProfileInfo { ProfileInfo { - name: name.to_string(), auth_method: if email.is_some() || user_name.is_some() { "oauth" } else { @@ -122,36 +66,9 @@ mod tests { } } - #[test] - fn resolve_profile_info_prefers_explicit_profile() { - let profile = resolve_profile_info_from_profiles( - Some("work"), - Some("other-org"), - vec![ - profile_info("other", Some("other-org"), None, None), - profile_info("work", Some("work-org"), None, None), - ], - ) - .expect("profile"); - - assert_eq!(profile.name, "work"); - } - - #[test] - fn resolve_profile_info_finds_profile_by_org_name() { - let profile = resolve_profile_info_from_profiles( - None, - Some("work-org"), - vec![profile_info("work", Some("work-org"), None, None)], - ) - .expect("profile"); - - assert_eq!(profile.name, "work"); - } - #[test] fn profile_author_slug_prefers_user_name() { - let profile = profile_info("work", None, Some("Alice Smith"), Some("alice@example.com")); + let profile = profile_info(None, Some("Alice Smith"), Some("alice@example.com")); assert_eq!( profile_author_slug(&profile).as_deref(), Some("alice-smith") @@ -160,13 +77,13 @@ mod tests { #[test] fn profile_author_slug_falls_back_to_email_local_part() { - let profile = profile_info("work", None, None, Some("alice.dev@example.com")); + let profile = profile_info(None, None, Some("alice.dev@example.com")); assert_eq!(profile_author_slug(&profile).as_deref(), Some("alice-dev")); } #[test] - fn profile_author_slug_ignores_internal_profile_name() { - let profile = profile_info("Work Profile", None, None, None); + fn profile_author_slug_returns_none_without_identity_or_org() { + let profile = profile_info(None, None, None); assert_eq!(profile_author_slug(&profile), None); } diff --git a/tests/cli.rs b/tests/cli.rs index a120455c..f40c0682 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -130,12 +130,12 @@ fn status_quiet_and_verbose_conflict() { } #[test] -fn status_json_keeps_local_org_when_global_profile_has_different_org() { +fn status_json_prefers_local_org_over_global_org() { let repo = make_git_repo(); fs::create_dir_all(repo.path().join(".bt")).expect("create local bt dir"); fs::write( repo.path().join(".bt/config.json"), - r#"{"profile":null,"org":"local-org","project":"local-project","project_id":null}"#, + r#"{"org":"local-org","project":"local-project","project_id":null}"#, ) .expect("write local config"); @@ -145,11 +145,9 @@ fn status_json_keeps_local_org_when_global_profile_has_different_org() { fs::create_dir_all(&global_bt_dir).expect("create global bt dir"); fs::write( global_bt_dir.join("config.json"), - r#"{"profile":"default-profile","org":"profile-org"}"#, + r#"{"org":"profile-org"}"#, ) .expect("write global config"); - write_auth_store(config_home.path(), &[("default-profile", "profile-org")]); - let mut cmd = bt_command(); clear_braintrust_auth_env(&mut cmd); cmd.current_dir(repo.path()) @@ -160,7 +158,6 @@ fn status_json_keeps_local_org_when_global_profile_has_different_org() { .success() .stdout(predicate::str::contains(r#""org":"local-org""#)) .stdout(predicate::str::contains(r#""project":"local-project""#)) - .stdout(predicate::str::contains(r#""profile""#).not()) .stdout(predicate::str::contains(r#""org":"profile-org""#).not()); } From 2a8191aa2dbaa5944e53eb0651ba863fe76b814d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 16 Jul 2026 13:43:02 -0700 Subject: [PATCH 03/24] chore: add compatibility code to prevent users from having to re-login to use the new auth storage format orgid,email --- src/auth.rs | 641 ++++++++++++++++++++++++++++++++++++++++----- src/config/mod.rs | 15 ++ src/init.rs | 3 + src/switch.rs | 9 +- tests/functions.rs | 19 ++ 5 files changed, 614 insertions(+), 73 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 8ab0df8c..9833bfad 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -58,6 +58,7 @@ pub struct ResolvedAuth { pub app_url: Option, pub org_name: Option, pub is_oauth: bool, + slot_key: Option, } #[derive(Debug, Clone)] @@ -179,7 +180,7 @@ pub(crate) async fn list_available_orgs_for_api_key( .collect()) } -#[derive(Debug, Clone, Serialize, Deserialize, Default)] +#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)] struct AuthStore { #[serde(default)] profiles: BTreeMap, @@ -191,7 +192,7 @@ struct SecretStore { secrets: BTreeMap, } -#[derive(Debug, Clone, Serialize, Deserialize, Default)] +#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)] struct AuthProfile { #[serde(default)] auth_kind: AuthKind, @@ -425,7 +426,16 @@ pub async fn login(base: &BaseArgs) -> Result { builder = builder.default_project(project); } let login = match builder.build().await { - Ok(client) => client.wait_for_login().await?, + Ok(client) => match client.wait_for_login().await { + Ok(login) => login, + Err(err) => { + let err: anyhow::Error = err.into(); + if !auth.is_oauth && is_unauthorized_auth_error(&err) { + return Err(err.context("API key is not valid")); + } + return Err(err); + } + }, Err(_err) if auth.is_oauth => { let org_name = auth.org_name.clone().unwrap_or_default(); let login = LoginState::new(); @@ -442,9 +452,17 @@ pub async fn login(base: &BaseArgs) -> Result { ); login } - Err(err) => return Err(err.into()), + Err(err) => { + let err: anyhow::Error = err.into(); + if is_unauthorized_auth_error(&err) { + return Err(err.context("API key is not valid")); + } + return Err(err); + } }; + reconcile_resolved_auth_slot(&auth, &login)?; + let api_url = login .api_url() .or(auth.api_url.clone()) @@ -696,6 +714,24 @@ fn has_cached_project_id(base: &BaseArgs) -> bool { .is_some_and(|project_id| !project_id.trim().is_empty()) } +fn is_unauthorized_auth_error(err: &anyhow::Error) -> bool { + err.chain().any(|source| { + if let Some(http_error) = source.downcast_ref::() { + return matches!(http_error.status.as_u16(), 401 | 403); + } + if let Some(sdk_error) = source.downcast_ref::() { + return matches!( + sdk_error, + braintrust_sdk_rust::BraintrustError::Api { + status: 401 | 403, + .. + } + ); + } + false + }) +} + fn maybe_warn_api_key_override(_base: &BaseArgs) {} fn resolve_cli_api_key_override(base: &BaseArgs) -> Option { @@ -755,6 +791,7 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { app_url: base.app_url.clone(), org_name: effective_org_name(base, &cfg_org).map(str::to_string), is_oauth: false, + slot_key: None, }); } @@ -789,6 +826,7 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { app_url: base.app_url.clone(), org_name: effective_org_name(base, &cfg_org).map(str::to_string), is_oauth: false, + slot_key: None, }); } @@ -804,6 +842,7 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { app_url: base.app_url.clone(), org_name: effective_org_name(base, &cfg_org).map(str::to_string), is_oauth: false, + slot_key: None, }) } @@ -825,6 +864,7 @@ fn resolve_preferred_api_key_auth( app_url: base.app_url.clone(), org_name: Some(org.to_string()), is_oauth: false, + slot_key: None, })); } @@ -843,9 +883,11 @@ fn resolve_api_key_profile_auth( cfg_org: &Option, profile_name: &str, ) -> Result { - let profile = store.profiles.get(profile_name).cloned().ok_or_else(|| { - anyhow::anyhow!("auth login '{profile_name}' not found; run `bt auth profiles`") - })?; + let profile = store + .profiles + .get(profile_name) + .cloned() + .ok_or_else(|| anyhow::anyhow!("saved auth login not found; run `bt auth profiles`"))?; let api_key = load_profile_secret_with_legacy( profile_name, profile.legacy_secret_key.as_deref(), @@ -868,51 +910,172 @@ fn resolve_api_key_profile_auth( .map(str::to_string) .or_else(|| profile.org_name.clone()), is_oauth: false, + slot_key: Some(profile_name.to_string()), }; maybe_rekey_api_key_profile_after_secret_load(store, profile_name, &api_key)?; Ok(resolved) } +fn replace_with_canonical_auth_profile( + store: &mut AuthStore, + current_key: &str, + mut profile: AuthProfile, +) -> bool { + let canonical_key = canonical_profile_key(current_key, &profile); + if canonical_key != current_key && profile.legacy_secret_key.is_none() { + // Keep the old key as a lazy keychain fallback. Secrets are relocated + // only when they are next saved, avoiding platform-specific migration + // work while auth.json is being upgraded. + profile.legacy_secret_key = Some(current_key.to_string()); + } + + let unchanged = canonical_key == current_key + && store + .profiles + .get(current_key) + .is_some_and(|existing| existing == &profile); + if unchanged { + return false; + } + + if canonical_key != current_key { + store.profiles.remove(current_key); + if let Some(existing) = store.profiles.get(&canonical_key) { + if !should_replace_migrated_profile(existing, &profile) { + return true; + } + } + } + store.profiles.insert(canonical_key, profile); + true +} + fn maybe_rekey_api_key_profile_after_secret_load( store: &mut AuthStore, profile_name: &str, api_key: &str, ) -> Result<()> { - let Some(profile) = store.profiles.get(profile_name).cloned() else { + let Some(mut profile) = store.profiles.get(profile_name).cloned() else { + return Ok(()); + }; + if profile.auth_kind != AuthKind::ApiKey + || profile + .org_id + .as_deref() + .map(str::trim) + .is_none_or(str::is_empty) + { + return Ok(()); + } + + profile.api_key_hash = Some(api_key_hash(api_key)); + if replace_with_canonical_auth_profile(store, profile_name, profile) { + save_auth_store(store)?; + } + Ok(()) +} + +async fn reconcile_oauth_slot_from_access_token( + store: &mut AuthStore, + slot_key: &str, + access_token: &str, + app_url: &str, +) -> Result<()> { + let Some(mut profile) = store.profiles.get(slot_key).cloned() else { return Ok(()); }; - if profile.auth_kind != AuthKind::ApiKey { + if profile.auth_kind != AuthKind::Oauth || profile.org_id.is_some() { return Ok(()); } - let Some(org_id) = profile - .org_id + + let Some(org_name) = profile + .org_name .as_deref() .map(str::trim) - .filter(|value| !value.is_empty()) + .filter(|org| !org.is_empty()) else { + profile.org_id = Some(String::new()); + if replace_with_canonical_auth_profile(store, slot_key, profile) { + save_auth_store(store)?; + } return Ok(()); }; - let hash = api_key_hash(api_key); - let new_key = api_key_slot_key(&hash, org_id); - let already_canonical = - profile_name == new_key && profile.api_key_hash.as_deref() == Some(&hash); - if already_canonical { + // A legacy auth entry only cached the org name. Resolve the stable ID from + // the newly refreshed token; a failed best-effort lookup must not turn a + // successful token refresh into a failed command. + let Ok(orgs) = fetch_login_orgs(access_token, app_url).await else { + return Ok(()); + }; + let Some(org) = find_login_org(&orgs, org_name) else { + return Ok(()); + }; + + profile.org_id = Some(org.id.clone()); + profile.org_name = Some(org.name.clone()); + let identity = decode_jwt_identity(access_token); + if identity.email.is_some() { + profile.email = identity.email; + } + if identity.name.is_some() { + profile.user_name = identity.name; + } + if replace_with_canonical_auth_profile(store, slot_key, profile) { + save_auth_store(store)?; + } + Ok(()) +} + +fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Result<()> { + let Some(slot_key) = auth.slot_key.as_deref() else { + return Ok(()); + }; + let mut store = load_auth_store()?; + let Some(mut profile) = store.profiles.get(slot_key).cloned() else { + return Ok(()); + }; + + let login_org_id = login.org_id().unwrap_or_default(); + let is_cross_org = profile.auth_kind == AuthKind::Oauth + && auth + .org_name + .as_deref() + .is_none_or(|org| org.trim().is_empty()); + if login_org_id.trim().is_empty() && !is_cross_org { + // The SDK's OAuth compatibility path does not always return org + // metadata. `bt auth profiles` performs the same reconciliation after + // its explicit credential verification request. return Ok(()); } - let mut updated = profile; - updated.api_key_hash = Some(hash); - if profile_name != new_key && updated.legacy_secret_key.is_none() { - updated.legacy_secret_key = Some(profile_name.to_string()); + profile.org_id = Some(login_org_id); + profile.org_name = if is_cross_org { + None + } else { + login + .org_name() + .filter(|org| !org.trim().is_empty()) + .or_else(|| auth.org_name.clone()) + }; + match profile.auth_kind { + AuthKind::ApiKey => { + let Some(api_key) = auth.api_key.as_deref() else { + return Ok(()); + }; + profile.api_key_hash = Some(api_key_hash(api_key)); + if profile.api_key_hint.is_none() { + profile.api_key_hint = Some(obscure_api_key(api_key)); + } + } + AuthKind::Oauth if profile.email.as_deref().is_none_or(str::is_empty) => return Ok(()), + AuthKind::Oauth => {} } - if profile_name != new_key { - store.profiles.remove(profile_name); + if replace_with_canonical_auth_profile(&mut store, slot_key, profile) { + save_auth_store(&store)?; } - store.profiles.insert(new_key, updated); - save_auth_store(store) + Ok(()) } async fn resolve_oauth_profile_auth( @@ -921,11 +1084,10 @@ async fn resolve_oauth_profile_auth( cfg_org: &Option, profile_name: &str, ) -> Result { - let profile = store - .profiles - .get(profile_name) - .cloned() - .ok_or_else(|| anyhow::anyhow!("oauth login '{profile_name}' not found"))?; + let profile = + store.profiles.get(profile_name).cloned().ok_or_else(|| { + anyhow::anyhow!("saved OAuth login not found; run `bt auth profiles`") + })?; let client_id = profile.oauth_client_id.as_deref().ok_or_else(|| { recoverable_auth_error( RecoverableAuthErrorKind::OauthClientId, @@ -951,6 +1113,7 @@ async fn resolve_oauth_profile_auth( app_url, org_name, is_oauth: true, + slot_key: Some(profile_name.to_string()), }; if let Some(cached_access_token) = load_valid_cached_oauth_access_token( @@ -972,8 +1135,9 @@ async fn resolve_oauth_profile_auth( ), ) })?; + let login_label = auth_slot_label(profile_name, &profile); let refreshed = - refresh_oauth_access_token(&api_url, &refresh_token, client_id, profile_name).await?; + refresh_oauth_access_token(&api_url, &refresh_token, client_id, &login_label).await?; save_profile_oauth_access_token(profile_name, &refreshed.access_token)?; let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { @@ -993,6 +1157,13 @@ async fn resolve_oauth_profile_auth( } } save_auth_store(store)?; + reconcile_oauth_slot_from_access_token( + store, + profile_name, + &refreshed.access_token, + auth.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + ) + .await?; auth.api_key = Some(refreshed.access_token); Ok(auth) } @@ -1053,13 +1224,11 @@ fn profile_identity_label(profile: &AuthProfile) -> Option { } } -fn auth_slot_label(name: &str, profile: &AuthProfile) -> String { +fn auth_slot_label(_name: &str, profile: &AuthProfile) -> String { let mut parts = vec![profile_org_label(profile)]; parts.push(auth_kind_label(profile.auth_kind).to_string()); if let Some(identity) = profile_identity_label(profile) { parts.push(identity); - } else if !name.contains("::") { - parts.push(name.to_string()); } parts.join(" — ") } @@ -1212,7 +1381,7 @@ fn profile_label_from_store(name: &str, store: &AuthStore) -> String { .profiles .get(name) .map(|profile| auth_slot_label(name, profile)) - .unwrap_or_else(|| name.to_string()) + .unwrap_or_else(|| "saved auth login".to_string()) } fn select_profile_from_store( @@ -1248,8 +1417,11 @@ fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec Result<()> { .get(profile_name.as_str()) .cloned() .ok_or_else(|| { - anyhow::anyhow!( - "OAuth login '{}' not found; run `bt auth profiles` to see available logins", - profile_name - ) + anyhow::anyhow!("OAuth login not found; run `bt auth profiles` to see available logins") })?; let api_url = profile @@ -1669,9 +1838,9 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { eprintln!("Cached access token expiry before refresh: unknown"); } + let login_label = auth_slot_label(&profile_name, &profile); let refreshed = - refresh_oauth_access_token(&api_url, &refresh_token, &client_id, profile_name.as_str()) - .await?; + refresh_oauth_access_token(&api_url, &refresh_token, &client_id, &login_label).await?; save_profile_oauth_access_token(profile_name.as_str(), &refreshed.access_token)?; let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { @@ -1693,6 +1862,13 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { } } save_auth_store(&store)?; + reconcile_oauth_slot_from_access_token( + &mut store, + profile_name.as_str(), + &refreshed.access_token, + profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + ) + .await?; if let Some(expires_at) = new_expires_at { let now = current_unix_timestamp(); @@ -1827,7 +2003,7 @@ fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) -> } async fn run_profiles(base: &BaseArgs, _args: AuthProfilesArgs) -> Result<()> { - let store = load_auth_store()?; + let mut store = load_auth_store()?; if store.profiles.is_empty() { return emit_result(base.json, serde_json::json!([]), || { println!("No saved auth logins. Run `bt auth login` to create one.") @@ -1835,6 +2011,7 @@ async fn run_profiles(base: &BaseArgs, _args: AuthProfilesArgs) -> Result<()> { } let verifications = verify_all_profiles_from_store(&store).await; + reconcile_verified_auth_slots(&mut store, &verifications)?; let all_network_errors = verifications .iter() .all(|v| v.status == "error" && !v.error.as_deref().unwrap_or("").contains("invalid")); @@ -1892,16 +2069,19 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result< let label = auth_slot_label(profile_name, &profile); if !force { - if let Some(term) = ui::prompt_term() { - let confirmed = Confirm::new() - .with_prompt(format!("Delete {label}?")) - .default(false) - .interact_on(&term)?; - if !confirmed { - return emit_result(base_json, auth_profile_json(&profile, "cancelled"), || { - eprintln!("Cancelled") - }); - } + let term = ui::prompt_term().ok_or_else(|| { + anyhow::anyhow!( + "logout confirmation requires an interactive terminal; rerun with --force" + ) + })?; + let confirmed = Confirm::new() + .with_prompt(format!("Delete {label}?")) + .default(false) + .interact_on(&term)?; + if !confirmed { + return emit_result(base_json, auth_profile_json(&profile, "cancelled"), || { + eprintln!("Cancelled") + }); } } @@ -2021,6 +2201,8 @@ fn load_credential_for_profile(name: &str, profile: &AuthProfile) -> CredentialL pub struct ProfileVerification { #[serde(skip_serializing)] pub name: String, + #[serde(skip_serializing)] + slot_hash: Option, pub auth: String, #[serde(skip_serializing_if = "Option::is_none")] pub org: Option, @@ -2054,6 +2236,7 @@ fn build_verification( }; ProfileVerification { name: name.to_string(), + slot_hash: None, auth: auth_kind.to_string(), org, org_id, @@ -2102,17 +2285,38 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifi }; match fetch_login_orgs(&credential, app_url).await { - Ok(_) => mk(ProfileStatus::Ok, jwt_id, hint), + Ok(orgs) => { + let mut verification = mk(ProfileStatus::Ok, jwt_id, hint); + if !is_cross_org_oauth_profile(profile) { + if let Some(org) = profile + .org_id + .as_deref() + .and_then(|id| find_login_org(&orgs, id)) + .or_else(|| { + profile + .org_name + .as_deref() + .and_then(|name| find_login_org(&orgs, name)) + }) + { + verification.org = Some(org.name.clone()); + verification.org_id = Some(org.id.clone()); + } + } + if profile.auth_kind == AuthKind::ApiKey { + verification.slot_hash = Some(api_key_hash(&credential)); + } + verification + } Err(e) => { - let msg = e.to_string(); - let status = if msg.contains("401") || msg.contains("Unauthorized") { + let status = if is_unauthorized_auth_error(&e) { if profile.auth_kind == AuthKind::Oauth { ProfileStatus::Expired } else { ProfileStatus::Error("invalid API key".to_string()) } } else { - ProfileStatus::Error(msg) + ProfileStatus::Error(e.to_string()) }; mk(status, None, hint) } @@ -2137,6 +2341,59 @@ async fn verify_all_profiles_from_store(store: &AuthStore) -> Vec Result<()> { + let mut changed = false; + for verification in verifications + .iter() + .filter(|verification| verification.status == "ok") + { + let Some(mut profile) = store.profiles.get(&verification.name).cloned() else { + continue; + }; + + if let Some(org_id) = verification.org_id.as_deref() { + profile.org_id = Some(org_id.to_string()); + profile.org_name = verification.org.clone(); + } else if is_cross_org_oauth_profile(&profile) { + profile.org_id = Some(String::new()); + profile.org_name = None; + } + + match profile.auth_kind { + AuthKind::ApiKey => { + let Some(hash) = verification.slot_hash.as_deref() else { + continue; + }; + if profile.org_id.as_deref().is_none_or(str::is_empty) { + continue; + } + profile.api_key_hash = Some(hash.to_string()); + } + AuthKind::Oauth => { + if let Some(email) = verification.user_email.as_deref() { + profile.email = Some(email.to_string()); + } + if let Some(user_name) = verification.user_name.as_deref() { + profile.user_name = Some(user_name.to_string()); + } + if profile.email.as_deref().is_none_or(str::is_empty) { + continue; + } + } + } + + changed |= replace_with_canonical_auth_profile(store, verification.name.as_str(), profile); + } + + if changed { + save_auth_store(store)?; + } + Ok(()) +} + fn format_verification_line(v: &ProfileVerification) -> String { let mut parts = vec![ v.org.clone().unwrap_or_else(|| "cross-org".to_string()), @@ -3559,13 +3816,29 @@ fn load_auth_store_from_path(path: &Path) -> Result { .with_context(|| format!("failed to read auth config {}", path.display()))?; let store: AuthStore = serde_json::from_str(&data) .with_context(|| format!("failed to parse auth config {}", path.display()))?; - Ok(migrate_auth_store(store)) + let migrated = migrate_auth_store(store.clone()); + if migrated != store { + save_auth_store_to_path(path, &migrated).with_context(|| { + format!("failed to persist auth config migration {}", path.display()) + })?; + } + Ok(migrated) } fn migrate_auth_store(store: AuthStore) -> AuthStore { let mut migrated = AuthStore::default(); for (old_key, mut profile) in store.profiles { normalize_profile_cached_fields_from_key(&old_key, &mut profile); + if profile.auth_kind == AuthKind::Oauth + && profile.org_id.is_none() + && profile + .org_name + .as_deref() + .is_none_or(|org| org.trim().is_empty()) + { + profile.org_id = Some(String::new()); + profile.org_name = None; + } let new_key = canonical_profile_key(&old_key, &profile); if new_key != old_key && profile.legacy_secret_key.is_none() { profile.legacy_secret_key = Some(old_key.clone()); @@ -3605,9 +3878,7 @@ fn normalize_profile_cached_fields_from_key(current_key: &str, profile: &mut Aut AuthKind::Oauth => { let key_matches_email = profile.email.as_deref() == Some(right); if key_matches_email || (profile.email.is_none() && right.contains('@')) { - if profile.org_id.is_none() { - profile.org_id = Some(left.to_string()); - } + profile.org_id = Some(left.to_string()); if profile.email.is_none() { profile.email = Some(right.to_string()); } @@ -3615,12 +3886,8 @@ fn normalize_profile_cached_fields_from_key(current_key: &str, profile: &mut Aut } AuthKind::ApiKey => { if looks_like_sha256_hex(left) && !right.trim().is_empty() { - if profile.api_key_hash.is_none() { - profile.api_key_hash = Some(left.to_string()); - } - if profile.org_id.is_none() { - profile.org_id = Some(right.to_string()); - } + profile.api_key_hash = Some(left.to_string()); + profile.org_id = Some(right.to_string()); } } } @@ -4238,6 +4505,105 @@ mod tests { assert_eq!(info.org_name, None); } + fn save_cached_oauth_login(store: &mut AuthStore, org_id: &str, org_name: &str) -> String { + let slot_key = oauth_slot_key(org_id, "user@example.test"); + store.profiles.insert( + slot_key.clone(), + AuthProfile { + auth_kind: AuthKind::Oauth, + api_url: Some("https://api.example.test".to_string()), + app_url: Some("https://www.example.test".to_string()), + org_id: Some(org_id.to_string()), + org_name: Some(org_name.to_string()), + oauth_client_id: Some("bt_cli".to_string()), + oauth_access_expires_at: Some(current_unix_timestamp() + 3600), + user_name: Some("Test User".to_string()), + email: Some("user@example.test".to_string()), + ..Default::default() + }, + ); + save_profile_secret_plaintext( + &oauth_access_secret_key(&slot_key), + "cached-oauth-access-token", + ) + .expect("save cached OAuth token"); + slot_key + } + + #[tokio::test] + async fn auth_precedence_keeps_env_api_key_below_oauth() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + save_cached_oauth_login(&mut store, "org_fake", "test-org"); + save_auth_store(&store).expect("save auth store"); + let mut base = make_base(); + base.org_name = Some("test-org".to_string()); + base.api_key = Some("environment-api-key".to_string()); + base.api_key_source = Some(crate::args::ArgValueSource::EnvVariable); + + let resolved = resolve_auth(&base).await.expect("resolve auth"); + + assert!(resolved.is_oauth); + assert_eq!( + resolved.api_key.as_deref(), + Some("cached-oauth-access-token") + ); + } + + #[tokio::test] + async fn auth_precedence_cli_api_key_overrides_oauth() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + save_cached_oauth_login(&mut store, "org_fake", "test-org"); + save_auth_store(&store).expect("save auth store"); + let mut base = make_base(); + base.org_name = Some("test-org".to_string()); + base.api_key = Some("command-line-api-key".to_string()); + base.api_key_source = Some(crate::args::ArgValueSource::CommandLine); + + let resolved = resolve_auth(&base).await.expect("resolve auth"); + + assert!(!resolved.is_oauth); + assert_eq!(resolved.api_key.as_deref(), Some("command-line-api-key")); + } + + #[tokio::test] + async fn auth_precedence_prefer_api_key_promotes_env_api_key() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + save_cached_oauth_login(&mut store, "org_fake", "test-org"); + save_auth_store(&store).expect("save auth store"); + let mut base = make_base(); + base.org_name = Some("test-org".to_string()); + base.api_key = Some("environment-api-key".to_string()); + base.api_key_source = Some(crate::args::ArgValueSource::EnvVariable); + base.prefer_api_key = true; + + let resolved = resolve_auth(&base).await.expect("resolve auth"); + + assert!(!resolved.is_oauth); + assert_eq!(resolved.api_key.as_deref(), Some("environment-api-key")); + } + + #[tokio::test] + async fn auth_precedence_prefer_api_key_falls_back_to_oauth_without_key() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + save_cached_oauth_login(&mut store, "org_fake", "test-org"); + save_auth_store(&store).expect("save auth store"); + let mut base = make_base(); + base.org_name = Some("test-org".to_string()); + base.prefer_api_key = true; + + let resolved = resolve_auth(&base).await.expect("resolve auth"); + + assert!(resolved.is_oauth); + assert_eq!( + resolved.api_key.as_deref(), + Some("cached-oauth-access-token") + ); + } + #[tokio::test] async fn active_auth_info_prefer_api_key_selects_stored_key_for_org() { let _env = TestEnv::new(None, None).await; @@ -4356,6 +4722,34 @@ mod tests { assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_work")); } + #[test] + fn migrate_auth_store_rekeys_cross_org_oauth_with_empty_org_id() { + let mut store = AuthStore::default(); + store.profiles.insert( + "legacy-cross-org".to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_name: None, + email: Some("user@example.test".to_string()), + oauth_client_id: Some("bt_cli_legacy".to_string()), + ..Default::default() + }, + ); + + let migrated = migrate_auth_store(store); + let profile = migrated + .profiles + .get(&oauth_slot_key("", "user@example.test")) + .expect("cross-org OAuth slot"); + + assert_eq!(profile.org_id.as_deref(), Some("")); + assert_eq!( + profile.legacy_secret_key.as_deref(), + Some("legacy-cross-org") + ); + assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_legacy")); + } + #[test] fn migrate_auth_store_rekeys_api_key_slots_by_hash_and_org() { let hash = api_key_hash("test-api-key"); @@ -4380,6 +4774,113 @@ mod tests { assert_eq!(profile.api_key_hint.as_deref(), Some("test-****i-key")); } + #[test] + fn load_auth_store_persists_canonical_slot_migration() { + let unique = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock") + .as_nanos(); + let dir = std::env::temp_dir().join(format!("bt-auth-migration-test-{unique}")); + fs::create_dir_all(&dir).expect("create dir"); + let path = dir.join("auth.json"); + + let mut store = AuthStore::default(); + store.profiles.insert( + "legacy-login".to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + email: Some("user@example.test".to_string()), + oauth_client_id: Some("bt_cli_legacy".to_string()), + ..Default::default() + }, + ); + save_auth_store_to_path(&path, &store).expect("save legacy store"); + + let loaded = load_auth_store_from_path(&path).expect("load and migrate"); + let persisted: AuthStore = + serde_json::from_str(&fs::read_to_string(&path).expect("read migrated store")) + .expect("parse migrated store"); + let slot_key = oauth_slot_key("org_fake", "user@example.test"); + + for migrated in [&loaded, &persisted] { + let profile = migrated + .profiles + .get(&slot_key) + .expect("canonical OAuth slot"); + assert_eq!(profile.legacy_secret_key.as_deref(), Some("legacy-login")); + assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_legacy")); + } + + let _ = fs::remove_dir_all(&dir); + } + + #[tokio::test] + async fn verified_legacy_slots_gain_stable_keys() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + store.profiles.insert( + "legacy-oauth".to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_name: Some("test-org".to_string()), + oauth_client_id: Some("bt_cli_legacy".to_string()), + ..Default::default() + }, + ); + store.profiles.insert( + "legacy-api-key".to_string(), + AuthProfile { + auth_kind: AuthKind::ApiKey, + org_name: Some("test-org".to_string()), + api_key_hint: Some("sk-****abcde".to_string()), + ..Default::default() + }, + ); + let hash = api_key_hash("test-api-key"); + let verifications = vec![ + ProfileVerification { + name: "legacy-oauth".to_string(), + slot_hash: None, + auth: "oauth".to_string(), + org: Some("test-org".to_string()), + org_id: Some("org_fake".to_string()), + user_name: Some("Test User".to_string()), + user_email: Some("user@example.test".to_string()), + api_key_hint: None, + status: "ok".to_string(), + error: None, + }, + ProfileVerification { + name: "legacy-api-key".to_string(), + slot_hash: Some(hash.clone()), + auth: "api_key".to_string(), + org: Some("test-org".to_string()), + org_id: Some("org_fake".to_string()), + user_name: None, + user_email: None, + api_key_hint: Some("sk-****abcde".to_string()), + status: "ok".to_string(), + error: None, + }, + ]; + + reconcile_verified_auth_slots(&mut store, &verifications) + .expect("reconcile verified slots"); + + let oauth = store + .profiles + .get(&oauth_slot_key("org_fake", "user@example.test")) + .expect("canonical OAuth slot"); + assert_eq!(oauth.legacy_secret_key.as_deref(), Some("legacy-oauth")); + let api_key = store + .profiles + .get(&api_key_slot_key(&hash, "org_fake")) + .expect("canonical API-key slot"); + assert_eq!(api_key.legacy_secret_key.as_deref(), Some("legacy-api-key")); + } + #[test] fn migrate_auth_store_dedupes_oauth_slots_by_latest_expiry() { let mut store = AuthStore::default(); @@ -4641,6 +5142,7 @@ mod tests { fn format_verification_line_ok_with_identity() { let v = ProfileVerification { name: "work".into(), + slot_hash: None, auth: "oauth".into(), org: Some("acme".into()), org_id: None, @@ -4660,6 +5162,7 @@ mod tests { fn format_verification_line_ok_with_api_key_hint() { let v = ProfileVerification { name: "work".into(), + slot_hash: Some(api_key_hash("test-api-key")), auth: "api_key".into(), org: Some("acme".into()), org_id: None, @@ -4679,6 +5182,7 @@ mod tests { fn format_verification_line_expired() { let v = ProfileVerification { name: "old".into(), + slot_hash: None, auth: "oauth".into(), org: None, org_id: None, @@ -4698,6 +5202,7 @@ mod tests { fn format_verification_line_error() { let v = ProfileVerification { name: "bad".into(), + slot_hash: None, auth: "api_key".into(), org: Some("corp".into()), org_id: None, diff --git a/src/config/mod.rs b/src/config/mod.rs index b6dc85fb..26493ada 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -522,6 +522,21 @@ mod tests { assert!(config.extra.contains_key("another")); } + #[test] + fn legacy_profile_key_is_ignored_and_not_persisted() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("config.json"); + fs::write(&path, r#"{"org":"test-org","profile":"legacy-login"}"#).unwrap(); + + let config = load_file(&path); + assert_eq!(config.org.as_deref(), Some("test-org")); + assert!(!config.extra.contains_key("profile")); + + save_file(&path, &config).unwrap(); + let persisted = fs::read_to_string(&path).unwrap(); + assert!(!persisted.contains("profile")); + } + #[test] fn unknown_keys_roundtrip_through_save() { let tmp = TempDir::new().unwrap(); diff --git a/src/init.rs b/src/init.rs index f9e1e03c..c03594a6 100644 --- a/src/init.rs +++ b/src/init.rs @@ -64,6 +64,9 @@ pub async fn run(base: BaseArgs, _args: InitArgs) -> Result<()> { bail!("--org and --project required in non-interactive mode"); } else { let mut login_base = base.clone(); + if login_base.org_name.is_none() { + login_base.org_name = config::load().ok().and_then(|cfg| cfg.org); + } if login_base.org_name.is_none() { login_base.org_name = select_saved_auth_org_for_init()?; } diff --git a/src/switch.rs b/src/switch.rs index 021ccc16..40fed377 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -57,11 +57,10 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { let mut login_base = base.clone(); if login_base.org_name.is_none() { - let saved_org = select_saved_auth_org_for_switch()?; - login_base.org_name = resolved_org - .clone() - .or_else(|| current_cfg.org.clone()) - .or(saved_org); + login_base.org_name = resolved_org.clone().or_else(|| current_cfg.org.clone()); + if login_base.org_name.is_none() { + login_base.org_name = select_saved_auth_org_for_switch()?; + } } let ctx = login(&login_base).await?; diff --git a/tests/functions.rs b/tests/functions.rs index f91c25d2..84a9973f 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -793,6 +793,25 @@ fn auth_logout_json_with_no_profiles_emits_empty_status() { assert_eq!(stdout, r#"{"status":"empty"}"#); } +#[test] +fn auth_logout_requires_force_without_an_interactive_terminal() { + let cwd = tempdir().expect("create temp cwd"); + let config_dir = tempdir().expect("create temp config dir"); + seed_api_key_profile(config_dir.path()); + + let output = auth_sub_command( + cwd.path(), + config_dir.path(), + &["logout", "--org", "test-org", "--no-input"], + ) + .output() + .expect("run non-interactive auth logout"); + + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr).contains("rerun with --force")); + assert!(config_dir.path().join("bt").join("auth.json").exists()); +} + #[test] fn auth_refresh_errors_when_no_oauth_login_exists_even_with_json() { // --json must not swallow a real error: refresh only applies to OAuth From fd70cdf352d6b255cb0fe6f97b40b420d3dc4198 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 16 Jul 2026 15:27:28 -0700 Subject: [PATCH 04/24] chore: remove dead code --- src/auth.rs | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 9833bfad..b7f9b3be 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -363,7 +363,6 @@ pub async fn login_read_only(base: &BaseArgs) -> Result { /// Build login context from stored auth without forcing a login validation request. /// Use for read-oriented flows where downstream API calls can surface auth errors. pub async fn fast_login(base: &BaseArgs) -> Result { - maybe_warn_api_key_override(base); let auth = resolve_auth(base).await?; let api_key = auth.api_key.clone().ok_or_else(|| { anyhow::anyhow!( @@ -397,7 +396,6 @@ pub async fn fast_login(base: &BaseArgs) -> Result { } pub async fn login(base: &BaseArgs) -> Result { - maybe_warn_api_key_override(base); let auth = resolve_auth(base).await?; let api_key = auth.api_key.clone().ok_or_else(|| { anyhow::anyhow!( @@ -732,8 +730,6 @@ fn is_unauthorized_auth_error(err: &anyhow::Error) -> bool { }) } -fn maybe_warn_api_key_override(_base: &BaseArgs) {} - fn resolve_cli_api_key_override(base: &BaseArgs) -> Option { if matches!( base.api_key_source, @@ -3623,13 +3619,6 @@ fn save_profile_oauth_refresh_token(profile_name: &str, refresh_token: &str) -> save_profile_secret(&key, refresh_token) } -#[cfg(test)] -#[allow(dead_code)] -fn load_profile_oauth_refresh_token(profile_name: &str) -> Result> { - let key = oauth_refresh_secret_key(profile_name); - load_profile_secret(&key) -} - fn load_profile_oauth_refresh_token_for_profile( profile_name: &str, profile: &AuthProfile, @@ -3652,13 +3641,6 @@ fn save_profile_oauth_access_token(profile_name: &str, access_token: &str) -> Re save_profile_secret(&key, access_token) } -#[cfg(test)] -#[allow(dead_code)] -fn load_profile_oauth_access_token(profile_name: &str) -> Result> { - let key = oauth_access_secret_key(profile_name); - load_profile_secret(&key) -} - fn load_profile_oauth_access_token_for_profile( profile_name: &str, profile: &AuthProfile, @@ -3789,12 +3771,6 @@ fn api_key_slot_key(api_key_hash: &str, org_id: &str) -> String { format!("{api_key_hash}::{org_id}") } -#[cfg(test)] -#[allow(dead_code)] -fn split_slot_key(slot_key: &str) -> Option<(&str, &str)> { - slot_key.split_once("::") -} - fn current_unix_timestamp() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) From b18e27ea8a8a94315f7aed89383b71161265fc0f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 16 Jul 2026 16:06:17 -0700 Subject: [PATCH 05/24] chore: shared helper for bt status/login to prevent drift remove dead code --- src/auth.rs | 357 +++++++++++++++++++++++++++++++--------------------- 1 file changed, 216 insertions(+), 141 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index b7f9b3be..7055a346 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -776,101 +776,111 @@ fn effective_org_name<'a>(base: &'a BaseArgs, cfg_org: &'a Option) -> Op .or_else(|| crate::config::trimmed_option(cfg_org.as_deref())) } -pub async fn resolve_auth(base: &BaseArgs) -> Result { - let mut store = load_auth_store()?; - let cfg_org = config_auth_context(base); - - if let Some(api_key) = resolve_cli_api_key_override(base) { - return Ok(ResolvedAuth { - api_key: Some(api_key), - api_url: base.api_url.clone(), - app_url: base.app_url.clone(), - org_name: effective_org_name(base, &cfg_org).map(str::to_string), - is_oauth: false, - slot_key: None, - }); - } - - if base.prefer_api_key { - if let Some(auth) = resolve_preferred_api_key_auth(base, &mut store, &cfg_org)? { - return Ok(auth); +/// The auth source selected by the precedence ladder, before any live +/// credential is fetched. `resolve_auth` turns this into a `ResolvedAuth` +/// (fetching/refreshing tokens); `active_auth_info` turns it into a +/// `ProfileInfo` for display. Both share [`resolve_auth_source`] so the +/// precedence order documented in the README lives in exactly one place. +#[derive(Debug, Clone, PartialEq, Eq)] +enum AuthSource { + CliApiKey(String), + EnvApiKey(String), + Oauth(String), + ApiKey(String), + None, +} + +/// Pure auth-source precedence ladder (README "Auth resolution order"): +/// 1. explicit `--api-key` +/// 2. `--prefer-api-key`: `BRAINTRUST_API_KEY` → stored API key → OAuth fallback +/// 3. stored OAuth login for the selected org +/// 4. stored API key login for the selected org +/// 5. `BRAINTRUST_API_KEY` +/// +/// The slot selectors return `Ok(None)` when no candidate matches (the ladder +/// continues) and may return `Err` for an ambiguous selection that neither +/// caller can resolve without prompting (the ladder stops). +fn resolve_auth_source( + prefer_api_key: bool, + cli_api_key: Option, + env_api_key: impl Fn() -> Option, + select_oauth: impl Fn() -> Result>, + select_api_key: impl Fn() -> Result>, +) -> Result { + if let Some(api_key) = cli_api_key { + return Ok(AuthSource::CliApiKey(api_key)); + } + + if prefer_api_key { + if let Some(api_key) = env_api_key() { + return Ok(AuthSource::EnvApiKey(api_key)); } - if let Some(profile_name) = - select_oauth_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? - { - return resolve_oauth_profile_auth(base, &mut store, &cfg_org, &profile_name).await; + if let Some(slot) = select_api_key()? { + return Ok(AuthSource::ApiKey(slot)); } - bail!("--prefer-api-key requires an API key or OAuth login for the selected org"); + if let Some(slot) = select_oauth()? { + return Ok(AuthSource::Oauth(slot)); + } + return Ok(AuthSource::None); } - if let Some(profile_name) = - select_oauth_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? - { - return resolve_oauth_profile_auth(base, &mut store, &cfg_org, &profile_name).await; + if let Some(slot) = select_oauth()? { + return Ok(AuthSource::Oauth(slot)); } - - if let Some(profile_name) = - select_api_key_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? - { - return resolve_api_key_profile_auth(base, &mut store, &cfg_org, &profile_name); + if let Some(slot) = select_api_key()? { + return Ok(AuthSource::ApiKey(slot)); } - - if let Some(api_key) = resolve_env_api_key(base) { - return Ok(ResolvedAuth { - api_key: Some(api_key), - api_url: base.api_url.clone(), - app_url: base.app_url.clone(), - org_name: effective_org_name(base, &cfg_org).map(str::to_string), - is_oauth: false, - slot_key: None, - }); + if let Some(api_key) = env_api_key() { + return Ok(AuthSource::EnvApiKey(api_key)); } - - if effective_org_name(base, &cfg_org).is_none() { - if let Some(err) = missing_org_for_stored_logins_error(&store) { - return Err(err); - } - } - - Ok(ResolvedAuth { - api_key: None, - api_url: base.api_url.clone(), - app_url: base.app_url.clone(), - org_name: effective_org_name(base, &cfg_org).map(str::to_string), - is_oauth: false, - slot_key: None, - }) + Ok(AuthSource::None) } -fn resolve_preferred_api_key_auth( - base: &BaseArgs, - store: &mut AuthStore, - cfg_org: &Option, -) -> Result> { - let Some(org) = effective_org_name(base, cfg_org) else { - bail!( - "--prefer-api-key requires an org; pass --org or run `bt switch /`" - ); - }; +pub async fn resolve_auth(base: &BaseArgs) -> Result { + let mut store = load_auth_store()?; + let cfg_org = config_auth_context(base); + let can_prompt = ui::can_prompt(); + + let source = resolve_auth_source( + base.prefer_api_key, + resolve_cli_api_key_override(base), + || resolve_env_api_key(base), + || select_oauth_profile_for_auth(base, &store, &cfg_org, can_prompt), + || select_api_key_profile_for_auth(base, &store, &cfg_org, can_prompt), + )?; - if let Some(api_key) = resolve_env_api_key(base) { - return Ok(Some(ResolvedAuth { + match source { + AuthSource::CliApiKey(api_key) | AuthSource::EnvApiKey(api_key) => Ok(ResolvedAuth { api_key: Some(api_key), api_url: base.api_url.clone(), app_url: base.app_url.clone(), - org_name: Some(org.to_string()), + org_name: effective_org_name(base, &cfg_org).map(str::to_string), is_oauth: false, slot_key: None, - })); - } - - if let Some(profile_name) = - select_api_key_profile_for_auth(base, store, cfg_org, ui::can_prompt())? - { - return resolve_api_key_profile_auth(base, store, cfg_org, &profile_name).map(Some); + }), + AuthSource::Oauth(slot) => { + resolve_oauth_profile_auth(base, &mut store, &cfg_org, &slot).await + } + AuthSource::ApiKey(slot) => resolve_api_key_profile_auth(base, &mut store, &cfg_org, &slot), + AuthSource::None => { + if base.prefer_api_key { + bail!("--prefer-api-key requires an API key or OAuth login for the selected org"); + } + if effective_org_name(base, &cfg_org).is_none() { + if let Some(err) = missing_org_for_stored_logins_error(&store) { + return Err(err); + } + } + Ok(ResolvedAuth { + api_key: None, + api_url: base.api_url.clone(), + app_url: base.app_url.clone(), + org_name: effective_org_name(base, &cfg_org).map(str::to_string), + is_oauth: false, + slot_key: None, + }) + } } - - Ok(None) } fn resolve_api_key_profile_auth( @@ -893,7 +903,7 @@ fn resolve_api_key_profile_auth( RecoverableAuthErrorKind::StoredCredential, format!( "no keychain credential found for auth login '{}'; re-run `bt auth login --org --api-key `", - auth_slot_label(profile_name, &profile) + auth_slot_label(&profile) ), ) })?; @@ -1089,7 +1099,7 @@ async fn resolve_oauth_profile_auth( RecoverableAuthErrorKind::OauthClientId, format!( "oauth login for '{}' is missing client_id; re-run `bt auth login --oauth --org `", - auth_slot_label(profile_name, &profile) + auth_slot_label(&profile) ), ) })?; @@ -1127,11 +1137,11 @@ async fn resolve_oauth_profile_auth( RecoverableAuthErrorKind::OauthRefreshToken, format!( "oauth refresh token missing for '{}'; re-run `bt auth login --oauth --org `", - auth_slot_label(profile_name, &profile) + auth_slot_label(&profile) ), ) })?; - let login_label = auth_slot_label(profile_name, &profile); + let login_label = auth_slot_label(&profile); let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, client_id, &login_label).await?; save_profile_oauth_access_token(profile_name, &refreshed.access_token)?; @@ -1220,7 +1230,7 @@ fn profile_identity_label(profile: &AuthProfile) -> Option { } } -fn auth_slot_label(_name: &str, profile: &AuthProfile) -> String { +fn auth_slot_label(profile: &AuthProfile) -> String { let mut parts = vec![profile_org_label(profile)]; parts.push(auth_kind_label(profile.auth_kind).to_string()); if let Some(identity) = profile_identity_label(profile) { @@ -1284,48 +1294,35 @@ fn ad_hoc_api_key_profile(org: Option<&str>, api_key: &str) -> ProfileInfo { } pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Option { - if let Some(api_key) = resolve_cli_api_key_override(base) { - return Some(ad_hoc_api_key_profile(org, &api_key)); - } - let store = load_auth_store().unwrap_or_default(); - if base.prefer_api_key { - let org = org?; - if let Some(api_key) = resolve_env_api_key(base) { - return Some(ad_hoc_api_key_profile(Some(org), &api_key)); - } - - let api_key_candidates = auth_profile_names_by_kind(&store, Some(org), AuthKind::ApiKey); - match api_key_candidates.as_slice() { - [name] => return profile_info_for_candidate(&store, name), - [] => {} - _ => return None, - } - - let oauth_candidates = auth_profile_names_by_kind(&store, Some(org), AuthKind::Oauth); - return match oauth_candidates.as_slice() { - [name] => profile_info_for_candidate(&store, name), - [] => None, - _ => None, - }; - } + // Display-side mirror of the interactive selectors: a lone candidate is + // chosen, no candidates continue the ladder, and an ambiguous set stops it + // (a real command would prompt or error, so status shows nothing here). + let select = |kind| match auth_profile_names_by_kind(&store, org, kind).as_slice() { + [] => Ok(None), + [name] => Ok(Some((*name).to_string())), + _ => bail!("multiple {kind:?} logins"), + }; - let oauth_candidates = auth_profile_names_by_kind(&store, org, AuthKind::Oauth); - match oauth_candidates.as_slice() { - [name] => return profile_info_for_candidate(&store, name), - [] => {} - _ => return None, - } + let source = resolve_auth_source( + base.prefer_api_key, + resolve_cli_api_key_override(base), + || resolve_env_api_key(base), + || select(AuthKind::Oauth), + || select(AuthKind::ApiKey), + ) + .ok()?; - let api_key_candidates = auth_profile_names_by_kind(&store, org, AuthKind::ApiKey); - match api_key_candidates.as_slice() { - [name] => return profile_info_for_candidate(&store, name), - [] => {} - _ => return None, + match source { + AuthSource::CliApiKey(api_key) | AuthSource::EnvApiKey(api_key) => { + Some(ad_hoc_api_key_profile(org, &api_key)) + } + AuthSource::Oauth(slot) | AuthSource::ApiKey(slot) => { + profile_info_for_candidate(&store, &slot) + } + AuthSource::None => None, } - - resolve_env_api_key(base).map(|api_key| ad_hoc_api_key_profile(org, &api_key)) } fn missing_org_for_stored_logins_error(store: &AuthStore) -> Option { @@ -1350,7 +1347,7 @@ fn missing_org_for_stored_logins_error(store: &AuthStore) -> Option>() .join(", "); let all_api_key = candidates @@ -1376,7 +1373,7 @@ fn profile_label_from_store(name: &str, store: &AuthStore) -> String { store .profiles .get(name) - .map(|profile| auth_slot_label(name, profile)) + .map(auth_slot_label) .unwrap_or_else(|| "saved auth login".to_string()) } @@ -1415,7 +1412,7 @@ fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec Result<()> { let selected_api_url = resolve_profile_api_url(base.api_url.clone(), Some(&selected_org), &login_orgs)?; - let _slot_key = commit_api_key_profile( + commit_api_key_profile( &api_key, selected_api_url.clone(), base.app_url.clone(), @@ -1648,7 +1645,7 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { let selected_api_url = resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)?; - let _slot_key = commit_oauth_profile( + commit_oauth_profile( &oauth_tokens, selected_api_url.clone(), app_url.clone(), @@ -1696,7 +1693,7 @@ pub(crate) fn commit_api_key_profile( app_url: Option, org_id: String, org_name: String, -) -> Result { +) -> Result<()> { let hash = api_key_hash(api_key); let slot_key = api_key_slot_key(&hash, &org_id); save_profile_secret(&slot_key, api_key)?; @@ -1706,7 +1703,7 @@ pub(crate) fn commit_api_key_profile( delete_legacy_profile_secrets(old_profile); } store.profiles.insert( - slot_key.clone(), + slot_key, AuthProfile { auth_kind: AuthKind::ApiKey, api_url: Some(api_url), @@ -1722,8 +1719,7 @@ pub(crate) fn commit_api_key_profile( legacy_secret_key: None, }, ); - save_auth_store(&store)?; - Ok(slot_key) + save_auth_store(&store) } fn commit_oauth_profile( @@ -1732,7 +1728,7 @@ fn commit_oauth_profile( app_url: String, client_id: String, selected_org: Option<&LoginOrgInfo>, -) -> Result { +) -> Result<()> { let refresh_token = tokens.refresh_token.as_ref().ok_or_else(|| { anyhow::anyhow!( "oauth token response did not include a refresh_token; cannot create persistent oauth login" @@ -1762,7 +1758,7 @@ fn commit_oauth_profile( delete_legacy_profile_secrets(old_profile); } store.profiles.insert( - slot_key.clone(), + slot_key, AuthProfile { auth_kind: AuthKind::Oauth, api_url: Some(api_url), @@ -1778,8 +1774,7 @@ fn commit_oauth_profile( legacy_secret_key: None, }, ); - save_auth_store(&store)?; - Ok(slot_key) + save_auth_store(&store) } async fn run_login_refresh(base: &BaseArgs) -> Result<()> { @@ -1806,7 +1801,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { let client_id = profile.oauth_client_id.clone().ok_or_else(|| { anyhow::anyhow!( "OAuth login for '{}' is missing client_id; re-run `bt auth login --oauth --org `", - auth_slot_label(&profile_name, &profile) + auth_slot_label(&profile) ) })?; let previous_expires_at = profile.oauth_access_expires_at; @@ -1815,14 +1810,14 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { || { anyhow::anyhow!( "OAuth refresh token missing for '{}'; re-run `bt auth login --oauth --org `", - auth_slot_label(&profile_name, &profile) + auth_slot_label(&profile) ) }, )?; eprintln!( "Refreshing OAuth token for {} (api_url: {api_url})", - auth_slot_label(&profile_name, &profile) + auth_slot_label(&profile) ); if let Some(expires_at) = previous_expires_at { let now = current_unix_timestamp(); @@ -1834,7 +1829,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { eprintln!("Cached access token expiry before refresh: unknown"); } - let login_label = auth_slot_label(&profile_name, &profile); + let login_label = auth_slot_label(&profile); let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &client_id, &login_label).await?; save_profile_oauth_access_token(profile_name.as_str(), &refreshed.access_token)?; @@ -2062,7 +2057,7 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result< let profile = store.profiles.get(profile_name).cloned().ok_or_else(|| { anyhow::anyhow!("auth login not found; run `bt auth profiles` to see available logins") })?; - let label = auth_slot_label(profile_name, &profile); + let label = auth_slot_label(&profile); if !force { let term = ui::prompt_term().ok_or_else(|| { @@ -2439,8 +2434,8 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { .collect(); println!("{}", serde_json::to_string(&output)?); } else { - for (name, profile) in &store.profiles { - println!(" {}", auth_slot_label(name, profile)); + for profile in store.profiles.values() { + println!(" {}", auth_slot_label(profile)); } } Ok(()) @@ -4952,6 +4947,86 @@ mod tests { } } + fn auth_source( + prefer_api_key: bool, + cli: Option<&str>, + env: Option<&str>, + oauth: Option<&str>, + api_key: Option<&str>, + ) -> AuthSource { + resolve_auth_source( + prefer_api_key, + cli.map(str::to_string), + || env.map(str::to_string), + || Ok(oauth.map(str::to_string)), + || Ok(api_key.map(str::to_string)), + ) + .expect("resolve auth source") + } + + #[test] + fn auth_source_cli_api_key_wins_over_everything() { + assert_eq!( + auth_source(false, Some("cli"), Some("env"), Some("oauth"), Some("ak")), + AuthSource::CliApiKey("cli".into()) + ); + assert_eq!( + auth_source(true, Some("cli"), Some("env"), Some("oauth"), Some("ak")), + AuthSource::CliApiKey("cli".into()) + ); + } + + #[test] + fn auth_source_default_order_is_oauth_then_api_key_then_env() { + assert_eq!( + auth_source(false, None, Some("env"), Some("oauth"), Some("ak")), + AuthSource::Oauth("oauth".into()) + ); + assert_eq!( + auth_source(false, None, Some("env"), None, Some("ak")), + AuthSource::ApiKey("ak".into()) + ); + assert_eq!( + auth_source(false, None, Some("env"), None, None), + AuthSource::EnvApiKey("env".into()) + ); + assert_eq!( + auth_source(false, None, None, None, None), + AuthSource::None + ); + } + + #[test] + fn auth_source_prefer_api_key_order_is_env_then_api_key_then_oauth() { + assert_eq!( + auth_source(true, None, Some("env"), Some("oauth"), Some("ak")), + AuthSource::EnvApiKey("env".into()) + ); + assert_eq!( + auth_source(true, None, None, Some("oauth"), Some("ak")), + AuthSource::ApiKey("ak".into()) + ); + // No env/stored API key, but OAuth for the org is available: fall back to it. + assert_eq!( + auth_source(true, None, None, Some("oauth"), None), + AuthSource::Oauth("oauth".into()) + ); + assert_eq!(auth_source(true, None, None, None, None), AuthSource::None); + } + + #[test] + fn auth_source_ambiguous_selection_stops_the_ladder() { + let err = resolve_auth_source( + false, + None, + || None, + || bail!("multiple oauth logins"), + || Ok(Some("ak".to_string())), + ) + .expect_err("ambiguous oauth should stop the ladder"); + assert!(err.to_string().contains("multiple oauth logins")); + } + #[tokio::test] async fn persist_post_login_context_clears_stale_project_for_org_only_login() { let _env = TestEnv::new(None, None).await; From fb1ac8c7912729c7d2073f2b6fb8ca087acf3f54 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Fri, 17 Jul 2026 11:26:00 -0700 Subject: [PATCH 06/24] chore(auth): remove client-id flag The field is needed by OAuth but it's unused both bt and the backend Therefore exposing it is useless Now bt_cli is hardcoded as its value --- src/auth.rs | 95 +++++++++------------------------------------------- tests/cli.rs | 9 +++++ 2 files changed, 24 insertions(+), 80 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 7055a346..a5ee496d 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -37,6 +37,7 @@ use crate::{ }; const KEYCHAIN_SERVICE: &str = "com.braintrust.bt.cli"; +const OAUTH_CLIENT_ID: &str = "bt_cli"; const OAUTH_SCOPE: &str = "mcp"; const OAUTH_CALLBACK_TIMEOUT: Duration = Duration::from_secs(300); const OAUTH_REFRESH_SAFETY_WINDOW_SECONDS: u64 = 60; @@ -79,7 +80,6 @@ pub struct AvailableOrg { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RecoverableAuthErrorKind { - OauthClientId, OauthRefreshToken, StoredCredential, } @@ -109,8 +109,7 @@ pub fn is_missing_credential_error(err: &anyhow::Error) -> bool { .is_some_and(|err| { matches!( err.kind, - RecoverableAuthErrorKind::OauthClientId - | RecoverableAuthErrorKind::OauthRefreshToken + RecoverableAuthErrorKind::OauthRefreshToken | RecoverableAuthErrorKind::StoredCredential ) }) @@ -205,8 +204,6 @@ struct AuthProfile { #[serde(default)] org_name: Option, #[serde(default)] - oauth_client_id: Option, - #[serde(default)] oauth_access_expires_at: Option, #[serde(default)] user_name: Option, @@ -313,10 +310,6 @@ struct AuthLoginArgs { #[arg(long)] oauth: bool, - /// OAuth client id (defaults to bt_cli) - #[arg(long, value_name = "CLIENT_ID")] - client_id: Option, - /// Do not try to open a browser automatically #[arg(long)] no_browser: bool, @@ -1094,15 +1087,6 @@ async fn resolve_oauth_profile_auth( store.profiles.get(profile_name).cloned().ok_or_else(|| { anyhow::anyhow!("saved OAuth login not found; run `bt auth profiles`") })?; - let client_id = profile.oauth_client_id.as_deref().ok_or_else(|| { - recoverable_auth_error( - RecoverableAuthErrorKind::OauthClientId, - format!( - "oauth login for '{}' is missing client_id; re-run `bt auth login --oauth --org `", - auth_slot_label(&profile) - ), - ) - })?; let api_url = base .api_url .clone() @@ -1142,8 +1126,7 @@ async fn resolve_oauth_profile_auth( ) })?; let login_label = auth_slot_label(&profile); - let refreshed = - refresh_oauth_access_token(&api_url, &refresh_token, client_id, &login_label).await?; + let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &login_label).await?; save_profile_oauth_access_token(profile_name, &refreshed.access_token)?; let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { @@ -1411,8 +1394,7 @@ fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec Result<()> { .app_url .clone() .unwrap_or_else(|| DEFAULT_APP_URL.to_string()); - let client_id = args - .client_id - .clone() - .unwrap_or_else(default_oauth_client_id); - let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256(); let state = generate_random_token(32)?; let callback_server = bind_oauth_callback_server()?; let redirect_uri = callback_server.redirect_uri(); - let oauth_client = build_oauth_client(&api_url, &client_id, Some(&redirect_uri))?; + let oauth_client = build_oauth_client(&api_url, Some(&redirect_uri))?; let (authorize_url, _) = oauth_client .authorize_url(|| CsrfToken::new(state.clone())) .add_scope(Scope::new(OAUTH_SCOPE.to_string())) @@ -1624,14 +1601,9 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { bail!("oauth state mismatch; please try again"); } - let oauth_tokens = exchange_oauth_authorization_code( - &api_url, - &client_id, - &redirect_uri, - &auth_code, - pkce_verifier, - ) - .await?; + let oauth_tokens = + exchange_oauth_authorization_code(&api_url, &redirect_uri, &auth_code, pkce_verifier) + .await?; let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; let selected_org = select_login_org( login_orgs.clone(), @@ -1649,7 +1621,6 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { &oauth_tokens, selected_api_url.clone(), app_url.clone(), - client_id.clone(), selected_org.as_ref(), )?; let context_update = persist_post_login_context( @@ -1710,7 +1681,6 @@ pub(crate) fn commit_api_key_profile( app_url, org_id: Some(org_id), org_name: Some(org_name), - oauth_client_id: None, oauth_access_expires_at: None, user_name: None, email: None, @@ -1726,7 +1696,6 @@ fn commit_oauth_profile( tokens: &OAuthTokenResponse, api_url: String, app_url: String, - client_id: String, selected_org: Option<&LoginOrgInfo>, ) -> Result<()> { let refresh_token = tokens.refresh_token.as_ref().ok_or_else(|| { @@ -1765,7 +1734,6 @@ fn commit_oauth_profile( app_url: Some(app_url), org_id: Some(org_id), org_name: selected_org.map(|org| org.name.clone()), - oauth_client_id: Some(client_id), oauth_access_expires_at, user_name: jwt_id.name, email: jwt_id.email, @@ -1798,12 +1766,6 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { .api_url .clone() .unwrap_or_else(|| DEFAULT_API_URL.to_string()); - let client_id = profile.oauth_client_id.clone().ok_or_else(|| { - anyhow::anyhow!( - "OAuth login for '{}' is missing client_id; re-run `bt auth login --oauth --org `", - auth_slot_label(&profile) - ) - })?; let previous_expires_at = profile.oauth_access_expires_at; let refresh_token = load_profile_oauth_refresh_token_for_profile(profile_name.as_str(), &profile)?.ok_or_else( @@ -1830,8 +1792,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { } let login_label = auth_slot_label(&profile); - let refreshed = - refresh_oauth_access_token(&api_url, &refresh_token, &client_id, &login_label).await?; + let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &login_label).await?; save_profile_oauth_access_token(profile_name.as_str(), &refreshed.access_token)?; let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { @@ -2664,10 +2625,6 @@ fn resolve_profile_api_url( ) } -fn default_oauth_client_id() -> String { - "bt_cli".to_string() -} - fn generate_random_token(num_bytes: usize) -> Result { let mut bytes = vec![0u8; num_bytes]; getrandom::fill(&mut bytes) @@ -3024,7 +2981,6 @@ fn is_ssh_session() -> bool { async fn exchange_oauth_authorization_code( api_url: &str, - client_id: &str, redirect_uri: &str, code: &str, code_verifier: PkceCodeVerifier, @@ -3040,7 +2996,7 @@ async fn exchange_oauth_authorization_code( api_url, &[ ("grant_type", "authorization_code"), - ("client_id", client_id), + ("client_id", OAUTH_CLIENT_ID), ("code", code), ("redirect_uri", redirect_uri), ("code_verifier", code_verifier.secret()), @@ -3077,7 +3033,6 @@ fn map_refresh_oauth_error( async fn refresh_oauth_access_token( api_url: &str, refresh_token: &str, - client_id: &str, auth_login: &str, ) -> Result { let http_client = build_http_client_from_builder( @@ -3091,7 +3046,7 @@ async fn refresh_oauth_access_token( .post(&token_url) .form(&[ ("grant_type", "refresh_token"), - ("client_id", client_id), + ("client_id", OAUTH_CLIENT_ID), ("refresh_token", refresh_token), ]) .send() @@ -3133,18 +3088,14 @@ async fn request_oauth_token( .context("failed to parse oauth token response") } -fn build_oauth_client( - api_url: &str, - client_id: &str, - redirect_uri: Option<&str>, -) -> Result { +fn build_oauth_client(api_url: &str, redirect_uri: Option<&str>) -> Result { let api_url = api_url.trim_end_matches('/'); let auth_url = AuthUrl::new(format!("{api_url}/oauth/authorize")) .context("failed to construct oauth authorize URL")?; let token_url = TokenUrl::new(format!("{api_url}/oauth/token")) .context("failed to construct oauth token URL")?; let client = BasicClient::new( - ClientId::new(client_id.to_string()), + ClientId::new(OAUTH_CLIENT_ID.to_string()), None, auth_url, Some(token_url), @@ -4282,7 +4233,6 @@ mod tests { app_url: Some((*app_url).to_string()), org_id: None, org_name: Some((*org_name).to_string()), - oauth_client_id: None, oauth_access_expires_at: None, user_name: None, email: None, @@ -4486,7 +4436,6 @@ mod tests { app_url: Some("https://www.example.test".to_string()), org_id: Some(org_id.to_string()), org_name: Some(org_name.to_string()), - oauth_client_id: Some("bt_cli".to_string()), oauth_access_expires_at: Some(current_unix_timestamp() + 3600), user_name: Some("Test User".to_string()), email: Some("user@example.test".to_string()), @@ -4656,7 +4605,6 @@ mod tests { api_url: Some("https://api.example.com".to_string()), app_url: Some("https://www.example.com".to_string()), org_name: Some("Example Org".to_string()), - oauth_client_id: None, oauth_access_expires_at: None, ..Default::default() }, @@ -4680,7 +4628,6 @@ mod tests { org_id: Some("org_fake".to_string()), org_name: Some("test-org".to_string()), email: Some("user@example.test".to_string()), - oauth_client_id: Some("bt_cli_work".to_string()), ..Default::default() }, ); @@ -4690,7 +4637,6 @@ mod tests { let profile = migrated.profiles.get(&key).expect("migrated profile"); assert_eq!(profile.legacy_secret_key.as_deref(), Some("work")); - assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_work")); } #[test] @@ -4702,7 +4648,6 @@ mod tests { auth_kind: AuthKind::Oauth, org_name: None, email: Some("user@example.test".to_string()), - oauth_client_id: Some("bt_cli_legacy".to_string()), ..Default::default() }, ); @@ -4718,7 +4663,6 @@ mod tests { profile.legacy_secret_key.as_deref(), Some("legacy-cross-org") ); - assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_legacy")); } #[test] @@ -4763,7 +4707,6 @@ mod tests { org_id: Some("org_fake".to_string()), org_name: Some("test-org".to_string()), email: Some("user@example.test".to_string()), - oauth_client_id: Some("bt_cli_legacy".to_string()), ..Default::default() }, ); @@ -4781,7 +4724,6 @@ mod tests { .get(&slot_key) .expect("canonical OAuth slot"); assert_eq!(profile.legacy_secret_key.as_deref(), Some("legacy-login")); - assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_legacy")); } let _ = fs::remove_dir_all(&dir); @@ -4796,7 +4738,6 @@ mod tests { AuthProfile { auth_kind: AuthKind::Oauth, org_name: Some("test-org".to_string()), - oauth_client_id: Some("bt_cli_legacy".to_string()), ..Default::default() }, ); @@ -4855,8 +4796,7 @@ mod tests { #[test] fn migrate_auth_store_dedupes_oauth_slots_by_latest_expiry() { let mut store = AuthStore::default(); - for (name, expires_at, client_id) in [("old", 10, "bt_cli_old"), ("new", 20, "bt_cli_new")] - { + for (name, expires_at) in [("old", 10), ("new", 20)] { store.profiles.insert( name.to_string(), AuthProfile { @@ -4864,7 +4804,6 @@ mod tests { org_id: Some("org_fake".to_string()), org_name: Some("test-org".to_string()), email: Some("user@example.test".to_string()), - oauth_client_id: Some(client_id.to_string()), oauth_access_expires_at: Some(expires_at), ..Default::default() }, @@ -4875,7 +4814,6 @@ mod tests { let key = oauth_slot_key("org_fake", "user@example.test"); assert_eq!(migrated.profiles.len(), 1); let profile = migrated.profiles.get(&key).expect("migrated profile"); - assert_eq!(profile.oauth_client_id.as_deref(), Some("bt_cli_new")); assert_eq!(profile.legacy_secret_key.as_deref(), Some("new")); } @@ -4990,10 +4928,7 @@ mod tests { auth_source(false, None, Some("env"), None, None), AuthSource::EnvApiKey("env".into()) ); - assert_eq!( - auth_source(false, None, None, None, None), - AuthSource::None - ); + assert_eq!(auth_source(false, None, None, None, None), AuthSource::None); } #[test] diff --git a/tests/cli.rs b/tests/cli.rs index f40c0682..084add7c 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -49,6 +49,15 @@ fn write_auth_store(config_home: &Path, profiles: &[(&str, &str)]) { fs::write(auth_dir.join("auth.json"), body).expect("write auth store"); } +#[test] +fn auth_login_does_not_expose_client_id() { + bt_command() + .args(["auth", "login", "--help"]) + .assert() + .success() + .stdout(predicate::str::contains("--client-id").not()); +} + #[test] fn global_quiet_flag_still_parses_for_other_commands() { bt_command().args(["status", "--quiet"]).assert().success(); From dedf063518f1c6ca070908e8ed7c5b9a6597663b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Fri, 17 Jul 2026 12:15:15 -0700 Subject: [PATCH 07/24] chore: rename bt auth profiles to bt auth logins --- README.md | 2 +- src/auth.rs | 31 ++++++++++++++++--------------- src/main.rs | 2 +- tests/functions.rs | 24 ++++++++++++------------ 4 files changed, 30 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 237e27d8..172785ea 100644 --- a/README.md +++ b/README.md @@ -323,7 +323,7 @@ Local version and pagination-key conversion helpers: - You can pass `--no-browser` to print the URL without auto-opening. - On remote/SSH hosts, paste the final callback URL from your local browser if localhost callback cannot be delivered. - List saved auth logins: - - `bt auth profiles` + - `bt auth logins` - Log out: - `bt auth logout --org myorg` - `bt auth logout --org myorg --api-key-hint sk-****abcde` diff --git a/src/auth.rs b/src/auth.rs index a5ee496d..c0dce814 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -280,7 +280,7 @@ struct OAuthErrorResponse { #[command(after_help = "\ Examples: bt auth login - bt auth profiles + bt auth logins bt auth refresh --org acme bt auth logout --org acme ")] @@ -296,13 +296,13 @@ enum AuthCommand { /// Force-refresh OAuth access token for the selected org Refresh, /// List saved auth logins and check connection status - Profiles(AuthProfilesArgs), + Logins(AuthLoginsArgs), /// Log out by removing a saved auth login Logout(AuthLogoutArgs), } #[derive(Debug, Clone, Args)] -struct AuthProfilesArgs {} +struct AuthLoginsArgs {} #[derive(Debug, Clone, Args)] struct AuthLoginArgs { @@ -335,7 +335,7 @@ pub async fn run(base: BaseArgs, args: AuthArgs) -> Result<()> { match args.command { AuthCommand::Login(login_args) => run_login_set(&base, login_args).await, AuthCommand::Refresh => run_login_refresh(&base).await, - AuthCommand::Profiles(profile_args) => run_profiles(&base, profile_args).await, + AuthCommand::Logins(logins_args) => run_logins(&base, logins_args).await, AuthCommand::Logout(logout_args) => run_login_logout(base, logout_args), } } @@ -886,7 +886,7 @@ fn resolve_api_key_profile_auth( .profiles .get(profile_name) .cloned() - .ok_or_else(|| anyhow::anyhow!("saved auth login not found; run `bt auth profiles`"))?; + .ok_or_else(|| anyhow::anyhow!("saved auth login not found; run `bt auth logins`"))?; let api_key = load_profile_secret_with_legacy( profile_name, profile.legacy_secret_key.as_deref(), @@ -1043,7 +1043,7 @@ fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Resu .is_none_or(|org| org.trim().is_empty()); if login_org_id.trim().is_empty() && !is_cross_org { // The SDK's OAuth compatibility path does not always return org - // metadata. `bt auth profiles` performs the same reconciliation after + // metadata. `bt auth logins` performs the same reconciliation after // its explicit credential verification request. return Ok(()); } @@ -1083,10 +1083,11 @@ async fn resolve_oauth_profile_auth( cfg_org: &Option, profile_name: &str, ) -> Result { - let profile = - store.profiles.get(profile_name).cloned().ok_or_else(|| { - anyhow::anyhow!("saved OAuth login not found; run `bt auth profiles`") - })?; + let profile = store + .profiles + .get(profile_name) + .cloned() + .ok_or_else(|| anyhow::anyhow!("saved OAuth login not found; run `bt auth logins`"))?; let api_url = base .api_url .clone() @@ -1751,7 +1752,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { let profile_name = select_oauth_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? .ok_or_else(|| { anyhow::anyhow!( - "no OAuth login selected; pass --org or run `bt auth profiles` to see available logins" + "no OAuth login selected; pass --org or run `bt auth logins` to see available logins" ) })?; let profile = store @@ -1759,7 +1760,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { .get(profile_name.as_str()) .cloned() .ok_or_else(|| { - anyhow::anyhow!("OAuth login not found; run `bt auth profiles` to see available logins") + anyhow::anyhow!("OAuth login not found; run `bt auth logins` to see available logins") })?; let api_url = profile @@ -1954,7 +1955,7 @@ fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) -> Ok(()) } -async fn run_profiles(base: &BaseArgs, _args: AuthProfilesArgs) -> Result<()> { +async fn run_logins(base: &BaseArgs, _args: AuthLoginsArgs) -> Result<()> { let mut store = load_auth_store()?; if store.profiles.is_empty() { return emit_result(base.json, serde_json::json!([]), || { @@ -2016,7 +2017,7 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result< let mut store = load_auth_store()?; let profile = store.profiles.get(profile_name).cloned().ok_or_else(|| { - anyhow::anyhow!("auth login not found; run `bt auth profiles` to see available logins") + anyhow::anyhow!("auth login not found; run `bt auth logins` to see available logins") })?; let label = auth_slot_label(&profile); @@ -2095,7 +2096,7 @@ fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { } let profile_name = match candidates.len() { - 0 => bail!("no matching auth login found; run `bt auth profiles` to see available logins"), + 0 => bail!("no matching auth login found; run `bt auth logins` to see available logins"), 1 => candidates[0].to_string(), _ if ui::can_prompt() => { select_profile_from_store("Select auth login to log out", &candidates, org, &store)? diff --git a/src/main.rs b/src/main.rs index 009cc335..7d0082a5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -501,7 +501,7 @@ fn looks_like_user_error(err: &anyhow::Error) -> bool { fn print_error(err: &anyhow::Error, code: ExitCode, missing_credential: bool) { eprintln!("error: {err}"); if code == ExitCode::Auth && !missing_credential { - eprintln!("Your credentials may be expired or invalid. For OAuth login, try `bt auth refresh --org `; if refresh fails, re-run `bt auth login --oauth --org `. Run `bt auth profiles` and `bt status` to inspect auth status."); + eprintln!("Your credentials may be expired or invalid. For OAuth login, try `bt auth refresh --org `; if refresh fails, re-run `bt auth login --oauth --org `. Run `bt auth logins` and `bt status` to inspect auth status."); } if code == ExitCode::Error { eprintln!("If this seems like a bug, file an issue at https://github.com/braintrustdata/bt/issues/new and include `bt --version`, `bt status --json`, and the command you ran."); diff --git a/tests/functions.rs b/tests/functions.rs index 84a9973f..e1dd66cc 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -190,8 +190,8 @@ fn sanitized_env_keys() -> &'static [&'static str] { ] } -fn auth_profiles_command(cwd: &Path, config_dir: &Path) -> Command { - auth_sub_command(cwd, config_dir, &["profiles"]) +fn auth_logins_command(cwd: &Path, config_dir: &Path) -> Command { + auth_sub_command(cwd, config_dir, &["logins"]) } #[derive(Debug, Clone)] @@ -704,14 +704,14 @@ fn functions_help_lists_push_and_pull() { } #[test] -fn auth_profiles_ignores_api_key_env_override() { +fn auth_logins_ignores_api_key_env_override() { let cwd = tempdir().expect("create temp cwd"); let config_dir = tempdir().expect("create temp config dir"); - let output = auth_profiles_command(cwd.path(), config_dir.path()) + let output = auth_logins_command(cwd.path(), config_dir.path()) .env("BRAINTRUST_API_KEY", "test-key") .output() - .expect("run bt auth profiles with api key env"); + .expect("run bt auth logins with api key env"); assert!(output.status.success()); let stdout = String::from_utf8_lossy(&output.stdout); @@ -720,15 +720,15 @@ fn auth_profiles_ignores_api_key_env_override() { } #[test] -fn auth_profiles_ignores_api_key_from_dotenv() { +fn auth_logins_ignores_api_key_from_dotenv() { let cwd = tempdir().expect("create temp cwd"); let config_dir = tempdir().expect("create temp config dir"); fs::write(cwd.path().join(".env"), "BRAINTRUST_API_KEY=test-key\n").expect("write .env"); - let output = auth_profiles_command(cwd.path(), config_dir.path()) + let output = auth_logins_command(cwd.path(), config_dir.path()) .env_remove("BRAINTRUST_API_KEY") .output() - .expect("run bt auth profiles with dotenv api key"); + .expect("run bt auth logins with dotenv api key"); assert!(output.status.success()); let stdout = String::from_utf8_lossy(&output.stdout); @@ -737,14 +737,14 @@ fn auth_profiles_ignores_api_key_from_dotenv() { } #[test] -fn auth_profiles_json_with_no_profiles_emits_empty_array() { +fn auth_logins_json_with_no_profiles_emits_empty_array() { let cwd = tempdir().expect("create temp cwd"); let config_dir = tempdir().expect("create temp config dir"); - let output = auth_profiles_command(cwd.path(), config_dir.path()) + let output = auth_logins_command(cwd.path(), config_dir.path()) .arg("--json") .output() - .expect("run bt auth profiles --json"); + .expect("run bt auth logins --json"); assert!(output.status.success()); let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); @@ -764,7 +764,7 @@ fn seed_api_key_profile(config_dir: &Path) { fn auth_sub_command(cwd: &Path, config_dir: &Path, sub: &[&str]) -> Command { // Shared builder for `bt auth ` so each auth subcommand inherits the - // same isolated config dir / env scrubbing as `auth profiles`. + // same isolated config dir / env scrubbing as `auth logins`. let mut cmd = Command::new(bt_binary_path()); cmd.arg("auth") .args(sub) From 780a30b4ffcc45fd64e77971c130d290038bdff4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Wed, 22 Jul 2026 14:13:34 -0700 Subject: [PATCH 08/24] =?UTF-8?q?Authentication=20=E2=80=94=20src/auth.rs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Preserves separate OAuth and API-key login identities. - Keeps multiple API keys for the same org distinct using key hints. - Implements the intended precedence: 1. Explicit --api-key 2. --prefer-api-key 3. OAuth 4. BRAINTRUST_API_KEY 5. Stored API key - Validates explicit/environment API keys against the requested org. - Does not fall back to OAuth when a selected key is invalid or belongs to another org. - Rejects API-key authentication in cross-org mode. - Pins the exact API-key login selected by bt switch or bt init. - Adds bt auth login --global/--local. - Adds filtering to bt auth logins --org/--prefer-api-key. - Changes bare bt auth logout to select from all saved logins instead of deleting the currently active one. - Adds bt auth logout --oauth and retains --api-key-hint. Config handling — src/config/mod.rs - Makes global/local merging org-safe: - A project is never inherited from another org. - project and project_id remain coupled. - Local project-without-org does not inherit the global org. - Stores cross-org as: ```json { "org": "" } ``` - Ignores the legacy profile field. - Preserves unknown config keys during updates. - Only treats .bt/config.json as a local config; a bare .bt directory is not one. - Adds separate filesystem discovery rules for bt init. bt switch — src/switch.rs - Starts with a saved-login picker: - OAuth entries collapse into org choices. - API keys remain separate and display their hints. - Cross-org OAuth is selectable. - Automatically selects a sole login/project. - Prompts for project when multiple projects exist. - Pins a selected API-key slot exactly. - Implements global/local scope rules: - --global does not search for local config. - --local requires an existing .bt/config.json. - Interactive ambiguity opens a scope picker. - Non-interactive ambiguity errors before network work. - Preserves unknown config fields. - Fixes the scope-picker ANSI corruption you reported by passing plain labels to Dialoguer. bt init — src/init.rs - Adds: - --here - -f/--force - Searches upward for .bt, .git, home, or filesystem root using the specified boundary rules. - Resolves destination errors before authentication. - Uses the same saved-login/project selection behavior as switch. - Excludes cross-org from the picker because init requires a project. - Writes org, project, and project_id. - Reports permission/write failures with the affected path. - Now exposes the real destination error instead of only saying: ```text could not resolve `bt init` destination ``` Cross-org CLI/status — src/args.rs, src/main.rs, src/status.rs - Normalizes these to canonical cross-org: ```sh --org cross-org --org "" ``` - Displays cross-org as cross-org in human and JSON status output. - Prevents stale projects from being combined with cross-org context. - Makes --prefer-api-key fail actionably from cross-org context. --- README.md | 44 +- src/args.rs | 47 +- src/auth.rs | 1434 +++++++++++++++++++++++-------------- src/config/mod.rs | 514 +++++++++---- src/datasets/pipeline.rs | 21 +- src/datasets/snapshots.rs | 4 +- src/eval.rs | 12 +- src/functions/push.rs | 43 +- src/init.rs | 135 ++-- src/main.rs | 15 +- src/setup/mod.rs | 21 +- src/status.rs | 350 ++++----- src/switch.rs | 447 ++++-------- src/traces.rs | 53 +- src/traces/waterfall.rs | 6 +- src/ui/mod.rs | 1 + src/ui/select.rs | 51 +- src/utils/mod.rs | 2 + src/utils/shell.rs | 10 + tests/functions.rs | 30 + 20 files changed, 1813 insertions(+), 1427 deletions(-) create mode 100644 src/utils/shell.rs diff --git a/README.md b/README.md index 172785ea..f7bbc605 100644 --- a/README.md +++ b/README.md @@ -316,7 +316,8 @@ Local version and pagination-key conversion helpers: - `bt auth login` - First prompt chooses: `OAuth (browser)` (default) or `API key`. - OAuth can be saved for an org or in cross-org mode. API-key logins are saved per org; if a key can access multiple orgs, `bt` uses a searchable org picker. - - After login, `bt` updates the active org context immediately. If `--project` is set, it also switches that project; otherwise it clears any stale default project for the new login. + - After login, `bt` updates the active org context immediately. If `--project` is set, it validates and saves that project's name and ID. Without `--project`, a same-org login preserves the existing project; changing orgs clears stale project context. + - Use `--global` or `--local` to choose the config scope. Without either flag, an existing local config causes an interactive scope picker (default: local); non-interactive runs must pass a scope. `--local` never creates `.bt`. - `bt` confirms the resolved API URL before saving. - Login with OAuth (browser-based, stores refresh token in secure credential store): - `bt auth login --oauth --org myorg` @@ -324,10 +325,14 @@ Local version and pagination-key conversion helpers: - On remote/SSH hosts, paste the final callback URL from your local browser if localhost callback cannot be delivered. - List saved auth logins: - `bt auth logins` + - `bt auth logins --org test-org` (matches stored org name or ID) + - `bt auth logins --prefer-api-key` (API-key logins only) + - Both filters can be combined. - Log out: - - `bt auth logout --org myorg` - - `bt auth logout --org myorg --api-key-hint sk-****abcde` - - `bt auth logout --force` (skip confirmation) + - `bt auth logout` — choose from all saved logins interactively + - `bt auth logout --org test-org --oauth` — filter to the org's OAuth login + - `bt auth logout --org test-org --api-key-hint sk-****abcde` — select an API-key login + - `bt auth logout --force` (skip confirmation after selecting a login) - Show current auth context: - `bt status` - Force-refresh OAuth access token for debugging: @@ -338,20 +343,39 @@ Auth resolution order for commands is: 1. Explicit `--api-key sk-...` 2. `--prefer-api-key` / `BRAINTRUST_PREFER_API_KEY` (uses `BRAINTRUST_API_KEY` first, then a stored API key for the selected org, then falls back to OAuth) 3. Stored OAuth login for the selected org (or cross-org OAuth when selected) -4. Stored API key login for the selected org when no OAuth login is available -5. `BRAINTRUST_API_KEY` +4. `BRAINTRUST_API_KEY` +5. Stored API key login for the selected org + +`--prefer-api-key` without `--org` targets the org shown by `bt status`. It cannot be used from cross-org context; pass a concrete `--org`. Once a key is selected, an invalid key or a key belonging to another requested org is an error and does not fall back to OAuth. Multiple keys in one org remain separate and are shown with key hints. On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Service daemon. On macOS, it uses the `security` keychain utility. If a secure store is unavailable, `bt` falls back to a plaintext secrets file with `0600` permissions. +## `bt init` + +`bt init` creates a project-local `.bt/config.json`. It walks upward to the first `.bt`, `.git`, home, or filesystem-root boundary. A git marker (directory or file) selects that repository root; reaching home/root or an existing `.bt` is an error. + +- `bt init --org test-org --project test-project` — initialize the containing repository +- `bt init --here` — create in the current directory without walking (including at home or `/`) +- `bt init --force` — overwrite an existing discovered `.bt/config.json`; it does not change discovery + +The saved context includes `org`, `project`, and `project_id`. Cross-org is excluded from the init picker because init always requires a project. + ## `bt switch` Interactively switch org and project context: -- `bt switch` — interactive picker for org and project -- `bt switch myproject` — switch to a project by name -- `bt switch myorg/myproject` — switch to a specific org and project +- `bt switch` — first choose a saved OAuth-backed org or a specific API-key login, then choose a project +- `bt switch myproject` — switch the current org to a project by name +- `bt switch test-org/test-project` — switch to a specific org and project +- `bt switch --org cross-org` — select cross-org OAuth and clear project context - `bt switch --global` — persist to global config (`~/.config/bt/config.json`) -- `bt switch --local` — persist to local config (`.bt/config.json`) +- `bt switch --local` — update an existing local config (`.bt/config.json`); it never creates one + +A sole login/project is selected automatically. Multiple API keys in one org remain separate picker entries with hints; OAuth accounts collapse to an org choice, so run `bt auth login` again to change OAuth accounts within that org. With an existing local config and no scope flag, interactive mode asks for global/local (default: local); non-interactive mode requires `--global` or `--local`. + +## Config context merging + +Global config is `~/.config/bt/config.json`; local config is the first discovered `.bt/config.json`. Local context wins, but a local org inherits the global project only when both configs select the same org. A local project with no org never inherits a global org. Cross-org is stored as `"org": ""`, treated as a distinct org, and rendered as `cross-org` by `bt status`. Legacy `profile` fields are ignored; unknown extra keys are preserved when context is updated. ## `bt status` diff --git a/src/args.rs b/src/args.rs index 930cfe9d..1770da97 100644 --- a/src/args.rs +++ b/src/args.rs @@ -10,7 +10,7 @@ pub enum ArgValueSource { EnvVariable, } -#[derive(Debug, Clone, Args)] +#[derive(Debug, Clone, Default, Args)] pub struct BaseArgs { /// Output as JSON #[arg(long, global = true)] @@ -39,7 +39,7 @@ pub struct BaseArgs { pub no_input: bool, /// Override active org (or via BRAINTRUST_ORG_NAME) - #[arg(short = 'o', long = "org", env = "BRAINTRUST_ORG_NAME", global = true)] + #[arg(short = 'o', long = "org", env = "BRAINTRUST_ORG_NAME", global = true, value_parser = parse_org_name)] pub org_name: Option, #[arg(skip)] @@ -65,6 +65,10 @@ pub struct BaseArgs { #[arg(skip)] pub api_key_source: Option, + /// Exact auth slot selected internally by switch/init. + #[arg(skip)] + pub pinned_auth_slot: Option, + /// Prefer API key credentials for the selected org when available. #[arg(long = "prefer-api-key", env = "BRAINTRUST_PREFER_API_KEY", global = true, value_parser = clap::builder::BoolishValueParser::new(), default_value_t = false)] pub prefer_api_key: bool, @@ -115,6 +119,19 @@ pub struct CLIArgs { pub base: BaseArgs, } +fn parse_org_name(value: &str) -> Result { + Ok(crate::config::normalize_org(value).to_string()) +} + +pub(crate) fn custom_api_without_app_url(api_url: Option<&str>, app_url: Option<&str>) -> bool { + app_url.is_none_or(|url| url.trim().is_empty()) + && api_url.is_some_and(|url| { + !url.trim() + .trim_end_matches('/') + .eq_ignore_ascii_case(DEFAULT_API_URL.trim_end_matches('/')) + }) +} + impl BaseArgs { pub fn ca_cert(&self) -> Option<&Path> { self.ca_cert.as_deref() @@ -124,3 +141,29 @@ impl BaseArgs { self.verbose && self.verbose_source.is_some() } } + +#[cfg(test)] +mod tests { + use super::{custom_api_without_app_url, parse_org_name, DEFAULT_API_URL}; + + #[test] + fn org_normalization() { + for (input, expected) in [ + ("cross-org", ""), + (" ", ""), + (" test-org ", "test-org"), + (" org_test_123 ", "org_test_123"), + ] { + assert_eq!(parse_org_name(input).unwrap(), expected); + } + assert!(custom_api_without_app_url( + Some("https://api.example.test"), + None + )); + assert!(!custom_api_without_app_url(Some(DEFAULT_API_URL), None)); + assert!(!custom_api_without_app_url( + Some("https://api.example.test"), + Some("https://app.example.test") + )); + } +} diff --git a/src/auth.rs b/src/auth.rs index c0dce814..a300a550 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -33,7 +33,8 @@ use crate::{ config, http::{build_http_client, build_http_client_from_builder, ApiClient}, projects::api, - switch, ui, + ui, + utils::shell_quote_arg, }; const KEYCHAIN_SERVICE: &str = "com.braintrust.bt.cli"; @@ -139,36 +140,19 @@ pub async fn list_available_orgs(base: &BaseArgs) -> Result> { .context("login state missing API key")?, }; - let mut orgs = fetch_login_orgs(&api_key, &app_url).await?; - orgs.sort_by(|a, b| { - a.name - .to_ascii_lowercase() - .cmp(&b.name.to_ascii_lowercase()) - .then_with(|| a.name.cmp(&b.name)) - }); - - Ok(orgs - .into_iter() - .map(|org| AvailableOrg { - id: org.id, - name: org.name, - api_url: org.api_url, - }) - .collect()) + available_orgs(&api_key, &app_url).await } pub(crate) async fn list_available_orgs_for_api_key( api_key: &str, app_url: &str, ) -> Result> { - let mut orgs = fetch_login_orgs(api_key, app_url).await?; - orgs.sort_by(|a, b| { - a.name - .to_ascii_lowercase() - .cmp(&b.name.to_ascii_lowercase()) - .then_with(|| a.name.cmp(&b.name)) - }); + available_orgs(api_key, app_url).await +} +async fn available_orgs(api_key: &str, app_url: &str) -> Result> { + let mut orgs = fetch_login_orgs(api_key, app_url).await?; + sort_login_orgs(&mut orgs); Ok(orgs .into_iter() .map(|org| AvailableOrg { @@ -279,10 +263,12 @@ struct OAuthErrorResponse { #[derive(Debug, Clone, Args)] #[command(after_help = "\ Examples: - bt auth login - bt auth logins - bt auth refresh --org acme - bt auth logout --org acme + bt auth login --global + bt auth login --oauth --org test-org --local + bt auth logins --org test-org --prefer-api-key + bt auth refresh --org test-org + bt auth logout + bt auth logout --org test-org --oauth ")] pub struct AuthArgs { #[command(subcommand)] @@ -313,10 +299,17 @@ struct AuthLoginArgs { /// Do not try to open a browser automatically #[arg(long)] no_browser: bool, + + #[command(flatten)] + scope: config::ScopeArgs, } #[derive(Debug, Clone, Args)] struct AuthLogoutArgs { + /// Only consider OAuth logins + #[arg(long, conflicts_with = "api_key_hint")] + oauth: bool, + /// API key hint to log out of when multiple API keys exist for an org #[arg(long = "api-key-hint", value_name = "HINT")] api_key_hint: Option, @@ -333,7 +326,10 @@ struct PostLoginContextUpdate { pub async fn run(base: BaseArgs, args: AuthArgs) -> Result<()> { match args.command { - AuthCommand::Login(login_args) => run_login_set(&base, login_args).await, + AuthCommand::Login(login_args) => { + login_args.scope.preflight(ui::can_prompt())?; + run_login_set(&base, login_args).await + } AuthCommand::Refresh => run_login_refresh(&base).await, AuthCommand::Logins(logins_args) => run_logins(&base, logins_args).await, AuthCommand::Logout(logout_args) => run_login_logout(base, logout_args), @@ -757,16 +753,16 @@ fn config_auth_context(base: &BaseArgs) -> Option { } fn config_auth_context_from_config(base: &BaseArgs, cfg: &crate::config::Config) -> Option { - if crate::config::trimmed_option(base.org_name.as_deref()).is_none() { - crate::config::trimmed_option(cfg.org.as_deref()).map(str::to_string) + if crate::config::org_option(base.org_name.as_deref()).is_none() { + crate::config::org_option(cfg.org.as_deref()).map(str::to_string) } else { None } } fn effective_org_name<'a>(base: &'a BaseArgs, cfg_org: &'a Option) -> Option<&'a str> { - crate::config::trimmed_option(base.org_name.as_deref()) - .or_else(|| crate::config::trimmed_option(cfg_org.as_deref())) + crate::config::org_option(base.org_name.as_deref()) + .or_else(|| crate::config::org_option(cfg_org.as_deref())) } /// The auth source selected by the precedence ladder, before any live @@ -787,8 +783,8 @@ enum AuthSource { /// 1. explicit `--api-key` /// 2. `--prefer-api-key`: `BRAINTRUST_API_KEY` → stored API key → OAuth fallback /// 3. stored OAuth login for the selected org -/// 4. stored API key login for the selected org -/// 5. `BRAINTRUST_API_KEY` +/// 4. `BRAINTRUST_API_KEY` +/// 5. stored API key login for the selected org /// /// The slot selectors return `Ok(None)` when no candidate matches (the ladder /// continues) and may return `Err` for an ambiguous selection that neither @@ -820,12 +816,12 @@ fn resolve_auth_source( if let Some(slot) = select_oauth()? { return Ok(AuthSource::Oauth(slot)); } - if let Some(slot) = select_api_key()? { - return Ok(AuthSource::ApiKey(slot)); - } if let Some(api_key) = env_api_key() { return Ok(AuthSource::EnvApiKey(api_key)); } + if let Some(slot) = select_api_key()? { + return Ok(AuthSource::ApiKey(slot)); + } Ok(AuthSource::None) } @@ -834,32 +830,33 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { let cfg_org = config_auth_context(base); let can_prompt = ui::can_prompt(); + let effective_org = effective_org_name(base, &cfg_org); + reject_cross_org_api_key_preference(base.prefer_api_key, effective_org, &store)?; + + if let Some(slot) = base.pinned_auth_slot.clone() { + return resolve_saved_auth_slot(base, &mut store, &None, &slot).await; + } + let source = resolve_auth_source( base.prefer_api_key, resolve_cli_api_key_override(base), || resolve_env_api_key(base), - || select_oauth_profile_for_auth(base, &store, &cfg_org, can_prompt), - || select_api_key_profile_for_auth(base, &store, &cfg_org, can_prompt), + || select_profile_for_auth(base, &store, &cfg_org, AuthKind::Oauth, can_prompt), + || select_profile_for_auth(base, &store, &cfg_org, AuthKind::ApiKey, can_prompt), )?; match source { - AuthSource::CliApiKey(api_key) | AuthSource::EnvApiKey(api_key) => Ok(ResolvedAuth { - api_key: Some(api_key), - api_url: base.api_url.clone(), - app_url: base.app_url.clone(), - org_name: effective_org_name(base, &cfg_org).map(str::to_string), - is_oauth: false, - slot_key: None, - }), - AuthSource::Oauth(slot) => { - resolve_oauth_profile_auth(base, &mut store, &cfg_org, &slot).await + AuthSource::CliApiKey(api_key) | AuthSource::EnvApiKey(api_key) => { + resolve_ad_hoc_api_key_auth(base, &cfg_org, api_key).await + } + AuthSource::Oauth(slot) | AuthSource::ApiKey(slot) => { + resolve_saved_auth_slot(base, &mut store, &cfg_org, &slot).await } - AuthSource::ApiKey(slot) => resolve_api_key_profile_auth(base, &mut store, &cfg_org, &slot), AuthSource::None => { if base.prefer_api_key { bail!("--prefer-api-key requires an API key or OAuth login for the selected org"); } - if effective_org_name(base, &cfg_org).is_none() { + if effective_org.is_none() { if let Some(err) = missing_org_for_stored_logins_error(&store) { return Err(err); } @@ -868,7 +865,7 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { api_key: None, api_url: base.api_url.clone(), app_url: base.app_url.clone(), - org_name: effective_org_name(base, &cfg_org).map(str::to_string), + org_name: effective_org.map(str::to_string), is_oauth: false, slot_key: None, }) @@ -876,6 +873,72 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { } } +async fn resolve_ad_hoc_api_key_auth( + base: &BaseArgs, + cfg_org: &Option, + api_key: String, +) -> Result { + let requested_org = effective_org_name(base, cfg_org); + if requested_org == Some("") { + bail!("API keys require a concrete org; rerun with --org "); + } + + let mut resolved_org = requested_org.map(str::to_string); + let mut resolved_api_url = base.api_url.clone(); + if let Some(requested_org) = requested_org { + if crate::args::custom_api_without_app_url(base.api_url.as_deref(), base.app_url.as_deref()) + { + bail!("API key organization validation with a custom API URL requires --app-url or BRAINTRUST_APP_URL"); + } + let app_url = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let orgs = fetch_login_orgs(&api_key, app_url).await.map_err(|err| { + if is_unauthorized_auth_error(&err) { + anyhow::anyhow!("API key is not valid") + } else { + err.context("failed to validate API key organization membership") + } + })?; + let selected_org = find_login_org(&orgs, requested_org).ok_or_else(|| { + let available = login_org_names(&orgs); + anyhow::anyhow!( + "API key does not belong to requested org '{requested_org}'. Available orgs for this key: {available}" + ) + })?; + resolved_org = Some(selected_org.name.clone()); + resolved_api_url = resolved_api_url.or_else(|| selected_org.api_url.clone()); + } + + Ok(ResolvedAuth { + api_key: Some(api_key), + api_url: resolved_api_url, + app_url: base.app_url.clone(), + org_name: resolved_org, + is_oauth: false, + slot_key: None, + }) +} + +async fn resolve_saved_auth_slot( + base: &BaseArgs, + store: &mut AuthStore, + cfg_org: &Option, + slot: &str, +) -> Result { + let kind = store + .profiles + .get(slot) + .map(|profile| profile.auth_kind) + .ok_or_else(|| { + anyhow::anyhow!( + "saved auth login not found; run `bt auth logins` to see available logins" + ) + })?; + match kind { + AuthKind::ApiKey => resolve_api_key_profile_auth(base, store, cfg_org, slot), + AuthKind::Oauth => resolve_oauth_profile_auth(base, store, cfg_org, slot).await, + } +} + fn resolve_api_key_profile_auth( base: &BaseArgs, store: &mut AuthStore, @@ -887,6 +950,15 @@ fn resolve_api_key_profile_auth( .get(profile_name) .cloned() .ok_or_else(|| anyhow::anyhow!("saved auth login not found; run `bt auth logins`"))?; + if let Some(requested_org) = effective_org_name(base, cfg_org) { + if !profile_matches_org_identifier(&profile, requested_org) { + bail!( + "stored API key for '{}' does not belong to requested org '{requested_org}'", + profile_org_label(&profile) + ); + } + } + let api_key = load_profile_secret_with_legacy( profile_name, profile.legacy_secret_key.as_deref(), @@ -1121,13 +1193,13 @@ async fn resolve_oauth_profile_auth( recoverable_auth_error( RecoverableAuthErrorKind::OauthRefreshToken, format!( - "oauth refresh token missing for '{}'; re-run `bt auth login --oauth --org `", - auth_slot_label(&profile) + "oauth refresh token missing for '{}'; re-run `{}`", + auth_slot_label(&profile), + oauth_reauth_command(&profile) ), ) })?; - let login_label = auth_slot_label(&profile); - let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &login_label).await?; + let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &profile).await?; save_profile_oauth_access_token(profile_name, &refreshed.access_token)?; let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { @@ -1158,10 +1230,8 @@ async fn resolve_oauth_profile_auth( Ok(auth) } -pub async fn resolved_auth_env(base: &BaseArgs) -> Result> { - let auth = resolve_auth(base).await?; +fn auth_env(auth: ResolvedAuth) -> Vec<(String, String)> { let mut envs = Vec::new(); - if let Some(api_key) = auth.api_key { envs.push(("BRAINTRUST_API_KEY".to_string(), api_key)); } @@ -1174,15 +1244,17 @@ pub async fn resolved_auth_env(base: &BaseArgs) -> Result> if let Some(org_name) = auth.org_name { envs.push(("BRAINTRUST_ORG_NAME".to_string(), org_name)); } - Ok(envs) + envs } pub async fn resolved_runner_env(base: &BaseArgs) -> Result> { - let mut envs = resolved_auth_env(base).await?; + let auth = resolve_auth(base).await?; + let resolved_org = auth.org_name.clone(); + let mut envs = auth_env(auth); let project = base .project .clone() - .or_else(|| crate::config::load().ok().and_then(|c| c.project)); + .or_else(|| crate::config::configured_project_for_context(base, resolved_org.as_deref())); if let Some(project) = project { envs.push(("BRAINTRUST_DEFAULT_PROJECT".to_string(), project)); } @@ -1193,27 +1265,53 @@ fn profile_matches_org_identifier(profile: &AuthProfile, org: &str) -> bool { profile.org_id.as_deref() == Some(org) || profile.org_name.as_deref() == Some(org) } -fn profile_org_label(profile: &AuthProfile) -> String { +fn profile_org(profile: &AuthProfile) -> &str { profile .org_name .as_deref() .filter(|org| !org.trim().is_empty()) - .map(str::to_string) - .unwrap_or_else(|| "cross-org".to_string()) + .or(profile + .org_id + .as_deref() + .filter(|org| !org.trim().is_empty())) + .unwrap_or("") } -fn profile_identity_label(profile: &AuthProfile) -> Option { - match profile.auth_kind { - AuthKind::Oauth => match (profile.user_name.as_deref(), profile.email.as_deref()) { - (Some(name), Some(email)) => Some(format!("{name} ({email})")), - (None, Some(email)) => Some(email.to_string()), - (Some(name), None) => Some(name.to_string()), - (None, None) => None, - }, - AuthKind::ApiKey => profile.api_key_hint.clone(), +fn profile_org_label(profile: &AuthProfile) -> String { + config::display_org(profile_org(profile)).to_string() +} + +fn oauth_reauth_command(profile: &AuthProfile) -> String { + format!( + "bt auth login --oauth --org {}", + shell_quote_arg(config::display_org(profile_org(profile))) + ) +} + +pub(crate) fn identity_label( + name: Option<&str>, + email: Option<&str>, + fallback: Option<&str>, +) -> Option { + match (name, email) { + (Some(name), Some(email)) => Some(format!("{name} ({email})")), + (Some(name), None) => Some(name.to_string()), + (None, Some(email)) => Some(email.to_string()), + (None, None) => fallback.map(str::to_string), } } +fn profile_identity_label(profile: &AuthProfile) -> Option { + let fallback = (profile.auth_kind == AuthKind::ApiKey) + .then_some(profile.api_key_hint.as_deref()) + .flatten(); + identity_label( + profile.user_name.as_deref(), + profile.email.as_deref(), + fallback, + ) +} + fn auth_slot_label(profile: &AuthProfile) -> String { let mut parts = vec![profile_org_label(profile)]; parts.push(auth_kind_label(profile.auth_kind).to_string()); @@ -1224,15 +1322,20 @@ fn auth_slot_label(profile: &AuthProfile) -> String { } fn is_cross_org_oauth_profile(profile: &AuthProfile) -> bool { - profile.auth_kind == AuthKind::Oauth - && profile - .org_id - .as_deref() - .is_none_or(|org_id| org_id.trim().is_empty()) - && profile - .org_name - .as_deref() - .is_none_or(|org_name| org_name.trim().is_empty()) + profile.auth_kind == AuthKind::Oauth && profile_org(profile).is_empty() +} + +fn reject_cross_org_api_key_preference( + prefer_api_key: bool, + org: Option<&str>, + store: &AuthStore, +) -> Result<()> { + let cross_org = org == Some("") + || (org.is_none() && store.profiles.values().any(is_cross_org_oauth_profile)); + if prefer_api_key && cross_org { + bail!("--prefer-api-key cannot be used from cross-org context; rerun with --org "); + } + Ok(()) } fn auth_profile_names_by_kind<'a>( @@ -1277,28 +1380,29 @@ fn ad_hoc_api_key_profile(org: Option<&str>, api_key: &str) -> ProfileInfo { } } -pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Option { +pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Result> { let store = load_auth_store().unwrap_or_default(); - // Display-side mirror of the interactive selectors: a lone candidate is - // chosen, no candidates continue the ladder, and an ambiguous set stops it - // (a real command would prompt or error, so status shows nothing here). + reject_cross_org_api_key_preference(base.prefer_api_key, org, &store)?; + let select = |kind| match auth_profile_names_by_kind(&store, org, kind).as_slice() { [] => Ok(None), [name] => Ok(Some((*name).to_string())), _ => bail!("multiple {kind:?} logins"), }; - let source = resolve_auth_source( + let source = match resolve_auth_source( base.prefer_api_key, resolve_cli_api_key_override(base), || resolve_env_api_key(base), || select(AuthKind::Oauth), || select(AuthKind::ApiKey), - ) - .ok()?; + ) { + Ok(source) => source, + Err(_) => return Ok(None), + }; - match source { + Ok(match source { AuthSource::CliApiKey(api_key) | AuthSource::EnvApiKey(api_key) => { Some(ad_hoc_api_key_profile(org, &api_key)) } @@ -1306,7 +1410,7 @@ pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Option None, - } + }) } fn missing_org_for_stored_logins_error(store: &AuthStore) -> Option { @@ -1314,15 +1418,7 @@ fn missing_org_for_stored_logins_error(store: &AuthStore) -> Option>(); if candidates.is_empty() { @@ -1387,6 +1483,53 @@ fn select_profile_from_store( Ok(names[idx].to_string()) } +fn saved_login_names(store: &AuthStore, include_cross_org: bool) -> Vec<&str> { + let mut oauth_orgs = BTreeSet::new(); + store + .profiles + .iter() + .filter(|(_, profile)| { + profile.auth_kind == AuthKind::ApiKey + || ((include_cross_org || !is_cross_org_oauth_profile(profile)) + && oauth_orgs.insert( + profile + .org_id + .as_deref() + .filter(|id| !id.is_empty()) + .unwrap_or_else(|| profile_org(profile)) + .to_ascii_lowercase(), + )) + }) + .map(|(name, _)| name.as_str()) + .collect() +} + +pub(crate) fn select_saved_login( + base: &mut BaseArgs, + current_org: Option<&str>, + include_cross_org: bool, +) -> Result { + let store = load_auth_store()?; + let names = saved_login_names(&store, include_cross_org); + let selected = match names.as_slice() { + [] => return Ok(false), + [name] => (*name).to_string(), + _ if ui::can_prompt() => { + select_profile_from_store("Select login", &names, current_org, &store)? + } + _ => { + bail!("multiple saved logins match; pass --org , or rerun interactively to choose") + } + }; + let profile = &store.profiles[&selected]; + if profile.auth_kind == AuthKind::ApiKey { + base.pinned_auth_slot = Some(selected); + } else { + base.org_name = Some(profile_org(profile).to_string()); + } + Ok(true) +} + fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec { names .iter() @@ -1402,26 +1545,20 @@ fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec, - can_prompt: bool, -) -> Result> { - let org = effective_org_name(base, cfg_org); - let candidates = auth_profile_names_by_kind(store, org, AuthKind::Oauth); - select_auth_profile_candidate("OAuth login", org, &candidates, store, can_prompt) -} - -fn select_api_key_profile_for_auth( +fn select_profile_for_auth( base: &BaseArgs, store: &AuthStore, cfg_org: &Option, + kind: AuthKind, can_prompt: bool, ) -> Result> { let org = effective_org_name(base, cfg_org); - let candidates = auth_profile_names_by_kind(store, org, AuthKind::ApiKey); - select_auth_profile_candidate("API key", org, &candidates, store, can_prompt) + let candidates = auth_profile_names_by_kind(store, org, kind); + let label = match kind { + AuthKind::Oauth => "OAuth login", + AuthKind::ApiKey => "API key", + }; + select_auth_profile_candidate(label, org, &candidates, store, can_prompt) } fn select_auth_profile_candidate( @@ -1458,6 +1595,11 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { if args.oauth { return run_login_oauth(base, args).await; } + if base.org_name.as_deref() == Some("") { + bail!( + "API-key login requires a concrete org; cross-org API keys do not exist. Use --oauth, or rerun with --org " + ); + } let has_explicit_api_key = base.api_key.as_ref().is_some_and(|k| !k.trim().is_empty()); if !has_explicit_api_key && ui::can_prompt() { @@ -1490,6 +1632,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { if requested_org_resolution == RequestedOrgResolution::SwitchToOauth { return run_login_oauth(base, args).await; } + let configured_org = config::load().ok().and_then(|cfg| cfg.org); let selected_org = select_login_org( login_orgs.clone(), match requested_org_resolution { @@ -1499,7 +1642,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { } RequestedOrgResolution::SwitchToOauth => unreachable!("handled above"), }, - None, + configured_org.as_deref(), interactive, base.verbose, false, @@ -1524,6 +1667,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { &selected_api_url, &login_app_url, Some(&selected_org), + &args.scope, ) .await .context("login succeeded, but failed to update active context")?; @@ -1606,10 +1750,11 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { exchange_oauth_authorization_code(&api_url, &redirect_uri, &auth_code, pkce_verifier) .await?; let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; + let configured_org = config::load().ok().and_then(|cfg| cfg.org); let selected_org = select_login_org( login_orgs.clone(), base.org_name.as_deref(), - None, + configured_org.as_deref(), ui::can_prompt(), base.verbose, true, @@ -1630,6 +1775,7 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { &selected_api_url, &app_url, selected_org.as_ref(), + &args.scope, ) .await .context("login succeeded, but failed to update active context")?; @@ -1749,8 +1895,14 @@ fn commit_oauth_profile( async fn run_login_refresh(base: &BaseArgs) -> Result<()> { let mut store = load_auth_store()?; let cfg_org = config_auth_context(base); - let profile_name = select_oauth_profile_for_auth(base, &store, &cfg_org, ui::can_prompt())? - .ok_or_else(|| { + let profile_name = select_profile_for_auth( + base, + &store, + &cfg_org, + AuthKind::Oauth, + ui::can_prompt(), + )? + .ok_or_else(|| { anyhow::anyhow!( "no OAuth login selected; pass --org or run `bt auth logins` to see available logins" ) @@ -1772,9 +1924,10 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { load_profile_oauth_refresh_token_for_profile(profile_name.as_str(), &profile)?.ok_or_else( || { anyhow::anyhow!( - "OAuth refresh token missing for '{}'; re-run `bt auth login --oauth --org `", - auth_slot_label(&profile) - ) + "OAuth refresh token missing for '{}'; re-run `{}`", + auth_slot_label(&profile), + oauth_reauth_command(&profile) + ) }, )?; @@ -1792,8 +1945,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { eprintln!("Cached access token expiry before refresh: unknown"); } - let login_label = auth_slot_label(&profile); - let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &login_label).await?; + let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &profile).await?; save_profile_oauth_access_token(profile_name.as_str(), &refreshed.access_token)?; let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { @@ -1909,7 +2061,7 @@ async fn resolve_post_login_project( let ctx = build_login_context_for_selected_org(credential, api_url, app_url, Some(selected_org)); let client = ApiClient::new(&ctx)?; - switch::validate_or_create_project(&client, project_name) + ui::select_or_create_project(&client, Some(project_name), None, None) .await .map(Some) } @@ -1920,23 +2072,27 @@ async fn persist_post_login_context( api_url: &str, app_url: &str, selected_org: Option<&LoginOrgInfo>, + scope: &config::ScopeArgs, ) -> Result { + // Scope is prompted last, after org (during login) and project. let project = resolve_post_login_project(base, credential, api_url, app_url, selected_org).await?; - let path = if ui::can_prompt() && config::local_path().is_some() { - switch::select_scope()?.0 - } else { - config::global_path()? - }; - + let (path, _) = scope.resolve(ui::can_prompt(), "Where to use this login")?; let mut cfg = config::load_file(&path); - switch::apply_switch_config( - &mut cfg, - selected_org.map(|org| org.name.as_str()), - project.as_ref(), - ); + let org = selected_org.map_or("", |org| org.name.as_str()); + let preserve_project = project.is_none() + && selected_org.is_some() + && config::org_option(cfg.org.as_deref()) == Some(org); + if !preserve_project { + cfg.set_context( + Some(org), + project + .as_ref() + .map(|project| (project.name.as_str(), project.id.as_str())), + ); + } config::save_file(&path, &cfg) - .context(format!("Could not save config to {}", path.display()))?; + .with_context(|| format!("Could not save config to {}", path.display()))?; Ok(PostLoginContextUpdate { display: format_post_login_context(selected_org, project.as_ref()), @@ -1955,22 +2111,49 @@ fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) -> Ok(()) } +fn filter_auth_store( + store: &AuthStore, + org: Option<&str>, + kind: Option, + api_key_hint: Option<&str>, +) -> AuthStore { + let mut filtered = store.clone(); + filtered.profiles.retain(|_, profile| { + org.is_none_or(|org| profile_matches_org_identifier(profile, org)) + && kind.is_none_or(|kind| profile.auth_kind == kind) + && api_key_hint.is_none_or(|hint| { + profile.auth_kind == AuthKind::ApiKey + && profile.api_key_hint.as_deref() == Some(hint.trim()) + }) + }); + filtered +} + async fn run_logins(base: &BaseArgs, _args: AuthLoginsArgs) -> Result<()> { let mut store = load_auth_store()?; - if store.profiles.is_empty() { + let has_filter = base.org_name.is_some() || base.prefer_api_key; + let filtered = filter_auth_store( + &store, + base.org_name.as_deref(), + base.prefer_api_key.then_some(AuthKind::ApiKey), + None, + ); + if filtered.profiles.is_empty() { return emit_result(base.json, serde_json::json!([]), || { - println!("No saved auth logins. Run `bt auth login` to create one.") + if store.profiles.is_empty() && !has_filter { + println!("No saved auth logins. Run `bt auth login` to create one."); + } }); } - let verifications = verify_all_profiles_from_store(&store).await; + let verifications = verify_all_profiles_from_store(&filtered).await; reconcile_verified_auth_slots(&mut store, &verifications)?; let all_network_errors = verifications .iter() .all(|v| v.status == "error" && !v.error.as_deref().unwrap_or("").contains("invalid")); if all_network_errors { eprintln!("Could not reach Braintrust API. Showing saved auth logins:"); - print_saved_profiles(&store, base.json)?; + print_saved_profiles(&filtered, base.json)?; return Ok(()); } @@ -2064,47 +2247,41 @@ fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { }); } - let cfg_org = config_auth_context(&base); - let org = effective_org_name(&base, &cfg_org); - let mut candidates: Vec<&str> = store + let requested_org = if matches!( + base.org_name_source, + Some(crate::args::ArgValueSource::CommandLine) + ) { + config::org_option(base.org_name.as_deref()) + } else { + None + }; + let filtered = filter_auth_store( + &store, + requested_org, + args.oauth.then_some(AuthKind::Oauth), + args.api_key_hint.as_deref(), + ); + let candidates = filtered .profiles - .iter() - .filter(|(_, profile)| org.is_none_or(|org| profile_matches_org_identifier(profile, org))) - .filter(|(_, profile)| { - args.api_key_hint.as_deref().is_none_or(|hint| { - profile.auth_kind == AuthKind::ApiKey - && profile.api_key_hint.as_deref() == Some(hint.trim()) - }) - }) - .map(|(name, _)| name.as_str()) - .collect(); - - if args.api_key_hint.is_none() && org.is_some() { - let oauth_candidates: Vec<&str> = candidates - .iter() - .copied() - .filter(|name| { - store - .profiles - .get(*name) - .is_some_and(|profile| profile.auth_kind == AuthKind::Oauth) - }) - .collect(); - if !oauth_candidates.is_empty() { - candidates = oauth_candidates; - } - } + .keys() + .map(String::as_str) + .collect::>(); + let cfg_org = config_auth_context(&base); + let current_org = effective_org_name(&base, &cfg_org); let profile_name = match candidates.len() { 0 => bail!("no matching auth login found; run `bt auth logins` to see available logins"), 1 => candidates[0].to_string(), - _ if ui::can_prompt() => { - select_profile_from_store("Select auth login to log out", &candidates, org, &store)? - } + _ if ui::can_prompt() => select_profile_from_store( + "Select auth login to log out", + &candidates, + current_org, + &filtered, + )?, _ => { - let labels = candidate_identities(&candidates, &store).join(", "); + let labels = candidate_identities(&candidates, &filtered).join(", "); bail!( - "multiple auth logins match: {labels}. Pass --org or --api-key-hint to disambiguate." + "multiple auth logins match: {labels}. Rerun interactively, or use --org with --oauth or --api-key-hint to disambiguate." ); } }; @@ -2290,10 +2467,20 @@ async fn verify_all_profiles_from_store(store: &AuthStore) -> Vec String { let mut parts = vec![ - v.org.clone().unwrap_or_else(|| "cross-org".to_string()), + config::display_org(v.org.as_deref().unwrap_or("")).to_string(), v.auth.clone(), ]; match v.status.as_str() { "ok" => { - let id = match (&v.user_name, &v.user_email) { - (Some(name), Some(email)) => Some(format!("{name} ({email})")), - (None, Some(email)) => Some(email.clone()), - _ => v.api_key_hint.clone(), - }; - if let Some(id) = id { + if let Some(id) = identity_label( + v.user_name.as_deref(), + v.user_email.as_deref(), + v.api_key_hint.as_deref(), + ) { parts.push(id); } } @@ -2377,12 +2563,26 @@ fn format_verification_line(v: &ProfileVerification) -> String { parts.join(" — ") } +fn profiles_grouped_by_org(store: &AuthStore) -> Vec<(&str, &AuthProfile)> { + let mut profiles = store + .profiles + .iter() + .map(|(name, profile)| (name.as_str(), profile)) + .collect::>(); + profiles.sort_by(|(a_name, a), (b_name, b)| { + profile_org(a) + .cmp(profile_org(b)) + .then_with(|| a_name.cmp(b_name)) + }); + profiles +} + fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { + let profiles = profiles_grouped_by_org(store); if json { - let output: Vec = store - .profiles - .values() - .map(|p| { + let output: Vec = profiles + .into_iter() + .map(|(_, p)| { serde_json::json!({ "auth": auth_kind_label(p.auth_kind), "org": p.org_name, @@ -2396,7 +2596,7 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { .collect(); println!("{}", serde_json::to_string(&output)?); } else { - for profile in store.profiles.values() { + for (_, profile) in profiles { println!(" {}", auth_slot_label(profile)); } } @@ -2444,12 +2644,11 @@ fn select_login_org( if orgs.is_empty() { bail!("no organizations found for this credential"); } - orgs.sort_by(|a, b| { - a.name - .to_ascii_lowercase() - .cmp(&b.name.to_ascii_lowercase()) - .then_with(|| a.name.cmp(&b.name)) - }); + sort_login_orgs(&mut orgs); + + if requested_org_name == Some("") { + return Ok(None); + } if let Some(name) = requested_org_name { return find_login_org(&orgs, name) @@ -2505,6 +2704,15 @@ fn select_login_org( )) } +fn sort_login_orgs(orgs: &mut [LoginOrgInfo]) { + orgs.sort_by(|a, b| { + a.name + .to_ascii_lowercase() + .cmp(&b.name.to_ascii_lowercase()) + .then_with(|| a.name.cmp(&b.name)) + }); +} + fn move_default_login_org_first( orgs: &mut Vec, default_org_name: Option<&str>, @@ -2536,13 +2744,18 @@ fn find_login_org_index(orgs: &[LoginOrgInfo], requested_org_name: &str) -> Opti }) } -fn missing_requested_org_error(orgs: &[LoginOrgInfo], requested_org_name: &str) -> anyhow::Error { - let available = orgs - .iter() +fn login_org_names(orgs: &[LoginOrgInfo]) -> String { + orgs.iter() .map(|org| org.name.as_str()) .collect::>() - .join(", "); - anyhow::anyhow!("org '{requested_org_name}' not found. Available: {available}") + .join(", ") +} + +fn missing_requested_org_error(orgs: &[LoginOrgInfo], requested_org_name: &str) -> anyhow::Error { + anyhow::anyhow!( + "org '{requested_org_name}' not found. Available: {}", + login_org_names(orgs) + ) } fn resolve_requested_org_for_api_key_login( @@ -3008,19 +3221,20 @@ async fn exchange_oauth_authorization_code( fn map_refresh_oauth_error( api_url: &str, - auth_login: &str, + profile: &AuthProfile, status: reqwest::StatusCode, body: &str, ) -> anyhow::Error { if let Ok(server_err) = serde_json::from_str::(body) { if matches!(server_err.error.as_deref(), Some("invalid_grant")) { let mut message = format!( - "oauth refresh token expired or was rejected for auth login '{auth_login}'" + "oauth refresh token expired or was rejected for auth login '{}'", + auth_slot_label(profile) ); if let Some(description) = server_err.error_description.as_deref() { message.push_str(&format!(" ({description})")); } - message.push_str("; re-run `bt auth login --oauth --org `"); + message.push_str(&format!("; re-run `{}`", oauth_reauth_command(profile))); return recoverable_auth_error(RecoverableAuthErrorKind::OauthRefreshToken, message); } } @@ -3034,7 +3248,7 @@ fn map_refresh_oauth_error( async fn refresh_oauth_access_token( api_url: &str, refresh_token: &str, - auth_login: &str, + profile: &AuthProfile, ) -> Result { let http_client = build_http_client_from_builder( reqwest::Client::builder() @@ -3056,7 +3270,7 @@ async fn refresh_oauth_access_token( if !response.status().is_success() { let status = response.status(); let body = response.text().await.unwrap_or_default(); - return Err(map_refresh_oauth_error(api_url, auth_login, status, &body)); + return Err(map_refresh_oauth_error(api_url, profile, status, &body)); } response @@ -3936,26 +4150,7 @@ mod tests { }; fn make_base() -> BaseArgs { - BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - no_color: false, - no_input: false, - project: None, - project_source: None, - org_name: None, - org_name_source: None, - api_key: None, - api_key_source: None, - prefer_api_key: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn auth_config(org: Option<&str>) -> crate::config::Config { @@ -4216,6 +4411,7 @@ mod tests { fn setup_global_config(project_id: Option<&str>, org: Option<&str>) { let cfg = crate::config::Config { org: org.map(str::to_string), + project: project_id.map(|_| "test-project".to_string()), project_id: project_id.map(str::to_string), ..crate::config::Config::default() }; @@ -4223,23 +4419,25 @@ mod tests { crate::config::save_global(&cfg).expect("save global config"); } - fn setup_auth_store_profiles(profiles: &[(&str, &str, &str, &str)]) { - let mut store = AuthStore::default(); + fn org_profile(kind: AuthKind, org_id: &str, org_name: &str) -> AuthProfile { + AuthProfile { + auth_kind: kind, + org_id: Some(org_id.into()), + org_name: Some(org_name.into()), + ..Default::default() + } + } + + fn setup_auth_store_profiles(profiles: &[(&str, &str, &str, &str)]) { + let mut store = AuthStore::default(); for (profile_name, org_name, api_url, app_url) in profiles { store.profiles.insert( (*profile_name).to_string(), AuthProfile { - auth_kind: AuthKind::ApiKey, api_url: Some((*api_url).to_string()), app_url: Some((*app_url).to_string()), - org_id: None, org_name: Some((*org_name).to_string()), - oauth_access_expires_at: None, - user_name: None, - email: None, - api_key_hash: None, - api_key_hint: None, - legacy_secret_key: None, + ..Default::default() }, ); } @@ -4294,22 +4492,26 @@ mod tests { #[test] fn invalid_grant_refresh_error_is_treated_as_recoverable() { + let profile = org_profile(AuthKind::Oauth, "org_test", "BT Staging"); + let command = oauth_reauth_command(&profile); + assert_eq!(command, "bt auth login --oauth --org 'BT Staging'"); let err = map_refresh_oauth_error( "https://api.example.com", - "work", + &profile, reqwest::StatusCode::BAD_REQUEST, r#"{"error":"invalid_grant","error_description":"refresh token expired"}"#, ); assert!(is_missing_credential_error(&err)); assert!(err.to_string().contains("refresh token expired")); + assert!(err.to_string().contains(&format!("re-run `{command}`"))); } #[test] fn nonrecoverable_refresh_errors_remain_nonrecoverable() { let err = map_refresh_oauth_error( "https://api.example.com", - "work", + &org_profile(AuthKind::Oauth, "org_test", "test-org"), reqwest::StatusCode::BAD_REQUEST, "unexpected response", ); @@ -4375,8 +4577,11 @@ mod tests { } async fn login_read_only_probe(&self, org_name: Option<&str>) -> Result { - self.login_read_only_with_base(base_args_for_path_probe(org_name)) - .await + let mut base = base_args_for_path_probe(org_name); + if let Some(org) = org_name.filter(|org| !org.trim().is_empty()) { + base.app_url = Some(spawn_api_key_login_server(org)); + } + self.login_read_only_with_base(base).await } } @@ -4388,6 +4593,41 @@ mod tests { } } + fn spawn_login_response_server(status: &str, body: String) -> String { + use std::io::{Read as _, Write as _}; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind login server"); + let address = listener.local_addr().expect("login server address"); + let status = status.to_string(); + std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept login request"); + let mut request = [0u8; 4096]; + let _ = stream.read(&mut request); + write!( + stream, + "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .expect("write login response"); + }); + format!("http://{address}") + } + + fn spawn_api_key_login_server(org_name: &str) -> String { + spawn_login_response_server( + "200 OK", + serde_json::json!({ + "org_info": [{ + "id": "org_test", + "name": org_name, + "api_url": "https://api.example.test" + }] + }) + .to_string(), + ) + } + #[test] fn default_app_url_is_www() { assert_eq!(DEFAULT_APP_URL, "https://www.braintrust.dev"); @@ -4400,11 +4640,8 @@ mod tests { store.profiles.insert( api_key_slot_key(&api_key_hash("test-api-key"), "org_fake"), AuthProfile { - auth_kind: AuthKind::ApiKey, - org_id: Some("org_fake".to_string()), - org_name: Some("test-org".to_string()), api_key_hint: Some("sk-****abcde".to_string()), - ..Default::default() + ..org_profile(AuthKind::ApiKey, "org_fake", "test-org") }, ); store.profiles.insert( @@ -4420,11 +4657,25 @@ mod tests { ); save_auth_store(&store).expect("save auth store"); - let info = active_auth_info(&make_base(), None).expect("active auth info"); + let info = active_auth_info(&make_base(), None) + .expect("resolve active auth") + .expect("active auth info"); assert_eq!(info.auth_method, "oauth"); assert_eq!(info.email.as_deref(), Some("user@example.test")); assert_eq!(info.org_name, None); + + let mut base = make_base(); + base.prefer_api_key = true; + assert!(resolve_auth(&base) + .await + .unwrap_err() + .to_string() + .contains("cross-org")); + assert!(active_auth_info(&base, None) + .unwrap_err() + .to_string() + .contains("cross-org")); } fn save_cached_oauth_login(store: &mut AuthStore, org_id: &str, org_name: &str) -> String { @@ -4432,15 +4683,12 @@ mod tests { store.profiles.insert( slot_key.clone(), AuthProfile { - auth_kind: AuthKind::Oauth, api_url: Some("https://api.example.test".to_string()), app_url: Some("https://www.example.test".to_string()), - org_id: Some(org_id.to_string()), - org_name: Some(org_name.to_string()), oauth_access_expires_at: Some(current_unix_timestamp() + 3600), user_name: Some("Test User".to_string()), email: Some("user@example.test".to_string()), - ..Default::default() + ..org_profile(AuthKind::Oauth, org_id, org_name) }, ); save_profile_secret_plaintext( @@ -4481,6 +4729,7 @@ mod tests { base.org_name = Some("test-org".to_string()); base.api_key = Some("command-line-api-key".to_string()); base.api_key_source = Some(crate::args::ArgValueSource::CommandLine); + base.app_url = Some(spawn_api_key_login_server("test-org")); let resolved = resolve_auth(&base).await.expect("resolve auth"); @@ -4488,6 +4737,46 @@ mod tests { assert_eq!(resolved.api_key.as_deref(), Some("command-line-api-key")); } + #[tokio::test] + async fn ad_hoc_api_key_validation_errors() { + for (app_url, expected) in [ + ( + spawn_api_key_login_server("different-org"), + "does not belong", + ), + ( + spawn_login_response_server("401 Unauthorized", "{}".into()), + "not valid", + ), + ] { + let mut base = make_base(); + base.org_name = Some("requested-org".into()); + base.app_url = Some(app_url); + assert!( + resolve_ad_hoc_api_key_auth(&base, &None, "selected-key".into()) + .await + .unwrap_err() + .to_string() + .contains(expected) + ); + } + } + + #[test] + fn selected_stored_api_key_wrong_org_fails_before_secret_lookup() { + let mut store = AuthStore::default(); + store.profiles.insert( + "stored-slot".into(), + org_profile(AuthKind::ApiKey, "org_actual", "actual-org"), + ); + let mut base = make_base(); + base.org_name = Some("requested-org".into()); + + let err = resolve_api_key_profile_auth(&base, &mut store, &None, "stored-slot") + .expect_err("wrong-org stored key must fail locally"); + assert!(err.to_string().contains("does not belong")); + } + #[tokio::test] async fn auth_precedence_prefer_api_key_promotes_env_api_key() { let _env = TestEnv::new(None, None).await; @@ -4499,6 +4788,7 @@ mod tests { base.api_key = Some("environment-api-key".to_string()); base.api_key_source = Some(crate::args::ArgValueSource::EnvVariable); base.prefer_api_key = true; + base.app_url = Some(spawn_api_key_login_server("test-org")); let resolved = resolve_auth(&base).await.expect("resolve auth"); @@ -4525,6 +4815,28 @@ mod tests { ); } + #[tokio::test] + async fn active_auth_info_hides_ambiguous_api_keys_instead_of_failing_status() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + for (slot, hint) in [("key-a", "sk-****aaaaa"), ("key-b", "sk-****bbbbb")] { + store.profiles.insert( + slot.into(), + AuthProfile { + api_key_hint: Some(hint.into()), + ..org_profile(AuthKind::ApiKey, "org_test", "test-org") + }, + ); + } + save_auth_store(&store).expect("save auth store"); + let mut base = make_base(); + base.org_name = Some("test-org".into()); + + assert!(active_auth_info(&base, Some("test-org")) + .expect("status auth lookup") + .is_none()); + } + #[tokio::test] async fn active_auth_info_prefer_api_key_selects_stored_key_for_org() { let _env = TestEnv::new(None, None).await; @@ -4532,29 +4844,25 @@ mod tests { store.profiles.insert( oauth_slot_key("org_fake", "user@example.test"), AuthProfile { - auth_kind: AuthKind::Oauth, - org_id: Some("org_fake".to_string()), - org_name: Some("test-org".to_string()), user_name: Some("Test User".to_string()), email: Some("user@example.test".to_string()), - ..Default::default() + ..org_profile(AuthKind::Oauth, "org_fake", "test-org") }, ); store.profiles.insert( api_key_slot_key(&api_key_hash("test-api-key"), "org_fake"), AuthProfile { - auth_kind: AuthKind::ApiKey, - org_id: Some("org_fake".to_string()), - org_name: Some("test-org".to_string()), api_key_hint: Some("sk-****abcde".to_string()), - ..Default::default() + ..org_profile(AuthKind::ApiKey, "org_fake", "test-org") }, ); save_auth_store(&store).expect("save auth store"); let mut base = make_base(); base.prefer_api_key = true; - let info = active_auth_info(&base, Some("test-org")).expect("active auth info"); + let info = active_auth_info(&base, Some("test-org")) + .expect("resolve active auth") + .expect("active auth info"); assert_eq!(info.auth_method, "api_key"); assert_eq!(info.api_key_hint.as_deref(), Some("sk-****abcde")); @@ -4566,13 +4874,7 @@ mod tests { let mut store = AuthStore::default(); store.profiles.insert( "work".to_string(), - AuthProfile { - auth_kind: AuthKind::ApiKey, - org_id: Some("org_fake".to_string()), - org_name: Some("test-org".to_string()), - api_key_hash: None, - ..Default::default() - }, + org_profile(AuthKind::ApiKey, "org_fake", "test-org"), ); maybe_rekey_api_key_profile_after_secret_load(&mut store, "work", "test-api-key") @@ -4625,11 +4927,8 @@ mod tests { store.profiles.insert( "work".to_string(), AuthProfile { - auth_kind: AuthKind::Oauth, - org_id: Some("org_fake".to_string()), - org_name: Some("test-org".to_string()), email: Some("user@example.test".to_string()), - ..Default::default() + ..org_profile(AuthKind::Oauth, "org_fake", "test-org") }, ); @@ -4916,18 +5215,18 @@ mod tests { } #[test] - fn auth_source_default_order_is_oauth_then_api_key_then_env() { + fn auth_source_default_order_is_oauth_then_env_then_stored_api_key() { assert_eq!( auth_source(false, None, Some("env"), Some("oauth"), Some("ak")), AuthSource::Oauth("oauth".into()) ); assert_eq!( auth_source(false, None, Some("env"), None, Some("ak")), - AuthSource::ApiKey("ak".into()) + AuthSource::EnvApiKey("env".into()) ); assert_eq!( - auth_source(false, None, Some("env"), None, None), - AuthSource::EnvApiKey("env".into()) + auth_source(false, None, None, None, Some("ak")), + AuthSource::ApiKey("ak".into()) ); assert_eq!(auth_source(false, None, None, None, None), AuthSource::None); } @@ -4963,32 +5262,118 @@ mod tests { assert!(err.to_string().contains("multiple oauth logins")); } + fn login_filter_store() -> AuthStore { + let mut store = AuthStore::default(); + for (slot, kind, suffix, hint) in [ + ("oauth-a", AuthKind::Oauth, "a", None), + ("key-a", AuthKind::ApiKey, "a", Some("sk-****aaaaa")), + ("key-b", AuthKind::ApiKey, "b", Some("sk-****bbbbb")), + ] { + store.profiles.insert( + slot.into(), + AuthProfile { + api_key_hint: hint.map(str::to_string), + ..org_profile( + kind, + &format!("org_test_{suffix}"), + &format!("test-org-{suffix}"), + ) + }, + ); + } + store.profiles.insert( + "cross".into(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some(String::new()), + ..Default::default() + }, + ); + store + } + + #[test] + fn login_and_logout_filters_compose() { + let store = login_filter_store(); + let matches = |org, kind, hint| { + filter_auth_store(&store, org, kind, hint) + .profiles + .into_keys() + .collect::>() + }; + assert_eq!( + matches(Some("test-org-a"), None, None), + ["key-a", "oauth-a"] + ); + assert_eq!(matches(Some("org_test_b"), None, None), ["key-b"]); + assert_eq!( + matches(Some("test-org-a"), Some(AuthKind::ApiKey), None), + ["key-a"] + ); + assert_eq!(matches(Some(""), None, None), ["cross"]); + assert!(matches(Some(""), Some(AuthKind::ApiKey), None).is_empty()); + assert_eq!(matches(None, None, None).len(), 4); + assert_eq!( + matches(Some("test-org-a"), Some(AuthKind::Oauth), None), + ["oauth-a"] + ); + assert_eq!(matches(None, None, Some("sk-****bbbbb")), ["key-b"]); + + let mut picker_store = store.clone(); + picker_store.profiles.insert( + "oauth-a-duplicate".into(), + picker_store.profiles["oauth-a"].clone(), + ); + assert_eq!(saved_login_names(&picker_store, true).len(), 4); + assert_eq!(saved_login_names(&picker_store, false).len(), 3); + } + #[tokio::test] - async fn persist_post_login_context_clears_stale_project_for_org_only_login() { + async fn post_login_context_preserves_only_same_org_projects() { let _env = TestEnv::new(None, None).await; - crate::config::save_global(&crate::config::Config { - org: Some("old-org".to_string()), - project: Some("stale-project".to_string()), - project_id: Some("proj_stale".to_string()), - ..Default::default() - }) - .expect("save initial config"); + let save = |org: &str| { + crate::config::save_global(&crate::config::Config { + org: Some(org.into()), + project: Some("test-project".into()), + project_id: Some("proj_test".into()), + ..Default::default() + }) + .unwrap(); + }; + let persist = |org: Option| async move { + persist_post_login_context( + &make_base(), + "test-credential", + "https://api.example.test", + "https://www.example.test", + org.as_ref(), + &config::ScopeArgs { + global: true, + local: false, + }, + ) + .await + .unwrap(); + crate::config::load_global().unwrap() + }; - let update = persist_post_login_context( - &make_base(), - "test-api-key", - "https://api.example.test", - "https://www.example.test", - Some(&login_org("org_123", "acme")), - ) - .await - .expect("persist context"); - let cfg = crate::config::load_global().expect("load global config"); + save("old-org"); + let cfg = persist(Some(login_org("org_test", "test-org"))).await; + assert_eq!((cfg.org.as_deref(), cfg.project), (Some("test-org"), None)); - assert_eq!(update.display, "acme"); - assert_eq!(cfg.org.as_deref(), Some("acme")); - assert_eq!(cfg.project, None); - assert_eq!(cfg.project_id, None); + save("test-org"); + let cfg = persist(Some(login_org("org_test", "test-org"))).await; + assert_eq!( + (cfg.project.as_deref(), cfg.project_id.as_deref()), + (Some("test-project"), Some("proj_test")) + ); + + save(""); + let cfg = persist(None).await; + assert_eq!( + (cfg.org.as_deref(), cfg.project, cfg.project_id), + (Some(""), None, None) + ); } #[tokio::test] @@ -5012,197 +5397,230 @@ mod tests { } #[test] - fn resolve_requested_org_for_api_key_login_keeps_matching_requested_org() { - let orgs = vec![login_org("org_1", "acme")]; - - let resolution = - resolve_requested_org_for_api_key_login(&orgs, Some("acme"), false, |_, _| { - panic!("prompt should not be called") - }) - .expect("resolve"); - - assert_eq!(resolution, RequestedOrgResolution::UseRequestedOrg); - } - - #[test] - fn resolve_requested_org_for_api_key_login_errors_without_prompt() { - let orgs = vec![login_org("org_1", "braintrustdata.com")]; - - let err = - resolve_requested_org_for_api_key_login(&orgs, Some("ced-test-1"), false, |_, _| { - panic!("prompt should not be called") - }) - .expect_err("should fail"); - - assert!(err - .to_string() - .contains("org 'ced-test-1' not found. Available: braintrustdata.com")); - } - - #[test] - fn resolve_requested_org_for_api_key_login_can_switch_to_oauth() { - let orgs = vec![login_org("org_1", "braintrustdata.com")]; - - let resolution = resolve_requested_org_for_api_key_login( - &orgs, - Some("ced-test-1"), - true, - |requested_org_name, available_orgs| { - assert_eq!(requested_org_name, "ced-test-1"); - assert_eq!(available_orgs.len(), 1); - Ok(ApiKeyOrgMismatchAction::UseOauth) - }, - ) - .expect("resolve"); - - assert_eq!(resolution, RequestedOrgResolution::SwitchToOauth); - } - - #[test] - fn resolve_requested_org_for_api_key_login_can_continue_with_api_key() { - let orgs = vec![login_org("org_1", "braintrustdata.com")]; - - let resolution = resolve_requested_org_for_api_key_login( + fn requested_api_key_org_resolution() { + fn no_prompt(_: &str, _: &[LoginOrgInfo]) -> Result { + panic!("prompt should not be called") + } + let orgs = vec![login_org("org_test", "test-org")]; + assert_eq!( + resolve_requested_org_for_api_key_login(&orgs, Some("test-org"), false, no_prompt) + .unwrap(), + RequestedOrgResolution::UseRequestedOrg + ); + assert!(resolve_requested_org_for_api_key_login( &orgs, - Some("ced-test-1"), - true, - |requested_org_name, available_orgs| { - assert_eq!(requested_org_name, "ced-test-1"); - assert_eq!(available_orgs.len(), 1); - Ok(ApiKeyOrgMismatchAction::UseApiKey) - }, + Some("other-org"), + false, + no_prompt ) - .expect("resolve"); + .unwrap_err() + .to_string() + .contains("org 'other-org' not found. Available: test-org")); - assert_eq!(resolution, RequestedOrgResolution::IgnoreRequestedOrg); - } - - #[test] - fn obscure_api_key_standard() { - assert_eq!(obscure_api_key("sk-LumEdp0BbLRzhJwO"), "sk-****zhJwO"); - } - - #[test] - fn obscure_api_key_short() { - assert_eq!(obscure_api_key("abc"), "****"); - } - - #[test] - fn obscure_api_key_no_dash() { - assert_eq!(obscure_api_key("abcdefghijklm"), "****ijklm"); - } - - #[test] - fn obscure_api_key_non_ascii() { - assert_eq!(obscure_api_key("sk-café-résumé-key"), "****"); - } - - #[test] - fn decode_jwt_identity_extracts_claims() { - let header = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"RS256"}"#); - let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(r#"{"name":"Alice","email":"alice@example.com"}"#); - let token = format!("{header}.{payload}.sig"); - let id = decode_jwt_identity(&token); - assert_eq!(id.name.as_deref(), Some("Alice")); - assert_eq!(id.email.as_deref(), Some("alice@example.com")); + for (action, expected) in [ + ( + ApiKeyOrgMismatchAction::UseOauth, + RequestedOrgResolution::SwitchToOauth, + ), + ( + ApiKeyOrgMismatchAction::UseApiKey, + RequestedOrgResolution::IgnoreRequestedOrg, + ), + ] { + let actual = resolve_requested_org_for_api_key_login( + &orgs, + Some("other-org"), + true, + |requested, available| { + assert_eq!(requested, "other-org"); + assert_eq!(available.len(), 1); + Ok(action) + }, + ) + .unwrap(); + assert_eq!(actual, expected); + } } #[test] - fn decode_jwt_identity_handles_missing_claims() { - let header = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"RS256"}"#); - let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"sub":"123"}"#); - let token = format!("{header}.{payload}.sig"); - let id = decode_jwt_identity(&token); - assert_eq!(id.name, None); - assert_eq!(id.email, None); + fn obscure_api_keys() { + for (key, expected) in [ + ("sk-LumEdp0BbLRzhJwO", "sk-****zhJwO"), + ("abc", "****"), + ("abcdefghijklm", "****ijklm"), + ("sk-café-résumé-key", "****"), + ] { + assert_eq!(obscure_api_key(key), expected); + } } #[test] - fn decode_jwt_identity_handles_garbage() { + fn decode_jwt_identity_handles_claims_and_invalid_tokens() { + let encode = |payload| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload); + let header = encode(r#"{"alg":"RS256"}"#); + for (payload, expected) in [ + ( + r#"{"name":"Test User","email":"user@example.test"}"#, + (Some("Test User"), Some("user@example.test")), + ), + (r#"{"sub":"123"}"#, (None, None)), + ] { + let id = decode_jwt_identity(&format!("{header}.{}.sig", encode(payload))); + assert_eq!((id.name.as_deref(), id.email.as_deref()), expected); + } let id = decode_jwt_identity("not-a-jwt"); - assert_eq!(id.name, None); - assert_eq!(id.email, None); + assert_eq!((id.name, id.email), (None, None)); } #[test] - fn format_verification_line_ok_with_identity() { - let v = ProfileVerification { - name: "work".into(), + fn auth_logins_are_grouped_by_org() { + let verification = |name: &str, org: Option<&str>| ProfileVerification { + name: name.into(), slot_hash: None, auth: "oauth".into(), - org: Some("acme".into()), - org_id: None, - user_name: Some("Alice".into()), - user_email: Some("alice@example.com".into()), - api_key_hint: None, - status: "ok".into(), - error: None, - }; - assert_eq!( - format_verification_line(&v), - "acme — oauth — Alice (alice@example.com)" - ); - } - - #[test] - fn format_verification_line_ok_with_api_key_hint() { - let v = ProfileVerification { - name: "work".into(), - slot_hash: Some(api_key_hash("test-api-key")), - auth: "api_key".into(), - org: Some("acme".into()), + org: org.map(str::to_string), org_id: None, user_name: None, user_email: None, - api_key_hint: Some("sk-****zhJwO".into()), + api_key_hint: None, status: "ok".into(), error: None, }; + let mut verifications = vec![ + verification("profile-z", Some("test-org-a")), + verification("profile-a", Some("test-org-b")), + verification("profile-m", Some("test-org-a")), + verification("profile-x", None), + ]; + + sort_profile_verifications(&mut verifications); + + let order = verifications + .iter() + .map(|v| (v.org.as_deref(), v.name.as_str())) + .collect::>(); assert_eq!( - format_verification_line(&v), - "acme — api_key — sk-****zhJwO" + order, + vec![ + (None, "profile-x"), + (Some("test-org-a"), "profile-m"), + (Some("test-org-a"), "profile-z"), + (Some("test-org-b"), "profile-a"), + ] ); } #[test] - fn format_verification_line_expired() { - let v = ProfileVerification { - name: "old".into(), - slot_hash: None, - auth: "oauth".into(), - org: None, - org_id: None, - user_name: None, - user_email: None, - api_key_hint: None, - status: "expired".into(), - error: None, - }; + fn saved_auth_logins_are_grouped_by_org() { + let mut store = AuthStore::default(); + for (name, org) in [ + ("profile-z", "test-org-a"), + ("profile-a", "test-org-b"), + ("profile-m", "test-org-a"), + ] { + store.profiles.insert( + name.into(), + AuthProfile { + org_name: Some(org.into()), + ..Default::default() + }, + ); + } + + let order = profiles_grouped_by_org(&store) + .into_iter() + .map(|(name, profile)| (profile_org(profile), name)) + .collect::>(); assert_eq!( - format_verification_line(&v), - "cross-org — oauth — token expired" + order, + vec![ + ("test-org-a", "profile-m"), + ("test-org-a", "profile-z"), + ("test-org-b", "profile-a"), + ] ); } #[test] - fn format_verification_line_error() { - let v = ProfileVerification { - name: "bad".into(), + fn verification_line_formatting() { + for (name, email, hint, expected) in [ + ( + Some("Test User"), + Some("user@example.test"), + None, + Some("Test User (user@example.test)"), + ), + ( + None, + Some("user@example.test"), + None, + Some("user@example.test"), + ), + (None, None, Some("sk-****abcde"), Some("sk-****abcde")), + (None, None, None, None), + ] { + assert_eq!(identity_label(name, email, hint).as_deref(), expected); + } + + let verification = |auth: &str, + org: Option<&str>, + identity: Option<&str>, + hint: Option<&str>, + status: &str, + error: Option<&str>| ProfileVerification { + name: "test-profile".into(), slot_hash: None, - auth: "api_key".into(), - org: Some("corp".into()), + auth: auth.into(), + org: org.map(str::to_string), org_id: None, - user_name: None, - user_email: None, - api_key_hint: None, - status: "error".into(), - error: Some("invalid API key".into()), + user_name: identity.map(str::to_string), + user_email: identity.map(|_| "user@example.test".into()), + api_key_hint: hint.map(str::to_string), + status: status.into(), + error: error.map(str::to_string), }; - assert_eq!( - format_verification_line(&v), - "corp — api_key — invalid API key" - ); + let cases = [ + ( + verification( + "oauth", + Some("test-org"), + Some("Test User"), + None, + "ok", + None, + ), + "test-org — oauth — Test User (user@example.test)", + ), + ( + verification( + "api_key", + Some("test-org"), + None, + Some("sk-****abcde"), + "ok", + None, + ), + "test-org — api_key — sk-****abcde", + ), + ( + verification("oauth", None, None, None, "expired", None), + "cross-org — oauth — token expired", + ), + ( + verification( + "api_key", + Some("test-org"), + None, + None, + "error", + Some("invalid API key"), + ), + "test-org — api_key — invalid API key", + ), + ]; + for (verification, expected) in cases { + assert_eq!(format_verification_line(&verification), expected); + } } #[tokio::test] @@ -5230,80 +5648,44 @@ mod tests { } } - #[tokio::test] - async fn oauth_callback_listener_responds_to_http_request() { + async fn assert_oauth_callback(stale_connection: bool, code: &str, state: &str) { use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; use tokio::net::TcpStream; - let callback_server = bind_oauth_callback_server().expect("bind callback server"); - let addr = format!("127.0.0.1:{}", callback_server.port); - let callback = tokio::spawn(wait_for_oauth_callback(callback_server)); - - let mut stream = TcpStream::connect(addr) - .await - .expect("connect to callback listener"); - stream - .write_all( - b"GET /callback?code=test-code&state=test-state HTTP/1.1\r\nHost: 127.0.0.1\r\nUser-Agent: test\r\n\r\n", - ) - .await - .expect("write callback request"); + let server = bind_oauth_callback_server().unwrap(); + let addr = format!("127.0.0.1:{}", server.port); + let callback = tokio::spawn(wait_for_oauth_callback(server)); + if stale_connection { + drop(TcpStream::connect(&addr).await.unwrap()); + } + let mut stream = TcpStream::connect(addr).await.unwrap(); + let request = + format!("GET /callback?code={code}&state={state} HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n"); + stream.write_all(request.as_bytes()).await.unwrap(); let mut response = vec![0u8; 4096]; - let bytes_read = tokio::time::timeout(Duration::from_secs(1), stream.read(&mut response)) - .await - .expect("callback response timed out") - .expect("read callback response"); - let response = String::from_utf8_lossy(&response[..bytes_read]); - - let params = callback + let read = tokio::time::timeout(Duration::from_secs(1), stream.read(&mut response)) .await - .expect("callback task") - .expect("callback params"); - assert_eq!(params.code.as_deref(), Some("test-code")); - assert_eq!(params.state.as_deref(), Some("test-state")); + .unwrap() + .unwrap(); + let params = callback.await.unwrap().unwrap(); + assert_eq!( + (params.code.as_deref(), params.state.as_deref()), + (Some(code), Some(state)) + ); + let response = String::from_utf8_lossy(&response[..read]); assert!(response.starts_with("HTTP/1.1 200 OK")); assert!(response.contains("Authorization Successful")); } #[tokio::test] - async fn oauth_callback_listener_ignores_empty_connection() { - use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; - use tokio::net::TcpStream; - - let callback_server = bind_oauth_callback_server().expect("bind callback server"); - let addr = format!("127.0.0.1:{}", callback_server.port); - let callback = tokio::spawn(wait_for_oauth_callback(callback_server)); - - let stale = TcpStream::connect(&addr) - .await - .expect("connect stale callback socket"); - drop(stale); - - let mut stream = TcpStream::connect(addr) - .await - .expect("connect to callback listener"); - stream - .write_all( - b"GET /callback?code=next-code&state=next-state HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - ) - .await - .expect("write callback request"); - - let mut response = vec![0u8; 4096]; - let bytes_read = tokio::time::timeout(Duration::from_secs(1), stream.read(&mut response)) - .await - .expect("callback response timed out") - .expect("read callback response"); - let response = String::from_utf8_lossy(&response[..bytes_read]); + async fn oauth_callback_listener_responds_to_http_request() { + assert_oauth_callback(false, "test-code", "test-state").await; + } - let params = callback - .await - .expect("callback task") - .expect("callback params"); - assert_eq!(params.code.as_deref(), Some("next-code")); - assert_eq!(params.state.as_deref(), Some("next-state")); - assert!(response.starts_with("HTTP/1.1 200 OK")); + #[tokio::test] + async fn oauth_callback_listener_ignores_empty_connection() { + assert_oauth_callback(true, "next-code", "next-state").await; } #[tokio::test] @@ -5314,13 +5696,13 @@ mod tests { #[tokio::test] async fn login_read_only_cached_project_id_and_org_uses_fast_path() { - let env = TestEnv::new(Some("proj_123"), None).await; + let env = TestEnv::new(Some("proj_123"), Some("test-org")).await; let ctx = env - .login_read_only_probe(Some("acme")) + .login_read_only_probe(Some("test-org")) .await .expect("fast path should succeed"); - assert_eq!(ctx.login.org_name().as_deref(), Some("acme")); + assert_eq!(ctx.login.org_name().as_deref(), Some("test-org")); assert_eq!(ctx.login.org_id().as_deref(), Some("")); assert_eq!(ctx.api_url, "not-a-valid-url"); } @@ -5332,9 +5714,13 @@ mod tests { } #[tokio::test] - async fn login_read_only_cached_project_id_but_whitespace_org_falls_back_to_login() { + async fn login_read_only_cached_project_id_but_whitespace_org_is_cross_org() { let env = TestEnv::new(Some("proj_123"), None).await; - assert_invalid_api_url(env.login_read_only_probe(Some(" ")).await); + let err = match env.login_read_only_probe(Some(" ")).await { + Ok(_) => panic!("whitespace org should be canonical cross-org"), + Err(err) => err, + }; + assert!(err.to_string().contains("concrete org")); } #[tokio::test] @@ -5376,19 +5762,21 @@ mod tests { #[tokio::test] async fn login_read_only_cached_project_id_and_org_uses_default_urls() { - let env = TestEnv::new(Some("proj_123"), None).await; + let env = TestEnv::new(Some("proj_123"), Some("test-org")).await; let mut base = make_base(); base.api_key = Some("test-api-key".into()); - base.org_name = Some("acme".into()); + base.org_name = Some("test-org".into()); + let app_url = spawn_api_key_login_server("test-org"); + base.app_url = Some(app_url.clone()); let ctx = env .login_read_only_with_base(base) .await .expect("fast path should succeed"); - assert_eq!(ctx.login.org_name().as_deref(), Some("acme")); - assert_eq!(ctx.api_url, DEFAULT_API_URL); - assert_eq!(ctx.app_url, DEFAULT_APP_URL); + assert_eq!(ctx.login.org_name().as_deref(), Some("test-org")); + assert_eq!(ctx.api_url, "https://api.example.test"); + assert_eq!(ctx.app_url, app_url); } #[tokio::test] diff --git a/src/config/mod.rs b/src/config/mod.rs index 26493ada..db36c628 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,4 +1,4 @@ -use anyhow::{anyhow, bail, Result}; +use anyhow::{anyhow, bail, Context, Result}; use clap::{Args, Subcommand}; use std::{ env, fs, @@ -70,17 +70,32 @@ impl Config { .collect() } - pub(crate) fn merge(&self, other: &Config) -> Config { + pub(crate) fn set_context(&mut self, org: Option<&str>, project: Option<(&str, &str)>) { + self.org = org_option(org).map(str::to_string); + (self.project, self.project_id) = project + .map(|(name, id)| (name.to_string(), id.to_string())) + .unzip(); + } + + pub(crate) fn merge(&self, local: &Config) -> Config { let mut extra = self.extra.clone(); - extra.extend(other.extra.clone()); - let project = other.project.clone().or_else(|| self.project.clone()); - let project_id = if other.project.is_some() { - other.project_id.clone() - } else { - self.project_id.clone() + extra.extend(local.extra.clone()); + let global_id = self.project.as_ref().and(self.project_id.clone()); + let (org, project, project_id) = match (&local.org, &local.project) { + (Some(org), Some(project)) => ( + Some(org.clone()), + Some(project.clone()), + local.project_id.clone(), + ), + (Some(org), None) if self.org.as_ref() == Some(org) => { + (Some(org.clone()), self.project.clone(), global_id) + } + (Some(org), None) => (Some(org.clone()), None, None), + (None, Some(project)) => (None, Some(project.clone()), local.project_id.clone()), + (None, None) => (self.org.clone(), self.project.clone(), global_id), }; Config { - org: other.org.clone().or_else(|| self.org.clone()), + org, project, project_id, extra, @@ -127,6 +142,14 @@ pub fn load_file(path: &Path) -> Config { config.extra.remove("profile"); + // Fold a literal "cross-org" to the canonical "" marker on load. + if let Some(org) = config.org.as_deref() { + let normalized = normalize_org(org); + if normalized != org { + config.org = Some(normalized.to_string()); + } + } + for key in config.extra.keys() { print_command_status( CommandStatus::Error, @@ -177,13 +200,37 @@ pub(crate) fn project_from_config_for_context( .flatten() } -fn config_matches_context(_base: &BaseArgs, cfg: &Config, resolved_org: Option<&str>) -> bool { - let cfg_org = trimmed_option(cfg.org.as_deref()); - let resolved_org = trimmed_option(resolved_org); +fn config_matches_context(base: &BaseArgs, cfg: &Config, resolved_org: Option<&str>) -> bool { + let cfg_org = org_option(cfg.org.as_deref()); + let requested_org = org_option(resolved_org).or_else(|| org_option(base.org_name.as_deref())); + + requested_org.is_none_or(|resolved| cfg_org == Some(resolved)) +} - cfg_org - .zip(resolved_org) - .is_none_or(|(cfg, resolved)| cfg == resolved) +/// Trim an org while preserving the empty cross-org marker. +pub(crate) fn org_option(value: Option<&str>) -> Option<&str> { + value.map(str::trim) +} + +/// Human-facing spelling of the empty cross-org marker. +pub(crate) const CROSS_ORG_ALIAS: &str = "cross-org"; + +/// Trim an org and fold the [`CROSS_ORG_ALIAS`] to the canonical `""` marker. +pub(crate) fn normalize_org(value: &str) -> &str { + let trimmed = value.trim(); + if trimmed == CROSS_ORG_ALIAS { + "" + } else { + trimmed + } +} + +pub(crate) fn display_org(org: &str) -> &str { + if org.is_empty() { + CROSS_ORG_ALIAS + } else { + org + } } pub(crate) fn trimmed_option(value: Option<&str>) -> Option<&str> { @@ -209,57 +256,110 @@ pub fn save_global(config: &Config) -> Result<()> { } pub fn find_local_config_dir() -> Option { - let home = dirs::home_dir(); - let mut current_dir = std::env::current_dir().ok()?; + find_local_config_dir_from(std::env::current_dir().ok()?, dirs::home_dir().as_deref()) +} + +enum ProjectBoundary { + Bt(PathBuf), + Git(PathBuf), + Home, + Root, +} - loop { - if current_dir.join(".bt").is_dir() { - return Some(current_dir.join(".bt")); +fn project_boundary(start: PathBuf, home: Option<&Path>) -> ProjectBoundary { + for dir in start.ancestors() { + if Some(dir) == home { + return ProjectBoundary::Home; } - if current_dir.join(".git").exists() { - return None; + if dir.parent().is_none() { + return ProjectBoundary::Root; } - if Some(¤t_dir) == home.as_ref() { - return None; + let bt = dir.join(".bt"); + if bt.is_dir() { + return ProjectBoundary::Bt(bt); } - if !current_dir.pop() { - return None; + if dir.join(".git").exists() { + return ProjectBoundary::Git(dir.to_path_buf()); } } + unreachable!("path ancestors always include a filesystem root") } -pub fn local_path() -> Option { - find_local_config_dir().map(|dir| dir.join("config.json")) -} - -pub enum WriteTarget { - Global(PathBuf), - Local(PathBuf), +fn find_local_config_dir_from(current_dir: PathBuf, home: Option<&Path>) -> Option { + match project_boundary(current_dir, home) { + ProjectBoundary::Bt(dir) if dir.join("config.json").is_file() => Some(dir), + _ => None, + } } -pub fn write_target() -> Result { - match local_path() { - Some(p) => Ok(WriteTarget::Local(p)), - None => Ok(WriteTarget::Global(global_path()?)), - } +pub fn local_path() -> Option { + find_local_config_dir().map(|dir| dir.join("config.json")) } /// Resolve which config file to write based on --global/--local flags. pub fn resolve_write_path(global: bool, local: bool) -> Result { if global { - global_path() - } else if local { - match local_path() { - Some(p) => Ok(p), - None => { - bail!("No local .bt directory found. Use bt init to initialize this directory.") - } + return global_path(); + } + match local_path() { + Some(path) => Ok(path), + None if local => { + bail!("No existing local .bt/config.json found. Run `bt init` first, or use --global.") } - } else { - match write_target()? { - WriteTarget::Local(p) | WriteTarget::Global(p) => Ok(p), + None => global_path(), + } +} + +/// Resolve the create/overwrite target for `bt init`. +pub fn init_target(here: bool, force: bool) -> Result { + init_target_from( + std::env::current_dir().context("could not read current directory")?, + dirs::home_dir().as_deref(), + here, + force, + ) +} + +fn init_target_from( + current_dir: PathBuf, + home: Option<&Path>, + here: bool, + force: bool, +) -> Result { + if here { + let path = current_dir.join(".bt/config.json"); + if path.exists() && !force { + bail!( + "{} already exists; rerun with --force to overwrite it", + path.display() + ); } + return Ok(path); + } + + let path = match project_boundary(current_dir, home) { + ProjectBoundary::Home => bail!( + "reached the home directory without finding a project git root; run `bt init` inside a repository, or pass --here" + ), + ProjectBoundary::Root => bail!( + "reached the filesystem root without finding a project git root; run `bt init` inside a repository, or pass --here" + ), + ProjectBoundary::Git(dir) => return Ok(dir.join(".bt/config.json")), + ProjectBoundary::Bt(dir) => dir.join("config.json"), + }; + if !path.is_file() { + bail!( + "found {} without config.json; remove the incomplete .bt directory, then rerun `bt init`", + path.parent().unwrap_or(&path).display() + ); } + if !force { + bail!( + "{} already exists; use `bt switch` to change it, or rerun with --force to overwrite it", + path.display() + ); + } + Ok(path) } pub fn local_save_path() -> Result { @@ -278,15 +378,53 @@ pub fn save_local(config: &Config, create_dir: bool) -> Result { // --- CLI commands --- -#[derive(Debug, Clone, Args)] +#[derive(Debug, Clone, Default, Args)] pub struct ScopeArgs { - /// Apply to global config (~/.config/bt/config.json) + /// Use global config (~/.config/bt/config.json) #[arg(long, short = 'g', conflicts_with = "local")] - global: bool, + pub(crate) global: bool, - /// Apply to local config (.bt/config.json) + /// Use local config (.bt/config.json) #[arg(long, short = 'l')] - local: bool, + pub(crate) local: bool, +} + +fn scope_labels(global: &Path, local: &Path) -> [String; 2] { + [ + format!("Global ({})", global.parent().unwrap_or(global).display()), + format!("Local ({})", local.parent().unwrap_or(local).display()), + ] +} + +type ResolvedScope = (PathBuf, &'static str); + +impl ScopeArgs { + pub(crate) fn preflight(&self, can_prompt: bool) -> Result<()> { + (!can_prompt) + .then(|| self.resolve(false, "")) + .transpose() + .map(drop) + } + + pub(crate) fn resolve(&self, can_prompt: bool, prompt: &str) -> Result { + if self.global || self.local { + let scope = if self.global { "global" } else { "local" }; + return resolve_write_path(self.global, self.local).map(|path| (path, scope)); + } + let Some(local) = local_path() else { + return Ok((global_path()?, "global")); + }; + if !can_prompt { + bail!("both global and local config scopes are available; pass --global or --local"); + } + let global = global_path()?; + let options = scope_labels(&global, &local); + Ok(if crate::ui::fuzzy_select(prompt, &options, 1)? == 0 { + (global, "global") + } else { + (local, "local") + }) + } } #[derive(Debug, Clone, Args)] @@ -367,82 +505,78 @@ mod tests { use tempfile::TempDir; #[test] - fn merge_other_takes_precedence() { - let base = Config { - org: Some("base-org".into()), - project: Some("base-proj".into()), - ..Default::default() - }; - let other = Config { - org: Some("other-org".into()), - project: Some("other-proj".into()), + fn merge_keeps_org_and_project_contexts_together() { + let c = |org: Option<&str>, project: Option<&str>, id: Option<&str>| Config { + org: org.map(str::to_string), + project: project.map(str::to_string), + project_id: id.map(str::to_string), ..Default::default() }; - let merged = base.merge(&other); - assert_eq!(merged.org, Some("other-org".into())); - assert_eq!(merged.project, Some("other-proj".into())); + let g = || c(Some("global"), Some("global-proj"), Some("proj_g")); + let cases = [ + (Config::default(), Config::default(), Config::default()), + ( + g(), + c(Some("other"), Some("other-proj"), None), + c(Some("other"), Some("other-proj"), None), + ), + ( + c(Some("base"), None, None), + c(None, Some("local"), None), + c(None, Some("local"), None), + ), + (g(), c(Some("global"), None, None), g()), + ( + g(), + c(Some("local"), None, None), + c(Some("local"), None, None), + ), + ( + g(), + c(None, Some("local"), Some("proj_l")), + c(None, Some("local"), Some("proj_l")), + ), + (g(), c(Some(""), None, None), c(Some(""), None, None)), + (g(), Config::default(), g()), + ]; + for (global, local, expected) in cases { + assert_eq!(global.merge(&local), expected); + } } #[test] - fn merge_self_fills_when_other_none() { - let base = Config { - org: Some("base-org".into()), - project: Some("base-proj".into()), - ..Default::default() - }; - let other = Config::default(); - let merged = base.merge(&other); - assert_eq!(merged.org, Some("base-org".into())); - assert_eq!(merged.project, Some("base-proj".into())); + fn scope_labels_are_plain_text() { + let labels = scope_labels( + Path::new("/home/test-user/.config/bt/config.json"), + Path::new("/work/test-project/.bt/config.json"), + ); + assert_eq!(labels[1], "Local (/work/test-project/.bt)"); + assert!(labels.iter().all(|label| !label.contains('\u{1b}'))); } #[test] - fn merge_both_none_stays_none() { - let base = Config::default(); - let other = Config::default(); - let merged = base.merge(&other); - assert_eq!(merged.org, None); - assert_eq!(merged.project, None); - } + fn option_helpers_handle_empty_values() { + for (input, org, trimmed) in [ + (None, None, None), + (Some(""), Some(""), None), + (Some(" "), Some(""), None), + (Some("test-org"), Some("test-org"), Some("test-org")), + ] { + assert_eq!(org_option(input), org); + assert_eq!(trimmed_option(input), trimmed); + } - #[test] - fn merge_partial_fill() { - let base = Config { - org: Some("base-org".into()), - project: None, - ..Default::default() - }; - let other = Config { - org: None, - project: Some("other-proj".into()), - ..Default::default() - }; - let merged = base.merge(&other); - assert_eq!(merged.org, Some("base-org".into())); - assert_eq!(merged.project, Some("other-proj".into())); + let mut cfg = Config::default(); + cfg.set_context(Some(" test-org "), Some(("test-project", "proj_test"))); + assert_eq!(cfg.org.as_deref(), Some("test-org")); + assert_eq!(cfg.project.as_deref(), Some("test-project")); + assert_eq!(cfg.project_id.as_deref(), Some("proj_test")); + cfg.set_context(Some(""), None); + assert_eq!((cfg.org.as_deref(), cfg.project), (Some(""), None)); } fn base_args() -> BaseArgs { - BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - no_color: false, - no_input: false, - org_name: None, - org_name_source: None, - project: None, - project_source: None, - api_key: None, - api_key_source: None, - prefer_api_key: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn config(org: Option<&str>, project: Option<&str>) -> Config { @@ -454,17 +588,18 @@ mod tests { } #[test] - fn project_config_matches_org() { + fn project_config_must_match_org_context() { let base = base_args(); - let cases = [ - (config(Some("acme"), Some("demo")), Some("demo")), - (config(Some("other"), Some("demo")), None), - (config(None, Some("demo")), Some("demo")), - ]; - - for (cfg, expected) in cases { + for (config_org, resolved_org, expected) in [ + (Some("test-org"), "test-org", Some("test-project")), + (Some("other-org"), "test-org", None), + (None, "test-org", None), + (Some(""), "test-org", None), + (Some(""), "", Some("test-project")), + ] { + let cfg = config(config_org, Some("test-project")); assert_eq!( - project_from_config_for_context(&base, &cfg, Some("acme")).as_deref(), + project_from_config_for_context(&base, &cfg, Some(resolved_org)).as_deref(), expected ); } @@ -537,6 +672,24 @@ mod tests { assert!(!persisted.contains("profile")); } + #[test] + fn load_folds_cross_org_alias_to_empty_marker() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("config.json"); + // Literal "cross-org" must load identically to the "" marker. + for spelling in [ + r#"{"org":"cross-org"}"#, + r#"{"org":" cross-org "}"#, + r#"{"org":""}"#, + ] { + fs::write(&path, spelling).unwrap(); + assert_eq!(load_file(&path).org.as_deref(), Some(""), "{spelling}"); + } + + fs::write(&path, r#"{"org":"test-org"}"#).unwrap(); + assert_eq!(load_file(&path).org.as_deref(), Some("test-org")); + } + #[test] fn unknown_keys_roundtrip_through_save() { let tmp = TempDir::new().unwrap(); @@ -569,4 +722,113 @@ mod tests { save_file(&path, &config).unwrap(); assert!(path.exists()); } + + #[test] + fn local_discovery_requires_config_json_and_stops_at_first_bt() { + let tmp = TempDir::new().unwrap(); + let repo = tmp.path().join("repo"); + let nested = repo.join("a").join("b"); + fs::create_dir_all(&nested).unwrap(); + fs::create_dir(repo.join(".git")).unwrap(); + fs::create_dir(repo.join(".bt")).unwrap(); + + assert_eq!(find_local_config_dir_from(nested.clone(), None), None); + + fs::write(repo.join(".bt/config.json"), "{}").unwrap(); + assert_eq!( + find_local_config_dir_from(nested, None), + Some(repo.join(".bt")) + ); + } + + #[test] + fn local_discovery_does_not_use_home_bt() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("home"); + fs::create_dir_all(home.join(".bt")).unwrap(); + fs::write(home.join(".bt/config.json"), "{}").unwrap(); + + assert_eq!( + find_local_config_dir_from(home.clone(), Some(home.as_path())), + None + ); + } + + #[test] + fn init_target_finds_nested_git_directory_or_file() { + for git_is_file in [false, true] { + let tmp = TempDir::new().unwrap(); + let repo = tmp.path().join("repo"); + let nested = repo.join("nested").join("deeper"); + fs::create_dir_all(&nested).unwrap(); + if git_is_file { + fs::write(repo.join(".git"), "gitdir: synthetic").unwrap(); + } else { + fs::create_dir(repo.join(".git")).unwrap(); + } + + assert_eq!( + init_target_from(nested, Some(tmp.path()), false, false).unwrap(), + repo.join(".bt/config.json") + ); + } + } + + #[test] + fn init_target_existing_bt_requires_force_and_existing_config() { + let tmp = TempDir::new().unwrap(); + let repo = tmp.path().join("repo"); + let nested = repo.join("nested"); + fs::create_dir_all(repo.join(".bt")).unwrap(); + fs::create_dir_all(&nested).unwrap(); + + assert!(init_target_from(nested.clone(), Some(tmp.path()), false, true).is_err()); + + let target = repo.join(".bt/config.json"); + fs::write(&target, "{}").unwrap(); + assert!(init_target_from(nested.clone(), Some(tmp.path()), false, false).is_err()); + assert_eq!( + init_target_from(nested, Some(tmp.path()), false, true).unwrap(), + target + ); + } + + #[test] + fn init_target_here_bypasses_home_boundary_and_honors_force() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("home"); + fs::create_dir_all(&home).unwrap(); + let target = home.join(".bt/config.json"); + + assert_eq!( + init_target_from(home.clone(), Some(home.as_path()), true, false).unwrap(), + target + ); + fs::create_dir_all(target.parent().unwrap()).unwrap(); + fs::write(&target, "{}").unwrap(); + assert!(init_target_from(home.clone(), Some(home.as_path()), true, false).is_err()); + assert_eq!( + init_target_from(home, Some(tmp.path()), true, true).unwrap(), + target + ); + } + + #[test] + fn init_target_home_wins_over_git_marker() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("home"); + fs::create_dir_all(home.join(".git")).unwrap(); + assert!(init_target_from(home.clone(), Some(home.as_path()), false, false).is_err()); + } + + #[cfg(unix)] + #[test] + fn init_target_here_bypasses_filesystem_root_boundary() { + let root = PathBuf::from("/"); + assert_eq!( + init_target_from(root.clone(), None, true, true).unwrap(), + root.join(".bt/config.json") + ); + assert!(init_target_from(root, None, false, false).is_err()); + } } diff --git a/src/datasets/pipeline.rs b/src/datasets/pipeline.rs index cd726927..a00b5618 100644 --- a/src/datasets/pipeline.rs +++ b/src/datasets/pipeline.rs @@ -2065,26 +2065,7 @@ mod tests { } fn test_base_args() -> BaseArgs { - BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - no_color: false, - no_input: false, - org_name: None, - org_name_source: None, - project: None, - project_source: None, - api_key: None, - api_key_source: None, - prefer_api_key: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn test_source() -> PipelineSourceInspect { diff --git a/src/datasets/snapshots.rs b/src/datasets/snapshots.rs index 047f2d3e..337dd74e 100644 --- a/src/datasets/snapshots.rs +++ b/src/datasets/snapshots.rs @@ -1093,7 +1093,9 @@ fn resolve_default_snapshot_author(base: &BaseArgs, ctx: &ResolvedContext) -> Op return None; } - let profile = auth::active_auth_info(base, Some(ctx.client.org_name()))?; + let profile = auth::active_auth_info(base, Some(ctx.client.org_name())) + .ok() + .flatten()?; profile_author_slug(&profile) } diff --git a/src/eval.rs b/src/eval.rs index d279951a..ab6bfc79 100644 --- a/src/eval.rs +++ b/src/eval.rs @@ -35,6 +35,7 @@ use crate::ui::{ summary_metric_unit, SummaryExperimentColumn, SummaryMetricCell, SummaryMetricKind, SummaryMetricRow, SummaryTableOptions, }; +use crate::utils::shell_quote_arg; const MAX_NAME_LENGTH: usize = 40; const WATCH_POLL_INTERVAL: Duration = Duration::from_millis(500); @@ -3401,17 +3402,6 @@ fn build_experiment_compare_command( )) } -fn shell_quote_arg(value: &str) -> String { - if value - .chars() - .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.' | '/' | ':' | '=')) - { - value.to_string() - } else { - format!("'{}'", value.replace('\'', "'\\''")) - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/functions/push.rs b/src/functions/push.rs index b4d1f1fb..958ded96 100644 --- a/src/functions/push.rs +++ b/src/functions/push.rs @@ -13,8 +13,7 @@ use indicatif::{ProgressBar, ProgressStyle}; use serde::Deserialize; use serde_json::{json, Map, Value}; -use crate::args::BaseArgs; -use crate::args::DEFAULT_API_URL; +use crate::args::{custom_api_without_app_url, BaseArgs}; use crate::auth::{list_available_orgs, resolve_auth}; use crate::config; @@ -270,21 +269,10 @@ pub async fn run(base: BaseArgs, args: PushArgs) -> Result<()> { ); } }; - let has_app_url = resolved_auth - .app_url - .as_deref() - .map(str::trim) - .is_some_and(|value| !value.is_empty()); - let custom_api_without_app_url = resolved_auth - .api_url - .as_deref() - .map(str::trim) - .map(|value| value.trim_end_matches('/')) - .is_some_and(|api_url| { - !api_url.eq_ignore_ascii_case(DEFAULT_API_URL.trim_end_matches('/')) - }) - && !has_app_url; - if custom_api_without_app_url { + if custom_api_without_app_url( + resolved_auth.api_url.as_deref(), + resolved_auth.app_url.as_deref(), + ) { return fail_push( &base, 0, @@ -4055,25 +4043,6 @@ mod tests { } fn test_base_args() -> BaseArgs { - BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - no_color: false, - no_input: false, - org_name: None, - org_name_source: None, - project: None, - project_source: None, - api_key: None, - api_key_source: None, - prefer_api_key: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } } diff --git a/src/init.rs b/src/init.rs index c03594a6..ae57e137 100644 --- a/src/init.rs +++ b/src/init.rs @@ -1,4 +1,4 @@ -use anyhow::{bail, Result}; +use anyhow::{bail, Context, Result}; use clap::Args; use crate::{ @@ -6,109 +6,90 @@ use crate::{ auth::{self, login}, config, http::ApiClient, - ui::{is_interactive, print_command_status, select_project, CommandStatus, ProjectSelectMode}, + ui::{print_command_status, select_or_create_project, CommandStatus}, }; #[derive(Debug, Clone, Args)] #[command(after_help = "\ Examples: bt init - bt init --org acme --project my-app + bt init --org test-org --project test-project + bt init --here + bt init --here --force ")] -pub struct InitArgs {} +pub struct InitArgs { + /// Create .bt/config.json in the current directory without searching upward. + /// + /// Bypasses the normal home and filesystem-root search boundaries, so it + /// also applies when the current directory is ~ or /. + #[arg(long)] + here: bool, -fn select_saved_auth_org_for_init() -> Result> { - let mut orgs = auth::list_profiles()? - .into_iter() - .filter_map(|profile| profile.org_name) - .filter(|org| !org.trim().is_empty()) - .collect::>(); - orgs.sort(); - orgs.dedup(); - - match orgs.len() { - 0 => Ok(None), - 1 => Ok(orgs.into_iter().next()), - _ => { - let labels = orgs.iter().map(String::as_str).collect::>(); - let idx = crate::ui::fuzzy_select("Select organization", &labels, 0)?; - Ok(Some(orgs[idx].clone())) - } - } + /// Overwrite an existing .bt/config.json. Does not change discovery. + #[arg(long, short = 'f')] + force: bool, } -pub async fn run(base: BaseArgs, _args: InitArgs) -> Result<()> { - let config_path = config::local_save_path()?; - if config_path.exists() { - if base.json { - let existing = config::load_file(&config_path); - let payload = serde_json::json!({ - "initialized": false, - "status": "already-initialized", - "org": existing.org, - "project": existing.project, - "path": config_path.display().to_string(), - }); - println!("{}", serde_json::to_string(&payload)?); - } else { - print_command_status(CommandStatus::Warning, "Already Initialized"); - } - return Ok(()); +pub async fn run(base: BaseArgs, args: InitArgs) -> Result<()> { + let config_path = config::init_target(args.here, args.force)?; + let current_cfg = config::load().unwrap_or_default(); + let mut login_base = base.clone(); + login_base.project = None; + login_base.project_source = None; + if login_base.org_name.is_none() + && !auth::select_saved_login(&mut login_base, current_cfg.org.as_deref(), false)? + { + bail!("no saved concrete-org login is available; run `bt auth login --org `"); } - eprintln!("Link to a Braintrust project..."); + let ctx = login(&login_base).await?; + let client = ApiClient::new(&ctx)?; + let org = client.org_name().to_string(); + if org.is_empty() { + bail!( + "cross-org mode has no project; `bt init` is project-scoped. Rerun with --org --project " + ); + } - let (org, project) = if let (Some(o), Some(p)) = (&base.org_name, &base.project) { - (o.clone(), p.clone()) - } else if !is_interactive() { - bail!("--org and --project required in non-interactive mode"); - } else { - let mut login_base = base.clone(); - if login_base.org_name.is_none() { - login_base.org_name = config::load().ok().and_then(|cfg| cfg.org); - } - if login_base.org_name.is_none() { - login_base.org_name = select_saved_auth_org_for_init()?; - } - let ctx = login(&login_base).await?; - let client = ApiClient::new(&ctx)?; + let project = select_or_create_project( + &client, + base.project.as_deref(), + None, + Some("Link to project"), + ) + .await?; + let mut cfg = config::Config::default(); + cfg.set_context( + Some(&org), + Some((project.name.as_str(), project.id.as_str())), + ); - let org = client.org_name().to_string(); - let project = select_project( - &client, - None, - Some("Link to project"), - ProjectSelectMode::ExistingOnly, + config::save_file(&config_path, &cfg).with_context(|| { + format!( + "authentication succeeded, but initialization failed: could not create or write {}; any credential updates remain saved", + config_path.display() ) - .await? - .name; - - (org, project) - }; - - let cfg = config::Config { - org: Some(org.clone()), - project: Some(project.clone()), - ..Default::default() - }; - - let written_path = config::save_local(&cfg, true)?; + })?; if base.json { let payload = serde_json::json!({ "initialized": true, "status": "created", "org": org, - "project": project, - "path": written_path.display().to_string(), + "project": project.name, + "project_id": project.id, + "path": config_path.display().to_string(), }); println!("{}", serde_json::to_string(&payload)?); } else { print_command_status( CommandStatus::Success, - &format!("Project linked to {org}/{project}"), + &format!("Project linked to {org}/{}", project.name), + ); + print_command_status( + CommandStatus::Success, + &format!("Created {}", config_path.display()), ); - print_command_status(CommandStatus::Success, "Created .bt/config.json"); } Ok(()) diff --git a/src/main.rs b/src/main.rs index 7d0082a5..035cc806 100644 --- a/src/main.rs +++ b/src/main.rs @@ -491,11 +491,16 @@ fn has_io_error(err: &anyhow::Error) -> bool { } fn looks_like_user_error(err: &anyhow::Error) -> bool { - let message = err.to_string().to_lowercase(); - message.contains("required") - || message.contains("use:") - || message.contains("not found") - || message.contains("invalid") + err.chain().any(|source| { + let message = source.to_string().to_lowercase(); + message.contains("required") + || message.contains("use:") + || message.contains("not found") + || message.contains("invalid") + || message.contains("already exists") + || message.contains("without finding") + || message.contains("without config.json") + }) } fn print_error(err: &anyhow::Error, code: ExitCode, missing_credential: bool) { diff --git a/src/setup/mod.rs b/src/setup/mod.rs index b98e79f4..17937d3e 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -5187,26 +5187,7 @@ mod tests { } fn make_base_args() -> BaseArgs { - BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - no_color: false, - no_input: false, - org_name: None, - org_name_source: None, - project: None, - project_source: None, - api_key: None, - api_key_source: None, - prefer_api_key: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn restore_env_var(key: &str, previous: Option) { diff --git a/src/status.rs b/src/status.rs index e6073031..756089ff 100644 --- a/src/status.rs +++ b/src/status.rs @@ -30,22 +30,14 @@ struct StatusOutput { source: Option, } -fn format_identity(p: &auth::ProfileInfo) -> Option { - if let Some(ref email) = p.email { - match p.user_name.as_deref() { - Some(name) => Some(format!("{name} ({email})")), - None => Some(email.clone()), - } - } else { - p.api_key_hint.clone() - } -} - fn format_auth(p: &auth::ProfileInfo) -> String { - match format_identity(p) { - Some(identity) => format!("{} — {identity}", p.auth_method), - None => p.auth_method.clone(), - } + auth::identity_label( + p.user_name.as_deref(), + p.email.as_deref(), + p.api_key_hint.as_deref(), + ) + .map(|identity| format!("{} — {identity}", p.auth_method)) + .unwrap_or_else(|| p.auth_method.clone()) } pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { @@ -66,7 +58,7 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { &global_path, ); let merged_cfg = global_cfg.merge(&local_cfg); - let auth_info = auth::active_auth_info(&base, org.as_deref()); + let auth_info = auth::active_auth_info(&base, org.as_deref())?; if base .project @@ -77,9 +69,10 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { project = config::project_from_config_for_context(&base, &merged_cfg, org.as_deref()); } + let display_org = org.as_deref().map(config::display_org); if base.json { let output = StatusOutput { - org, + org: display_org.map(str::to_string), project, user_name: auth_info.as_ref().and_then(|p| p.user_name.clone()), user_email: auth_info.as_ref().and_then(|p| p.email.clone()), @@ -92,16 +85,7 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { } if base.verbose { - let org_display = if org.is_none() - && auth_info - .as_ref() - .is_some_and(|p| p.auth_method == "oauth" && p.org_name.is_none()) - { - "cross-org" - } else { - org.as_deref().unwrap_or("(unset)") - }; - println!("org: {org_display}"); + println!("org: {}", display_org.unwrap_or("(unset)")); println!("project: {}", project.as_deref().unwrap_or("(unset)")); if let Some(ref p) = auth_info { println!("auth: {}", format_auth(p)); @@ -109,11 +93,11 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { if let Some(src) = source { println!("source: {src}"); } - } else if org.is_some() { - let scope = match (&org, &project) { - (Some(o), Some(p)) => format!("{o}/{p}"), - (Some(o), None) => o.to_string(), - _ => unreachable!(), + } else if let Some(org) = display_org { + let scope = match (org, project.as_deref()) { + ("cross-org", _) => org.to_string(), + (org, Some(project)) => format!("{org}/{project}"), + (org, None) => org.to_string(), }; println!("{scope}"); let auth_line = match &auth_info { @@ -182,20 +166,19 @@ pub(crate) fn resolve_config( cli_project, env_project, } = overrides; - let env_org = env_org.filter(|s| !s.is_empty()); + // `Some("")` is the canonical cross-org marker for both CLI and env + // sources, so org overrides must not filter empty strings. let env_project = env_project.filter(|s| !s.is_empty()); - + let merged = global.merge(local); let org = cli_org .clone() .or_else(|| env_org.clone()) - .or_else(|| local.org.clone()) - .or_else(|| global.org.clone()); + .or_else(|| merged.org.clone()); let project = cli_project .clone() .or_else(|| env_project.clone()) - .or_else(|| local.project.clone()) - .or_else(|| global.project.clone()); + .or_else(|| merged.project.clone()); let source = if cli_org.is_some() || cli_project.is_some() { Some("cli".to_string()) @@ -230,199 +213,108 @@ mod tests { } #[test] - fn cli_overrides_everything() { - let global = config(Some("global-org"), Some("global-proj")); - let local = config(Some("local-org"), Some("local-proj")); - let local_path = Some(PathBuf::from("/project/.bt/config.json")); - let global_path = Some(PathBuf::from("/home/.bt/config.json")); - - let (org, project, source) = resolve_config( - ConfigOverrides { - cli_org: s("cli-org"), - cli_project: s("cli-proj"), - ..Default::default() - }, - &global, - &local, - &local_path, - &global_path, - ); - - assert_eq!(org, s("cli-org")); - assert_eq!(project, s("cli-proj")); - assert_eq!(source, s("cli")); - } - - #[test] - fn env_overrides_config_below_cli() { - let global = config(Some("global-org"), Some("global-proj")); - let local = config(Some("local-org"), Some("local-proj")); - let local_path = Some(PathBuf::from("/project/.bt/config.json")); - let global_path = Some(PathBuf::from("/home/.bt/config.json")); - - let (org, project, source) = resolve_config( - ConfigOverrides { - env_org: s("env-org"), - env_project: s("env-proj"), - ..Default::default() - }, - &global, - &local, - &local_path, - &global_path, - ); - - assert_eq!(org, s("env-org")); - assert_eq!(project, s("env-proj")); - assert_eq!(source, s("env")); - } - - #[test] - fn local_overrides_global() { - let global = config(Some("global-org"), Some("global-proj")); - let local = config(Some("local-org"), Some("local-proj")); - let local_path = Some(PathBuf::from("/project/.bt/config.json")); - let global_path = Some(PathBuf::from("/home/.bt/config.json")); - - let (org, project, source) = resolve_config( - ConfigOverrides::default(), - &global, - &local, - &local_path, - &global_path, - ); - - assert_eq!(org, s("local-org")); - assert_eq!(project, s("local-proj")); - assert_eq!(source, s("/project/.bt/config.json")); - } - - #[test] - fn global_used_when_local_empty() { - let global = config(Some("global-org"), Some("global-proj")); - let local = config(None, None); - let local_path = Some(PathBuf::from("/project/.bt/config.json")); - let global_path = Some(PathBuf::from("/home/.bt/config.json")); - - let (org, project, source) = resolve_config( - ConfigOverrides::default(), - &global, - &local, - &local_path, - &global_path, - ); - - assert_eq!(org, s("global-org")); - assert_eq!(project, s("global-proj")); - assert_eq!(source, s("/home/.bt/config.json")); - } - - #[test] - fn no_source_when_all_empty() { - let global = config(None, None); - let local = config(None, None); - let local_path = Some(PathBuf::from("/project/.bt/config.json")); - let global_path = Some(PathBuf::from("/home/.bt/config.json")); - - let (org, project, source) = resolve_config( - ConfigOverrides::default(), - &global, - &local, - &local_path, - &global_path, - ); - - assert_eq!(org, None); - assert_eq!(project, None); - assert_eq!(source, None); - } - - #[test] - fn mixed_sources_org_cli_project_local() { - let global = config(Some("global-org"), Some("global-proj")); - let local = config(None, Some("local-proj")); - let local_path = Some(PathBuf::from("/project/.bt/config.json")); - let global_path = Some(PathBuf::from("/home/.bt/config.json")); - - let (org, project, source) = resolve_config( - ConfigOverrides { - cli_org: s("cli-org"), - ..Default::default() - }, - &global, - &local, - &local_path, - &global_path, - ); - - assert_eq!(org, s("cli-org")); - assert_eq!(project, s("local-proj")); - assert_eq!(source, s("cli")); - } - - #[test] - fn values_cascade_across_layers() { - let global = config(Some("global-org"), None); - let local = config(None, Some("local-proj")); + fn config_precedence_and_context_safety() { + let both = || config(Some("global-org"), Some("global-proj")); + let cases = [ + ( + "cli", + ConfigOverrides { + cli_org: s("cli-org"), + cli_project: s("cli-proj"), + ..Default::default() + }, + both(), + config(Some("local-org"), Some("local-proj")), + (Some("cli-org"), Some("cli-proj"), Some("cli")), + ), + ( + "env", + ConfigOverrides { + env_org: s("env-org"), + env_project: s("env-proj"), + ..Default::default() + }, + both(), + config(Some("local-org"), Some("local-proj")), + (Some("env-org"), Some("env-proj"), Some("env")), + ), + ( + "local", + ConfigOverrides::default(), + both(), + config(Some("local-org"), Some("local-proj")), + ( + Some("local-org"), + Some("local-proj"), + Some("/project/.bt/config.json"), + ), + ), + ( + "global", + ConfigOverrides::default(), + both(), + config(None, None), + ( + Some("global-org"), + Some("global-proj"), + Some("/home/.bt/config.json"), + ), + ), + ( + "empty", + ConfigOverrides::default(), + config(None, None), + config(None, None), + (None, None, None), + ), + ( + "mixed cli/local", + ConfigOverrides { + cli_org: s("cli-org"), + ..Default::default() + }, + both(), + config(None, Some("local-proj")), + (Some("cli-org"), Some("local-proj"), Some("cli")), + ), + ( + "local project does not inherit global org", + ConfigOverrides::default(), + config(Some("global-org"), None), + config(None, Some("local-proj")), + (None, Some("local-proj"), Some("/project/.bt/config.json")), + ), + ( + "local cross-org", + ConfigOverrides::default(), + both(), + config(Some(""), None), + (Some(""), None, Some("/project/.bt/config.json")), + ), + ( + "env cross-org", + ConfigOverrides { + env_org: s(""), + ..Default::default() + }, + both(), + config(None, None), + (Some(""), Some("global-proj"), Some("env")), + ), + ]; let local_path = Some(PathBuf::from("/project/.bt/config.json")); let global_path = Some(PathBuf::from("/home/.bt/config.json")); - - let (org, project, source) = resolve_config( - ConfigOverrides::default(), - &global, - &local, - &local_path, - &global_path, - ); - - assert_eq!(org, s("global-org")); - assert_eq!(project, s("local-proj")); - assert_eq!(source, s("/project/.bt/config.json")); - } - - fn profile( - user_name: Option<&str>, - email: Option<&str>, - api_key_hint: Option<&str>, - ) -> auth::ProfileInfo { - auth::ProfileInfo { - auth_method: if api_key_hint.is_some() { - "api_key" - } else { - "oauth" - } - .into(), - org_name: None, - user_name: user_name.map(Into::into), - email: email.map(Into::into), - api_key_hint: api_key_hint.map(Into::into), + for (name, overrides, global, local, expected) in cases { + let actual = resolve_config(overrides, &global, &local, &local_path, &global_path); + assert_eq!( + actual, + ( + expected.0.map(str::to_string), + expected.1.map(str::to_string), + expected.2.map(str::to_string), + ), + "{name}" + ); } } - - #[test] - fn format_identity_name_and_email() { - let p = profile(Some("Alice"), Some("alice@example.com"), None); - assert_eq!( - format_identity(&p), - Some("Alice (alice@example.com)".into()) - ); - } - - #[test] - fn format_identity_email_only() { - let p = profile(None, Some("alice@example.com"), None); - assert_eq!(format_identity(&p), Some("alice@example.com".into())); - } - - #[test] - fn format_identity_api_key_hint() { - let p = profile(None, None, Some("sk-****zhJwO")); - assert_eq!(format_identity(&p), Some("sk-****zhJwO".into())); - } - - #[test] - fn format_identity_none() { - let p = profile(None, None, None); - assert_eq!(format_identity(&p), None); - } } diff --git a/src/switch.rs b/src/switch.rs index 40fed377..05419919 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -1,30 +1,24 @@ use anyhow::{bail, Context, Result}; use clap::Args; -use dialoguer::{console, theme::ColorfulTheme, Select}; use crate::args::BaseArgs; use crate::auth::{self, login}; use crate::config; use crate::http::ApiClient; -use crate::projects::api; -use crate::ui::{ - is_interactive, print_command_status, select_project, with_spinner, CommandStatus, -}; +use crate::ui::{can_prompt, print_command_status, select_or_create_project, CommandStatus}; #[derive(Debug, Clone, Args)] #[command(after_help = "\ Examples: bt switch - bt switch my-project - bt switch personal-org/my-project + bt switch test-project + bt switch test-org/test-project + bt switch --org cross-org ")] pub struct SwitchArgs { - /// Force set global config value - #[arg(long, short = 'g', conflicts_with = "local")] - global: bool, - /// Force set local config value - #[arg(long, short = 'l')] - local: bool, + #[command(flatten)] + scope: config::ScopeArgs, + /// Target: project name or org/project #[arg(value_name = "TARGET")] target: Option, @@ -51,66 +45,70 @@ impl SwitchArgs { } pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { - let current_cfg = config::load().unwrap_or_default(); + args.scope.preflight(can_prompt())?; + let current_cfg = if args.scope.global { + config::load_global().unwrap_or_default() + } else { + config::load().unwrap_or_default() + }; let (resolved_org, resolved_project) = args.resolve_target(&base); - let mut interactive = false; + let bare_switch = resolved_org.is_none() && resolved_project.is_none(); let mut login_base = base.clone(); - if login_base.org_name.is_none() { - login_base.org_name = resolved_org.clone().or_else(|| current_cfg.org.clone()); - if login_base.org_name.is_none() { - login_base.org_name = select_saved_auth_org_for_switch()?; - } + login_base.org_name = resolved_org.clone(); + login_base.project = None; + login_base.project_source = None; + + if login_base.org_name.is_none() && !bare_switch { + login_base.org_name = current_cfg.org.clone(); + } + if login_base.org_name.is_none() + && !auth::select_saved_login(&mut login_base, current_cfg.org.as_deref(), true)? + { + bail!("no saved auth logins found; run `bt auth login` to create one"); + } + + if login_base.org_name.as_deref() == Some("") && resolved_project.is_some() { + bail!( + "cross-org mode cannot have a default project; rerun with --org --project " + ); } let ctx = login(&login_base).await?; let client = ApiClient::new(&ctx)?; let org_name = client.org_name().to_string(); - let project = match resolved_project { - Some(p) => validate_or_create_project(&client, &p).await?, - None => { - if !is_interactive() { - bail!("target required. Use: bt switch or bt switch /"); - } - interactive = true; - select_project( + let project = if org_name.is_empty() { + None + } else { + Some( + select_or_create_project( &client, + resolved_project.as_deref(), + current_cfg.project.as_deref(), None, - None, - crate::ui::ProjectSelectMode::ExistingOnly, ) - .await? - } - }; - - let (path, scope) = if args.local { - ( - config::local_path().ok_or_else(|| { - anyhow::anyhow!( - "No local .bt directory found. Use bt init to initialize this directory." - ) - })?, - "local", + .await?, ) - } else if args.global { - (config::global_path()?, "global") - } else if interactive && config::local_path().is_some() { - select_scope()? - } else { - (config::global_path()?, "global") }; + // Scope is prompted last, after org and project. + let (path, scope) = args.scope.resolve(can_prompt(), "Save to")?; let mut cfg = config::load_file(&path); - apply_switch_config(&mut cfg, Some(&org_name), Some(&project)); + cfg.set_context( + Some(&org_name), + project + .as_ref() + .map(|project| (project.name.as_str(), project.id.as_str())), + ); config::save_file(&path, &cfg) - .context(format!("Could not save config to {}", path.display()))?; + .with_context(|| format!("Could not save config to {}", path.display()))?; if base.json { let payload = serde_json::json!({ - "org": org_name, - "project": project.name, - "project_id": project.id, + "org": config::display_org(&org_name), + "project": project.as_ref().map(|p| p.name.clone()), + "project_id": project.as_ref().map(|p| p.id.clone()), "scope": scope, "path": path.display().to_string(), }); @@ -118,7 +116,10 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { return Ok(()); } - let display = format!("{org_name}/{}", project.name); + let display = project + .as_ref() + .map(|project| format!("{org_name}/{}", project.name)) + .unwrap_or_else(|| config::display_org(&org_name).to_string()); print_command_status(CommandStatus::Success, &format!("Switched to {display}")); if base.verbose { eprintln!("Wrote to {}", path.display()); @@ -127,281 +128,75 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { Ok(()) } -fn select_saved_auth_org_for_switch() -> Result> { - if !is_interactive() { - return Ok(None); - } - - let mut orgs = auth::list_profiles()? - .into_iter() - .filter_map(|profile| profile.org_name) - .filter(|org| !org.trim().is_empty()) - .collect::>(); - orgs.sort(); - orgs.dedup(); - - match orgs.len() { - 0 => Ok(None), - 1 => Ok(orgs.into_iter().next()), - _ => { - let labels = orgs.iter().map(String::as_str).collect::>(); - let idx = crate::ui::fuzzy_select("Select organization", &labels, 0)?; - Ok(Some(orgs[idx].clone())) - } - } -} - -pub(crate) fn select_scope() -> Result<(std::path::PathBuf, &'static str)> { - let global = config::global_path()?; - let local = config::local_path().unwrap(); - let options = [ - format!( - "Global ({})", - console::style( - dirs::home_dir() - .and_then(|home| global - .parent() - .unwrap() - .strip_prefix(&home) - .ok() - .map(|rel| format!("~/{}", rel.display()))) - .unwrap_or_else(|| global.parent().unwrap().display().to_string()) - ) - .dim() - ), - format!( - "Local ({})", - console::style( - local - .parent() - .and_then(|bt| { - let bt_name = bt.file_name()?; - let parent_name = bt.parent()?.file_name()?; - Some(format!( - "{}/{}", - parent_name.to_string_lossy(), - bt_name.to_string_lossy() - )) - }) - .unwrap_or_else(|| local.parent().unwrap().display().to_string()), - ) - .dim() - ), - ]; - let idx = Select::with_theme(&ColorfulTheme::default()) - .with_prompt("Save to") - .items(&options) - .default(1) - .interact()?; - if idx == 0 { - Ok((global, "global")) - } else { - Ok((local, "local")) - } -} - -pub(crate) async fn validate_or_create_project( - client: &ApiClient, - name: &str, -) -> Result { - let exists = with_spinner("Loading project...", api::get_project_by_name(client, name)).await?; - - if let Some(project) = exists { - return Ok(project); - } - - if !is_interactive() { - bail!("project '{name}' not found"); - } - - let create = dialoguer::Confirm::new() - .with_prompt(format!("Project '{name}' not found. Create it?")) - .default(false) - .interact()?; - - if create { - with_spinner("Creating project...", api::create_project(client, name)).await - } else { - bail!("project '{name}' not found"); - } -} - -pub(crate) fn apply_switch_config( - cfg: &mut config::Config, - org_name: Option<&str>, - project: Option<&api::Project>, -) { - cfg.org = config::trimmed_option(org_name).map(str::to_string); - match project { - Some(project) => { - cfg.project = Some(project.name.clone()); - cfg.project_id = Some(project.id.clone()); - } - None => { - cfg.project = None; - cfg.project_id = None; - } - } -} - #[cfg(test)] mod tests { use super::*; - fn switch_args(target: Option<&str>) -> SwitchArgs { - SwitchArgs { - global: false, - local: false, - target: target.map(String::from), - } - } - - fn base_args(org: Option<&str>, project: Option<&str>) -> BaseArgs { - BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - no_color: false, - no_input: false, - org_name: org.map(String::from), - org_name_source: None, - project: project.map(String::from), - project_source: None, - api_key: None, - api_key_source: None, - prefer_api_key: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } - } - - // --- resolve_target tests (unchanged) --- - - #[test] - fn no_args_returns_none() { - let args = switch_args(None); - let base = base_args(None, None); - assert_eq!(args.resolve_target(&base), (None, None)); - } - - #[test] - fn positional_org_project() { - let args = switch_args(Some("myorg/proj")); - let base = base_args(None, None); - assert_eq!( - args.resolve_target(&base), - (Some("myorg".into()), Some("proj".into())) - ); - } - - #[test] - fn positional_project_only() { - let args = switch_args(Some("proj")); - let base = base_args(None, None); - assert_eq!(args.resolve_target(&base), (None, Some("proj".into()))); - } - - #[test] - fn slash_with_empty_org() { - let args = switch_args(Some("/project")); - let base = base_args(None, None); - assert_eq!(args.resolve_target(&base), (None, Some("project".into()))); - } - - #[test] - fn slash_with_empty_project() { - let args = switch_args(Some("org/")); - let base = base_args(None, None); - assert_eq!(args.resolve_target(&base), (Some("org".into()), None)); - } - - #[test] - fn flag_org_only() { - let args = switch_args(None); - let base = base_args(Some("x"), None); - assert_eq!(args.resolve_target(&base), (Some("x".into()), None)); - } - - #[test] - fn flag_project_only() { - let args = switch_args(None); - let base = base_args(None, Some("y")); - assert_eq!(args.resolve_target(&base), (None, Some("y".into()))); - } - - #[test] - fn flags_only() { - let args = switch_args(None); - let base = base_args(Some("a"), Some("b")); - assert_eq!( - args.resolve_target(&base), - (Some("a".into()), Some("b".into())) - ); - } - #[test] - fn flag_overrides_positional_project() { - let args = switch_args(Some("myorg/proj")); - let base = base_args(None, Some("foo")); - assert_eq!( - args.resolve_target(&base), - (Some("myorg".into()), Some("foo".into())) - ); - } - - #[test] - fn flag_org_with_positional_project() { - let args = switch_args(Some("proj")); - let base = base_args(Some("bar"), None); - assert_eq!( - args.resolve_target(&base), - (Some("bar".into()), Some("proj".into())) - ); - } - - #[test] - fn flag_override_both() { - let args = switch_args(Some("myorg/proj")); - let base = base_args(Some("x"), Some("y")); - assert_eq!( - args.resolve_target(&base), - (Some("x".into()), Some("y".into())) - ); - } - - #[test] - fn apply_switch_config_sets_project_id_with_project_name_and_org() { - let mut cfg = config::Config::default(); - let project = api::Project { - id: "proj_123".to_string(), - name: "my-project".to_string(), - org_id: "org_123".to_string(), - description: None, - }; - - apply_switch_config(&mut cfg, Some("acme-org"), Some(&project)); - - assert_eq!(cfg.org.as_deref(), Some("acme-org")); - assert_eq!(cfg.project.as_deref(), Some("my-project")); - assert_eq!(cfg.project_id.as_deref(), Some("proj_123")); - } - - #[test] - fn apply_switch_config_clears_project_and_org_when_context_is_org_only() { - let mut cfg = config::Config { - org: Some("old-org".to_string()), - project: Some("stale-project".to_string()), - project_id: Some("proj_stale".to_string()), - ..Default::default() - }; - - apply_switch_config(&mut cfg, None, None); - - assert_eq!(cfg.org, None); - assert_eq!(cfg.project, None); - assert_eq!(cfg.project_id, None); + fn resolve_target_combines_positionals_and_flags() { + for (target, org, project, expected) in [ + (None, None, None, (None, None)), + ( + Some("test-org/test-project"), + None, + None, + (Some("test-org"), Some("test-project")), + ), + ( + Some("test-project"), + None, + None, + (None, Some("test-project")), + ), + ( + Some("/test-project"), + None, + None, + (None, Some("test-project")), + ), + (Some("test-org/"), None, None, (Some("test-org"), None)), + (None, Some("test-org"), None, (Some("test-org"), None)), + ( + None, + None, + Some("test-project"), + (None, Some("test-project")), + ), + ( + None, + Some("test-org"), + Some("test-project"), + (Some("test-org"), Some("test-project")), + ), + ( + Some("old-org/old-project"), + None, + Some("test-project"), + (Some("old-org"), Some("test-project")), + ), + ( + Some("old-project"), + Some("test-org"), + None, + (Some("test-org"), Some("old-project")), + ), + ( + Some("old-org/old-project"), + Some("test-org"), + Some("test-project"), + (Some("test-org"), Some("test-project")), + ), + ] { + let args = SwitchArgs { + scope: config::ScopeArgs::default(), + target: target.map(str::to_string), + }; + let base = BaseArgs { + org_name: org.map(str::to_string), + project: project.map(str::to_string), + ..Default::default() + }; + let actual = args.resolve_target(&base); + assert_eq!((actual.0.as_deref(), actual.1.as_deref()), expected); + } } } diff --git a/src/traces.rs b/src/traces.rs index 396b4de0..97a881c9 100644 --- a/src/traces.rs +++ b/src/traces.rs @@ -5260,7 +5260,7 @@ fn logs_hints( has_more: bool, object_ref: &str, next_cursor: Option<&str>, - profile: Option<&str>, + org: Option<&str>, limit: usize, ) -> Vec { let mut hints = vec![ @@ -5277,7 +5277,7 @@ fn logs_hints( if let Some(cursor) = next_cursor { hints.push(format!( "Next page: bt view logs{} --object-ref {object_ref} --cursor {cursor}{limit_suffix}", - org_flag_suffix(profile) + org_flag_suffix(org) )); } else { hints.push(format!( @@ -5293,13 +5293,13 @@ fn trace_hints( object_ref: &str, trace_id: &str, next_cursor: Option<&str>, - profile: Option<&str>, + org: Option<&str>, limit: usize, ) -> Vec { let mut hints = vec![ format!( "Trace fetch returns truncated span rows; use `bt view span{} --object-ref {object_ref} --id ` for full single-span payloads.", - org_flag_suffix(profile) + org_flag_suffix(org) ), "Write output to a file for long traces.".to_string(), ]; @@ -5312,7 +5312,7 @@ fn trace_hints( if let Some(cursor) = next_cursor { hints.push(format!( "Next page: bt view trace{} --object-ref {object_ref} --trace-id {trace_id} --cursor {cursor}{limit_suffix}", - org_flag_suffix(profile) + org_flag_suffix(org) )); } else { hints.push(format!( @@ -5337,21 +5337,21 @@ fn detail_view_trace_url_value(view: DetailView) -> Option<&'static str> { } } -fn thread_hints(profile: Option<&str>) -> Vec { +fn thread_hints(org: Option<&str>) -> Vec { vec![ format!( "Open an interactive thread view with `bt view thread{} --trace-id `.", - org_flag_suffix(profile) + org_flag_suffix(org) ), "Use `--non-interactive` for a compact text transcript.".to_string(), ] } -fn waterfall_hints(profile: Option<&str>, has_more: bool) -> Vec { +fn waterfall_hints(org: Option<&str>, has_more: bool) -> Vec { let mut hints = vec![ format!( "Render an agent-readable trace report with `bt view waterfall{} --trace-id `.", - org_flag_suffix(profile) + org_flag_suffix(org) ), "Use `--json` for computed offsets, token counts, costs, cache metrics, and raw ids." .to_string(), @@ -5369,7 +5369,7 @@ fn print_logs_text( rows: &[TraceSummaryRow], list_mode: ListMode, object_ref: &str, - profile: Option<&str>, + org: Option<&str>, limit: usize, preview_length: usize, next_cursor: Option<&str>, @@ -5406,7 +5406,7 @@ fn print_logs_text( println!("next_cursor: {cursor}"); println!( "next: bt view logs{} --object-ref {object_ref} --cursor {cursor} --non-interactive{limit_suffix}", - org_flag_suffix(profile) + org_flag_suffix(org) ); } else { println!("No additional rows."); @@ -5513,7 +5513,7 @@ fn print_thread_text( target: &ResolvedTraceCommandTarget, messages: &[Value], summary: &ThreadSummary, - profile: Option<&str>, + org: Option<&str>, ) { println!( "bt view thread: project={} root_span_id={} messages={}", @@ -5552,7 +5552,7 @@ fn print_thread_text( println!( "\njson: bt view thread --json{} --trace-id {}", - org_flag_suffix(profile), + org_flag_suffix(org), target.root_span_id ); } @@ -5561,7 +5561,7 @@ fn print_trace_text( trace_id: &str, rows: &[Map], object_ref: &str, - profile: Option<&str>, + org: Option<&str>, limit: usize, preview_length: usize, next_cursor: Option<&str>, @@ -5589,7 +5589,7 @@ fn print_trace_text( } println!( "\nTrace output is truncated. Re-run with --json for the full trace (`bt view trace --json{0} --object-ref {1} --trace-id {2}`), or fetch a single span with `bt view span{0} --object-ref {1} --id `.", - org_flag_suffix(profile), + org_flag_suffix(org), object_ref, trace_id ); @@ -5602,7 +5602,7 @@ fn print_trace_text( println!("next_cursor: {cursor}"); println!( "next: bt view trace{} --object-ref {} --trace-id {} --cursor {} --non-interactive{limit_suffix}", - org_flag_suffix(profile), + org_flag_suffix(org), object_ref, trace_id, cursor @@ -6701,26 +6701,7 @@ mod tests { use serde_json::json; fn base_args() -> BaseArgs { - BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - no_color: false, - no_input: false, - org_name: None, - org_name_source: None, - project: None, - project_source: None, - api_key: None, - api_key_source: None, - prefer_api_key: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn parsed_url_with_org(org: &str) -> ParsedTraceUrl { diff --git a/src/traces/waterfall.rs b/src/traces/waterfall.rs index 84caefa2..374556f8 100644 --- a/src/traces/waterfall.rs +++ b/src/traces/waterfall.rs @@ -140,7 +140,7 @@ pub(super) fn build_waterfall_view( pub(super) fn print_waterfall_text( target: &ResolvedTraceCommandTarget, waterfall: &WaterfallView, - profile: Option<&str>, + org: Option<&str>, limit: usize, has_more: bool, ) { @@ -197,13 +197,13 @@ pub(super) fn print_waterfall_text( } println!( "\nspan detail: bt view span{} --object-ref {} --id ", - org_flag_suffix(profile), + org_flag_suffix(org), format_object_ref_arg(&target.object_ref) ); println!("Use the inline `id=` value from a span row."); println!( "json: bt view waterfall --json{} --trace-id {}", - org_flag_suffix(profile), + org_flag_suffix(org), target.root_span_id ); } diff --git a/src/ui/mod.rs b/src/ui/mod.rs index 3f51e64d..3b0f28b4 100644 --- a/src/ui/mod.rs +++ b/src/ui/mod.rs @@ -56,6 +56,7 @@ pub use ratatui_table::{ box_with_title, render_experiment_summary_table, summary_metric_unit, SummaryExperimentColumn, SummaryMetricCell, SummaryMetricKind, SummaryMetricRow, SummaryTableOptions, }; +pub(crate) use select::select_or_create_project; pub use select::{fuzzy_select, select_project, ProjectSelectMode}; pub use spinner::{with_spinner, with_spinner_visible}; diff --git a/src/ui/select.rs b/src/ui/select.rs index bcfd02e2..1e0bcebe 100644 --- a/src/ui/select.rs +++ b/src/ui/select.rs @@ -2,7 +2,7 @@ use std::io::IsTerminal; use anyhow::{bail, Result}; use dialoguer::console::{style, Key, Term}; -use dialoguer::{theme::ColorfulTheme, FuzzySelect, Input}; +use dialoguer::{theme::ColorfulTheme, Confirm, FuzzySelect, Input}; use fuzzy_matcher::{skim::SkimMatcherV2, FuzzyMatcher}; const MAX_VISIBLE_ITEMS: usize = 12; @@ -246,6 +246,41 @@ fn fuzzy_select_with_pinned_first( } } +/// Resolve an explicitly named project, optionally creating it, or select an existing one. +pub(crate) async fn select_or_create_project( + client: &ApiClient, + requested: Option<&str>, + current: Option<&str>, + select_label: Option<&str>, +) -> Result { + let Some(name) = requested else { + return select_project( + client, + current, + select_label, + ProjectSelectMode::ExistingOnly, + ) + .await; + }; + if let Some(project) = + with_spinner("Loading project...", api::get_project_by_name(client, name)).await? + { + return Ok(project); + } + let Some(term) = super::prompt_term() else { + bail!("project '{name}' not found"); + }; + if Confirm::new() + .with_prompt(format!("Project '{name}' not found. Create it?")) + .default(false) + .interact_on(&term)? + { + with_spinner("Creating project...", api::create_project(client, name)).await + } else { + bail!("project '{name}' not found") + } +} + /// Interactive selector for project data. pub async fn select_project( client: &ApiClient, @@ -258,6 +293,20 @@ pub async fn select_project( let label = select_label.unwrap_or("Select project"); + if mode == ProjectSelectMode::ExistingOnly { + match projects.len() { + 0 => bail!("organization '{}' has no projects", client.org_name()), + 1 => return Ok(projects.remove(0)), + _ if !super::can_prompt() => { + bail!( + "organization '{}' has multiple projects; pass --project ", + client.org_name() + ) + } + _ => {} + } + } + if mode_allows_create(mode) { let names = project_display_names(&projects, mode); let default_sel = default_project_selection(&projects, current, mode)?; diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 52a4f874..6c945ed9 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -6,6 +6,7 @@ mod ids; mod json_object; mod plurals; mod profile; +mod shell; pub(crate) use app_url::{app_project_url, app_project_url_with_encoded_path}; pub use duration::parse_duration_to_seconds; @@ -15,3 +16,4 @@ pub(crate) use ids::new_uuid_id; pub(crate) use json_object::lookup_object_path; pub use plurals::pluralize; pub(crate) use profile::{profile_author_slug, sanitize_name_segment}; +pub(crate) use shell::quote_arg as shell_quote_arg; diff --git a/src/utils/shell.rs b/src/utils/shell.rs new file mode 100644 index 00000000..5d4f0a2e --- /dev/null +++ b/src/utils/shell.rs @@ -0,0 +1,10 @@ +pub(crate) fn quote_arg(value: &str) -> String { + if value + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.' | '/' | ':' | '=')) + { + value.to_string() + } else { + format!("'{}'", value.replace('\'', "'\\''")) + } +} diff --git a/tests/functions.rs b/tests/functions.rs index e1dd66cc..7467d0d5 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -812,6 +812,36 @@ fn auth_logout_requires_force_without_an_interactive_terminal() { assert!(config_dir.path().join("bt").join("auth.json").exists()); } +#[test] +fn auth_logout_bare_does_not_default_to_current_login() { + let cwd = tempdir().expect("create temp cwd"); + let config_dir = tempdir().expect("create temp config dir"); + let bt_config_dir = config_dir.path().join("bt"); + fs::create_dir_all(&bt_config_dir).expect("create bt config dir"); + fs::write(bt_config_dir.join("config.json"), r#"{"org":"test-org-a"}"#) + .expect("write active config"); + fs::write( + bt_config_dir.join("auth.json"), + r#"{"profiles":{"test-profile-a":{"auth_kind":"api_key","org_id":"org_test_a","org_name":"test-org-a","api_key_hint":"sk-****aaaaa"},"test-profile-b":{"auth_kind":"api_key","org_id":"org_test_b","org_name":"test-org-b","api_key_hint":"sk-****bbbbb"}}}"#, + ) + .expect("write auth logins"); + + let output = auth_sub_command( + cwd.path(), + config_dir.path(), + &["logout", "--no-input", "--force"], + ) + .output() + .expect("run bare non-interactive auth logout"); + + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr).contains("multiple auth logins match")); + let remaining = + fs::read_to_string(bt_config_dir.join("auth.json")).expect("read remaining auth logins"); + assert!(remaining.contains("test-org-a")); + assert!(remaining.contains("test-org-b")); +} + #[test] fn auth_refresh_errors_when_no_oauth_login_exists_even_with_json() { // --json must not swallow a real error: refresh only applies to OAuth From 40df301c843624d647df15ab4790362d38db4473 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 23 Jul 2026 13:18:34 -0700 Subject: [PATCH 09/24] bug fix, api url choice for cross org when there are multiple possible api urls --- src/auth.rs | 95 ++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 87 insertions(+), 8 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index a300a550..072d0e4d 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1011,12 +1011,20 @@ fn replace_with_canonical_auth_profile( } if canonical_key != current_key { - store.profiles.remove(current_key); - if let Some(existing) = store.profiles.get(&canonical_key) { - if !should_replace_migrated_profile(existing, &profile) { + // Two entries collapse onto the same slot (same OAuth org+email, or the + // same API key+org). Keep the usable one and delete the loser's secrets + // so we never orphan a credential in the keychain, and never drop the + // entry that still holds a working refresh token. + if let Some(existing) = store.profiles.get(&canonical_key).cloned() { + if should_replace_canonical_profile(&canonical_key, &existing, &profile) { + delete_all_profile_secrets(&canonical_key, &existing); + } else { + delete_all_profile_secrets(current_key, &profile); + store.profiles.remove(current_key); return true; } } + store.profiles.remove(current_key); } store.profiles.insert(canonical_key, profile); true @@ -1463,6 +1471,16 @@ fn select_profile_from_store( current: Option<&str>, store: &AuthStore, ) -> Result { + // Surface cross-org OAuth first so it is a predictable top-of-list choice + // rather than falling wherever its slot key happens to sort. The stable + // sort preserves the existing order of the remaining entries. + let mut names: Vec<&str> = names.to_vec(); + names.sort_by_key(|name| { + !store + .profiles + .get(*name) + .is_some_and(is_cross_org_oauth_profile) + }); let labels: Vec = names .iter() .map(|name| profile_label_from_store(name, store)) @@ -1651,8 +1669,12 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { let selected_org = selected_org.ok_or_else(|| { anyhow::anyhow!("API-key login requires an org; pass --org or rerun interactively") })?; - let selected_api_url = - resolve_profile_api_url(base.api_url.clone(), Some(&selected_org), &login_orgs)?; + let selected_api_url = resolve_profile_api_url( + base.api_url.clone(), + Some(&selected_org), + &login_orgs, + ui::can_prompt(), + )?; commit_api_key_profile( &api_key, @@ -1760,8 +1782,12 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { true, explicitly_quiet(base), )?; - let selected_api_url = - resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)?; + let selected_api_url = resolve_profile_api_url( + base.api_url.clone(), + selected_org.as_ref(), + &login_orgs, + ui::can_prompt(), + )?; commit_oauth_profile( &oauth_tokens, @@ -2812,6 +2838,7 @@ fn resolve_profile_api_url( explicit_api_url: Option, selected_org: Option<&LoginOrgInfo>, orgs: &[LoginOrgInfo], + can_prompt: bool, ) -> Result { if let Some(api_url) = explicit_api_url { return Ok(api_url); @@ -2834,8 +2861,18 @@ fn resolve_profile_api_url( .unwrap_or_else(|| DEFAULT_API_URL.to_string())); } + // A cross-org login spans orgs on different data planes. Let the user pick + // which API URL to store rather than failing outright. + if can_prompt { + let idx = ui::fuzzy_select("Select API URL", &api_urls, 0)?; + return Ok(api_urls + .into_iter() + .nth(idx) + .expect("selected API URL should be in range")); + } + bail!( - "multiple organizations expose different API URLs; choose an organization or pass --api-url explicitly" + "multiple organizations expose different API URLs; pass --org to pick one, or --api-url explicitly" ) } @@ -3834,6 +3871,22 @@ fn delete_legacy_profile_secrets(profile: &AuthProfile) { } } +/// Delete every secret a profile could reference: those stored under its own +/// slot key and those under its lazy legacy fallback key. Used when a duplicate +/// login is discarded during canonicalization so nothing is orphaned. +fn delete_all_profile_secrets(slot_key: &str, profile: &AuthProfile) { + match profile.auth_kind { + AuthKind::ApiKey => { + let _ = delete_profile_secret(slot_key); + } + AuthKind::Oauth => { + let _ = delete_profile_oauth_refresh_token(slot_key); + let _ = delete_profile_oauth_access_token(slot_key); + } + } + delete_legacy_profile_secrets(profile); +} + fn load_valid_cached_oauth_access_token( profile_name: &str, profile: &AuthProfile, @@ -4002,6 +4055,32 @@ fn should_replace_migrated_profile(existing: &AuthProfile, candidate: &AuthProfi } } +/// Runtime variant of [`should_replace_migrated_profile`] that can read the +/// keychain: when two OAuth logins collapse onto the same slot, keep whichever +/// still has a loadable refresh token (cached access-token expiry is unrelated +/// to which refresh token is live). Falls back to the pure expiry heuristic +/// when both or neither can refresh. +fn should_replace_canonical_profile( + slot_key: &str, + existing: &AuthProfile, + candidate: &AuthProfile, +) -> bool { + if let (AuthKind::Oauth, AuthKind::Oauth) = (existing.auth_kind, candidate.auth_kind) { + let has_refresh = |profile: &AuthProfile| { + matches!( + load_profile_oauth_refresh_token_for_profile(slot_key, profile), + Ok(Some(_)) + ) + }; + match (has_refresh(existing), has_refresh(candidate)) { + (false, true) => return true, + (true, false) => return false, + _ => {} + } + } + should_replace_migrated_profile(existing, candidate) +} + fn looks_like_sha256_hex(value: &str) -> bool { value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) } From 03a8fc3e08fca6b758afdfcf560c65fc37574775 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 23 Jul 2026 13:55:46 -0700 Subject: [PATCH 10/24] merge multiple old oauth logins into one when new bt is used --- src/auth.rs | 12 ++++++++++++ src/init.rs | 4 +++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/auth.rs b/src/auth.rs index 072d0e4d..7fb421b9 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1049,6 +1049,18 @@ fn maybe_rekey_api_key_profile_after_secret_load( } profile.api_key_hash = Some(api_key_hash(api_key)); + + // If the secret still lives only in the lazy legacy keychain slot, relocate + // it to the canonical slot now (mirroring the OAuth refresh path) so future + // resolves stop paying a permanent miss+fallback and deleting the old-named + // keychain item can't orphan the login. + if profile.legacy_secret_key.is_some() { + let canonical_key = canonical_profile_key(profile_name, &profile); + save_profile_secret(&canonical_key, api_key)?; + delete_legacy_profile_secrets(&profile); + profile.legacy_secret_key = None; + } + if replace_with_canonical_auth_profile(store, profile_name, profile) { save_auth_store(store)?; } diff --git a/src/init.rs b/src/init.rs index ae57e137..9449495e 100644 --- a/src/init.rs +++ b/src/init.rs @@ -58,7 +58,9 @@ pub async fn run(base: BaseArgs, args: InitArgs) -> Result<()> { Some("Link to project"), ) .await?; - let mut cfg = config::Config::default(); + // Load any existing file (only reachable via --force) so unknown passthrough + // keys are preserved, matching switch/config-set/post-login writers. + let mut cfg = config::load_file(&config_path); cfg.set_context( Some(&org), Some((project.name.as_str(), project.id.as_str())), From 40c1fb27478c0a990786826f88ba190cee289c91 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 23 Jul 2026 14:54:08 -0700 Subject: [PATCH 11/24] fix: temp file used to wrtie config.json now has unique name before it was always the same name which could lead to corruption if multiple bt used it at the same time --- src/auth.rs | 90 +++++++++++++++++---------------------------------- src/switch.rs | 10 +++++- 2 files changed, 39 insertions(+), 61 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 7fb421b9..83d7dbfa 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -3540,40 +3540,26 @@ fn load_secret_store() -> Result { fn save_secret_store(store: &SecretStore) -> Result<()> { let path = secret_store_path()?; - if let Some(parent) = path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create directory {}", parent.display()))?; - } + let parent = path.parent().unwrap_or_else(|| Path::new(".")); + fs::create_dir_all(parent) + .with_context(|| format!("failed to create directory {}", parent.display()))?; let data = serde_json::to_string_pretty(store).context("failed to serialize secret store")?; - let temp_path = path.with_extension("tmp"); - let mut file = fs::File::create(&temp_path) - .with_context(|| format!("failed to write temp secret store {}", temp_path.display()))?; + // A uniquely-named temp file (created `0600` by `tempfile`) prevents two + // concurrent `bt` writers from sharing one `.tmp` inode and renaming + // interleaved bytes over the store, and closes the umask window that a + // truncate-then-chmod on a fixed name would leave open. + let mut file = tempfile::NamedTempFile::new_in(parent) + .with_context(|| format!("failed to create temp secret store in {}", parent.display()))?; file.write_all(data.as_bytes()) - .with_context(|| format!("failed to write temp secret store {}", temp_path.display()))?; + .context("failed to write temp secret store")?; file.write_all(b"\n") - .with_context(|| format!("failed to write temp secret store {}", temp_path.display()))?; - file.sync_all() - .with_context(|| format!("failed to flush temp secret store {}", temp_path.display()))?; - - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - fs::set_permissions(&temp_path, fs::Permissions::from_mode(0o600)).with_context(|| { - format!( - "failed to set permissions on temp secret store {}", - temp_path.display() - ) - })?; - } - - fs::rename(&temp_path, &path).with_context(|| { - format!( - "failed to move temp secret store {} to {}", - temp_path.display(), - path.display() - ) - })?; + .context("failed to write temp secret store")?; + file.as_file() + .sync_all() + .context("failed to flush temp secret store")?; + file.persist(&path) + .with_context(|| format!("failed to move temp secret store to {}", path.display()))?; #[cfg(unix)] { @@ -4153,40 +4139,24 @@ fn save_auth_store(store: &AuthStore) -> Result<()> { } fn save_auth_store_to_path(path: &Path, store: &AuthStore) -> Result<()> { - if let Some(parent) = path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("failed to create directory {}", parent.display()))?; - } + let parent = path.parent().unwrap_or_else(|| Path::new(".")); + fs::create_dir_all(parent) + .with_context(|| format!("failed to create directory {}", parent.display()))?; let data = serde_json::to_string_pretty(store).context("failed to serialize auth config")?; - let temp_path = path.with_extension("tmp"); - let mut file = fs::File::create(&temp_path) - .with_context(|| format!("failed to write temp auth config {}", temp_path.display()))?; + // Unique temp name (see `save_secret_store`): keeps concurrent writers from + // colliding on a shared `.tmp` inode and publishing a corrupt store. + let mut file = tempfile::NamedTempFile::new_in(parent) + .with_context(|| format!("failed to create temp auth config in {}", parent.display()))?; file.write_all(data.as_bytes()) - .with_context(|| format!("failed to write temp auth config {}", temp_path.display()))?; + .context("failed to write temp auth config")?; file.write_all(b"\n") - .with_context(|| format!("failed to write temp auth config {}", temp_path.display()))?; - file.sync_all() - .with_context(|| format!("failed to flush temp auth config {}", temp_path.display()))?; - - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - fs::set_permissions(&temp_path, fs::Permissions::from_mode(0o600)).with_context(|| { - format!( - "failed to set permissions on temp auth config {}", - temp_path.display() - ) - })?; - } - - fs::rename(&temp_path, path).with_context(|| { - format!( - "failed to move temp auth config {} to {}", - temp_path.display(), - path.display() - ) - })?; + .context("failed to write temp auth config")?; + file.as_file() + .sync_all() + .context("failed to flush temp auth config")?; + file.persist(path) + .with_context(|| format!("failed to move temp auth config to {}", path.display()))?; #[cfg(unix)] { diff --git a/src/switch.rs b/src/switch.rs index 05419919..b66647c1 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -30,7 +30,7 @@ impl SwitchArgs { None => (None, None), Some(t) if t.contains('/') => { let parts: Vec<&str> = t.splitn(2, '/').collect(); - let o = (!parts[0].is_empty()).then(|| parts[0].to_string()); + let o = (!parts[0].is_empty()).then(|| config::normalize_org(parts[0]).to_string()); let p = (!parts[1].is_empty()).then(|| parts[1].to_string()); (o, p) } @@ -154,6 +154,14 @@ mod tests { (None, Some("test-project")), ), (Some("test-org/"), None, None, (Some("test-org"), None)), + // Positional "cross-org" folds to the "" marker, matching --org. + ( + Some("cross-org/test-project"), + None, + None, + (Some(""), Some("test-project")), + ), + (Some("cross-org/"), None, None, (Some(""), None)), (None, Some("test-org"), None, (Some("test-org"), None)), ( None, From fa2b5866375e62fd1df2f65c89c808758bc63982 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 23 Jul 2026 15:21:25 -0700 Subject: [PATCH 12/24] fix(datasets): default name used org with BRAINTRUST_API_KEY `bt datasets snapshots create my-dataset` would use the name of the OAuth login when the credentials used were BRAINTRUST_API_KEY, ie not OAuth --- src/datasets/snapshots.rs | 14 -------------- src/utils/profile.rs | 13 +++++++++++-- 2 files changed, 11 insertions(+), 16 deletions(-) diff --git a/src/datasets/snapshots.rs b/src/datasets/snapshots.rs index 337dd74e..cdb4cdd2 100644 --- a/src/datasets/snapshots.rs +++ b/src/datasets/snapshots.rs @@ -1089,10 +1089,6 @@ fn print_restore_preview( } fn resolve_default_snapshot_author(base: &BaseArgs, ctx: &ResolvedContext) -> Option { - if api_key_override_active(base) { - return None; - } - let profile = auth::active_auth_info(base, Some(ctx.client.org_name())) .ok() .flatten()?; @@ -1104,16 +1100,6 @@ fn default_snapshot_name(author: &str, now: DateTime) -> String { format!("{author}-{}", now.format("%Y%m%d-%H%M%Sz")) } -fn api_key_override_active(base: &BaseArgs) -> bool { - matches!( - base.api_key_source, - Some(crate::args::ArgValueSource::CommandLine) - ) && base - .api_key - .as_deref() - .is_some_and(|value| !value.trim().is_empty()) -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/utils/profile.rs b/src/utils/profile.rs index 613710a0..ebd80bb3 100644 --- a/src/utils/profile.rs +++ b/src/utils/profile.rs @@ -1,10 +1,12 @@ use crate::auth::ProfileInfo; +/// Slug for the authenticated human (name, else email local part). Org is +/// deliberately not a fallback: a bare API key has no author, so callers +/// substitute a generic placeholder rather than name the snapshot after the org. pub(crate) fn profile_author_slug(profile: &ProfileInfo) -> Option { [ profile.user_name.as_deref(), profile.email.as_deref().and_then(email_local_part), - profile.org_name.as_deref(), ] .into_iter() .flatten() @@ -82,11 +84,18 @@ mod tests { } #[test] - fn profile_author_slug_returns_none_without_identity_or_org() { + fn profile_author_slug_returns_none_without_identity() { let profile = profile_info(None, None, None); assert_eq!(profile_author_slug(&profile), None); } + #[test] + fn profile_author_slug_ignores_org_name() { + // A bare API key has an org but no human identity — not an author. + let profile = profile_info(Some("test-org"), None, None); + assert_eq!(profile_author_slug(&profile), None); + } + #[test] fn sanitize_name_segment_collapses_non_alnum() { assert_eq!( From bf9a0efa618d1aa8dbd9dd9dc3c1b5630904c811 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 23 Jul 2026 16:17:01 -0700 Subject: [PATCH 13/24] fix: detect symlinked home ; hide api keys better --- src/auth.rs | 10 +++++++++- src/config/mod.rs | 8 +++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 83d7dbfa..13c4482c 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -3954,8 +3954,13 @@ pub fn obscure_api_key(key: &str) -> String { if !key.is_ascii() || key.len() <= 8 { return "****".to_string(); } - let prefix_end = key.find('-').map(|i| i + 1).unwrap_or(0); let suffix_start = key.len().saturating_sub(5); + let prefix_end = key.find('-').map(|i| i + 1).unwrap_or(0); + // A late first dash can push the prefix up to (or past) the suffix, leaving + // no masked middle and revealing the whole key. Fully mask instead. + if prefix_end >= suffix_start { + return "****".to_string(); + } format!("{}****{}", &key[..prefix_end], &key[suffix_start..]) } @@ -5509,6 +5514,9 @@ mod tests { ("sk-LumEdp0BbLRzhJwO", "sk-****zhJwO"), ("abc", "****"), ("abcdefghijklm", "****ijklm"), + // Late first dash leaves no maskable middle: fully mask rather than + // reveal the whole key (would otherwise be "abcdefg-****g-hij"). + ("abcdefg-hij", "****"), ("sk-café-résumé-key", "****"), ] { assert_eq!(obscure_api_key(key), expected); diff --git a/src/config/mod.rs b/src/config/mod.rs index db36c628..62bb4a7b 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -267,8 +267,14 @@ enum ProjectBoundary { } fn project_boundary(start: PathBuf, home: Option<&Path>) -> ProjectBoundary { + // `current_dir()` is the physical path (symlinks resolved) while `$HOME` may + // not be, so also compare canonicalized forms — exact equality alone can + // walk straight past a symlinked home boundary. + let home_canon = home.and_then(|h| fs::canonicalize(h).ok()); for dir in start.ancestors() { - if Some(dir) == home { + let at_home = + Some(dir) == home || (home_canon.is_some() && fs::canonicalize(dir).ok() == home_canon); + if at_home { return ProjectBoundary::Home; } if dir.parent().is_none() { From 773405d1024ce28c8791f2a1f4b881c4b3ae4e03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 23 Jul 2026 18:03:54 -0700 Subject: [PATCH 14/24] chore: failing to update auth.json no longer crashes the whole bt command --- src/auth.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 13c4482c..92c392bc 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -4011,9 +4011,15 @@ fn load_auth_store_from_path(path: &Path) -> Result { .with_context(|| format!("failed to parse auth config {}", path.display()))?; let migrated = migrate_auth_store(store.clone()); if migrated != store { - save_auth_store_to_path(path, &migrated).with_context(|| { - format!("failed to persist auth config migration {}", path.display()) - })?; + // The migrated store is already usable in memory, so a failed write-back + // must not break read-only commands (`bt status`, `bt auth logins`). Warn + // and proceed; the next writable run retries the migration. + if let Err(err) = save_auth_store_to_path(path, &migrated) { + eprintln!( + "warning: Migrating {} to use the new format failed. Please delete this file and login again. ({err})", + path.display() + ); + } } Ok(migrated) } From 36772034c5a9660c9d338aabe0a8dd1c508299d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Fri, 24 Jul 2026 13:19:57 -0700 Subject: [PATCH 15/24] fix(auth): prune orphaned auth secrets after migration Always show current auth in bt status even if it's from API key in env --- src/auth.rs | 102 ++++++++++++++++++++++++++++++++++++++++++++++++-- src/status.rs | 41 ++++++++++---------- 2 files changed, 120 insertions(+), 23 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 92c392bc..78f9ccb2 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -4014,11 +4014,15 @@ fn load_auth_store_from_path(path: &Path) -> Result { // The migrated store is already usable in memory, so a failed write-back // must not break read-only commands (`bt status`, `bt auth logins`). Warn // and proceed; the next writable run retries the migration. - if let Err(err) = save_auth_store_to_path(path, &migrated) { - eprintln!( + match save_auth_store_to_path(path, &migrated) { + // Only prune once the collapsed store is durably on disk; otherwise + // the on-disk file still references the dropped duplicate and the + // next load must be able to retry the migration. + Ok(()) => prune_orphaned_migration_secrets(&store, &migrated), + Err(err) => eprintln!( "warning: Migrating {} to use the new format failed. Please delete this file and login again. ({err})", path.display() - ); + ), } } Ok(migrated) @@ -4054,6 +4058,48 @@ fn migrate_auth_store(store: AuthStore) -> AuthStore { migrated } +/// Secret slots left dangling after migration collapsed duplicate logins onto a +/// shared canonical key. A surviving login keeps its secret under its +/// `legacy_secret_key` (until it is lazily relocated) or, absent one, under its +/// own slot key; any pre-migration key outside that referenced set belonged to a +/// dropped duplicate and can be deleted. Pure so it stays unit-testable; the +/// caller performs the keychain I/O. +fn orphaned_migration_secret_keys<'a>( + before: &'a AuthStore, + after: &AuthStore, +) -> Vec<(&'a str, AuthKind)> { + let referenced: BTreeSet<&str> = after + .profiles + .iter() + .map(|(slot, profile)| { + profile + .legacy_secret_key + .as_deref() + .unwrap_or(slot.as_str()) + }) + .collect(); + before + .profiles + .iter() + .filter(|(key, _)| !referenced.contains(key.as_str())) + .map(|(key, profile)| (key.as_str(), profile.auth_kind)) + .collect() +} + +fn prune_orphaned_migration_secrets(before: &AuthStore, after: &AuthStore) { + for (key, auth_kind) in orphaned_migration_secret_keys(before, after) { + match auth_kind { + AuthKind::ApiKey => { + let _ = delete_profile_secret(key); + } + AuthKind::Oauth => { + let _ = delete_profile_oauth_refresh_token(key); + let _ = delete_profile_oauth_access_token(key); + } + } + } +} + fn should_replace_migrated_profile(existing: &AuthProfile, candidate: &AuthProfile) -> bool { match (existing.auth_kind, candidate.auth_kind) { (AuthKind::Oauth, AuthKind::Oauth) => { @@ -5189,6 +5235,56 @@ mod tests { assert_eq!(profile.legacy_secret_key.as_deref(), Some("new")); } + #[test] + fn migration_reports_dropped_duplicate_secret_as_orphan() { + // Two legacy OAuth entries for the same org+email collapse onto one + // canonical slot. The survivor's secret stays reachable (via its + // legacy_secret_key), while the dropped duplicate's key must be reported + // as an orphan so its keychain secret can be deleted. + let mut store = AuthStore::default(); + for (name, expires_at) in [("old", 10), ("new", 20)] { + store.profiles.insert( + name.to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + email: Some("user@example.test".to_string()), + oauth_access_expires_at: Some(expires_at), + ..Default::default() + }, + ); + } + + let migrated = migrate_auth_store(store.clone()); + let orphans = orphaned_migration_secret_keys(&store, &migrated); + + assert_eq!(orphans, vec![("old", AuthKind::Oauth)]); + // The survivor "new" is referenced through the canonical slot's + // legacy_secret_key and must never be pruned. + assert!(!orphans.iter().any(|(key, _)| *key == "new")); + } + + #[test] + fn migration_without_collapse_reports_no_orphans() { + // A single entry that merely gets rekeyed keeps its secret under the + // legacy key, so nothing is orphaned. + let mut store = AuthStore::default(); + store.profiles.insert( + "test-org".to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + org_id: Some("org_fake".to_string()), + org_name: Some("test-org".to_string()), + email: Some("user@example.test".to_string()), + ..Default::default() + }, + ); + + let migrated = migrate_auth_store(store.clone()); + assert!(orphaned_migration_secret_keys(&store, &migrated).is_empty()); + } + #[test] fn config_auth_context_returns_config_org() { let base = make_base(); diff --git a/src/status.rs b/src/status.rs index 756089ff..b8a8ccbf 100644 --- a/src/status.rs +++ b/src/status.rs @@ -93,28 +93,29 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { if let Some(src) = source { println!("source: {src}"); } - } else if let Some(org) = display_org { - let scope = match (org, project.as_deref()) { - ("cross-org", _) => org.to_string(), - (org, Some(project)) => format!("{org}/{project}"), - (org, None) => org.to_string(), - }; - println!("{scope}"); - let auth_line = match &auth_info { - Some(p) => format!(" auth: {}", format_auth(p)), - None => " auth: (none)".to_string(), + } else { + // Plain one-liner. Always surface the active auth — even when no org is + // configured (an env-only API key, or a cross-org OAuth login) — instead + // of hiding it behind --verbose. + let cross_org_oauth = auth_info + .as_ref() + .is_some_and(|p| p.auth_method == "oauth" && p.org_name.is_none()); + let header = match display_org { + Some("cross-org") => "cross-org".to_string(), + Some(org) => match project.as_deref() { + Some(project) => format!("{org}/{project}"), + None => org.to_string(), + }, + None if cross_org_oauth => "cross-org".to_string(), + None if auth_info.is_some() => "No default org".to_string(), + None => "No org/project configured. Run `bt switch` to set one.".to_string(), }; - println!("{auth_line}"); - } else if auth_info - .as_ref() - .is_some_and(|p| p.auth_method == "oauth" && p.org_name.is_none()) - { - println!("cross-org"); - if let Some(p) = &auth_info { - println!(" auth: {}", format_auth(p)); + println!("{header}"); + match &auth_info { + Some(p) => println!(" auth: {}", format_auth(p)), + None if display_org.is_some() => println!(" auth: (none)"), + None => {} } - } else { - println!("No org/project configured. Run `bt switch` to set one."); } Ok(()) From 7d4fa22604e25b53798f217920cb572b0e704ad2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Tue, 28 Jul 2026 14:04:45 -0700 Subject: [PATCH 16/24] feat(login): separate auth and org MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OAuth logins are instance-scoped by normalized app_url; API keys remain org-scoped.gc Added OAuth migration and duplicate pruning.gc Added org_id, app_url, and api_url config support with correct precedence and context coupling.gc Reworked bt auth login, bt init, and bt switch.gc Removed cross-org/profile selection behavior.gc Updated login filtering, logout, status, setup compatibility, URL-derived context hints, tests, and README.gc diff --git a/README.md b/README.mdgc index f7bbc60..88de925 100644gc --- a/README.mdgc +++ b/README.mdgc @@ -139,7 +139,7 @@ Remove-Item -Recurse -Force (Join-Path $env:APPDATA "bt") -ErrorAction SilentlyCgc | ------------- | ------------------------------------------------------------------ |gc | `bt init` | Initialize `.bt/` config directory and link to a project |gc | `bt auth` | Authenticate with Braintrust |gc -| `bt switch` | Switch org and project context |gc +| `bt switch` | Switch instance, org, and project context |gc | `bt status` | Show current org and project context |gc | `bt datasets` | Manage datasets and dataset pipelines |gc | `bt eval` | Run eval files (Unix only) |gc @@ -312,43 +312,40 @@ Local version and pagination-key conversion helpers:gc gc ## `bt auth`gc gc -- Authenticate interactively (prompts for auth method and organization):gc +- Authenticate interactively:gc - `bt auth login`gc - - First prompt chooses: `OAuth (browser)` (default) or `API key`.gc - - OAuth can be saved for an org or in cross-org mode. API-key logins are saved per org; if a key can access multiple orgs, `bt` uses a searchable org picker.gc - - After login, `bt` updates the active org context immediately. If `--project` is set, it validates and saves that project's name and ID. Without `--project`, a same-org login preserves the existing project; changing orgs clears stale project context.gc + - First choose `OAuth (browser)` (default) or `API key`, then choose an organization and config scope.gc + - OAuth is stored once per Braintrust instance, identified by app URL, and can authenticate every organization available to that user in the instance.gc + - API-key logins remain organization-scoped; multiple keys for one organization remain distinct.gc + - Login writes `org`, `org_id`, `project`, `project_id`, `app_url`, and `api_url` to the selected config scope. A same-context login preserves the existing project when no project is requested.gc - Use `--global` or `--local` to choose the config scope. Without either flag, an existing local config causes an interactive scope picker (default: local); non-interactive runs must pass a scope. `--local` never creates `.bt`.gc - - `bt` confirms the resolved API URL before saving.gc -- Login with OAuth (browser-based, stores refresh token in secure credential store):gc - - `bt auth login --oauth --org myorg`gc - - You can pass `--no-browser` to print the URL without auto-opening.gc - - On remote/SSH hosts, paste the final callback URL from your local browser if localhost callback cannot be delivered.gc +- Login with OAuth:gc + - `bt auth login --oauth --org test-org`gc + - You can pass `--no-browser` to print the URL without opening it automatically.gc + - On remote/SSH hosts, paste the final callback URL if the localhost callback cannot be delivered.gc - List saved auth logins:gc - `bt auth logins`gc - - `bt auth logins --org test-org` (matches stored org name or ID)gc - - `bt auth logins --prefer-api-key` (API-key logins only)gc - - Both filters can be combined.gc + - `bt auth logins --org test-org` dynamically lists only credentials that can use that organization.gc + - `bt auth logins --prefer-api-key` lists API-key logins only.gc - Log out:gc - `bt auth logout` — choose from all saved logins interactivelygc - - `bt auth logout --org test-org --oauth` — filter to the org's OAuth logingc - - `bt auth logout --org test-org --api-key-hint sk-****abcde` — select an API-key logingc - - `bt auth logout --force` (skip confirmation after selecting a login)gc -- Show current auth context:gc - - `bt status`gc -- Force-refresh OAuth access token for debugging:gc - - `bt auth refresh --org myorg`gc + - `bt auth logout --app-url https://www.example.test --oauth`gc + - `bt auth logout --org test-org --api-key-hint sk-****abcde`gc + - `bt auth logout --force` — skip confirmationgc +- Force-refresh the OAuth login for the selected instance:gc + - `bt auth refresh --app-url https://www.example.test`gc gc Auth resolution order for commands is:gc gc 1. Explicit `--api-key sk-...`gc -2. `--prefer-api-key` / `BRAINTRUST_PREFER_API_KEY` (uses `BRAINTRUST_API_KEY` first, then a stored API key for the selected org, then falls back to OAuth)gc -3. Stored OAuth login for the selected org (or cross-org OAuth when selected)gc +2. `--prefer-api-key` / `BRAINTRUST_PREFER_API_KEY` (`BRAINTRUST_API_KEY`, then a matching stored API key, then matching OAuth)gc +3. OAuth for the selected Braintrust instance when it can access the selected organizationgc 4. `BRAINTRUST_API_KEY`gc -5. Stored API key login for the selected orggc +5. A matching stored API keygc gc -`--prefer-api-key` without `--org` targets the org shown by `bt status`. It cannot be used from cross-org context; pass a concrete `--org`. Once a key is selected, an invalid key or a key belonging to another requested org is an error and does not fall back to OAuth. Multiple keys in one org remain separate and are shown with key hints.gc +OAuth credentials are matched by app URL. API-key credentials are matched by app URL, API URL, and organization. Explicit flags override environment variables, which override local config, global config, and finally the built-in Braintrust URLs.gc gc -On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Service daemon. On macOS, it uses the `security` keychain utility. If a secure store is unavailable, `bt` falls back to a plaintext secrets file with `0600` permissions.gc +On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Service daemon. On macOS, it uses the `security` keychain utility. If secure storage is unavailable, `bt` falls back to a plaintext secrets file with `0600` permissions.gc gc ## `bt init`gc gc @@ -358,24 +355,23 @@ On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Segc - `bt init --here` — create in the current directory without walking (including at home or `/`)gc - `bt init --force` — overwrite an existing discovered `.bt/config.json`; it does not change discoverygc gc -The saved context includes `org`, `project`, and `project_id`. Cross-org is excluded from the init picker because init always requires a project.gc +The saved context includes the Braintrust instance URLs, organization name and ID, and project name and ID.gc gc ## `bt switch`gc gc -Interactively switch org and project context:gc +`bt switch` changes context without selecting a credential. It chooses a Braintrust instance, discovers the organizations available through that instance's credentials, and then chooses a project.gc gc -- `bt switch` — first choose a saved OAuth-backed org or a specific API-key login, then choose a projectgc -- `bt switch myproject` — switch the current org to a project by namegc -- `bt switch test-org/test-project` — switch to a specific org and projectgc -- `bt switch --org cross-org` — select cross-org OAuth and clear project contextgc +- `bt switch`gc +- `bt switch test-project`gc +- `bt switch test-org/test-project`gc - `bt switch --global` — persist to global config (`~/.config/bt/config.json`)gc - `bt switch --local` — update an existing local config (`.bt/config.json`); it never creates onegc gc -A sole login/project is selected automatically. Multiple API keys in one org remain separate picker entries with hints; OAuth accounts collapse to an org choice, so run `bt auth login` again to change OAuth accounts within that org. With an existing local config and no scope flag, interactive mode asks for global/local (default: local); non-interactive mode requires `--global` or `--local`.gc +A sole instance, organization, or project is selected automatically. With an existing local config and no scope flag, interactive mode asks for global/local (default: local); non-interactive mode requires `--global` or `--local`.gc gc ## Config context merginggc gc -Global config is `~/.config/bt/config.json`; local config is the first discovered `.bt/config.json`. Local context wins, but a local org inherits the global project only when both configs select the same org. A local project with no org never inherits a global org. Cross-org is stored as `"org": ""`, treated as a distinct org, and rendered as `cross-org` by `bt status`. Legacy `profile` fields are ignored; unknown extra keys are preserved when context is updated.gc +Global config is `~/.config/bt/config.json`; local config is the first discovered `.bt/config.json`. Both use the fields `org`, `org_id`, `project`, `project_id`, `app_url`, and `api_url`. Local values win. Organization IDs stay coupled to organization names, and organization/project context is inherited only within the same app URL. Legacy `profile` fields and obsolete empty cross-org contexts are ignored; unknown extra keys are preserved during updates.gc gc ## `bt status`gc gc diff --git a/src/args.rs b/src/args.rsgc index 1770da9..cfe4e4c 100644gc --- a/src/args.rsgc +++ b/src/args.rsgc @@ -45,6 +45,10 @@ pub struct BaseArgs {gc #[arg(skip)]gc pub org_name_source: Option,gc gc + /// Stable org ID resolved from config or internal context selection.gc + #[arg(skip)]gc + pub org_id: Option,gc +gc /// Override active projectgc #[arg(gc short = 'p',gc @@ -65,10 +69,6 @@ pub struct BaseArgs {gc #[arg(skip)]gc pub api_key_source: Option,gc gc - /// Exact auth slot selected internally by switch/init.gc - #[arg(skip)]gc - pub pinned_auth_slot: Option,gc -gc /// Prefer API key credentials for the selected org when available.gc #[arg(long = "prefer-api-key", env = "BRAINTRUST_PREFER_API_KEY", global = true, value_parser = clap::builder::BoolishValueParser::new(), default_value_t = false)]gc pub prefer_api_key: bool,gc @@ -82,6 +82,9 @@ pub struct BaseArgs {gc )]gc pub api_url: Option,gc gc + #[arg(skip)]gc + pub api_url_source: Option,gc +gc /// Override app URL (or via BRAINTRUST_APP_URL)gc #[arg(gc long,gc @@ -91,6 +94,9 @@ pub struct BaseArgs {gc )]gc pub app_url: Option,gc gc + #[arg(skip)]gc + pub app_url_source: Option,gc +gc /// Path to a PEM-encoded CA bundle used for HTTPS requests.gc #[arg(gc long = "ca-cert",gc @@ -120,7 +126,11 @@ pub struct CLIArgs {gc }gc gc fn parse_org_name(value: &str) -> Result {gc - Ok(crate::config::normalize_org(value).to_string())gc + let value = value.trim();gc + if value.is_empty() {gc + return Err("organization cannot be empty".to_string());gc + }gc + Ok(value.to_string())gc }gc gc pub(crate) fn custom_api_without_app_url(api_url: Option<&str>, app_url: Option<&str>) -> bool {gc @@ -149,13 +159,13 @@ mod tests {gc #[test]gc fn org_normalization() {gc for (input, expected) in [gc - ("cross-org", ""),gc - (" ", ""),gc + ("cross-org", "cross-org"),gc (" test-org ", "test-org"),gc (" org_test_123 ", "org_test_123"),gc ] {gc assert_eq!(parse_org_name(input).unwrap(), expected);gc }gc + assert!(parse_org_name(" ").is_err());gc assert!(custom_api_without_app_url(gc Some("https://api.example.test"),gc Nonegc diff --git a/src/auth.rs b/src/auth.rsgc index 78f9ccb..f422677 100644gc --- a/src/auth.rsgc +++ b/src/auth.rsgc @@ -59,6 +59,7 @@ pub struct ResolvedAuth {gc pub api_url: Option,gc pub app_url: Option,gc pub org_name: Option,gc + pub org_id: Option,gc pub is_oauth: bool,gc slot_key: Option,gc }gc @@ -72,6 +73,11 @@ pub struct ProfileInfo {gc pub api_key_hint: Option,gc }gc gc +#[derive(Debug, Clone, PartialEq, Eq)]gc +pub struct AvailableInstance {gc + pub app_url: String,gc +}gc +gc #[derive(Debug, Clone, PartialEq, Eq)]gc pub struct AvailableOrg {gc pub id: String,gc @@ -82,6 +88,7 @@ pub struct AvailableOrg {gc #[derive(Debug, Clone, Copy, PartialEq, Eq)]gc enum RecoverableAuthErrorKind {gc OauthRefreshToken,gc + OauthOrgAccess,gc StoredCredential,gc }gc gc @@ -103,6 +110,14 @@ fn recoverable_auth_error(kind: RecoverableAuthErrorKind, message: String) -> angc anyhow::Error::new(RecoverableAuthError { kind, message })gc }gc gc +fn is_oauth_org_access_error(err: &anyhow::Error) -> bool {gc + err.chain().any(|source| {gc + sourcegc + .downcast_ref::()gc + .is_some_and(|err| err.kind == RecoverableAuthErrorKind::OauthOrgAccess)gc + })gc +}gc +gc pub fn is_missing_credential_error(err: &anyhow::Error) -> bool {gc err.chain().any(|source| {gc sourcegc @@ -126,6 +141,14 @@ pub fn list_profiles() -> Result> {gc .collect())gc }gc gc +pub(crate) fn has_oauth_login_for_instance(base: &BaseArgs) -> Result {gc + let store = load_auth_store()?;gc + Ok(storegc + .profilesgc + .values()gc + .any(|profile| profile.auth_kind == AuthKind::Oauth && profile_matches_urls(base, profile)))gc +}gc +gc pub async fn list_available_orgs(base: &BaseArgs) -> Result> {gc let resolved = resolve_auth(base).await?;gc let app_url = resolvedgc @@ -163,6 +186,133 @@ async fn available_orgs(api_key: &str, app_url: &str) -> Result Result> {gc + let store = load_auth_store()?;gc + let constrain_app = matches!(gc + base.app_url_source,gc + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable)gc + );gc + let requested_app = constrain_appgc + .then_some(base.app_url.as_deref())gc + .flatten()gc + .map(canonical_url);gc + let mut apps = storegc + .profilesgc + .values()gc + .map(profile_app_url)gc + .filter(|app| requested_app.is_none_or(|requested| canonical_url(app) == requested))gc + .map(|app| canonical_url(app).to_string())gc + .collect::>();gc +gc + if basegc + .api_keygc + .as_deref()gc + .is_some_and(|key| !key.trim().is_empty())gc + {gc + let app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + if requested_app.is_none_or(|requested| canonical_url(app) == requested) {gc + apps.insert(canonical_url(app).to_string());gc + }gc + }gc +gc + if apps.is_empty() && constrain_app {gc + bail!(gc + "no credentials found for app URL '{}'; run `bt auth login --app-url {}`",gc + base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc + shell_quote_arg(base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL))gc + );gc + }gc + Ok(appsgc + .into_iter()gc + .map(|app_url| AvailableInstance { app_url })gc + .collect())gc +}gc +gc +pub(crate) async fn available_orgs_for_instance(gc + base: &BaseArgs,gc + app_url: &str,gc +) -> Result> {gc + let mut store = load_auth_store()?;gc + let explicit_api = matches!(gc + base.api_url_source,gc + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable)gc + )gc + .then(|| base.api_url.as_deref())gc + .flatten();gc + let mut orgs = BTreeMap::::new();gc + let matching = storegc + .profilesgc + .iter()gc + .filter(|(_, profile)| canonical_url(profile_app_url(profile)) == canonical_url(app_url))gc + .filter(|(_, profile)| {gc + profile.auth_kind == AuthKind::Oauthgc + || explicit_apigc + .is_none_or(|url| canonical_url(url) == canonical_url(profile_api_url(profile)))gc + })gc + .map(|(slot, profile)| (slot.clone(), profile.clone()))gc + .collect::>();gc +gc + for (slot, profile) in matching {gc + match profile.auth_kind {gc + AuthKind::Oauth => {gc + let mut oauth_base = base.clone();gc + oauth_base.app_url = Some(app_url.to_string());gc + if oauth_base.api_url_source.is_none() {gc + oauth_base.api_url = profile.api_url.clone();gc + }gc + let token = load_oauth_access_token(&oauth_base, &mut store, &slot).await?;gc + for org in fetch_login_orgs(&token, app_url).await? {gc + orgs.insert(gc + org.id.clone(),gc + AvailableOrg {gc + id: org.id,gc + name: org.name,gc + api_url: org.api_url,gc + },gc + );gc + }gc + }gc + AuthKind::ApiKey => {gc + if let (Some(id), Some(name)) = (profile.org_id, profile.org_name) {gc + orgs.entry(id.clone()).or_insert(AvailableOrg {gc + id,gc + name,gc + api_url: profile.api_url,gc + });gc + }gc + }gc + }gc + }gc +gc + if let Some(api_key) = base.api_key.as_deref().filter(|key| !key.trim().is_empty()) {gc + let base_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + if canonical_url(base_app) == canonical_url(app_url) {gc + for org in fetch_login_orgs(api_key, app_url).await? {gc + orgs.insert(gc + org.id.clone(),gc + AvailableOrg {gc + id: org.id,gc + name: org.name,gc + api_url: org.api_url,gc + },gc + );gc + }gc + }gc + }gc +gc + let mut orgs = orgs.into_values().collect::>();gc + orgs.sort_by(|a, b| {gc + a.namegc + .to_ascii_lowercase()gc + .cmp(&b.name.to_ascii_lowercase())gc + .then_with(|| a.name.cmp(&b.name))gc + });gc + if orgs.is_empty() {gc + bail!("no organizations are available for Braintrust instance '{app_url}'");gc + }gc + Ok(orgs)gc +}gc +gc #[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]gc struct AuthStore {gc #[serde(default)]gc @@ -279,7 +429,7 @@ pub struct AuthArgs {gc enum AuthCommand {gc /// Authenticate with Braintrust (OAuth or API key)gc Login(AuthLoginArgs),gc - /// Force-refresh OAuth access token for the selected orggc + /// Force-refresh the OAuth access token for the selected instancegc Refresh,gc /// List saved auth logins and check connection statusgc Logins(AuthLoginsArgs),gc @@ -332,7 +482,7 @@ pub async fn run(base: BaseArgs, args: AuthArgs) -> Result<()> {gc }gc AuthCommand::Refresh => run_login_refresh(&base).await,gc AuthCommand::Logins(logins_args) => run_logins(&base, logins_args).await,gc - AuthCommand::Logout(logout_args) => run_login_logout(base, logout_args),gc + AuthCommand::Logout(logout_args) => run_login_logout(base, logout_args).await,gc }gc }gc gc @@ -371,7 +521,7 @@ pub async fn fast_login(base: &BaseArgs) -> Result {gc let login = LoginState::new();gc login.set(gc api_key,gc - String::new(),gc + auth.org_id.clone().unwrap_or_default(),gc org_name,gc api_url.clone(),gc app_url.clone(),gc @@ -428,7 +578,7 @@ pub async fn login(base: &BaseArgs) -> Result {gc let login = LoginState::new();gc login.set(gc api_key.clone(),gc - String::new(),gc + auth.org_id.clone().unwrap_or_default(),gc org_name,gc auth.api_urlgc .clone()gc @@ -752,8 +902,20 @@ fn config_auth_context(base: &BaseArgs) -> Option {gc config_auth_context_from_config(base, &cfg)gc }gc gc +fn configured_org_for_app_url(app_url: &str) -> Option {gc + let cfg = crate::config::load().ok()?;gc + let config_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + crate::config::urls_equal(app_url, config_app)gc + .then_some(cfg.org)gc + .flatten()gc +}gc +gc fn config_auth_context_from_config(base: &BaseArgs, cfg: &crate::config::Config) -> Option {gc - if crate::config::org_option(base.org_name.as_deref()).is_none() {gc + let base_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let config_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + if crate::config::org_option(base.org_name.as_deref()).is_none()gc + && crate::config::urls_equal(base_app, config_app)gc + {gc crate::config::org_option(cfg.org.as_deref()).map(str::to_string)gc } else {gc Nonegc @@ -831,11 +993,6 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result {gc let can_prompt = ui::can_prompt();gc gc let effective_org = effective_org_name(base, &cfg_org);gc - reject_cross_org_api_key_preference(base.prefer_api_key, effective_org, &store)?;gc -gc - if let Some(slot) = base.pinned_auth_slot.clone() {gc - return resolve_saved_auth_slot(base, &mut store, &None, &slot).await;gc - }gc gc let source = resolve_auth_source(gc base.prefer_api_key,gc @@ -849,13 +1006,49 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result {gc AuthSource::CliApiKey(api_key) | AuthSource::EnvApiKey(api_key) => {gc resolve_ad_hoc_api_key_auth(base, &cfg_org, api_key).awaitgc }gc - AuthSource::Oauth(slot) | AuthSource::ApiKey(slot) => {gc + AuthSource::Oauth(slot) => {gc + match resolve_saved_auth_slot(base, &mut store, &cfg_org, &slot).await {gc + Ok(auth) => Ok(auth),gc + Err(err) if is_oauth_org_access_error(&err) && !base.prefer_api_key => {gc + if let Some(api_key) = resolve_env_api_key(base) {gc + return resolve_ad_hoc_api_key_auth(base, &cfg_org, api_key).await;gc + }gc + if let Some(api_key_slot) = select_profile_for_auth(gc + base,gc + &store,gc + &cfg_org,gc + AuthKind::ApiKey,gc + can_prompt,gc + )? {gc + return resolve_saved_auth_slot(base, &mut store, &cfg_org, &api_key_slot)gc + .await;gc + }gc + Err(err)gc + }gc + Err(err) => Err(err),gc + }gc + }gc + AuthSource::ApiKey(slot) => {gc resolve_saved_auth_slot(base, &mut store, &cfg_org, &slot).awaitgc }gc AuthSource::None => {gc if base.prefer_api_key {gc bail!("--prefer-api-key requires an API key or OAuth login for the selected org");gc }gc + if !store.profiles.is_empty()gc + && !storegc + .profilesgc + .values()gc + .any(|profile| profile_matches_urls(base, profile))gc + {gc + let app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let api = base.api_url.as_deref().unwrap_or(DEFAULT_API_URL);gc + bail!(gc + "no credentials match app URL '{}' and API URL '{}'; run `bt auth login` with these URLs",gc + app,gc + apigc + );gc + }gc if effective_org.is_none() {gc if let Some(err) = missing_org_for_stored_logins_error(&store) {gc return Err(err);gc @@ -866,6 +1059,7 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result {gc api_url: base.api_url.clone(),gc app_url: base.app_url.clone(),gc org_name: effective_org.map(str::to_string),gc + org_id: base.org_id.clone(),gc is_oauth: false,gc slot_key: None,gc })gc @@ -884,6 +1078,7 @@ async fn resolve_ad_hoc_api_key_auth(gc }gc gc let mut resolved_org = requested_org.map(str::to_string);gc + let mut resolved_org_id = base.org_id.clone();gc let mut resolved_api_url = base.api_url.clone();gc if let Some(requested_org) = requested_org {gc if crate::args::custom_api_without_app_url(base.api_url.as_deref(), base.app_url.as_deref())gc @@ -905,6 +1100,7 @@ async fn resolve_ad_hoc_api_key_auth(gc )gc })?;gc resolved_org = Some(selected_org.name.clone());gc + resolved_org_id = Some(selected_org.id.clone());gc resolved_api_url = resolved_api_url.or_else(|| selected_org.api_url.clone());gc }gc gc @@ -913,6 +1109,7 @@ async fn resolve_ad_hoc_api_key_auth(gc api_url: resolved_api_url,gc app_url: base.app_url.clone(),gc org_name: resolved_org,gc + org_id: resolved_org_id,gc is_oauth: false,gc slot_key: None,gc })gc @@ -980,6 +1177,7 @@ fn resolve_api_key_profile_auth(gc org_name: effective_org_name(base, cfg_org)gc .map(str::to_string)gc .or_else(|| profile.org_name.clone()),gc + org_id: profile.org_id.clone().or_else(|| base.org_id.clone()),gc is_oauth: false,gc slot_key: Some(profile_name.to_string()),gc };gc @@ -1067,57 +1265,6 @@ fn maybe_rekey_api_key_profile_after_secret_load(gc Ok(())gc }gc gc -async fn reconcile_oauth_slot_from_access_token(gc - store: &mut AuthStore,gc - slot_key: &str,gc - access_token: &str,gc - app_url: &str,gc -) -> Result<()> {gc - let Some(mut profile) = store.profiles.get(slot_key).cloned() else {gc - return Ok(());gc - };gc - if profile.auth_kind != AuthKind::Oauth || profile.org_id.is_some() {gc - return Ok(());gc - }gc -gc - let Some(org_name) = profilegc - .org_namegc - .as_deref()gc - .map(str::trim)gc - .filter(|org| !org.is_empty())gc - else {gc - profile.org_id = Some(String::new());gc - if replace_with_canonical_auth_profile(store, slot_key, profile) {gc - save_auth_store(store)?;gc - }gc - return Ok(());gc - };gc -gc - // A legacy auth entry only cached the org name. Resolve the stable ID fromgc - // the newly refreshed token; a failed best-effort lookup must not turn agc - // successful token refresh into a failed command.gc - let Ok(orgs) = fetch_login_orgs(access_token, app_url).await else {gc - return Ok(());gc - };gc - let Some(org) = find_login_org(&orgs, org_name) else {gc - return Ok(());gc - };gc -gc - profile.org_id = Some(org.id.clone());gc - profile.org_name = Some(org.name.clone());gc - let identity = decode_jwt_identity(access_token);gc - if identity.email.is_some() {gc - profile.email = identity.email;gc - }gc - if identity.name.is_some() {gc - profile.user_name = identity.name;gc - }gc - if replace_with_canonical_auth_profile(store, slot_key, profile) {gc - save_auth_store(store)?;gc - }gc - Ok(())gc -}gc -gc fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Result<()> {gc let Some(slot_key) = auth.slot_key.as_deref() else {gc return Ok(());gc @@ -1127,40 +1274,25 @@ fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Resugc return Ok(());gc };gc gc + if profile.auth_kind == AuthKind::Oauth {gc + return Ok(());gc + }gc let login_org_id = login.org_id().unwrap_or_default();gc - let is_cross_org = profile.auth_kind == AuthKind::Oauthgc - && authgc - .org_namegc - .as_deref()gc - .is_none_or(|org| org.trim().is_empty());gc - if login_org_id.trim().is_empty() && !is_cross_org {gc - // The SDK's OAuth compatibility path does not always return orggc - // metadata. `bt auth logins` performs the same reconciliation aftergc - // its explicit credential verification request.gc + if login_org_id.trim().is_empty() {gc return Ok(());gc }gc gc profile.org_id = Some(login_org_id);gc - profile.org_name = if is_cross_org {gc - Nonegc - } else {gc - logingc - .org_name()gc - .filter(|org| !org.trim().is_empty())gc - .or_else(|| auth.org_name.clone())gc + profile.org_name = logingc + .org_name()gc + .filter(|org| !org.trim().is_empty())gc + .or_else(|| auth.org_name.clone());gc + let Some(api_key) = auth.api_key.as_deref() else {gc + return Ok(());gc };gc - match profile.auth_kind {gc - AuthKind::ApiKey => {gc - let Some(api_key) = auth.api_key.as_deref() else {gc - return Ok(());gc - };gc - profile.api_key_hash = Some(api_key_hash(api_key));gc - if profile.api_key_hint.is_none() {gc - profile.api_key_hint = Some(obscure_api_key(api_key));gc - }gc - }gc - AuthKind::Oauth if profile.email.as_deref().is_none_or(str::is_empty) => return Ok(()),gc - AuthKind::Oauth => {}gc + profile.api_key_hash = Some(api_key_hash(api_key));gc + if profile.api_key_hint.is_none() {gc + profile.api_key_hint = Some(obscure_api_key(api_key));gc }gc gc if replace_with_canonical_auth_profile(&mut store, slot_key, profile) {gc @@ -1169,43 +1301,22 @@ fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Resugc Ok(())gc }gc gc -async fn resolve_oauth_profile_auth(gc +async fn load_oauth_access_token(gc base: &BaseArgs,gc store: &mut AuthStore,gc - cfg_org: &Option,gc profile_name: &str,gc -) -> Result {gc +) -> Result {gc let profile = storegc .profilesgc .get(profile_name)gc .cloned()gc .ok_or_else(|| anyhow::anyhow!("saved OAuth login not found; run `bt auth logins`"))?;gc - let api_url = basegc - .api_urlgc - .clone()gc - .or_else(|| profile.api_url.clone())gc - .unwrap_or_else(|| DEFAULT_API_URL.to_string());gc - let app_url = base.app_url.clone().or_else(|| profile.app_url.clone());gc - let org_name = effective_org_name(base, cfg_org)gc - .map(str::to_string)gc - .or_else(|| profile.org_name.clone());gc -gc - let mut auth = ResolvedAuth {gc - api_key: None,gc - api_url: Some(api_url.clone()),gc - app_url,gc - org_name,gc - is_oauth: true,gc - slot_key: Some(profile_name.to_string()),gc - };gc -gc - if let Some(cached_access_token) = load_valid_cached_oauth_access_token(gc + if let Some(cached) = load_valid_cached_oauth_access_token(gc profile_name,gc &profile,gc profile.oauth_access_expires_at,gc )? {gc - auth.api_key = Some(cached_access_token);gc - return Ok(auth);gc + return Ok(cached);gc }gc gc let refresh_token = load_profile_oauth_refresh_token_for_profile(profile_name, &profile)?gc @@ -1219,7 +1330,11 @@ async fn resolve_oauth_profile_auth(gc ),gc )gc })?;gc - let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &profile).await?;gc + let api_url = basegc + .api_urlgc + .as_deref()gc + .unwrap_or_else(|| profile_api_url(&profile));gc + let refreshed = refresh_oauth_access_token(api_url, &refresh_token, &profile).await?;gc save_profile_oauth_access_token(profile_name, &refreshed.access_token)?;gc let mut refresh_rotated = false;gc if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() {gc @@ -1239,14 +1354,70 @@ async fn resolve_oauth_profile_auth(gc }gc }gc save_auth_store(store)?;gc - reconcile_oauth_slot_from_access_token(gc - store,gc - profile_name,gc - &refreshed.access_token,gc - auth.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc - )gc - .await?;gc - auth.api_key = Some(refreshed.access_token);gc + Ok(refreshed.access_token)gc +}gc +gc +async fn resolve_oauth_profile_auth(gc + base: &BaseArgs,gc + store: &mut AuthStore,gc + cfg_org: &Option,gc + profile_name: &str,gc +) -> Result {gc + let profile = storegc + .profilesgc + .get(profile_name)gc + .cloned()gc + .ok_or_else(|| anyhow::anyhow!("saved OAuth login not found; run `bt auth logins`"))?;gc + let access_token = load_oauth_access_token(base, store, profile_name).await?;gc + let auth = ResolvedAuth {gc + api_key: Some(access_token),gc + api_url: Some(gc + base.api_urlgc + .clone()gc + .or_else(|| profile.api_url.clone())gc + .unwrap_or_else(|| DEFAULT_API_URL.to_string()),gc + ),gc + app_url: Some(gc + base.app_urlgc + .clone()gc + .or_else(|| profile.app_url.clone())gc + .unwrap_or_else(|| DEFAULT_APP_URL.to_string()),gc + ),gc + org_name: effective_org_name(base, cfg_org).map(str::to_string),gc + org_id: base.org_id.clone(),gc + is_oauth: true,gc + slot_key: Some(profile_name.to_string()),gc + };gc + resolve_oauth_org_context(auth).awaitgc +}gc +gc +async fn resolve_oauth_org_context(mut auth: ResolvedAuth) -> Result {gc + let requested_org = auth.org_name.as_deref().ok_or_else(|| {gc + recoverable_auth_error(gc + RecoverableAuthErrorKind::OauthOrgAccess,gc + "an active organization is required; run `bt switch` or pass --org ".to_string(),gc + )gc + })?;gc + let credential = authgc + .api_keygc + .as_deref()gc + .context("OAuth access token is missing")?;gc + let app_url = auth.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let orgs = fetch_login_orgs(credential, app_url).await?;gc + let selected = find_login_org(&orgs, requested_org).ok_or_else(|| {gc + recoverable_auth_error(gc + RecoverableAuthErrorKind::OauthOrgAccess,gc + format!(gc + "OAuth login for '{}' cannot access organization '{requested_org}'",gc + canonical_url(app_url)gc + ),gc + )gc + })?;gc + auth.org_name = Some(selected.name.clone());gc + auth.org_id = Some(selected.id.clone());gc + if auth.api_url.is_none() {gc + auth.api_url = selected.api_url.clone();gc + }gc Ok(auth)gc }gc gc @@ -1281,6 +1452,44 @@ pub async fn resolved_runner_env(base: &BaseArgs) -> Result &str {gc + url.trim().trim_end_matches('/')gc +}gc +gc +fn profile_app_url(profile: &AuthProfile) -> &str {gc + profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL)gc +}gc +gc +fn profile_api_url(profile: &AuthProfile) -> &str {gc + profile.api_url.as_deref().unwrap_or(DEFAULT_API_URL)gc +}gc +gc +fn profile_matches_urls(base: &BaseArgs, profile: &AuthProfile) -> bool {gc + let app_url = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + if canonical_url(app_url) != canonical_url(profile_app_url(profile)) {gc + return false;gc + }gc + profile.auth_kind == AuthKind::Oauthgc + || canonical_url(base.api_url.as_deref().unwrap_or(DEFAULT_API_URL))gc + == canonical_url(profile_api_url(profile))gc +}gc +gc +/// Match only URL filters the caller actually supplied. Listing and logout usegc +/// this variant so an absent filter means "all instances", while command authgc +/// uses [`profile_matches_urls`] and therefore honors the built-in URL defaults.gc +fn profile_matches_url_filters(base: &BaseArgs, profile: &AuthProfile) -> bool {gc + let app_matches = basegc + .app_urlgc + .as_deref()gc + .is_none_or(|url| canonical_url(url) == canonical_url(profile_app_url(profile)));gc + app_matchesgc + && (profile.auth_kind == AuthKind::Oauthgc + || basegc + .api_urlgc + .as_deref()gc + .is_none_or(|url| canonical_url(url) == canonical_url(profile_api_url(profile))))gc +}gc +gc fn profile_matches_org_identifier(profile: &AuthProfile, org: &str) -> bool {gc profile.org_id.as_deref() == Some(org) || profile.org_name.as_deref() == Some(org)gc }gc @@ -1298,13 +1507,13 @@ fn profile_org(profile: &AuthProfile) -> &str {gc }gc gc fn profile_org_label(profile: &AuthProfile) -> String {gc - config::display_org(profile_org(profile)).to_string()gc + profile_org(profile).to_string()gc }gc gc fn oauth_reauth_command(profile: &AuthProfile) -> String {gc format!(gc - "bt auth login --oauth --org {}",gc - shell_quote_arg(config::display_org(profile_org(profile)))gc + "bt auth login --oauth --app-url {}",gc + shell_quote_arg(profile_app_url(profile))gc )gc }gc gc @@ -1333,32 +1542,18 @@ fn profile_identity_label(profile: &AuthProfile) -> Option {gc }gc gc fn auth_slot_label(profile: &AuthProfile) -> String {gc - let mut parts = vec![profile_org_label(profile)];gc - parts.push(auth_kind_label(profile.auth_kind).to_string());gc + let mut parts = match profile.auth_kind {gc + AuthKind::Oauth => vec![profile_app_url(profile).to_string(), "oauth".to_string()],gc + AuthKind::ApiKey => vec![profile_org_label(profile), "api_key".to_string()],gc + };gc if let Some(identity) = profile_identity_label(profile) {gc parts.push(identity);gc }gc parts.join(" — ")gc }gc gc -fn is_cross_org_oauth_profile(profile: &AuthProfile) -> bool {gc - profile.auth_kind == AuthKind::Oauth && profile_org(profile).is_empty()gc -}gc -gc -fn reject_cross_org_api_key_preference(gc - prefer_api_key: bool,gc - org: Option<&str>,gc - store: &AuthStore,gc -) -> Result<()> {gc - let cross_org = org == Some("")gc - || (org.is_none() && store.profiles.values().any(is_cross_org_oauth_profile));gc - if prefer_api_key && cross_org {gc - bail!("--prefer-api-key cannot be used from cross-org context; rerun with --org ");gc - }gc - Ok(())gc -}gc -gc fn auth_profile_names_by_kind<'a>(gc + base: &BaseArgs,gc store: &'a AuthStore,gc org: Option<&str>,gc kind: AuthKind,gc @@ -1367,10 +1562,10 @@ fn auth_profile_names_by_kind<'a>(gc .profilesgc .iter()gc .filter(|(_, profile)| profile.auth_kind == kind)gc - .filter(|(_, profile)| match org {gc - Some(org) => profile_matches_org_identifier(profile, org),gc - None if kind == AuthKind::Oauth => is_cross_org_oauth_profile(profile),gc - None => false,gc + .filter(|(_, profile)| profile_matches_urls(base, profile))gc + .filter(|(_, profile)| {gc + kind == AuthKind::Oauthgc + || org.is_some_and(|org| profile_matches_org_identifier(profile, org))gc })gc .map(|(name, _)| name.as_str())gc .collect()gc @@ -1403,9 +1598,7 @@ fn ad_hoc_api_key_profile(org: Option<&str>, api_key: &str) -> ProfileInfo {gc pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Result> {gc let store = load_auth_store().unwrap_or_default();gc gc - reject_cross_org_api_key_preference(base.prefer_api_key, org, &store)?;gc -gc - let select = |kind| match auth_profile_names_by_kind(&store, org, kind).as_slice() {gc + let select = |kind| match auth_profile_names_by_kind(base, &store, org, kind).as_slice() {gc [] => Ok(None),gc [name] => Ok(Some((*name).to_string())),gc _ => bail!("multiple {kind:?} logins"),gc @@ -1434,13 +1627,7 @@ pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Result Option {gc - let candidates = storegc - .profilesgc - .iter()gc - .filter(|(_, profile)| {gc - !is_cross_org_oauth_profile(profile) && !profile_org(profile).is_empty()gc - })gc - .collect::>();gc + let candidates = store.profiles.iter().collect::>();gc if candidates.is_empty() {gc return None;gc }gc @@ -1483,16 +1670,7 @@ fn select_profile_from_store(gc current: Option<&str>,gc store: &AuthStore,gc ) -> Result {gc - // Surface cross-org OAuth first so it is a predictable top-of-list choicegc - // rather than falling wherever its slot key happens to sort. The stablegc - // sort preserves the existing order of the remaining entries.gc - let mut names: Vec<&str> = names.to_vec();gc - names.sort_by_key(|name| {gc - !storegc - .profilesgc - .get(*name)gc - .is_some_and(is_cross_org_oauth_profile)gc - });gc + let names: Vec<&str> = names.to_vec();gc let labels: Vec = namesgc .iter()gc .map(|name| profile_label_from_store(name, store))gc @@ -1513,53 +1691,6 @@ fn select_profile_from_store(gc Ok(names[idx].to_string())gc }gc gc -fn saved_login_names(store: &AuthStore, include_cross_org: bool) -> Vec<&str> {gc - let mut oauth_orgs = BTreeSet::new();gc - storegc - .profilesgc - .iter()gc - .filter(|(_, profile)| {gc - profile.auth_kind == AuthKind::ApiKeygc - || ((include_cross_org || !is_cross_org_oauth_profile(profile))gc - && oauth_orgs.insert(gc - profilegc - .org_idgc - .as_deref()gc - .filter(|id| !id.is_empty())gc - .unwrap_or_else(|| profile_org(profile))gc - .to_ascii_lowercase(),gc - ))gc - })gc - .map(|(name, _)| name.as_str())gc - .collect()gc -}gc -gc -pub(crate) fn select_saved_login(gc - base: &mut BaseArgs,gc - current_org: Option<&str>,gc - include_cross_org: bool,gc -) -> Result {gc - let store = load_auth_store()?;gc - let names = saved_login_names(&store, include_cross_org);gc - let selected = match names.as_slice() {gc - [] => return Ok(false),gc - [name] => (*name).to_string(),gc - _ if ui::can_prompt() => {gc - select_profile_from_store("Select login", &names, current_org, &store)?gc - }gc - _ => {gc - bail!("multiple saved logins match; pass --org , or rerun interactively to choose")gc - }gc - };gc - let profile = &store.profiles[&selected];gc - if profile.auth_kind == AuthKind::ApiKey {gc - base.pinned_auth_slot = Some(selected);gc - } else {gc - base.org_name = Some(profile_org(profile).to_string());gc - }gc - Ok(true)gc -}gc -gc fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec {gc namesgc .iter()gc @@ -1583,12 +1714,18 @@ fn select_profile_for_auth(gc can_prompt: bool,gc ) -> Result> {gc let org = effective_org_name(base, cfg_org);gc - let candidates = auth_profile_names_by_kind(store, org, kind);gc + let candidates = auth_profile_names_by_kind(base, store, org, kind);gc let label = match kind {gc AuthKind::Oauth => "OAuth login",gc AuthKind::ApiKey => "API key",gc };gc - select_auth_profile_candidate(label, org, &candidates, store, can_prompt)gc + select_auth_profile_candidate(gc + label,gc + org,gc + &candidates,gc + store,gc + can_prompt && kind == AuthKind::ApiKey,gc + )gc }gc gc fn select_auth_profile_candidate(gc @@ -1609,13 +1746,18 @@ fn select_auth_profile_candidate(gc }gc _ => {gc let identities = candidate_identities(candidates, store).join(", ");gc + if kind_label == "OAuth login" {gc + bail!(gc + "multiple Braintrust OAuth instances are available: {identities}. Run `bt switch` or pass --app-url ."gc + );gc + }gc if let Some(org) = org {gc bail!(gc "multiple {kind_label} logins for org '{org}': {identities}. Rerun interactively or remove one with `bt auth logout`."gc );gc }gc bail!(gc - "multiple cross-org {kind_label} logins available: {identities}. Rerun interactively or remove one with `bt auth logout`."gc + "multiple {kind_label} logins available: {identities}. Pass --app-url , rerun interactively, or remove one with `bt auth logout`."gc );gc }gc }gc @@ -1625,10 +1767,12 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> {gc if args.oauth {gc return run_login_oauth(base, args).await;gc }gc - if base.org_name.as_deref() == Some("") {gc - bail!(gc - "API-key login requires a concrete org; cross-org API keys do not exist. Use --oauth, or rerun with --org "gc - );gc + if basegc + .org_namegc + .as_deref()gc + .is_some_and(|org| org.trim().is_empty())gc + {gc + bail!("API-key login requires a non-empty organization");gc }gc gc let has_explicit_api_key = base.api_key.as_ref().is_some_and(|k| !k.trim().is_empty());gc @@ -1662,7 +1806,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> {gc if requested_org_resolution == RequestedOrgResolution::SwitchToOauth {gc return run_login_oauth(base, args).await;gc }gc - let configured_org = config::load().ok().and_then(|cfg| cfg.org);gc + let configured_org = configured_org_for_app_url(&login_app_url);gc let selected_org = select_login_org(gc login_orgs.clone(),gc match requested_org_resolution {gc @@ -1675,7 +1819,6 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> {gc configured_org.as_deref(),gc interactive,gc base.verbose,gc - false,gc explicitly_quiet(base),gc )?;gc let selected_org = selected_org.ok_or_else(|| {gc @@ -1691,7 +1834,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> {gc commit_api_key_profile(gc &api_key,gc selected_api_url.clone(),gc - base.app_url.clone(),gc + Some(login_app_url.clone()),gc selected_org.id.clone(),gc selected_org.name.clone(),gc )?;gc @@ -1784,48 +1927,46 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> {gc exchange_oauth_authorization_code(&api_url, &redirect_uri, &auth_code, pkce_verifier)gc .await?;gc let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?;gc - let configured_org = config::load().ok().and_then(|cfg| cfg.org);gc + let configured_org = configured_org_for_app_url(&app_url);gc let selected_org = select_login_org(gc login_orgs.clone(),gc base.org_name.as_deref(),gc configured_org.as_deref(),gc ui::can_prompt(),gc base.verbose,gc - true,gc explicitly_quiet(base),gc )?;gc + let selected_org = selected_org.ok_or_else(|| {gc + anyhow::anyhow!(gc + "OAuth login requires an organization; pass --org or rerun interactively"gc + )gc + })?;gc let selected_api_url = resolve_profile_api_url(gc base.api_url.clone(),gc - selected_org.as_ref(),gc + Some(&selected_org),gc &login_orgs,gc ui::can_prompt(),gc )?;gc gc - commit_oauth_profile(gc - &oauth_tokens,gc - selected_api_url.clone(),gc - app_url.clone(),gc - selected_org.as_ref(),gc - )?;gc + commit_oauth_profile(&oauth_tokens, api_url.clone(), app_url.clone())?;gc let context_update = persist_post_login_context(gc base,gc &oauth_tokens.access_token,gc &selected_api_url,gc &app_url,gc - selected_org.as_ref(),gc + Some(&selected_org),gc &args.scope,gc )gc .awaitgc .context("login succeeded, but failed to update active context")?;gc gc - let human = format_login_success(selected_org.as_ref(), &selected_api_url);gc + let human = format_login_success(Some(&selected_org), &selected_api_url);gc emit_result(gc base.json,gc serde_json::json!({gc "auth": "oauth",gc - "org": selected_org.as_ref().map(|org| org.name.clone()),gc - "org_id": selected_org.as_ref().map(|org| org.id.clone()),gc - "cross_org": selected_org.is_none(),gc + "org": selected_org.name,gc + "org_id": selected_org.id,gc "api_url": selected_api_url,gc "app_url": app_url,gc "status": "ok",gc @@ -1881,7 +2022,6 @@ fn commit_oauth_profile(gc tokens: &OAuthTokenResponse,gc api_url: String,gc app_url: String,gc - selected_org: Option<&LoginOrgInfo>,gc ) -> Result<()> {gc let refresh_token = tokens.refresh_token.as_ref().ok_or_else(|| {gc anyhow::anyhow!(gc @@ -1891,7 +2031,7 @@ fn commit_oauth_profile(gc gc let oauth_access_expires_at = determine_oauth_access_expiry_epoch(tokens);gc let jwt_id = decode_jwt_identity(&tokens.access_token);gc - let email = jwt_idgc + let _email = jwt_idgc .emailgc .clone()gc .filter(|email| !email.trim().is_empty())gc @@ -1900,25 +2040,25 @@ fn commit_oauth_profile(gc "oauth token did not include an email; cannot create persistent oauth login"gc )gc })?;gc - let org_id = selected_org.map(|org| org.id.clone()).unwrap_or_default();gc - let slot_key = oauth_slot_key(&org_id, &email);gc + let app_url = canonical_url(&app_url).to_string();gc + let slot_key = oauth_slot_key(&app_url);gc gc + let mut store = load_auth_store()?;gc + if let Some(old_profile) = store.profiles.get(&slot_key) {gc + delete_all_profile_secrets(&slot_key, old_profile);gc + }gc save_profile_oauth_refresh_token(&slot_key, refresh_token)?;gc save_profile_oauth_access_token(&slot_key, &tokens.access_token)?;gc let _ = delete_profile_secret(&slot_key);gc gc - let mut store = load_auth_store()?;gc - if let Some(old_profile) = store.profiles.get(&slot_key) {gc - delete_legacy_profile_secrets(old_profile);gc - }gc store.profiles.insert(gc slot_key,gc AuthProfile {gc auth_kind: AuthKind::Oauth,gc api_url: Some(api_url),gc app_url: Some(app_url),gc - org_id: Some(org_id),gc - org_name: selected_org.map(|org| org.name.clone()),gc + org_id: None,gc + org_name: None,gc oauth_access_expires_at,gc user_name: jwt_id.name,gc email: jwt_id.email,gc @@ -1942,7 +2082,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> {gc )?gc .ok_or_else(|| {gc anyhow::anyhow!(gc - "no OAuth login selected; pass --org or run `bt auth logins` to see available logins"gc + "no OAuth login selected; pass --app-url or run `bt auth logins` to see available logins"gc )gc })?;gc let profile = storegc @@ -1953,9 +2093,10 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> {gc anyhow::anyhow!("OAuth login not found; run `bt auth logins` to see available logins")gc })?;gc gc - let api_url = profilegc + let api_url = basegc .api_urlgc .clone()gc + .or_else(|| profile.api_url.clone())gc .unwrap_or_else(|| DEFAULT_API_URL.to_string());gc let previous_expires_at = profile.oauth_access_expires_at;gc let refresh_token =gc @@ -2005,14 +2146,6 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> {gc }gc }gc save_auth_store(&store)?;gc - reconcile_oauth_slot_from_access_token(gc - &mut store,gc - profile_name.as_str(),gc - &refreshed.access_token,gc - profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc - )gc - .await?;gc -gc if let Some(expires_at) = new_expires_at {gc let now = current_unix_timestamp();gc let remaining = expires_at.saturating_sub(now);gc @@ -2030,8 +2163,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> {gc base.json,gc serde_json::json!({gc "auth": "oauth",gc - "org": profile.org_name,gc - "org_id": profile.org_id.filter(|org_id| !org_id.trim().is_empty()),gc + "app_url": profile.app_url,gc "user_email": profile.email,gc "access_expires_at": new_expires_at,gc "refresh_token_rotated": refresh_rotated,gc @@ -2042,10 +2174,9 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> {gc }gc gc fn format_login_success(selected_org: Option<&LoginOrgInfo>, api_url: &str) -> String {gc - match selected_org {gc - Some(org) => format!("Logged in as {} (api: {api_url})", org.name),gc - None => format!("Logged in (cross-org, api: {api_url})"),gc - }gc + selected_orggc + .map(|org| format!("Logged in as {} (api: {api_url})", org.name))gc + .unwrap_or_else(|| format!("Logged in (api: {api_url})"))gc }gc gc fn build_login_context_for_selected_org(gc @@ -2076,7 +2207,7 @@ fn format_post_login_context(gc match (selected_org, project) {gc (Some(org), Some(project)) => format!("{}/{}", org.name, project.name),gc (Some(org), None) => org.name.clone(),gc - (None, _) => "cross-org mode".to_string(),gc + (None, _) => "Braintrust".to_string(),gc }gc }gc gc @@ -2091,11 +2222,8 @@ async fn resolve_post_login_project(gc return Ok(None);gc };gc gc - let selected_org = selected_org.ok_or_else(|| {gc - anyhow::anyhow!(gc - "cannot set a default project in cross-org mode; rerun `bt auth login --org --project `"gc - )gc - })?;gc + let selected_org = selected_orggc + .ok_or_else(|| anyhow::anyhow!("an organization is required to select a project"))?;gc let ctx =gc build_login_context_for_selected_org(credential, api_url, app_url, Some(selected_org));gc let client = ApiClient::new(&ctx)?;gc @@ -2117,23 +2245,35 @@ async fn persist_post_login_context(gc resolve_post_login_project(base, credential, api_url, app_url, selected_org).await?;gc let (path, _) = scope.resolve(ui::can_prompt(), "Where to use this login")?;gc let mut cfg = config::load_file(&path);gc - let org = selected_org.map_or("", |org| org.name.as_str());gc + let selected_org = selected_orggc + .ok_or_else(|| anyhow::anyhow!("an organization is required to update config"))?;gc let preserve_project = project.is_none()gc - && selected_org.is_some()gc - && config::org_option(cfg.org.as_deref()) == Some(org);gc - if !preserve_project {gc - cfg.set_context(gc - Some(org),gc - projectgc - .as_ref()gc - .map(|project| (project.name.as_str(), project.id.as_str())),gc - );gc - }gc + && config::org_option(cfg.org.as_deref()) == Some(selected_org.name.as_str())gc + && cfg.org_id.as_deref() == Some(selected_org.id.as_str())gc + && cfggc + .app_urlgc + .as_deref()gc + .is_some_and(|url| config::urls_equal(url, app_url));gc + let selected_project = if preserve_project {gc + cfg.project.clone().zip(cfg.project_id.clone())gc + } else {gc + projectgc + .as_ref()gc + .map(|project| (project.name.clone(), project.id.clone()))gc + };gc + cfg.set_context(gc + (selected_org.name.as_str(), selected_org.id.as_str()),gc + selected_projectgc + .as_ref()gc + .map(|(name, id)| (name.as_str(), id.as_str())),gc + app_url,gc + api_url,gc + );gc config::save_file(&path, &cfg)gc .with_context(|| format!("Could not save config to {}", path.display()))?;gc gc Ok(PostLoginContextUpdate {gc - display: format_post_login_context(selected_org, project.as_ref()),gc + display: format_post_login_context(Some(selected_org), project.as_ref()),gc path,gc })gc }gc @@ -2150,14 +2290,14 @@ fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) ->gc }gc gc fn filter_auth_store(gc + base: &BaseArgs,gc store: &AuthStore,gc - org: Option<&str>,gc kind: Option,gc api_key_hint: Option<&str>,gc ) -> AuthStore {gc let mut filtered = store.clone();gc filtered.profiles.retain(|_, profile| {gc - org.is_none_or(|org| profile_matches_org_identifier(profile, org))gc + profile_matches_url_filters(base, profile)gc && kind.is_none_or(|kind| profile.auth_kind == kind)gc && api_key_hint.is_none_or(|hint| {gc profile.auth_kind == AuthKind::ApiKeygc @@ -2167,15 +2307,68 @@ fn filter_auth_store(gc filteredgc }gc gc +async fn filter_auth_store_for_org(gc + base: &BaseArgs,gc + store: &mut AuthStore,gc + candidates: AuthStore,gc + org: Option<&str>,gc +) -> Result {gc + let Some(org) = org else {gc + return Ok(candidates);gc + };gc + let mut filtered = AuthStore::default();gc + for (slot, profile) in candidates.profiles {gc + let matches = match profile.auth_kind {gc + AuthKind::ApiKey => profile_matches_org_identifier(&profile, org),gc + AuthKind::Oauth => {gc + let mut oauth_base = base.clone();gc + oauth_base.app_url = Some(profile_app_url(&profile).to_string());gc + if oauth_base.api_url_source.is_none() {gc + oauth_base.api_url = profile.api_url.clone();gc + }gc + let token = load_oauth_access_token(&oauth_base, store, &slot).await?;gc + let orgs = fetch_login_orgs(&token, profile_app_url(&profile)).await?;gc + find_login_org(&orgs, org).is_some()gc + }gc + };gc + if matches {gc + filtered.profiles.insert(slot, profile);gc + }gc + }gc + Ok(filtered)gc +}gc +gc async fn run_logins(base: &BaseArgs, _args: AuthLoginsArgs) -> Result<()> {gc let mut store = load_auth_store()?;gc - let has_filter = base.org_name.is_some() || base.prefer_api_key;gc - let filtered = filter_auth_store(gc + let requested_org = matches!(gc + base.org_name_source,gc + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable)gc + )gc + .then(|| base.org_name.as_deref())gc + .flatten();gc + let has_url_filter = matches!(gc + base.app_url_source,gc + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable)gc + ) || matches!(gc + base.api_url_source,gc + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable)gc + );gc + let has_filter = requested_org.is_some() || base.prefer_api_key || has_url_filter;gc + let mut filter_base = base.clone();gc + if filter_base.app_url_source.is_none() {gc + filter_base.app_url = None;gc + }gc + if filter_base.api_url_source.is_none() {gc + filter_base.api_url = None;gc + }gc + let candidates = filter_auth_store(gc + &filter_base,gc &store,gc - base.org_name.as_deref(),gc base.prefer_api_key.then_some(AuthKind::ApiKey),gc None,gc );gc + let filtered =gc + filter_auth_store_for_org(&filter_base, &mut store, candidates, requested_org).await?;gc if filtered.profiles.is_empty() {gc return emit_result(base.json, serde_json::json!([]), || {gc if store.profiles.is_empty() && !has_filter {gc @@ -2226,6 +2419,8 @@ fn auth_profile_json(profile: &AuthProfile, status: &str) -> serde_json::Value {gc "user_name": profile.user_name,gc "user_email": profile.email,gc "api_key_hint": profile.api_key_hint,gc + "app_url": profile.app_url,gc + "api_url": profile.api_url,gc "status": status,gc })gc }gc @@ -2277,7 +2472,7 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result Result<()> {gc +async fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> {gc let store = load_auth_store()?;gc if store.profiles.is_empty() {gc return emit_result(base.json, serde_json::json!({ "status": "empty" }), || {gc @@ -2287,18 +2482,29 @@ fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> {gc gc let requested_org = if matches!(gc base.org_name_source,gc - Some(crate::args::ArgValueSource::CommandLine)gc + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable)gc ) {gc config::org_option(base.org_name.as_deref())gc } else {gc Nonegc };gc - let filtered = filter_auth_store(gc + let mut filter_base = base.clone();gc + if filter_base.app_url_source.is_none() {gc + filter_base.app_url = None;gc + }gc + if filter_base.api_url_source.is_none() {gc + filter_base.api_url = None;gc + }gc + let candidates = filter_auth_store(gc + &filter_base,gc &store,gc - requested_org,gc args.oauth.then_some(AuthKind::Oauth),gc args.api_key_hint.as_deref(),gc );gc + let mut mutable_store = store.clone();gc + let filtered =gc + filter_auth_store_for_org(&filter_base, &mut mutable_store, candidates, requested_org)gc + .await?;gc let candidates = filteredgc .profilesgc .keys()gc @@ -2319,7 +2525,7 @@ fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> {gc _ => {gc let labels = candidate_identities(&candidates, &filtered).join(", ");gc bail!(gc - "multiple auth logins match: {labels}. Rerun interactively, or use --org with --oauth or --api-key-hint to disambiguate."gc + "multiple auth logins match: {labels}. Rerun interactively, use --app-url with --oauth, or use --org with --api-key-hint ."gc );gc }gc };gc @@ -2382,6 +2588,10 @@ pub struct ProfileVerification {gc pub user_email: Option,gc #[serde(skip_serializing_if = "Option::is_none")]gc pub api_key_hint: Option,gc + #[serde(skip_serializing_if = "Option::is_none")]gc + pub app_url: Option,gc + #[serde(skip_serializing_if = "Option::is_none")]gc + pub api_url: Option,gc pub status: String,gc #[serde(skip_serializing_if = "Option::is_none")]gc pub error: Option,gc @@ -2389,9 +2599,7 @@ pub struct ProfileVerification {gc gc fn build_verification(gc name: &str,gc - auth_kind: &str,gc - org: Option,gc - org_id: Option,gc + profile: &AuthProfile,gc jwt_id: Option,gc api_key_hint: Option,gc status: ProfileStatus,gc @@ -2405,12 +2613,17 @@ fn build_verification(gc ProfileVerification {gc name: name.to_string(),gc slot_hash: None,gc - auth: auth_kind.to_string(),gc - org,gc - org_id,gc + auth: auth_kind_label(profile.auth_kind).to_string(),gc + org: profile.org_name.clone(),gc + org_id: profilegc + .org_idgc + .clone()gc + .filter(|org_id| !org_id.trim().is_empty()),gc user_name: jwt_id.as_ref().and_then(|j| j.name.clone()),gc user_email: jwt_id.as_ref().and_then(|j| j.email.clone()),gc api_key_hint,gc + app_url: profile.app_url.clone(),gc + api_url: profile.api_url.clone(),gc status: status_str.to_string(),gc error,gc }gc @@ -2418,20 +2631,8 @@ fn build_verification(gc gc async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerification {gc let app_url = profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc - let auth_kind = auth_kind_label(profile.auth_kind);gc let mk = |status, jwt_id: Option, hint: Option| {gc - build_verification(gc - name,gc - auth_kind,gc - profile.org_name.clone(),gc - profilegc - .org_idgc - .clone()gc - .filter(|org_id| !org_id.trim().is_empty()),gc - jwt_id,gc - hint,gc - status,gc - )gc + build_verification(name, profile, jwt_id, hint, status)gc };gc gc let credential = match load_credential_for_profile(name, profile) {gc @@ -2455,7 +2656,7 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifigc match fetch_login_orgs(&credential, app_url).await {gc Ok(orgs) => {gc let mut verification = mk(ProfileStatus::Ok, jwt_id, hint);gc - if !is_cross_org_oauth_profile(profile) {gc + if profile.auth_kind == AuthKind::ApiKey {gc if let Some(org) = profilegc .org_idgc .as_deref()gc @@ -2470,8 +2671,6 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifigc verification.org = Some(org.name.clone());gc verification.org_id = Some(org.id.clone());gc }gc - }gc - if profile.auth_kind == AuthKind::ApiKey {gc verification.slot_hash = Some(api_key_hash(&credential));gc }gc verificationgc @@ -2532,11 +2731,13 @@ fn reconcile_verified_auth_slots(gc continue;gc };gc gc - if let Some(org_id) = verification.org_id.as_deref() {gc - profile.org_id = Some(org_id.to_string());gc - profile.org_name = verification.org.clone();gc - } else if is_cross_org_oauth_profile(&profile) {gc - profile.org_id = Some(String::new());gc + if profile.auth_kind == AuthKind::ApiKey {gc + if let Some(org_id) = verification.org_id.as_deref() {gc + profile.org_id = Some(org_id.to_string());gc + profile.org_name = verification.org.clone();gc + }gc + } else {gc + profile.org_id = None;gc profile.org_name = None;gc }gc gc @@ -2573,10 +2774,14 @@ fn reconcile_verified_auth_slots(gc }gc gc fn format_verification_line(v: &ProfileVerification) -> String {gc - let mut parts = vec![gc - config::display_org(v.org.as_deref().unwrap_or("")).to_string(),gc - v.auth.clone(),gc - ];gc + let subject = if v.auth == "oauth" {gc + v.app_urlgc + .clone()gc + .unwrap_or_else(|| DEFAULT_APP_URL.to_string())gc + } else {gc + v.org.clone().unwrap_or_else(|| "(unknown org)".to_string())gc + };gc + let mut parts = vec![subject, v.auth.clone()];gc match v.status.as_str() {gc "ok" => {gc if let Some(id) = identity_label(gc @@ -2628,6 +2833,8 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> {gc "user_name": p.user_name,gc "user_email": p.email,gc "api_key_hint": p.api_key_hint,gc + "app_url": p.app_url,gc + "api_url": p.api_url,gc "status": "unchecked"gc })gc })gc @@ -2676,7 +2883,6 @@ fn select_login_org(gc default_org_name: Option<&str>,gc interactive: bool,gc verbose: bool,gc - allow_cross_org: bool,gc quiet_requested: bool,gc ) -> Result> {gc if orgs.is_empty() {gc @@ -2684,10 +2890,6 @@ fn select_login_org(gc }gc sort_login_orgs(&mut orgs);gc gc - if requested_org_name == Some("") {gc - return Ok(None);gc - }gc -gc if let Some(name) = requested_org_name {gc return find_login_org(&orgs, name)gc .cloned()gc @@ -2700,44 +2902,30 @@ fn select_login_org(gc }gc gc if !interactive {gc - if allow_cross_org {gc - bail!(gc - "organization selection required in non-interactive mode; pass --org or rerun interactively to choose cross-org mode"gc - );gc - }gc return Ok(None);gc }gc gc - let default_org_matched = move_default_login_org_first(&mut orgs, default_org_name);gc - let offset = if allow_cross_org { 1 } else { 0 };gc - let mut labels: Vec = Vec::new();gc - if allow_cross_org {gc - labels.push(gc - "No default org (cross-org mode; pass --org or BRAINTRUST_ORG_NAME when needed)"gc - .to_string(),gc - );gc - }gc - labels.extend(orgs.iter().map(|org| {gc - if verbose {gc - let api_url = org.api_url.as_deref().unwrap_or(DEFAULT_API_URL);gc - format!("{} [{}] ({})", org.name, org.id, api_url)gc - } else {gc - org.name.clone()gc - }gc - }));gc + move_default_login_org_first(&mut orgs, default_org_name);gc + let labels: Vec = orgsgc + .iter()gc + .map(|org| {gc + if verbose {gc + let api_url = org.api_url.as_deref().unwrap_or(DEFAULT_API_URL);gc + format!("{} [{}] ({})", org.name, org.id, api_url)gc + } else {gc + org.name.clone()gc + }gc + })gc + .collect();gc let label_refs: Vec<&str> = labels.iter().map(String::as_str).collect();gc if !quiet_requested {gc eprintln!("\n\nA Braintrust organization is usually a team or a company.");gc }gc - let default = if default_org_matched { offset } else { 0 };gc - let selection = ui::fuzzy_select("Select organization", &label_refs, default)?;gc - if allow_cross_org && selection == 0 {gc - return Ok(None);gc - }gc + let selection = ui::fuzzy_select("Select organization", &label_refs, 0)?;gc gc Ok(Some(gc orgs.into_iter()gc - .nth(selection - offset)gc + .nth(selection)gc .expect("selected index should be in range"),gc ))gc }gc @@ -2855,8 +3043,11 @@ fn resolve_profile_api_url(gc if let Some(api_url) = explicit_api_url {gc return Ok(api_url);gc }gc - if let Some(api_url) = selected_org.and_then(|org| org.api_url.clone()) {gc - return Ok(api_url);gc + if let Some(selected_org) = selected_org {gc + return Ok(selected_orggc + .api_urlgc + .clone()gc + .unwrap_or_else(|| DEFAULT_API_URL.to_string()));gc }gc gc let mut api_urls = orgsgc @@ -2873,8 +3064,6 @@ fn resolve_profile_api_url(gc .unwrap_or_else(|| DEFAULT_API_URL.to_string()));gc }gc gc - // A cross-org login spans orgs on different data planes. Let the user pickgc - // which API URL to store rather than failing outright.gc if can_prompt {gc let idx = ui::fuzzy_select("Select API URL", &api_urls, 0)?;gc return Ok(api_urlsgc @@ -3980,8 +4169,8 @@ fn api_key_hash(api_key: &str) -> String {gc sha256_hex(api_key)gc }gc gc -fn oauth_slot_key(org_id: &str, email: &str) -> String {gc - format!("{org_id}::{email}")gc +fn oauth_slot_key(app_url: &str) -> String {gc + format!("oauth::{}", sha256_hex(canonical_url(app_url)))gc }gc gc fn api_key_slot_key(api_key_hash: &str, org_id: &str) -> String {gc @@ -4030,29 +4219,38 @@ fn load_auth_store_from_path(path: &Path) -> Result {gc gc fn migrate_auth_store(store: AuthStore) -> AuthStore {gc let mut migrated = AuthStore::default();gc + let mut oauth_refresh_usable = BTreeMap::::new();gc for (old_key, mut profile) in store.profiles {gc normalize_profile_cached_fields_from_key(&old_key, &mut profile);gc - if profile.auth_kind == AuthKind::Oauthgc - && profile.org_id.is_none()gc - && profilegc - .org_namegc - .as_deref()gc - .is_none_or(|org| org.trim().is_empty())gc - {gc - profile.org_id = Some(String::new());gc + let refresh_usable = profile.auth_kind == AuthKind::Oauthgc + && matches!(gc + load_profile_oauth_refresh_token_for_profile(&old_key, &profile),gc + Ok(Some(_))gc + );gc + if profile.auth_kind == AuthKind::Oauth {gc + profile.org_id = None;gc profile.org_name = None;gc + profile.app_url = Some(canonical_url(profile_app_url(&profile)).to_string());gc }gc let new_key = canonical_profile_key(&old_key, &profile);gc if new_key != old_key && profile.legacy_secret_key.is_none() {gc profile.legacy_secret_key = Some(old_key.clone());gc }gc - if migratedgc - .profilesgc - .get(&new_key)gc - .is_some_and(|existing| !should_replace_migrated_profile(existing, &profile))gc - {gc - continue;gc + if let Some(existing) = migrated.profiles.get(&new_key) {gc + let existing_usable = oauth_refresh_usable.get(&new_key).copied().unwrap_or(false);gc + let replace = match (existing.auth_kind, profile.auth_kind) {gc + (AuthKind::Oauth, AuthKind::Oauth) => {gc + (refresh_usable && !existing_usable)gc + || (refresh_usable == existing_usablegc + && should_replace_migrated_profile(existing, &profile))gc + }gc + _ => false,gc + };gc + if !replace {gc + continue;gc + }gc }gc + oauth_refresh_usable.insert(new_key.clone(), refresh_usable);gc migrated.profiles.insert(new_key, profile);gc }gc migratedgc @@ -4166,17 +4364,7 @@ fn normalize_profile_cached_fields_from_key(current_key: &str, profile: &mut Autgc gc fn canonical_profile_key(current_key: &str, profile: &AuthProfile) -> String {gc match profile.auth_kind {gc - AuthKind::Oauth => {gc - let Some(email) = profile.email.as_deref().filter(|value| !value.is_empty()) else {gc - return current_key.to_string();gc - };gc - let org_id = profile.org_id.as_deref().unwrap_or_default();gc - if profile.org_id.is_some() || profile.org_name.is_none() {gc - oauth_slot_key(org_id, email)gc - } else {gc - current_key.to_string()gc - }gc - }gc + AuthKind::Oauth => oauth_slot_key(profile_app_url(profile)),gc AuthKind::ApiKey => match (gc profilegc .api_key_hashgc @@ -4537,6 +4725,13 @@ mod tests {gc crate::config::save_global(&cfg).expect("save global config");gc }gc gc + fn set_global_config_urls(app_url: &str, api_url: Option<&str>) {gc + let mut cfg = crate::config::load_global().expect("load global config");gc + cfg.app_url = Some(app_url.to_string());gc + cfg.api_url = api_url.map(str::to_string);gc + crate::config::save_global(&cfg).expect("save global config URLs");gc + }gc +gc fn org_profile(kind: AuthKind, org_id: &str, org_name: &str) -> AuthProfile {gc AuthProfile {gc auth_kind: kind,gc @@ -4612,7 +4807,10 @@ mod tests {gc fn invalid_grant_refresh_error_is_treated_as_recoverable() {gc let profile = org_profile(AuthKind::Oauth, "org_test", "BT Staging");gc let command = oauth_reauth_command(&profile);gc - assert_eq!(command, "bt auth login --oauth --org 'BT Staging'");gc + assert_eq!(gc + command,gc + "bt auth login --oauth --app-url https://www.braintrust.dev"gc + );gc let err = map_refresh_oauth_error(gc "https://api.example.com",gc &profile,gc @@ -4751,62 +4949,20 @@ mod tests {gc assert_eq!(DEFAULT_APP_URL, "https://www.braintrust.dev");gc }gc gc - #[tokio::test]gc - async fn active_auth_info_without_org_uses_cross_org_oauth_only() {gc - let _env = TestEnv::new(None, None).await;gc - let mut store = AuthStore::default();gc - store.profiles.insert(gc - api_key_slot_key(&api_key_hash("test-api-key"), "org_fake"),gc - AuthProfile {gc - api_key_hint: Some("sk-****abcde".to_string()),gc - ..org_profile(AuthKind::ApiKey, "org_fake", "test-org")gc - },gc - );gc - store.profiles.insert(gc - oauth_slot_key("", "user@example.test"),gc - AuthProfile {gc - auth_kind: AuthKind::Oauth,gc - org_id: Some(String::new()),gc - org_name: None,gc - user_name: Some("Test User".to_string()),gc - email: Some("user@example.test".to_string()),gc - ..Default::default()gc - },gc - );gc - save_auth_store(&store).expect("save auth store");gc -gc - let info = active_auth_info(&make_base(), None)gc - .expect("resolve active auth")gc - .expect("active auth info");gc -gc - assert_eq!(info.auth_method, "oauth");gc - assert_eq!(info.email.as_deref(), Some("user@example.test"));gc - assert_eq!(info.org_name, None);gc -gc - let mut base = make_base();gc - base.prefer_api_key = true;gc - assert!(resolve_auth(&base)gc - .awaitgc - .unwrap_err()gc - .to_string()gc - .contains("cross-org"));gc - assert!(active_auth_info(&base, None)gc - .unwrap_err()gc - .to_string()gc - .contains("cross-org"));gc - }gc -gc - fn save_cached_oauth_login(store: &mut AuthStore, org_id: &str, org_name: &str) -> String {gc - let slot_key = oauth_slot_key(org_id, "user@example.test");gc + fn save_cached_oauth_login(store: &mut AuthStore, app_url: &str) -> String {gc + let slot_key = oauth_slot_key(app_url);gc store.profiles.insert(gc slot_key.clone(),gc AuthProfile {gc api_url: Some("https://api.example.test".to_string()),gc - app_url: Some("https://www.example.test".to_string()),gc + app_url: Some(app_url.to_string()),gc oauth_access_expires_at: Some(current_unix_timestamp() + 3600),gc user_name: Some("Test User".to_string()),gc email: Some("user@example.test".to_string()),gc - ..org_profile(AuthKind::Oauth, org_id, org_name)gc + auth_kind: AuthKind::Oauth,gc + org_id: None,gc + org_name: None,gc + ..Default::default()gc },gc );gc save_profile_secret_plaintext(gc @@ -4821,10 +4977,12 @@ mod tests {gc async fn auth_precedence_keeps_env_api_key_below_oauth() {gc let _env = TestEnv::new(None, None).await;gc let mut store = AuthStore::default();gc - save_cached_oauth_login(&mut store, "org_fake", "test-org");gc + let app_url = spawn_api_key_login_server("test-org");gc + save_cached_oauth_login(&mut store, &app_url);gc save_auth_store(&store).expect("save auth store");gc let mut base = make_base();gc base.org_name = Some("test-org".to_string());gc + base.app_url = Some(app_url);gc base.api_key = Some("environment-api-key".to_string());gc base.api_key_source = Some(crate::args::ArgValueSource::EnvVariable);gc gc @@ -4841,13 +4999,14 @@ mod tests {gc async fn auth_precedence_cli_api_key_overrides_oauth() {gc let _env = TestEnv::new(None, None).await;gc let mut store = AuthStore::default();gc - save_cached_oauth_login(&mut store, "org_fake", "test-org");gc + let app_url = spawn_api_key_login_server("test-org");gc + save_cached_oauth_login(&mut store, &app_url);gc save_auth_store(&store).expect("save auth store");gc let mut base = make_base();gc base.org_name = Some("test-org".to_string());gc base.api_key = Some("command-line-api-key".to_string());gc base.api_key_source = Some(crate::args::ArgValueSource::CommandLine);gc - base.app_url = Some(spawn_api_key_login_server("test-org"));gc + base.app_url = Some(app_url);gc gc let resolved = resolve_auth(&base).await.expect("resolve auth");gc gc @@ -4899,14 +5058,15 @@ mod tests {gc async fn auth_precedence_prefer_api_key_promotes_env_api_key() {gc let _env = TestEnv::new(None, None).await;gc let mut store = AuthStore::default();gc - save_cached_oauth_login(&mut store, "org_fake", "test-org");gc + let app_url = spawn_api_key_login_server("test-org");gc + save_cached_oauth_login(&mut store, &app_url);gc save_auth_store(&store).expect("save auth store");gc let mut base = make_base();gc base.org_name = Some("test-org".to_string());gc base.api_key = Some("environment-api-key".to_string());gc base.api_key_source = Some(crate::args::ArgValueSource::EnvVariable);gc base.prefer_api_key = true;gc - base.app_url = Some(spawn_api_key_login_server("test-org"));gc + base.app_url = Some(app_url);gc gc let resolved = resolve_auth(&base).await.expect("resolve auth");gc gc @@ -4918,10 +5078,12 @@ mod tests {gc async fn auth_precedence_prefer_api_key_falls_back_to_oauth_without_key() {gc let _env = TestEnv::new(None, None).await;gc let mut store = AuthStore::default();gc - save_cached_oauth_login(&mut store, "org_fake", "test-org");gc + let app_url = spawn_api_key_login_server("test-org");gc + save_cached_oauth_login(&mut store, &app_url);gc save_auth_store(&store).expect("save auth store");gc let mut base = make_base();gc base.org_name = Some("test-org".to_string());gc + base.app_url = Some(app_url);gc base.prefer_api_key = true;gc gc let resolved = resolve_auth(&base).await.expect("resolve auth");gc @@ -4960,11 +5122,13 @@ mod tests {gc let _env = TestEnv::new(None, None).await;gc let mut store = AuthStore::default();gc store.profiles.insert(gc - oauth_slot_key("org_fake", "user@example.test"),gc + oauth_slot_key(DEFAULT_APP_URL),gc AuthProfile {gc + auth_kind: AuthKind::Oauth,gc + app_url: Some(DEFAULT_APP_URL.to_string()),gc user_name: Some("Test User".to_string()),gc email: Some("user@example.test".to_string()),gc - ..org_profile(AuthKind::Oauth, "org_fake", "test-org")gc + ..Default::default()gc },gc );gc store.profiles.insert(gc @@ -5051,20 +5215,23 @@ mod tests {gc );gc gc let migrated = migrate_auth_store(store);gc - let key = oauth_slot_key("org_fake", "user@example.test");gc + let key = oauth_slot_key(DEFAULT_APP_URL);gc let profile = migrated.profiles.get(&key).expect("migrated profile");gc gc assert_eq!(profile.legacy_secret_key.as_deref(), Some("work"));gc + assert_eq!(profile.org_id, None);gc + assert_eq!(profile.org_name, None);gc }gc gc #[test]gc - fn migrate_auth_store_rekeys_cross_org_oauth_with_empty_org_id() {gc + fn migrate_auth_store_purges_old_oauth_org_scope() {gc let mut store = AuthStore::default();gc store.profiles.insert(gc - "legacy-cross-org".to_string(),gc + "legacy-login".to_string(),gc AuthProfile {gc auth_kind: AuthKind::Oauth,gc - org_name: None,gc + org_id: Some("org_old".to_string()),gc + org_name: Some("old-org".to_string()),gc email: Some("user@example.test".to_string()),gc ..Default::default()gc },gc @@ -5073,14 +5240,12 @@ mod tests {gc let migrated = migrate_auth_store(store);gc let profile = migratedgc .profilesgc - .get(&oauth_slot_key("", "user@example.test"))gc - .expect("cross-org OAuth slot");gc + .get(&oauth_slot_key(DEFAULT_APP_URL))gc + .expect("instance OAuth slot");gc gc - assert_eq!(profile.org_id.as_deref(), Some(""));gc - assert_eq!(gc - profile.legacy_secret_key.as_deref(),gc - Some("legacy-cross-org")gc - );gc + assert_eq!(profile.org_id, None);gc + assert_eq!(profile.org_name, None);gc + assert_eq!(profile.legacy_secret_key.as_deref(), Some("legacy-login"));gc }gc gc #[test]gc @@ -5134,7 +5299,7 @@ mod tests {gc let persisted: AuthStore =gc serde_json::from_str(&fs::read_to_string(&path).expect("read migrated store"))gc .expect("parse migrated store");gc - let slot_key = oauth_slot_key("org_fake", "user@example.test");gc + let slot_key = oauth_slot_key(DEFAULT_APP_URL);gc gc for migrated in [&loaded, &persisted] {gc let profile = migratedgc @@ -5179,6 +5344,8 @@ mod tests {gc user_name: Some("Test User".to_string()),gc user_email: Some("user@example.test".to_string()),gc api_key_hint: None,gc + app_url: Some(DEFAULT_APP_URL.to_string()),gc + api_url: None,gc status: "ok".to_string(),gc error: None,gc },gc @@ -5191,6 +5358,8 @@ mod tests {gc user_name: None,gc user_email: None,gc api_key_hint: Some("sk-****abcde".to_string()),gc + app_url: None,gc + api_url: None,gc status: "ok".to_string(),gc error: None,gc },gc @@ -5201,7 +5370,7 @@ mod tests {gc gc let oauth = storegc .profilesgc - .get(&oauth_slot_key("org_fake", "user@example.test"))gc + .get(&oauth_slot_key(DEFAULT_APP_URL))gc .expect("canonical OAuth slot");gc assert_eq!(oauth.legacy_secret_key.as_deref(), Some("legacy-oauth"));gc let api_key = storegc @@ -5229,12 +5398,40 @@ mod tests {gc }gc gc let migrated = migrate_auth_store(store);gc - let key = oauth_slot_key("org_fake", "user@example.test");gc + let key = oauth_slot_key(DEFAULT_APP_URL);gc assert_eq!(migrated.profiles.len(), 1);gc let profile = migrated.profiles.get(&key).expect("migrated profile");gc assert_eq!(profile.legacy_secret_key.as_deref(), Some("new"));gc }gc gc + #[tokio::test]gc + async fn migrate_auth_store_prefers_loadable_refresh_token_before_expiry() {gc + let _env = TestEnv::new(None, None).await;gc + let mut store = AuthStore::default();gc + for (name, expires_at) in [("usable-old", 10), ("missing-new", 20)] {gc + store.profiles.insert(gc + name.to_string(),gc + AuthProfile {gc + auth_kind: AuthKind::Oauth,gc + oauth_access_expires_at: Some(expires_at),gc + ..Default::default()gc + },gc + );gc + }gc + save_profile_secret_plaintext(gc + &oauth_refresh_secret_key("usable-old"),gc + "test-refresh-token",gc + )gc + .expect("save refresh token");gc +gc + let migrated = migrate_auth_store(store);gc + let profile = migratedgc + .profilesgc + .get(&oauth_slot_key(DEFAULT_APP_URL))gc + .expect("migrated OAuth profile");gc + assert_eq!(profile.legacy_secret_key.as_deref(), Some("usable-old"));gc + }gc +gc #[test]gc fn migration_reports_dropped_duplicate_secret_as_orphan() {gc // Two legacy OAuth entries for the same org+email collapse onto onegc @@ -5293,6 +5490,12 @@ mod tests {gc let org = config_auth_context_from_config(&base, &cfg);gc gc assert_eq!(org.as_deref(), Some("local-org"));gc +gc + let other_instance = BaseArgs {gc + app_url: Some("https://other.example.test".into()),gc + ..basegc + };gc + assert_eq!(config_auth_context_from_config(&other_instance, &cfg), None);gc }gc gc #[test]gc @@ -5430,80 +5633,143 @@ mod tests {gc assert!(err.to_string().contains("multiple oauth logins"));gc }gc gc - fn login_filter_store() -> AuthStore {gc + #[tokio::test]gc + async fn available_instances_dedupes_logins_by_app_url() {gc + let _env = TestEnv::new(None, None).await;gc let mut store = AuthStore::default();gc - for (slot, kind, suffix, hint) in [gc - ("oauth-a", AuthKind::Oauth, "a", None),gc - ("key-a", AuthKind::ApiKey, "a", Some("sk-****aaaaa")),gc - ("key-b", AuthKind::ApiKey, "b", Some("sk-****bbbbb")),gc + for (slot, kind, app_url) in [gc + ("oauth-a", AuthKind::Oauth, "https://one.example.test/"),gc + ("key-a", AuthKind::ApiKey, "https://one.example.test"),gc + ("key-b", AuthKind::ApiKey, "https://two.example.test"),gc ] {gc store.profiles.insert(gc slot.into(),gc AuthProfile {gc - api_key_hint: hint.map(str::to_string),gc - ..org_profile(gc - kind,gc - &format!("org_test_{suffix}"),gc - &format!("test-org-{suffix}"),gc - )gc + auth_kind: kind,gc + app_url: Some(app_url.into()),gc + org_id: (kind == AuthKind::ApiKey).then(|| format!("org_{slot}")),gc + org_name: (kind == AuthKind::ApiKey).then(|| format!("org-{slot}")),gc + ..Default::default()gc },gc );gc }gc + save_auth_store(&store).expect("save auth store");gc +gc + let instances = available_instances(&BaseArgs::default()).expect("list instances");gc + assert_eq!(gc + instancesgc + .iter()gc + .map(|instance| instance.app_url.as_str())gc + .collect::>(),gc + ["https://one.example.test", "https://two.example.test"]gc + );gc +gc + let filtered = available_instances(&BaseArgs {gc + app_url: Some("https://two.example.test/".into()),gc + app_url_source: Some(crate::args::ArgValueSource::CommandLine),gc + ..Default::default()gc + })gc + .expect("filter instances");gc + assert_eq!(filtered.len(), 1);gc + assert_eq!(filtered[0].app_url, "https://two.example.test");gc + }gc +gc + #[tokio::test]gc + async fn org_filter_includes_oauth_login_when_discovered_membership_matches() {gc + let _env = TestEnv::new(None, None).await;gc + let app_url = spawn_api_key_login_server("test-org");gc + let mut store = AuthStore::default();gc + let oauth_slot = save_cached_oauth_login(&mut store, &app_url);gc store.profiles.insert(gc - "cross".into(),gc + "other-key".into(),gc AuthProfile {gc - auth_kind: AuthKind::Oauth,gc - org_id: Some(String::new()),gc - ..Default::default()gc + app_url: Some(app_url.clone()),gc + api_url: Some("https://api.example.test".into()),gc + ..org_profile(AuthKind::ApiKey, "org_other", "other-org")gc },gc );gc - storegc + save_auth_store(&store).expect("save auth store");gc + let base = BaseArgs::default();gc + let candidates = filter_auth_store(&base, &store, None, None);gc + let filtered = filter_auth_store_for_org(&base, &mut store, candidates, Some("test-org"))gc + .awaitgc + .expect("filter by org");gc + assert_eq!(gc + filtered.profiles.into_keys().collect::>(),gc + [oauth_slot]gc + );gc }gc gc #[test]gc - fn login_and_logout_filters_compose() {gc - let store = login_filter_store();gc - let matches = |org, kind, hint| {gc - filter_auth_store(&store, org, kind, hint)gc - .profilesgc - .into_keys()gc - .collect::>()gc + fn command_auth_uses_builtin_url_defaults_when_urls_are_unset() {gc + let default_oauth = AuthProfile {gc + auth_kind: AuthKind::Oauth,gc + app_url: Some(DEFAULT_APP_URL.into()),gc + ..Default::default()gc };gc - assert_eq!(gc - matches(Some("test-org-a"), None, None),gc - ["key-a", "oauth-a"]gc - );gc - assert_eq!(matches(Some("org_test_b"), None, None), ["key-b"]);gc - assert_eq!(gc - matches(Some("test-org-a"), Some(AuthKind::ApiKey), None),gc - ["key-a"]gc - );gc - assert_eq!(matches(Some(""), None, None), ["cross"]);gc - assert!(matches(Some(""), Some(AuthKind::ApiKey), None).is_empty());gc - assert_eq!(matches(None, None, None).len(), 4);gc - assert_eq!(gc - matches(Some("test-org-a"), Some(AuthKind::Oauth), None),gc - ["oauth-a"]gc - );gc - assert_eq!(matches(None, None, Some("sk-****bbbbb")), ["key-b"]);gc + let custom_oauth = AuthProfile {gc + auth_kind: AuthKind::Oauth,gc + app_url: Some("https://www.example.test".into()),gc + ..Default::default()gc + };gc + assert!(profile_matches_urls(&BaseArgs::default(), &default_oauth));gc + assert!(!profile_matches_urls(&BaseArgs::default(), &custom_oauth));gc + assert!(profile_matches_url_filters(gc + &BaseArgs::default(),gc + &custom_oauthgc + ));gc + }gc gc - let mut picker_store = store.clone();gc - picker_store.profiles.insert(gc - "oauth-a-duplicate".into(),gc - picker_store.profiles["oauth-a"].clone(),gc + #[test]gc + fn login_filter_matches_instance_and_auth_kind() {gc + let mut store = AuthStore::default();gc + store.profiles.insert(gc + "oauth".into(),gc + AuthProfile {gc + auth_kind: AuthKind::Oauth,gc + app_url: Some("https://www.example.test".into()),gc + ..Default::default()gc + },gc + );gc + store.profiles.insert(gc + "key".into(),gc + AuthProfile {gc + app_url: Some("https://www.example.test".into()),gc + api_url: Some("https://api.example.test".into()),gc + api_key_hint: Some("sk-****abcde".into()),gc + ..org_profile(AuthKind::ApiKey, "org_test", "test-org")gc + },gc );gc - assert_eq!(saved_login_names(&picker_store, true).len(), 4);gc - assert_eq!(saved_login_names(&picker_store, false).len(), 3);gc + let base = BaseArgs {gc + app_url: Some("https://www.example.test/".into()),gc + api_url: Some("https://api.example.test/".into()),gc + ..Default::default()gc + };gc + let filtered = filter_auth_store(&base, &store, None, None);gc + assert_eq!(filtered.profiles.len(), 2);gc + let filtered = filter_auth_store(&base, &store, Some(AuthKind::ApiKey), None);gc + assert_eq!(filtered.profiles.into_keys().collect::>(), ["key"]);gc +gc + let mismatched_api = BaseArgs {gc + app_url: base.app_url.clone(),gc + api_url: Some("https://other-api.example.test".into()),gc + ..Default::default()gc + };gc + let filtered = filter_auth_store(&mismatched_api, &store, None, None);gc + assert_eq!(filtered.profiles.into_keys().collect::>(), ["oauth"]);gc }gc gc #[tokio::test]gc async fn post_login_context_preserves_only_same_org_projects() {gc let _env = TestEnv::new(None, None).await;gc - let save = |org: &str| {gc + let save = |org: &str, org_id: &str| {gc crate::config::save_global(&crate::config::Config {gc org: Some(org.into()),gc + org_id: Some(org_id.into()),gc project: Some("test-project".into()),gc project_id: Some("proj_test".into()),gc + app_url: Some("https://www.example.test".into()),gc + api_url: Some("https://api.example.test".into()),gc ..Default::default()gc })gc .unwrap();gc @@ -5525,27 +5791,20 @@ mod tests {gc crate::config::load_global().unwrap()gc };gc gc - save("old-org");gc + save("old-org", "org_old");gc let cfg = persist(Some(login_org("org_test", "test-org"))).await;gc assert_eq!((cfg.org.as_deref(), cfg.project), (Some("test-org"), None));gc gc - save("test-org");gc + save("test-org", "org_test");gc let cfg = persist(Some(login_org("org_test", "test-org"))).await;gc assert_eq!(gc (cfg.project.as_deref(), cfg.project_id.as_deref()),gc (Some("test-project"), Some("proj_test"))gc );gc -gc - save("");gc - let cfg = persist(None).await;gc - assert_eq!(gc - (cfg.org.as_deref(), cfg.project, cfg.project_id),gc - (Some(""), None, None)gc - );gc }gc gc #[tokio::test]gc - async fn resolve_post_login_project_rejects_cross_org_default_project() {gc + async fn resolve_post_login_project_requires_an_org() {gc let mut base = make_base();gc base.project = Some("demo-project".to_string());gc gc @@ -5557,11 +5816,9 @@ mod tests {gc None,gc )gc .awaitgc - .expect_err("cross-org project selection should fail");gc + .expect_err("missing org should fail");gc gc - assert!(errgc - .to_string()gc - .contains("cannot set a default project in cross-org mode"));gc + assert!(err.to_string().contains("organization is required"));gc }gc gc #[test]gc @@ -5654,6 +5911,8 @@ mod tests {gc user_name: None,gc user_email: None,gc api_key_hint: None,gc + app_url: Some(DEFAULT_APP_URL.to_string()),gc + api_url: None,gc status: "ok".into(),gc error: None,gc };gc @@ -5747,6 +6006,8 @@ mod tests {gc user_name: identity.map(str::to_string),gc user_email: identity.map(|_| "user@example.test".into()),gc api_key_hint: hint.map(str::to_string),gc + app_url: (auth == "oauth").then(|| DEFAULT_APP_URL.to_string()),gc + api_url: None,gc status: status.into(),gc error: error.map(str::to_string),gc };gc @@ -5760,7 +6021,7 @@ mod tests {gc "ok",gc None,gc ),gc - "test-org — oauth — Test User (user@example.test)",gc + "https://www.braintrust.dev — oauth — Test User (user@example.test)",gc ),gc (gc verification(gc @@ -5775,7 +6036,7 @@ mod tests {gc ),gc (gc verification("oauth", None, None, None, "expired", None),gc - "cross-org — oauth — token expired",gc + "https://www.braintrust.dev — oauth — token expired",gc ),gc (gc verification(gc @@ -5868,13 +6129,16 @@ mod tests {gc #[tokio::test]gc async fn login_read_only_cached_project_id_and_org_uses_fast_path() {gc let env = TestEnv::new(Some("proj_123"), Some("test-org")).await;gc + let mut base = base_args_for_path_probe(Some("test-org"));gc + base.app_url = Some(spawn_api_key_login_server("test-org"));gc + set_global_config_urls(base.app_url.as_deref().unwrap(), None);gc let ctx = envgc - .login_read_only_probe(Some("test-org"))gc + .login_read_only_with_base(base)gc .awaitgc .expect("fast path should succeed");gc gc assert_eq!(ctx.login.org_name().as_deref(), Some("test-org"));gc - assert_eq!(ctx.login.org_id().as_deref(), Some(""));gc + assert_eq!(ctx.login.org_id().as_deref(), Some("org_test"));gc assert_eq!(ctx.api_url, "not-a-valid-url");gc }gc gc @@ -5884,16 +6148,6 @@ mod tests {gc assert_invalid_api_url(env.login_read_only_probe(None).await);gc }gc gc - #[tokio::test]gc - async fn login_read_only_cached_project_id_but_whitespace_org_is_cross_org() {gc - let env = TestEnv::new(Some("proj_123"), None).await;gc - let err = match env.login_read_only_probe(Some(" ")).await {gc - Ok(_) => panic!("whitespace org should be canonical cross-org"),gc - Err(err) => err,gc - };gc - assert!(err.to_string().contains("concrete org"));gc - }gc -gc #[tokio::test]gc async fn login_read_only_whitespace_project_id_is_treated_as_not_cached() {gc let env = TestEnv::new(Some(" "), None).await; // has_cached_project_id => falsegc @@ -5919,9 +6173,12 @@ mod tests {gc ]);gc save_profile_secret_plaintext("acme-profile", "acme-secret").expect("save acme secret");gc save_profile_secret_plaintext("other-profile", "other-secret").expect("save other secret");gc + set_global_config_urls("https://www.acme.example", Some("https://api.acme.example"));gc + let mut base = make_base();gc + crate::config::apply_base_config(&mut base);gc gc let ctx = envgc - .login_read_only_with_base(make_base())gc + .login_read_only_with_base(base)gc .awaitgc .expect("fast path should succeed with cfg org");gc gc @@ -5939,6 +6196,7 @@ mod tests {gc base.org_name = Some("test-org".into());gc let app_url = spawn_api_key_login_server("test-org");gc base.app_url = Some(app_url.clone());gc + set_global_config_urls(&app_url, None);gc gc let ctx = envgc .login_read_only_with_base(base)gc diff --git a/src/config/mod.rs b/src/config/mod.rsgc index 62bb4a7..9d732d5 100644gc --- a/src/config/mod.rsgc +++ b/src/config/mod.rsgc @@ -8,7 +8,7 @@ use std::{gc gc use serde::{Deserialize, Serialize};gc gc -use crate::args::BaseArgs;gc +use crate::args::{BaseArgs, DEFAULT_APP_URL};gc use crate::ui::{print_command_status, CommandStatus};gc gc mod get;gc @@ -19,32 +19,73 @@ mod set;gc #[serde(default)]gc pub struct Config {gc pub org: Option,gc + pub org_id: Option,gc pub project: Option,gc pub project_id: Option,gc + pub app_url: Option,gc + pub api_url: Option,gc #[serde(flatten)]gc pub extra: serde_json::Map,gc }gc gc -pub const KNOWN_KEYS: &[&str] = &["org", "project", "project_id"];gc +pub const KNOWN_KEYS: &[&str] = &[gc + "org",gc + "org_id",gc + "project",gc + "project_id",gc + "app_url",gc + "api_url",gc +];gc gc impl Config {gc pub fn get_field(&self, key: &str) -> Option<&str> {gc match key {gc "org" => self.org.as_deref(),gc + "org_id" => self.org_id.as_deref(),gc "project" => self.project.as_deref(),gc "project_id" => self.project_id.as_deref(),gc + "app_url" => self.app_url.as_deref(),gc + "api_url" => self.api_url.as_deref(),gc _ => None,gc }gc }gc gc pub fn set_field(&mut self, key: &str, value: String) -> bool {gc match key {gc - "org" => self.org = Some(value),gc + "org" => {gc + let value = value.trim().to_string();gc + if self.org.as_ref() != Some(&value) {gc + self.org_id = None;gc + self.project = None;gc + self.project_id = None;gc + }gc + self.org = (!value.is_empty()).then_some(value);gc + }gc + "org_id" => self.org_id = self.org.as_ref().map(|_| value),gc "project" => {gc self.project = Some(value);gc self.project_id = None;gc }gc "project_id" => self.project_id = Some(value),gc + "app_url" => {gc + let value = value.trim().to_string();gc + let previous = self.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let next = if !value.is_empty() {gc + value.as_str()gc + } else {gc + DEFAULT_APP_URLgc + };gc + if !urls_equal(previous, next) {gc + self.org = None;gc + self.org_id = None;gc + self.project = None;gc + self.project_id = None;gc + }gc + self.app_url = (!value.is_empty()).then_some(value);gc + }gc + "api_url" => {gc + self.api_url = trimmed_option(Some(&value)).map(str::to_string);gc + }gc _ => return false,gc }gc truegc @@ -52,12 +93,32 @@ impl Config {gc gc pub fn unset_field(&mut self, key: &str) -> bool {gc match key {gc - "org" => self.org = None,gc + "org" => {gc + self.org = None;gc + self.org_id = None;gc + self.project = None;gc + self.project_id = None;gc + }gc + "org_id" => self.org_id = None,gc "project" => {gc self.project = None;gc self.project_id = None;gc }gc "project_id" => self.project_id = None,gc + "app_url" => {gc + if selfgc + .app_urlgc + .as_deref()gc + .is_some_and(|url| !urls_equal(url, DEFAULT_APP_URL))gc + {gc + self.org = None;gc + self.org_id = None;gc + self.project = None;gc + self.project_id = None;gc + }gc + self.app_url = None;gc + }gc + "api_url" => self.api_url = None,gc _ => return false,gc }gc truegc @@ -70,39 +131,105 @@ impl Config {gc .collect()gc }gc gc - pub(crate) fn set_context(&mut self, org: Option<&str>, project: Option<(&str, &str)>) {gc - self.org = org_option(org).map(str::to_string);gc + pub(crate) fn set_context(gc + &mut self,gc + org: (&str, &str),gc + project: Option<(&str, &str)>,gc + app_url: &str,gc + api_url: &str,gc + ) {gc + self.org = Some(org.0.trim().to_string());gc + self.org_id = Some(org.1.trim().to_string());gc (self.project, self.project_id) = projectgc .map(|(name, id)| (name.to_string(), id.to_string()))gc .unzip();gc + self.app_url = Some(app_url.to_string());gc + self.api_url = Some(api_url.to_string());gc }gc gc pub(crate) fn merge(&self, local: &Config) -> Config {gc let mut extra = self.extra.clone();gc extra.extend(local.extra.clone());gc - let global_id = self.project.as_ref().and(self.project_id.clone());gc - let (org, project, project_id) = match (&local.org, &local.project) {gc +gc + let app_url = local.app_url.clone().or_else(|| self.app_url.clone());gc + let api_url = local.api_url.clone().or_else(|| self.api_url.clone());gc + let global_app = self.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let merged_app = app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let same_instance = urls_equal(global_app, merged_app);gc + let same_org = same_instance && local.org == self.org;gc + let global_project_id = self.project.as_ref().and(self.project_id.clone());gc +gc + let (org, org_id, project, project_id) = match (&local.org, &local.project) {gc (Some(org), Some(project)) => (gc Some(org.clone()),gc + local.org_id.clone(),gc Some(project.clone()),gc local.project_id.clone(),gc ),gc - (Some(org), None) if self.org.as_ref() == Some(org) => {gc - (Some(org.clone()), self.project.clone(), global_id)gc - }gc - (Some(org), None) => (Some(org.clone()), None, None),gc - (None, Some(project)) => (None, Some(project.clone()), local.project_id.clone()),gc - (None, None) => (self.org.clone(), self.project.clone(), global_id),gc + (Some(org), None) if same_org => (gc + Some(org.clone()),gc + local.org_id.clone().or_else(|| self.org_id.clone()),gc + self.project.clone(),gc + global_project_id,gc + ),gc + (Some(org), None) => (Some(org.clone()), local.org_id.clone(), None, None),gc + (None, Some(project)) => (None, None, Some(project.clone()), local.project_id.clone()),gc + (None, None) if same_instance => (gc + self.org.clone(),gc + self.org_id.clone(),gc + self.project.clone(),gc + global_project_id,gc + ),gc + (None, None) => (None, None, None, None),gc };gc Config {gc org,gc + org_id,gc project,gc project_id,gc + app_url,gc + api_url,gc extra,gc }gc }gc }gc gc +pub(crate) fn urls_equal(left: &str, right: &str) -> bool {gc + left.trim().trim_end_matches('/') == right.trim().trim_end_matches('/')gc +}gc +gc +/// Apply config-file URL and org-ID fallbacks after clap has resolved CLI/env.gc +pub fn apply_base_config(base: &mut BaseArgs) {gc + let cfg = load().unwrap_or_default();gc + apply_config_to_base(base, &cfg);gc +}gc +gc +fn apply_config_to_base(base: &mut BaseArgs, cfg: &Config) {gc + if base.app_url.is_none() {gc + base.app_url = cfg.app_url.clone();gc + }gc +gc + if base.api_url.is_none() {gc + base.api_url = cfg.api_url.clone();gc + }gc +gc + let effective_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let config_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let same_instance = urls_equal(effective_app, config_app);gc + if base.org_name_source.is_none() {gc + if base.org_name.is_none() && same_instance {gc + base.org_name = cfg.org.clone();gc + }gc + if same_instance && base.org_name == cfg.org {gc + base.org_id = cfg.org_id.clone();gc + } else {gc + base.org_id = None;gc + }gc + } else {gc + base.org_id = None;gc + }gc +}gc +gc pub fn global_config_dir() -> Result {gc if let Some(xdg) = env::var_os("XDG_CONFIG_HOME") {gc return Ok(PathBuf::from(xdg).join("bt"));gc @@ -142,13 +269,17 @@ pub fn load_file(path: &Path) -> Config {gc gc config.extra.remove("profile");gc gc - // Fold a literal "cross-org" to the canonical "" marker on load.gc - if let Some(org) = config.org.as_deref() {gc - let normalized = normalize_org(org);gc - if normalized != org {gc - config.org = Some(normalized.to_string());gc - }gc + config.org = trimmed_option(config.org.as_deref()).map(str::to_string);gc + config.org_id = configgc + .orggc + .as_ref()gc + .and(trimmed_option(config.org_id.as_deref()).map(str::to_string));gc + if config.org.is_none() {gc + config.project = None;gc + config.project_id = None;gc }gc + config.app_url = trimmed_option(config.app_url.as_deref()).map(str::to_string);gc + config.api_url = trimmed_option(config.api_url.as_deref()).map(str::to_string);gc gc for key in config.extra.keys() {gc print_command_status(gc @@ -201,36 +332,20 @@ pub(crate) fn project_from_config_for_context(gc }gc gc fn config_matches_context(base: &BaseArgs, cfg: &Config, resolved_org: Option<&str>) -> bool {gc + let requested_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let cfg_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + if !urls_equal(requested_app, cfg_app) {gc + return false;gc + }gc +gc let cfg_org = org_option(cfg.org.as_deref());gc let requested_org = org_option(resolved_org).or_else(|| org_option(base.org_name.as_deref()));gc gc requested_org.is_none_or(|resolved| cfg_org == Some(resolved))gc }gc gc -/// Trim an org while preserving the empty cross-org marker.gc pub(crate) fn org_option(value: Option<&str>) -> Option<&str> {gc - value.map(str::trim)gc -}gc -gc -/// Human-facing spelling of the empty cross-org marker.gc -pub(crate) const CROSS_ORG_ALIAS: &str = "cross-org";gc -gc -/// Trim an org and fold the [`CROSS_ORG_ALIAS`] to the canonical `""` marker.gc -pub(crate) fn normalize_org(value: &str) -> &str {gc - let trimmed = value.trim();gc - if trimmed == CROSS_ORG_ALIAS {gc - ""gc - } else {gc - trimmedgc - }gc -}gc -gc -pub(crate) fn display_org(org: &str) -> &str {gc - if org.is_empty() {gc - CROSS_ORG_ALIASgc - } else {gc - orggc - }gc + trimmed_option(value)gc }gc gc pub(crate) fn trimmed_option(value: Option<&str>) -> Option<&str> {gc @@ -451,14 +566,14 @@ enum ConfigCommands {gc },gc /// Get a config valuegc Get {gc - /// Config key (org, project, project_id)gc + /// Config key (org, org_id, project, project_id, app_url, api_url)gc key: String,gc #[command(flatten)]gc scope: ScopeArgs,gc },gc /// Set a config valuegc Set {gc - /// Config key (org, project, project_id)gc + /// Config key (org, org_id, project, project_id, app_url, api_url)gc key: String,gc /// Value to setgc value: String,gc @@ -467,7 +582,7 @@ enum ConfigCommands {gc },gc /// Remove a config valuegc Unset {gc - /// Config key (org, project, project_id)gc + /// Config key (org, org_id, project, project_id, app_url, api_url)gc key: String,gc #[command(flatten)]gc scope: ScopeArgs,gc @@ -550,6 +665,131 @@ mod tests {gc }gc }gc gc + #[test]gc + fn merge_inherits_context_only_within_the_same_instance() {gc + let global = Config {gc + org: Some("test-org".into()),gc + org_id: Some("org_test".into()),gc + project: Some("test-project".into()),gc + project_id: Some("proj_test".into()),gc + app_url: Some("https://www.example.test".into()),gc + api_url: Some("https://api.example.test".into()),gc + ..Default::default()gc + };gc + let same_instance = Config {gc + app_url: Some("https://www.example.test/".into()),gc + api_url: Some("https://proxy.example.test".into()),gc + ..Default::default()gc + };gc + let merged = global.merge(&same_instance);gc + assert_eq!(merged.org.as_deref(), Some("test-org"));gc + assert_eq!(merged.org_id.as_deref(), Some("org_test"));gc + assert_eq!(merged.project_id.as_deref(), Some("proj_test"));gc + assert_eq!(gc + merged.api_url.as_deref(),gc + Some("https://proxy.example.test")gc + );gc +gc + let other_instance = Config {gc + app_url: Some("https://self-hosted.example.test".into()),gc + ..Default::default()gc + };gc + let merged = global.merge(&other_instance);gc + assert_eq!(merged.org, None);gc + assert_eq!(merged.org_id, None);gc + assert_eq!(merged.project, None);gc + assert_eq!(merged.app_url, other_instance.app_url);gc + }gc +gc + #[test]gc + fn config_fills_urls_and_coupled_org_id_without_overriding_cli() {gc + let cfg = Config {gc + org: Some("config-org".into()),gc + org_id: Some("org_config".into()),gc + app_url: Some("https://www.example.test".into()),gc + api_url: Some("https://api.example.test".into()),gc + ..Default::default()gc + };gc + let mut base = BaseArgs::default();gc + apply_config_to_base(&mut base, &cfg);gc + assert_eq!(base.org_name.as_deref(), Some("config-org"));gc + assert_eq!(base.org_id.as_deref(), Some("org_config"));gc + assert_eq!(base.app_url, cfg.app_url);gc + assert_eq!(base.api_url, cfg.api_url);gc +gc + let mut base = BaseArgs {gc + org_name: Some("cli-org".into()),gc + org_name_source: Some(crate::args::ArgValueSource::CommandLine),gc + app_url: Some("https://cli.example.test".into()),gc + ..Default::default()gc + };gc + apply_config_to_base(&mut base, &cfg);gc + assert_eq!(base.org_name.as_deref(), Some("cli-org"));gc + assert_eq!(base.org_id, None);gc + assert_eq!(base.app_url.as_deref(), Some("https://cli.example.test"));gc + assert_eq!(base.api_url, cfg.api_url);gc +gc + let mut same_instance = BaseArgs {gc + app_url: Some("https://www.example.test/".into()),gc + ..Default::default()gc + };gc + apply_config_to_base(&mut same_instance, &cfg);gc + assert_eq!(same_instance.api_url, cfg.api_url);gc +gc + let mut other_instance = BaseArgs {gc + app_url: Some("https://other.example.test".into()),gc + ..Default::default()gc + };gc + apply_config_to_base(&mut other_instance, &cfg);gc + assert_eq!(other_instance.org_name, None);gc + assert_eq!(other_instance.org_id, None);gc + }gc +gc + #[test]gc + fn configured_project_does_not_cross_instance_boundaries() {gc + let cfg = Config {gc + org: Some("test-org".into()),gc + project: Some("test-project".into()),gc + app_url: Some("https://www.example.test".into()),gc + ..Default::default()gc + };gc + let matching = BaseArgs {gc + org_name: Some("test-org".into()),gc + app_url: Some("https://www.example.test/".into()),gc + ..Default::default()gc + };gc + assert_eq!(gc + project_from_config_for_context(&matching, &cfg, Some("test-org")).as_deref(),gc + Some("test-project")gc + );gc +gc + let other = BaseArgs {gc + app_url: Some("https://other.example.test".into()),gc + ..matchinggc + };gc + assert_eq!(gc + project_from_config_for_context(&other, &cfg, Some("test-org")),gc + Nonegc + );gc + }gc +gc + #[test]gc + fn changing_app_url_clears_coupled_context() {gc + let mut cfg = Config {gc + org: Some("test-org".into()),gc + org_id: Some("org_test".into()),gc + project: Some("test-project".into()),gc + project_id: Some("proj_test".into()),gc + app_url: Some("https://www.example.test".into()),gc + ..Default::default()gc + };gc + assert!(cfg.set_field("app_url", "https://other.example.test".into()));gc + assert_eq!(cfg.org, None);gc + assert_eq!(cfg.org_id, None);gc + assert_eq!(cfg.project, None);gc + assert_eq!(cfg.project_id, None);gc + }gc +gc #[test]gc fn scope_labels_are_plain_text() {gc let labels = scope_labels(gc @@ -564,8 +804,8 @@ mod tests {gc fn option_helpers_handle_empty_values() {gc for (input, org, trimmed) in [gc (None, None, None),gc - (Some(""), Some(""), None),gc - (Some(" "), Some(""), None),gc + (Some(""), None, None),gc + (Some(" "), None, None),gc (Some("test-org"), Some("test-org"), Some("test-org")),gc ] {gc assert_eq!(org_option(input), org);gc @@ -573,12 +813,16 @@ mod tests {gc }gc gc let mut cfg = Config::default();gc - cfg.set_context(Some(" test-org "), Some(("test-project", "proj_test")));gc + cfg.set_context(gc + ("test-org", "org_test"),gc + Some(("test-project", "proj_test")),gc + "https://www.example.test",gc + "https://api.example.test",gc + );gc assert_eq!(cfg.org.as_deref(), Some("test-org"));gc + assert_eq!(cfg.org_id.as_deref(), Some("org_test"));gc assert_eq!(cfg.project.as_deref(), Some("test-project"));gc assert_eq!(cfg.project_id.as_deref(), Some("proj_test"));gc - cfg.set_context(Some(""), None);gc - assert_eq!((cfg.org.as_deref(), cfg.project), (Some(""), None));gc }gc gc fn base_args() -> BaseArgs {gc @@ -679,21 +923,19 @@ mod tests {gc }gc gc #[test]gc - fn load_folds_cross_org_alias_to_empty_marker() {gc + fn load_purges_obsolete_empty_org_context() {gc let tmp = TempDir::new().unwrap();gc let path = tmp.path().join("config.json");gc - // Literal "cross-org" must load identically to the "" marker.gc - for spelling in [gc - r#"{"org":"cross-org"}"#,gc - r#"{"org":" cross-org "}"#,gc - r#"{"org":""}"#,gc - ] {gc - fs::write(&path, spelling).unwrap();gc - assert_eq!(load_file(&path).org.as_deref(), Some(""), "{spelling}");gc - }gc -gc - fs::write(&path, r#"{"org":"test-org"}"#).unwrap();gc - assert_eq!(load_file(&path).org.as_deref(), Some("test-org"));gc + fs::write(gc + &path,gc + r#"{"org":"","org_id":"org_old","project":"old","project_id":"proj_old"}"#,gc + )gc + .unwrap();gc + let loaded = load_file(&path);gc + assert_eq!(loaded.org, None);gc + assert_eq!(loaded.org_id, None);gc + assert_eq!(loaded.project, None);gc + assert_eq!(loaded.project_id, None);gc }gc gc #[test]gc diff --git a/src/experiments/mod.rs b/src/experiments/mod.rsgc index 341ec84..7ee4bec 100644gc --- a/src/experiments/mod.rsgc +++ b/src/experiments/mod.rsgc @@ -219,11 +219,7 @@ fn apply_experiment_url_hints_to_base(gc }gc }gc gc - let has_org_override = basegc - .org_namegc - .as_deref()gc - .map(str::trim)gc - .is_some_and(|v| !v.is_empty());gc + let has_org_override = base.org_name_source.is_some();gc if !has_org_override {gc if let Some(org) = parsed_urlgc .orggc @@ -232,6 +228,7 @@ fn apply_experiment_url_hints_to_base(gc .filter(|v| !v.is_empty())gc {gc base.org_name = Some(org.to_string());gc + base.org_id = None;gc }gc }gc gc @@ -377,6 +374,32 @@ mod tests {gc );gc }gc gc + #[test]gc + fn comparison_url_org_overrides_config_but_not_cli() {gc + let parsed = ParsedExperimentCompareUrl {gc + org: Some("url-org".to_string()),gc + project: None,gc + base_experiment: None,gc + comparison_experiment: None,gc + };gc + let config_base = BaseArgs {gc + org_name: Some("config-org".to_string()),gc + org_id: Some("org_config".to_string()),gc + ..Default::default()gc + };gc + let updated = apply_experiment_url_hints_to_base(config_base, Some(&parsed));gc + assert_eq!(updated.org_name.as_deref(), Some("url-org"));gc + assert_eq!(updated.org_id, None);gc +gc + let cli_base = BaseArgs {gc + org_name: Some("cli-org".to_string()),gc + org_name_source: Some(crate::args::ArgValueSource::CommandLine),gc + ..Default::default()gc + };gc + let updated = apply_experiment_url_hints_to_base(cli_base, Some(&parsed));gc + assert_eq!(updated.org_name.as_deref(), Some("cli-org"));gc + }gc +gc #[test]gc fn compare_startup_url_uses_url_like_positional_arg() {gc let args = ExperimentsArgs {gc diff --git a/src/init.rs b/src/init.rsgc index 9449495..00bd6e6 100644gc --- a/src/init.rsgc +++ b/src/init.rsgc @@ -1,12 +1,10 @@gc -use anyhow::{bail, Context, Result};gc +use anyhow::{Context, Result};gc use clap::Args;gc gc use crate::{gc - args::BaseArgs,gc - auth::{self, login},gc - config,gc - http::ApiClient,gc - ui::{print_command_status, select_or_create_project, CommandStatus},gc + args::{ArgValueSource, BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL},gc + config, switch,gc + ui::{print_command_status, CommandStatus},gc };gc gc #[derive(Debug, Clone, Args)]gc @@ -33,37 +31,44 @@ pub struct InitArgs {gc pub async fn run(base: BaseArgs, args: InitArgs) -> Result<()> {gc let config_path = config::init_target(args.here, args.force)?;gc let current_cfg = config::load().unwrap_or_default();gc - let mut login_base = base.clone();gc - login_base.project = None;gc - login_base.project_source = None;gc - if login_base.org_name.is_none()gc - && !auth::select_saved_login(&mut login_base, current_cfg.org.as_deref(), false)?gc - {gc - bail!("no saved concrete-org login is available; run `bt auth login --org `");gc - }gc -gc - let ctx = login(&login_base).await?;gc - let client = ApiClient::new(&ctx)?;gc - let org = client.org_name().to_string();gc - if org.is_empty() {gc - bail!(gc - "cross-org mode has no project; `bt init` is project-scoped. Rerun with --org --project "gc - );gc - }gc -gc - let project = select_or_create_project(gc - &client,gc + let requested_org = matches!(gc + base.org_name_source,gc + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable)gc + )gc + .then(|| base.org_name.as_deref())gc + .flatten();gc + let (instance, org, project) = switch::select_context(gc + &base,gc + requested_org,gc base.project.as_deref(),gc - None,gc + ¤t_cfg,gc Some("Link to project"),gc )gc .await?;gc - // Load any existing file (only reachable via --force) so unknown passthroughgc - // keys are preserved, matching switch/config-set/post-login writers.gc + let api_url = if matches!(gc + base.api_url_source,gc + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable)gc + ) {gc + base.api_url.clone()gc + } else {gc + org.api_url.clone().or_else(|| {gc + config::urls_equal(gc + current_cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc + &instance.app_url,gc + )gc + .then(|| current_cfg.api_url.clone())gc + .flatten()gc + })gc + }gc + .unwrap_or_else(|| DEFAULT_API_URL.to_string());gc +gc + // With --force, preserve unknown passthrough keys from the old file.gc let mut cfg = config::load_file(&config_path);gc cfg.set_context(gc - Some(&org),gc + (org.name.as_str(), org.id.as_str()),gc Some((project.name.as_str(), project.id.as_str())),gc + &instance.app_url,gc + &api_url,gc );gc gc config::save_file(&config_path, &cfg).with_context(|| {gc @@ -77,16 +82,19 @@ pub async fn run(base: BaseArgs, args: InitArgs) -> Result<()> {gc let payload = serde_json::json!({gc "initialized": true,gc "status": "created",gc - "org": org,gc + "org": org.name,gc + "org_id": org.id,gc "project": project.name,gc "project_id": project.id,gc + "app_url": instance.app_url,gc + "api_url": api_url,gc "path": config_path.display().to_string(),gc });gc println!("{}", serde_json::to_string(&payload)?);gc } else {gc print_command_status(gc CommandStatus::Success,gc - &format!("Project linked to {org}/{}", project.name),gc + &format!("Project linked to {}/{}", org.name, project.name),gc );gc print_command_status(gc CommandStatus::Success,gc diff --git a/src/main.rs b/src/main.rsgc index 035cc80..0fb1106 100644gc --- a/src/main.rsgc +++ b/src/main.rsgc @@ -58,7 +58,7 @@ const HELP_TEMPLATE: &str = "\gc Coregc init Initialize .bt config directory and filesgc auth Authenticate bt with Braintrustgc - switch Switch org and project contextgc + switch Switch instance, org, and project contextgc view View logs, traces, and spansgc gc Projects & resourcesgc @@ -161,7 +161,7 @@ enum Commands {gc Sync(CLIArgs),gc /// Local utility commandsgc Util(CLIArgs),gc - /// Switch org and project contextgc + /// Switch instance, org, and project contextgc Switch(CLIArgs),gc /// Show current org and project contextgc Status(CLIArgs),gc @@ -296,6 +296,7 @@ fn try_main() -> Result<()> {gc let matches = Cli::command().get_matches_from(&argv);gc let mut cli = Cli::from_arg_matches(&matches).expect("clap matches should parse");gc apply_base_arg_sources(&matches, cli.command.base_mut());gc + config::apply_base_config(cli.command.base_mut());gc apply_base_output_defaults(&mut cli.command);gc configure_output(cli.command.base());gc apply_runtime_env_overrides(cli.command.base());gc @@ -349,6 +350,8 @@ fn apply_base_arg_sources(matches: &ArgMatches, base: &mut BaseArgs) {gc base.org_name_source = find_value_source(matches, "org_name").and_then(map_value_source);gc base.project_source = find_value_source(matches, "project").and_then(map_value_source);gc base.api_key_source = find_value_source(matches, "api_key").and_then(map_value_source);gc + base.api_url_source = find_value_source(matches, "api_url").and_then(map_value_source);gc + base.app_url_source = find_value_source(matches, "app_url").and_then(map_value_source);gc }gc gc fn apply_base_output_defaults(command: &mut Commands) {gc diff --git a/src/setup/mod.rs b/src/setup/mod.rsgc index 17937d3..bea36fa 100644gc --- a/src/setup/mod.rsgc +++ b/src/setup/mod.rsgc @@ -1179,6 +1179,10 @@ fn should_print_agent_selection_intro(gc gc fn apply_setup_config_fallbacks(base: &mut BaseArgs) {gc let cfg = config::load().unwrap_or_default();gc + let same_instance = config::urls_equal(gc + base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc + cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc + );gc gc if basegc .org_namegc @@ -1186,11 +1190,16 @@ fn apply_setup_config_fallbacks(base: &mut BaseArgs) {gc .map(str::trim)gc .is_none_or(str::is_empty)gc {gc - base.org_name = cfggc - .orggc - .as_deref()gc - .map(|value| value.trim().to_string())gc - .filter(|value| !value.is_empty());gc + if same_instance {gc + base.org_name = cfggc + .orggc + .as_deref()gc + .map(|value| value.trim().to_string())gc + .filter(|value| !value.is_empty());gc + base.org_id = base.org_name.as_ref().and(cfg.org_id.clone());gc + } else {gc + base.org_id = None;gc + }gc }gc gc if basegc @@ -1692,7 +1701,8 @@ async fn ensure_org_or_setup_browser_auth(gc auth_base.org_name = Some(org_name.to_string());gc let has_saved_auth_for_org = profilesgc .iter()gc - .any(|profile| profile.org_name.as_deref() == Some(org_name));gc + .any(|profile| profile.org_name.as_deref() == Some(org_name))gc + || auth::has_oauth_login_for_instance(&auth_base)?;gc gc if has_saved_auth_for_org {gc match auth::login(&auth_base).await {gc diff --git a/src/status.rs b/src/status.rsgc index b8a8ccb..b673899 100644gc --- a/src/status.rsgc +++ b/src/status.rsgc @@ -2,7 +2,7 @@ use anyhow::Result;gc use clap::Args;gc use serde::Serialize;gc gc -use crate::args::BaseArgs;gc +use crate::args::{BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL};gc use crate::auth;gc use crate::config;gc gc @@ -18,7 +18,11 @@ pub struct StatusArgs {}gc #[derive(Serialize)]gc struct StatusOutput {gc org: Option,gc + org_id: Option,gc project: Option,gc + project_id: Option,gc + app_url: Option,gc + api_url: Option,gc #[serde(skip_serializing_if = "Option::is_none")]gc user_name: Option,gc #[serde(skip_serializing_if = "Option::is_none")]gc @@ -69,11 +73,38 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> {gc project = config::project_from_config_for_context(&base, &merged_cfg, org.as_deref());gc }gc gc - let display_org = org.as_deref().map(config::display_org);gc + let org_id = if base.org_name_source.is_some() {gc + Nonegc + } else {gc + base.org_id.clone()gc + };gc + let configured_project =gc + config::project_from_config_for_context(&base, &merged_cfg, org.as_deref());gc + let project_id = (base.project_source.is_none()gc + && configured_project.is_some()gc + && configured_project == project)gc + .then(|| merged_cfg.project_id.clone())gc + .flatten();gc + let app_url = Some(gc + base.app_urlgc + .clone()gc + .or_else(|| merged_cfg.app_url.clone())gc + .unwrap_or_else(|| DEFAULT_APP_URL.to_string()),gc + );gc + let api_url = Some(gc + base.api_urlgc + .clone()gc + .or_else(|| merged_cfg.api_url.clone())gc + .unwrap_or_else(|| DEFAULT_API_URL.to_string()),gc + );gc if base.json {gc let output = StatusOutput {gc - org: display_org.map(str::to_string),gc + org: org.clone(),gc + org_id,gc project,gc + project_id,gc + app_url,gc + api_url,gc user_name: auth_info.as_ref().and_then(|p| p.user_name.clone()),gc user_email: auth_info.as_ref().and_then(|p| p.email.clone()),gc api_key_hint: auth_info.as_ref().and_then(|p| p.api_key_hint.clone()),gc @@ -85,8 +116,12 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> {gc }gc gc if base.verbose {gc - println!("org: {}", display_org.unwrap_or("(unset)"));gc + println!("org: {}", org.as_deref().unwrap_or("(unset)"));gc + println!("org_id: {}", org_id.as_deref().unwrap_or("(unset)"));gc println!("project: {}", project.as_deref().unwrap_or("(unset)"));gc + println!("project_id: {}", project_id.as_deref().unwrap_or("(unset)"));gc + println!("app_url: {}", app_url.as_deref().unwrap_or(DEFAULT_APP_URL));gc + println!("api_url: {}", api_url.as_deref().unwrap_or(DEFAULT_API_URL));gc if let Some(ref p) = auth_info {gc println!("auth: {}", format_auth(p));gc }gc @@ -94,26 +129,18 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> {gc println!("source: {src}");gc }gc } else {gc - // Plain one-liner. Always surface the active auth — even when no org isgc - // configured (an env-only API key, or a cross-org OAuth login) — insteadgc - // of hiding it behind --verbose.gc - let cross_org_oauth = auth_infogc - .as_ref()gc - .is_some_and(|p| p.auth_method == "oauth" && p.org_name.is_none());gc - let header = match display_org {gc - Some("cross-org") => "cross-org".to_string(),gc + let header = match org.as_deref() {gc Some(org) => match project.as_deref() {gc Some(project) => format!("{org}/{project}"),gc None => org.to_string(),gc },gc - None if cross_org_oauth => "cross-org".to_string(),gc None if auth_info.is_some() => "No default org".to_string(),gc None => "No org/project configured. Run `bt switch` to set one.".to_string(),gc };gc println!("{header}");gc match &auth_info {gc Some(p) => println!(" auth: {}", format_auth(p)),gc - None if display_org.is_some() => println!(" auth: (none)"),gc + None if org.is_some() => println!(" auth: (none)"),gc None => {}gc }gc }gc @@ -127,6 +154,10 @@ pub(crate) struct ConfigOverrides {gc env_org: Option,gc cli_project: Option,gc env_project: Option,gc + cli_app_url: Option,gc + env_app_url: Option,gc + cli_api_url: Option,gc + env_api_url: Option,gc }gc gc impl ConfigOverrides {gc @@ -143,12 +174,26 @@ impl ConfigOverrides {gc Some(ArgValueSource::EnvVariable) => (None, base.project.clone()),gc None => (None, None),gc };gc + let (cli_app_url, env_app_url) = match base.app_url_source {gc + Some(ArgValueSource::CommandLine) => (base.app_url.clone(), None),gc + Some(ArgValueSource::EnvVariable) => (None, base.app_url.clone()),gc + None => (None, None),gc + };gc + let (cli_api_url, env_api_url) = match base.api_url_source {gc + Some(ArgValueSource::CommandLine) => (base.api_url.clone(), None),gc + Some(ArgValueSource::EnvVariable) => (None, base.api_url.clone()),gc + None => (None, None),gc + };gc gc Self {gc cli_org,gc env_org,gc cli_project,gc env_project,gc + cli_app_url,gc + env_app_url,gc + cli_api_url,gc + env_api_url,gc }gc }gc }gc @@ -166,28 +211,48 @@ pub(crate) fn resolve_config(gc env_org,gc cli_project,gc env_project,gc + cli_app_url,gc + env_app_url,gc + cli_api_url,gc + env_api_url,gc } = overrides;gc - // `Some("")` is the canonical cross-org marker for both CLI and envgc - // sources, so org overrides must not filter empty strings.gc let env_project = env_project.filter(|s| !s.is_empty());gc let merged = global.merge(local);gc - let org = cli_orggc - .clone()gc - .or_else(|| env_org.clone())gc - .or_else(|| merged.org.clone());gc + let app_override = cli_app_url.as_deref().or(env_app_url.as_deref());gc + let config_app = merged.app_url.as_deref().unwrap_or(DEFAULT_APP_URL);gc + let same_instance = app_override.is_none_or(|app| config::urls_equal(app, config_app));gc + let config_org = same_instance.then(|| merged.org.clone()).flatten();gc + let config_project = same_instance.then(|| merged.project.clone()).flatten();gc + let org = cli_org.clone().or_else(|| env_org.clone()).or(config_org);gc gc let project = cli_projectgc .clone()gc .or_else(|| env_project.clone())gc - .or_else(|| merged.project.clone());gc + .or(config_project);gc gc - let source = if cli_org.is_some() || cli_project.is_some() {gc + let source = if cli_org.is_some()gc + || cli_project.is_some()gc + || cli_app_url.is_some()gc + || cli_api_url.is_some()gc + {gc Some("cli".to_string())gc - } else if env_org.is_some() || env_project.is_some() {gc + } else if env_org.is_some()gc + || env_project.is_some()gc + || env_app_url.is_some()gc + || env_api_url.is_some()gc + {gc Some("env".to_string())gc - } else if local.org.is_some() || local.project.is_some() {gc + } else if local.org.is_some()gc + || local.project.is_some()gc + || local.app_url.is_some()gc + || local.api_url.is_some()gc + {gc local_path.as_ref().map(|p| p.display().to_string())gc - } else if global.org.is_some() || global.project.is_some() {gc + } else if global.org.is_some()gc + || global.project.is_some()gc + || global.app_url.is_some()gc + || global.api_url.is_some()gc + {gc global_path.as_ref().map(|p| p.display().to_string())gc } else {gc Nonegc @@ -268,6 +333,16 @@ mod tests {gc config(None, None),gc (None, None, None),gc ),gc + (gc + "app override does not inherit another instance's context",gc + ConfigOverrides {gc + cli_app_url: s("https://other.example.test"),gc + ..Default::default()gc + },gc + both(),gc + config(None, None),gc + (None, None, Some("cli")),gc + ),gc (gc "mixed cli/local",gc ConfigOverrides {gc @@ -285,23 +360,6 @@ mod tests {gc config(None, Some("local-proj")),gc (None, Some("local-proj"), Some("/project/.bt/config.json")),gc ),gc - (gc - "local cross-org",gc - ConfigOverrides::default(),gc - both(),gc - config(Some(""), None),gc - (Some(""), None, Some("/project/.bt/config.json")),gc - ),gc - (gc - "env cross-org",gc - ConfigOverrides {gc - env_org: s(""),gc - ..Default::default()gc - },gc - both(),gc - config(None, None),gc - (Some(""), Some("global-proj"), Some("env")),gc - ),gc ];gc let local_path = Some(PathBuf::from("/project/.bt/config.json"));gc let global_path = Some(PathBuf::from("/home/.bt/config.json"));gc diff --git a/src/switch.rs b/src/switch.rsgc index b66647c..5ae7353 100644gc --- a/src/switch.rsgc +++ b/src/switch.rsgc @@ -1,8 +1,8 @@gc use anyhow::{bail, Context, Result};gc use clap::Args;gc gc -use crate::args::BaseArgs;gc -use crate::auth::{self, login};gc +use crate::args::{ArgValueSource, BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL};gc +use crate::auth::{self, login, AvailableInstance, AvailableOrg};gc use crate::config;gc use crate::http::ApiClient;gc use crate::ui::{can_prompt, print_command_status, select_or_create_project, CommandStatus};gc @@ -13,7 +13,6 @@ Examples:gc bt switchgc bt switch test-projectgc bt switch test-org/test-projectgc - bt switch --org cross-orggc ")]gc pub struct SwitchArgs {gc #[command(flatten)]gc @@ -28,87 +27,207 @@ impl SwitchArgs {gc fn resolve_target(&self, base: &BaseArgs) -> (Option, Option) {gc let (pos_org, pos_project) = match &self.target {gc None => (None, None),gc - Some(t) if t.contains('/') => {gc - let parts: Vec<&str> = t.splitn(2, '/').collect();gc - let o = (!parts[0].is_empty()).then(|| config::normalize_org(parts[0]).to_string());gc - let p = (!parts[1].is_empty()).then(|| parts[1].to_string());gc - (o, p)gc + Some(target) if target.contains('/') => {gc + let parts: Vec<&str> = target.splitn(2, '/').collect();gc + let org = (!parts[0].trim().is_empty()).then(|| parts[0].trim().to_string());gc + let project = (!parts[1].trim().is_empty()).then(|| parts[1].trim().to_string());gc + (org, project)gc }gc - Some(t) => (None, Some(t.clone())),gc + Some(target) => (None, Some(target.clone())),gc };gc gc - let org = base.org_name.clone().or(pos_org);gc - let project = base.project.clone().or(pos_project);gc -gc - (org, project)gc + (gc + base.org_namegc + .as_ref()gc + .filter(|_| {gc + matches!(gc + base.org_name_source,gc + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable)gc + )gc + })gc + .cloned()gc + .or(pos_org),gc + base.project.clone().or(pos_project),gc + )gc }gc }gc gc -pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> {gc - args.scope.preflight(can_prompt())?;gc - let current_cfg = if args.scope.global {gc - config::load_global().unwrap_or_default()gc - } else {gc - config::load().unwrap_or_default()gc - };gc - let (resolved_org, resolved_project) = args.resolve_target(&base);gc - let bare_switch = resolved_org.is_none() && resolved_project.is_none();gc +fn find_org<'a>(orgs: &'a [AvailableOrg], identifier: &str) -> Option<&'a AvailableOrg> {gc + orgs.iter()gc + .find(|org| org.id == identifier || org.name == identifier)gc + .or_else(|| {gc + let lowered = identifier.to_ascii_lowercase();gc + orgs.iter()gc + .find(|org| org.name.to_ascii_lowercase() == lowered)gc + })gc +}gc gc - let mut login_base = base.clone();gc - login_base.org_name = resolved_org.clone();gc - login_base.project = None;gc - login_base.project_source = None;gc +fn select_instance(gc + instances: &[AvailableInstance],gc + current_app_url: Option<&str>,gc +) -> Result {gc + match instances {gc + [] => bail!("no saved auth logins found; run `bt auth login` to create one"),gc + [instance] => Ok(instance.clone()),gc + _ if can_prompt() => {gc + let labels = instancesgc + .iter()gc + .map(|instance| instance.app_url.as_str())gc + .collect::>();gc + let default = current_app_urlgc + .and_then(|current| {gc + instancesgc + .iter()gc + .position(|instance| config::urls_equal(&instance.app_url, current))gc + })gc + .unwrap_or(0);gc + let idx = crate::ui::fuzzy_select("Select Braintrust instance", &labels, default)?;gc + Ok(instances[idx].clone())gc + }gc + _ => bail!(gc + "multiple Braintrust instances are available; pass --app-url or rerun interactively"gc + ),gc + }gc +}gc gc - if login_base.org_name.is_none() && !bare_switch {gc - login_base.org_name = current_cfg.org.clone();gc +fn select_org(gc + orgs: &[AvailableOrg],gc + requested: Option<&str>,gc + current: Option<&str>,gc +) -> Result {gc + if let Some(requested) = requested {gc + return find_org(orgs, requested)gc + .cloned()gc + .ok_or_else(|| anyhow::anyhow!("organization '{requested}' is not available"));gc }gc - if login_base.org_name.is_none()gc - && !auth::select_saved_login(&mut login_base, current_cfg.org.as_deref(), true)?gc - {gc - bail!("no saved auth logins found; run `bt auth login` to create one");gc + match orgs {gc + [] => bail!("no organizations are available for the selected Braintrust instance"),gc + [org] => Ok(org.clone()),gc + _ if can_prompt() => {gc + let labels = orgs.iter().map(|org| org.name.as_str()).collect::>();gc + let default = currentgc + .and_then(|current| {gc + orgs.iter()gc + .position(|org| org.id == current || org.name == current)gc + })gc + .unwrap_or(0);gc + let idx = crate::ui::fuzzy_select("Select organization", &labels, default)?;gc + Ok(orgs[idx].clone())gc + }gc + _ => bail!("organization selection requires an interactive terminal; pass --org "),gc }gc +}gc gc - if login_base.org_name.as_deref() == Some("") && resolved_project.is_some() {gc - bail!(gc - "cross-org mode cannot have a default project; rerun with --org --project "gc - );gc - }gc +pub(crate) async fn select_context(gc + base: &BaseArgs,gc + requested_org: Option<&str>,gc + requested_project: Option<&str>,gc + current_cfg: &config::Config,gc + project_prompt: Option<&str>,gc +) -> Result<(gc + AvailableInstance,gc + AvailableOrg,gc + crate::projects::api::Project,gc +)> {gc + let instances = auth::available_instances(base)?;gc + let instance = select_instance(&instances, current_cfg.app_url.as_deref())?;gc + let orgs = auth::available_orgs_for_instance(base, &instance.app_url).await?;gc + let same_current_instance = config::urls_equal(gc + current_cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc + &instance.app_url,gc + );gc + let current_org = same_current_instancegc + .then(|| current_cfg.org_id.as_deref().or(current_cfg.org.as_deref()))gc + .flatten();gc + let org = select_org(&orgs, requested_org, current_org)?;gc +gc + let explicit_api_url = matches!(gc + base.api_url_source,gc + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable)gc + )gc + .then(|| base.api_url.clone())gc + .flatten();gc + let current_api_url = same_current_instancegc + .then(|| current_cfg.api_url.clone())gc + .flatten();gc + let api_url = explicit_api_urlgc + .or_else(|| org.api_url.clone())gc + .or(current_api_url)gc + .unwrap_or_else(|| DEFAULT_API_URL.to_string());gc +gc + let mut login_base = base.clone();gc + login_base.app_url = Some(instance.app_url.clone());gc + login_base.api_url = Some(api_url);gc + login_base.org_name = Some(org.name.clone());gc + login_base.org_id = Some(org.id.clone());gc + login_base.project = None;gc + login_base.project_source = None;gc gc let ctx = login(&login_base).await?;gc let client = ApiClient::new(&ctx)?;gc - let org_name = client.org_name().to_string();gc + let current_project = (same_current_instancegc + && current_cfg.org_id.as_deref() == Some(org.id.as_str()))gc + .then_some(current_cfg.project.as_deref())gc + .flatten();gc + let project =gc + select_or_create_project(&client, requested_project, current_project, project_prompt)gc + .await?;gc gc - let project = if org_name.is_empty() {gc - Nonegc + Ok((instance, org, project))gc +}gc +gc +pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> {gc + args.scope.preflight(can_prompt())?;gc + let current_cfg = if args.scope.global {gc + config::load_global().unwrap_or_default()gc } else {gc - Some(gc - select_or_create_project(gc - &client,gc - resolved_project.as_deref(),gc - current_cfg.project.as_deref(),gc - None,gc - )gc - .await?,gc - )gc + config::load().unwrap_or_default()gc };gc + let (requested_org, requested_project) = args.resolve_target(&base);gc + let (instance, org, project) = select_context(gc + &base,gc + requested_org.as_deref(),gc + requested_project.as_deref(),gc + ¤t_cfg,gc + None,gc + )gc + .await?;gc + let api_url = if matches!(gc + base.api_url_source,gc + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable)gc + ) {gc + base.api_url.clone()gc + } else {gc + org.api_url.clone().or_else(|| {gc + config::urls_equal(gc + current_cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL),gc + &instance.app_url,gc + )gc + .then(|| current_cfg.api_url.clone())gc + .flatten()gc + })gc + }gc + .unwrap_or_else(|| DEFAULT_API_URL.to_string());gc gc - // Scope is prompted last, after org and project.gc let (path, scope) = args.scope.resolve(can_prompt(), "Save to")?;gc let mut cfg = config::load_file(&path);gc cfg.set_context(gc - Some(&org_name),gc - projectgc - .as_ref()gc - .map(|project| (project.name.as_str(), project.id.as_str())),gc + (org.name.as_str(), org.id.as_str()),gc + Some((project.name.as_str(), project.id.as_str())),gc + &instance.app_url,gc + &api_url,gc );gc config::save_file(&path, &cfg)gc .with_context(|| format!("Could not save config to {}", path.display()))?;gc gc if base.json {gc let payload = serde_json::json!({gc - "org": config::display_org(&org_name),gc - "project": project.as_ref().map(|p| p.name.clone()),gc - "project_id": project.as_ref().map(|p| p.id.clone()),gc + "org": org.name,gc + "org_id": org.id,gc + "project": project.name,gc + "project_id": project.id,gc + "app_url": instance.app_url,gc + "api_url": api_url,gc "scope": scope,gc "path": path.display().to_string(),gc });gc @@ -116,11 +235,10 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> {gc return Ok(());gc }gc gc - let display = projectgc - .as_ref()gc - .map(|project| format!("{org_name}/{}", project.name))gc - .unwrap_or_else(|| config::display_org(&org_name).to_string());gc - print_command_status(CommandStatus::Success, &format!("Switched to {display}"));gc + print_command_status(gc + CommandStatus::Success,gc + &format!("Switched to {}/{}", org.name, project.name),gc + );gc if base.verbose {gc eprintln!("Wrote to {}", path.display());gc }gc @@ -131,80 +249,27 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> {gc #[cfg(test)]gc mod tests {gc use super::*;gc +gc #[test]gc fn resolve_target_combines_positionals_and_flags() {gc - for (target, org, project, expected) in [gc - (None, None, None, (None, None)),gc - (gc - Some("test-org/test-project"),gc - None,gc - None,gc - (Some("test-org"), Some("test-project")),gc - ),gc - (gc - Some("test-project"),gc - None,gc - None,gc - (None, Some("test-project")),gc - ),gc - (gc - Some("/test-project"),gc - None,gc - None,gc - (None, Some("test-project")),gc - ),gc - (Some("test-org/"), None, None, (Some("test-org"), None)),gc - // Positional "cross-org" folds to the "" marker, matching --org.gc - (gc - Some("cross-org/test-project"),gc - None,gc - None,gc - (Some(""), Some("test-project")),gc - ),gc - (Some("cross-org/"), None, None, (Some(""), None)),gc - (None, Some("test-org"), None, (Some("test-org"), None)),gc - (gc - None,gc - None,gc - Some("test-project"),gc - (None, Some("test-project")),gc - ),gc - (gc - None,gc - Some("test-org"),gc - Some("test-project"),gc - (Some("test-org"), Some("test-project")),gc - ),gc - (gc - Some("old-org/old-project"),gc - None,gc - Some("test-project"),gc - (Some("old-org"), Some("test-project")),gc - ),gc - (gc - Some("old-project"),gc - Some("test-org"),gc - None,gc - (Some("test-org"), Some("old-project")),gc - ),gc - (gc - Some("old-org/old-project"),gc - Some("test-org"),gc - Some("test-project"),gc - (Some("test-org"), Some("test-project")),gc - ),gc - ] {gc - let args = SwitchArgs {gc - scope: config::ScopeArgs::default(),gc - target: target.map(str::to_string),gc - };gc - let base = BaseArgs {gc - org_name: org.map(str::to_string),gc - project: project.map(str::to_string),gc - ..Default::default()gc - };gc - let actual = args.resolve_target(&base);gc - assert_eq!((actual.0.as_deref(), actual.1.as_deref()), expected);gc - }gc + let args = SwitchArgs {gc + scope: config::ScopeArgs::default(),gc + target: Some("test-org/test-project".to_string()),gc + };gc + let base = BaseArgs::default();gc + let actual = args.resolve_target(&base);gc + assert_eq!(actual.0.as_deref(), Some("test-org"));gc + assert_eq!(actual.1.as_deref(), Some("test-project"));gc + }gc +gc + #[test]gc + fn find_org_matches_name_id_and_case() {gc + let orgs = vec![AvailableOrg {gc + id: "org_test".to_string(),gc + name: "test-org".to_string(),gc + api_url: None,gc + }];gc + assert!(find_org(&orgs, "org_test").is_some());gc + assert!(find_org(&orgs, "TEST-ORG").is_some());gc }gc }gc diff --git a/src/traces.rs b/src/traces.rsgc index 97a881c..c525012 100644gc --- a/src/traces.rsgc +++ b/src/traces.rsgc @@ -5728,14 +5728,11 @@ fn apply_url_hints_to_base(mut base: BaseArgs, parsed_url: Option<&ParsedTraceUrgc return base;gc };gc gc - let has_org_override = basegc - .org_namegc - .as_deref()gc - .map(str::trim)gc - .is_some_and(|v| !v.is_empty());gc + let has_org_override = base.org_name_source.is_some();gc gc if !has_org_override {gc base.org_name = Some(url_org.to_string());gc + base.org_id = None;gc }gc basegc }gc @@ -6925,19 +6922,23 @@ mod tests {gc }gc gc #[test]gc - fn apply_url_hints_infers_org_from_url() {gc - let base = base_args();gc + fn apply_url_hints_override_config_org_and_clear_its_id() {gc + let mut base = base_args();gc + base.org_name = Some("config-org".to_string());gc + base.org_id = Some("org_config".to_string());gc let parsed = parsed_url_with_org("Lovable");gc gc let updated = apply_url_hints_for_test(base, Some(&parsed));gc gc assert_eq!(updated.org_name.as_deref(), Some("Lovable"));gc + assert_eq!(updated.org_id, None);gc }gc gc #[test]gc fn apply_url_hints_preserves_explicit_org() {gc let mut base = base_args();gc base.org_name = Some("explicit-org".to_string());gc + base.org_name_source = Some(crate::args::ArgValueSource::CommandLine);gc let parsed = parsed_url_with_org("Lovable");gc gc let updated = apply_url_hints_for_test(base, Some(&parsed));gc diff --git a/tests/eval_dev_server.rs b/tests/eval_dev_server.rsgc index 227221d..3a5095f 100644gc --- a/tests/eval_dev_server.rsgc +++ b/tests/eval_dev_server.rsgc @@ -71,7 +71,7 @@ fn start_mock_auth_server() -> (u16, thread::JoinHandle<()>) {gc .expect("set mock listener blocking");gc gc let handle = thread::spawn(move || {gc - let response_body = r#"{"org_info": [{"name": "test-org"}]}"#;gc + let response_body = r#"{"org_info": [{"id": "org_test", "name": "test-org"}]}"#;gc let http_response = format!(gc "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",gc response_body.len(),gc --- README.md | 62 +- src/args.rs | 24 +- src/auth.rs | 1334 +++++++++++++++++++++++--------------- src/config/mod.rs | 370 +++++++++-- src/experiments/mod.rs | 33 +- src/init.rs | 72 +- src/main.rs | 7 +- src/setup/mod.rs | 22 +- src/status.rs | 142 ++-- src/switch.rs | 335 ++++++---- src/traces.rs | 15 +- tests/eval_dev_server.rs | 2 +- 12 files changed, 1546 insertions(+), 872 deletions(-) diff --git a/README.md b/README.md index f7bbc605..88de9252 100644 --- a/README.md +++ b/README.md @@ -139,7 +139,7 @@ Remove-Item -Recurse -Force (Join-Path $env:APPDATA "bt") -ErrorAction SilentlyC | ------------- | ------------------------------------------------------------------ | | `bt init` | Initialize `.bt/` config directory and link to a project | | `bt auth` | Authenticate with Braintrust | -| `bt switch` | Switch org and project context | +| `bt switch` | Switch instance, org, and project context | | `bt status` | Show current org and project context | | `bt datasets` | Manage datasets and dataset pipelines | | `bt eval` | Run eval files (Unix only) | @@ -312,43 +312,40 @@ Local version and pagination-key conversion helpers: ## `bt auth` -- Authenticate interactively (prompts for auth method and organization): +- Authenticate interactively: - `bt auth login` - - First prompt chooses: `OAuth (browser)` (default) or `API key`. - - OAuth can be saved for an org or in cross-org mode. API-key logins are saved per org; if a key can access multiple orgs, `bt` uses a searchable org picker. - - After login, `bt` updates the active org context immediately. If `--project` is set, it validates and saves that project's name and ID. Without `--project`, a same-org login preserves the existing project; changing orgs clears stale project context. + - First choose `OAuth (browser)` (default) or `API key`, then choose an organization and config scope. + - OAuth is stored once per Braintrust instance, identified by app URL, and can authenticate every organization available to that user in the instance. + - API-key logins remain organization-scoped; multiple keys for one organization remain distinct. + - Login writes `org`, `org_id`, `project`, `project_id`, `app_url`, and `api_url` to the selected config scope. A same-context login preserves the existing project when no project is requested. - Use `--global` or `--local` to choose the config scope. Without either flag, an existing local config causes an interactive scope picker (default: local); non-interactive runs must pass a scope. `--local` never creates `.bt`. - - `bt` confirms the resolved API URL before saving. -- Login with OAuth (browser-based, stores refresh token in secure credential store): - - `bt auth login --oauth --org myorg` - - You can pass `--no-browser` to print the URL without auto-opening. - - On remote/SSH hosts, paste the final callback URL from your local browser if localhost callback cannot be delivered. +- Login with OAuth: + - `bt auth login --oauth --org test-org` + - You can pass `--no-browser` to print the URL without opening it automatically. + - On remote/SSH hosts, paste the final callback URL if the localhost callback cannot be delivered. - List saved auth logins: - `bt auth logins` - - `bt auth logins --org test-org` (matches stored org name or ID) - - `bt auth logins --prefer-api-key` (API-key logins only) - - Both filters can be combined. + - `bt auth logins --org test-org` dynamically lists only credentials that can use that organization. + - `bt auth logins --prefer-api-key` lists API-key logins only. - Log out: - `bt auth logout` — choose from all saved logins interactively - - `bt auth logout --org test-org --oauth` — filter to the org's OAuth login - - `bt auth logout --org test-org --api-key-hint sk-****abcde` — select an API-key login - - `bt auth logout --force` (skip confirmation after selecting a login) -- Show current auth context: - - `bt status` -- Force-refresh OAuth access token for debugging: - - `bt auth refresh --org myorg` + - `bt auth logout --app-url https://www.example.test --oauth` + - `bt auth logout --org test-org --api-key-hint sk-****abcde` + - `bt auth logout --force` — skip confirmation +- Force-refresh the OAuth login for the selected instance: + - `bt auth refresh --app-url https://www.example.test` Auth resolution order for commands is: 1. Explicit `--api-key sk-...` -2. `--prefer-api-key` / `BRAINTRUST_PREFER_API_KEY` (uses `BRAINTRUST_API_KEY` first, then a stored API key for the selected org, then falls back to OAuth) -3. Stored OAuth login for the selected org (or cross-org OAuth when selected) +2. `--prefer-api-key` / `BRAINTRUST_PREFER_API_KEY` (`BRAINTRUST_API_KEY`, then a matching stored API key, then matching OAuth) +3. OAuth for the selected Braintrust instance when it can access the selected organization 4. `BRAINTRUST_API_KEY` -5. Stored API key login for the selected org +5. A matching stored API key -`--prefer-api-key` without `--org` targets the org shown by `bt status`. It cannot be used from cross-org context; pass a concrete `--org`. Once a key is selected, an invalid key or a key belonging to another requested org is an error and does not fall back to OAuth. Multiple keys in one org remain separate and are shown with key hints. +OAuth credentials are matched by app URL. API-key credentials are matched by app URL, API URL, and organization. Explicit flags override environment variables, which override local config, global config, and finally the built-in Braintrust URLs. -On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Service daemon. On macOS, it uses the `security` keychain utility. If a secure store is unavailable, `bt` falls back to a plaintext secrets file with `0600` permissions. +On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Service daemon. On macOS, it uses the `security` keychain utility. If secure storage is unavailable, `bt` falls back to a plaintext secrets file with `0600` permissions. ## `bt init` @@ -358,24 +355,23 @@ On Linux, secure storage uses `secret-tool` (libsecret) with a running Secret Se - `bt init --here` — create in the current directory without walking (including at home or `/`) - `bt init --force` — overwrite an existing discovered `.bt/config.json`; it does not change discovery -The saved context includes `org`, `project`, and `project_id`. Cross-org is excluded from the init picker because init always requires a project. +The saved context includes the Braintrust instance URLs, organization name and ID, and project name and ID. ## `bt switch` -Interactively switch org and project context: +`bt switch` changes context without selecting a credential. It chooses a Braintrust instance, discovers the organizations available through that instance's credentials, and then chooses a project. -- `bt switch` — first choose a saved OAuth-backed org or a specific API-key login, then choose a project -- `bt switch myproject` — switch the current org to a project by name -- `bt switch test-org/test-project` — switch to a specific org and project -- `bt switch --org cross-org` — select cross-org OAuth and clear project context +- `bt switch` +- `bt switch test-project` +- `bt switch test-org/test-project` - `bt switch --global` — persist to global config (`~/.config/bt/config.json`) - `bt switch --local` — update an existing local config (`.bt/config.json`); it never creates one -A sole login/project is selected automatically. Multiple API keys in one org remain separate picker entries with hints; OAuth accounts collapse to an org choice, so run `bt auth login` again to change OAuth accounts within that org. With an existing local config and no scope flag, interactive mode asks for global/local (default: local); non-interactive mode requires `--global` or `--local`. +A sole instance, organization, or project is selected automatically. With an existing local config and no scope flag, interactive mode asks for global/local (default: local); non-interactive mode requires `--global` or `--local`. ## Config context merging -Global config is `~/.config/bt/config.json`; local config is the first discovered `.bt/config.json`. Local context wins, but a local org inherits the global project only when both configs select the same org. A local project with no org never inherits a global org. Cross-org is stored as `"org": ""`, treated as a distinct org, and rendered as `cross-org` by `bt status`. Legacy `profile` fields are ignored; unknown extra keys are preserved when context is updated. +Global config is `~/.config/bt/config.json`; local config is the first discovered `.bt/config.json`. Both use the fields `org`, `org_id`, `project`, `project_id`, `app_url`, and `api_url`. Local values win. Organization IDs stay coupled to organization names, and organization/project context is inherited only within the same app URL. Legacy `profile` fields and obsolete empty cross-org contexts are ignored; unknown extra keys are preserved during updates. ## `bt status` diff --git a/src/args.rs b/src/args.rs index 1770da97..cfe4e4ca 100644 --- a/src/args.rs +++ b/src/args.rs @@ -45,6 +45,10 @@ pub struct BaseArgs { #[arg(skip)] pub org_name_source: Option, + /// Stable org ID resolved from config or internal context selection. + #[arg(skip)] + pub org_id: Option, + /// Override active project #[arg( short = 'p', @@ -65,10 +69,6 @@ pub struct BaseArgs { #[arg(skip)] pub api_key_source: Option, - /// Exact auth slot selected internally by switch/init. - #[arg(skip)] - pub pinned_auth_slot: Option, - /// Prefer API key credentials for the selected org when available. #[arg(long = "prefer-api-key", env = "BRAINTRUST_PREFER_API_KEY", global = true, value_parser = clap::builder::BoolishValueParser::new(), default_value_t = false)] pub prefer_api_key: bool, @@ -82,6 +82,9 @@ pub struct BaseArgs { )] pub api_url: Option, + #[arg(skip)] + pub api_url_source: Option, + /// Override app URL (or via BRAINTRUST_APP_URL) #[arg( long, @@ -91,6 +94,9 @@ pub struct BaseArgs { )] pub app_url: Option, + #[arg(skip)] + pub app_url_source: Option, + /// Path to a PEM-encoded CA bundle used for HTTPS requests. #[arg( long = "ca-cert", @@ -120,7 +126,11 @@ pub struct CLIArgs { } fn parse_org_name(value: &str) -> Result { - Ok(crate::config::normalize_org(value).to_string()) + let value = value.trim(); + if value.is_empty() { + return Err("organization cannot be empty".to_string()); + } + Ok(value.to_string()) } pub(crate) fn custom_api_without_app_url(api_url: Option<&str>, app_url: Option<&str>) -> bool { @@ -149,13 +159,13 @@ mod tests { #[test] fn org_normalization() { for (input, expected) in [ - ("cross-org", ""), - (" ", ""), + ("cross-org", "cross-org"), (" test-org ", "test-org"), (" org_test_123 ", "org_test_123"), ] { assert_eq!(parse_org_name(input).unwrap(), expected); } + assert!(parse_org_name(" ").is_err()); assert!(custom_api_without_app_url( Some("https://api.example.test"), None diff --git a/src/auth.rs b/src/auth.rs index 78f9ccb2..f422677b 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -59,6 +59,7 @@ pub struct ResolvedAuth { pub api_url: Option, pub app_url: Option, pub org_name: Option, + pub org_id: Option, pub is_oauth: bool, slot_key: Option, } @@ -72,6 +73,11 @@ pub struct ProfileInfo { pub api_key_hint: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AvailableInstance { + pub app_url: String, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct AvailableOrg { pub id: String, @@ -82,6 +88,7 @@ pub struct AvailableOrg { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RecoverableAuthErrorKind { OauthRefreshToken, + OauthOrgAccess, StoredCredential, } @@ -103,6 +110,14 @@ fn recoverable_auth_error(kind: RecoverableAuthErrorKind, message: String) -> an anyhow::Error::new(RecoverableAuthError { kind, message }) } +fn is_oauth_org_access_error(err: &anyhow::Error) -> bool { + err.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|err| err.kind == RecoverableAuthErrorKind::OauthOrgAccess) + }) +} + pub fn is_missing_credential_error(err: &anyhow::Error) -> bool { err.chain().any(|source| { source @@ -126,6 +141,14 @@ pub fn list_profiles() -> Result> { .collect()) } +pub(crate) fn has_oauth_login_for_instance(base: &BaseArgs) -> Result { + let store = load_auth_store()?; + Ok(store + .profiles + .values() + .any(|profile| profile.auth_kind == AuthKind::Oauth && profile_matches_urls(base, profile))) +} + pub async fn list_available_orgs(base: &BaseArgs) -> Result> { let resolved = resolve_auth(base).await?; let app_url = resolved @@ -163,6 +186,133 @@ async fn available_orgs(api_key: &str, app_url: &str) -> Result Result> { + let store = load_auth_store()?; + let constrain_app = matches!( + base.app_url_source, + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable) + ); + let requested_app = constrain_app + .then_some(base.app_url.as_deref()) + .flatten() + .map(canonical_url); + let mut apps = store + .profiles + .values() + .map(profile_app_url) + .filter(|app| requested_app.is_none_or(|requested| canonical_url(app) == requested)) + .map(|app| canonical_url(app).to_string()) + .collect::>(); + + if base + .api_key + .as_deref() + .is_some_and(|key| !key.trim().is_empty()) + { + let app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + if requested_app.is_none_or(|requested| canonical_url(app) == requested) { + apps.insert(canonical_url(app).to_string()); + } + } + + if apps.is_empty() && constrain_app { + bail!( + "no credentials found for app URL '{}'; run `bt auth login --app-url {}`", + base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + shell_quote_arg(base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL)) + ); + } + Ok(apps + .into_iter() + .map(|app_url| AvailableInstance { app_url }) + .collect()) +} + +pub(crate) async fn available_orgs_for_instance( + base: &BaseArgs, + app_url: &str, +) -> Result> { + let mut store = load_auth_store()?; + let explicit_api = matches!( + base.api_url_source, + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable) + ) + .then(|| base.api_url.as_deref()) + .flatten(); + let mut orgs = BTreeMap::::new(); + let matching = store + .profiles + .iter() + .filter(|(_, profile)| canonical_url(profile_app_url(profile)) == canonical_url(app_url)) + .filter(|(_, profile)| { + profile.auth_kind == AuthKind::Oauth + || explicit_api + .is_none_or(|url| canonical_url(url) == canonical_url(profile_api_url(profile))) + }) + .map(|(slot, profile)| (slot.clone(), profile.clone())) + .collect::>(); + + for (slot, profile) in matching { + match profile.auth_kind { + AuthKind::Oauth => { + let mut oauth_base = base.clone(); + oauth_base.app_url = Some(app_url.to_string()); + if oauth_base.api_url_source.is_none() { + oauth_base.api_url = profile.api_url.clone(); + } + let token = load_oauth_access_token(&oauth_base, &mut store, &slot).await?; + for org in fetch_login_orgs(&token, app_url).await? { + orgs.insert( + org.id.clone(), + AvailableOrg { + id: org.id, + name: org.name, + api_url: org.api_url, + }, + ); + } + } + AuthKind::ApiKey => { + if let (Some(id), Some(name)) = (profile.org_id, profile.org_name) { + orgs.entry(id.clone()).or_insert(AvailableOrg { + id, + name, + api_url: profile.api_url, + }); + } + } + } + } + + if let Some(api_key) = base.api_key.as_deref().filter(|key| !key.trim().is_empty()) { + let base_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + if canonical_url(base_app) == canonical_url(app_url) { + for org in fetch_login_orgs(api_key, app_url).await? { + orgs.insert( + org.id.clone(), + AvailableOrg { + id: org.id, + name: org.name, + api_url: org.api_url, + }, + ); + } + } + } + + let mut orgs = orgs.into_values().collect::>(); + orgs.sort_by(|a, b| { + a.name + .to_ascii_lowercase() + .cmp(&b.name.to_ascii_lowercase()) + .then_with(|| a.name.cmp(&b.name)) + }); + if orgs.is_empty() { + bail!("no organizations are available for Braintrust instance '{app_url}'"); + } + Ok(orgs) +} + #[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)] struct AuthStore { #[serde(default)] @@ -279,7 +429,7 @@ pub struct AuthArgs { enum AuthCommand { /// Authenticate with Braintrust (OAuth or API key) Login(AuthLoginArgs), - /// Force-refresh OAuth access token for the selected org + /// Force-refresh the OAuth access token for the selected instance Refresh, /// List saved auth logins and check connection status Logins(AuthLoginsArgs), @@ -332,7 +482,7 @@ pub async fn run(base: BaseArgs, args: AuthArgs) -> Result<()> { } AuthCommand::Refresh => run_login_refresh(&base).await, AuthCommand::Logins(logins_args) => run_logins(&base, logins_args).await, - AuthCommand::Logout(logout_args) => run_login_logout(base, logout_args), + AuthCommand::Logout(logout_args) => run_login_logout(base, logout_args).await, } } @@ -371,7 +521,7 @@ pub async fn fast_login(base: &BaseArgs) -> Result { let login = LoginState::new(); login.set( api_key, - String::new(), + auth.org_id.clone().unwrap_or_default(), org_name, api_url.clone(), app_url.clone(), @@ -428,7 +578,7 @@ pub async fn login(base: &BaseArgs) -> Result { let login = LoginState::new(); login.set( api_key.clone(), - String::new(), + auth.org_id.clone().unwrap_or_default(), org_name, auth.api_url .clone() @@ -752,8 +902,20 @@ fn config_auth_context(base: &BaseArgs) -> Option { config_auth_context_from_config(base, &cfg) } +fn configured_org_for_app_url(app_url: &str) -> Option { + let cfg = crate::config::load().ok()?; + let config_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + crate::config::urls_equal(app_url, config_app) + .then_some(cfg.org) + .flatten() +} + fn config_auth_context_from_config(base: &BaseArgs, cfg: &crate::config::Config) -> Option { - if crate::config::org_option(base.org_name.as_deref()).is_none() { + let base_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let config_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + if crate::config::org_option(base.org_name.as_deref()).is_none() + && crate::config::urls_equal(base_app, config_app) + { crate::config::org_option(cfg.org.as_deref()).map(str::to_string) } else { None @@ -831,11 +993,6 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { let can_prompt = ui::can_prompt(); let effective_org = effective_org_name(base, &cfg_org); - reject_cross_org_api_key_preference(base.prefer_api_key, effective_org, &store)?; - - if let Some(slot) = base.pinned_auth_slot.clone() { - return resolve_saved_auth_slot(base, &mut store, &None, &slot).await; - } let source = resolve_auth_source( base.prefer_api_key, @@ -849,13 +1006,49 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { AuthSource::CliApiKey(api_key) | AuthSource::EnvApiKey(api_key) => { resolve_ad_hoc_api_key_auth(base, &cfg_org, api_key).await } - AuthSource::Oauth(slot) | AuthSource::ApiKey(slot) => { + AuthSource::Oauth(slot) => { + match resolve_saved_auth_slot(base, &mut store, &cfg_org, &slot).await { + Ok(auth) => Ok(auth), + Err(err) if is_oauth_org_access_error(&err) && !base.prefer_api_key => { + if let Some(api_key) = resolve_env_api_key(base) { + return resolve_ad_hoc_api_key_auth(base, &cfg_org, api_key).await; + } + if let Some(api_key_slot) = select_profile_for_auth( + base, + &store, + &cfg_org, + AuthKind::ApiKey, + can_prompt, + )? { + return resolve_saved_auth_slot(base, &mut store, &cfg_org, &api_key_slot) + .await; + } + Err(err) + } + Err(err) => Err(err), + } + } + AuthSource::ApiKey(slot) => { resolve_saved_auth_slot(base, &mut store, &cfg_org, &slot).await } AuthSource::None => { if base.prefer_api_key { bail!("--prefer-api-key requires an API key or OAuth login for the selected org"); } + if !store.profiles.is_empty() + && !store + .profiles + .values() + .any(|profile| profile_matches_urls(base, profile)) + { + let app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let api = base.api_url.as_deref().unwrap_or(DEFAULT_API_URL); + bail!( + "no credentials match app URL '{}' and API URL '{}'; run `bt auth login` with these URLs", + app, + api + ); + } if effective_org.is_none() { if let Some(err) = missing_org_for_stored_logins_error(&store) { return Err(err); @@ -866,6 +1059,7 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { api_url: base.api_url.clone(), app_url: base.app_url.clone(), org_name: effective_org.map(str::to_string), + org_id: base.org_id.clone(), is_oauth: false, slot_key: None, }) @@ -884,6 +1078,7 @@ async fn resolve_ad_hoc_api_key_auth( } let mut resolved_org = requested_org.map(str::to_string); + let mut resolved_org_id = base.org_id.clone(); let mut resolved_api_url = base.api_url.clone(); if let Some(requested_org) = requested_org { if crate::args::custom_api_without_app_url(base.api_url.as_deref(), base.app_url.as_deref()) @@ -905,6 +1100,7 @@ async fn resolve_ad_hoc_api_key_auth( ) })?; resolved_org = Some(selected_org.name.clone()); + resolved_org_id = Some(selected_org.id.clone()); resolved_api_url = resolved_api_url.or_else(|| selected_org.api_url.clone()); } @@ -913,6 +1109,7 @@ async fn resolve_ad_hoc_api_key_auth( api_url: resolved_api_url, app_url: base.app_url.clone(), org_name: resolved_org, + org_id: resolved_org_id, is_oauth: false, slot_key: None, }) @@ -980,6 +1177,7 @@ fn resolve_api_key_profile_auth( org_name: effective_org_name(base, cfg_org) .map(str::to_string) .or_else(|| profile.org_name.clone()), + org_id: profile.org_id.clone().or_else(|| base.org_id.clone()), is_oauth: false, slot_key: Some(profile_name.to_string()), }; @@ -1067,57 +1265,6 @@ fn maybe_rekey_api_key_profile_after_secret_load( Ok(()) } -async fn reconcile_oauth_slot_from_access_token( - store: &mut AuthStore, - slot_key: &str, - access_token: &str, - app_url: &str, -) -> Result<()> { - let Some(mut profile) = store.profiles.get(slot_key).cloned() else { - return Ok(()); - }; - if profile.auth_kind != AuthKind::Oauth || profile.org_id.is_some() { - return Ok(()); - } - - let Some(org_name) = profile - .org_name - .as_deref() - .map(str::trim) - .filter(|org| !org.is_empty()) - else { - profile.org_id = Some(String::new()); - if replace_with_canonical_auth_profile(store, slot_key, profile) { - save_auth_store(store)?; - } - return Ok(()); - }; - - // A legacy auth entry only cached the org name. Resolve the stable ID from - // the newly refreshed token; a failed best-effort lookup must not turn a - // successful token refresh into a failed command. - let Ok(orgs) = fetch_login_orgs(access_token, app_url).await else { - return Ok(()); - }; - let Some(org) = find_login_org(&orgs, org_name) else { - return Ok(()); - }; - - profile.org_id = Some(org.id.clone()); - profile.org_name = Some(org.name.clone()); - let identity = decode_jwt_identity(access_token); - if identity.email.is_some() { - profile.email = identity.email; - } - if identity.name.is_some() { - profile.user_name = identity.name; - } - if replace_with_canonical_auth_profile(store, slot_key, profile) { - save_auth_store(store)?; - } - Ok(()) -} - fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Result<()> { let Some(slot_key) = auth.slot_key.as_deref() else { return Ok(()); @@ -1127,40 +1274,25 @@ fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Resu return Ok(()); }; + if profile.auth_kind == AuthKind::Oauth { + return Ok(()); + } let login_org_id = login.org_id().unwrap_or_default(); - let is_cross_org = profile.auth_kind == AuthKind::Oauth - && auth - .org_name - .as_deref() - .is_none_or(|org| org.trim().is_empty()); - if login_org_id.trim().is_empty() && !is_cross_org { - // The SDK's OAuth compatibility path does not always return org - // metadata. `bt auth logins` performs the same reconciliation after - // its explicit credential verification request. + if login_org_id.trim().is_empty() { return Ok(()); } profile.org_id = Some(login_org_id); - profile.org_name = if is_cross_org { - None - } else { - login - .org_name() - .filter(|org| !org.trim().is_empty()) - .or_else(|| auth.org_name.clone()) + profile.org_name = login + .org_name() + .filter(|org| !org.trim().is_empty()) + .or_else(|| auth.org_name.clone()); + let Some(api_key) = auth.api_key.as_deref() else { + return Ok(()); }; - match profile.auth_kind { - AuthKind::ApiKey => { - let Some(api_key) = auth.api_key.as_deref() else { - return Ok(()); - }; - profile.api_key_hash = Some(api_key_hash(api_key)); - if profile.api_key_hint.is_none() { - profile.api_key_hint = Some(obscure_api_key(api_key)); - } - } - AuthKind::Oauth if profile.email.as_deref().is_none_or(str::is_empty) => return Ok(()), - AuthKind::Oauth => {} + profile.api_key_hash = Some(api_key_hash(api_key)); + if profile.api_key_hint.is_none() { + profile.api_key_hint = Some(obscure_api_key(api_key)); } if replace_with_canonical_auth_profile(&mut store, slot_key, profile) { @@ -1169,43 +1301,22 @@ fn reconcile_resolved_auth_slot(auth: &ResolvedAuth, login: &LoginState) -> Resu Ok(()) } -async fn resolve_oauth_profile_auth( +async fn load_oauth_access_token( base: &BaseArgs, store: &mut AuthStore, - cfg_org: &Option, profile_name: &str, -) -> Result { +) -> Result { let profile = store .profiles .get(profile_name) .cloned() .ok_or_else(|| anyhow::anyhow!("saved OAuth login not found; run `bt auth logins`"))?; - let api_url = base - .api_url - .clone() - .or_else(|| profile.api_url.clone()) - .unwrap_or_else(|| DEFAULT_API_URL.to_string()); - let app_url = base.app_url.clone().or_else(|| profile.app_url.clone()); - let org_name = effective_org_name(base, cfg_org) - .map(str::to_string) - .or_else(|| profile.org_name.clone()); - - let mut auth = ResolvedAuth { - api_key: None, - api_url: Some(api_url.clone()), - app_url, - org_name, - is_oauth: true, - slot_key: Some(profile_name.to_string()), - }; - - if let Some(cached_access_token) = load_valid_cached_oauth_access_token( + if let Some(cached) = load_valid_cached_oauth_access_token( profile_name, &profile, profile.oauth_access_expires_at, )? { - auth.api_key = Some(cached_access_token); - return Ok(auth); + return Ok(cached); } let refresh_token = load_profile_oauth_refresh_token_for_profile(profile_name, &profile)? @@ -1219,7 +1330,11 @@ async fn resolve_oauth_profile_auth( ), ) })?; - let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, &profile).await?; + let api_url = base + .api_url + .as_deref() + .unwrap_or_else(|| profile_api_url(&profile)); + let refreshed = refresh_oauth_access_token(api_url, &refresh_token, &profile).await?; save_profile_oauth_access_token(profile_name, &refreshed.access_token)?; let mut refresh_rotated = false; if let Some(next_refresh_token) = refreshed.refresh_token.as_ref() { @@ -1239,14 +1354,70 @@ async fn resolve_oauth_profile_auth( } } save_auth_store(store)?; - reconcile_oauth_slot_from_access_token( - store, - profile_name, - &refreshed.access_token, - auth.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), - ) - .await?; - auth.api_key = Some(refreshed.access_token); + Ok(refreshed.access_token) +} + +async fn resolve_oauth_profile_auth( + base: &BaseArgs, + store: &mut AuthStore, + cfg_org: &Option, + profile_name: &str, +) -> Result { + let profile = store + .profiles + .get(profile_name) + .cloned() + .ok_or_else(|| anyhow::anyhow!("saved OAuth login not found; run `bt auth logins`"))?; + let access_token = load_oauth_access_token(base, store, profile_name).await?; + let auth = ResolvedAuth { + api_key: Some(access_token), + api_url: Some( + base.api_url + .clone() + .or_else(|| profile.api_url.clone()) + .unwrap_or_else(|| DEFAULT_API_URL.to_string()), + ), + app_url: Some( + base.app_url + .clone() + .or_else(|| profile.app_url.clone()) + .unwrap_or_else(|| DEFAULT_APP_URL.to_string()), + ), + org_name: effective_org_name(base, cfg_org).map(str::to_string), + org_id: base.org_id.clone(), + is_oauth: true, + slot_key: Some(profile_name.to_string()), + }; + resolve_oauth_org_context(auth).await +} + +async fn resolve_oauth_org_context(mut auth: ResolvedAuth) -> Result { + let requested_org = auth.org_name.as_deref().ok_or_else(|| { + recoverable_auth_error( + RecoverableAuthErrorKind::OauthOrgAccess, + "an active organization is required; run `bt switch` or pass --org ".to_string(), + ) + })?; + let credential = auth + .api_key + .as_deref() + .context("OAuth access token is missing")?; + let app_url = auth.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let orgs = fetch_login_orgs(credential, app_url).await?; + let selected = find_login_org(&orgs, requested_org).ok_or_else(|| { + recoverable_auth_error( + RecoverableAuthErrorKind::OauthOrgAccess, + format!( + "OAuth login for '{}' cannot access organization '{requested_org}'", + canonical_url(app_url) + ), + ) + })?; + auth.org_name = Some(selected.name.clone()); + auth.org_id = Some(selected.id.clone()); + if auth.api_url.is_none() { + auth.api_url = selected.api_url.clone(); + } Ok(auth) } @@ -1281,6 +1452,44 @@ pub async fn resolved_runner_env(base: &BaseArgs) -> Result &str { + url.trim().trim_end_matches('/') +} + +fn profile_app_url(profile: &AuthProfile) -> &str { + profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL) +} + +fn profile_api_url(profile: &AuthProfile) -> &str { + profile.api_url.as_deref().unwrap_or(DEFAULT_API_URL) +} + +fn profile_matches_urls(base: &BaseArgs, profile: &AuthProfile) -> bool { + let app_url = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + if canonical_url(app_url) != canonical_url(profile_app_url(profile)) { + return false; + } + profile.auth_kind == AuthKind::Oauth + || canonical_url(base.api_url.as_deref().unwrap_or(DEFAULT_API_URL)) + == canonical_url(profile_api_url(profile)) +} + +/// Match only URL filters the caller actually supplied. Listing and logout use +/// this variant so an absent filter means "all instances", while command auth +/// uses [`profile_matches_urls`] and therefore honors the built-in URL defaults. +fn profile_matches_url_filters(base: &BaseArgs, profile: &AuthProfile) -> bool { + let app_matches = base + .app_url + .as_deref() + .is_none_or(|url| canonical_url(url) == canonical_url(profile_app_url(profile))); + app_matches + && (profile.auth_kind == AuthKind::Oauth + || base + .api_url + .as_deref() + .is_none_or(|url| canonical_url(url) == canonical_url(profile_api_url(profile)))) +} + fn profile_matches_org_identifier(profile: &AuthProfile, org: &str) -> bool { profile.org_id.as_deref() == Some(org) || profile.org_name.as_deref() == Some(org) } @@ -1298,13 +1507,13 @@ fn profile_org(profile: &AuthProfile) -> &str { } fn profile_org_label(profile: &AuthProfile) -> String { - config::display_org(profile_org(profile)).to_string() + profile_org(profile).to_string() } fn oauth_reauth_command(profile: &AuthProfile) -> String { format!( - "bt auth login --oauth --org {}", - shell_quote_arg(config::display_org(profile_org(profile))) + "bt auth login --oauth --app-url {}", + shell_quote_arg(profile_app_url(profile)) ) } @@ -1333,32 +1542,18 @@ fn profile_identity_label(profile: &AuthProfile) -> Option { } fn auth_slot_label(profile: &AuthProfile) -> String { - let mut parts = vec![profile_org_label(profile)]; - parts.push(auth_kind_label(profile.auth_kind).to_string()); + let mut parts = match profile.auth_kind { + AuthKind::Oauth => vec![profile_app_url(profile).to_string(), "oauth".to_string()], + AuthKind::ApiKey => vec![profile_org_label(profile), "api_key".to_string()], + }; if let Some(identity) = profile_identity_label(profile) { parts.push(identity); } parts.join(" — ") } -fn is_cross_org_oauth_profile(profile: &AuthProfile) -> bool { - profile.auth_kind == AuthKind::Oauth && profile_org(profile).is_empty() -} - -fn reject_cross_org_api_key_preference( - prefer_api_key: bool, - org: Option<&str>, - store: &AuthStore, -) -> Result<()> { - let cross_org = org == Some("") - || (org.is_none() && store.profiles.values().any(is_cross_org_oauth_profile)); - if prefer_api_key && cross_org { - bail!("--prefer-api-key cannot be used from cross-org context; rerun with --org "); - } - Ok(()) -} - fn auth_profile_names_by_kind<'a>( + base: &BaseArgs, store: &'a AuthStore, org: Option<&str>, kind: AuthKind, @@ -1367,10 +1562,10 @@ fn auth_profile_names_by_kind<'a>( .profiles .iter() .filter(|(_, profile)| profile.auth_kind == kind) - .filter(|(_, profile)| match org { - Some(org) => profile_matches_org_identifier(profile, org), - None if kind == AuthKind::Oauth => is_cross_org_oauth_profile(profile), - None => false, + .filter(|(_, profile)| profile_matches_urls(base, profile)) + .filter(|(_, profile)| { + kind == AuthKind::Oauth + || org.is_some_and(|org| profile_matches_org_identifier(profile, org)) }) .map(|(name, _)| name.as_str()) .collect() @@ -1403,9 +1598,7 @@ fn ad_hoc_api_key_profile(org: Option<&str>, api_key: &str) -> ProfileInfo { pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Result> { let store = load_auth_store().unwrap_or_default(); - reject_cross_org_api_key_preference(base.prefer_api_key, org, &store)?; - - let select = |kind| match auth_profile_names_by_kind(&store, org, kind).as_slice() { + let select = |kind| match auth_profile_names_by_kind(base, &store, org, kind).as_slice() { [] => Ok(None), [name] => Ok(Some((*name).to_string())), _ => bail!("multiple {kind:?} logins"), @@ -1434,13 +1627,7 @@ pub(crate) fn active_auth_info(base: &BaseArgs, org: Option<&str>) -> Result Option { - let candidates = store - .profiles - .iter() - .filter(|(_, profile)| { - !is_cross_org_oauth_profile(profile) && !profile_org(profile).is_empty() - }) - .collect::>(); + let candidates = store.profiles.iter().collect::>(); if candidates.is_empty() { return None; } @@ -1483,16 +1670,7 @@ fn select_profile_from_store( current: Option<&str>, store: &AuthStore, ) -> Result { - // Surface cross-org OAuth first so it is a predictable top-of-list choice - // rather than falling wherever its slot key happens to sort. The stable - // sort preserves the existing order of the remaining entries. - let mut names: Vec<&str> = names.to_vec(); - names.sort_by_key(|name| { - !store - .profiles - .get(*name) - .is_some_and(is_cross_org_oauth_profile) - }); + let names: Vec<&str> = names.to_vec(); let labels: Vec = names .iter() .map(|name| profile_label_from_store(name, store)) @@ -1513,53 +1691,6 @@ fn select_profile_from_store( Ok(names[idx].to_string()) } -fn saved_login_names(store: &AuthStore, include_cross_org: bool) -> Vec<&str> { - let mut oauth_orgs = BTreeSet::new(); - store - .profiles - .iter() - .filter(|(_, profile)| { - profile.auth_kind == AuthKind::ApiKey - || ((include_cross_org || !is_cross_org_oauth_profile(profile)) - && oauth_orgs.insert( - profile - .org_id - .as_deref() - .filter(|id| !id.is_empty()) - .unwrap_or_else(|| profile_org(profile)) - .to_ascii_lowercase(), - )) - }) - .map(|(name, _)| name.as_str()) - .collect() -} - -pub(crate) fn select_saved_login( - base: &mut BaseArgs, - current_org: Option<&str>, - include_cross_org: bool, -) -> Result { - let store = load_auth_store()?; - let names = saved_login_names(&store, include_cross_org); - let selected = match names.as_slice() { - [] => return Ok(false), - [name] => (*name).to_string(), - _ if ui::can_prompt() => { - select_profile_from_store("Select login", &names, current_org, &store)? - } - _ => { - bail!("multiple saved logins match; pass --org , or rerun interactively to choose") - } - }; - let profile = &store.profiles[&selected]; - if profile.auth_kind == AuthKind::ApiKey { - base.pinned_auth_slot = Some(selected); - } else { - base.org_name = Some(profile_org(profile).to_string()); - } - Ok(true) -} - fn candidate_identities<'a>(names: &[&'a str], store: &'a AuthStore) -> Vec { names .iter() @@ -1583,12 +1714,18 @@ fn select_profile_for_auth( can_prompt: bool, ) -> Result> { let org = effective_org_name(base, cfg_org); - let candidates = auth_profile_names_by_kind(store, org, kind); + let candidates = auth_profile_names_by_kind(base, store, org, kind); let label = match kind { AuthKind::Oauth => "OAuth login", AuthKind::ApiKey => "API key", }; - select_auth_profile_candidate(label, org, &candidates, store, can_prompt) + select_auth_profile_candidate( + label, + org, + &candidates, + store, + can_prompt && kind == AuthKind::ApiKey, + ) } fn select_auth_profile_candidate( @@ -1609,13 +1746,18 @@ fn select_auth_profile_candidate( } _ => { let identities = candidate_identities(candidates, store).join(", "); + if kind_label == "OAuth login" { + bail!( + "multiple Braintrust OAuth instances are available: {identities}. Run `bt switch` or pass --app-url ." + ); + } if let Some(org) = org { bail!( "multiple {kind_label} logins for org '{org}': {identities}. Rerun interactively or remove one with `bt auth logout`." ); } bail!( - "multiple cross-org {kind_label} logins available: {identities}. Rerun interactively or remove one with `bt auth logout`." + "multiple {kind_label} logins available: {identities}. Pass --app-url , rerun interactively, or remove one with `bt auth logout`." ); } } @@ -1625,10 +1767,12 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { if args.oauth { return run_login_oauth(base, args).await; } - if base.org_name.as_deref() == Some("") { - bail!( - "API-key login requires a concrete org; cross-org API keys do not exist. Use --oauth, or rerun with --org " - ); + if base + .org_name + .as_deref() + .is_some_and(|org| org.trim().is_empty()) + { + bail!("API-key login requires a non-empty organization"); } let has_explicit_api_key = base.api_key.as_ref().is_some_and(|k| !k.trim().is_empty()); @@ -1662,7 +1806,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { if requested_org_resolution == RequestedOrgResolution::SwitchToOauth { return run_login_oauth(base, args).await; } - let configured_org = config::load().ok().and_then(|cfg| cfg.org); + let configured_org = configured_org_for_app_url(&login_app_url); let selected_org = select_login_org( login_orgs.clone(), match requested_org_resolution { @@ -1675,7 +1819,6 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { configured_org.as_deref(), interactive, base.verbose, - false, explicitly_quiet(base), )?; let selected_org = selected_org.ok_or_else(|| { @@ -1691,7 +1834,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { commit_api_key_profile( &api_key, selected_api_url.clone(), - base.app_url.clone(), + Some(login_app_url.clone()), selected_org.id.clone(), selected_org.name.clone(), )?; @@ -1784,48 +1927,46 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { exchange_oauth_authorization_code(&api_url, &redirect_uri, &auth_code, pkce_verifier) .await?; let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; - let configured_org = config::load().ok().and_then(|cfg| cfg.org); + let configured_org = configured_org_for_app_url(&app_url); let selected_org = select_login_org( login_orgs.clone(), base.org_name.as_deref(), configured_org.as_deref(), ui::can_prompt(), base.verbose, - true, explicitly_quiet(base), )?; + let selected_org = selected_org.ok_or_else(|| { + anyhow::anyhow!( + "OAuth login requires an organization; pass --org or rerun interactively" + ) + })?; let selected_api_url = resolve_profile_api_url( base.api_url.clone(), - selected_org.as_ref(), + Some(&selected_org), &login_orgs, ui::can_prompt(), )?; - commit_oauth_profile( - &oauth_tokens, - selected_api_url.clone(), - app_url.clone(), - selected_org.as_ref(), - )?; + commit_oauth_profile(&oauth_tokens, api_url.clone(), app_url.clone())?; let context_update = persist_post_login_context( base, &oauth_tokens.access_token, &selected_api_url, &app_url, - selected_org.as_ref(), + Some(&selected_org), &args.scope, ) .await .context("login succeeded, but failed to update active context")?; - let human = format_login_success(selected_org.as_ref(), &selected_api_url); + let human = format_login_success(Some(&selected_org), &selected_api_url); emit_result( base.json, serde_json::json!({ "auth": "oauth", - "org": selected_org.as_ref().map(|org| org.name.clone()), - "org_id": selected_org.as_ref().map(|org| org.id.clone()), - "cross_org": selected_org.is_none(), + "org": selected_org.name, + "org_id": selected_org.id, "api_url": selected_api_url, "app_url": app_url, "status": "ok", @@ -1881,7 +2022,6 @@ fn commit_oauth_profile( tokens: &OAuthTokenResponse, api_url: String, app_url: String, - selected_org: Option<&LoginOrgInfo>, ) -> Result<()> { let refresh_token = tokens.refresh_token.as_ref().ok_or_else(|| { anyhow::anyhow!( @@ -1891,7 +2031,7 @@ fn commit_oauth_profile( let oauth_access_expires_at = determine_oauth_access_expiry_epoch(tokens); let jwt_id = decode_jwt_identity(&tokens.access_token); - let email = jwt_id + let _email = jwt_id .email .clone() .filter(|email| !email.trim().is_empty()) @@ -1900,25 +2040,25 @@ fn commit_oauth_profile( "oauth token did not include an email; cannot create persistent oauth login" ) })?; - let org_id = selected_org.map(|org| org.id.clone()).unwrap_or_default(); - let slot_key = oauth_slot_key(&org_id, &email); + let app_url = canonical_url(&app_url).to_string(); + let slot_key = oauth_slot_key(&app_url); + let mut store = load_auth_store()?; + if let Some(old_profile) = store.profiles.get(&slot_key) { + delete_all_profile_secrets(&slot_key, old_profile); + } save_profile_oauth_refresh_token(&slot_key, refresh_token)?; save_profile_oauth_access_token(&slot_key, &tokens.access_token)?; let _ = delete_profile_secret(&slot_key); - let mut store = load_auth_store()?; - if let Some(old_profile) = store.profiles.get(&slot_key) { - delete_legacy_profile_secrets(old_profile); - } store.profiles.insert( slot_key, AuthProfile { auth_kind: AuthKind::Oauth, api_url: Some(api_url), app_url: Some(app_url), - org_id: Some(org_id), - org_name: selected_org.map(|org| org.name.clone()), + org_id: None, + org_name: None, oauth_access_expires_at, user_name: jwt_id.name, email: jwt_id.email, @@ -1942,7 +2082,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { )? .ok_or_else(|| { anyhow::anyhow!( - "no OAuth login selected; pass --org or run `bt auth logins` to see available logins" + "no OAuth login selected; pass --app-url or run `bt auth logins` to see available logins" ) })?; let profile = store @@ -1953,9 +2093,10 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { anyhow::anyhow!("OAuth login not found; run `bt auth logins` to see available logins") })?; - let api_url = profile + let api_url = base .api_url .clone() + .or_else(|| profile.api_url.clone()) .unwrap_or_else(|| DEFAULT_API_URL.to_string()); let previous_expires_at = profile.oauth_access_expires_at; let refresh_token = @@ -2005,14 +2146,6 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { } } save_auth_store(&store)?; - reconcile_oauth_slot_from_access_token( - &mut store, - profile_name.as_str(), - &refreshed.access_token, - profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), - ) - .await?; - if let Some(expires_at) = new_expires_at { let now = current_unix_timestamp(); let remaining = expires_at.saturating_sub(now); @@ -2030,8 +2163,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { base.json, serde_json::json!({ "auth": "oauth", - "org": profile.org_name, - "org_id": profile.org_id.filter(|org_id| !org_id.trim().is_empty()), + "app_url": profile.app_url, "user_email": profile.email, "access_expires_at": new_expires_at, "refresh_token_rotated": refresh_rotated, @@ -2042,10 +2174,9 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { } fn format_login_success(selected_org: Option<&LoginOrgInfo>, api_url: &str) -> String { - match selected_org { - Some(org) => format!("Logged in as {} (api: {api_url})", org.name), - None => format!("Logged in (cross-org, api: {api_url})"), - } + selected_org + .map(|org| format!("Logged in as {} (api: {api_url})", org.name)) + .unwrap_or_else(|| format!("Logged in (api: {api_url})")) } fn build_login_context_for_selected_org( @@ -2076,7 +2207,7 @@ fn format_post_login_context( match (selected_org, project) { (Some(org), Some(project)) => format!("{}/{}", org.name, project.name), (Some(org), None) => org.name.clone(), - (None, _) => "cross-org mode".to_string(), + (None, _) => "Braintrust".to_string(), } } @@ -2091,11 +2222,8 @@ async fn resolve_post_login_project( return Ok(None); }; - let selected_org = selected_org.ok_or_else(|| { - anyhow::anyhow!( - "cannot set a default project in cross-org mode; rerun `bt auth login --org --project `" - ) - })?; + let selected_org = selected_org + .ok_or_else(|| anyhow::anyhow!("an organization is required to select a project"))?; let ctx = build_login_context_for_selected_org(credential, api_url, app_url, Some(selected_org)); let client = ApiClient::new(&ctx)?; @@ -2117,23 +2245,35 @@ async fn persist_post_login_context( resolve_post_login_project(base, credential, api_url, app_url, selected_org).await?; let (path, _) = scope.resolve(ui::can_prompt(), "Where to use this login")?; let mut cfg = config::load_file(&path); - let org = selected_org.map_or("", |org| org.name.as_str()); + let selected_org = selected_org + .ok_or_else(|| anyhow::anyhow!("an organization is required to update config"))?; let preserve_project = project.is_none() - && selected_org.is_some() - && config::org_option(cfg.org.as_deref()) == Some(org); - if !preserve_project { - cfg.set_context( - Some(org), - project - .as_ref() - .map(|project| (project.name.as_str(), project.id.as_str())), - ); - } + && config::org_option(cfg.org.as_deref()) == Some(selected_org.name.as_str()) + && cfg.org_id.as_deref() == Some(selected_org.id.as_str()) + && cfg + .app_url + .as_deref() + .is_some_and(|url| config::urls_equal(url, app_url)); + let selected_project = if preserve_project { + cfg.project.clone().zip(cfg.project_id.clone()) + } else { + project + .as_ref() + .map(|project| (project.name.clone(), project.id.clone())) + }; + cfg.set_context( + (selected_org.name.as_str(), selected_org.id.as_str()), + selected_project + .as_ref() + .map(|(name, id)| (name.as_str(), id.as_str())), + app_url, + api_url, + ); config::save_file(&path, &cfg) .with_context(|| format!("Could not save config to {}", path.display()))?; Ok(PostLoginContextUpdate { - display: format_post_login_context(selected_org, project.as_ref()), + display: format_post_login_context(Some(selected_org), project.as_ref()), path, }) } @@ -2150,14 +2290,14 @@ fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) -> } fn filter_auth_store( + base: &BaseArgs, store: &AuthStore, - org: Option<&str>, kind: Option, api_key_hint: Option<&str>, ) -> AuthStore { let mut filtered = store.clone(); filtered.profiles.retain(|_, profile| { - org.is_none_or(|org| profile_matches_org_identifier(profile, org)) + profile_matches_url_filters(base, profile) && kind.is_none_or(|kind| profile.auth_kind == kind) && api_key_hint.is_none_or(|hint| { profile.auth_kind == AuthKind::ApiKey @@ -2167,15 +2307,68 @@ fn filter_auth_store( filtered } +async fn filter_auth_store_for_org( + base: &BaseArgs, + store: &mut AuthStore, + candidates: AuthStore, + org: Option<&str>, +) -> Result { + let Some(org) = org else { + return Ok(candidates); + }; + let mut filtered = AuthStore::default(); + for (slot, profile) in candidates.profiles { + let matches = match profile.auth_kind { + AuthKind::ApiKey => profile_matches_org_identifier(&profile, org), + AuthKind::Oauth => { + let mut oauth_base = base.clone(); + oauth_base.app_url = Some(profile_app_url(&profile).to_string()); + if oauth_base.api_url_source.is_none() { + oauth_base.api_url = profile.api_url.clone(); + } + let token = load_oauth_access_token(&oauth_base, store, &slot).await?; + let orgs = fetch_login_orgs(&token, profile_app_url(&profile)).await?; + find_login_org(&orgs, org).is_some() + } + }; + if matches { + filtered.profiles.insert(slot, profile); + } + } + Ok(filtered) +} + async fn run_logins(base: &BaseArgs, _args: AuthLoginsArgs) -> Result<()> { let mut store = load_auth_store()?; - let has_filter = base.org_name.is_some() || base.prefer_api_key; - let filtered = filter_auth_store( + let requested_org = matches!( + base.org_name_source, + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable) + ) + .then(|| base.org_name.as_deref()) + .flatten(); + let has_url_filter = matches!( + base.app_url_source, + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable) + ) || matches!( + base.api_url_source, + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable) + ); + let has_filter = requested_org.is_some() || base.prefer_api_key || has_url_filter; + let mut filter_base = base.clone(); + if filter_base.app_url_source.is_none() { + filter_base.app_url = None; + } + if filter_base.api_url_source.is_none() { + filter_base.api_url = None; + } + let candidates = filter_auth_store( + &filter_base, &store, - base.org_name.as_deref(), base.prefer_api_key.then_some(AuthKind::ApiKey), None, ); + let filtered = + filter_auth_store_for_org(&filter_base, &mut store, candidates, requested_org).await?; if filtered.profiles.is_empty() { return emit_result(base.json, serde_json::json!([]), || { if store.profiles.is_empty() && !has_filter { @@ -2226,6 +2419,8 @@ fn auth_profile_json(profile: &AuthProfile, status: &str) -> serde_json::Value { "user_name": profile.user_name, "user_email": profile.email, "api_key_hint": profile.api_key_hint, + "app_url": profile.app_url, + "api_url": profile.api_url, "status": status, }) } @@ -2277,7 +2472,7 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result< }) } -fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { +async fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { let store = load_auth_store()?; if store.profiles.is_empty() { return emit_result(base.json, serde_json::json!({ "status": "empty" }), || { @@ -2287,18 +2482,29 @@ fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { let requested_org = if matches!( base.org_name_source, - Some(crate::args::ArgValueSource::CommandLine) + Some(crate::args::ArgValueSource::CommandLine | crate::args::ArgValueSource::EnvVariable) ) { config::org_option(base.org_name.as_deref()) } else { None }; - let filtered = filter_auth_store( + let mut filter_base = base.clone(); + if filter_base.app_url_source.is_none() { + filter_base.app_url = None; + } + if filter_base.api_url_source.is_none() { + filter_base.api_url = None; + } + let candidates = filter_auth_store( + &filter_base, &store, - requested_org, args.oauth.then_some(AuthKind::Oauth), args.api_key_hint.as_deref(), ); + let mut mutable_store = store.clone(); + let filtered = + filter_auth_store_for_org(&filter_base, &mut mutable_store, candidates, requested_org) + .await?; let candidates = filtered .profiles .keys() @@ -2319,7 +2525,7 @@ fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { _ => { let labels = candidate_identities(&candidates, &filtered).join(", "); bail!( - "multiple auth logins match: {labels}. Rerun interactively, or use --org with --oauth or --api-key-hint to disambiguate." + "multiple auth logins match: {labels}. Rerun interactively, use --app-url with --oauth, or use --org with --api-key-hint ." ); } }; @@ -2382,6 +2588,10 @@ pub struct ProfileVerification { pub user_email: Option, #[serde(skip_serializing_if = "Option::is_none")] pub api_key_hint: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub app_url: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub api_url: Option, pub status: String, #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, @@ -2389,9 +2599,7 @@ pub struct ProfileVerification { fn build_verification( name: &str, - auth_kind: &str, - org: Option, - org_id: Option, + profile: &AuthProfile, jwt_id: Option, api_key_hint: Option, status: ProfileStatus, @@ -2405,12 +2613,17 @@ fn build_verification( ProfileVerification { name: name.to_string(), slot_hash: None, - auth: auth_kind.to_string(), - org, - org_id, + auth: auth_kind_label(profile.auth_kind).to_string(), + org: profile.org_name.clone(), + org_id: profile + .org_id + .clone() + .filter(|org_id| !org_id.trim().is_empty()), user_name: jwt_id.as_ref().and_then(|j| j.name.clone()), user_email: jwt_id.as_ref().and_then(|j| j.email.clone()), api_key_hint, + app_url: profile.app_url.clone(), + api_url: profile.api_url.clone(), status: status_str.to_string(), error, } @@ -2418,20 +2631,8 @@ fn build_verification( async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerification { let app_url = profile.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); - let auth_kind = auth_kind_label(profile.auth_kind); let mk = |status, jwt_id: Option, hint: Option| { - build_verification( - name, - auth_kind, - profile.org_name.clone(), - profile - .org_id - .clone() - .filter(|org_id| !org_id.trim().is_empty()), - jwt_id, - hint, - status, - ) + build_verification(name, profile, jwt_id, hint, status) }; let credential = match load_credential_for_profile(name, profile) { @@ -2455,7 +2656,7 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifi match fetch_login_orgs(&credential, app_url).await { Ok(orgs) => { let mut verification = mk(ProfileStatus::Ok, jwt_id, hint); - if !is_cross_org_oauth_profile(profile) { + if profile.auth_kind == AuthKind::ApiKey { if let Some(org) = profile .org_id .as_deref() @@ -2470,8 +2671,6 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifi verification.org = Some(org.name.clone()); verification.org_id = Some(org.id.clone()); } - } - if profile.auth_kind == AuthKind::ApiKey { verification.slot_hash = Some(api_key_hash(&credential)); } verification @@ -2532,11 +2731,13 @@ fn reconcile_verified_auth_slots( continue; }; - if let Some(org_id) = verification.org_id.as_deref() { - profile.org_id = Some(org_id.to_string()); - profile.org_name = verification.org.clone(); - } else if is_cross_org_oauth_profile(&profile) { - profile.org_id = Some(String::new()); + if profile.auth_kind == AuthKind::ApiKey { + if let Some(org_id) = verification.org_id.as_deref() { + profile.org_id = Some(org_id.to_string()); + profile.org_name = verification.org.clone(); + } + } else { + profile.org_id = None; profile.org_name = None; } @@ -2573,10 +2774,14 @@ fn reconcile_verified_auth_slots( } fn format_verification_line(v: &ProfileVerification) -> String { - let mut parts = vec![ - config::display_org(v.org.as_deref().unwrap_or("")).to_string(), - v.auth.clone(), - ]; + let subject = if v.auth == "oauth" { + v.app_url + .clone() + .unwrap_or_else(|| DEFAULT_APP_URL.to_string()) + } else { + v.org.clone().unwrap_or_else(|| "(unknown org)".to_string()) + }; + let mut parts = vec![subject, v.auth.clone()]; match v.status.as_str() { "ok" => { if let Some(id) = identity_label( @@ -2628,6 +2833,8 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { "user_name": p.user_name, "user_email": p.email, "api_key_hint": p.api_key_hint, + "app_url": p.app_url, + "api_url": p.api_url, "status": "unchecked" }) }) @@ -2676,7 +2883,6 @@ fn select_login_org( default_org_name: Option<&str>, interactive: bool, verbose: bool, - allow_cross_org: bool, quiet_requested: bool, ) -> Result> { if orgs.is_empty() { @@ -2684,10 +2890,6 @@ fn select_login_org( } sort_login_orgs(&mut orgs); - if requested_org_name == Some("") { - return Ok(None); - } - if let Some(name) = requested_org_name { return find_login_org(&orgs, name) .cloned() @@ -2700,44 +2902,30 @@ fn select_login_org( } if !interactive { - if allow_cross_org { - bail!( - "organization selection required in non-interactive mode; pass --org or rerun interactively to choose cross-org mode" - ); - } return Ok(None); } - let default_org_matched = move_default_login_org_first(&mut orgs, default_org_name); - let offset = if allow_cross_org { 1 } else { 0 }; - let mut labels: Vec = Vec::new(); - if allow_cross_org { - labels.push( - "No default org (cross-org mode; pass --org or BRAINTRUST_ORG_NAME when needed)" - .to_string(), - ); - } - labels.extend(orgs.iter().map(|org| { - if verbose { - let api_url = org.api_url.as_deref().unwrap_or(DEFAULT_API_URL); - format!("{} [{}] ({})", org.name, org.id, api_url) - } else { - org.name.clone() - } - })); + move_default_login_org_first(&mut orgs, default_org_name); + let labels: Vec = orgs + .iter() + .map(|org| { + if verbose { + let api_url = org.api_url.as_deref().unwrap_or(DEFAULT_API_URL); + format!("{} [{}] ({})", org.name, org.id, api_url) + } else { + org.name.clone() + } + }) + .collect(); let label_refs: Vec<&str> = labels.iter().map(String::as_str).collect(); if !quiet_requested { eprintln!("\n\nA Braintrust organization is usually a team or a company."); } - let default = if default_org_matched { offset } else { 0 }; - let selection = ui::fuzzy_select("Select organization", &label_refs, default)?; - if allow_cross_org && selection == 0 { - return Ok(None); - } + let selection = ui::fuzzy_select("Select organization", &label_refs, 0)?; Ok(Some( orgs.into_iter() - .nth(selection - offset) + .nth(selection) .expect("selected index should be in range"), )) } @@ -2855,8 +3043,11 @@ fn resolve_profile_api_url( if let Some(api_url) = explicit_api_url { return Ok(api_url); } - if let Some(api_url) = selected_org.and_then(|org| org.api_url.clone()) { - return Ok(api_url); + if let Some(selected_org) = selected_org { + return Ok(selected_org + .api_url + .clone() + .unwrap_or_else(|| DEFAULT_API_URL.to_string())); } let mut api_urls = orgs @@ -2873,8 +3064,6 @@ fn resolve_profile_api_url( .unwrap_or_else(|| DEFAULT_API_URL.to_string())); } - // A cross-org login spans orgs on different data planes. Let the user pick - // which API URL to store rather than failing outright. if can_prompt { let idx = ui::fuzzy_select("Select API URL", &api_urls, 0)?; return Ok(api_urls @@ -3980,8 +4169,8 @@ fn api_key_hash(api_key: &str) -> String { sha256_hex(api_key) } -fn oauth_slot_key(org_id: &str, email: &str) -> String { - format!("{org_id}::{email}") +fn oauth_slot_key(app_url: &str) -> String { + format!("oauth::{}", sha256_hex(canonical_url(app_url))) } fn api_key_slot_key(api_key_hash: &str, org_id: &str) -> String { @@ -4030,29 +4219,38 @@ fn load_auth_store_from_path(path: &Path) -> Result { fn migrate_auth_store(store: AuthStore) -> AuthStore { let mut migrated = AuthStore::default(); + let mut oauth_refresh_usable = BTreeMap::::new(); for (old_key, mut profile) in store.profiles { normalize_profile_cached_fields_from_key(&old_key, &mut profile); - if profile.auth_kind == AuthKind::Oauth - && profile.org_id.is_none() - && profile - .org_name - .as_deref() - .is_none_or(|org| org.trim().is_empty()) - { - profile.org_id = Some(String::new()); + let refresh_usable = profile.auth_kind == AuthKind::Oauth + && matches!( + load_profile_oauth_refresh_token_for_profile(&old_key, &profile), + Ok(Some(_)) + ); + if profile.auth_kind == AuthKind::Oauth { + profile.org_id = None; profile.org_name = None; + profile.app_url = Some(canonical_url(profile_app_url(&profile)).to_string()); } let new_key = canonical_profile_key(&old_key, &profile); if new_key != old_key && profile.legacy_secret_key.is_none() { profile.legacy_secret_key = Some(old_key.clone()); } - if migrated - .profiles - .get(&new_key) - .is_some_and(|existing| !should_replace_migrated_profile(existing, &profile)) - { - continue; + if let Some(existing) = migrated.profiles.get(&new_key) { + let existing_usable = oauth_refresh_usable.get(&new_key).copied().unwrap_or(false); + let replace = match (existing.auth_kind, profile.auth_kind) { + (AuthKind::Oauth, AuthKind::Oauth) => { + (refresh_usable && !existing_usable) + || (refresh_usable == existing_usable + && should_replace_migrated_profile(existing, &profile)) + } + _ => false, + }; + if !replace { + continue; + } } + oauth_refresh_usable.insert(new_key.clone(), refresh_usable); migrated.profiles.insert(new_key, profile); } migrated @@ -4166,17 +4364,7 @@ fn normalize_profile_cached_fields_from_key(current_key: &str, profile: &mut Aut fn canonical_profile_key(current_key: &str, profile: &AuthProfile) -> String { match profile.auth_kind { - AuthKind::Oauth => { - let Some(email) = profile.email.as_deref().filter(|value| !value.is_empty()) else { - return current_key.to_string(); - }; - let org_id = profile.org_id.as_deref().unwrap_or_default(); - if profile.org_id.is_some() || profile.org_name.is_none() { - oauth_slot_key(org_id, email) - } else { - current_key.to_string() - } - } + AuthKind::Oauth => oauth_slot_key(profile_app_url(profile)), AuthKind::ApiKey => match ( profile .api_key_hash @@ -4537,6 +4725,13 @@ mod tests { crate::config::save_global(&cfg).expect("save global config"); } + fn set_global_config_urls(app_url: &str, api_url: Option<&str>) { + let mut cfg = crate::config::load_global().expect("load global config"); + cfg.app_url = Some(app_url.to_string()); + cfg.api_url = api_url.map(str::to_string); + crate::config::save_global(&cfg).expect("save global config URLs"); + } + fn org_profile(kind: AuthKind, org_id: &str, org_name: &str) -> AuthProfile { AuthProfile { auth_kind: kind, @@ -4612,7 +4807,10 @@ mod tests { fn invalid_grant_refresh_error_is_treated_as_recoverable() { let profile = org_profile(AuthKind::Oauth, "org_test", "BT Staging"); let command = oauth_reauth_command(&profile); - assert_eq!(command, "bt auth login --oauth --org 'BT Staging'"); + assert_eq!( + command, + "bt auth login --oauth --app-url https://www.braintrust.dev" + ); let err = map_refresh_oauth_error( "https://api.example.com", &profile, @@ -4751,62 +4949,20 @@ mod tests { assert_eq!(DEFAULT_APP_URL, "https://www.braintrust.dev"); } - #[tokio::test] - async fn active_auth_info_without_org_uses_cross_org_oauth_only() { - let _env = TestEnv::new(None, None).await; - let mut store = AuthStore::default(); - store.profiles.insert( - api_key_slot_key(&api_key_hash("test-api-key"), "org_fake"), - AuthProfile { - api_key_hint: Some("sk-****abcde".to_string()), - ..org_profile(AuthKind::ApiKey, "org_fake", "test-org") - }, - ); - store.profiles.insert( - oauth_slot_key("", "user@example.test"), - AuthProfile { - auth_kind: AuthKind::Oauth, - org_id: Some(String::new()), - org_name: None, - user_name: Some("Test User".to_string()), - email: Some("user@example.test".to_string()), - ..Default::default() - }, - ); - save_auth_store(&store).expect("save auth store"); - - let info = active_auth_info(&make_base(), None) - .expect("resolve active auth") - .expect("active auth info"); - - assert_eq!(info.auth_method, "oauth"); - assert_eq!(info.email.as_deref(), Some("user@example.test")); - assert_eq!(info.org_name, None); - - let mut base = make_base(); - base.prefer_api_key = true; - assert!(resolve_auth(&base) - .await - .unwrap_err() - .to_string() - .contains("cross-org")); - assert!(active_auth_info(&base, None) - .unwrap_err() - .to_string() - .contains("cross-org")); - } - - fn save_cached_oauth_login(store: &mut AuthStore, org_id: &str, org_name: &str) -> String { - let slot_key = oauth_slot_key(org_id, "user@example.test"); + fn save_cached_oauth_login(store: &mut AuthStore, app_url: &str) -> String { + let slot_key = oauth_slot_key(app_url); store.profiles.insert( slot_key.clone(), AuthProfile { api_url: Some("https://api.example.test".to_string()), - app_url: Some("https://www.example.test".to_string()), + app_url: Some(app_url.to_string()), oauth_access_expires_at: Some(current_unix_timestamp() + 3600), user_name: Some("Test User".to_string()), email: Some("user@example.test".to_string()), - ..org_profile(AuthKind::Oauth, org_id, org_name) + auth_kind: AuthKind::Oauth, + org_id: None, + org_name: None, + ..Default::default() }, ); save_profile_secret_plaintext( @@ -4821,10 +4977,12 @@ mod tests { async fn auth_precedence_keeps_env_api_key_below_oauth() { let _env = TestEnv::new(None, None).await; let mut store = AuthStore::default(); - save_cached_oauth_login(&mut store, "org_fake", "test-org"); + let app_url = spawn_api_key_login_server("test-org"); + save_cached_oauth_login(&mut store, &app_url); save_auth_store(&store).expect("save auth store"); let mut base = make_base(); base.org_name = Some("test-org".to_string()); + base.app_url = Some(app_url); base.api_key = Some("environment-api-key".to_string()); base.api_key_source = Some(crate::args::ArgValueSource::EnvVariable); @@ -4841,13 +4999,14 @@ mod tests { async fn auth_precedence_cli_api_key_overrides_oauth() { let _env = TestEnv::new(None, None).await; let mut store = AuthStore::default(); - save_cached_oauth_login(&mut store, "org_fake", "test-org"); + let app_url = spawn_api_key_login_server("test-org"); + save_cached_oauth_login(&mut store, &app_url); save_auth_store(&store).expect("save auth store"); let mut base = make_base(); base.org_name = Some("test-org".to_string()); base.api_key = Some("command-line-api-key".to_string()); base.api_key_source = Some(crate::args::ArgValueSource::CommandLine); - base.app_url = Some(spawn_api_key_login_server("test-org")); + base.app_url = Some(app_url); let resolved = resolve_auth(&base).await.expect("resolve auth"); @@ -4899,14 +5058,15 @@ mod tests { async fn auth_precedence_prefer_api_key_promotes_env_api_key() { let _env = TestEnv::new(None, None).await; let mut store = AuthStore::default(); - save_cached_oauth_login(&mut store, "org_fake", "test-org"); + let app_url = spawn_api_key_login_server("test-org"); + save_cached_oauth_login(&mut store, &app_url); save_auth_store(&store).expect("save auth store"); let mut base = make_base(); base.org_name = Some("test-org".to_string()); base.api_key = Some("environment-api-key".to_string()); base.api_key_source = Some(crate::args::ArgValueSource::EnvVariable); base.prefer_api_key = true; - base.app_url = Some(spawn_api_key_login_server("test-org")); + base.app_url = Some(app_url); let resolved = resolve_auth(&base).await.expect("resolve auth"); @@ -4918,10 +5078,12 @@ mod tests { async fn auth_precedence_prefer_api_key_falls_back_to_oauth_without_key() { let _env = TestEnv::new(None, None).await; let mut store = AuthStore::default(); - save_cached_oauth_login(&mut store, "org_fake", "test-org"); + let app_url = spawn_api_key_login_server("test-org"); + save_cached_oauth_login(&mut store, &app_url); save_auth_store(&store).expect("save auth store"); let mut base = make_base(); base.org_name = Some("test-org".to_string()); + base.app_url = Some(app_url); base.prefer_api_key = true; let resolved = resolve_auth(&base).await.expect("resolve auth"); @@ -4960,11 +5122,13 @@ mod tests { let _env = TestEnv::new(None, None).await; let mut store = AuthStore::default(); store.profiles.insert( - oauth_slot_key("org_fake", "user@example.test"), + oauth_slot_key(DEFAULT_APP_URL), AuthProfile { + auth_kind: AuthKind::Oauth, + app_url: Some(DEFAULT_APP_URL.to_string()), user_name: Some("Test User".to_string()), email: Some("user@example.test".to_string()), - ..org_profile(AuthKind::Oauth, "org_fake", "test-org") + ..Default::default() }, ); store.profiles.insert( @@ -5051,20 +5215,23 @@ mod tests { ); let migrated = migrate_auth_store(store); - let key = oauth_slot_key("org_fake", "user@example.test"); + let key = oauth_slot_key(DEFAULT_APP_URL); let profile = migrated.profiles.get(&key).expect("migrated profile"); assert_eq!(profile.legacy_secret_key.as_deref(), Some("work")); + assert_eq!(profile.org_id, None); + assert_eq!(profile.org_name, None); } #[test] - fn migrate_auth_store_rekeys_cross_org_oauth_with_empty_org_id() { + fn migrate_auth_store_purges_old_oauth_org_scope() { let mut store = AuthStore::default(); store.profiles.insert( - "legacy-cross-org".to_string(), + "legacy-login".to_string(), AuthProfile { auth_kind: AuthKind::Oauth, - org_name: None, + org_id: Some("org_old".to_string()), + org_name: Some("old-org".to_string()), email: Some("user@example.test".to_string()), ..Default::default() }, @@ -5073,14 +5240,12 @@ mod tests { let migrated = migrate_auth_store(store); let profile = migrated .profiles - .get(&oauth_slot_key("", "user@example.test")) - .expect("cross-org OAuth slot"); + .get(&oauth_slot_key(DEFAULT_APP_URL)) + .expect("instance OAuth slot"); - assert_eq!(profile.org_id.as_deref(), Some("")); - assert_eq!( - profile.legacy_secret_key.as_deref(), - Some("legacy-cross-org") - ); + assert_eq!(profile.org_id, None); + assert_eq!(profile.org_name, None); + assert_eq!(profile.legacy_secret_key.as_deref(), Some("legacy-login")); } #[test] @@ -5134,7 +5299,7 @@ mod tests { let persisted: AuthStore = serde_json::from_str(&fs::read_to_string(&path).expect("read migrated store")) .expect("parse migrated store"); - let slot_key = oauth_slot_key("org_fake", "user@example.test"); + let slot_key = oauth_slot_key(DEFAULT_APP_URL); for migrated in [&loaded, &persisted] { let profile = migrated @@ -5179,6 +5344,8 @@ mod tests { user_name: Some("Test User".to_string()), user_email: Some("user@example.test".to_string()), api_key_hint: None, + app_url: Some(DEFAULT_APP_URL.to_string()), + api_url: None, status: "ok".to_string(), error: None, }, @@ -5191,6 +5358,8 @@ mod tests { user_name: None, user_email: None, api_key_hint: Some("sk-****abcde".to_string()), + app_url: None, + api_url: None, status: "ok".to_string(), error: None, }, @@ -5201,7 +5370,7 @@ mod tests { let oauth = store .profiles - .get(&oauth_slot_key("org_fake", "user@example.test")) + .get(&oauth_slot_key(DEFAULT_APP_URL)) .expect("canonical OAuth slot"); assert_eq!(oauth.legacy_secret_key.as_deref(), Some("legacy-oauth")); let api_key = store @@ -5229,12 +5398,40 @@ mod tests { } let migrated = migrate_auth_store(store); - let key = oauth_slot_key("org_fake", "user@example.test"); + let key = oauth_slot_key(DEFAULT_APP_URL); assert_eq!(migrated.profiles.len(), 1); let profile = migrated.profiles.get(&key).expect("migrated profile"); assert_eq!(profile.legacy_secret_key.as_deref(), Some("new")); } + #[tokio::test] + async fn migrate_auth_store_prefers_loadable_refresh_token_before_expiry() { + let _env = TestEnv::new(None, None).await; + let mut store = AuthStore::default(); + for (name, expires_at) in [("usable-old", 10), ("missing-new", 20)] { + store.profiles.insert( + name.to_string(), + AuthProfile { + auth_kind: AuthKind::Oauth, + oauth_access_expires_at: Some(expires_at), + ..Default::default() + }, + ); + } + save_profile_secret_plaintext( + &oauth_refresh_secret_key("usable-old"), + "test-refresh-token", + ) + .expect("save refresh token"); + + let migrated = migrate_auth_store(store); + let profile = migrated + .profiles + .get(&oauth_slot_key(DEFAULT_APP_URL)) + .expect("migrated OAuth profile"); + assert_eq!(profile.legacy_secret_key.as_deref(), Some("usable-old")); + } + #[test] fn migration_reports_dropped_duplicate_secret_as_orphan() { // Two legacy OAuth entries for the same org+email collapse onto one @@ -5293,6 +5490,12 @@ mod tests { let org = config_auth_context_from_config(&base, &cfg); assert_eq!(org.as_deref(), Some("local-org")); + + let other_instance = BaseArgs { + app_url: Some("https://other.example.test".into()), + ..base + }; + assert_eq!(config_auth_context_from_config(&other_instance, &cfg), None); } #[test] @@ -5430,80 +5633,143 @@ mod tests { assert!(err.to_string().contains("multiple oauth logins")); } - fn login_filter_store() -> AuthStore { + #[tokio::test] + async fn available_instances_dedupes_logins_by_app_url() { + let _env = TestEnv::new(None, None).await; let mut store = AuthStore::default(); - for (slot, kind, suffix, hint) in [ - ("oauth-a", AuthKind::Oauth, "a", None), - ("key-a", AuthKind::ApiKey, "a", Some("sk-****aaaaa")), - ("key-b", AuthKind::ApiKey, "b", Some("sk-****bbbbb")), + for (slot, kind, app_url) in [ + ("oauth-a", AuthKind::Oauth, "https://one.example.test/"), + ("key-a", AuthKind::ApiKey, "https://one.example.test"), + ("key-b", AuthKind::ApiKey, "https://two.example.test"), ] { store.profiles.insert( slot.into(), AuthProfile { - api_key_hint: hint.map(str::to_string), - ..org_profile( - kind, - &format!("org_test_{suffix}"), - &format!("test-org-{suffix}"), - ) + auth_kind: kind, + app_url: Some(app_url.into()), + org_id: (kind == AuthKind::ApiKey).then(|| format!("org_{slot}")), + org_name: (kind == AuthKind::ApiKey).then(|| format!("org-{slot}")), + ..Default::default() }, ); } + save_auth_store(&store).expect("save auth store"); + + let instances = available_instances(&BaseArgs::default()).expect("list instances"); + assert_eq!( + instances + .iter() + .map(|instance| instance.app_url.as_str()) + .collect::>(), + ["https://one.example.test", "https://two.example.test"] + ); + + let filtered = available_instances(&BaseArgs { + app_url: Some("https://two.example.test/".into()), + app_url_source: Some(crate::args::ArgValueSource::CommandLine), + ..Default::default() + }) + .expect("filter instances"); + assert_eq!(filtered.len(), 1); + assert_eq!(filtered[0].app_url, "https://two.example.test"); + } + + #[tokio::test] + async fn org_filter_includes_oauth_login_when_discovered_membership_matches() { + let _env = TestEnv::new(None, None).await; + let app_url = spawn_api_key_login_server("test-org"); + let mut store = AuthStore::default(); + let oauth_slot = save_cached_oauth_login(&mut store, &app_url); store.profiles.insert( - "cross".into(), + "other-key".into(), AuthProfile { - auth_kind: AuthKind::Oauth, - org_id: Some(String::new()), - ..Default::default() + app_url: Some(app_url.clone()), + api_url: Some("https://api.example.test".into()), + ..org_profile(AuthKind::ApiKey, "org_other", "other-org") }, ); - store + save_auth_store(&store).expect("save auth store"); + let base = BaseArgs::default(); + let candidates = filter_auth_store(&base, &store, None, None); + let filtered = filter_auth_store_for_org(&base, &mut store, candidates, Some("test-org")) + .await + .expect("filter by org"); + assert_eq!( + filtered.profiles.into_keys().collect::>(), + [oauth_slot] + ); } #[test] - fn login_and_logout_filters_compose() { - let store = login_filter_store(); - let matches = |org, kind, hint| { - filter_auth_store(&store, org, kind, hint) - .profiles - .into_keys() - .collect::>() + fn command_auth_uses_builtin_url_defaults_when_urls_are_unset() { + let default_oauth = AuthProfile { + auth_kind: AuthKind::Oauth, + app_url: Some(DEFAULT_APP_URL.into()), + ..Default::default() }; - assert_eq!( - matches(Some("test-org-a"), None, None), - ["key-a", "oauth-a"] - ); - assert_eq!(matches(Some("org_test_b"), None, None), ["key-b"]); - assert_eq!( - matches(Some("test-org-a"), Some(AuthKind::ApiKey), None), - ["key-a"] - ); - assert_eq!(matches(Some(""), None, None), ["cross"]); - assert!(matches(Some(""), Some(AuthKind::ApiKey), None).is_empty()); - assert_eq!(matches(None, None, None).len(), 4); - assert_eq!( - matches(Some("test-org-a"), Some(AuthKind::Oauth), None), - ["oauth-a"] - ); - assert_eq!(matches(None, None, Some("sk-****bbbbb")), ["key-b"]); + let custom_oauth = AuthProfile { + auth_kind: AuthKind::Oauth, + app_url: Some("https://www.example.test".into()), + ..Default::default() + }; + assert!(profile_matches_urls(&BaseArgs::default(), &default_oauth)); + assert!(!profile_matches_urls(&BaseArgs::default(), &custom_oauth)); + assert!(profile_matches_url_filters( + &BaseArgs::default(), + &custom_oauth + )); + } - let mut picker_store = store.clone(); - picker_store.profiles.insert( - "oauth-a-duplicate".into(), - picker_store.profiles["oauth-a"].clone(), + #[test] + fn login_filter_matches_instance_and_auth_kind() { + let mut store = AuthStore::default(); + store.profiles.insert( + "oauth".into(), + AuthProfile { + auth_kind: AuthKind::Oauth, + app_url: Some("https://www.example.test".into()), + ..Default::default() + }, + ); + store.profiles.insert( + "key".into(), + AuthProfile { + app_url: Some("https://www.example.test".into()), + api_url: Some("https://api.example.test".into()), + api_key_hint: Some("sk-****abcde".into()), + ..org_profile(AuthKind::ApiKey, "org_test", "test-org") + }, ); - assert_eq!(saved_login_names(&picker_store, true).len(), 4); - assert_eq!(saved_login_names(&picker_store, false).len(), 3); + let base = BaseArgs { + app_url: Some("https://www.example.test/".into()), + api_url: Some("https://api.example.test/".into()), + ..Default::default() + }; + let filtered = filter_auth_store(&base, &store, None, None); + assert_eq!(filtered.profiles.len(), 2); + let filtered = filter_auth_store(&base, &store, Some(AuthKind::ApiKey), None); + assert_eq!(filtered.profiles.into_keys().collect::>(), ["key"]); + + let mismatched_api = BaseArgs { + app_url: base.app_url.clone(), + api_url: Some("https://other-api.example.test".into()), + ..Default::default() + }; + let filtered = filter_auth_store(&mismatched_api, &store, None, None); + assert_eq!(filtered.profiles.into_keys().collect::>(), ["oauth"]); } #[tokio::test] async fn post_login_context_preserves_only_same_org_projects() { let _env = TestEnv::new(None, None).await; - let save = |org: &str| { + let save = |org: &str, org_id: &str| { crate::config::save_global(&crate::config::Config { org: Some(org.into()), + org_id: Some(org_id.into()), project: Some("test-project".into()), project_id: Some("proj_test".into()), + app_url: Some("https://www.example.test".into()), + api_url: Some("https://api.example.test".into()), ..Default::default() }) .unwrap(); @@ -5525,27 +5791,20 @@ mod tests { crate::config::load_global().unwrap() }; - save("old-org"); + save("old-org", "org_old"); let cfg = persist(Some(login_org("org_test", "test-org"))).await; assert_eq!((cfg.org.as_deref(), cfg.project), (Some("test-org"), None)); - save("test-org"); + save("test-org", "org_test"); let cfg = persist(Some(login_org("org_test", "test-org"))).await; assert_eq!( (cfg.project.as_deref(), cfg.project_id.as_deref()), (Some("test-project"), Some("proj_test")) ); - - save(""); - let cfg = persist(None).await; - assert_eq!( - (cfg.org.as_deref(), cfg.project, cfg.project_id), - (Some(""), None, None) - ); } #[tokio::test] - async fn resolve_post_login_project_rejects_cross_org_default_project() { + async fn resolve_post_login_project_requires_an_org() { let mut base = make_base(); base.project = Some("demo-project".to_string()); @@ -5557,11 +5816,9 @@ mod tests { None, ) .await - .expect_err("cross-org project selection should fail"); + .expect_err("missing org should fail"); - assert!(err - .to_string() - .contains("cannot set a default project in cross-org mode")); + assert!(err.to_string().contains("organization is required")); } #[test] @@ -5654,6 +5911,8 @@ mod tests { user_name: None, user_email: None, api_key_hint: None, + app_url: Some(DEFAULT_APP_URL.to_string()), + api_url: None, status: "ok".into(), error: None, }; @@ -5747,6 +6006,8 @@ mod tests { user_name: identity.map(str::to_string), user_email: identity.map(|_| "user@example.test".into()), api_key_hint: hint.map(str::to_string), + app_url: (auth == "oauth").then(|| DEFAULT_APP_URL.to_string()), + api_url: None, status: status.into(), error: error.map(str::to_string), }; @@ -5760,7 +6021,7 @@ mod tests { "ok", None, ), - "test-org — oauth — Test User (user@example.test)", + "https://www.braintrust.dev — oauth — Test User (user@example.test)", ), ( verification( @@ -5775,7 +6036,7 @@ mod tests { ), ( verification("oauth", None, None, None, "expired", None), - "cross-org — oauth — token expired", + "https://www.braintrust.dev — oauth — token expired", ), ( verification( @@ -5868,13 +6129,16 @@ mod tests { #[tokio::test] async fn login_read_only_cached_project_id_and_org_uses_fast_path() { let env = TestEnv::new(Some("proj_123"), Some("test-org")).await; + let mut base = base_args_for_path_probe(Some("test-org")); + base.app_url = Some(spawn_api_key_login_server("test-org")); + set_global_config_urls(base.app_url.as_deref().unwrap(), None); let ctx = env - .login_read_only_probe(Some("test-org")) + .login_read_only_with_base(base) .await .expect("fast path should succeed"); assert_eq!(ctx.login.org_name().as_deref(), Some("test-org")); - assert_eq!(ctx.login.org_id().as_deref(), Some("")); + assert_eq!(ctx.login.org_id().as_deref(), Some("org_test")); assert_eq!(ctx.api_url, "not-a-valid-url"); } @@ -5884,16 +6148,6 @@ mod tests { assert_invalid_api_url(env.login_read_only_probe(None).await); } - #[tokio::test] - async fn login_read_only_cached_project_id_but_whitespace_org_is_cross_org() { - let env = TestEnv::new(Some("proj_123"), None).await; - let err = match env.login_read_only_probe(Some(" ")).await { - Ok(_) => panic!("whitespace org should be canonical cross-org"), - Err(err) => err, - }; - assert!(err.to_string().contains("concrete org")); - } - #[tokio::test] async fn login_read_only_whitespace_project_id_is_treated_as_not_cached() { let env = TestEnv::new(Some(" "), None).await; // has_cached_project_id => false @@ -5919,9 +6173,12 @@ mod tests { ]); save_profile_secret_plaintext("acme-profile", "acme-secret").expect("save acme secret"); save_profile_secret_plaintext("other-profile", "other-secret").expect("save other secret"); + set_global_config_urls("https://www.acme.example", Some("https://api.acme.example")); + let mut base = make_base(); + crate::config::apply_base_config(&mut base); let ctx = env - .login_read_only_with_base(make_base()) + .login_read_only_with_base(base) .await .expect("fast path should succeed with cfg org"); @@ -5939,6 +6196,7 @@ mod tests { base.org_name = Some("test-org".into()); let app_url = spawn_api_key_login_server("test-org"); base.app_url = Some(app_url.clone()); + set_global_config_urls(&app_url, None); let ctx = env .login_read_only_with_base(base) diff --git a/src/config/mod.rs b/src/config/mod.rs index 62bb4a7b..9d732d52 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -8,7 +8,7 @@ use std::{ use serde::{Deserialize, Serialize}; -use crate::args::BaseArgs; +use crate::args::{BaseArgs, DEFAULT_APP_URL}; use crate::ui::{print_command_status, CommandStatus}; mod get; @@ -19,32 +19,73 @@ mod set; #[serde(default)] pub struct Config { pub org: Option, + pub org_id: Option, pub project: Option, pub project_id: Option, + pub app_url: Option, + pub api_url: Option, #[serde(flatten)] pub extra: serde_json::Map, } -pub const KNOWN_KEYS: &[&str] = &["org", "project", "project_id"]; +pub const KNOWN_KEYS: &[&str] = &[ + "org", + "org_id", + "project", + "project_id", + "app_url", + "api_url", +]; impl Config { pub fn get_field(&self, key: &str) -> Option<&str> { match key { "org" => self.org.as_deref(), + "org_id" => self.org_id.as_deref(), "project" => self.project.as_deref(), "project_id" => self.project_id.as_deref(), + "app_url" => self.app_url.as_deref(), + "api_url" => self.api_url.as_deref(), _ => None, } } pub fn set_field(&mut self, key: &str, value: String) -> bool { match key { - "org" => self.org = Some(value), + "org" => { + let value = value.trim().to_string(); + if self.org.as_ref() != Some(&value) { + self.org_id = None; + self.project = None; + self.project_id = None; + } + self.org = (!value.is_empty()).then_some(value); + } + "org_id" => self.org_id = self.org.as_ref().map(|_| value), "project" => { self.project = Some(value); self.project_id = None; } "project_id" => self.project_id = Some(value), + "app_url" => { + let value = value.trim().to_string(); + let previous = self.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let next = if !value.is_empty() { + value.as_str() + } else { + DEFAULT_APP_URL + }; + if !urls_equal(previous, next) { + self.org = None; + self.org_id = None; + self.project = None; + self.project_id = None; + } + self.app_url = (!value.is_empty()).then_some(value); + } + "api_url" => { + self.api_url = trimmed_option(Some(&value)).map(str::to_string); + } _ => return false, } true @@ -52,12 +93,32 @@ impl Config { pub fn unset_field(&mut self, key: &str) -> bool { match key { - "org" => self.org = None, + "org" => { + self.org = None; + self.org_id = None; + self.project = None; + self.project_id = None; + } + "org_id" => self.org_id = None, "project" => { self.project = None; self.project_id = None; } "project_id" => self.project_id = None, + "app_url" => { + if self + .app_url + .as_deref() + .is_some_and(|url| !urls_equal(url, DEFAULT_APP_URL)) + { + self.org = None; + self.org_id = None; + self.project = None; + self.project_id = None; + } + self.app_url = None; + } + "api_url" => self.api_url = None, _ => return false, } true @@ -70,39 +131,105 @@ impl Config { .collect() } - pub(crate) fn set_context(&mut self, org: Option<&str>, project: Option<(&str, &str)>) { - self.org = org_option(org).map(str::to_string); + pub(crate) fn set_context( + &mut self, + org: (&str, &str), + project: Option<(&str, &str)>, + app_url: &str, + api_url: &str, + ) { + self.org = Some(org.0.trim().to_string()); + self.org_id = Some(org.1.trim().to_string()); (self.project, self.project_id) = project .map(|(name, id)| (name.to_string(), id.to_string())) .unzip(); + self.app_url = Some(app_url.to_string()); + self.api_url = Some(api_url.to_string()); } pub(crate) fn merge(&self, local: &Config) -> Config { let mut extra = self.extra.clone(); extra.extend(local.extra.clone()); - let global_id = self.project.as_ref().and(self.project_id.clone()); - let (org, project, project_id) = match (&local.org, &local.project) { + + let app_url = local.app_url.clone().or_else(|| self.app_url.clone()); + let api_url = local.api_url.clone().or_else(|| self.api_url.clone()); + let global_app = self.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let merged_app = app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let same_instance = urls_equal(global_app, merged_app); + let same_org = same_instance && local.org == self.org; + let global_project_id = self.project.as_ref().and(self.project_id.clone()); + + let (org, org_id, project, project_id) = match (&local.org, &local.project) { (Some(org), Some(project)) => ( Some(org.clone()), + local.org_id.clone(), Some(project.clone()), local.project_id.clone(), ), - (Some(org), None) if self.org.as_ref() == Some(org) => { - (Some(org.clone()), self.project.clone(), global_id) - } - (Some(org), None) => (Some(org.clone()), None, None), - (None, Some(project)) => (None, Some(project.clone()), local.project_id.clone()), - (None, None) => (self.org.clone(), self.project.clone(), global_id), + (Some(org), None) if same_org => ( + Some(org.clone()), + local.org_id.clone().or_else(|| self.org_id.clone()), + self.project.clone(), + global_project_id, + ), + (Some(org), None) => (Some(org.clone()), local.org_id.clone(), None, None), + (None, Some(project)) => (None, None, Some(project.clone()), local.project_id.clone()), + (None, None) if same_instance => ( + self.org.clone(), + self.org_id.clone(), + self.project.clone(), + global_project_id, + ), + (None, None) => (None, None, None, None), }; Config { org, + org_id, project, project_id, + app_url, + api_url, extra, } } } +pub(crate) fn urls_equal(left: &str, right: &str) -> bool { + left.trim().trim_end_matches('/') == right.trim().trim_end_matches('/') +} + +/// Apply config-file URL and org-ID fallbacks after clap has resolved CLI/env. +pub fn apply_base_config(base: &mut BaseArgs) { + let cfg = load().unwrap_or_default(); + apply_config_to_base(base, &cfg); +} + +fn apply_config_to_base(base: &mut BaseArgs, cfg: &Config) { + if base.app_url.is_none() { + base.app_url = cfg.app_url.clone(); + } + + if base.api_url.is_none() { + base.api_url = cfg.api_url.clone(); + } + + let effective_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let config_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let same_instance = urls_equal(effective_app, config_app); + if base.org_name_source.is_none() { + if base.org_name.is_none() && same_instance { + base.org_name = cfg.org.clone(); + } + if same_instance && base.org_name == cfg.org { + base.org_id = cfg.org_id.clone(); + } else { + base.org_id = None; + } + } else { + base.org_id = None; + } +} + pub fn global_config_dir() -> Result { if let Some(xdg) = env::var_os("XDG_CONFIG_HOME") { return Ok(PathBuf::from(xdg).join("bt")); @@ -142,13 +269,17 @@ pub fn load_file(path: &Path) -> Config { config.extra.remove("profile"); - // Fold a literal "cross-org" to the canonical "" marker on load. - if let Some(org) = config.org.as_deref() { - let normalized = normalize_org(org); - if normalized != org { - config.org = Some(normalized.to_string()); - } + config.org = trimmed_option(config.org.as_deref()).map(str::to_string); + config.org_id = config + .org + .as_ref() + .and(trimmed_option(config.org_id.as_deref()).map(str::to_string)); + if config.org.is_none() { + config.project = None; + config.project_id = None; } + config.app_url = trimmed_option(config.app_url.as_deref()).map(str::to_string); + config.api_url = trimmed_option(config.api_url.as_deref()).map(str::to_string); for key in config.extra.keys() { print_command_status( @@ -201,36 +332,20 @@ pub(crate) fn project_from_config_for_context( } fn config_matches_context(base: &BaseArgs, cfg: &Config, resolved_org: Option<&str>) -> bool { + let requested_app = base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let cfg_app = cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + if !urls_equal(requested_app, cfg_app) { + return false; + } + let cfg_org = org_option(cfg.org.as_deref()); let requested_org = org_option(resolved_org).or_else(|| org_option(base.org_name.as_deref())); requested_org.is_none_or(|resolved| cfg_org == Some(resolved)) } -/// Trim an org while preserving the empty cross-org marker. pub(crate) fn org_option(value: Option<&str>) -> Option<&str> { - value.map(str::trim) -} - -/// Human-facing spelling of the empty cross-org marker. -pub(crate) const CROSS_ORG_ALIAS: &str = "cross-org"; - -/// Trim an org and fold the [`CROSS_ORG_ALIAS`] to the canonical `""` marker. -pub(crate) fn normalize_org(value: &str) -> &str { - let trimmed = value.trim(); - if trimmed == CROSS_ORG_ALIAS { - "" - } else { - trimmed - } -} - -pub(crate) fn display_org(org: &str) -> &str { - if org.is_empty() { - CROSS_ORG_ALIAS - } else { - org - } + trimmed_option(value) } pub(crate) fn trimmed_option(value: Option<&str>) -> Option<&str> { @@ -451,14 +566,14 @@ enum ConfigCommands { }, /// Get a config value Get { - /// Config key (org, project, project_id) + /// Config key (org, org_id, project, project_id, app_url, api_url) key: String, #[command(flatten)] scope: ScopeArgs, }, /// Set a config value Set { - /// Config key (org, project, project_id) + /// Config key (org, org_id, project, project_id, app_url, api_url) key: String, /// Value to set value: String, @@ -467,7 +582,7 @@ enum ConfigCommands { }, /// Remove a config value Unset { - /// Config key (org, project, project_id) + /// Config key (org, org_id, project, project_id, app_url, api_url) key: String, #[command(flatten)] scope: ScopeArgs, @@ -550,6 +665,131 @@ mod tests { } } + #[test] + fn merge_inherits_context_only_within_the_same_instance() { + let global = Config { + org: Some("test-org".into()), + org_id: Some("org_test".into()), + project: Some("test-project".into()), + project_id: Some("proj_test".into()), + app_url: Some("https://www.example.test".into()), + api_url: Some("https://api.example.test".into()), + ..Default::default() + }; + let same_instance = Config { + app_url: Some("https://www.example.test/".into()), + api_url: Some("https://proxy.example.test".into()), + ..Default::default() + }; + let merged = global.merge(&same_instance); + assert_eq!(merged.org.as_deref(), Some("test-org")); + assert_eq!(merged.org_id.as_deref(), Some("org_test")); + assert_eq!(merged.project_id.as_deref(), Some("proj_test")); + assert_eq!( + merged.api_url.as_deref(), + Some("https://proxy.example.test") + ); + + let other_instance = Config { + app_url: Some("https://self-hosted.example.test".into()), + ..Default::default() + }; + let merged = global.merge(&other_instance); + assert_eq!(merged.org, None); + assert_eq!(merged.org_id, None); + assert_eq!(merged.project, None); + assert_eq!(merged.app_url, other_instance.app_url); + } + + #[test] + fn config_fills_urls_and_coupled_org_id_without_overriding_cli() { + let cfg = Config { + org: Some("config-org".into()), + org_id: Some("org_config".into()), + app_url: Some("https://www.example.test".into()), + api_url: Some("https://api.example.test".into()), + ..Default::default() + }; + let mut base = BaseArgs::default(); + apply_config_to_base(&mut base, &cfg); + assert_eq!(base.org_name.as_deref(), Some("config-org")); + assert_eq!(base.org_id.as_deref(), Some("org_config")); + assert_eq!(base.app_url, cfg.app_url); + assert_eq!(base.api_url, cfg.api_url); + + let mut base = BaseArgs { + org_name: Some("cli-org".into()), + org_name_source: Some(crate::args::ArgValueSource::CommandLine), + app_url: Some("https://cli.example.test".into()), + ..Default::default() + }; + apply_config_to_base(&mut base, &cfg); + assert_eq!(base.org_name.as_deref(), Some("cli-org")); + assert_eq!(base.org_id, None); + assert_eq!(base.app_url.as_deref(), Some("https://cli.example.test")); + assert_eq!(base.api_url, cfg.api_url); + + let mut same_instance = BaseArgs { + app_url: Some("https://www.example.test/".into()), + ..Default::default() + }; + apply_config_to_base(&mut same_instance, &cfg); + assert_eq!(same_instance.api_url, cfg.api_url); + + let mut other_instance = BaseArgs { + app_url: Some("https://other.example.test".into()), + ..Default::default() + }; + apply_config_to_base(&mut other_instance, &cfg); + assert_eq!(other_instance.org_name, None); + assert_eq!(other_instance.org_id, None); + } + + #[test] + fn configured_project_does_not_cross_instance_boundaries() { + let cfg = Config { + org: Some("test-org".into()), + project: Some("test-project".into()), + app_url: Some("https://www.example.test".into()), + ..Default::default() + }; + let matching = BaseArgs { + org_name: Some("test-org".into()), + app_url: Some("https://www.example.test/".into()), + ..Default::default() + }; + assert_eq!( + project_from_config_for_context(&matching, &cfg, Some("test-org")).as_deref(), + Some("test-project") + ); + + let other = BaseArgs { + app_url: Some("https://other.example.test".into()), + ..matching + }; + assert_eq!( + project_from_config_for_context(&other, &cfg, Some("test-org")), + None + ); + } + + #[test] + fn changing_app_url_clears_coupled_context() { + let mut cfg = Config { + org: Some("test-org".into()), + org_id: Some("org_test".into()), + project: Some("test-project".into()), + project_id: Some("proj_test".into()), + app_url: Some("https://www.example.test".into()), + ..Default::default() + }; + assert!(cfg.set_field("app_url", "https://other.example.test".into())); + assert_eq!(cfg.org, None); + assert_eq!(cfg.org_id, None); + assert_eq!(cfg.project, None); + assert_eq!(cfg.project_id, None); + } + #[test] fn scope_labels_are_plain_text() { let labels = scope_labels( @@ -564,8 +804,8 @@ mod tests { fn option_helpers_handle_empty_values() { for (input, org, trimmed) in [ (None, None, None), - (Some(""), Some(""), None), - (Some(" "), Some(""), None), + (Some(""), None, None), + (Some(" "), None, None), (Some("test-org"), Some("test-org"), Some("test-org")), ] { assert_eq!(org_option(input), org); @@ -573,12 +813,16 @@ mod tests { } let mut cfg = Config::default(); - cfg.set_context(Some(" test-org "), Some(("test-project", "proj_test"))); + cfg.set_context( + ("test-org", "org_test"), + Some(("test-project", "proj_test")), + "https://www.example.test", + "https://api.example.test", + ); assert_eq!(cfg.org.as_deref(), Some("test-org")); + assert_eq!(cfg.org_id.as_deref(), Some("org_test")); assert_eq!(cfg.project.as_deref(), Some("test-project")); assert_eq!(cfg.project_id.as_deref(), Some("proj_test")); - cfg.set_context(Some(""), None); - assert_eq!((cfg.org.as_deref(), cfg.project), (Some(""), None)); } fn base_args() -> BaseArgs { @@ -679,21 +923,19 @@ mod tests { } #[test] - fn load_folds_cross_org_alias_to_empty_marker() { + fn load_purges_obsolete_empty_org_context() { let tmp = TempDir::new().unwrap(); let path = tmp.path().join("config.json"); - // Literal "cross-org" must load identically to the "" marker. - for spelling in [ - r#"{"org":"cross-org"}"#, - r#"{"org":" cross-org "}"#, - r#"{"org":""}"#, - ] { - fs::write(&path, spelling).unwrap(); - assert_eq!(load_file(&path).org.as_deref(), Some(""), "{spelling}"); - } - - fs::write(&path, r#"{"org":"test-org"}"#).unwrap(); - assert_eq!(load_file(&path).org.as_deref(), Some("test-org")); + fs::write( + &path, + r#"{"org":"","org_id":"org_old","project":"old","project_id":"proj_old"}"#, + ) + .unwrap(); + let loaded = load_file(&path); + assert_eq!(loaded.org, None); + assert_eq!(loaded.org_id, None); + assert_eq!(loaded.project, None); + assert_eq!(loaded.project_id, None); } #[test] diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index 341ec84d..7ee4bec8 100644 --- a/src/experiments/mod.rs +++ b/src/experiments/mod.rs @@ -219,11 +219,7 @@ fn apply_experiment_url_hints_to_base( } } - let has_org_override = base - .org_name - .as_deref() - .map(str::trim) - .is_some_and(|v| !v.is_empty()); + let has_org_override = base.org_name_source.is_some(); if !has_org_override { if let Some(org) = parsed_url .org @@ -232,6 +228,7 @@ fn apply_experiment_url_hints_to_base( .filter(|v| !v.is_empty()) { base.org_name = Some(org.to_string()); + base.org_id = None; } } @@ -377,6 +374,32 @@ mod tests { ); } + #[test] + fn comparison_url_org_overrides_config_but_not_cli() { + let parsed = ParsedExperimentCompareUrl { + org: Some("url-org".to_string()), + project: None, + base_experiment: None, + comparison_experiment: None, + }; + let config_base = BaseArgs { + org_name: Some("config-org".to_string()), + org_id: Some("org_config".to_string()), + ..Default::default() + }; + let updated = apply_experiment_url_hints_to_base(config_base, Some(&parsed)); + assert_eq!(updated.org_name.as_deref(), Some("url-org")); + assert_eq!(updated.org_id, None); + + let cli_base = BaseArgs { + org_name: Some("cli-org".to_string()), + org_name_source: Some(crate::args::ArgValueSource::CommandLine), + ..Default::default() + }; + let updated = apply_experiment_url_hints_to_base(cli_base, Some(&parsed)); + assert_eq!(updated.org_name.as_deref(), Some("cli-org")); + } + #[test] fn compare_startup_url_uses_url_like_positional_arg() { let args = ExperimentsArgs { diff --git a/src/init.rs b/src/init.rs index 9449495e..00bd6e65 100644 --- a/src/init.rs +++ b/src/init.rs @@ -1,12 +1,10 @@ -use anyhow::{bail, Context, Result}; +use anyhow::{Context, Result}; use clap::Args; use crate::{ - args::BaseArgs, - auth::{self, login}, - config, - http::ApiClient, - ui::{print_command_status, select_or_create_project, CommandStatus}, + args::{ArgValueSource, BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL}, + config, switch, + ui::{print_command_status, CommandStatus}, }; #[derive(Debug, Clone, Args)] @@ -33,37 +31,44 @@ pub struct InitArgs { pub async fn run(base: BaseArgs, args: InitArgs) -> Result<()> { let config_path = config::init_target(args.here, args.force)?; let current_cfg = config::load().unwrap_or_default(); - let mut login_base = base.clone(); - login_base.project = None; - login_base.project_source = None; - if login_base.org_name.is_none() - && !auth::select_saved_login(&mut login_base, current_cfg.org.as_deref(), false)? - { - bail!("no saved concrete-org login is available; run `bt auth login --org `"); - } - - let ctx = login(&login_base).await?; - let client = ApiClient::new(&ctx)?; - let org = client.org_name().to_string(); - if org.is_empty() { - bail!( - "cross-org mode has no project; `bt init` is project-scoped. Rerun with --org --project " - ); - } - - let project = select_or_create_project( - &client, + let requested_org = matches!( + base.org_name_source, + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable) + ) + .then(|| base.org_name.as_deref()) + .flatten(); + let (instance, org, project) = switch::select_context( + &base, + requested_org, base.project.as_deref(), - None, + ¤t_cfg, Some("Link to project"), ) .await?; - // Load any existing file (only reachable via --force) so unknown passthrough - // keys are preserved, matching switch/config-set/post-login writers. + let api_url = if matches!( + base.api_url_source, + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable) + ) { + base.api_url.clone() + } else { + org.api_url.clone().or_else(|| { + config::urls_equal( + current_cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + &instance.app_url, + ) + .then(|| current_cfg.api_url.clone()) + .flatten() + }) + } + .unwrap_or_else(|| DEFAULT_API_URL.to_string()); + + // With --force, preserve unknown passthrough keys from the old file. let mut cfg = config::load_file(&config_path); cfg.set_context( - Some(&org), + (org.name.as_str(), org.id.as_str()), Some((project.name.as_str(), project.id.as_str())), + &instance.app_url, + &api_url, ); config::save_file(&config_path, &cfg).with_context(|| { @@ -77,16 +82,19 @@ pub async fn run(base: BaseArgs, args: InitArgs) -> Result<()> { let payload = serde_json::json!({ "initialized": true, "status": "created", - "org": org, + "org": org.name, + "org_id": org.id, "project": project.name, "project_id": project.id, + "app_url": instance.app_url, + "api_url": api_url, "path": config_path.display().to_string(), }); println!("{}", serde_json::to_string(&payload)?); } else { print_command_status( CommandStatus::Success, - &format!("Project linked to {org}/{}", project.name), + &format!("Project linked to {}/{}", org.name, project.name), ); print_command_status( CommandStatus::Success, diff --git a/src/main.rs b/src/main.rs index 035cc806..0fb11065 100644 --- a/src/main.rs +++ b/src/main.rs @@ -58,7 +58,7 @@ const HELP_TEMPLATE: &str = "\ Core init Initialize .bt config directory and files auth Authenticate bt with Braintrust - switch Switch org and project context + switch Switch instance, org, and project context view View logs, traces, and spans Projects & resources @@ -161,7 +161,7 @@ enum Commands { Sync(CLIArgs), /// Local utility commands Util(CLIArgs), - /// Switch org and project context + /// Switch instance, org, and project context Switch(CLIArgs), /// Show current org and project context Status(CLIArgs), @@ -296,6 +296,7 @@ fn try_main() -> Result<()> { let matches = Cli::command().get_matches_from(&argv); let mut cli = Cli::from_arg_matches(&matches).expect("clap matches should parse"); apply_base_arg_sources(&matches, cli.command.base_mut()); + config::apply_base_config(cli.command.base_mut()); apply_base_output_defaults(&mut cli.command); configure_output(cli.command.base()); apply_runtime_env_overrides(cli.command.base()); @@ -349,6 +350,8 @@ fn apply_base_arg_sources(matches: &ArgMatches, base: &mut BaseArgs) { base.org_name_source = find_value_source(matches, "org_name").and_then(map_value_source); base.project_source = find_value_source(matches, "project").and_then(map_value_source); base.api_key_source = find_value_source(matches, "api_key").and_then(map_value_source); + base.api_url_source = find_value_source(matches, "api_url").and_then(map_value_source); + base.app_url_source = find_value_source(matches, "app_url").and_then(map_value_source); } fn apply_base_output_defaults(command: &mut Commands) { diff --git a/src/setup/mod.rs b/src/setup/mod.rs index 17937d3e..bea36fa1 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -1179,6 +1179,10 @@ fn should_print_agent_selection_intro( fn apply_setup_config_fallbacks(base: &mut BaseArgs) { let cfg = config::load().unwrap_or_default(); + let same_instance = config::urls_equal( + base.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + ); if base .org_name @@ -1186,11 +1190,16 @@ fn apply_setup_config_fallbacks(base: &mut BaseArgs) { .map(str::trim) .is_none_or(str::is_empty) { - base.org_name = cfg - .org - .as_deref() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()); + if same_instance { + base.org_name = cfg + .org + .as_deref() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + base.org_id = base.org_name.as_ref().and(cfg.org_id.clone()); + } else { + base.org_id = None; + } } if base @@ -1692,7 +1701,8 @@ async fn ensure_org_or_setup_browser_auth( auth_base.org_name = Some(org_name.to_string()); let has_saved_auth_for_org = profiles .iter() - .any(|profile| profile.org_name.as_deref() == Some(org_name)); + .any(|profile| profile.org_name.as_deref() == Some(org_name)) + || auth::has_oauth_login_for_instance(&auth_base)?; if has_saved_auth_for_org { match auth::login(&auth_base).await { diff --git a/src/status.rs b/src/status.rs index b8a8ccbf..b6738994 100644 --- a/src/status.rs +++ b/src/status.rs @@ -2,7 +2,7 @@ use anyhow::Result; use clap::Args; use serde::Serialize; -use crate::args::BaseArgs; +use crate::args::{BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL}; use crate::auth; use crate::config; @@ -18,7 +18,11 @@ pub struct StatusArgs {} #[derive(Serialize)] struct StatusOutput { org: Option, + org_id: Option, project: Option, + project_id: Option, + app_url: Option, + api_url: Option, #[serde(skip_serializing_if = "Option::is_none")] user_name: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -69,11 +73,38 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { project = config::project_from_config_for_context(&base, &merged_cfg, org.as_deref()); } - let display_org = org.as_deref().map(config::display_org); + let org_id = if base.org_name_source.is_some() { + None + } else { + base.org_id.clone() + }; + let configured_project = + config::project_from_config_for_context(&base, &merged_cfg, org.as_deref()); + let project_id = (base.project_source.is_none() + && configured_project.is_some() + && configured_project == project) + .then(|| merged_cfg.project_id.clone()) + .flatten(); + let app_url = Some( + base.app_url + .clone() + .or_else(|| merged_cfg.app_url.clone()) + .unwrap_or_else(|| DEFAULT_APP_URL.to_string()), + ); + let api_url = Some( + base.api_url + .clone() + .or_else(|| merged_cfg.api_url.clone()) + .unwrap_or_else(|| DEFAULT_API_URL.to_string()), + ); if base.json { let output = StatusOutput { - org: display_org.map(str::to_string), + org: org.clone(), + org_id, project, + project_id, + app_url, + api_url, user_name: auth_info.as_ref().and_then(|p| p.user_name.clone()), user_email: auth_info.as_ref().and_then(|p| p.email.clone()), api_key_hint: auth_info.as_ref().and_then(|p| p.api_key_hint.clone()), @@ -85,8 +116,12 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { } if base.verbose { - println!("org: {}", display_org.unwrap_or("(unset)")); + println!("org: {}", org.as_deref().unwrap_or("(unset)")); + println!("org_id: {}", org_id.as_deref().unwrap_or("(unset)")); println!("project: {}", project.as_deref().unwrap_or("(unset)")); + println!("project_id: {}", project_id.as_deref().unwrap_or("(unset)")); + println!("app_url: {}", app_url.as_deref().unwrap_or(DEFAULT_APP_URL)); + println!("api_url: {}", api_url.as_deref().unwrap_or(DEFAULT_API_URL)); if let Some(ref p) = auth_info { println!("auth: {}", format_auth(p)); } @@ -94,26 +129,18 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { println!("source: {src}"); } } else { - // Plain one-liner. Always surface the active auth — even when no org is - // configured (an env-only API key, or a cross-org OAuth login) — instead - // of hiding it behind --verbose. - let cross_org_oauth = auth_info - .as_ref() - .is_some_and(|p| p.auth_method == "oauth" && p.org_name.is_none()); - let header = match display_org { - Some("cross-org") => "cross-org".to_string(), + let header = match org.as_deref() { Some(org) => match project.as_deref() { Some(project) => format!("{org}/{project}"), None => org.to_string(), }, - None if cross_org_oauth => "cross-org".to_string(), None if auth_info.is_some() => "No default org".to_string(), None => "No org/project configured. Run `bt switch` to set one.".to_string(), }; println!("{header}"); match &auth_info { Some(p) => println!(" auth: {}", format_auth(p)), - None if display_org.is_some() => println!(" auth: (none)"), + None if org.is_some() => println!(" auth: (none)"), None => {} } } @@ -127,6 +154,10 @@ pub(crate) struct ConfigOverrides { env_org: Option, cli_project: Option, env_project: Option, + cli_app_url: Option, + env_app_url: Option, + cli_api_url: Option, + env_api_url: Option, } impl ConfigOverrides { @@ -143,12 +174,26 @@ impl ConfigOverrides { Some(ArgValueSource::EnvVariable) => (None, base.project.clone()), None => (None, None), }; + let (cli_app_url, env_app_url) = match base.app_url_source { + Some(ArgValueSource::CommandLine) => (base.app_url.clone(), None), + Some(ArgValueSource::EnvVariable) => (None, base.app_url.clone()), + None => (None, None), + }; + let (cli_api_url, env_api_url) = match base.api_url_source { + Some(ArgValueSource::CommandLine) => (base.api_url.clone(), None), + Some(ArgValueSource::EnvVariable) => (None, base.api_url.clone()), + None => (None, None), + }; Self { cli_org, env_org, cli_project, env_project, + cli_app_url, + env_app_url, + cli_api_url, + env_api_url, } } } @@ -166,28 +211,48 @@ pub(crate) fn resolve_config( env_org, cli_project, env_project, + cli_app_url, + env_app_url, + cli_api_url, + env_api_url, } = overrides; - // `Some("")` is the canonical cross-org marker for both CLI and env - // sources, so org overrides must not filter empty strings. let env_project = env_project.filter(|s| !s.is_empty()); let merged = global.merge(local); - let org = cli_org - .clone() - .or_else(|| env_org.clone()) - .or_else(|| merged.org.clone()); + let app_override = cli_app_url.as_deref().or(env_app_url.as_deref()); + let config_app = merged.app_url.as_deref().unwrap_or(DEFAULT_APP_URL); + let same_instance = app_override.is_none_or(|app| config::urls_equal(app, config_app)); + let config_org = same_instance.then(|| merged.org.clone()).flatten(); + let config_project = same_instance.then(|| merged.project.clone()).flatten(); + let org = cli_org.clone().or_else(|| env_org.clone()).or(config_org); let project = cli_project .clone() .or_else(|| env_project.clone()) - .or_else(|| merged.project.clone()); + .or(config_project); - let source = if cli_org.is_some() || cli_project.is_some() { + let source = if cli_org.is_some() + || cli_project.is_some() + || cli_app_url.is_some() + || cli_api_url.is_some() + { Some("cli".to_string()) - } else if env_org.is_some() || env_project.is_some() { + } else if env_org.is_some() + || env_project.is_some() + || env_app_url.is_some() + || env_api_url.is_some() + { Some("env".to_string()) - } else if local.org.is_some() || local.project.is_some() { + } else if local.org.is_some() + || local.project.is_some() + || local.app_url.is_some() + || local.api_url.is_some() + { local_path.as_ref().map(|p| p.display().to_string()) - } else if global.org.is_some() || global.project.is_some() { + } else if global.org.is_some() + || global.project.is_some() + || global.app_url.is_some() + || global.api_url.is_some() + { global_path.as_ref().map(|p| p.display().to_string()) } else { None @@ -268,6 +333,16 @@ mod tests { config(None, None), (None, None, None), ), + ( + "app override does not inherit another instance's context", + ConfigOverrides { + cli_app_url: s("https://other.example.test"), + ..Default::default() + }, + both(), + config(None, None), + (None, None, Some("cli")), + ), ( "mixed cli/local", ConfigOverrides { @@ -285,23 +360,6 @@ mod tests { config(None, Some("local-proj")), (None, Some("local-proj"), Some("/project/.bt/config.json")), ), - ( - "local cross-org", - ConfigOverrides::default(), - both(), - config(Some(""), None), - (Some(""), None, Some("/project/.bt/config.json")), - ), - ( - "env cross-org", - ConfigOverrides { - env_org: s(""), - ..Default::default() - }, - both(), - config(None, None), - (Some(""), Some("global-proj"), Some("env")), - ), ]; let local_path = Some(PathBuf::from("/project/.bt/config.json")); let global_path = Some(PathBuf::from("/home/.bt/config.json")); diff --git a/src/switch.rs b/src/switch.rs index b66647c1..5ae73538 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -1,8 +1,8 @@ use anyhow::{bail, Context, Result}; use clap::Args; -use crate::args::BaseArgs; -use crate::auth::{self, login}; +use crate::args::{ArgValueSource, BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL}; +use crate::auth::{self, login, AvailableInstance, AvailableOrg}; use crate::config; use crate::http::ApiClient; use crate::ui::{can_prompt, print_command_status, select_or_create_project, CommandStatus}; @@ -13,7 +13,6 @@ Examples: bt switch bt switch test-project bt switch test-org/test-project - bt switch --org cross-org ")] pub struct SwitchArgs { #[command(flatten)] @@ -28,87 +27,207 @@ impl SwitchArgs { fn resolve_target(&self, base: &BaseArgs) -> (Option, Option) { let (pos_org, pos_project) = match &self.target { None => (None, None), - Some(t) if t.contains('/') => { - let parts: Vec<&str> = t.splitn(2, '/').collect(); - let o = (!parts[0].is_empty()).then(|| config::normalize_org(parts[0]).to_string()); - let p = (!parts[1].is_empty()).then(|| parts[1].to_string()); - (o, p) + Some(target) if target.contains('/') => { + let parts: Vec<&str> = target.splitn(2, '/').collect(); + let org = (!parts[0].trim().is_empty()).then(|| parts[0].trim().to_string()); + let project = (!parts[1].trim().is_empty()).then(|| parts[1].trim().to_string()); + (org, project) } - Some(t) => (None, Some(t.clone())), + Some(target) => (None, Some(target.clone())), }; - let org = base.org_name.clone().or(pos_org); - let project = base.project.clone().or(pos_project); - - (org, project) + ( + base.org_name + .as_ref() + .filter(|_| { + matches!( + base.org_name_source, + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable) + ) + }) + .cloned() + .or(pos_org), + base.project.clone().or(pos_project), + ) } } -pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { - args.scope.preflight(can_prompt())?; - let current_cfg = if args.scope.global { - config::load_global().unwrap_or_default() - } else { - config::load().unwrap_or_default() - }; - let (resolved_org, resolved_project) = args.resolve_target(&base); - let bare_switch = resolved_org.is_none() && resolved_project.is_none(); +fn find_org<'a>(orgs: &'a [AvailableOrg], identifier: &str) -> Option<&'a AvailableOrg> { + orgs.iter() + .find(|org| org.id == identifier || org.name == identifier) + .or_else(|| { + let lowered = identifier.to_ascii_lowercase(); + orgs.iter() + .find(|org| org.name.to_ascii_lowercase() == lowered) + }) +} - let mut login_base = base.clone(); - login_base.org_name = resolved_org.clone(); - login_base.project = None; - login_base.project_source = None; +fn select_instance( + instances: &[AvailableInstance], + current_app_url: Option<&str>, +) -> Result { + match instances { + [] => bail!("no saved auth logins found; run `bt auth login` to create one"), + [instance] => Ok(instance.clone()), + _ if can_prompt() => { + let labels = instances + .iter() + .map(|instance| instance.app_url.as_str()) + .collect::>(); + let default = current_app_url + .and_then(|current| { + instances + .iter() + .position(|instance| config::urls_equal(&instance.app_url, current)) + }) + .unwrap_or(0); + let idx = crate::ui::fuzzy_select("Select Braintrust instance", &labels, default)?; + Ok(instances[idx].clone()) + } + _ => bail!( + "multiple Braintrust instances are available; pass --app-url or rerun interactively" + ), + } +} - if login_base.org_name.is_none() && !bare_switch { - login_base.org_name = current_cfg.org.clone(); +fn select_org( + orgs: &[AvailableOrg], + requested: Option<&str>, + current: Option<&str>, +) -> Result { + if let Some(requested) = requested { + return find_org(orgs, requested) + .cloned() + .ok_or_else(|| anyhow::anyhow!("organization '{requested}' is not available")); } - if login_base.org_name.is_none() - && !auth::select_saved_login(&mut login_base, current_cfg.org.as_deref(), true)? - { - bail!("no saved auth logins found; run `bt auth login` to create one"); + match orgs { + [] => bail!("no organizations are available for the selected Braintrust instance"), + [org] => Ok(org.clone()), + _ if can_prompt() => { + let labels = orgs.iter().map(|org| org.name.as_str()).collect::>(); + let default = current + .and_then(|current| { + orgs.iter() + .position(|org| org.id == current || org.name == current) + }) + .unwrap_or(0); + let idx = crate::ui::fuzzy_select("Select organization", &labels, default)?; + Ok(orgs[idx].clone()) + } + _ => bail!("organization selection requires an interactive terminal; pass --org "), } +} - if login_base.org_name.as_deref() == Some("") && resolved_project.is_some() { - bail!( - "cross-org mode cannot have a default project; rerun with --org --project " - ); - } +pub(crate) async fn select_context( + base: &BaseArgs, + requested_org: Option<&str>, + requested_project: Option<&str>, + current_cfg: &config::Config, + project_prompt: Option<&str>, +) -> Result<( + AvailableInstance, + AvailableOrg, + crate::projects::api::Project, +)> { + let instances = auth::available_instances(base)?; + let instance = select_instance(&instances, current_cfg.app_url.as_deref())?; + let orgs = auth::available_orgs_for_instance(base, &instance.app_url).await?; + let same_current_instance = config::urls_equal( + current_cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + &instance.app_url, + ); + let current_org = same_current_instance + .then(|| current_cfg.org_id.as_deref().or(current_cfg.org.as_deref())) + .flatten(); + let org = select_org(&orgs, requested_org, current_org)?; + + let explicit_api_url = matches!( + base.api_url_source, + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable) + ) + .then(|| base.api_url.clone()) + .flatten(); + let current_api_url = same_current_instance + .then(|| current_cfg.api_url.clone()) + .flatten(); + let api_url = explicit_api_url + .or_else(|| org.api_url.clone()) + .or(current_api_url) + .unwrap_or_else(|| DEFAULT_API_URL.to_string()); + + let mut login_base = base.clone(); + login_base.app_url = Some(instance.app_url.clone()); + login_base.api_url = Some(api_url); + login_base.org_name = Some(org.name.clone()); + login_base.org_id = Some(org.id.clone()); + login_base.project = None; + login_base.project_source = None; let ctx = login(&login_base).await?; let client = ApiClient::new(&ctx)?; - let org_name = client.org_name().to_string(); + let current_project = (same_current_instance + && current_cfg.org_id.as_deref() == Some(org.id.as_str())) + .then_some(current_cfg.project.as_deref()) + .flatten(); + let project = + select_or_create_project(&client, requested_project, current_project, project_prompt) + .await?; - let project = if org_name.is_empty() { - None + Ok((instance, org, project)) +} + +pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { + args.scope.preflight(can_prompt())?; + let current_cfg = if args.scope.global { + config::load_global().unwrap_or_default() } else { - Some( - select_or_create_project( - &client, - resolved_project.as_deref(), - current_cfg.project.as_deref(), - None, - ) - .await?, - ) + config::load().unwrap_or_default() }; + let (requested_org, requested_project) = args.resolve_target(&base); + let (instance, org, project) = select_context( + &base, + requested_org.as_deref(), + requested_project.as_deref(), + ¤t_cfg, + None, + ) + .await?; + let api_url = if matches!( + base.api_url_source, + Some(ArgValueSource::CommandLine | ArgValueSource::EnvVariable) + ) { + base.api_url.clone() + } else { + org.api_url.clone().or_else(|| { + config::urls_equal( + current_cfg.app_url.as_deref().unwrap_or(DEFAULT_APP_URL), + &instance.app_url, + ) + .then(|| current_cfg.api_url.clone()) + .flatten() + }) + } + .unwrap_or_else(|| DEFAULT_API_URL.to_string()); - // Scope is prompted last, after org and project. let (path, scope) = args.scope.resolve(can_prompt(), "Save to")?; let mut cfg = config::load_file(&path); cfg.set_context( - Some(&org_name), - project - .as_ref() - .map(|project| (project.name.as_str(), project.id.as_str())), + (org.name.as_str(), org.id.as_str()), + Some((project.name.as_str(), project.id.as_str())), + &instance.app_url, + &api_url, ); config::save_file(&path, &cfg) .with_context(|| format!("Could not save config to {}", path.display()))?; if base.json { let payload = serde_json::json!({ - "org": config::display_org(&org_name), - "project": project.as_ref().map(|p| p.name.clone()), - "project_id": project.as_ref().map(|p| p.id.clone()), + "org": org.name, + "org_id": org.id, + "project": project.name, + "project_id": project.id, + "app_url": instance.app_url, + "api_url": api_url, "scope": scope, "path": path.display().to_string(), }); @@ -116,11 +235,10 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { return Ok(()); } - let display = project - .as_ref() - .map(|project| format!("{org_name}/{}", project.name)) - .unwrap_or_else(|| config::display_org(&org_name).to_string()); - print_command_status(CommandStatus::Success, &format!("Switched to {display}")); + print_command_status( + CommandStatus::Success, + &format!("Switched to {}/{}", org.name, project.name), + ); if base.verbose { eprintln!("Wrote to {}", path.display()); } @@ -131,80 +249,27 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { #[cfg(test)] mod tests { use super::*; + #[test] fn resolve_target_combines_positionals_and_flags() { - for (target, org, project, expected) in [ - (None, None, None, (None, None)), - ( - Some("test-org/test-project"), - None, - None, - (Some("test-org"), Some("test-project")), - ), - ( - Some("test-project"), - None, - None, - (None, Some("test-project")), - ), - ( - Some("/test-project"), - None, - None, - (None, Some("test-project")), - ), - (Some("test-org/"), None, None, (Some("test-org"), None)), - // Positional "cross-org" folds to the "" marker, matching --org. - ( - Some("cross-org/test-project"), - None, - None, - (Some(""), Some("test-project")), - ), - (Some("cross-org/"), None, None, (Some(""), None)), - (None, Some("test-org"), None, (Some("test-org"), None)), - ( - None, - None, - Some("test-project"), - (None, Some("test-project")), - ), - ( - None, - Some("test-org"), - Some("test-project"), - (Some("test-org"), Some("test-project")), - ), - ( - Some("old-org/old-project"), - None, - Some("test-project"), - (Some("old-org"), Some("test-project")), - ), - ( - Some("old-project"), - Some("test-org"), - None, - (Some("test-org"), Some("old-project")), - ), - ( - Some("old-org/old-project"), - Some("test-org"), - Some("test-project"), - (Some("test-org"), Some("test-project")), - ), - ] { - let args = SwitchArgs { - scope: config::ScopeArgs::default(), - target: target.map(str::to_string), - }; - let base = BaseArgs { - org_name: org.map(str::to_string), - project: project.map(str::to_string), - ..Default::default() - }; - let actual = args.resolve_target(&base); - assert_eq!((actual.0.as_deref(), actual.1.as_deref()), expected); - } + let args = SwitchArgs { + scope: config::ScopeArgs::default(), + target: Some("test-org/test-project".to_string()), + }; + let base = BaseArgs::default(); + let actual = args.resolve_target(&base); + assert_eq!(actual.0.as_deref(), Some("test-org")); + assert_eq!(actual.1.as_deref(), Some("test-project")); + } + + #[test] + fn find_org_matches_name_id_and_case() { + let orgs = vec![AvailableOrg { + id: "org_test".to_string(), + name: "test-org".to_string(), + api_url: None, + }]; + assert!(find_org(&orgs, "org_test").is_some()); + assert!(find_org(&orgs, "TEST-ORG").is_some()); } } diff --git a/src/traces.rs b/src/traces.rs index 97a881c9..c5250124 100644 --- a/src/traces.rs +++ b/src/traces.rs @@ -5728,14 +5728,11 @@ fn apply_url_hints_to_base(mut base: BaseArgs, parsed_url: Option<&ParsedTraceUr return base; }; - let has_org_override = base - .org_name - .as_deref() - .map(str::trim) - .is_some_and(|v| !v.is_empty()); + let has_org_override = base.org_name_source.is_some(); if !has_org_override { base.org_name = Some(url_org.to_string()); + base.org_id = None; } base } @@ -6925,19 +6922,23 @@ mod tests { } #[test] - fn apply_url_hints_infers_org_from_url() { - let base = base_args(); + fn apply_url_hints_override_config_org_and_clear_its_id() { + let mut base = base_args(); + base.org_name = Some("config-org".to_string()); + base.org_id = Some("org_config".to_string()); let parsed = parsed_url_with_org("Lovable"); let updated = apply_url_hints_for_test(base, Some(&parsed)); assert_eq!(updated.org_name.as_deref(), Some("Lovable")); + assert_eq!(updated.org_id, None); } #[test] fn apply_url_hints_preserves_explicit_org() { let mut base = base_args(); base.org_name = Some("explicit-org".to_string()); + base.org_name_source = Some(crate::args::ArgValueSource::CommandLine); let parsed = parsed_url_with_org("Lovable"); let updated = apply_url_hints_for_test(base, Some(&parsed)); diff --git a/tests/eval_dev_server.rs b/tests/eval_dev_server.rs index 227221d7..3a5095f9 100644 --- a/tests/eval_dev_server.rs +++ b/tests/eval_dev_server.rs @@ -71,7 +71,7 @@ fn start_mock_auth_server() -> (u16, thread::JoinHandle<()>) { .expect("set mock listener blocking"); let handle = thread::spawn(move || { - let response_body = r#"{"org_info": [{"name": "test-org"}]}"#; + let response_body = r#"{"org_info": [{"id": "org_test", "name": "test-org"}]}"#; let http_response = format!( "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", response_body.len(), From 177488bdcd11f4de9aa3d6c43e7cddfb09d9419c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 9 Jul 2026 22:22:15 -0700 Subject: [PATCH 17/24] draft --- src/main.rs | 7 +++++++ src/sql.rs | 13 +++++++++++++ src/traces/waterfall.rs | 12 ++---------- src/utils/mod.rs | 2 ++ 4 files changed, 24 insertions(+), 10 deletions(-) diff --git a/src/main.rs b/src/main.rs index 0fb11065..cef9f2c1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,7 @@ mod args; mod auth; #[allow(dead_code)] mod config; +mod cost; mod datasets; mod env; #[cfg(unix)] @@ -75,6 +76,7 @@ Data & evaluation datasets Manage datasets eval Run eval files sql Run SQL queries against Braintrust + cost Estimate LLM cost for Braintrust resources sync Synchronize project logs between Braintrust and local NDJSON files Additional @@ -129,6 +131,8 @@ enum Commands { Docs(CLIArgs), /// Run SQL queries against Braintrust Sql(CLIArgs), + /// Estimate LLM cost for Braintrust resources + Cost(CLIArgs), /// Authenticate bt with Braintrust Auth(CLIArgs), /// View logs, traces, and spans @@ -176,6 +180,7 @@ impl Commands { Commands::Setup(cmd) => &cmd.base, Commands::Docs(cmd) => &cmd.base, Commands::Sql(cmd) => &cmd.base, + Commands::Cost(cmd) => &cmd.base, Commands::Auth(cmd) => &cmd.base, Commands::View(cmd) => &cmd.base, #[cfg(unix)] @@ -203,6 +208,7 @@ impl Commands { Commands::Setup(cmd) => &mut cmd.base, Commands::Docs(cmd) => &mut cmd.base, Commands::Sql(cmd) => &mut cmd.base, + Commands::Cost(cmd) => &mut cmd.base, Commands::Auth(cmd) => &mut cmd.base, Commands::View(cmd) => &mut cmd.base, #[cfg(unix)] @@ -311,6 +317,7 @@ fn try_main() -> Result<()> { Commands::View(cmd) => traces::run(cmd.base, cmd.args).await?, Commands::Init(cmd) => init::run(cmd.base, cmd.args).await?, Commands::Sql(cmd) => sql::run(cmd.base, cmd.args).await?, + Commands::Cost(cmd) => cost::run(cmd.base, cmd.args).await?, Commands::Setup(cmd) => setup::run_setup_top(cmd.base, cmd.args).await?, Commands::Docs(cmd) => setup::run_docs_top(cmd.base, cmd.args).await?, #[cfg(unix)] diff --git a/src/sql.rs b/src/sql.rs index 208535a7..8b7f974d 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -437,6 +437,19 @@ fn format_response(response: &SqlResponse, json_output: bool) -> Result } } +/// Run a read-only BTQL/SQL query and return just the result rows. +/// +/// Thin wrapper over [`execute_query`] so other commands (for example +/// `bt cost`) can reuse the `/btql` request body and headers without +/// duplicating them. +pub(crate) async fn run_btql_rows( + client: &ApiClient, + query: &str, + lint_mode: &str, +) -> Result>> { + Ok(execute_query(client, query, lint_mode).await?.data) +} + async fn execute_query(client: &ApiClient, query: &str, lint_mode: &str) -> Result { let body = query_body(query, lint_mode); let headers = org_headers(client); diff --git a/src/traces/waterfall.rs b/src/traces/waterfall.rs index 374556f8..244365eb 100644 --- a/src/traces/waterfall.rs +++ b/src/traces/waterfall.rs @@ -3,6 +3,8 @@ use std::cmp::Ordering; use serde::Serialize; use serde_json::{Map, Value}; +use crate::utils::format_cost; + use super::{ build_span_entries, extract_duration_seconds, extract_model_name, extract_parent_span_id, extract_span_end_seconds, extract_span_name_and_type, extract_start_time, @@ -604,16 +606,6 @@ fn format_percent(value: f64) -> String { format!("{value:.1}%") } -fn format_cost(cost: f64) -> String { - if cost > 0.0 && cost < 0.001 { - "<$0.001".to_string() - } else if cost < 1.0 { - format!("${cost:.3}") - } else { - format!("${cost:.2}") - } -} - #[cfg(test)] mod tests { use serde_json::{json, Map}; diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 6c945ed9..a71fc735 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,5 +1,6 @@ mod app_url; mod duration; +mod format; mod fs_atomic; mod git; mod ids; @@ -10,6 +11,7 @@ mod shell; pub(crate) use app_url::{app_project_url, app_project_url_with_encoded_path}; pub use duration::parse_duration_to_seconds; +pub(crate) use format::format_cost; pub use fs_atomic::{write_bytes_atomic, write_text_atomic}; pub use git::GitRepo; pub(crate) use ids::new_uuid_id; From c4aac031caaf696a94afc1e34a57c5c593b69b42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Thu, 9 Jul 2026 22:25:11 -0700 Subject: [PATCH 18/24] draft --- src/cost/experiments.rs | 435 ++++++++++++++++++++++++++++++++++++++++ src/cost/mod.rs | 38 ++++ src/utils/format.rs | 38 ++++ 3 files changed, 511 insertions(+) create mode 100644 src/cost/experiments.rs create mode 100644 src/cost/mod.rs create mode 100644 src/utils/format.rs diff --git a/src/cost/experiments.rs b/src/cost/experiments.rs new file mode 100644 index 00000000..5f8eb3d1 --- /dev/null +++ b/src/cost/experiments.rs @@ -0,0 +1,435 @@ +use std::collections::HashMap; +use std::fmt::Write as _; + +use anyhow::Result; +use dialoguer::console; +use serde::Serialize; +use serde_json::{Map, Value}; + +use crate::{ + experiments::api::list_experiments, + sql::run_btql_rows, + ui::{apply_column_padding, header, print_with_pager, styled_table, truncate, with_spinner}, + utils::{format_cost, pluralize}, +}; + +use super::ResolvedContext; + +/// Per-experiment cost stats returned by the aggregate query. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +struct CostStats { + /// `SUM(estimated_cost())` — null when nothing could be priced. + cost: Option, + /// Spans with a non-null `estimated_cost()` (the ones feeding the sum). + priced_spans: u64, + /// Non-scorer spans that clearly involved an LLM (llm span or token + /// metrics) but could not be priced — the "we don't know" signal. + unpriced_llm_spans: u64, +} + +/// `estimated_cost()` returns NULL for spans it cannot price (no logged cost and +/// no model-registry match) and for scorer spans. Comparing priced vs. unpriced +/// LLM spans lets us report cost coverage instead of a silently-wrong `$0`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +enum Coverage { + /// Every LLM span was priced. + Full, + /// Some, but not all, LLM spans were priced. + Partial, + /// LLM spans exist but none could be priced (e.g. missing model pricing). + Unknown, + /// No priced spans and no detectable unpriced LLM activity — no cost data. + None, +} + +impl Coverage { + fn classify(priced_spans: u64, unpriced_llm_spans: u64) -> Self { + match (priced_spans, unpriced_llm_spans) { + (0, 0) => Coverage::None, + (0, _) => Coverage::Unknown, + (_, 0) => Coverage::Full, + (_, _) => Coverage::Partial, + } + } +} + +#[derive(Debug, Clone, Serialize)] +struct ExperimentCostRow { + id: String, + name: String, + created: Option, + /// Estimated cost in USD. `null` when we could not price anything. + cost: Option, + priced_spans: u64, + unpriced_llm_spans: u64, + coverage: Coverage, +} + +pub(crate) async fn run(ctx: &ResolvedContext, json: bool) -> Result<()> { + let project_name = &ctx.project.name; + let experiments = with_spinner( + "Loading experiments...", + list_experiments(&ctx.client, project_name), + ) + .await?; + + if experiments.is_empty() { + if json { + println!("[]"); + } else { + println!("No experiments found in {project_name}"); + } + return Ok(()); + } + + let experiment_ids: Vec = experiments.iter().map(|e| e.id.clone()).collect(); + let query = build_cost_query(&experiment_ids); + let result_rows = with_spinner( + "Estimating cost...", + run_btql_rows(&ctx.client, &query, "strict"), + ) + .await?; + + let stats_by_id = index_cost_rows(&result_rows); + + let mut rows: Vec = experiments + .iter() + .map(|exp| { + let stats = stats_by_id.get(&exp.id).copied().unwrap_or_default(); + let coverage = Coverage::classify(stats.priced_spans, stats.unpriced_llm_spans); + ExperimentCostRow { + id: exp.id.clone(), + name: exp.name.clone(), + created: exp.created.clone(), + cost: stats.cost, + priced_spans: stats.priced_spans, + unpriced_llm_spans: stats.unpriced_llm_spans, + coverage, + } + }) + .collect(); + + // Highest cost first; experiments with no known cost sort last, tie-broken by name. + rows.sort_by(|a, b| { + let a_key = a.cost.unwrap_or(f64::NEG_INFINITY); + let b_key = b.cost.unwrap_or(f64::NEG_INFINITY); + b_key + .partial_cmp(&a_key) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.name.cmp(&b.name)) + }); + + if json { + println!("{}", serde_json::to_string(&rows)?); + return Ok(()); + } + + print_table(ctx, &rows)?; + Ok(()) +} + +fn build_cost_query(experiment_ids: &[String]) -> String { + let id_list = experiment_ids + .iter() + .map(|id| format!("'{}'", id.replace('\'', "''"))) + .collect::>() + .join(", "); + + // `SUM(estimated_cost())` matches the web UI's experiment cost (no span-type + // filter; scorer spans are already unpriced by `estimated_cost()`). + // `unpriced_llm_spans` counts non-scorer spans that had LLM activity (llm + // span type or token metrics) yet could not be priced. + format!( + "SELECT \ + experiment_id, \ + SUM(estimated_cost()) AS cost, \ + COUNT(estimated_cost()) AS priced_spans, \ + COUNT(CASE \ + WHEN (span_attributes.purpose IS NULL OR span_attributes.purpose != 'scorer') \ + AND estimated_cost() IS NULL \ + AND (span_attributes.type = 'llm' \ + OR metrics.prompt_tokens IS NOT NULL \ + OR metrics.completion_tokens IS NOT NULL) \ + THEN 1 END) AS unpriced_llm_spans \ + FROM experiment({id_list}) \ + GROUP BY experiment_id" + ) +} + +fn index_cost_rows(result_rows: &[Map]) -> HashMap { + let mut stats = HashMap::new(); + for row in result_rows { + let Some(id) = row.get("experiment_id").and_then(Value::as_str) else { + continue; + }; + stats.insert( + id.to_string(), + CostStats { + cost: value_as_opt_f64(row.get("cost")), + priced_spans: value_as_u64(row.get("priced_spans")), + unpriced_llm_spans: value_as_u64(row.get("unpriced_llm_spans")), + }, + ); + } + stats +} + +fn print_table(ctx: &ResolvedContext, rows: &[ExperimentCostRow]) -> Result<()> { + let mut output = String::new(); + let count = format!( + "{} {}", + rows.len(), + pluralize(rows.len(), "experiment", None) + ); + writeln!( + output, + "{} in {} {} {}\n", + console::style(count), + console::style(ctx.client.org_name()).bold(), + console::style("/").dim().bold(), + console::style(&ctx.project.name).bold() + )?; + + let mut table = styled_table(); + table.set_header(vec![ + header("Name"), + header("Created"), + header("Cost"), + header("Coverage"), + ]); + apply_column_padding(&mut table, (0, 4)); + + for row in rows { + let created = row + .created + .as_deref() + .map(|c| truncate(c, 10)) + .unwrap_or_else(|| "-".to_string()); + table.add_row(vec![ + truncate(&row.name, 60), + created, + cost_display(row.cost, row.coverage), + coverage_display(row.priced_spans, row.unpriced_llm_spans, row.coverage), + ]); + } + + write!(output, "{table}")?; + + let priced_total: f64 = rows.iter().filter_map(|r| r.cost).sum(); + let partial = rows + .iter() + .filter(|r| r.coverage == Coverage::Partial) + .count(); + let unknown = rows + .iter() + .filter(|r| r.coverage == Coverage::Unknown) + .count(); + let no_data = rows.iter().filter(|r| r.coverage == Coverage::None).count(); + + write!( + output, + "\n\nTotal (priced): {}", + console::style(format_cost(priced_total)).bold() + )?; + if partial > 0 { + write!( + output, + " {}", + console::style(format!("{partial} partial")).yellow() + )?; + } + if unknown > 0 { + write!( + output, + " {}", + console::style(format!("{unknown} unknown")).yellow() + )?; + } + if no_data > 0 { + write!( + output, + " {}", + console::style(format!("{no_data} without LLM cost data")).dim() + )?; + } + output.push('\n'); + + print_with_pager(&output)?; + Ok(()) +} + +fn cost_display(cost: Option, coverage: Coverage) -> String { + match coverage { + Coverage::Full => format_cost(cost.unwrap_or(0.0)), + Coverage::Partial => format!("~{}", format_cost(cost.unwrap_or(0.0))), + Coverage::Unknown => "unknown".to_string(), + Coverage::None => "n/a".to_string(), + } +} + +fn coverage_display(priced_spans: u64, unpriced_llm_spans: u64, coverage: Coverage) -> String { + match coverage { + Coverage::None => "—".to_string(), + _ => format!("{}/{}", priced_spans, priced_spans + unpriced_llm_spans), + } +} + +fn value_as_opt_f64(value: Option<&Value>) -> Option { + match value { + Some(Value::Number(n)) => n.as_f64(), + Some(Value::String(s)) => s.parse().ok(), + _ => None, + } +} + +fn value_as_u64(value: Option<&Value>) -> u64 { + match value { + Some(Value::Number(n)) => n + .as_u64() + .or_else(|| n.as_f64().map(|f| f.max(0.0) as u64)) + .unwrap_or(0), + Some(Value::String(s)) => s.parse().unwrap_or(0), + _ => 0, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn build_cost_query_lists_ids_and_matches_ui_sum() { + let query = build_cost_query(&["exp-a".to_string(), "exp-b".to_string()]); + assert!( + query.contains("FROM experiment('exp-a', 'exp-b')"), + "{query}" + ); + assert!(query.contains("SUM(estimated_cost()) AS cost"), "{query}"); + assert!( + query.contains("COUNT(estimated_cost()) AS priced_spans"), + "{query}" + ); + assert!(query.contains("AS unpriced_llm_spans"), "{query}"); + assert!( + query.contains("span_attributes.purpose != 'scorer'"), + "{query}" + ); + assert!(query.contains("span_attributes.type = 'llm'"), "{query}"); + assert!( + query.contains("metrics.prompt_tokens IS NOT NULL"), + "{query}" + ); + assert!(query.contains("GROUP BY experiment_id"), "{query}"); + // The cost sum must not be restricted by span type, to match the web UI. + assert!( + !query.contains("WHERE"), + "cost sum should not be filtered: {query}" + ); + } + + #[test] + fn build_cost_query_escapes_single_quotes() { + let query = build_cost_query(&["ex'p".to_string()]); + assert!(query.contains("experiment('ex''p')"), "{query}"); + } + + #[test] + fn classify_covers_all_cases() { + assert_eq!(Coverage::classify(0, 0), Coverage::None); + assert_eq!(Coverage::classify(0, 5), Coverage::Unknown); + assert_eq!(Coverage::classify(3, 2), Coverage::Partial); + assert_eq!(Coverage::classify(5, 0), Coverage::Full); + } + + #[test] + fn cost_and_coverage_display_reflect_classification() { + assert_eq!(cost_display(Some(1.5), Coverage::Full), "$1.50"); + assert_eq!(cost_display(Some(1.5), Coverage::Partial), "~$1.50"); + assert_eq!(cost_display(None, Coverage::Unknown), "unknown"); + assert_eq!(cost_display(None, Coverage::None), "n/a"); + + assert_eq!(coverage_display(5, 0, Coverage::Full), "5/5"); + assert_eq!(coverage_display(3, 2, Coverage::Partial), "3/5"); + assert_eq!(coverage_display(0, 4, Coverage::Unknown), "0/4"); + assert_eq!(coverage_display(0, 0, Coverage::None), "—"); + } + + #[test] + fn index_cost_rows_parses_null_cost_and_counts() { + let rows = vec![ + json!({ + "experiment_id": "exp-a", + "cost": 1.25, + "priced_spans": 8, + "unpriced_llm_spans": 2 + }) + .as_object() + .unwrap() + .clone(), + json!({ + "experiment_id": "exp-b", + "cost": null, + "priced_spans": 0, + "unpriced_llm_spans": 0 + }) + .as_object() + .unwrap() + .clone(), + ]; + let stats = index_cost_rows(&rows); + assert_eq!( + stats.get("exp-a"), + Some(&CostStats { + cost: Some(1.25), + priced_spans: 8, + unpriced_llm_spans: 2 + }) + ); + assert_eq!( + stats.get("exp-b"), + Some(&CostStats { + cost: None, + priced_spans: 0, + unpriced_llm_spans: 0 + }) + ); + } + + #[test] + fn experiment_cost_row_json_shape() { + let row = ExperimentCostRow { + id: "exp-a".to_string(), + name: "baseline".to_string(), + created: Some("2024-01-02T03:04:05Z".to_string()), + cost: Some(1.25), + priced_spans: 8, + unpriced_llm_spans: 2, + coverage: Coverage::Partial, + }; + let value = serde_json::to_value(&row).unwrap(); + assert_eq!(value["id"], "exp-a"); + assert_eq!(value["name"], "baseline"); + assert_eq!(value["cost"], 1.25); + assert_eq!(value["priced_spans"], 8); + assert_eq!(value["unpriced_llm_spans"], 2); + assert_eq!(value["coverage"], "partial"); + } + + #[test] + fn experiment_cost_row_json_null_cost_for_no_data() { + let row = ExperimentCostRow { + id: "exp-z".to_string(), + name: "empty".to_string(), + created: None, + cost: None, + priced_spans: 0, + unpriced_llm_spans: 0, + coverage: Coverage::None, + }; + let value = serde_json::to_value(&row).unwrap(); + assert!(value["cost"].is_null()); + assert_eq!(value["coverage"], "none"); + } +} diff --git a/src/cost/mod.rs b/src/cost/mod.rs new file mode 100644 index 00000000..f163d282 --- /dev/null +++ b/src/cost/mod.rs @@ -0,0 +1,38 @@ +use anyhow::Result; +use clap::{Args, Subcommand}; + +use crate::{args::BaseArgs, project_context::resolve_project_command_context_with_auth_mode}; + +mod experiments; + +pub(crate) use crate::project_context::ProjectContext as ResolvedContext; + +#[derive(Debug, Clone, Args)] +#[command( + about = "Estimate LLM cost for Braintrust resources", + after_help = "\ +Examples: + bt cost experiments Estimate LLM cost per experiment in the active project + bt cost experiments --json Emit structured per-experiment cost rows +" +)] +pub struct CostArgs { + #[command(subcommand)] + command: Option, +} + +#[derive(Debug, Clone, Subcommand)] +enum CostCommands { + /// Estimate LLM cost per experiment in the active project + Experiments, +} + +pub async fn run(base: BaseArgs, args: CostArgs) -> Result<()> { + // Every `bt cost` command is read-only: it only issues `GET /v1/experiment` + // and `POST /btql` (a read query). + let ctx = resolve_project_command_context_with_auth_mode(&base, true).await?; + + match args.command { + None | Some(CostCommands::Experiments) => experiments::run(&ctx, base.json).await, + } +} diff --git a/src/utils/format.rs b/src/utils/format.rs new file mode 100644 index 00000000..9df72d5b --- /dev/null +++ b/src/utils/format.rs @@ -0,0 +1,38 @@ +/// Format a USD cost value for display. +/// +/// Small non-zero values are shown as `<$0.001`, values under a dollar use +/// three decimal places, and larger values use two. +pub(crate) fn format_cost(cost: f64) -> String { + if cost > 0.0 && cost < 0.001 { + "<$0.001".to_string() + } else if cost < 1.0 { + format!("${cost:.3}") + } else { + format!("${cost:.2}") + } +} + +#[cfg(test)] +mod tests { + use super::format_cost; + + #[test] + fn formats_sub_millicent_costs() { + assert_eq!(format_cost(0.0004), "<$0.001"); + } + + #[test] + fn formats_sub_dollar_costs_with_three_decimals() { + assert_eq!(format_cost(0.123), "$0.123"); + } + + #[test] + fn formats_larger_costs_with_two_decimals() { + assert_eq!(format_cost(12.3456), "$12.35"); + } + + #[test] + fn formats_zero_as_three_decimals() { + assert_eq!(format_cost(0.0), "$0.000"); + } +} From 9c721caec9afcec2b8a9071bceb754c2e81d9d8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Mon, 20 Jul 2026 18:25:28 -0700 Subject: [PATCH 19/24] logs --- src/cost/logs.rs | 796 ++++++++++++++++++++++++++++++++++++++++++++ src/cost/mod.rs | 36 +- src/cost/pricing.rs | 535 +++++++++++++++++++++++++++++ tests/cli.rs | 13 + 4 files changed, 1373 insertions(+), 7 deletions(-) create mode 100644 src/cost/logs.rs create mode 100644 src/cost/pricing.rs diff --git a/src/cost/logs.rs b/src/cost/logs.rs new file mode 100644 index 00000000..bda05991 --- /dev/null +++ b/src/cost/logs.rs @@ -0,0 +1,796 @@ +use std::cmp::Ordering; +use std::fmt::Write as _; +use std::path::PathBuf; + +use anyhow::{bail, Context, Result}; +use chrono::{DateTime, Duration, Timelike, Utc}; +use clap::{builder::BoolishValueParser, Args}; +use dialoguer::console; +use serde::Serialize; +use serde_json::{Map, Value}; + +use crate::{ + sql::run_btql_rows, + ui::{apply_column_padding, header, print_with_pager, styled_table, truncate, with_spinner}, + utils::{format_cost, parse_duration_to_seconds, pluralize}, +}; + +use super::{ + pricing::{format_timestamp, parse_timestamp, PriceBook, TokenUsage}, + ResolvedContext, +}; + +#[derive(Debug, Clone, Args)] +#[command(after_help = "\ +Pricing file (USD per 1 million tokens): + version = 1 + + [models.\"custom-chat-model\"] + aliases = [\"custom-deployment-name\"] + + [[models.\"custom-chat-model\".rates]] + effective_from = \"2025-01-01T00:00:00Z\" + effective_until = \"2025-06-01T00:00:00Z\" + input_usd_per_1m_tokens = 3.0 + cached_input_usd_per_1m_tokens = 0.3 + cache_write_usd_per_1m_tokens = 3.75 + cache_write_5m_usd_per_1m_tokens = 3.75 + cache_write_1h_usd_per_1m_tokens = 6.0 + output_usd_per_1m_tokens = 15.0 + +Repeat [[models.\"...\".rates]] for historical prices. Bounds are [from, until). +When effective_until is omitted, a rate ends at the next effective_from, or never. +Braintrust's effective estimated cost takes precedence; file rates fill unpriced token spans.")] +pub(crate) struct LogsArgs { + /// Relative time window ending at --until + #[arg(long, env = "BRAINTRUST_COST_WINDOW", default_value = "7d")] + window: String, + + /// Absolute inclusive lower bound (RFC 3339 or YYYY-MM-DD); overrides --window + #[arg(long, env = "BRAINTRUST_COST_SINCE")] + since: Option, + + /// Absolute exclusive upper bound (RFC 3339 or YYYY-MM-DD); defaults to now + #[arg(long, env = "BRAINTRUST_COST_UNTIL")] + until: Option, + + /// TOML file containing historical per-model token prices + #[arg(long, env = "BRAINTRUST_COST_PRICING_FILE", value_name = "PATH")] + pricing_file: Option, + + /// Exclude spans whose purpose is scorer + #[arg( + long, + env = "BRAINTRUST_COST_EXCLUDE_SCORERS", + value_parser = BoolishValueParser::new(), + default_value_t = false + )] + exclude_scorers: bool, +} + +impl Default for LogsArgs { + fn default() -> Self { + Self { + window: "7d".to_string(), + since: None, + until: None, + pricing_file: None, + exclude_scorers: false, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct TimeRange { + since: DateTime, + until: DateTime, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct TimeSegment { + since: DateTime, + until: DateTime, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +enum Coverage { + Full, + Partial, + Unknown, + None, +} + +impl Coverage { + fn classify(priced_spans: u64, unpriced_token_spans: u64) -> Self { + match (priced_spans, unpriced_token_spans) { + (0, 0) => Coverage::None, + (0, _) => Coverage::Unknown, + (_, 0) => Coverage::Full, + (_, _) => Coverage::Partial, + } + } +} + +#[derive(Debug, Clone, Serialize)] +struct LogsCostRow { + model: Option, + purpose: Option, + candidate_spans: u64, + cost: Option, + braintrust_cost: Option, + file_cost: Option, + braintrust_priced_spans: u64, + file_priced_spans: u64, + unpriced_token_spans: u64, + no_usage_spans: u64, + coverage: Coverage, +} + +impl LogsCostRow { + fn priced_spans(&self) -> u64 { + self.braintrust_priced_spans + self.file_priced_spans + } +} + +#[derive(Debug, Clone, Serialize)] +struct LogsCostTotals { + candidate_spans: u64, + cost: Option, + braintrust_cost: Option, + file_cost: Option, + braintrust_priced_spans: u64, + file_priced_spans: u64, + unpriced_token_spans: u64, + no_usage_spans: u64, + coverage: Coverage, +} + +#[derive(Debug, Serialize)] +struct LogsCostOutput<'a> { + project: &'a str, + org: &'a str, + currency: &'static str, + since: String, + until: String, + excludes_scorers: bool, + pricing_file: Option, + rows: &'a [LogsCostRow], + totals: &'a LogsCostTotals, +} + +pub(crate) async fn run(ctx: &ResolvedContext, args: LogsArgs, json: bool) -> Result<()> { + let now = Utc::now() + .with_nanosecond(0) + .expect("zero nanoseconds is a valid timestamp"); + let range = resolve_time_range(&args, now)?; + let price_book = args + .pricing_file + .as_deref() + .map(PriceBook::load) + .transpose()?; + let segments = build_time_segments(range, price_book.as_ref()); + let query = build_cost_query(&ctx.project.id, range, &segments, args.exclude_scorers); + let result_rows = with_spinner( + "Estimating log cost...", + run_btql_rows(&ctx.client, &query, "default"), + ) + .await?; + + let mut rows = build_cost_rows(&result_rows, &segments, price_book.as_ref())?; + rows.sort_by(compare_cost_rows); + let totals = calculate_totals(&rows); + + if json { + let output = LogsCostOutput { + project: &ctx.project.name, + org: ctx.client.org_name(), + currency: "USD", + since: format_timestamp(range.since), + until: format_timestamp(range.until), + excludes_scorers: args.exclude_scorers, + pricing_file: args + .pricing_file + .as_ref() + .map(|path| path.display().to_string()), + rows: &rows, + totals: &totals, + }; + println!("{}", serde_json::to_string(&output)?); + return Ok(()); + } + + print_table(ctx, range, args.pricing_file.as_ref(), &rows, &totals)?; + Ok(()) +} + +fn resolve_time_range(args: &LogsArgs, now: DateTime) -> Result { + let until = args + .until + .as_deref() + .map(parse_timestamp) + .transpose() + .context("invalid --until")? + .unwrap_or(now); + let since = match args.since.as_deref() { + Some(since) => parse_timestamp(since).context("invalid --since")?, + None => { + let seconds = parse_duration_to_seconds(&args.window) + .with_context(|| format!("invalid --window '{}'", args.window))?; + if seconds == 0 { + bail!("--window must be greater than zero"); + } + let seconds = i64::try_from(seconds).context("--window is too large")?; + until + .checked_sub_signed(Duration::seconds(seconds)) + .context("--window produces a timestamp outside the supported range")? + } + }; + if since >= until { + bail!("--since must be earlier than --until"); + } + Ok(TimeRange { since, until }) +} + +fn build_time_segments(range: TimeRange, price_book: Option<&PriceBook>) -> Vec { + let mut boundaries = vec![range.since, range.until]; + if let Some(price_book) = price_book { + boundaries.extend(price_book.boundaries_between(range.since, range.until)); + } + boundaries.sort_unstable(); + boundaries.dedup(); + boundaries + .windows(2) + .map(|bounds| TimeSegment { + since: bounds[0], + until: bounds[1], + }) + .collect() +} + +fn build_cost_query( + project_id: &str, + range: TimeRange, + segments: &[TimeSegment], + exclude_scorers: bool, +) -> String { + let prompt_tokens = "COALESCE(metrics.prompt_tokens, 0)"; + let completion_tokens = "COALESCE(metrics.completion_tokens, 0)"; + let cached_tokens = "COALESCE(metrics.prompt_cached_tokens, 0)"; + let generic_write_tokens = "COALESCE(metrics.prompt_cache_creation_tokens, 0)"; + let write_5m_tokens = "COALESCE(metrics.prompt_cache_creation_5m_tokens, 0)"; + let write_1h_tokens = "COALESCE(metrics.prompt_cache_creation_1h_tokens, 0)"; + let split_write_tokens = format!("({write_5m_tokens} + {write_1h_tokens})"); + let effective_write_tokens = format!("GREATEST({generic_write_tokens}, {split_write_tokens})"); + let uncached_input_tokens = + format!("GREATEST(0, {prompt_tokens} - {cached_tokens} - {effective_write_tokens})"); + let token_activity = [ + "metrics.prompt_tokens IS NOT NULL", + "metrics.completion_tokens IS NOT NULL", + "metrics.prompt_cached_tokens IS NOT NULL", + "metrics.prompt_cache_creation_tokens IS NOT NULL", + "metrics.prompt_cache_creation_5m_tokens IS NOT NULL", + "metrics.prompt_cache_creation_1h_tokens IS NOT NULL", + ] + .join(" OR "); + + let mut select_fields = vec![ + "metadata.model AS model".to_string(), + "span_attributes.purpose AS purpose".to_string(), + format!( + "COUNT(CASE WHEN estimated_cost() IS NOT NULL OR metadata.model IS NOT NULL OR ({token_activity}) THEN 1 END) AS candidate_spans" + ), + "COUNT(estimated_cost()) AS braintrust_priced_spans".to_string(), + "SUM(estimated_cost()) AS braintrust_cost".to_string(), + format!( + "COUNT(CASE WHEN estimated_cost() IS NULL AND ({token_activity}) THEN 1 END) AS unpriced_token_spans" + ), + format!( + "COUNT(CASE WHEN estimated_cost() IS NULL AND metadata.model IS NOT NULL AND NOT ({token_activity}) THEN 1 END) AS no_usage_spans" + ), + ]; + + for (index, segment) in segments.iter().enumerate() { + let condition = format!( + "estimated_cost() IS NULL AND created >= {} AND created < {} AND ({token_activity})", + sql_quote(&format_timestamp(segment.since)), + sql_quote(&format_timestamp(segment.until)), + ); + let split_complete = format!("{split_write_tokens} >= {generic_write_tokens}"); + let fallback_write_tokens = + format!("CASE WHEN {split_complete} THEN 0 ELSE {effective_write_tokens} END"); + let split_5m = format!("CASE WHEN {split_complete} THEN {write_5m_tokens} ELSE 0 END"); + let split_1h = format!("CASE WHEN {split_complete} THEN {write_1h_tokens} ELSE 0 END"); + + select_fields.extend([ + format!("COUNT(CASE WHEN {condition} THEN 1 END) AS p{index}_spans"), + format!( + "SUM(CASE WHEN {condition} THEN {uncached_input_tokens} ELSE 0 END) AS p{index}_uncached_input_tokens" + ), + format!( + "SUM(CASE WHEN {condition} THEN {cached_tokens} ELSE 0 END) AS p{index}_cached_input_tokens" + ), + format!( + "SUM(CASE WHEN {condition} THEN {effective_write_tokens} ELSE 0 END) AS p{index}_effective_cache_write_tokens" + ), + format!( + "SUM(CASE WHEN {condition} THEN {split_5m} ELSE 0 END) AS p{index}_split_cache_write_5m_tokens" + ), + format!( + "SUM(CASE WHEN {condition} THEN {split_1h} ELSE 0 END) AS p{index}_split_cache_write_1h_tokens" + ), + format!( + "SUM(CASE WHEN {condition} THEN {fallback_write_tokens} ELSE 0 END) AS p{index}_fallback_cache_write_tokens" + ), + format!( + "SUM(CASE WHEN {condition} THEN {completion_tokens} ELSE 0 END) AS p{index}_output_tokens" + ), + ]); + } + + let scorer_filter = if exclude_scorers { + "\n AND (span_attributes.purpose IS NULL OR span_attributes.purpose != 'scorer')" + } else { + "" + }; + format!( + "SELECT\n {}\nFROM project_logs({}, shape => 'spans')\nWHERE created >= {}\n AND created < {}{}\nGROUP BY metadata.model, span_attributes.purpose", + select_fields.join(",\n "), + sql_quote(project_id), + sql_quote(&format_timestamp(range.since)), + sql_quote(&format_timestamp(range.until)), + scorer_filter, + ) +} + +fn build_cost_rows( + result_rows: &[Map], + segments: &[TimeSegment], + price_book: Option<&PriceBook>, +) -> Result> { + result_rows + .iter() + .map(|row| build_cost_row(row, segments, price_book)) + .filter_map(|result| match result { + Ok(row) if row.candidate_spans == 0 => None, + other => Some(other), + }) + .collect() +} + +fn build_cost_row( + row: &Map, + segments: &[TimeSegment], + price_book: Option<&PriceBook>, +) -> Result { + let model = row.get("model").and_then(Value::as_str).map(str::to_string); + let purpose = row + .get("purpose") + .and_then(Value::as_str) + .map(str::to_string); + let candidate_spans = value_as_u64(row.get("candidate_spans")); + let braintrust_priced_spans = value_as_u64(row.get("braintrust_priced_spans")); + let braintrust_cost = value_as_opt_f64(row.get("braintrust_cost")); + let expected_unpriced_token_spans = value_as_u64(row.get("unpriced_token_spans")); + let no_usage_spans = value_as_u64(row.get("no_usage_spans")); + + let mut file_cost = 0.0; + let mut file_priced_spans = 0_u64; + let mut unpriced_token_spans = 0_u64; + let mut segmented_token_spans = 0_u64; + + for (index, segment) in segments.iter().enumerate() { + let usage = token_usage_from_row(row, index); + segmented_token_spans = segmented_token_spans.saturating_add(usage.spans); + let rates = model.as_deref().and_then(|model| { + price_book.and_then(|price_book| price_book.rate_at(model, segment.since)) + }); + if let Some(rates) = rates { + file_cost += usage.cost(rates); + file_priced_spans = file_priced_spans.saturating_add(usage.spans); + } else { + unpriced_token_spans = unpriced_token_spans.saturating_add(usage.spans); + } + } + + if segmented_token_spans != expected_unpriced_token_spans { + bail!( + "cost query returned inconsistent token coverage for model {}: expected {}, got {}", + model.as_deref().unwrap_or(""), + expected_unpriced_token_spans, + segmented_token_spans + ); + } + + let priced_spans = braintrust_priced_spans.saturating_add(file_priced_spans); + let cost = (priced_spans > 0).then(|| braintrust_cost.unwrap_or(0.0) + file_cost); + Ok(LogsCostRow { + model, + purpose, + candidate_spans, + cost, + braintrust_cost: (braintrust_priced_spans > 0).then(|| braintrust_cost.unwrap_or(0.0)), + file_cost: (file_priced_spans > 0).then_some(file_cost), + braintrust_priced_spans, + file_priced_spans, + unpriced_token_spans, + no_usage_spans, + coverage: Coverage::classify(priced_spans, unpriced_token_spans), + }) +} + +fn token_usage_from_row(row: &Map, index: usize) -> TokenUsage { + let value = |suffix: &str| value_as_u64(row.get(&format!("p{index}_{suffix}"))); + TokenUsage { + spans: value("spans"), + uncached_input_tokens: value("uncached_input_tokens"), + cached_input_tokens: value("cached_input_tokens"), + effective_cache_write_tokens: value("effective_cache_write_tokens"), + split_cache_write_5m_tokens: value("split_cache_write_5m_tokens"), + split_cache_write_1h_tokens: value("split_cache_write_1h_tokens"), + fallback_cache_write_tokens: value("fallback_cache_write_tokens"), + output_tokens: value("output_tokens"), + } +} + +fn calculate_totals(rows: &[LogsCostRow]) -> LogsCostTotals { + let braintrust_priced_spans = rows.iter().map(|row| row.braintrust_priced_spans).sum(); + let file_priced_spans = rows.iter().map(|row| row.file_priced_spans).sum(); + let unpriced_token_spans = rows.iter().map(|row| row.unpriced_token_spans).sum(); + let priced_spans = braintrust_priced_spans + file_priced_spans; + let braintrust_cost = (braintrust_priced_spans > 0) + .then(|| rows.iter().filter_map(|row| row.braintrust_cost).sum()); + let file_cost = + (file_priced_spans > 0).then(|| rows.iter().filter_map(|row| row.file_cost).sum()); + LogsCostTotals { + candidate_spans: rows.iter().map(|row| row.candidate_spans).sum(), + cost: (priced_spans > 0).then(|| braintrust_cost.unwrap_or(0.0) + file_cost.unwrap_or(0.0)), + braintrust_cost, + file_cost, + braintrust_priced_spans, + file_priced_spans, + unpriced_token_spans, + no_usage_spans: rows.iter().map(|row| row.no_usage_spans).sum(), + coverage: Coverage::classify(priced_spans, unpriced_token_spans), + } +} + +fn compare_cost_rows(left: &LogsCostRow, right: &LogsCostRow) -> Ordering { + let left_cost = left.cost.unwrap_or(f64::NEG_INFINITY); + let right_cost = right.cost.unwrap_or(f64::NEG_INFINITY); + right_cost + .partial_cmp(&left_cost) + .unwrap_or(Ordering::Equal) + .then_with(|| left.model.cmp(&right.model)) + .then_with(|| left.purpose.cmp(&right.purpose)) +} + +fn print_table( + ctx: &ResolvedContext, + range: TimeRange, + pricing_file: Option<&PathBuf>, + rows: &[LogsCostRow], + totals: &LogsCostTotals, +) -> Result<()> { + let mut output = String::new(); + let count = format!( + "{} {}", + rows.len(), + pluralize(rows.len(), "cost group", None) + ); + writeln!( + output, + "{} in {} {} {}", + console::style(count), + console::style(ctx.client.org_name()).bold(), + console::style("/").dim().bold(), + console::style(&ctx.project.name).bold() + )?; + writeln!( + output, + "{} {} {}", + console::style(format_timestamp(range.since)).dim(), + console::style("to").dim(), + console::style(format_timestamp(range.until)).dim() + )?; + if let Some(path) = pricing_file { + writeln!(output, "Pricing: {}", console::style(path.display()).dim())?; + } + output.push('\n'); + + let mut table = styled_table(); + table.set_header(vec![ + header("Model"), + header("Purpose"), + header("Cost"), + header("Coverage"), + header("Sources"), + ]); + apply_column_padding(&mut table, (0, 4)); + + for row in rows { + table.add_row(vec![ + truncate(row.model.as_deref().unwrap_or(""), 50), + truncate(row.purpose.as_deref().unwrap_or("default"), 20), + cost_display(row.cost, row.coverage), + coverage_display(row.priced_spans(), row.unpriced_token_spans, row.coverage), + source_display(row), + ]); + } + write!(output, "{table}")?; + + let total_cost = totals + .cost + .map(format_cost) + .unwrap_or_else(|| "n/a".to_string()); + let total_prefix = if totals.unpriced_token_spans > 0 { + "~" + } else { + "" + }; + write!( + output, + "\n\nTotal: {}", + console::style(format!("{total_prefix}{total_cost}")).bold() + )?; + if let Some(file_cost) = totals.file_cost { + write!( + output, + " {}", + console::style(format!("{} from pricing file", format_cost(file_cost))).dim() + )?; + } + if totals.unpriced_token_spans > 0 { + write!( + output, + " {}", + console::style(format!( + "{} unpriced token {}", + totals.unpriced_token_spans, + pluralize(totals.unpriced_token_spans as usize, "span", None) + )) + .yellow() + )?; + } + if totals.no_usage_spans > 0 { + write!( + output, + " {}", + console::style(format!("{} without cost or usage", totals.no_usage_spans)).dim() + )?; + } + output.push('\n'); + + print_with_pager(&output)?; + Ok(()) +} + +fn cost_display(cost: Option, coverage: Coverage) -> String { + match (cost, coverage) { + (Some(cost), Coverage::Partial) => format!("~{}", format_cost(cost)), + (Some(cost), _) => format_cost(cost), + (None, Coverage::Unknown) => "unknown".to_string(), + (None, _) => "n/a".to_string(), + } +} + +fn coverage_display(priced_spans: u64, unpriced_token_spans: u64, coverage: Coverage) -> String { + match coverage { + Coverage::None => "—".to_string(), + _ => format!("{}/{}", priced_spans, priced_spans + unpriced_token_spans), + } +} + +fn source_display(row: &LogsCostRow) -> String { + let mut parts = Vec::new(); + if row.braintrust_priced_spans > 0 { + parts.push(format!("Braintrust {}", row.braintrust_priced_spans)); + } + if row.file_priced_spans > 0 { + parts.push(format!("file {}", row.file_priced_spans)); + } + if row.no_usage_spans > 0 { + parts.push(format!("no usage {}", row.no_usage_spans)); + } + if parts.is_empty() { + "—".to_string() + } else { + parts.join(", ") + } +} + +fn value_as_opt_f64(value: Option<&Value>) -> Option { + match value { + Some(Value::Number(number)) => number.as_f64(), + Some(Value::String(value)) => value.parse().ok(), + _ => None, + } + .filter(|value| value.is_finite() && *value >= 0.0) +} + +fn value_as_u64(value: Option<&Value>) -> u64 { + match value { + Some(Value::Number(number)) => number + .as_u64() + .or_else(|| number.as_f64().map(|value| value.max(0.0) as u64)) + .unwrap_or(0), + Some(Value::String(value)) => value.parse().unwrap_or(0), + _ => 0, + } +} + +fn sql_quote(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn range() -> TimeRange { + TimeRange { + since: parse_timestamp("2025-01-01").unwrap(), + until: parse_timestamp("2025-03-01").unwrap(), + } + } + + #[test] + fn query_is_timestamp_bounded_and_uses_spans() { + let query = build_cost_query( + "test-project-id", + range(), + &[TimeSegment { + since: range().since, + until: range().until, + }], + false, + ); + assert!( + query.contains("FROM project_logs('test-project-id', shape => 'spans')"), + "{query}" + ); + assert!( + query.contains("created >= '2025-01-01T00:00:00Z'"), + "{query}" + ); + assert!( + query.contains("created < '2025-03-01T00:00:00Z'"), + "{query}" + ); + assert!( + query.contains("SUM(estimated_cost()) AS braintrust_cost"), + "{query}" + ); + assert!(query.contains("p0_uncached_input_tokens"), "{query}"); + assert!(!query.contains("purpose != 'scorer'"), "{query}"); + } + + #[test] + fn query_can_exclude_scorers() { + let query = build_cost_query("test-project-id", range(), &[], true); + assert!( + query + .contains("span_attributes.purpose IS NULL OR span_attributes.purpose != 'scorer'"), + "{query}" + ); + } + + #[test] + fn time_range_uses_until_as_window_anchor() { + let args = LogsArgs { + window: "1d".to_string(), + until: Some("2025-02-02T12:00:00Z".to_string()), + ..LogsArgs::default() + }; + let resolved = resolve_time_range(&args, parse_timestamp("2030-01-01").unwrap()).unwrap(); + assert_eq!( + resolved.since, + parse_timestamp("2025-02-01T12:00:00Z").unwrap() + ); + assert_eq!( + resolved.until, + parse_timestamp("2025-02-02T12:00:00Z").unwrap() + ); + } + + #[test] + fn historical_file_rates_are_applied_to_the_matching_segment() { + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + file.path(), + r#" +version = 1 + +[models."test-model"] +[[models."test-model".rates]] +effective_from = "2025-01-01" +input_usd_per_1m_tokens = 1.0 +output_usd_per_1m_tokens = 2.0 + +[[models."test-model".rates]] +effective_from = "2025-02-01" +input_usd_per_1m_tokens = 3.0 +output_usd_per_1m_tokens = 4.0 +"#, + ) + .unwrap(); + let price_book = PriceBook::load(file.path()).unwrap(); + let segments = build_time_segments(range(), Some(&price_book)); + assert_eq!(segments.len(), 2); + + let row = json!({ + "model": "test-model", + "purpose": null, + "candidate_spans": 3, + "braintrust_priced_spans": 1, + "braintrust_cost": 0.5, + "unpriced_token_spans": 2, + "no_usage_spans": 0, + "p0_spans": 1, + "p0_uncached_input_tokens": 1000000, + "p0_cached_input_tokens": 0, + "p0_effective_cache_write_tokens": 0, + "p0_split_cache_write_5m_tokens": 0, + "p0_split_cache_write_1h_tokens": 0, + "p0_fallback_cache_write_tokens": 0, + "p0_output_tokens": 1000000, + "p1_spans": 1, + "p1_uncached_input_tokens": 1000000, + "p1_cached_input_tokens": 0, + "p1_effective_cache_write_tokens": 0, + "p1_split_cache_write_5m_tokens": 0, + "p1_split_cache_write_1h_tokens": 0, + "p1_fallback_cache_write_tokens": 0, + "p1_output_tokens": 1000000 + }) + .as_object() + .unwrap() + .clone(); + let cost_row = build_cost_row(&row, &segments, Some(&price_book)).unwrap(); + assert_eq!(cost_row.braintrust_cost, Some(0.5)); + assert_eq!(cost_row.file_cost, Some(10.0)); + assert_eq!(cost_row.cost, Some(10.5)); + assert_eq!(cost_row.file_priced_spans, 2); + assert_eq!(cost_row.coverage, Coverage::Full); + } + + #[test] + fn missing_historical_rate_remains_unpriced() { + let row = json!({ + "model": null, + "candidate_spans": 1, + "braintrust_priced_spans": 0, + "braintrust_cost": null, + "unpriced_token_spans": 1, + "no_usage_spans": 0, + "p0_spans": 1, + "p0_uncached_input_tokens": 10, + "p0_cached_input_tokens": 0, + "p0_effective_cache_write_tokens": 0, + "p0_split_cache_write_5m_tokens": 0, + "p0_split_cache_write_1h_tokens": 0, + "p0_fallback_cache_write_tokens": 0, + "p0_output_tokens": 5 + }) + .as_object() + .unwrap() + .clone(); + let segments = vec![TimeSegment { + since: range().since, + until: range().until, + }]; + let cost_row = build_cost_row(&row, &segments, None).unwrap(); + assert_eq!(cost_row.cost, None); + assert_eq!(cost_row.unpriced_token_spans, 1); + assert_eq!(cost_row.coverage, Coverage::Unknown); + } + + #[test] + fn partial_cost_is_marked_approximate() { + assert_eq!(cost_display(Some(1.25), Coverage::Partial), "~$1.25"); + assert_eq!(cost_display(None, Coverage::Unknown), "unknown"); + assert_eq!(coverage_display(3, 2, Coverage::Partial), "3/5"); + } +} diff --git a/src/cost/mod.rs b/src/cost/mod.rs index f163d282..90545ec7 100644 --- a/src/cost/mod.rs +++ b/src/cost/mod.rs @@ -1,9 +1,13 @@ +use std::io::{self, Write as _}; + use anyhow::Result; -use clap::{Args, Subcommand}; +use clap::{Args, CommandFactory, Subcommand}; use crate::{args::BaseArgs, project_context::resolve_project_command_context_with_auth_mode}; mod experiments; +mod logs; +mod pricing; pub(crate) use crate::project_context::ProjectContext as ResolvedContext; @@ -12,8 +16,10 @@ pub(crate) use crate::project_context::ProjectContext as ResolvedContext; about = "Estimate LLM cost for Braintrust resources", after_help = "\ Examples: - bt cost experiments Estimate LLM cost per experiment in the active project - bt cost experiments --json Emit structured per-experiment cost rows + bt cost experiments Estimate LLM cost per experiment + bt cost logs Estimate log cost over the last 7 days + bt cost logs --window 30d Estimate log cost over the last 30 days + bt cost logs --pricing-file prices.toml Price otherwise-unpriced token usage " )] pub struct CostArgs { @@ -25,14 +31,30 @@ pub struct CostArgs { enum CostCommands { /// Estimate LLM cost per experiment in the active project Experiments, + /// Estimate cost for logs in the active project + Logs(logs::LogsArgs), } pub async fn run(base: BaseArgs, args: CostArgs) -> Result<()> { - // Every `bt cost` command is read-only: it only issues `GET /v1/experiment` - // and `POST /btql` (a read query). + let Some(command) = args.command else { + return print_help(); + }; let ctx = resolve_project_command_context_with_auth_mode(&base, true).await?; - match args.command { - None | Some(CostCommands::Experiments) => experiments::run(&ctx, base.json).await, + match command { + CostCommands::Experiments => experiments::run(&ctx, base.json).await, + CostCommands::Logs(log_args) => logs::run(&ctx, log_args, base.json).await, } } + +fn print_help() -> Result<()> { + let mut root = crate::Cli::command(); + let command = root + .find_subcommand_mut("cost") + .expect("cost is a registered subcommand"); + command.set_bin_name("bt cost"); + let mut stdout = io::stdout().lock(); + command.write_help(&mut stdout)?; + writeln!(stdout)?; + Ok(()) +} diff --git a/src/cost/pricing.rs b/src/cost/pricing.rs new file mode 100644 index 00000000..e3e3ff5e --- /dev/null +++ b/src/cost/pricing.rs @@ -0,0 +1,535 @@ +use std::collections::{BTreeMap, HashMap}; +use std::fs; +use std::path::Path; + +use anyhow::{bail, Context, Result}; +use chrono::{DateTime, NaiveDate, Utc}; +use serde::Deserialize; + +const SUPPORTED_PRICING_VERSION: u32 = 1; + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(super) struct TokenRates { + pub input_usd_per_1m_tokens: f64, + pub cached_input_usd_per_1m_tokens: Option, + pub cache_write_usd_per_1m_tokens: Option, + pub cache_write_5m_usd_per_1m_tokens: Option, + pub cache_write_1h_usd_per_1m_tokens: Option, + pub output_usd_per_1m_tokens: f64, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(super) struct TokenUsage { + pub spans: u64, + pub uncached_input_tokens: u64, + pub cached_input_tokens: u64, + pub effective_cache_write_tokens: u64, + pub split_cache_write_5m_tokens: u64, + pub split_cache_write_1h_tokens: u64, + pub fallback_cache_write_tokens: u64, + pub output_tokens: u64, +} + +impl TokenUsage { + pub fn cost(self, rates: &TokenRates) -> f64 { + let per_million = 1_000_000.0; + let input_cost = + self.uncached_input_tokens as f64 * rates.input_usd_per_1m_tokens / per_million; + let cached_input_cost = self.cached_input_tokens as f64 + * rates + .cached_input_usd_per_1m_tokens + .unwrap_or(rates.input_usd_per_1m_tokens) + / per_million; + let generic_cache_write_rate = rates + .cache_write_usd_per_1m_tokens + .unwrap_or(rates.input_usd_per_1m_tokens); + let cache_write_cost = match ( + rates.cache_write_5m_usd_per_1m_tokens, + rates.cache_write_1h_usd_per_1m_tokens, + ) { + (Some(rate_5m), Some(rate_1h)) => { + (self.split_cache_write_5m_tokens as f64 * rate_5m + + self.split_cache_write_1h_tokens as f64 * rate_1h + + self.fallback_cache_write_tokens as f64 * generic_cache_write_rate) + / per_million + } + _ => self.effective_cache_write_tokens as f64 * generic_cache_write_rate / per_million, + }; + let output_cost = self.output_tokens as f64 * rates.output_usd_per_1m_tokens / per_million; + + input_cost + cached_input_cost + cache_write_cost + output_cost + } +} + +#[derive(Debug, Clone)] +struct PriceInterval { + effective_from: DateTime, + effective_until: Option>, + rates: TokenRates, +} + +impl PriceInterval { + fn contains(&self, timestamp: DateTime) -> bool { + timestamp >= self.effective_from + && self + .effective_until + .is_none_or(|effective_until| timestamp < effective_until) + } +} + +#[derive(Debug, Clone)] +struct ModelPriceHistory { + intervals: Vec, +} + +#[derive(Debug, Clone, Default)] +pub(super) struct PriceBook { + histories: Vec, + model_lookup: HashMap, +} + +impl PriceBook { + pub fn load(path: &Path) -> Result { + let contents = fs::read_to_string(path) + .with_context(|| format!("failed to read pricing file {}", path.display()))?; + let raw: RawPricingFile = toml::from_str(&contents) + .with_context(|| format!("failed to parse pricing file {}", path.display()))?; + Self::from_raw(raw).with_context(|| format!("invalid pricing file {}", path.display())) + } + + fn from_raw(raw: RawPricingFile) -> Result { + if raw.version != SUPPORTED_PRICING_VERSION { + bail!( + "unsupported pricing file version {}; expected {}", + raw.version, + SUPPORTED_PRICING_VERSION + ); + } + if raw.models.is_empty() { + bail!("pricing file must define at least one model under [models]"); + } + + let mut histories = Vec::with_capacity(raw.models.len()); + let mut model_lookup = HashMap::new(); + + for (model_name, raw_model) in raw.models { + let model_name = model_name.trim(); + if model_name.is_empty() { + bail!("model names cannot be empty"); + } + if raw_model.rates.is_empty() { + bail!("models.{model_name}.rates must contain at least one historical rate"); + } + + let mut parsed_rates = raw_model + .rates + .into_iter() + .enumerate() + .map(|(index, rate)| parse_rate(model_name, index, rate)) + .collect::>>()?; + parsed_rates.sort_by_key(|rate| rate.effective_from); + + for rates in parsed_rates.windows(2) { + let current = &rates[0]; + let next = &rates[1]; + if current.effective_from == next.effective_from { + bail!( + "models.{model_name}.rates has duplicate effective_from {}", + format_timestamp(current.effective_from) + ); + } + if current + .effective_until + .is_some_and(|effective_until| effective_until > next.effective_from) + { + bail!( + "models.{model_name}.rates has overlapping intervals at {}", + format_timestamp(next.effective_from) + ); + } + } + + let mut intervals = Vec::with_capacity(parsed_rates.len()); + for index in 0..parsed_rates.len() { + let next_start = parsed_rates.get(index + 1).map(|rate| rate.effective_from); + let rate = &parsed_rates[index]; + intervals.push(PriceInterval { + effective_from: rate.effective_from, + effective_until: rate.effective_until.or(next_start), + rates: rate.rates, + }); + } + + let history_index = histories.len(); + histories.push(ModelPriceHistory { intervals }); + insert_model_lookup(&mut model_lookup, model_name, history_index)?; + for alias in raw_model.aliases { + let alias = alias.trim(); + if alias.is_empty() { + bail!("models.{model_name}.aliases cannot contain an empty name"); + } + insert_model_lookup(&mut model_lookup, alias, history_index)?; + } + } + + Ok(Self { + histories, + model_lookup, + }) + } + + pub fn rate_at(&self, model: &str, timestamp: DateTime) -> Option<&TokenRates> { + let history_index = self.model_lookup.get(&normalize_model_name(model))?; + self.histories[*history_index] + .intervals + .iter() + .find(|interval| interval.contains(timestamp)) + .map(|interval| &interval.rates) + } + + pub fn boundaries_between( + &self, + since: DateTime, + until: DateTime, + ) -> Vec> { + let mut boundaries = Vec::new(); + for history in &self.histories { + for interval in &history.intervals { + if interval.effective_from > since && interval.effective_from < until { + boundaries.push(interval.effective_from); + } + if let Some(effective_until) = interval.effective_until { + if effective_until > since && effective_until < until { + boundaries.push(effective_until); + } + } + } + } + boundaries.sort_unstable(); + boundaries.dedup(); + boundaries + } +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawPricingFile { + version: u32, + models: BTreeMap, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawModelPricing { + #[serde(default)] + aliases: Vec, + rates: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawTokenRates { + effective_from: String, + effective_until: Option, + input_usd_per_1m_tokens: f64, + cached_input_usd_per_1m_tokens: Option, + cache_write_usd_per_1m_tokens: Option, + cache_write_5m_usd_per_1m_tokens: Option, + cache_write_1h_usd_per_1m_tokens: Option, + output_usd_per_1m_tokens: f64, +} + +struct ParsedTokenRates { + effective_from: DateTime, + effective_until: Option>, + rates: TokenRates, +} + +fn parse_rate(model_name: &str, index: usize, raw: RawTokenRates) -> Result { + let path = format!("models.{model_name}.rates[{index}]"); + let effective_from = parse_timestamp(&raw.effective_from) + .with_context(|| format!("{path}.effective_from is invalid"))?; + let effective_until = raw + .effective_until + .as_deref() + .map(parse_timestamp) + .transpose() + .with_context(|| format!("{path}.effective_until is invalid"))?; + if effective_until.is_some_and(|until| until <= effective_from) { + bail!("{path}.effective_until must be later than effective_from"); + } + + let rates = TokenRates { + input_usd_per_1m_tokens: raw.input_usd_per_1m_tokens, + cached_input_usd_per_1m_tokens: raw.cached_input_usd_per_1m_tokens, + cache_write_usd_per_1m_tokens: raw.cache_write_usd_per_1m_tokens, + cache_write_5m_usd_per_1m_tokens: raw.cache_write_5m_usd_per_1m_tokens, + cache_write_1h_usd_per_1m_tokens: raw.cache_write_1h_usd_per_1m_tokens, + output_usd_per_1m_tokens: raw.output_usd_per_1m_tokens, + }; + validate_rates(&path, &rates)?; + + Ok(ParsedTokenRates { + effective_from, + effective_until, + rates, + }) +} + +fn validate_rates(path: &str, rates: &TokenRates) -> Result<()> { + let values = [ + ( + "input_usd_per_1m_tokens", + Some(rates.input_usd_per_1m_tokens), + ), + ( + "cached_input_usd_per_1m_tokens", + rates.cached_input_usd_per_1m_tokens, + ), + ( + "cache_write_usd_per_1m_tokens", + rates.cache_write_usd_per_1m_tokens, + ), + ( + "cache_write_5m_usd_per_1m_tokens", + rates.cache_write_5m_usd_per_1m_tokens, + ), + ( + "cache_write_1h_usd_per_1m_tokens", + rates.cache_write_1h_usd_per_1m_tokens, + ), + ( + "output_usd_per_1m_tokens", + Some(rates.output_usd_per_1m_tokens), + ), + ]; + for (name, value) in values { + if value.is_some_and(|value| !value.is_finite() || value < 0.0) { + bail!("{path}.{name} must be a finite, non-negative number"); + } + } + + if rates.cache_write_5m_usd_per_1m_tokens.is_some() + != rates.cache_write_1h_usd_per_1m_tokens.is_some() + { + bail!( + "{path} must specify both cache_write_5m_usd_per_1m_tokens and cache_write_1h_usd_per_1m_tokens, or neither" + ); + } + Ok(()) +} + +fn insert_model_lookup( + model_lookup: &mut HashMap, + model_name: &str, + history_index: usize, +) -> Result<()> { + let normalized = normalize_model_name(model_name); + if model_lookup.insert(normalized, history_index).is_some() { + bail!("model name or alias '{model_name}' is defined more than once"); + } + Ok(()) +} + +fn normalize_model_name(model: &str) -> String { + model.trim().to_ascii_lowercase() +} + +pub(super) fn parse_timestamp(value: &str) -> Result> { + let value = value.trim(); + if let Ok(timestamp) = DateTime::parse_from_rfc3339(value) { + return Ok(timestamp.with_timezone(&Utc)); + } + if let Ok(date) = NaiveDate::parse_from_str(value, "%Y-%m-%d") { + return Ok(date + .and_hms_opt(0, 0, 0) + .expect("midnight is a valid time") + .and_utc()); + } + bail!("expected RFC 3339 timestamp or YYYY-MM-DD date, got '{value}'") +} + +pub(super) fn format_timestamp(timestamp: DateTime) -> String { + timestamp.to_rfc3339_opts(chrono::SecondsFormat::AutoSi, true) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn parse_book(input: &str) -> Result { + let raw: RawPricingFile = toml::from_str(input)?; + PriceBook::from_raw(raw) + } + + #[test] + fn historical_rates_apply_at_inclusive_start_and_exclusive_end() { + let book = parse_book( + r#" +version = 1 + +[models."test-model"] +aliases = ["test-deployment"] + +[[models."test-model".rates]] +effective_from = "2025-01-01" +input_usd_per_1m_tokens = 1.0 +output_usd_per_1m_tokens = 2.0 + +[[models."test-model".rates]] +effective_from = "2025-06-01T00:00:00Z" +input_usd_per_1m_tokens = 3.0 +output_usd_per_1m_tokens = 4.0 +"#, + ) + .expect("valid pricing file"); + + let before_change = parse_timestamp("2025-05-31T23:59:59Z").unwrap(); + let at_change = parse_timestamp("2025-06-01T00:00:00Z").unwrap(); + assert_eq!( + book.rate_at("test-model", before_change) + .unwrap() + .input_usd_per_1m_tokens, + 1.0 + ); + assert_eq!( + book.rate_at("TEST-DEPLOYMENT", at_change) + .unwrap() + .input_usd_per_1m_tokens, + 3.0 + ); + } + + #[test] + fn explicit_end_can_leave_an_unpriced_gap() { + let book = parse_book( + r#" +version = 1 + +[models."test-model"] +[[models."test-model".rates]] +effective_from = "2025-01-01" +effective_until = "2025-02-01" +input_usd_per_1m_tokens = 1.0 +output_usd_per_1m_tokens = 2.0 + +[[models."test-model".rates]] +effective_from = "2025-03-01" +input_usd_per_1m_tokens = 3.0 +output_usd_per_1m_tokens = 4.0 +"#, + ) + .expect("valid pricing file"); + + assert!(book + .rate_at("test-model", parse_timestamp("2025-02-15").unwrap()) + .is_none()); + } + + #[test] + fn rejects_overlapping_intervals() { + let error = parse_book( + r#" +version = 1 + +[models."test-model"] +[[models."test-model".rates]] +effective_from = "2025-01-01" +effective_until = "2025-07-01" +input_usd_per_1m_tokens = 1.0 +output_usd_per_1m_tokens = 2.0 + +[[models."test-model".rates]] +effective_from = "2025-06-01" +input_usd_per_1m_tokens = 3.0 +output_usd_per_1m_tokens = 4.0 +"#, + ) + .expect_err("overlap must fail"); + assert!( + error.to_string().contains("overlapping intervals"), + "{error}" + ); + } + + #[test] + fn rejects_only_one_ttl_cache_write_rate() { + let error = parse_book( + r#" +version = 1 + +[models."test-model"] +[[models."test-model".rates]] +effective_from = "2025-01-01" +input_usd_per_1m_tokens = 1.0 +cache_write_5m_usd_per_1m_tokens = 1.25 +output_usd_per_1m_tokens = 2.0 +"#, + ) + .expect_err("incomplete TTL rates must fail"); + assert!(error.to_string().contains("must specify both"), "{error}"); + } + + #[test] + fn computes_generic_and_split_cache_costs() { + let usage = TokenUsage { + spans: 1, + uncached_input_tokens: 1_000_000, + cached_input_tokens: 1_000_000, + effective_cache_write_tokens: 3_000_000, + split_cache_write_5m_tokens: 1_000_000, + split_cache_write_1h_tokens: 1_000_000, + fallback_cache_write_tokens: 1_000_000, + output_tokens: 1_000_000, + }; + let split_rates = TokenRates { + input_usd_per_1m_tokens: 1.0, + cached_input_usd_per_1m_tokens: Some(0.1), + cache_write_usd_per_1m_tokens: Some(1.25), + cache_write_5m_usd_per_1m_tokens: Some(1.25), + cache_write_1h_usd_per_1m_tokens: Some(2.0), + output_usd_per_1m_tokens: 3.0, + }; + assert_eq!(usage.cost(&split_rates), 8.6); + + let generic_rates = TokenRates { + cache_write_5m_usd_per_1m_tokens: None, + cache_write_1h_usd_per_1m_tokens: None, + ..split_rates + }; + assert_eq!(usage.cost(&generic_rates), 7.85); + } + + #[test] + fn boundaries_include_historical_changes_and_explicit_ends() { + let book = parse_book( + r#" +version = 1 + +[models."test-model"] +[[models."test-model".rates]] +effective_from = "2025-01-01" +effective_until = "2025-03-01" +input_usd_per_1m_tokens = 1.0 +output_usd_per_1m_tokens = 2.0 + +[[models."test-model".rates]] +effective_from = "2025-04-01" +input_usd_per_1m_tokens = 3.0 +output_usd_per_1m_tokens = 4.0 +"#, + ) + .unwrap(); + let boundaries = book.boundaries_between( + parse_timestamp("2025-02-01").unwrap(), + parse_timestamp("2025-05-01").unwrap(), + ); + assert_eq!( + boundaries, + vec![ + parse_timestamp("2025-03-01").unwrap(), + parse_timestamp("2025-04-01").unwrap() + ] + ); + } +} diff --git a/tests/cli.rs b/tests/cli.rs index 084add7c..33c36d12 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -49,6 +49,19 @@ fn write_auth_store(config_home: &Path, profiles: &[(&str, &str)]) { fs::write(auth_dir.join("auth.json"), body).expect("write auth store"); } +#[test] +fn cost_without_subcommand_prints_help_without_authentication() { + let mut command = bt_command(); + clear_braintrust_auth_env(&mut command); + command + .arg("cost") + .assert() + .success() + .stdout(predicate::str::contains("Usage: bt cost")) + .stdout(predicate::str::contains("experiments")) + .stdout(predicate::str::contains("logs")); +} + #[test] fn auth_login_does_not_expose_client_id() { bt_command() From 3809041cc56ced24f1cb6a3ab2b76be50e5210f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Wed, 22 Jul 2026 10:43:41 -0700 Subject: [PATCH 20/24] draft2 --- Cargo.lock | 52 +++ Cargo.toml | 1 + src/cost/experiments.rs | 435 ---------------------- src/cost/logs.rs | 796 ---------------------------------------- src/cost/mod.rs | 590 +++++++++++++++++++++++++++-- src/cost/pricing.rs | 30 +- src/main.rs | 4 +- src/sql.rs | 101 ++++- tests/cli.rs | 10 +- 9 files changed, 727 insertions(+), 1292 deletions(-) delete mode 100644 src/cost/experiments.rs delete mode 100644 src/cost/logs.rs diff --git a/Cargo.lock b/Cargo.lock index cefa77c6..224ea879 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -534,6 +534,7 @@ dependencies = [ "sha2", "strip-ansi-escapes", "tempfile", + "textplots", "tokio", "toml", "unicode-width 0.1.14", @@ -548,6 +549,12 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" + [[package]] name = "bytes" version = "1.11.1" @@ -662,6 +669,16 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "colored" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" +dependencies = [ + "lazy_static", + "windows-sys 0.59.0", +] + [[package]] name = "comfy-table" version = "7.2.2" @@ -1055,6 +1072,16 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "drawille" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e64e461c3f1e69d99372620640b3fd5f0309eeda2e26e4af69f6760c0e1df845" +dependencies = [ + "colored", + "fnv", +] + [[package]] name = "dunce" version = "1.0.5" @@ -1789,6 +1816,12 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4345964bb142484797b161f473a503a434de77149dd8c7427788c6e13379388" +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "leb128fmt" version = "0.1.0" @@ -2548,6 +2581,15 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "rgb" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +dependencies = [ + "bytemuck", +] + [[package]] name = "ring" version = "0.17.14" @@ -3100,6 +3142,16 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" +[[package]] +name = "textplots" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f7657a0066c9f9663659db0665319adff8b0943305fc73eddf1010e5a2072b1" +dependencies = [ + "drawille", + "rgb", +] + [[package]] name = "thiserror" version = "1.0.69" diff --git a/Cargo.toml b/Cargo.toml index a7da2549..7471e0c9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -52,6 +52,7 @@ glob = "0.3" flate2 = "1.1.2" tempfile = "3" uuid = { version = "1.21.0", features = ["v4"] } +textplots = "0.8.7" [profile.dist] inherits = "release" diff --git a/src/cost/experiments.rs b/src/cost/experiments.rs deleted file mode 100644 index 5f8eb3d1..00000000 --- a/src/cost/experiments.rs +++ /dev/null @@ -1,435 +0,0 @@ -use std::collections::HashMap; -use std::fmt::Write as _; - -use anyhow::Result; -use dialoguer::console; -use serde::Serialize; -use serde_json::{Map, Value}; - -use crate::{ - experiments::api::list_experiments, - sql::run_btql_rows, - ui::{apply_column_padding, header, print_with_pager, styled_table, truncate, with_spinner}, - utils::{format_cost, pluralize}, -}; - -use super::ResolvedContext; - -/// Per-experiment cost stats returned by the aggregate query. -#[derive(Debug, Clone, Copy, Default, PartialEq)] -struct CostStats { - /// `SUM(estimated_cost())` — null when nothing could be priced. - cost: Option, - /// Spans with a non-null `estimated_cost()` (the ones feeding the sum). - priced_spans: u64, - /// Non-scorer spans that clearly involved an LLM (llm span or token - /// metrics) but could not be priced — the "we don't know" signal. - unpriced_llm_spans: u64, -} - -/// `estimated_cost()` returns NULL for spans it cannot price (no logged cost and -/// no model-registry match) and for scorer spans. Comparing priced vs. unpriced -/// LLM spans lets us report cost coverage instead of a silently-wrong `$0`. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] -#[serde(rename_all = "snake_case")] -enum Coverage { - /// Every LLM span was priced. - Full, - /// Some, but not all, LLM spans were priced. - Partial, - /// LLM spans exist but none could be priced (e.g. missing model pricing). - Unknown, - /// No priced spans and no detectable unpriced LLM activity — no cost data. - None, -} - -impl Coverage { - fn classify(priced_spans: u64, unpriced_llm_spans: u64) -> Self { - match (priced_spans, unpriced_llm_spans) { - (0, 0) => Coverage::None, - (0, _) => Coverage::Unknown, - (_, 0) => Coverage::Full, - (_, _) => Coverage::Partial, - } - } -} - -#[derive(Debug, Clone, Serialize)] -struct ExperimentCostRow { - id: String, - name: String, - created: Option, - /// Estimated cost in USD. `null` when we could not price anything. - cost: Option, - priced_spans: u64, - unpriced_llm_spans: u64, - coverage: Coverage, -} - -pub(crate) async fn run(ctx: &ResolvedContext, json: bool) -> Result<()> { - let project_name = &ctx.project.name; - let experiments = with_spinner( - "Loading experiments...", - list_experiments(&ctx.client, project_name), - ) - .await?; - - if experiments.is_empty() { - if json { - println!("[]"); - } else { - println!("No experiments found in {project_name}"); - } - return Ok(()); - } - - let experiment_ids: Vec = experiments.iter().map(|e| e.id.clone()).collect(); - let query = build_cost_query(&experiment_ids); - let result_rows = with_spinner( - "Estimating cost...", - run_btql_rows(&ctx.client, &query, "strict"), - ) - .await?; - - let stats_by_id = index_cost_rows(&result_rows); - - let mut rows: Vec = experiments - .iter() - .map(|exp| { - let stats = stats_by_id.get(&exp.id).copied().unwrap_or_default(); - let coverage = Coverage::classify(stats.priced_spans, stats.unpriced_llm_spans); - ExperimentCostRow { - id: exp.id.clone(), - name: exp.name.clone(), - created: exp.created.clone(), - cost: stats.cost, - priced_spans: stats.priced_spans, - unpriced_llm_spans: stats.unpriced_llm_spans, - coverage, - } - }) - .collect(); - - // Highest cost first; experiments with no known cost sort last, tie-broken by name. - rows.sort_by(|a, b| { - let a_key = a.cost.unwrap_or(f64::NEG_INFINITY); - let b_key = b.cost.unwrap_or(f64::NEG_INFINITY); - b_key - .partial_cmp(&a_key) - .unwrap_or(std::cmp::Ordering::Equal) - .then_with(|| a.name.cmp(&b.name)) - }); - - if json { - println!("{}", serde_json::to_string(&rows)?); - return Ok(()); - } - - print_table(ctx, &rows)?; - Ok(()) -} - -fn build_cost_query(experiment_ids: &[String]) -> String { - let id_list = experiment_ids - .iter() - .map(|id| format!("'{}'", id.replace('\'', "''"))) - .collect::>() - .join(", "); - - // `SUM(estimated_cost())` matches the web UI's experiment cost (no span-type - // filter; scorer spans are already unpriced by `estimated_cost()`). - // `unpriced_llm_spans` counts non-scorer spans that had LLM activity (llm - // span type or token metrics) yet could not be priced. - format!( - "SELECT \ - experiment_id, \ - SUM(estimated_cost()) AS cost, \ - COUNT(estimated_cost()) AS priced_spans, \ - COUNT(CASE \ - WHEN (span_attributes.purpose IS NULL OR span_attributes.purpose != 'scorer') \ - AND estimated_cost() IS NULL \ - AND (span_attributes.type = 'llm' \ - OR metrics.prompt_tokens IS NOT NULL \ - OR metrics.completion_tokens IS NOT NULL) \ - THEN 1 END) AS unpriced_llm_spans \ - FROM experiment({id_list}) \ - GROUP BY experiment_id" - ) -} - -fn index_cost_rows(result_rows: &[Map]) -> HashMap { - let mut stats = HashMap::new(); - for row in result_rows { - let Some(id) = row.get("experiment_id").and_then(Value::as_str) else { - continue; - }; - stats.insert( - id.to_string(), - CostStats { - cost: value_as_opt_f64(row.get("cost")), - priced_spans: value_as_u64(row.get("priced_spans")), - unpriced_llm_spans: value_as_u64(row.get("unpriced_llm_spans")), - }, - ); - } - stats -} - -fn print_table(ctx: &ResolvedContext, rows: &[ExperimentCostRow]) -> Result<()> { - let mut output = String::new(); - let count = format!( - "{} {}", - rows.len(), - pluralize(rows.len(), "experiment", None) - ); - writeln!( - output, - "{} in {} {} {}\n", - console::style(count), - console::style(ctx.client.org_name()).bold(), - console::style("/").dim().bold(), - console::style(&ctx.project.name).bold() - )?; - - let mut table = styled_table(); - table.set_header(vec![ - header("Name"), - header("Created"), - header("Cost"), - header("Coverage"), - ]); - apply_column_padding(&mut table, (0, 4)); - - for row in rows { - let created = row - .created - .as_deref() - .map(|c| truncate(c, 10)) - .unwrap_or_else(|| "-".to_string()); - table.add_row(vec![ - truncate(&row.name, 60), - created, - cost_display(row.cost, row.coverage), - coverage_display(row.priced_spans, row.unpriced_llm_spans, row.coverage), - ]); - } - - write!(output, "{table}")?; - - let priced_total: f64 = rows.iter().filter_map(|r| r.cost).sum(); - let partial = rows - .iter() - .filter(|r| r.coverage == Coverage::Partial) - .count(); - let unknown = rows - .iter() - .filter(|r| r.coverage == Coverage::Unknown) - .count(); - let no_data = rows.iter().filter(|r| r.coverage == Coverage::None).count(); - - write!( - output, - "\n\nTotal (priced): {}", - console::style(format_cost(priced_total)).bold() - )?; - if partial > 0 { - write!( - output, - " {}", - console::style(format!("{partial} partial")).yellow() - )?; - } - if unknown > 0 { - write!( - output, - " {}", - console::style(format!("{unknown} unknown")).yellow() - )?; - } - if no_data > 0 { - write!( - output, - " {}", - console::style(format!("{no_data} without LLM cost data")).dim() - )?; - } - output.push('\n'); - - print_with_pager(&output)?; - Ok(()) -} - -fn cost_display(cost: Option, coverage: Coverage) -> String { - match coverage { - Coverage::Full => format_cost(cost.unwrap_or(0.0)), - Coverage::Partial => format!("~{}", format_cost(cost.unwrap_or(0.0))), - Coverage::Unknown => "unknown".to_string(), - Coverage::None => "n/a".to_string(), - } -} - -fn coverage_display(priced_spans: u64, unpriced_llm_spans: u64, coverage: Coverage) -> String { - match coverage { - Coverage::None => "—".to_string(), - _ => format!("{}/{}", priced_spans, priced_spans + unpriced_llm_spans), - } -} - -fn value_as_opt_f64(value: Option<&Value>) -> Option { - match value { - Some(Value::Number(n)) => n.as_f64(), - Some(Value::String(s)) => s.parse().ok(), - _ => None, - } -} - -fn value_as_u64(value: Option<&Value>) -> u64 { - match value { - Some(Value::Number(n)) => n - .as_u64() - .or_else(|| n.as_f64().map(|f| f.max(0.0) as u64)) - .unwrap_or(0), - Some(Value::String(s)) => s.parse().unwrap_or(0), - _ => 0, - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn build_cost_query_lists_ids_and_matches_ui_sum() { - let query = build_cost_query(&["exp-a".to_string(), "exp-b".to_string()]); - assert!( - query.contains("FROM experiment('exp-a', 'exp-b')"), - "{query}" - ); - assert!(query.contains("SUM(estimated_cost()) AS cost"), "{query}"); - assert!( - query.contains("COUNT(estimated_cost()) AS priced_spans"), - "{query}" - ); - assert!(query.contains("AS unpriced_llm_spans"), "{query}"); - assert!( - query.contains("span_attributes.purpose != 'scorer'"), - "{query}" - ); - assert!(query.contains("span_attributes.type = 'llm'"), "{query}"); - assert!( - query.contains("metrics.prompt_tokens IS NOT NULL"), - "{query}" - ); - assert!(query.contains("GROUP BY experiment_id"), "{query}"); - // The cost sum must not be restricted by span type, to match the web UI. - assert!( - !query.contains("WHERE"), - "cost sum should not be filtered: {query}" - ); - } - - #[test] - fn build_cost_query_escapes_single_quotes() { - let query = build_cost_query(&["ex'p".to_string()]); - assert!(query.contains("experiment('ex''p')"), "{query}"); - } - - #[test] - fn classify_covers_all_cases() { - assert_eq!(Coverage::classify(0, 0), Coverage::None); - assert_eq!(Coverage::classify(0, 5), Coverage::Unknown); - assert_eq!(Coverage::classify(3, 2), Coverage::Partial); - assert_eq!(Coverage::classify(5, 0), Coverage::Full); - } - - #[test] - fn cost_and_coverage_display_reflect_classification() { - assert_eq!(cost_display(Some(1.5), Coverage::Full), "$1.50"); - assert_eq!(cost_display(Some(1.5), Coverage::Partial), "~$1.50"); - assert_eq!(cost_display(None, Coverage::Unknown), "unknown"); - assert_eq!(cost_display(None, Coverage::None), "n/a"); - - assert_eq!(coverage_display(5, 0, Coverage::Full), "5/5"); - assert_eq!(coverage_display(3, 2, Coverage::Partial), "3/5"); - assert_eq!(coverage_display(0, 4, Coverage::Unknown), "0/4"); - assert_eq!(coverage_display(0, 0, Coverage::None), "—"); - } - - #[test] - fn index_cost_rows_parses_null_cost_and_counts() { - let rows = vec![ - json!({ - "experiment_id": "exp-a", - "cost": 1.25, - "priced_spans": 8, - "unpriced_llm_spans": 2 - }) - .as_object() - .unwrap() - .clone(), - json!({ - "experiment_id": "exp-b", - "cost": null, - "priced_spans": 0, - "unpriced_llm_spans": 0 - }) - .as_object() - .unwrap() - .clone(), - ]; - let stats = index_cost_rows(&rows); - assert_eq!( - stats.get("exp-a"), - Some(&CostStats { - cost: Some(1.25), - priced_spans: 8, - unpriced_llm_spans: 2 - }) - ); - assert_eq!( - stats.get("exp-b"), - Some(&CostStats { - cost: None, - priced_spans: 0, - unpriced_llm_spans: 0 - }) - ); - } - - #[test] - fn experiment_cost_row_json_shape() { - let row = ExperimentCostRow { - id: "exp-a".to_string(), - name: "baseline".to_string(), - created: Some("2024-01-02T03:04:05Z".to_string()), - cost: Some(1.25), - priced_spans: 8, - unpriced_llm_spans: 2, - coverage: Coverage::Partial, - }; - let value = serde_json::to_value(&row).unwrap(); - assert_eq!(value["id"], "exp-a"); - assert_eq!(value["name"], "baseline"); - assert_eq!(value["cost"], 1.25); - assert_eq!(value["priced_spans"], 8); - assert_eq!(value["unpriced_llm_spans"], 2); - assert_eq!(value["coverage"], "partial"); - } - - #[test] - fn experiment_cost_row_json_null_cost_for_no_data() { - let row = ExperimentCostRow { - id: "exp-z".to_string(), - name: "empty".to_string(), - created: None, - cost: None, - priced_spans: 0, - unpriced_llm_spans: 0, - coverage: Coverage::None, - }; - let value = serde_json::to_value(&row).unwrap(); - assert!(value["cost"].is_null()); - assert_eq!(value["coverage"], "none"); - } -} diff --git a/src/cost/logs.rs b/src/cost/logs.rs deleted file mode 100644 index bda05991..00000000 --- a/src/cost/logs.rs +++ /dev/null @@ -1,796 +0,0 @@ -use std::cmp::Ordering; -use std::fmt::Write as _; -use std::path::PathBuf; - -use anyhow::{bail, Context, Result}; -use chrono::{DateTime, Duration, Timelike, Utc}; -use clap::{builder::BoolishValueParser, Args}; -use dialoguer::console; -use serde::Serialize; -use serde_json::{Map, Value}; - -use crate::{ - sql::run_btql_rows, - ui::{apply_column_padding, header, print_with_pager, styled_table, truncate, with_spinner}, - utils::{format_cost, parse_duration_to_seconds, pluralize}, -}; - -use super::{ - pricing::{format_timestamp, parse_timestamp, PriceBook, TokenUsage}, - ResolvedContext, -}; - -#[derive(Debug, Clone, Args)] -#[command(after_help = "\ -Pricing file (USD per 1 million tokens): - version = 1 - - [models.\"custom-chat-model\"] - aliases = [\"custom-deployment-name\"] - - [[models.\"custom-chat-model\".rates]] - effective_from = \"2025-01-01T00:00:00Z\" - effective_until = \"2025-06-01T00:00:00Z\" - input_usd_per_1m_tokens = 3.0 - cached_input_usd_per_1m_tokens = 0.3 - cache_write_usd_per_1m_tokens = 3.75 - cache_write_5m_usd_per_1m_tokens = 3.75 - cache_write_1h_usd_per_1m_tokens = 6.0 - output_usd_per_1m_tokens = 15.0 - -Repeat [[models.\"...\".rates]] for historical prices. Bounds are [from, until). -When effective_until is omitted, a rate ends at the next effective_from, or never. -Braintrust's effective estimated cost takes precedence; file rates fill unpriced token spans.")] -pub(crate) struct LogsArgs { - /// Relative time window ending at --until - #[arg(long, env = "BRAINTRUST_COST_WINDOW", default_value = "7d")] - window: String, - - /// Absolute inclusive lower bound (RFC 3339 or YYYY-MM-DD); overrides --window - #[arg(long, env = "BRAINTRUST_COST_SINCE")] - since: Option, - - /// Absolute exclusive upper bound (RFC 3339 or YYYY-MM-DD); defaults to now - #[arg(long, env = "BRAINTRUST_COST_UNTIL")] - until: Option, - - /// TOML file containing historical per-model token prices - #[arg(long, env = "BRAINTRUST_COST_PRICING_FILE", value_name = "PATH")] - pricing_file: Option, - - /// Exclude spans whose purpose is scorer - #[arg( - long, - env = "BRAINTRUST_COST_EXCLUDE_SCORERS", - value_parser = BoolishValueParser::new(), - default_value_t = false - )] - exclude_scorers: bool, -} - -impl Default for LogsArgs { - fn default() -> Self { - Self { - window: "7d".to_string(), - since: None, - until: None, - pricing_file: None, - exclude_scorers: false, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct TimeRange { - since: DateTime, - until: DateTime, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct TimeSegment { - since: DateTime, - until: DateTime, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] -#[serde(rename_all = "snake_case")] -enum Coverage { - Full, - Partial, - Unknown, - None, -} - -impl Coverage { - fn classify(priced_spans: u64, unpriced_token_spans: u64) -> Self { - match (priced_spans, unpriced_token_spans) { - (0, 0) => Coverage::None, - (0, _) => Coverage::Unknown, - (_, 0) => Coverage::Full, - (_, _) => Coverage::Partial, - } - } -} - -#[derive(Debug, Clone, Serialize)] -struct LogsCostRow { - model: Option, - purpose: Option, - candidate_spans: u64, - cost: Option, - braintrust_cost: Option, - file_cost: Option, - braintrust_priced_spans: u64, - file_priced_spans: u64, - unpriced_token_spans: u64, - no_usage_spans: u64, - coverage: Coverage, -} - -impl LogsCostRow { - fn priced_spans(&self) -> u64 { - self.braintrust_priced_spans + self.file_priced_spans - } -} - -#[derive(Debug, Clone, Serialize)] -struct LogsCostTotals { - candidate_spans: u64, - cost: Option, - braintrust_cost: Option, - file_cost: Option, - braintrust_priced_spans: u64, - file_priced_spans: u64, - unpriced_token_spans: u64, - no_usage_spans: u64, - coverage: Coverage, -} - -#[derive(Debug, Serialize)] -struct LogsCostOutput<'a> { - project: &'a str, - org: &'a str, - currency: &'static str, - since: String, - until: String, - excludes_scorers: bool, - pricing_file: Option, - rows: &'a [LogsCostRow], - totals: &'a LogsCostTotals, -} - -pub(crate) async fn run(ctx: &ResolvedContext, args: LogsArgs, json: bool) -> Result<()> { - let now = Utc::now() - .with_nanosecond(0) - .expect("zero nanoseconds is a valid timestamp"); - let range = resolve_time_range(&args, now)?; - let price_book = args - .pricing_file - .as_deref() - .map(PriceBook::load) - .transpose()?; - let segments = build_time_segments(range, price_book.as_ref()); - let query = build_cost_query(&ctx.project.id, range, &segments, args.exclude_scorers); - let result_rows = with_spinner( - "Estimating log cost...", - run_btql_rows(&ctx.client, &query, "default"), - ) - .await?; - - let mut rows = build_cost_rows(&result_rows, &segments, price_book.as_ref())?; - rows.sort_by(compare_cost_rows); - let totals = calculate_totals(&rows); - - if json { - let output = LogsCostOutput { - project: &ctx.project.name, - org: ctx.client.org_name(), - currency: "USD", - since: format_timestamp(range.since), - until: format_timestamp(range.until), - excludes_scorers: args.exclude_scorers, - pricing_file: args - .pricing_file - .as_ref() - .map(|path| path.display().to_string()), - rows: &rows, - totals: &totals, - }; - println!("{}", serde_json::to_string(&output)?); - return Ok(()); - } - - print_table(ctx, range, args.pricing_file.as_ref(), &rows, &totals)?; - Ok(()) -} - -fn resolve_time_range(args: &LogsArgs, now: DateTime) -> Result { - let until = args - .until - .as_deref() - .map(parse_timestamp) - .transpose() - .context("invalid --until")? - .unwrap_or(now); - let since = match args.since.as_deref() { - Some(since) => parse_timestamp(since).context("invalid --since")?, - None => { - let seconds = parse_duration_to_seconds(&args.window) - .with_context(|| format!("invalid --window '{}'", args.window))?; - if seconds == 0 { - bail!("--window must be greater than zero"); - } - let seconds = i64::try_from(seconds).context("--window is too large")?; - until - .checked_sub_signed(Duration::seconds(seconds)) - .context("--window produces a timestamp outside the supported range")? - } - }; - if since >= until { - bail!("--since must be earlier than --until"); - } - Ok(TimeRange { since, until }) -} - -fn build_time_segments(range: TimeRange, price_book: Option<&PriceBook>) -> Vec { - let mut boundaries = vec![range.since, range.until]; - if let Some(price_book) = price_book { - boundaries.extend(price_book.boundaries_between(range.since, range.until)); - } - boundaries.sort_unstable(); - boundaries.dedup(); - boundaries - .windows(2) - .map(|bounds| TimeSegment { - since: bounds[0], - until: bounds[1], - }) - .collect() -} - -fn build_cost_query( - project_id: &str, - range: TimeRange, - segments: &[TimeSegment], - exclude_scorers: bool, -) -> String { - let prompt_tokens = "COALESCE(metrics.prompt_tokens, 0)"; - let completion_tokens = "COALESCE(metrics.completion_tokens, 0)"; - let cached_tokens = "COALESCE(metrics.prompt_cached_tokens, 0)"; - let generic_write_tokens = "COALESCE(metrics.prompt_cache_creation_tokens, 0)"; - let write_5m_tokens = "COALESCE(metrics.prompt_cache_creation_5m_tokens, 0)"; - let write_1h_tokens = "COALESCE(metrics.prompt_cache_creation_1h_tokens, 0)"; - let split_write_tokens = format!("({write_5m_tokens} + {write_1h_tokens})"); - let effective_write_tokens = format!("GREATEST({generic_write_tokens}, {split_write_tokens})"); - let uncached_input_tokens = - format!("GREATEST(0, {prompt_tokens} - {cached_tokens} - {effective_write_tokens})"); - let token_activity = [ - "metrics.prompt_tokens IS NOT NULL", - "metrics.completion_tokens IS NOT NULL", - "metrics.prompt_cached_tokens IS NOT NULL", - "metrics.prompt_cache_creation_tokens IS NOT NULL", - "metrics.prompt_cache_creation_5m_tokens IS NOT NULL", - "metrics.prompt_cache_creation_1h_tokens IS NOT NULL", - ] - .join(" OR "); - - let mut select_fields = vec![ - "metadata.model AS model".to_string(), - "span_attributes.purpose AS purpose".to_string(), - format!( - "COUNT(CASE WHEN estimated_cost() IS NOT NULL OR metadata.model IS NOT NULL OR ({token_activity}) THEN 1 END) AS candidate_spans" - ), - "COUNT(estimated_cost()) AS braintrust_priced_spans".to_string(), - "SUM(estimated_cost()) AS braintrust_cost".to_string(), - format!( - "COUNT(CASE WHEN estimated_cost() IS NULL AND ({token_activity}) THEN 1 END) AS unpriced_token_spans" - ), - format!( - "COUNT(CASE WHEN estimated_cost() IS NULL AND metadata.model IS NOT NULL AND NOT ({token_activity}) THEN 1 END) AS no_usage_spans" - ), - ]; - - for (index, segment) in segments.iter().enumerate() { - let condition = format!( - "estimated_cost() IS NULL AND created >= {} AND created < {} AND ({token_activity})", - sql_quote(&format_timestamp(segment.since)), - sql_quote(&format_timestamp(segment.until)), - ); - let split_complete = format!("{split_write_tokens} >= {generic_write_tokens}"); - let fallback_write_tokens = - format!("CASE WHEN {split_complete} THEN 0 ELSE {effective_write_tokens} END"); - let split_5m = format!("CASE WHEN {split_complete} THEN {write_5m_tokens} ELSE 0 END"); - let split_1h = format!("CASE WHEN {split_complete} THEN {write_1h_tokens} ELSE 0 END"); - - select_fields.extend([ - format!("COUNT(CASE WHEN {condition} THEN 1 END) AS p{index}_spans"), - format!( - "SUM(CASE WHEN {condition} THEN {uncached_input_tokens} ELSE 0 END) AS p{index}_uncached_input_tokens" - ), - format!( - "SUM(CASE WHEN {condition} THEN {cached_tokens} ELSE 0 END) AS p{index}_cached_input_tokens" - ), - format!( - "SUM(CASE WHEN {condition} THEN {effective_write_tokens} ELSE 0 END) AS p{index}_effective_cache_write_tokens" - ), - format!( - "SUM(CASE WHEN {condition} THEN {split_5m} ELSE 0 END) AS p{index}_split_cache_write_5m_tokens" - ), - format!( - "SUM(CASE WHEN {condition} THEN {split_1h} ELSE 0 END) AS p{index}_split_cache_write_1h_tokens" - ), - format!( - "SUM(CASE WHEN {condition} THEN {fallback_write_tokens} ELSE 0 END) AS p{index}_fallback_cache_write_tokens" - ), - format!( - "SUM(CASE WHEN {condition} THEN {completion_tokens} ELSE 0 END) AS p{index}_output_tokens" - ), - ]); - } - - let scorer_filter = if exclude_scorers { - "\n AND (span_attributes.purpose IS NULL OR span_attributes.purpose != 'scorer')" - } else { - "" - }; - format!( - "SELECT\n {}\nFROM project_logs({}, shape => 'spans')\nWHERE created >= {}\n AND created < {}{}\nGROUP BY metadata.model, span_attributes.purpose", - select_fields.join(",\n "), - sql_quote(project_id), - sql_quote(&format_timestamp(range.since)), - sql_quote(&format_timestamp(range.until)), - scorer_filter, - ) -} - -fn build_cost_rows( - result_rows: &[Map], - segments: &[TimeSegment], - price_book: Option<&PriceBook>, -) -> Result> { - result_rows - .iter() - .map(|row| build_cost_row(row, segments, price_book)) - .filter_map(|result| match result { - Ok(row) if row.candidate_spans == 0 => None, - other => Some(other), - }) - .collect() -} - -fn build_cost_row( - row: &Map, - segments: &[TimeSegment], - price_book: Option<&PriceBook>, -) -> Result { - let model = row.get("model").and_then(Value::as_str).map(str::to_string); - let purpose = row - .get("purpose") - .and_then(Value::as_str) - .map(str::to_string); - let candidate_spans = value_as_u64(row.get("candidate_spans")); - let braintrust_priced_spans = value_as_u64(row.get("braintrust_priced_spans")); - let braintrust_cost = value_as_opt_f64(row.get("braintrust_cost")); - let expected_unpriced_token_spans = value_as_u64(row.get("unpriced_token_spans")); - let no_usage_spans = value_as_u64(row.get("no_usage_spans")); - - let mut file_cost = 0.0; - let mut file_priced_spans = 0_u64; - let mut unpriced_token_spans = 0_u64; - let mut segmented_token_spans = 0_u64; - - for (index, segment) in segments.iter().enumerate() { - let usage = token_usage_from_row(row, index); - segmented_token_spans = segmented_token_spans.saturating_add(usage.spans); - let rates = model.as_deref().and_then(|model| { - price_book.and_then(|price_book| price_book.rate_at(model, segment.since)) - }); - if let Some(rates) = rates { - file_cost += usage.cost(rates); - file_priced_spans = file_priced_spans.saturating_add(usage.spans); - } else { - unpriced_token_spans = unpriced_token_spans.saturating_add(usage.spans); - } - } - - if segmented_token_spans != expected_unpriced_token_spans { - bail!( - "cost query returned inconsistent token coverage for model {}: expected {}, got {}", - model.as_deref().unwrap_or(""), - expected_unpriced_token_spans, - segmented_token_spans - ); - } - - let priced_spans = braintrust_priced_spans.saturating_add(file_priced_spans); - let cost = (priced_spans > 0).then(|| braintrust_cost.unwrap_or(0.0) + file_cost); - Ok(LogsCostRow { - model, - purpose, - candidate_spans, - cost, - braintrust_cost: (braintrust_priced_spans > 0).then(|| braintrust_cost.unwrap_or(0.0)), - file_cost: (file_priced_spans > 0).then_some(file_cost), - braintrust_priced_spans, - file_priced_spans, - unpriced_token_spans, - no_usage_spans, - coverage: Coverage::classify(priced_spans, unpriced_token_spans), - }) -} - -fn token_usage_from_row(row: &Map, index: usize) -> TokenUsage { - let value = |suffix: &str| value_as_u64(row.get(&format!("p{index}_{suffix}"))); - TokenUsage { - spans: value("spans"), - uncached_input_tokens: value("uncached_input_tokens"), - cached_input_tokens: value("cached_input_tokens"), - effective_cache_write_tokens: value("effective_cache_write_tokens"), - split_cache_write_5m_tokens: value("split_cache_write_5m_tokens"), - split_cache_write_1h_tokens: value("split_cache_write_1h_tokens"), - fallback_cache_write_tokens: value("fallback_cache_write_tokens"), - output_tokens: value("output_tokens"), - } -} - -fn calculate_totals(rows: &[LogsCostRow]) -> LogsCostTotals { - let braintrust_priced_spans = rows.iter().map(|row| row.braintrust_priced_spans).sum(); - let file_priced_spans = rows.iter().map(|row| row.file_priced_spans).sum(); - let unpriced_token_spans = rows.iter().map(|row| row.unpriced_token_spans).sum(); - let priced_spans = braintrust_priced_spans + file_priced_spans; - let braintrust_cost = (braintrust_priced_spans > 0) - .then(|| rows.iter().filter_map(|row| row.braintrust_cost).sum()); - let file_cost = - (file_priced_spans > 0).then(|| rows.iter().filter_map(|row| row.file_cost).sum()); - LogsCostTotals { - candidate_spans: rows.iter().map(|row| row.candidate_spans).sum(), - cost: (priced_spans > 0).then(|| braintrust_cost.unwrap_or(0.0) + file_cost.unwrap_or(0.0)), - braintrust_cost, - file_cost, - braintrust_priced_spans, - file_priced_spans, - unpriced_token_spans, - no_usage_spans: rows.iter().map(|row| row.no_usage_spans).sum(), - coverage: Coverage::classify(priced_spans, unpriced_token_spans), - } -} - -fn compare_cost_rows(left: &LogsCostRow, right: &LogsCostRow) -> Ordering { - let left_cost = left.cost.unwrap_or(f64::NEG_INFINITY); - let right_cost = right.cost.unwrap_or(f64::NEG_INFINITY); - right_cost - .partial_cmp(&left_cost) - .unwrap_or(Ordering::Equal) - .then_with(|| left.model.cmp(&right.model)) - .then_with(|| left.purpose.cmp(&right.purpose)) -} - -fn print_table( - ctx: &ResolvedContext, - range: TimeRange, - pricing_file: Option<&PathBuf>, - rows: &[LogsCostRow], - totals: &LogsCostTotals, -) -> Result<()> { - let mut output = String::new(); - let count = format!( - "{} {}", - rows.len(), - pluralize(rows.len(), "cost group", None) - ); - writeln!( - output, - "{} in {} {} {}", - console::style(count), - console::style(ctx.client.org_name()).bold(), - console::style("/").dim().bold(), - console::style(&ctx.project.name).bold() - )?; - writeln!( - output, - "{} {} {}", - console::style(format_timestamp(range.since)).dim(), - console::style("to").dim(), - console::style(format_timestamp(range.until)).dim() - )?; - if let Some(path) = pricing_file { - writeln!(output, "Pricing: {}", console::style(path.display()).dim())?; - } - output.push('\n'); - - let mut table = styled_table(); - table.set_header(vec![ - header("Model"), - header("Purpose"), - header("Cost"), - header("Coverage"), - header("Sources"), - ]); - apply_column_padding(&mut table, (0, 4)); - - for row in rows { - table.add_row(vec![ - truncate(row.model.as_deref().unwrap_or(""), 50), - truncate(row.purpose.as_deref().unwrap_or("default"), 20), - cost_display(row.cost, row.coverage), - coverage_display(row.priced_spans(), row.unpriced_token_spans, row.coverage), - source_display(row), - ]); - } - write!(output, "{table}")?; - - let total_cost = totals - .cost - .map(format_cost) - .unwrap_or_else(|| "n/a".to_string()); - let total_prefix = if totals.unpriced_token_spans > 0 { - "~" - } else { - "" - }; - write!( - output, - "\n\nTotal: {}", - console::style(format!("{total_prefix}{total_cost}")).bold() - )?; - if let Some(file_cost) = totals.file_cost { - write!( - output, - " {}", - console::style(format!("{} from pricing file", format_cost(file_cost))).dim() - )?; - } - if totals.unpriced_token_spans > 0 { - write!( - output, - " {}", - console::style(format!( - "{} unpriced token {}", - totals.unpriced_token_spans, - pluralize(totals.unpriced_token_spans as usize, "span", None) - )) - .yellow() - )?; - } - if totals.no_usage_spans > 0 { - write!( - output, - " {}", - console::style(format!("{} without cost or usage", totals.no_usage_spans)).dim() - )?; - } - output.push('\n'); - - print_with_pager(&output)?; - Ok(()) -} - -fn cost_display(cost: Option, coverage: Coverage) -> String { - match (cost, coverage) { - (Some(cost), Coverage::Partial) => format!("~{}", format_cost(cost)), - (Some(cost), _) => format_cost(cost), - (None, Coverage::Unknown) => "unknown".to_string(), - (None, _) => "n/a".to_string(), - } -} - -fn coverage_display(priced_spans: u64, unpriced_token_spans: u64, coverage: Coverage) -> String { - match coverage { - Coverage::None => "—".to_string(), - _ => format!("{}/{}", priced_spans, priced_spans + unpriced_token_spans), - } -} - -fn source_display(row: &LogsCostRow) -> String { - let mut parts = Vec::new(); - if row.braintrust_priced_spans > 0 { - parts.push(format!("Braintrust {}", row.braintrust_priced_spans)); - } - if row.file_priced_spans > 0 { - parts.push(format!("file {}", row.file_priced_spans)); - } - if row.no_usage_spans > 0 { - parts.push(format!("no usage {}", row.no_usage_spans)); - } - if parts.is_empty() { - "—".to_string() - } else { - parts.join(", ") - } -} - -fn value_as_opt_f64(value: Option<&Value>) -> Option { - match value { - Some(Value::Number(number)) => number.as_f64(), - Some(Value::String(value)) => value.parse().ok(), - _ => None, - } - .filter(|value| value.is_finite() && *value >= 0.0) -} - -fn value_as_u64(value: Option<&Value>) -> u64 { - match value { - Some(Value::Number(number)) => number - .as_u64() - .or_else(|| number.as_f64().map(|value| value.max(0.0) as u64)) - .unwrap_or(0), - Some(Value::String(value)) => value.parse().unwrap_or(0), - _ => 0, - } -} - -fn sql_quote(value: &str) -> String { - format!("'{}'", value.replace('\'', "''")) -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - - fn range() -> TimeRange { - TimeRange { - since: parse_timestamp("2025-01-01").unwrap(), - until: parse_timestamp("2025-03-01").unwrap(), - } - } - - #[test] - fn query_is_timestamp_bounded_and_uses_spans() { - let query = build_cost_query( - "test-project-id", - range(), - &[TimeSegment { - since: range().since, - until: range().until, - }], - false, - ); - assert!( - query.contains("FROM project_logs('test-project-id', shape => 'spans')"), - "{query}" - ); - assert!( - query.contains("created >= '2025-01-01T00:00:00Z'"), - "{query}" - ); - assert!( - query.contains("created < '2025-03-01T00:00:00Z'"), - "{query}" - ); - assert!( - query.contains("SUM(estimated_cost()) AS braintrust_cost"), - "{query}" - ); - assert!(query.contains("p0_uncached_input_tokens"), "{query}"); - assert!(!query.contains("purpose != 'scorer'"), "{query}"); - } - - #[test] - fn query_can_exclude_scorers() { - let query = build_cost_query("test-project-id", range(), &[], true); - assert!( - query - .contains("span_attributes.purpose IS NULL OR span_attributes.purpose != 'scorer'"), - "{query}" - ); - } - - #[test] - fn time_range_uses_until_as_window_anchor() { - let args = LogsArgs { - window: "1d".to_string(), - until: Some("2025-02-02T12:00:00Z".to_string()), - ..LogsArgs::default() - }; - let resolved = resolve_time_range(&args, parse_timestamp("2030-01-01").unwrap()).unwrap(); - assert_eq!( - resolved.since, - parse_timestamp("2025-02-01T12:00:00Z").unwrap() - ); - assert_eq!( - resolved.until, - parse_timestamp("2025-02-02T12:00:00Z").unwrap() - ); - } - - #[test] - fn historical_file_rates_are_applied_to_the_matching_segment() { - let file = tempfile::NamedTempFile::new().unwrap(); - std::fs::write( - file.path(), - r#" -version = 1 - -[models."test-model"] -[[models."test-model".rates]] -effective_from = "2025-01-01" -input_usd_per_1m_tokens = 1.0 -output_usd_per_1m_tokens = 2.0 - -[[models."test-model".rates]] -effective_from = "2025-02-01" -input_usd_per_1m_tokens = 3.0 -output_usd_per_1m_tokens = 4.0 -"#, - ) - .unwrap(); - let price_book = PriceBook::load(file.path()).unwrap(); - let segments = build_time_segments(range(), Some(&price_book)); - assert_eq!(segments.len(), 2); - - let row = json!({ - "model": "test-model", - "purpose": null, - "candidate_spans": 3, - "braintrust_priced_spans": 1, - "braintrust_cost": 0.5, - "unpriced_token_spans": 2, - "no_usage_spans": 0, - "p0_spans": 1, - "p0_uncached_input_tokens": 1000000, - "p0_cached_input_tokens": 0, - "p0_effective_cache_write_tokens": 0, - "p0_split_cache_write_5m_tokens": 0, - "p0_split_cache_write_1h_tokens": 0, - "p0_fallback_cache_write_tokens": 0, - "p0_output_tokens": 1000000, - "p1_spans": 1, - "p1_uncached_input_tokens": 1000000, - "p1_cached_input_tokens": 0, - "p1_effective_cache_write_tokens": 0, - "p1_split_cache_write_5m_tokens": 0, - "p1_split_cache_write_1h_tokens": 0, - "p1_fallback_cache_write_tokens": 0, - "p1_output_tokens": 1000000 - }) - .as_object() - .unwrap() - .clone(); - let cost_row = build_cost_row(&row, &segments, Some(&price_book)).unwrap(); - assert_eq!(cost_row.braintrust_cost, Some(0.5)); - assert_eq!(cost_row.file_cost, Some(10.0)); - assert_eq!(cost_row.cost, Some(10.5)); - assert_eq!(cost_row.file_priced_spans, 2); - assert_eq!(cost_row.coverage, Coverage::Full); - } - - #[test] - fn missing_historical_rate_remains_unpriced() { - let row = json!({ - "model": null, - "candidate_spans": 1, - "braintrust_priced_spans": 0, - "braintrust_cost": null, - "unpriced_token_spans": 1, - "no_usage_spans": 0, - "p0_spans": 1, - "p0_uncached_input_tokens": 10, - "p0_cached_input_tokens": 0, - "p0_effective_cache_write_tokens": 0, - "p0_split_cache_write_5m_tokens": 0, - "p0_split_cache_write_1h_tokens": 0, - "p0_fallback_cache_write_tokens": 0, - "p0_output_tokens": 5 - }) - .as_object() - .unwrap() - .clone(); - let segments = vec![TimeSegment { - since: range().since, - until: range().until, - }]; - let cost_row = build_cost_row(&row, &segments, None).unwrap(); - assert_eq!(cost_row.cost, None); - assert_eq!(cost_row.unpriced_token_spans, 1); - assert_eq!(cost_row.coverage, Coverage::Unknown); - } - - #[test] - fn partial_cost_is_marked_approximate() { - assert_eq!(cost_display(Some(1.25), Coverage::Partial), "~$1.25"); - assert_eq!(cost_display(None, Coverage::Unknown), "unknown"); - assert_eq!(coverage_display(3, 2, Coverage::Partial), "3/5"); - } -} diff --git a/src/cost/mod.rs b/src/cost/mod.rs index 90545ec7..937ec15e 100644 --- a/src/cost/mod.rs +++ b/src/cost/mod.rs @@ -1,60 +1,578 @@ -use std::io::{self, Write as _}; +//! `bt cost` — estimate LLM spend across a project's logs, topics, experiments, +//! and playgrounds. +//! +//! Cost is one query: `SUM(cost) GROUP BY WHERE AND