From c4c3ca7ea1a953a7a8c4d0a6ec8d089c48ee98c3 Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Fri, 31 Jul 2026 12:47:49 -0700 Subject: [PATCH 1/7] simplify profiles --- src/auth.rs | 884 ++++++++++++++++++++++++----------------- src/config/mod.rs | 25 +- src/init.rs | 7 +- src/projects/create.rs | 1 + src/setup/mod.rs | 74 +--- src/switch.rs | 242 +---------- src/traces.rs | 35 +- src/utils/profile.rs | 10 +- 8 files changed, 570 insertions(+), 708 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index a915a636..18813a2c 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -48,6 +48,7 @@ pub struct LoginContext { pub login: LoginState, pub api_url: String, pub app_url: String, + pub profile: Option, } #[derive(Debug, Clone)] @@ -57,6 +58,7 @@ pub struct ResolvedAuth { pub app_url: Option, pub org_name: Option, pub is_oauth: bool, + pub profile: Option, } #[derive(Debug, Clone)] @@ -142,7 +144,7 @@ pub fn list_profiles() -> Result> { .iter() .map(|(name, p)| ProfileInfo { name: name.clone(), - org_name: p.org_name.clone(), + org_name: p.org_constraint().map(str::to_string), user_name: p.user_name.clone(), email: p.email.clone(), api_key_hint: p.api_key_hint.clone(), @@ -158,89 +160,34 @@ pub(crate) fn list_stored_profiles() -> Result> { .map(|(name, profile)| StoredProfileInfo { name: name.clone(), is_oauth: profile.auth_kind == AuthKind::Oauth, - org_name: profile.org_name.clone(), + org_name: profile.org_constraint().map(str::to_string), }) .collect()) } -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 profiles for org '{identifier}': {}. Use --profile 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()) - } - } +pub fn select_profile_interactive(current: Option<&str>) -> Result> { + select_profile_for_app_interactive(current, None) } -pub fn select_profile_interactive(current: Option<&str>) -> Result> { - let profiles = list_profiles()?; - if profiles.is_empty() { +fn select_profile_for_app_interactive( + current: Option<&str>, + app_url: Option<&str>, +) -> Result> { + let store = load_auth_store()?; + let names = store + .profiles + .iter() + .filter(|(_, profile)| profile_matches_requested_app_url(profile, app_url)) + .map(|(name, _)| name.as_str()) + .collect::>(); + if names.is_empty() { bail!("no auth profiles found. Run `bt auth login` to create one."); } - if profiles.len() == 1 { - return Ok(Some(profiles[0].name.clone())); + if names.len() == 1 { + return Ok(Some(names[0].to_string())); } - 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 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 org", &labels, default)?; - Ok(Some(profiles[idx].name.clone())) + select_profile_from_store("Select profile", &names, current, &store).map(Some) } pub async fn list_available_orgs(base: &BaseArgs) -> Result> { @@ -313,13 +260,22 @@ struct SecretStore { struct AuthProfile { #[serde(default)] auth_kind: AuthKind, + // Legacy API URL. For OAuth profiles this is also the token endpoint used + // to refresh credentials. It must not be used to infer an app URL or as a + // data-plane default for commands; that URL is resolved from app_url, + // credential, and org at login time. #[serde(default)] api_url: Option, #[serde(default)] app_url: Option, + // An org constraint for API-key profiles. Older versions also populated + // this as a selected org, so it is only a constraint when org_bound was + // explicitly written after checking the credential against app_url. #[serde(default)] org_name: Option, #[serde(default)] + org_bound: Option, + #[serde(default)] oauth_client_id: Option, #[serde(default)] oauth_access_expires_at: Option, @@ -328,6 +284,8 @@ struct AuthProfile { #[serde(default)] email: Option, #[serde(default)] + user_id: Option, + #[serde(default)] api_key_hint: Option, } @@ -339,6 +297,14 @@ enum AuthKind { Oauth, } +impl AuthProfile { + fn org_constraint(&self) -> Option<&str> { + (self.auth_kind == AuthKind::ApiKey && self.org_bound == Some(true)) + .then(|| crate::config::trimmed_option(self.org_name.as_deref())) + .flatten() + } +} + #[derive(Debug, Clone, Deserialize)] struct ApiKeyLoginResponse { org_info: Vec, @@ -456,7 +422,7 @@ pub async fn run(base: BaseArgs, args: AuthArgs) -> Result<()> { } pub async fn login_read_only(base: &BaseArgs) -> Result { - if !has_cached_project_id(base) { + if !has_cached_project_id(base) || base.api_url.is_none() { return login(base).await; } @@ -501,6 +467,7 @@ pub async fn fast_login(base: &BaseArgs) -> Result { login, api_url, app_url, + profile: auth.profile, }) } @@ -533,29 +500,7 @@ pub async fn login(base: &BaseArgs) -> Result { if let Some(project) = &project { builder = builder.default_project(project); } - 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}"))?; - let login = LoginState::new(); - login.set( - api_key.clone(), - String::new(), - org_name, - auth.api_url - .clone() - .unwrap_or_else(|| DEFAULT_API_URL.to_string()), - auth.app_url - .clone() - .unwrap_or_else(|| DEFAULT_APP_URL.to_string()), - ); - login - } - Err(err) => return Err(err.into()), - }; + let login = builder.build().await?.wait_for_login().await?; let api_url = login .api_url() @@ -571,6 +516,7 @@ pub async fn login(base: &BaseArgs) -> Result { login, api_url, app_url, + profile: auth.profile, }; maybe_warn_ai_provider_key_staleness(base, &ctx).await; Ok(ctx) @@ -882,11 +828,17 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { 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 configured_profile_matches_app = store.profiles.get(&profile).is_none_or(|stored| { + profile_matches_requested_app_url(stored, base.app_url.as_deref()) + }); + if configured_profile_matches_app { + auth_base.profile = Some(profile); + } } if let Some(profile_name) = - maybe_select_profile_for_auth(&auth_base, &store, &cfg_org, ui::can_prompt())? + maybe_select_profile_for_auth(&auth_base, &mut store, cfg_org.as_deref(), ui::can_prompt()) + .await? { auth_base.profile = Some(profile_name); } @@ -901,12 +853,10 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { return Ok(auth); } - let effective_org = auth_base.org_name.as_deref().or(cfg_org.as_deref()); let profile_name = auth_base .profile .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()) }) @@ -917,55 +867,64 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { ) })? .to_string(); + auth.api_key = Some( + resolve_oauth_profile_credential(&profile_name, &mut store, auth.api_url.as_deref()) + .await?, + ); + Ok(auth) +} + +async fn resolve_oauth_profile_credential( + profile_name: &str, + store: &mut AuthStore, + api_url_override: Option<&str>, +) -> Result { let profile = store .profiles - .get(profile_name.as_str()) + .get(profile_name) + .cloned() .ok_or_else(|| anyhow::anyhow!("profile '{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 {}`", - shell_quote_arg(&profile_name) + shell_quote_arg(profile_name) ), ) })?; - let cached_expires_at = profile.oauth_access_expires_at; - let api_url = auth - .api_url - .clone() - .unwrap_or_else(|| DEFAULT_API_URL.to_string()); - if let Some(cached_access_token) = - load_valid_cached_oauth_access_token(&profile_name, cached_expires_at)? + load_valid_cached_oauth_access_token(profile_name, profile.oauth_access_expires_at)? { - auth.api_key = Some(cached_access_token); - return Ok(auth); + return Ok(cached_access_token); } - let refresh_token = load_profile_oauth_refresh_token(&profile_name)?.ok_or_else(|| { + 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 {}`", - shell_quote_arg(&profile_name) + shell_quote_arg(profile_name) ), ) })?; + let api_url = api_url_override + .map(str::to_string) + .or_else(|| profile.api_url.clone()) + .unwrap_or_else(|| DEFAULT_API_URL.to_string()); 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)?; 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)?; } } - if let Some(profile) = store.profiles.get_mut(&profile_name) { + if let Some(profile) = store.profiles.get_mut(profile_name) { profile.oauth_access_expires_at = determine_oauth_access_expiry_epoch(&refreshed); } - save_auth_store(&store)?; - auth.api_key = Some(refreshed.access_token); - Ok(auth) + save_auth_store(store)?; + Ok(refreshed.access_token) } pub async fn resolved_auth_env(base: &BaseArgs) -> Result> { @@ -999,46 +958,11 @@ 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(), - ); - } - - let matches: Vec<&str> = store - .profiles - .iter() - .filter(|(_, p)| p.org_name.as_deref() == Some(org)) - .map(|(name, _)| name.as_str()) - .collect(); - - match matches.len() { - 0 => None, - 1 => Some(matches[0]), - _ => None, - } -} - -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)) - .map(|(name, _)| name.as_str()) - .collect() -} - fn profile_label_from_store(name: &str, store: &AuthStore) -> String { match store .profiles .get(name) - .and_then(|profile| profile.org_name.as_deref()) + .and_then(AuthProfile::org_constraint) { Some(org) if org != name => format!("{} (profile: {})", org, name), _ => name.to_string(), @@ -1062,7 +986,7 @@ fn select_profile_from_store( || store .profiles .get(*name) - .and_then(|profile| profile.org_name.as_deref()) + .and_then(AuthProfile::org_constraint) == Some(current) }) }) @@ -1071,10 +995,10 @@ fn select_profile_from_store( Ok(names[idx].to_string()) } -fn maybe_select_profile_for_auth( +async fn maybe_select_profile_for_auth( base: &BaseArgs, - store: &AuthStore, - cfg_org: &Option, + store: &mut AuthStore, + cfg_org: Option<&str>, can_prompt: bool, ) -> Result> { if resolve_api_key_override(base).is_some() { @@ -1090,46 +1014,114 @@ fn maybe_select_profile_for_auth( 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); + let names: Vec = store + .profiles + .iter() + .filter(|(_, profile)| profile_matches_requested_app_url(profile, base.app_url.as_deref())) + .map(|(name, _)| name.clone()) + .collect(); + if names.len() <= 1 { + return Ok(names.first().cloned()); + } + + let effective_org = base.org_name.as_deref().or(cfg_org); + let names = if let Some(org) = effective_org { + let mut matches = Vec::new(); + let mut failures = Vec::new(); + for name in names { + match profile_can_access_org(&name, store, base, org).await { + Ok(true) => matches.push(name), + Ok(false) => {} + Err(err) => failures.push(format!("{name}: {err}")), + } } - - if !can_prompt { - bail!( - "multiple profiles for org '{org}': {}. Use --profile to disambiguate.", - matching_profiles.join(", ") - ); + match matches.len() { + 0 => { + let detail = if failures.is_empty() { + String::new() + } else { + format!(" Could not verify: {}.", failures.join("; ")) + }; + bail!( + "no auth profile can access org '{org}' on app URL '{}'.{detail} Run `bt auth login` or pass --profile .", + normalized_app_url(base.app_url.as_deref()) + ); + } + 1 => return Ok(matches.into_iter().next()), + _ => matches, } + } else { + names + }; - 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(); + let name_refs = names.iter().map(String::as_str).collect::>(); if !can_prompt { + let scope = effective_org + .map(|org| format!(" that can access org '{org}'")) + .unwrap_or_default(); bail!( - "multiple auth profiles available: {}. Pass --profile , set BRAINTRUST_PROFILE, or configure an org.", + "profile selection required in non-interactive mode: multiple auth profiles{scope} available: {}. Pass --profile or set BRAINTRUST_PROFILE.", names.join(", ") ); } - select_profile_from_store("Select org", &names, None, store).map(Some) + select_profile_from_store("Select profile", &name_refs, None, store).map(Some) +} + +async fn profile_can_access_org( + profile_name: &str, + store: &mut AuthStore, + base: &BaseArgs, + org_name: &str, +) -> Result { + let profile = store + .profiles + .get(profile_name) + .cloned() + .ok_or_else(|| anyhow::anyhow!("profile '{profile_name}' not found"))?; + if let Some(constraint) = profile.org_constraint() { + return Ok(constraint == org_name); + } + + let credential = match profile.auth_kind { + AuthKind::ApiKey => load_profile_secret(profile_name)?.ok_or_else(|| { + recoverable_auth_error( + RecoverableAuthErrorKind::StoredCredential, + format!("no keychain credential found for profile '{profile_name}'"), + ) + })?, + AuthKind::Oauth => { + resolve_oauth_profile_credential(profile_name, store, base.api_url.as_deref()).await? + } + }; + let app_url = base + .app_url + .as_deref() + .or(profile.app_url.as_deref()) + .unwrap_or(DEFAULT_APP_URL); + let orgs = fetch_login_orgs(&credential, app_url).await?; + if profile.auth_kind == AuthKind::ApiKey { + if let Some(constraint) = single_org_api_key_constraint(&credential, &orgs) { + if let Some(stored) = store.profiles.get_mut(profile_name) { + stored.org_name = Some(constraint.name.clone()); + stored.org_bound = Some(true); + } + save_auth_store(store)?; + } + } + Ok(orgs.iter().any(|org| org.name == org_name)) +} + +fn normalized_app_url(value: Option<&str>) -> &str { + crate::config::trimmed_option(value) + .unwrap_or(DEFAULT_APP_URL) + .trim_end_matches('/') +} + +fn profile_matches_requested_app_url(profile: &AuthProfile, requested: Option<&str>) -> bool { + requested.is_none_or(|requested| { + normalized_app_url(profile.app_url.as_deref()) == normalized_app_url(Some(requested)) + }) } fn resolve_auth_from_store_with_secret_lookup( @@ -1148,6 +1140,7 @@ where app_url: base.app_url.clone(), org_name: base.org_name.clone().or_else(|| cfg_org.clone()), is_oauth: false, + profile: None, }); } @@ -1157,12 +1150,8 @@ where .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 { @@ -1176,6 +1165,13 @@ where shell_quote_arg(profile_name) ) })?; + if !profile_matches_requested_app_url(profile, base.app_url.as_deref()) { + bail!( + "profile '{profile_name}' belongs to app URL '{}', but '{}' was requested", + normalized_app_url(profile.app_url.as_deref()), + normalized_app_url(base.app_url.as_deref()) + ); + } let is_oauth = profile.auth_kind == AuthKind::Oauth; let api_key = if is_oauth { None @@ -1191,16 +1187,25 @@ where })?) }; + let requested_org = base.org_name.clone().or_else(|| cfg_org.clone()); + let org_name = match (requested_org, profile.org_constraint()) { + (Some(requested), Some(constraint)) if requested != constraint => { + bail!( + "profile '{profile_name}' uses a credential bound to org '{constraint}', but org '{requested}' was requested" + ); + } + (Some(requested), _) => Some(requested), + (None, Some(constraint)) => Some(constraint.to_string()), + (None, None) => None, + }; + return Ok(ResolvedAuth { api_key, - api_url: base.api_url.clone().or_else(|| profile.api_url.clone()), + api_url: base.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()), + org_name, is_oauth, + profile: Some(profile_name.to_string()), }); } @@ -1210,6 +1215,7 @@ where app_url: base.app_url.clone(), org_name: base.org_name.clone().or_else(|| cfg_org.clone()), is_oauth: false, + profile: None, }) } @@ -1240,6 +1246,7 @@ 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 org_constraint = single_org_api_key_constraint(&api_key, &login_orgs).cloned(); let store = load_auth_store()?; let requested_org_resolution = resolve_requested_org_for_api_key_login( &login_orgs, @@ -1271,8 +1278,8 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { 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, + org_constraint.as_ref().map(|org| org.name.as_str()), + &login_app_url, &store, )?; if should_confirm_overwrite { @@ -1282,9 +1289,8 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { commit_api_key_profile( &profile_name, &api_key, - selected_api_url.clone(), - base.app_url.clone(), - selected_org.as_ref().map(|org| org.name.clone()), + Some(login_app_url.clone()), + org_constraint.as_ref().map(|org| org.name.clone()), )?; let context_update = persist_post_login_context( base, @@ -1420,10 +1426,9 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { commit_oauth_profile( &profile_name, &oauth_tokens, - selected_api_url.clone(), + api_url, app_url.clone(), client_id.clone(), - selected_org.as_ref().map(|org| org.name.clone()), )?; let context_update = persist_post_login_context( base, @@ -1463,7 +1468,6 @@ 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<()> { @@ -1472,17 +1476,20 @@ pub(crate) fn commit_api_key_profile( let _ = delete_profile_oauth_access_token(profile_name); let mut store = load_auth_store()?; + let org_bound = org_name.is_some(); store.profiles.insert( profile_name.to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some(api_url), - app_url, + api_url: None, + app_url: Some(app_url.unwrap_or_else(|| DEFAULT_APP_URL.to_string())), org_name, + org_bound: Some(org_bound), oauth_client_id: None, oauth_access_expires_at: None, user_name: None, email: None, + user_id: None, api_key_hint: Some(obscure_api_key(api_key)), }, ); @@ -1495,7 +1502,6 @@ fn commit_oauth_profile( api_url: String, app_url: String, client_id: String, - org_name: Option, ) -> Result<()> { let refresh_token = tokens.refresh_token.as_ref().ok_or_else(|| { anyhow::anyhow!( @@ -1516,11 +1522,13 @@ fn commit_oauth_profile( auth_kind: AuthKind::Oauth, api_url: Some(api_url), app_url: Some(app_url), - org_name, + org_name: None, + org_bound: Some(false), oauth_client_id: Some(client_id), oauth_access_expires_at, user_name: jwt_id.name, email: jwt_id.email, + user_id: jwt_id.subject, api_key_hint: None, }, ); @@ -1626,10 +1634,9 @@ fn resolve_selected_profile_name_for_debug( } } - 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")); - } + let (configured_profile, _) = config_auth_context(base); + if let Some(profile_name) = configured_profile { + return Ok((profile_name, "config")); } if store.profiles.len() == 1 { @@ -1643,7 +1650,7 @@ fn resolve_selected_profile_name_for_debug( } } - bail!("no profile selected; pass --profile , set BRAINTRUST_PROFILE, or configure an org") + bail!("no profile selected; pass --profile or set BRAINTRUST_PROFILE") } fn resolve_profile_name( @@ -1683,11 +1690,11 @@ fn default_login_org_name( let stored_org_name = store .profiles .get(profile_name) - .and_then(|profile| profile.org_name.as_deref()) + .and_then(AuthProfile::org_constraint) .map(str::trim) .filter(|org_name| !org_name.is_empty()); - Some(stored_org_name.unwrap_or(profile_name).to_string()) + stored_org_name.map(str::to_string) } fn default_profile_name(suggested_org_name: Option<&str>) -> String { @@ -1711,24 +1718,23 @@ fn next_available_profile_name(base_name: &str, store: &AuthStore) -> String { fn resolve_api_key_login_profile_name( explicit_profile: Option<&str>, - suggested_org_name: Option<&str>, - selected_api_url: &str, + org_constraint: Option<&str>, + app_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 profile_name = resolve_profile_name(Some(profile_name), org_constraint)?; 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) + !profile_matches_api_key_login_target(profile, app_url, org_constraint) }); 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 - }); + let default_name = default_profile_name(org_constraint); + let has_matching_api_key_profile = store + .profiles + .values() + .any(|profile| profile_matches_api_key_login_target(profile, app_url, org_constraint)); if has_matching_api_key_profile { return Ok((next_available_profile_name(&default_name, store), false)); @@ -1742,38 +1748,25 @@ fn resolve_api_key_login_profile_name( fn resolve_oauth_login_profile_name( explicit_profile: Option<&str>, - suggested_org_name: Option<&str>, - selected_api_url: &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, - ) - }); + let profile_name = resolve_profile_name(Some(profile_name), None)?; + let should_confirm_overwrite = store + .profiles + .get(&profile_name) + .is_some_and(|profile| !profile_matches_oauth_login_target(profile, app_url, 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, - ) - }) + .filter(|(_, profile)| profile_matches_oauth_login_target(profile, app_url, jwt_id)) .max_by(|(left_name, left), (right_name, right)| { left.oauth_access_expires_at .unwrap_or_default() @@ -1786,36 +1779,47 @@ fn resolve_oauth_login_profile_name( return Ok((profile_name, false)); } - let default_name = default_profile_name(suggested_org_name); - Ok(( - default_name.clone(), - store.profiles.contains_key(&default_name), - )) + let default_name = default_oauth_profile_name(jwt_id); + Ok((next_available_profile_name(&default_name, store), false)) } fn profile_matches_api_key_login_target( profile: &AuthProfile, - selected_api_url: &str, - suggested_org_name: Option<&str>, + app_url: &str, + org_constraint: 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 + && normalized_app_url(profile.app_url.as_deref()) == normalized_app_url(Some(app_url)) + && profile.org_constraint() == org_constraint } 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 + && normalized_app_url(profile.app_url.as_deref()) == normalized_app_url(Some(app_url)) + && match (profile.user_id.as_deref(), jwt_id.subject.as_deref()) { + (Some(stored), Some(current)) => stored == current, + _ => profile.user_name == jwt_id.name && profile.email == jwt_id.email, + } +} + +fn default_oauth_profile_name(jwt_id: &JwtIdentity) -> String { + let candidate = jwt_id + .email + .as_deref() + .and_then(|email| { + email + .split_once('@') + .map(|(local, _)| local) + .or(Some(email)) + }) + .or(jwt_id.name.as_deref()); + candidate + .and_then(crate::utils::sanitize_name_segment) + .unwrap_or_else(|| "profile".to_string()) } fn confirm_profile_overwrite(profile_name: &str) -> Result<()> { @@ -1870,6 +1874,7 @@ fn build_login_context_for_selected_org( login, api_url: api_url.to_string(), app_url: app_url.to_string(), + profile: None, } } @@ -2188,7 +2193,7 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifi build_verification( name, auth_kind, - profile.org_name.clone(), + profile.org_constraint().map(str::to_string), jwt_id, hint, status, @@ -2208,6 +2213,19 @@ async fn verify_profile_full(name: &str, profile: &AuthProfile) -> ProfileVerifi }; match fetch_login_orgs(&credential, app_url).await { + Ok(orgs) + if profile + .org_constraint() + .is_some_and(|constraint| !orgs.iter().any(|org| org.name == constraint)) => + { + mk( + ProfileStatus::Error( + "credential no longer has access to its bound org".to_string(), + ), + jwt_id, + hint, + ) + } Ok(_) => mk(ProfileStatus::Ok, jwt_id, hint), Err(e) => { let msg = e.to_string(); @@ -2279,7 +2297,7 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { serde_json::json!({ "name": name, "auth": match p.auth_kind { AuthKind::ApiKey => "api_key", AuthKind::Oauth => "oauth" }, - "org": p.org_name, + "org": p.org_constraint(), "user_name": p.user_name, "user_email": p.email, "api_key_hint": p.api_key_hint, @@ -2295,8 +2313,7 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { AuthKind::Oauth => "oauth", }; let org = profile - .org_name - .as_deref() + .org_constraint() .map(|o| format!(" org={o}")) .unwrap_or_default(); let id = match (profile.user_name.as_deref(), profile.email.as_deref()) { @@ -2343,6 +2360,13 @@ async fn fetch_login_orgs(api_key: &str, app_url: &str) -> Result( + credential: &str, + orgs: &'a [LoginOrgInfo], +) -> Option<&'a LoginOrgInfo> { + (credential.trim().starts_with("sk-") && orgs.len() == 1).then(|| &orgs[0]) +} + fn select_login_org( mut orgs: Vec, requested_org_name: Option<&str>, @@ -3528,6 +3552,7 @@ fn decode_jwt_payload(token: &str) -> Option { } struct JwtIdentity { + subject: Option, name: Option, email: Option, } @@ -3536,6 +3561,10 @@ fn decode_jwt_identity(token: &str) -> JwtIdentity { let extract = || -> Option { let payload = decode_jwt_payload(token)?; Some(JwtIdentity { + subject: payload + .get("sub") + .and_then(|v| v.as_str()) + .map(String::from), name: payload .get("name") .and_then(|v| v.as_str()) @@ -3547,6 +3576,7 @@ fn decode_jwt_identity(token: &str) -> JwtIdentity { }) }; extract().unwrap_or(JwtIdentity { + subject: None, name: None, email: None, }) @@ -3976,10 +4006,12 @@ mod tests { api_url: Some((*api_url).to_string()), app_url: Some((*app_url).to_string()), org_name: Some((*org_name).to_string()), + org_bound: Some(true), oauth_client_id: None, oauth_access_expires_at: None, user_name: None, email: None, + user_id: None, api_key_hint: None, }, ); @@ -4160,6 +4192,7 @@ 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()), + org_bound: Some(true), oauth_client_id: None, oauth_access_expires_at: None, ..Default::default() @@ -4187,6 +4220,7 @@ 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()), + org_bound: Some(true), oauth_client_id: None, oauth_access_expires_at: None, ..Default::default() @@ -4201,7 +4235,7 @@ mod tests { ) .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.api_url, None); assert_eq!(resolved.org_name.as_deref(), Some("Example Org")); assert!(!resolved.is_oauth); } @@ -4279,6 +4313,7 @@ mod tests { api_url: Some("https://api.example.com".to_string()), app_url: None, org_name: Some("Example Org".to_string()), + org_bound: Some(true), oauth_client_id: None, oauth_access_expires_at: None, ..Default::default() @@ -4312,6 +4347,7 @@ mod tests { api_url: Some("https://api.example.com".to_string()), app_url: None, org_name: Some("Example Org".to_string()), + org_bound: Some(true), oauth_client_id: None, oauth_access_expires_at: None, ..Default::default() @@ -4344,6 +4380,7 @@ mod tests { api_url: Some("https://api.example.com".to_string()), app_url: None, org_name: Some("Example Org".to_string()), + org_bound: Some(true), oauth_client_id: None, oauth_access_expires_at: None, ..Default::default() @@ -4362,9 +4399,10 @@ mod tests { } #[test] - fn resolve_auth_marks_oauth_profiles() { + fn resolve_auth_keeps_oauth_profile_and_org_orthogonal() { let mut base = make_base(); base.profile = Some("work".to_string()); + base.org_name = Some("Other Org".to_string()); let mut store = AuthStore::default(); store.profiles.insert( @@ -4390,7 +4428,7 @@ mod tests { assert!(resolved.is_oauth); assert_eq!(resolved.api_key, None); - assert_eq!(resolved.org_name.as_deref(), Some("Example Org")); + assert_eq!(resolved.org_name.as_deref(), Some("Other Org")); } #[test] @@ -4432,67 +4470,8 @@ 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() { + #[tokio::test] + async fn profile_selection_requires_choice_when_multiple_profiles_without_prompt() { let base = make_base(); let mut store = AuthStore::default(); store.profiles.insert( @@ -4510,7 +4489,8 @@ mod tests { }, ); - let err = maybe_select_profile_for_auth(&base, &store, &None, false) + let err = maybe_select_profile_for_auth(&base, &mut store, None, false) + .await .expect_err("selection should be required"); assert!(err.to_string().contains("multiple auth profiles available")); @@ -4519,8 +4499,8 @@ mod tests { assert!(err.to_string().contains("--profile ")); } - #[test] - fn profile_selection_requires_choice_for_ambiguous_org_without_prompt() { + #[tokio::test] + async fn org_does_not_disambiguate_profiles() { let mut base = make_base(); base.org_name = Some("acme".into()); @@ -4529,6 +4509,7 @@ mod tests { "work-1".into(), AuthProfile { org_name: Some("acme".into()), + org_bound: Some(true), ..Default::default() }, ); @@ -4536,20 +4517,47 @@ mod tests { "work-2".into(), AuthProfile { org_name: Some("acme".into()), + org_bound: Some(true), ..Default::default() }, ); - let err = maybe_select_profile_for_auth(&base, &store, &None, false) - .expect_err("org selection should be required"); + let err = maybe_select_profile_for_auth(&base, &mut store, None, false) + .await + .expect_err("profile selection should be required"); - assert!(err.to_string().contains("multiple profiles for org 'acme'")); + assert!(err + .to_string() + .contains("multiple auth profiles that can access org 'acme' available")); 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() { + #[tokio::test] + async fn org_selects_unique_compatible_profile() { + let mut base = make_base(); + base.org_name = Some("target-org".into()); + let mut store = AuthStore::default(); + for (name, org) in [("target", "target-org"), ("other", "other-org")] { + store.profiles.insert( + name.into(), + AuthProfile { + auth_kind: AuthKind::ApiKey, + org_name: Some(org.into()), + org_bound: Some(true), + ..Default::default() + }, + ); + } + + let selected = maybe_select_profile_for_auth(&base, &mut store, None, false) + .await + .expect("select compatible profile"); + assert_eq!(selected.as_deref(), Some("target")); + } + + #[tokio::test] + async fn profile_selection_skips_when_api_key_override_is_active() { let mut base = make_base(); base.api_key = Some("explicit-key".into()); @@ -4559,12 +4567,63 @@ mod tests { .insert("alpha".into(), AuthProfile::default()); store.profiles.insert("beta".into(), AuthProfile::default()); - let selection = maybe_select_profile_for_auth(&base, &store, &None, false) + let selection = maybe_select_profile_for_auth(&base, &mut store, None, false) + .await .expect("api key override should skip profile selection"); assert_eq!(selection, None); } + #[tokio::test] + async fn app_url_selects_the_unique_matching_profile() { + let mut base = make_base(); + base.app_url = Some("https://app.two.example/".into()); + let mut store = AuthStore::default(); + store.profiles.insert( + "one".into(), + AuthProfile { + app_url: Some("https://app.one.example".into()), + ..Default::default() + }, + ); + store.profiles.insert( + "two".into(), + AuthProfile { + app_url: Some("https://app.two.example".into()), + ..Default::default() + }, + ); + + let selected = maybe_select_profile_for_auth(&base, &mut store, None, false) + .await + .expect("select"); + assert_eq!(selected.as_deref(), Some("two")); + } + + #[test] + fn explicit_profile_cannot_be_used_with_another_app_url() { + let mut base = make_base(); + base.profile = Some("work".into()); + base.app_url = Some("https://app.other.example".into()); + let mut store = AuthStore::default(); + store.profiles.insert( + "work".into(), + AuthProfile { + app_url: Some("https://app.work.example".into()), + ..Default::default() + }, + ); + + let err = resolve_auth_from_store_with_secret_lookup( + &base, + &store, + |_| Ok(Some("profile-key".into())), + &None, + ) + .expect_err("app URL mismatch should fail"); + assert!(err.to_string().contains("belongs to app URL")); + } + #[test] fn resolve_auth_uses_org_to_find_profile() { let mut base = make_base(); @@ -4618,7 +4677,7 @@ mod tests { } #[test] - fn resolve_auth_config_org_overrides_profile_org() { + fn resolve_auth_rejects_config_org_outside_api_key_constraint() { let mut base = make_base(); base.profile = Some("default-profile".to_string()); @@ -4627,20 +4686,47 @@ mod tests { "default-profile".into(), AuthProfile { org_name: Some("profile-org".into()), + org_bound: Some(true), ..Default::default() }, ); let cfg_org = Some("local-org".to_string()); - let resolved = resolve_auth_from_store_with_secret_lookup( + let err = 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")); + .expect_err("org constraint should be enforced"); + assert!(err.to_string().contains("bound to org 'profile-org'")); + assert!(err.to_string().contains("org 'local-org' was requested")); + } + + #[test] + fn legacy_selected_org_is_not_treated_as_a_verified_constraint() { + let mut base = make_base(); + base.profile = Some("legacy-profile".to_string()); + let mut store = AuthStore::default(); + store.profiles.insert( + "legacy-profile".into(), + AuthProfile { + auth_kind: AuthKind::ApiKey, + org_name: Some("legacy-selected-org".into()), + org_bound: None, + ..Default::default() + }, + ); + + let resolved = resolve_auth_from_store_with_secret_lookup( + &base, + &store, + |_| Ok(Some("sk-test-key".into())), + &None, + ) + .expect("resolve legacy profile"); + + assert_eq!(resolved.org_name, None); } #[test] @@ -4660,7 +4746,7 @@ mod tests { } #[test] - fn resolve_auth_explicit_profile_overrides_org_resolution() { + fn resolve_auth_explicit_api_key_profile_rejects_other_org() { let mut base = make_base(); base.profile = Some("other".into()); base.org_name = Some("acme-corp".into()); @@ -4677,20 +4763,20 @@ mod tests { "other".into(), AuthProfile { org_name: Some("other-org".into()), + org_bound: Some(true), api_url: Some("https://api.other.com".into()), ..Default::default() }, ); - let resolved = resolve_auth_from_store_with_secret_lookup( + let err = 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")); + .expect_err("org constraint should be enforced"); + assert!(err.to_string().contains("bound to org 'other-org'")); } #[test] @@ -4700,8 +4786,9 @@ mod tests { "acme".into(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.acme.example".into()), + app_url: Some("https://app.acme.example".into()), org_name: Some("acme".into()), + org_bound: Some(true), ..Default::default() }, ); @@ -4709,7 +4796,7 @@ mod tests { let (profile_name, should_confirm) = resolve_api_key_login_profile_name( None, Some("acme"), - "https://api.acme.example", + "https://app.acme.example", &store, ) .expect("resolve"); @@ -4725,8 +4812,9 @@ mod tests { "work".into(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.test.example".into()), + app_url: Some("https://app.test.example".into()), org_name: Some("test-org".into()), + org_bound: Some(true), ..Default::default() }, ); @@ -4734,7 +4822,7 @@ mod tests { let (profile_name, should_confirm) = resolve_api_key_login_profile_name( Some("work"), Some("test-org"), - "https://api.test.example", + "https://app.test.example", &store, ) .expect("resolve"); @@ -4750,8 +4838,9 @@ mod tests { "work".into(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.test.example".into()), + app_url: Some("https://app.test.example".into()), org_name: Some("test-org".into()), + org_bound: Some(true), ..Default::default() }, ); @@ -4759,7 +4848,7 @@ mod tests { let (profile_name, should_confirm) = resolve_api_key_login_profile_name( Some("work"), Some("other-org"), - "https://api.test.example", + "https://app.test.example", &store, ) .expect("resolve"); @@ -4775,6 +4864,7 @@ mod tests { "work".into(), AuthProfile { org_name: Some("acme".into()), + org_bound: Some(true), ..Default::default() }, ); @@ -4786,13 +4876,10 @@ mod tests { } #[test] - fn default_login_org_name_falls_back_to_profile_name() { + fn default_login_org_name_does_not_treat_profile_name_as_org() { let store = AuthStore::default(); - assert_eq!( - default_login_org_name(&store, Some(" acme "), None).as_deref(), - Some("acme") - ); + assert_eq!(default_login_org_name(&store, Some(" acme "), None), None); } #[test] @@ -4865,6 +4952,7 @@ mod tests { ); let jwt_id = JwtIdentity { + subject: None, name: Some("Alice".into()), email: Some("alice@example.com".into()), }; @@ -4898,6 +4986,7 @@ mod tests { }, ); let jwt_id = JwtIdentity { + subject: None, name: Some("Test User".into()), email: Some("user@test.example".into()), }; @@ -4917,7 +5006,7 @@ mod tests { } #[test] - fn resolve_oauth_login_profile_name_confirms_explicit_different_target() { + fn resolve_oauth_login_profile_name_ignores_org_change_for_same_identity() { let mut store = AuthStore::default(); store.profiles.insert( "work".into(), @@ -4932,6 +5021,7 @@ mod tests { }, ); let jwt_id = JwtIdentity { + subject: None, name: Some("Test User".into()), email: Some("user@test.example".into()), }; @@ -4947,7 +5037,30 @@ mod tests { .expect("resolve"); assert_eq!(profile_name, "work"); - assert!(should_confirm); + assert!(!should_confirm); + } + + #[test] + fn oauth_profile_identity_prefers_stable_subject_claim() { + let profile = AuthProfile { + auth_kind: AuthKind::Oauth, + app_url: Some("https://app.test.example".into()), + user_id: Some("user-one".into()), + user_name: Some("Test User".into()), + email: Some("user@test.example".into()), + ..Default::default() + }; + let other_user = JwtIdentity { + subject: Some("user-two".into()), + name: Some("Test User".into()), + email: Some("user@test.example".into()), + }; + + assert!(!profile_matches_oauth_login_target( + &profile, + "https://app.test.example", + &other_user + )); } fn login_org(id: &str, name: &str) -> LoginOrgInfo { @@ -4958,6 +5071,22 @@ mod tests { } } + #[test] + fn only_single_org_sk_credentials_get_an_org_constraint() { + let one_org = vec![login_org("org_1", "test-org")]; + assert_eq!( + single_org_api_key_constraint("sk-test-key", &one_org).map(|org| org.name.as_str()), + Some("test-org") + ); + assert!(single_org_api_key_constraint("oauth-token", &one_org).is_none()); + + let multiple_orgs = vec![ + login_org("org_1", "test-org"), + login_org("org_2", "other-org"), + ]; + assert!(single_org_api_key_constraint("sk-test-key", &multiple_orgs).is_none()); + } + #[tokio::test] async fn persist_post_login_context_clears_stale_project_for_org_only_login() { let _env = TestEnv::new(None, None).await; @@ -5102,6 +5231,7 @@ mod tests { .encode(r#"{"name":"Alice","email":"alice@example.com"}"#); let token = format!("{header}.{payload}.sig"); let id = decode_jwt_identity(&token); + assert_eq!(id.subject, None); assert_eq!(id.name.as_deref(), Some("Alice")); assert_eq!(id.email.as_deref(), Some("alice@example.com")); } @@ -5112,6 +5242,7 @@ mod tests { 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.subject.as_deref(), Some("123")); assert_eq!(id.name, None); assert_eq!(id.email, None); } @@ -5119,6 +5250,7 @@ mod tests { #[test] fn decode_jwt_identity_handles_garbage() { let id = decode_jwt_identity("not-a-jwt"); + assert_eq!(id.subject, None); assert_eq!(id.name, None); assert_eq!(id.email, None); } @@ -5331,7 +5463,8 @@ mod tests { } #[tokio::test] - async fn login_read_only_cached_project_id_and_config_org_uses_fast_path() { + async fn login_read_only_cached_project_id_and_active_profile_uses_explicit_api_url_fast_path() + { let env = TestEnv::new(Some("proj_123"), Some("acme-org")).await; setup_auth_store_profiles(&[ ( @@ -5347,11 +5480,17 @@ mod tests { "https://www.other.example", ), ]); + let mut cfg = crate::config::load_global().expect("load global config"); + cfg.profile = Some("acme-profile".to_string()); + crate::config::save_global(&cfg).expect("save active profile"); save_profile_secret_plaintext("acme-profile", "acme-secret").expect("save acme secret"); save_profile_secret_plaintext("other-profile", "other-secret").expect("save other secret"); + let mut base = make_base(); + base.profile = Some("acme-profile".to_string()); + base.api_url = Some("https://api.acme.example".to_string()); 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"); @@ -5362,10 +5501,11 @@ mod tests { } #[tokio::test] - async fn login_read_only_cached_project_id_and_org_uses_default_urls() { + async fn login_read_only_cached_project_id_and_org_uses_explicit_default_api_url() { let env = TestEnv::new(Some("proj_123"), None).await; let mut base = make_base(); base.api_key = Some("test-api-key".into()); + base.api_url = Some(DEFAULT_API_URL.to_string()); base.org_name = Some("acme".into()); let ctx = env diff --git a/src/config/mod.rs b/src/config/mod.rs index 779499d5..e1e4bf9b 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -180,21 +180,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> { @@ -466,13 +458,16 @@ mod tests { } #[test] - fn project_config_matches_explicit_profile_or_legacy_org() { + fn project_config_matches_org_independently_of_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/init.rs b/src/init.rs index e9976c1b..087a4eec 100644 --- a/src/init.rs +++ b/src/init.rs @@ -38,8 +38,8 @@ pub async fn run(base: BaseArgs, _args: InitArgs) -> Result<()> { eprintln!("Link to a Braintrust project..."); - let (org, project) = if let (Some(o), Some(p)) = (&base.org_name, &base.project) { - (o.clone(), p.clone()) + let (org, project, profile) = if let (Some(o), Some(p)) = (&base.org_name, &base.project) { + (o.clone(), p.clone(), base.profile.clone()) } else if !is_interactive() { bail!("--org and --project required in non-interactive mode"); } else { @@ -62,10 +62,11 @@ pub async fn run(base: BaseArgs, _args: InitArgs) -> Result<()> { .await? .name; - (org, project) + (org, project, ctx.profile) }; let cfg = config::Config { + profile, org: Some(org.clone()), project: Some(project.clone()), ..Default::default() diff --git a/src/projects/create.rs b/src/projects/create.rs index b41c6637..b676e75c 100644 --- a/src/projects/create.rs +++ b/src/projects/create.rs @@ -189,6 +189,7 @@ mod tests { login, api_url: server.base_url.clone(), app_url: "https://app.example.com".to_string(), + profile: None, }) .expect("build client"); (server, client) diff --git a/src/setup/mod.rs b/src/setup/mod.rs index 45ebfbcd..d0d9fac0 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -1469,13 +1469,14 @@ async fn run_setup_browser_auth( }; let stored_profiles = auth::list_profiles()?; let profile_name = setup_browser_profile_name(profile_name, &org.name, &stored_profiles); + let org_constraint = (completed.api_key.trim().starts_with("sk-") && available_orgs.len() == 1) + .then(|| org.name.clone()); auth::commit_api_key_profile( &profile_name, &completed.api_key, - login.api_url.clone(), Some(login.app_url.clone()), - Some(org.name.clone()), + org_constraint, ) .context("failed to save Braintrust auth profile after browser setup")?; @@ -1522,7 +1523,7 @@ fn setup_browser_profile_name( fn resolve_profile_name_for_setup( base: &BaseArgs, profiles: &[auth::ProfileInfo], - prompt_for_choice: bool, + _prompt_for_choice: bool, ) -> Result> { if let Some(profile_name) = base .profile @@ -1538,46 +1539,10 @@ fn resolve_profile_name_for_setup( ); } - 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 mut matches = profiles - .iter() - .filter(|profile| profile.org_name.as_deref() == Some(org_name)) - .map(|profile| profile.name.clone()) - .collect::>(); - matches.sort(); - - return match matches.len() { - 0 => Ok(None), - 1 => Ok(Some(matches.remove(0))), - _ 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(", ") - ), - }; - } - 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 profile selected")) - } else { - Ok(None) - } + Ok(None) } fn find_http_error(err: &anyhow::Error) -> Option<&crate::http::HttpError> { @@ -1646,6 +1611,7 @@ fn build_api_key_login_context( login, api_url, app_url, + profile: None, } } @@ -1758,13 +1724,17 @@ async fn ensure_profile_or_setup_browser_auth( auth_base.api_key = None; auth_base.api_key_source = None; - if let Some(profile_name) = selected_profile { + if let Some(profile_name) = selected_profile.as_ref() { auth_base.profile = Some(profile_name.clone()); + } + if selected_profile.is_some() || !profiles.is_empty() { match auth::login(&auth_base).await { Ok(ctx) => { - base.profile = auth_base.profile.clone(); - let is_oauth = auth::resolve_auth(&auth_base).await?.is_oauth; + base.profile = ctx.profile.clone(); + let mut resolved_base = auth_base.clone(); + resolved_base.profile = ctx.profile.clone(); + let is_oauth = auth::resolve_auth(&resolved_base).await?.is_oauth; return Ok(SetupAuthLogin { login: ctx, is_oauth, @@ -1773,10 +1743,8 @@ 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...", - profile_name, err - ); + let profile = selected_profile.as_deref().unwrap_or("selected profile"); + eprintln!(" Profile '{profile}' credentials inaccessible ({err}). Re-authenticating in the browser..."); } if !can_prompt { bail!( @@ -1785,7 +1753,7 @@ async fn ensure_profile_or_setup_browser_auth( } return run_setup_browser_auth( base, - Some(&profile_name), + selected_profile.as_deref(), project_name, project_was_explicit, requested_org, @@ -2043,17 +2011,12 @@ async fn ensure_setup_auth( &[], )?; 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()), + (api_key.trim().starts_with("sk-") && available_orgs.len() == 1) + .then(|| org.name.clone()), )?; return build_setup_auth_context(base, client, false, needs_api_key, None).await; } @@ -5474,6 +5437,7 @@ mod tests { login, api_url, app_url, + profile: None, } } diff --git a/src/switch.rs b/src/switch.rs index 35495bb6..1d3edb8e 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -3,7 +3,7 @@ use clap::Args; use dialoguer::{console, theme::ColorfulTheme, Select}; use crate::args::BaseArgs; -use crate::auth::{self, login}; +use crate::auth::login; use crate::config; use crate::http::ApiClient; use crate::projects::api; @@ -59,54 +59,21 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { .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, - )?, + if resolved_project.is_none() && is_interactive() { + interactive = true; + } + let requested_profile = if has_api_key_override { + None + } else { + config::trimmed_option(base.profile.as_deref()) + .or_else(|| config::trimmed_option(current_cfg.profile.as_deref())) + .map(str::to_string) }; - // 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(); + login_base.profile = requested_profile; + login_base.org_name = resolved_org.clone().or_else(|| current_cfg.org.clone()); + login_base.project = resolved_project.clone(); let ctx = login(&login_base).await?; let client = ApiClient::new(&ctx)?; @@ -147,9 +114,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); + let config_profile = ctx.profile.as_deref().map(str::to_string); apply_switch_config( &mut cfg, config_profile.as_deref(), @@ -278,35 +243,9 @@ pub(crate) fn apply_switch_config( } } -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 { @@ -339,16 +278,6 @@ mod tests { } } - fn profile_info(name: &str, org_name: Option<&str>) -> ProfileInfo { - ProfileInfo { - name: name.to_string(), - org_name: org_name.map(String::from), - user_name: None, - email: None, - api_key_hint: None, - } - } - // --- resolve_target tests (unchanged) --- #[test] @@ -443,99 +372,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(); @@ -590,52 +426,4 @@ mod tests { 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 6ef7ff1c..1ec1c80d 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}; @@ -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; }; @@ -5756,11 +5742,6 @@ where if !has_org_override && !has_profile_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 } @@ -6963,16 +6944,14 @@ mod tests { } #[test] - fn apply_url_hints_infers_profile_from_url_org() { + fn apply_url_hints_keeps_profile_independent_from_url_org() { 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_to_base(base, Some(&parsed)); assert_eq!(updated.org_name.as_deref(), Some("Lovable")); - assert_eq!(updated.profile.as_deref(), Some("lovable-profile")); + assert_eq!(updated.profile, None); } #[test] @@ -6981,9 +6960,7 @@ mod tests { base.profile = Some("explicit-profile".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_to_base(base, Some(&parsed)); assert_eq!(updated.profile.as_deref(), Some("explicit-profile")); assert!(updated.org_name.is_none()); diff --git a/src/utils/profile.rs b/src/utils/profile.rs index b97812b2..51cd0f16 100644 --- a/src/utils/profile.rs +++ b/src/utils/profile.rs @@ -24,12 +24,6 @@ fn resolve_profile_info_from_profiles( } 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)) @@ -40,7 +34,9 @@ fn resolve_profile_info_from_profiles( .into_iter() .find(|profile| profile.name == profile_name); } - return None; + if profiles.len() != 1 { + return None; + } } if profiles.len() == 1 { From 650c0f16130568cf23f1053f1a44f9f6ffe5f04f Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Fri, 31 Jul 2026 12:50:04 -0700 Subject: [PATCH 2/7] simplify --- src/auth.rs | 71 ++++++++++++++++++++++++++++++++++------------------- 1 file changed, 46 insertions(+), 25 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 18813a2c..60482690 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -260,12 +260,11 @@ struct SecretStore { struct AuthProfile { #[serde(default)] auth_kind: AuthKind, - // Legacy API URL. For OAuth profiles this is also the token endpoint used - // to refresh credentials. It must not be used to infer an app URL or as a - // data-plane default for commands; that URL is resolved from app_url, - // credential, and org at login time. - #[serde(default)] - api_url: Option, + // OAuth API base used to refresh credentials. Keep the legacy serialized + // key for compatibility with existing profiles. This must not be used as + // a command data-plane default or to infer an app URL. + #[serde(default, rename = "api_url")] + oauth_api_url: Option, #[serde(default)] app_url: Option, // An org constraint for API-key profiles. Older versions also populated @@ -910,7 +909,7 @@ async fn resolve_oauth_profile_credential( })?; let api_url = api_url_override .map(str::to_string) - .or_else(|| profile.api_url.clone()) + .or_else(|| profile.oauth_api_url.clone()) .unwrap_or_else(|| DEFAULT_API_URL.to_string()); let refreshed = refresh_oauth_access_token(&api_url, &refresh_token, client_id, profile_name).await?; @@ -1481,7 +1480,7 @@ pub(crate) fn commit_api_key_profile( profile_name.to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: None, + oauth_api_url: None, app_url: Some(app_url.unwrap_or_else(|| DEFAULT_APP_URL.to_string())), org_name, org_bound: Some(org_bound), @@ -1520,7 +1519,7 @@ fn commit_oauth_profile( profile_name.to_string(), AuthProfile { auth_kind: AuthKind::Oauth, - api_url: Some(api_url), + oauth_api_url: Some(api_url), app_url: Some(app_url), org_name: None, org_bound: Some(false), @@ -1549,7 +1548,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { } let api_url = profile - .api_url + .oauth_api_url .clone() .unwrap_or_else(|| DEFAULT_API_URL.to_string()); let client_id = profile.oauth_client_id.clone().ok_or_else(|| { @@ -4003,7 +4002,7 @@ mod tests { (*profile_name).to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some((*api_url).to_string()), + oauth_api_url: Some((*api_url).to_string()), app_url: Some((*app_url).to_string()), org_name: Some((*org_name).to_string()), org_bound: Some(true), @@ -4189,7 +4188,7 @@ mod tests { "work".to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), + 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()), org_bound: Some(true), @@ -4217,7 +4216,7 @@ mod tests { "work".to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), + 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()), org_bound: Some(true), @@ -4274,7 +4273,7 @@ mod tests { "work".to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), + oauth_api_url: Some("https://api.example.com".to_string()), app_url: None, org_name: None, oauth_client_id: None, @@ -4310,7 +4309,7 @@ mod tests { "work".to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), + oauth_api_url: Some("https://api.example.com".to_string()), app_url: None, org_name: Some("Example Org".to_string()), org_bound: Some(true), @@ -4344,7 +4343,7 @@ mod tests { "work".to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), + oauth_api_url: Some("https://api.example.com".to_string()), app_url: None, org_name: Some("Example Org".to_string()), org_bound: Some(true), @@ -4377,7 +4376,7 @@ mod tests { "work".to_string(), AuthProfile { auth_kind: AuthKind::ApiKey, - api_url: Some("https://api.example.com".to_string()), + oauth_api_url: Some("https://api.example.com".to_string()), app_url: None, org_name: Some("Example Org".to_string()), org_bound: Some(true), @@ -4409,7 +4408,7 @@ mod tests { "work".to_string(), AuthProfile { auth_kind: AuthKind::Oauth, - api_url: Some("https://api.example.com".to_string()), + 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()), @@ -4634,7 +4633,7 @@ mod tests { "work".into(), AuthProfile { org_name: Some("acme-corp".into()), - api_url: Some("https://api.acme.com".into()), + oauth_api_url: Some("https://api.acme.com".into()), ..Default::default() }, ); @@ -4659,7 +4658,7 @@ mod tests { "work".into(), AuthProfile { org_name: Some("acme-corp".into()), - api_url: Some("https://api.acme.com".into()), + oauth_api_url: Some("https://api.acme.com".into()), ..Default::default() }, ); @@ -4764,7 +4763,7 @@ mod tests { AuthProfile { org_name: Some("other-org".into()), org_bound: Some(true), - api_url: Some("https://api.other.com".into()), + oauth_api_url: Some("https://api.other.com".into()), ..Default::default() }, ); @@ -4928,7 +4927,7 @@ mod tests { "older".into(), AuthProfile { auth_kind: AuthKind::Oauth, - api_url: Some("https://api.acme.example".into()), + 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), @@ -4941,7 +4940,7 @@ mod tests { "newer".into(), AuthProfile { auth_kind: AuthKind::Oauth, - api_url: Some("https://api.acme.example".into()), + 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), @@ -4977,7 +4976,7 @@ mod tests { "work".into(), AuthProfile { auth_kind: AuthKind::Oauth, - api_url: Some("https://api.test.example".into()), + 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()), @@ -5012,7 +5011,7 @@ mod tests { "work".into(), AuthProfile { auth_kind: AuthKind::Oauth, - api_url: Some("https://api.test.example".into()), + 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()), @@ -5063,6 +5062,28 @@ mod tests { )); } + #[test] + fn oauth_api_url_keeps_legacy_serialized_key() { + let profile: AuthProfile = serde_json::from_value(serde_json::json!({ + "auth_kind": "oauth", + "api_url": "https://api.test.example" + })) + .expect("deserialize legacy profile"); + assert_eq!( + profile.oauth_api_url.as_deref(), + Some("https://api.test.example") + ); + + let serialized = serde_json::to_value(profile).expect("serialize profile"); + assert_eq!( + serialized + .get("api_url") + .and_then(serde_json::Value::as_str), + Some("https://api.test.example") + ); + assert!(serialized.get("oauth_api_url").is_none()); + } + fn login_org(id: &str, name: &str) -> LoginOrgInfo { LoginOrgInfo { id: id.to_string(), From a93ceaa42d0dcd8bfad17c6b5964e7f90ac6231b Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Fri, 31 Jul 2026 13:58:36 -0700 Subject: [PATCH 3/7] simplify further --- src/auth.rs | 152 +++++++++++++++++++++++++++++++++++----------------- 1 file changed, 102 insertions(+), 50 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index 60482690..dfaac805 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1396,28 +1396,15 @@ 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(), - ui::can_prompt(), - base.verbose, - true, - explicitly_quiet(base), - )?; - let selected_api_url = - resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)?; + let selected_org = select_explicit_oauth_login_org(&login_orgs, base.org_name.as_deref())?; + let selected_api_url = if selected_org.is_some() { + resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)? + } else { + api_url.clone() + }; 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, - )?; + let (profile_name, should_confirm_overwrite) = + resolve_oauth_login_profile_name(base.profile.as_deref(), &app_url, &jwt_id, &store)?; if should_confirm_overwrite { confirm_profile_overwrite(&profile_name)?; } @@ -1429,15 +1416,20 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { app_url.clone(), client_id.clone(), )?; - let context_update = persist_post_login_context( - base, - &profile_name, - &oauth_tokens.access_token, - &selected_api_url, - &app_url, - selected_org.as_ref(), - ) - .await + let context_update = match selected_org.as_ref() { + Some(org) => { + persist_post_login_context( + base, + &profile_name, + &oauth_tokens.access_token, + &selected_api_url, + &app_url, + Some(org), + ) + .await + } + None => persist_identity_login_context(&profile_name), + } .context("login succeeded, but failed to update active context")?; let human = format_login_success(&selected_org, &profile_name, &selected_api_url); @@ -1453,10 +1445,12 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { }), || { ui::print_command_status(ui::CommandStatus::Success, &human); - ui::print_command_status( - ui::CommandStatus::Success, - &format!("Switched to {}", context_update.display), - ); + let context_status = if selected_org.is_some() { + format!("Switched to {}", context_update.display) + } else { + format!("Using profile '{profile_name}' (organization unchanged)") + }; + ui::print_command_status(ui::CommandStatus::Success, &context_status); if base.verbose { eprintln!("Wrote to {}", context_update.path.display()); } @@ -1747,8 +1741,6 @@ fn resolve_api_key_login_profile_name( 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, @@ -1851,7 +1843,7 @@ fn format_login_success( "Logged in as {} (profile: {profile_name}, api: {api_url})", org.name ), - None => format!("Logged in (cross-org, profile: {profile_name}, api: {api_url})"), + None => format!("Logged in (profile: {profile_name}, api: {api_url})"), } } @@ -1944,6 +1936,24 @@ async fn persist_post_login_context( }) } +fn persist_identity_login_context(profile_name: &str) -> Result { + let path = if ui::can_prompt() && config::local_path().is_some() { + switch::select_scope()?.0 + } else { + config::global_path()? + }; + + let mut cfg = config::load_file(&path); + cfg.profile = Some(profile_name.to_string()); + config::save_file(&path, &cfg) + .context(format!("Could not save config to {}", path.display()))?; + + Ok(PostLoginContextUpdate { + display: profile_name.to_string(), + path, + }) +} + /// 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 { @@ -2366,6 +2376,19 @@ fn single_org_api_key_constraint<'a>( (credential.trim().starts_with("sk-") && orgs.len() == 1).then(|| &orgs[0]) } +fn select_explicit_oauth_login_org( + orgs: &[LoginOrgInfo], + requested_org_name: Option<&str>, +) -> Result> { + requested_org_name + .map(|org_name| { + find_login_org(orgs, org_name) + .cloned() + .ok_or_else(|| missing_requested_org_error(orgs, org_name)) + }) + .transpose() +} + fn select_login_org( mut orgs: Vec, requested_org_name: Option<&str>, @@ -4955,15 +4978,9 @@ mod tests { 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"); + let (profile_name, should_confirm) = + resolve_oauth_login_profile_name(None, "https://www.acme.example", &jwt_id, &store) + .expect("resolve"); assert_eq!(profile_name, "newer"); assert!(!should_confirm); @@ -4992,8 +5009,6 @@ mod tests { 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, @@ -5027,8 +5042,6 @@ mod tests { 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, @@ -5108,6 +5121,45 @@ mod tests { assert!(single_org_api_key_constraint("sk-test-key", &multiple_orgs).is_none()); } + #[test] + fn oauth_login_only_selects_an_explicit_org() { + let orgs = vec![ + login_org("org_1", "test-org"), + login_org("org_2", "other-org"), + ]; + + assert!(select_explicit_oauth_login_org(&orgs, None) + .expect("no org selection") + .is_none()); + assert_eq!( + select_explicit_oauth_login_org(&orgs, Some("other-org")) + .expect("explicit org selection") + .map(|org| org.id), + Some("org_2".to_string()) + ); + } + + #[tokio::test] + async fn identity_login_preserves_existing_org_and_project() { + let _env = TestEnv::new(None, None).await; + crate::config::save_global(&crate::config::Config { + profile: Some("old-profile".to_string()), + org: Some("test-org".to_string()), + project: Some("test-project".to_string()), + project_id: Some("proj_test".to_string()), + ..Default::default() + }) + .expect("save initial config"); + + persist_identity_login_context("new-profile").expect("persist identity"); + let cfg = crate::config::load_global().expect("load global config"); + + assert_eq!(cfg.profile.as_deref(), Some("new-profile")); + 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")); + } + #[tokio::test] async fn persist_post_login_context_clears_stale_project_for_org_only_login() { let _env = TestEnv::new(None, None).await; From 4b8fd609620453a5ea32d74b9291bfbd1a7e45ee Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Fri, 31 Jul 2026 15:12:20 -0700 Subject: [PATCH 4/7] cedric --- src/auth.rs | 42 ++++++++++++++++++++++------ src/switch.rs | 72 ++++++++++++++++++++++++++++++++++++++++++++---- src/ui/select.rs | 35 ++++++++++++++++++++++- 3 files changed, 133 insertions(+), 16 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index dfaac805..12f432ed 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -835,9 +835,14 @@ pub async fn resolve_auth(base: &BaseArgs) -> Result { } } - if let Some(profile_name) = - maybe_select_profile_for_auth(&auth_base, &mut store, cfg_org.as_deref(), ui::can_prompt()) - .await? + if let Some(profile_name) = maybe_select_profile_for_auth( + &auth_base, + &mut store, + cfg_org.as_deref(), + ui::can_prompt(), + None, + ) + .await? { auth_base.profile = Some(profile_name); } @@ -999,6 +1004,7 @@ async fn maybe_select_profile_for_auth( store: &mut AuthStore, cfg_org: Option<&str>, can_prompt: bool, + current_profile: Option<&str>, ) -> Result> { if resolve_api_key_override(base).is_some() { return Ok(None); @@ -1064,7 +1070,25 @@ async fn maybe_select_profile_for_auth( ); } - select_profile_from_store("Select profile", &name_refs, None, store).map(Some) + select_profile_from_store("Select profile", &name_refs, current_profile, store).map(Some) +} + +pub(crate) async fn select_compatible_profile_interactive( + base: &BaseArgs, + current_profile: Option<&str>, +) -> Result> { + let mut selection_base = base.clone(); + selection_base.profile = None; + selection_base.profile_explicit = false; + let mut store = load_auth_store()?; + maybe_select_profile_for_auth( + &selection_base, + &mut store, + None, + ui::can_prompt(), + current_profile, + ) + .await } async fn profile_can_access_org( @@ -4511,7 +4535,7 @@ mod tests { }, ); - let err = maybe_select_profile_for_auth(&base, &mut store, None, false) + let err = maybe_select_profile_for_auth(&base, &mut store, None, false, None) .await .expect_err("selection should be required"); @@ -4544,7 +4568,7 @@ mod tests { }, ); - let err = maybe_select_profile_for_auth(&base, &mut store, None, false) + let err = maybe_select_profile_for_auth(&base, &mut store, None, false, None) .await .expect_err("profile selection should be required"); @@ -4572,7 +4596,7 @@ mod tests { ); } - let selected = maybe_select_profile_for_auth(&base, &mut store, None, false) + let selected = maybe_select_profile_for_auth(&base, &mut store, None, false, None) .await .expect("select compatible profile"); assert_eq!(selected.as_deref(), Some("target")); @@ -4589,7 +4613,7 @@ mod tests { .insert("alpha".into(), AuthProfile::default()); store.profiles.insert("beta".into(), AuthProfile::default()); - let selection = maybe_select_profile_for_auth(&base, &mut store, None, false) + let selection = maybe_select_profile_for_auth(&base, &mut store, None, false, None) .await .expect("api key override should skip profile selection"); @@ -4616,7 +4640,7 @@ mod tests { }, ); - let selected = maybe_select_profile_for_auth(&base, &mut store, None, false) + let selected = maybe_select_profile_for_auth(&base, &mut store, None, false, None) .await .expect("select"); assert_eq!(selected.as_deref(), Some("two")); diff --git a/src/switch.rs b/src/switch.rs index 1d3edb8e..3916b878 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -3,12 +3,12 @@ use clap::Args; use dialoguer::{console, theme::ColorfulTheme, Select}; use crate::args::BaseArgs; -use crate::auth::login; +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, + fuzzy_select, is_interactive, print_command_status, select_project, with_spinner, CommandStatus, }; #[derive(Debug, Clone, Args)] @@ -64,15 +64,31 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { } let requested_profile = if has_api_key_override { None + } else if base.profile.is_some() { + base.profile.clone() + } else if interactive { + let mut profile_base = base.clone(); + profile_base.org_name = resolved_org.clone(); + auth::select_compatible_profile_interactive(&profile_base, current_cfg.profile.as_deref()) + .await? + } else { + config::trimmed_option(current_cfg.profile.as_deref()).map(str::to_string) + }; + + let selected_org = if resolved_org.is_some() { + resolved_org.clone() + } else if interactive { + let mut org_base = base.clone(); + org_base.profile = requested_profile.clone(); + org_base.org_name = None; + Some(select_org_for_switch(&org_base, current_cfg.org.as_deref()).await?) } else { - config::trimmed_option(base.profile.as_deref()) - .or_else(|| config::trimmed_option(current_cfg.profile.as_deref())) - .map(str::to_string) + current_cfg.org.clone() }; let mut login_base = base.clone(); login_base.profile = requested_profile; - login_base.org_name = resolved_org.clone().or_else(|| current_cfg.org.clone()); + login_base.org_name = selected_org; login_base.project = resolved_project.clone(); let ctx = login(&login_base).await?; @@ -146,6 +162,30 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { Ok(()) } +async fn select_org_for_switch(base: &BaseArgs, current_org: Option<&str>) -> Result { + let orgs = auth::list_available_orgs(base).await?; + if orgs.len() == 1 { + return Ok(orgs[0].name.clone()); + } + if orgs.is_empty() { + bail!("no organizations available for the selected profile"); + } + + let labels = orgs.iter().map(|org| org.name.as_str()).collect::>(); + let default = default_org_selection(&orgs, current_org); + let selected = fuzzy_select("Select organization", &labels, default)?; + Ok(orgs[selected].name.clone()) +} + +fn default_org_selection(orgs: &[auth::AvailableOrg], current_org: Option<&str>) -> usize { + current_org + .and_then(|current| { + orgs.iter() + .position(|org| org.name.eq_ignore_ascii_case(current)) + }) + .unwrap_or(0) +} + pub(crate) fn select_scope() -> Result<(std::path::PathBuf, &'static str)> { let global = config::global_path()?; let local = config::local_path().unwrap(); @@ -247,6 +287,14 @@ pub(crate) fn apply_switch_config( mod tests { use super::*; + fn available_org(name: &str) -> auth::AvailableOrg { + auth::AvailableOrg { + id: format!("id-{name}"), + name: name.to_string(), + api_url: None, + } + } + fn switch_args(target: Option<&str>) -> SwitchArgs { SwitchArgs { global: false, @@ -255,6 +303,18 @@ mod tests { } } + #[test] + fn org_picker_defaults_to_current_org() { + let orgs = vec![available_org("alpha"), available_org("beta")]; + assert_eq!(default_org_selection(&orgs, Some("BETA")), 1); + } + + #[test] + fn org_picker_defaults_to_first_when_current_org_is_unavailable() { + let orgs = vec![available_org("alpha"), available_org("beta")]; + assert_eq!(default_org_selection(&orgs, Some("missing")), 0); + } + fn base_args(org: Option<&str>, project: Option<&str>) -> BaseArgs { BaseArgs { json: false, diff --git a/src/ui/select.rs b/src/ui/select.rs index bcfd02e2..665a5396 100644 --- a/src/ui/select.rs +++ b/src/ui/select.rs @@ -256,6 +256,10 @@ pub async fn select_project( let mut projects = with_spinner("Loading projects...", api::list_projects(client)).await?; projects.sort_by(|a, b| a.name.cmp(&b.name)); + if let Some(project) = take_only_existing_project(&mut projects, mode) { + return Ok(project); + } + let label = select_label.unwrap_or("Select project"); if mode_allows_create(mode) { @@ -358,6 +362,14 @@ fn mode_allows_create(mode: ProjectSelectMode) -> bool { matches!(mode, ProjectSelectMode::AllowCreateWithDefaultProjectNote) } +fn take_only_existing_project( + projects: &mut Vec, + mode: ProjectSelectMode, +) -> Option { + (matches!(mode, ProjectSelectMode::ExistingOnly) && projects.len() == 1) + .then(|| projects.remove(0)) +} + fn default_new_project_name() -> String { let output = std::process::Command::new("whoami").output(); let user = output @@ -374,7 +386,7 @@ fn default_new_project_name() -> String { mod tests { use super::{ default_new_project_name, default_project_selection, project_display_names, - project_selection_labels, ProjectSelectMode, + project_selection_labels, take_only_existing_project, ProjectSelectMode, }; use crate::projects::api::Project; @@ -414,6 +426,27 @@ mod tests { assert!(err.to_string().contains("no projects found")); } + #[test] + fn existing_only_auto_selects_a_sole_project() { + let mut projects = vec![project("only")]; + let selected = take_only_existing_project(&mut projects, ProjectSelectMode::ExistingOnly) + .expect("sole project"); + + assert_eq!(selected.name, "only"); + assert!(projects.is_empty()); + } + + #[test] + fn create_mode_keeps_a_sole_project_in_the_picker() { + let mut projects = vec![project("only")]; + assert!(take_only_existing_project( + &mut projects, + ProjectSelectMode::AllowCreateWithDefaultProjectNote + ) + .is_none()); + assert_eq!(projects.len(), 1); + } + #[test] fn project_selection_labels_returns_project_names() { let labels = project_selection_labels(&[project("alpha")]); From 35b749f79b6a7f2d23cb2cd9c927461bc5df91dd Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Fri, 31 Jul 2026 16:02:35 -0700 Subject: [PATCH 5/7] rm more --- README.md | 34 ++-- skills/shared/braintrust-cli-body.md | 2 +- src/auth.rs | 280 +++++++++------------------ src/eval.rs | 2 +- src/main.rs | 24 ++- src/setup/mod.rs | 25 ++- src/status.rs | 74 ++++++- src/utils/profile.rs | 3 + tests/cli.rs | 27 +++ tests/functions.rs | 137 ++----------- 10 files changed, 263 insertions(+), 345 deletions(-) diff --git a/README.md b/README.md index cfa19e70..3a80617e 100644 --- a/README.md +++ b/README.md @@ -138,7 +138,8 @@ Remove-Item -Recurse -Force (Join-Path $env:APPDATA "bt") -ErrorAction SilentlyC | Command | Description | | ------------- | ------------------------------------------------------------------ | | `bt init` | Initialize `.bt/` config directory and link to a project | -| `bt auth` | Authenticate with Braintrust | +| `bt login` | Log in to Braintrust or refresh an OAuth login | +| `bt logout` | Remove a saved Braintrust login | | `bt switch` | Switch org and project context | | `bt status` | Show current org and project context | | `bt datasets` | Manage datasets and dataset pipelines | @@ -310,36 +311,33 @@ Local version and pagination-key conversion helpers: - `bt util version inspect p07639577379371417602` - `bt util version inspect p07639577379371417602 --utc` -## `bt auth` +## `bt login` and `bt logout` -- Authenticate interactively (prompts for auth method, profile name defaults to org name): - - `bt auth login` +- Authenticate interactively: + - `bt 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. - - `bt` confirms the resolved API URL before saving. + - Login stores an identity and its app URL. OAuth login does not select an organization unless `--org` is passed explicitly. + - Use `bt switch` to select the active profile, organization, and project context. - Login with OAuth (browser-based, stores refresh token in secure credential store): - - `bt auth login --oauth --profile work` + - `bt login --oauth --profile work` - 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: - - `bt auth profiles` +- List profiles and the current context: + - `bt status --all` - Log out (remove a saved profile): - - `bt auth logout` - - `bt auth logout --force` (skip confirmation) -- Show current auth source/profile: - - `bt auth status` + - `bt logout` + - `bt logout --force` (skip confirmation) - Force-refresh OAuth access token for debugging: - - `bt auth refresh --profile work` + - `bt login --refresh --profile work` 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) +4. Compatible profile for the selected app URL and organization +5. Single-profile auto-select (if only one compatible profile exists) +6. Interactive profile picker (if multiple compatible profiles exist and a TTY is available) 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. diff --git a/skills/shared/braintrust-cli-body.md b/skills/shared/braintrust-cli-body.md index 01603c6f..52294804 100644 --- a/skills/shared/braintrust-cli-body.md +++ b/skills/shared/braintrust-cli-body.md @@ -21,4 +21,4 @@ Use the Braintrust `bt` CLI for projects, traces, prompts, and sync workflows. ## Guardrails - Prefer `bt` commands over direct API calls when both can accomplish the task. -- Respect existing login/profile settings from `bt auth`. +- Respect existing login/profile settings from `bt login`. diff --git a/src/auth.rs b/src/auth.rs index 12f432ed..a3c67d70 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -16,7 +16,7 @@ use anyhow::{bail, Context, Result}; use base64::Engine as _; use braintrust_sdk_rust::{BraintrustClient, LoginState}; use chrono::{DateTime, Months, Utc}; -use clap::{Args, Subcommand}; +use clap::Args; use crossterm::event::{self, Event, KeyCode, KeyEventKind}; use dialoguer::{Confirm, Input, Password}; use oauth2::basic::BasicClient; @@ -64,6 +64,9 @@ pub struct ResolvedAuth { #[derive(Debug, Clone)] pub struct ProfileInfo { pub name: String, + pub auth: String, + pub app_url: String, + pub oauth_api_url: Option, pub org_name: Option, pub user_name: Option, pub email: Option, @@ -144,6 +147,18 @@ pub fn list_profiles() -> Result> { .iter() .map(|(name, p)| ProfileInfo { name: name.clone(), + auth: match p.auth_kind { + AuthKind::ApiKey => "api_key", + AuthKind::Oauth => "oauth", + } + .to_string(), + app_url: p + .app_url + .clone() + .unwrap_or_else(|| DEFAULT_APP_URL.to_string()), + oauth_api_url: (p.auth_kind == AuthKind::Oauth) + .then(|| p.oauth_api_url.clone()) + .flatten(), org_name: p.org_constraint().map(str::to_string), user_name: p.user_name.clone(), email: p.email.clone(), @@ -181,7 +196,7 @@ fn select_profile_for_app_interactive( .map(|(name, _)| name.as_str()) .collect::>(); if names.is_empty() { - bail!("no auth profiles found. Run `bt auth login` to create one."); + bail!("no auth profiles found. Run `bt login` to create one."); } if names.len() == 1 { return Ok(Some(names[0].to_string())); @@ -349,43 +364,15 @@ struct OAuthErrorResponse { } #[derive(Debug, Clone, Args)] -#[command(after_help = "\ -Examples: - bt auth login - bt auth profiles - bt auth refresh - bt auth logout --profile work -")] -pub struct AuthArgs { - #[command(subcommand)] - command: AuthCommand, -} - -#[derive(Debug, Clone, Subcommand)] -enum AuthCommand { - /// Authenticate with Braintrust (OAuth or API key) - Login(AuthLoginArgs), - /// Force-refresh OAuth access token for a profile - Refresh, - /// List auth profiles and check connection status - Profiles(AuthProfilesArgs), - /// Log out by removing a saved profile - Logout(AuthLogoutArgs), -} - -#[derive(Debug, Clone, Args)] -struct AuthProfilesArgs { - /// Only show the profile with this name - #[arg(long, value_name = "NAME")] - profile: Option, -} - -#[derive(Debug, Clone, Args)] -struct AuthLoginArgs { +pub struct LoginArgs { /// Use OAuth login instead of API key login #[arg(long)] oauth: bool, + /// Force-refresh OAuth credentials for the selected profile + #[arg(long, conflicts_with_all = ["oauth", "client_id", "no_browser"])] + refresh: bool, + /// OAuth client id (defaults to bt_cli_) #[arg(long, value_name = "CLIENT_ID")] client_id: Option, @@ -396,11 +383,7 @@ struct AuthLoginArgs { } #[derive(Debug, Clone, Args)] -struct AuthLogoutArgs { - /// Profile name to log out of (interactive picker if omitted) - #[arg(long)] - profile: Option, - +pub struct LogoutArgs { /// Skip confirmation prompt #[arg(long, short = 'f')] force: bool, @@ -411,15 +394,18 @@ struct PostLoginContextUpdate { path: PathBuf, } -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::Logout(logout_args) => run_login_logout(base, logout_args), +pub async fn run_login_command(base: BaseArgs, args: LoginArgs) -> Result<()> { + if args.refresh { + run_login_refresh(&base).await + } else { + run_login_set(&base, args).await } } +pub fn run_logout_command(base: BaseArgs, args: LogoutArgs) -> Result<()> { + run_login_logout(base, args) +} + pub async fn login_read_only(base: &BaseArgs) -> Result { if !has_cached_project_id(base) || base.api_url.is_none() { return login(base).await; @@ -440,7 +426,7 @@ pub async fn fast_login(base: &BaseArgs) -> Result { let auth = resolve_auth(base).await?; let api_key = auth.api_key.clone().ok_or_else(|| { anyhow::anyhow!( - "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt auth login`" + "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt login`" ) })?; let org_name = auth.org_name.clone().unwrap_or_default(); @@ -475,7 +461,7 @@ pub async fn login(base: &BaseArgs) -> Result { let auth = resolve_auth(base).await?; let api_key = auth.api_key.clone().ok_or_else(|| { anyhow::anyhow!( - "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt auth login`" + "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt login`" ) })?; @@ -892,7 +878,7 @@ async fn resolve_oauth_profile_credential( recoverable_auth_error( RecoverableAuthErrorKind::OauthClientId, format!( - "oauth profile '{profile_name}' is missing client_id; re-run `bt auth login --oauth --profile {}`", + "oauth profile '{profile_name}' is missing client_id; re-run `bt login --oauth --profile {}`", shell_quote_arg(profile_name) ), ) @@ -907,7 +893,7 @@ async fn resolve_oauth_profile_credential( recoverable_auth_error( RecoverableAuthErrorKind::OauthRefreshToken, format!( - "oauth refresh token missing for profile '{profile_name}'; re-run `bt auth login --oauth --profile {}`", + "oauth refresh token missing for profile '{profile_name}'; re-run `bt login --oauth --profile {}`", shell_quote_arg(profile_name) ), ) @@ -1048,7 +1034,7 @@ async fn maybe_select_profile_for_auth( format!(" Could not verify: {}.", failures.join("; ")) }; bail!( - "no auth profile can access org '{org}' on app URL '{}'.{detail} Run `bt auth login` or pass --profile .", + "no auth profile can access org '{org}' on app URL '{}'.{detail} Run `bt login` or pass --profile .", normalized_app_url(base.app_url.as_deref()) ); } @@ -1184,7 +1170,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 '{profile_name}' not found; run `bt status --all` or `bt login --profile {}`", shell_quote_arg(profile_name) ) })?; @@ -1203,7 +1189,7 @@ where recoverable_auth_error( RecoverableAuthErrorKind::StoredCredential, format!( - "no keychain credential found for profile '{profile_name}'; re-run `bt auth login --profile {}`", + "no keychain credential found for profile '{profile_name}'; re-run `bt login --profile {}`", shell_quote_arg(profile_name) ), ) @@ -1242,7 +1228,7 @@ where }) } -async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { +async fn run_login_set(base: &BaseArgs, args: LoginArgs) -> Result<()> { if args.oauth { return run_login_oauth(base, args).await; } @@ -1350,7 +1336,7 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { ) } -async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { +async fn run_login_oauth(base: &BaseArgs, args: LoginArgs) -> Result<()> { let api_url = base .api_url .clone() @@ -1561,7 +1547,7 @@ async fn run_login_refresh(base: &BaseArgs) -> Result<()> { .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" + "profile '{profile_name}' uses api key auth; `bt login --refresh` only applies to oauth profiles" ); } @@ -1571,14 +1557,14 @@ 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 {}`", + "oauth profile '{profile_name}' is missing client_id; re-run `bt login --oauth --profile {}`", shell_quote_arg(&profile_name) ) })?; 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 {}`", + "oauth refresh token missing for profile '{profile_name}'; re-run `bt login --oauth --profile {}`", shell_quote_arg(&profile_name) ) })?; @@ -1917,7 +1903,7 @@ async fn resolve_post_login_project( 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 `" + "cannot set a default project in cross-org mode; rerun `bt login --org --project `" ) })?; let ctx = @@ -1988,7 +1974,7 @@ fn profile_not_found_err(name: &str, store: &AuthStore) -> anyhow::Error { format!(": {}", available.join(", ")) }; anyhow::anyhow!( - "profile '{name}' not found; run `bt auth profiles` to see available profiles{suffix}" + "profile '{name}' not found; run `bt status --all` to see available profiles{suffix}" ) } @@ -2003,62 +1989,6 @@ fn emit_result(json: bool, payload: serde_json::Value, human: impl FnOnce()) -> Ok(()) } -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() { - return emit_result(base.json, serde_json::json!([]), || { - println!("No saved profiles. Run `bt auth login` to create one.") - }); - } - - let verifications = verify_all_profiles_from_store(&filtered_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)?; - return Ok(()); - } - - if base.json { - println!("{}", serde_json::to_string(&verifications)?); - return Ok(()); - } - - for v in &verifications { - let cmd_status = match v.status.as_str() { - "ok" => crate::ui::CommandStatus::Success, - "expired" => crate::ui::CommandStatus::Warning, - _ => crate::ui::CommandStatus::Error, - }; - crate::ui::print_command_status(cmd_status, &format_verification_line(v)); - } - - if base.verbose { - if let Ok(path) = auth_store_path() { - eprintln!("\nCredentials: {}", path.display()); - } - } - - Ok(()) -} - fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result<()> { let profile_name = profile_name.trim(); if profile_name.is_empty() { @@ -2110,7 +2040,7 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result< ) } -fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { +fn run_login_logout(base: BaseArgs, args: LogoutArgs) -> Result<()> { let store = load_auth_store()?; if store.profiles.is_empty() { return emit_result(base.json, serde_json::json!({ "status": "empty" }), || { @@ -2118,7 +2048,7 @@ fn run_login_logout(base: BaseArgs, args: AuthLogoutArgs) -> Result<()> { }); } - let profile_name = if let Some(p) = args.profile.or(base.profile) { + let profile_name = if let Some(p) = base.profile { let p = p.trim().to_string(); if !store.profiles.contains_key(&p) { return Err(profile_not_found_err(&p, &store)); @@ -2177,6 +2107,9 @@ fn load_credential_for_profile(name: &str, profile: &AuthProfile) -> CredentialL pub struct ProfileVerification { pub name: String, pub auth: String, + pub app_url: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub api_url: Option, #[serde(skip_serializing_if = "Option::is_none")] pub org: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -2192,12 +2125,15 @@ pub struct ProfileVerification { fn build_verification( name: &str, - auth_kind: &str, - org: Option, + profile: &AuthProfile, jwt_id: Option, api_key_hint: Option, status: ProfileStatus, ) -> ProfileVerification { + let auth_kind = match profile.auth_kind { + AuthKind::ApiKey => "api_key", + AuthKind::Oauth => "oauth", + }; let (status_str, error) = match &status { ProfileStatus::Ok => ("ok", None), ProfileStatus::Expired => ("expired", None), @@ -2207,7 +2143,14 @@ fn build_verification( ProfileVerification { name: name.to_string(), auth: auth_kind.to_string(), - org, + app_url: profile + .app_url + .clone() + .unwrap_or_else(|| DEFAULT_APP_URL.to_string()), + api_url: (profile.auth_kind == AuthKind::Oauth) + .then_some(profile.oauth_api_url.clone()) + .flatten(), + org: profile.org_constraint().map(str::to_string), 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, @@ -2218,19 +2161,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 = match profile.auth_kind { - AuthKind::ApiKey => "api_key", - AuthKind::Oauth => "oauth", - }; let mk = |status, jwt_id: Option, hint: Option| { - build_verification( - name, - auth_kind, - profile.org_constraint().map(str::to_string), - jwt_id, - hint, - status, - ) + build_verification(name, profile, jwt_id, hint, status) }; let credential = match load_credential_for_profile(name, profile) { @@ -2294,8 +2226,20 @@ async fn verify_all_profiles_from_store(store: &AuthStore) -> Vec String { - let mut parts = vec![v.name.clone(), v.auth.clone()]; +pub(crate) async fn profile_verifications() -> Result> { + let store = load_auth_store()?; + Ok(verify_all_profiles_from_store(&store).await) +} + +pub(crate) fn credentials_path() -> Result { + auth_store_path() +} + +pub(crate) fn format_verification_line(v: &ProfileVerification) -> String { + let mut parts = vec![v.name.clone(), v.app_url.clone(), v.auth.clone()]; + if let Some(ref api_url) = v.api_url { + parts.push(format!("api: {api_url}")); + } if let Some(ref org) = v.org { parts.push(format!("org: {org}")); } @@ -2321,49 +2265,6 @@ fn format_verification_line(v: &ProfileVerification) -> String { parts.join(" — ") } -fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> { - if json { - let output: Vec = store - .profiles - .iter() - .map(|(name, p)| { - serde_json::json!({ - "name": name, - "auth": match p.auth_kind { AuthKind::ApiKey => "api_key", AuthKind::Oauth => "oauth" }, - "org": p.org_constraint(), - "user_name": p.user_name, - "user_email": p.email, - "api_key_hint": p.api_key_hint, - "status": "unchecked" - }) - }) - .collect(); - 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_constraint() - .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}"); - } - } - Ok(()) -} - async fn fetch_login_orgs(api_key: &str, app_url: &str) -> Result> { let login_url = format!("{}/api/apikey/login", app_url.trim_end_matches('/')); let client = build_http_client(crate::http::DEFAULT_HTTP_TIMEOUT) @@ -3016,7 +2917,7 @@ fn map_refresh_oauth_error( message.push_str(&format!(" ({description})")); } message.push_str(&format!( - "; re-run `bt auth login --oauth --profile {}`", + "; re-run `bt login --oauth --profile {}`", shell_quote_arg(profile_name) )); return recoverable_auth_error(RecoverableAuthErrorKind::OauthRefreshToken, message); @@ -4124,7 +4025,7 @@ mod tests { assert!(err.to_string().contains("refresh token expired")); assert!(err .to_string() - .contains("re-run `bt auth login --oauth --profile 'test profile'`")); + .contains("re-run `bt login --oauth --profile 'test profile'`")); } #[test] @@ -5357,6 +5258,8 @@ mod tests { let v = ProfileVerification { name: "work".into(), auth: "oauth".into(), + app_url: "https://app.test.example".into(), + api_url: Some("https://api.test.example".into()), org: Some("acme".into()), user_name: Some("Alice".into()), user_email: Some("alice@example.com".into()), @@ -5366,7 +5269,7 @@ mod tests { }; assert_eq!( format_verification_line(&v), - "work — oauth — org: acme — Alice (alice@example.com)" + "work — https://app.test.example — oauth — api: https://api.test.example — org: acme — Alice (alice@example.com)" ); } @@ -5375,6 +5278,8 @@ mod tests { let v = ProfileVerification { name: "work".into(), auth: "api_key".into(), + app_url: "https://app.test.example".into(), + api_url: None, org: Some("acme".into()), user_name: None, user_email: None, @@ -5384,7 +5289,7 @@ mod tests { }; assert_eq!( format_verification_line(&v), - "work — api_key — org: acme — sk-****zhJwO" + "work — https://app.test.example — api_key — org: acme — sk-****zhJwO" ); } @@ -5393,6 +5298,8 @@ mod tests { let v = ProfileVerification { name: "old".into(), auth: "oauth".into(), + app_url: "https://app.test.example".into(), + api_url: None, org: None, user_name: None, user_email: None, @@ -5400,7 +5307,10 @@ mod tests { status: "expired".into(), error: None, }; - assert_eq!(format_verification_line(&v), "old — oauth — token expired"); + assert_eq!( + format_verification_line(&v), + "old — https://app.test.example — oauth — token expired" + ); } #[test] @@ -5408,6 +5318,8 @@ mod tests { let v = ProfileVerification { name: "bad".into(), auth: "api_key".into(), + app_url: "https://app.test.example".into(), + api_url: None, org: Some("corp".into()), user_name: None, user_email: None, @@ -5417,7 +5329,7 @@ mod tests { }; assert_eq!( format_verification_line(&v), - "bad — api_key — org: corp — invalid API key" + "bad — https://app.test.example — api_key — org: corp — invalid API key" ); } diff --git a/src/eval.rs b/src/eval.rs index ebc4f4dc..38d73e27 100644 --- a/src/eval.rs +++ b/src/eval.rs @@ -2905,7 +2905,7 @@ impl EvalUi { self.record_deferred_error(message); } if show_hint { - let hint = "Hint: pass --api-key, set BRAINTRUST_API_KEY, run `bt auth login`/`bt auth login --oauth`, or use --no-send-logs for local evals."; + let hint = "Hint: pass --api-key, set BRAINTRUST_API_KEY, run `bt login`/`bt login --oauth`, or use --no-send-logs for local evals."; if self.verbose { let _ = self.progress.println(hint.dark_grey().to_string()); } else { diff --git a/src/main.rs b/src/main.rs index 5f12e60f..24341fc6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -57,7 +57,8 @@ const HELP_TEMPLATE: &str = "\ Core init Initialize .bt config directory and files - auth Authenticate bt with Braintrust + login Log in to Braintrust + logout Remove a saved Braintrust login switch Switch org and project context view View logs, traces, and spans @@ -80,7 +81,7 @@ Data & evaluation Additional docs Manage workflow docs for coding agents setup Configure Braintrust setup flows - status Show current org and project context + status Show current identity, org, and project context update Update bt in-place Flags @@ -129,8 +130,10 @@ enum Commands { Docs(CLIArgs), /// Run SQL queries against Braintrust Sql(CLIArgs), - /// Authenticate bt with Braintrust - Auth(CLIArgs), + /// Log in to Braintrust + Login(CLIArgs), + /// Remove a saved Braintrust login + Logout(CLIArgs), /// View logs, traces, and spans View(CLIArgs), #[cfg(unix)] @@ -163,7 +166,7 @@ enum Commands { Util(CLIArgs), /// Switch org and project context Switch(CLIArgs), - /// Show current org and project context + /// Show current identity, org, and project context Status(CLIArgs), // /// View and modify config // Config(CLIArgs), @@ -176,7 +179,8 @@ impl Commands { Commands::Setup(cmd) => &cmd.base, Commands::Docs(cmd) => &cmd.base, Commands::Sql(cmd) => &cmd.base, - Commands::Auth(cmd) => &cmd.base, + Commands::Login(cmd) => &cmd.base, + Commands::Logout(cmd) => &cmd.base, Commands::View(cmd) => &cmd.base, #[cfg(unix)] Commands::Eval(cmd) => &cmd.base, @@ -203,7 +207,8 @@ impl Commands { Commands::Setup(cmd) => &mut cmd.base, Commands::Docs(cmd) => &mut cmd.base, Commands::Sql(cmd) => &mut cmd.base, - Commands::Auth(cmd) => &mut cmd.base, + Commands::Login(cmd) => &mut cmd.base, + Commands::Logout(cmd) => &mut cmd.base, Commands::View(cmd) => &mut cmd.base, #[cfg(unix)] Commands::Eval(cmd) => &mut cmd.base, @@ -307,7 +312,8 @@ fn try_main() -> Result<()> { let command_result: Result<()> = runtime.block_on(async move { match cli.command { - Commands::Auth(cmd) => auth::run(cmd.base, cmd.args).await?, + Commands::Login(cmd) => auth::run_login_command(cmd.base, cmd.args).await?, + Commands::Logout(cmd) => auth::run_logout_command(cmd.base, cmd.args)?, 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?, @@ -500,7 +506,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 profiles, try `bt login --refresh --profile `; if refresh fails, re-run `bt login --oauth --profile `. Run `bt status --all` to inspect profile 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 d0d9fac0..075bf35e 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -41,7 +41,13 @@ const SETUP_WIZARD_POLL_PATH: &str = "/api/cli/wizard-session/poll"; const SETUP_WIZARD_POLL_INTERVAL: Duration = Duration::from_secs(2); const SETUP_WIZARD_MAX_CONSECUTIVE_POLL_FAILURES: usize = 30; const README_AGENT_SECTION_MARKERS: &[&str] = &[ - "bt eval", "bt sql", "bt view", "bt auth", "bt setup", "bt docs", + "bt eval", + "bt sql", + "bt view", + "bt login", + "bt logout", + "bt setup", + "bt docs", ]; const ALL_AGENTS: [Agent; 7] = [ Agent::Claude, @@ -1535,7 +1541,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" + "profile '{profile_name}' not found; run `bt status --all` to see available profiles" ); } @@ -5742,6 +5748,9 @@ mod tests { let profiles = vec![ auth::ProfileInfo { name: "zeta".to_string(), + auth: "api_key".to_string(), + app_url: "https://app.test.example".to_string(), + oauth_api_url: None, org_name: Some("Zeta Org".to_string()), user_name: None, email: None, @@ -5749,6 +5758,9 @@ mod tests { }, auth::ProfileInfo { name: "alpha".to_string(), + auth: "api_key".to_string(), + app_url: "https://app.test.example".to_string(), + oauth_api_url: None, org_name: Some("Alpha Org".to_string()), user_name: None, email: None, @@ -5777,6 +5789,9 @@ mod tests { base.profile = Some("missing".to_string()); let profiles = vec![auth::ProfileInfo { name: "work".to_string(), + auth: "api_key".to_string(), + app_url: "https://app.test.example".to_string(), + oauth_api_url: None, org_name: Some("Acme".to_string()), user_name: None, email: None, @@ -5793,6 +5808,9 @@ mod tests { let profiles = vec![ auth::ProfileInfo { name: "Acme".to_string(), + auth: "api_key".to_string(), + app_url: "https://app.test.example".to_string(), + oauth_api_url: None, org_name: Some("Acme".to_string()), user_name: None, email: None, @@ -5800,6 +5818,9 @@ mod tests { }, auth::ProfileInfo { name: "Acme-2".to_string(), + auth: "api_key".to_string(), + app_url: "https://app.test.example".to_string(), + oauth_api_url: None, org_name: Some("Acme".to_string()), user_name: None, email: None, diff --git a/src/status.rs b/src/status.rs index 93b555ed..d67bb138 100644 --- a/src/status.rs +++ b/src/status.rs @@ -10,23 +10,38 @@ use crate::{config, utils::resolve_profile_info}; #[command(after_help = "\ Examples: bt status + bt status --all bt status --json bt status --verbose ")] -pub struct StatusArgs {} +pub struct StatusArgs { + /// Include all saved login profiles and their connection status + #[arg(long)] + all: bool, +} #[derive(Serialize)] struct StatusOutput { org: Option, project: Option, + #[serde(skip_serializing_if = "Option::is_none")] + project_id: Option, profile: Option, #[serde(skip_serializing_if = "Option::is_none")] + auth: Option, + #[serde(skip_serializing_if = "Option::is_none")] + app_url: Option, + #[serde(skip_serializing_if = "Option::is_none")] + api_url: 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, source: Option, + #[serde(skip_serializing_if = "Option::is_none")] + profiles: Option>, } fn format_identity(p: &auth::ProfileInfo) -> Option { @@ -40,7 +55,13 @@ fn format_identity(p: &auth::ProfileInfo) -> Option { } } -pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { +fn format_auth(p: &auth::ProfileInfo) -> String { + format_identity(p) + .map(|identity| format!("{} — {identity}", p.auth)) + .unwrap_or_else(|| p.auth.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(); let local_path = config::local_path(); @@ -101,28 +122,64 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { config::project_from_config_for_context(&project_base, &merged_cfg, org.as_deref()); } + let project_id = (project.as_deref() == merged_cfg.project.as_deref()) + .then(|| merged_cfg.project_id.clone()) + .flatten(); + let profiles = if args.all { + Some(auth::profile_verifications().await?) + } else { + None + }; + if base.json { let output = StatusOutput { org, project, + project_id, profile: profile_info.as_ref().map(|p| p.name.clone()), + auth: profile_info.as_ref().map(|p| p.auth.clone()), + app_url: profile_info.as_ref().map(|p| p.app_url.clone()), + api_url: profile_info.as_ref().and_then(|p| p.oauth_api_url.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()), source, + profiles, }; println!("{}", serde_json::to_string(&output)?); return Ok(()); } + if let Some(profiles) = profiles.as_ref() { + if profiles.is_empty() { + eprintln!("No saved profiles. Run `bt login` to create one."); + } else { + for profile in profiles { + let status = match profile.status.as_str() { + "ok" => crate::ui::CommandStatus::Success, + "expired" => crate::ui::CommandStatus::Warning, + _ => crate::ui::CommandStatus::Error, + }; + crate::ui::print_command_status(status, &auth::format_verification_line(profile)); + } + if let Ok(path) = auth::credentials_path() { + eprintln!("\nCredentials: {}\n", path.display()); + } + } + } + if base.verbose { println!("org: {}", org.as_deref().unwrap_or("(unset)")); println!("project: {}", project.as_deref().unwrap_or("(unset)")); + println!("project_id: {}", project_id.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}"); - } + println!("app_url: {}", p.app_url); + println!( + "api_url: {}", + p.oauth_api_url.as_deref().unwrap_or("(derived per org)") + ); + println!("auth: {}", format_auth(p)); } if let Some(src) = source { println!("source: {src}"); @@ -345,6 +402,13 @@ mod tests { ) -> auth::ProfileInfo { auth::ProfileInfo { name: name.into(), + auth: if api_key_hint.is_some() { + "api_key".into() + } else { + "oauth".into() + }, + app_url: "https://app.test.example".into(), + oauth_api_url: None, org_name: None, user_name: user_name.map(Into::into), email: email.map(Into::into), diff --git a/src/utils/profile.rs b/src/utils/profile.rs index 51cd0f16..f9f2bc1a 100644 --- a/src/utils/profile.rs +++ b/src/utils/profile.rs @@ -101,6 +101,9 @@ mod tests { ) -> ProfileInfo { ProfileInfo { name: name.to_string(), + auth: "oauth".to_string(), + app_url: "https://app.test.example".to_string(), + oauth_api_url: None, org_name: org_name.map(ToOwned::to_owned), user_name: user_name.map(ToOwned::to_owned), email: email.map(ToOwned::to_owned), diff --git a/tests/cli.rs b/tests/cli.rs index acb09bfd..8d0c7fb7 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -109,6 +109,33 @@ fn top_level_help_shows_update_not_self() { .stdout(predicate::str::contains("self Self-management commands").not()); } +#[test] +fn status_all_json_includes_profile_urls() { + let home = tempfile::tempdir().expect("home tempdir"); + let config_home = tempfile::tempdir().expect("config tempdir"); + let auth_dir = config_home.path().join("bt"); + fs::create_dir_all(&auth_dir).expect("create auth dir"); + fs::write( + auth_dir.join("auth.json"), + r#"{"profiles":{"test-profile":{"auth_kind":"oauth","api_url":"https://oauth-api.test.example","app_url":"https://app.test.example","oauth_client_id":"bt_cli_test"}}}"#, + ) + .expect("write auth store"); + + let mut cmd = bt_command(); + clear_braintrust_auth_env(&mut cmd); + cmd.env("HOME", home.path()) + .env("XDG_CONFIG_HOME", config_home.path()) + .args(["status", "--all", "--json"]) + .assert() + .success() + .stdout(predicate::str::contains( + "\"app_url\":\"https://app.test.example\"", + )) + .stdout(predicate::str::contains( + "\"api_url\":\"https://oauth-api.test.example\"", + )); +} + #[test] fn topics_report_help_accepts_global_org_short_conflict_free() { bt_command() diff --git a/tests/functions.rs b/tests/functions.rs index 03f70300..19e14d42 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -190,10 +190,6 @@ fn sanitized_env_keys() -> &'static [&'static str] { ] } -fn auth_profiles_command(cwd: &Path, config_dir: &Path) -> Command { - auth_sub_command(cwd, config_dir, &["profiles"]) -} - #[derive(Debug, Clone)] struct MockProject { id: String, @@ -701,78 +697,6 @@ fn functions_help_lists_push_and_pull() { assert!(stdout.contains("pull")); } -#[test] -fn auth_profiles_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()) - .env("BRAINTRUST_API_KEY", "test-key") - .output() - .expect("run bt auth profiles with api key env"); - - 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("Auth source: --api-key/BRAINTRUST_API_KEY override")); - assert!(!stderr.contains("pass --prefer-profile or unset BRAINTRUST_API_KEY")); -} - -#[test] -fn auth_profiles_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()) - .env_remove("BRAINTRUST_API_KEY") - .output() - .expect("run bt auth profiles with dotenv api key"); - - 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("Auth source: --api-key/BRAINTRUST_API_KEY override")); - assert!(!stderr.contains("pass --prefer-profile or unset BRAINTRUST_API_KEY")); -} - -#[test] -fn auth_profiles_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()) - .arg("--json") - .output() - .expect("run bt auth profiles --json"); - - assert!(output.status.success()); - let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); - 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. fn seed_api_key_profile(config_dir: &Path) { @@ -784,62 +708,25 @@ fn seed_api_key_profile(config_dir: &Path) { .expect("write auth.json"); } -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`. - let mut cmd = Command::new(bt_binary_path()); - cmd.arg("auth") - .args(sub) - .current_dir(cwd) - .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") - .env_remove("BRAINTRUST_ENV_FILE"); - cmd -} - #[test] -fn auth_logout_json_with_no_profiles_emits_empty_status() { +fn root_login_refresh_uses_selected_profile() { let cwd = tempdir().expect("create temp cwd"); let config_dir = tempdir().expect("create temp config dir"); - - let output = auth_sub_command(cwd.path(), config_dir.path(), &["logout", "--json"]) - .output() - .expect("run bt auth logout --json"); - - assert!(output.status.success()); - let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); - assert_eq!(stdout, r#"{"status":"empty"}"#); -} - -#[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. - 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(), - &["refresh", "--profile", "test-profile", "--json"], - ) - .output() - .expect("run bt auth refresh --profile test-profile --json"); + let mut cmd = Command::new(bt_binary_path()); + cmd.args(["login", "--refresh", "--profile", "test-profile", "--json"]) + .current_dir(cwd.path()) + .env("XDG_CONFIG_HOME", config_dir.path()) + .env("APPDATA", config_dir.path()) + .env("BRAINTRUST_NO_COLOR", "1") + .env_remove("BRAINTRUST_API_KEY") + .env_remove("BRAINTRUST_ORG_NAME"); + let output = cmd.output().expect("run bt login --refresh"); assert!(!output.status.success()); - let stderr = String::from_utf8_lossy(&output.stderr); - assert!( - stderr.contains("only applies to oauth profiles"), - "expected oauth-only refresh hint, got: {stderr}" - ); - assert!(String::from_utf8_lossy(&output.stdout).trim().is_empty()); + assert!(String::from_utf8_lossy(&output.stderr) + .contains("`bt login --refresh` only applies to oauth profiles")); } #[test] From 99eb39ccb90aa622d0681889ec16c4869f8e658c Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Fri, 31 Jul 2026 18:49:47 -0700 Subject: [PATCH 6/7] cedric --- src/auth.rs | 375 ++++++++++++++++++---------------------------------- 1 file changed, 130 insertions(+), 245 deletions(-) diff --git a/src/auth.rs b/src/auth.rs index a3c67d70..30e01df4 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -29,10 +29,8 @@ use tokio::sync::oneshot; use crate::{ args::{BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL}, - config, - http::{build_http_client, build_http_client_from_builder, ApiClient}, - projects::api, - switch, ui, + http::{build_http_client, build_http_client_from_builder}, + ui, }; const KEYCHAIN_SERVICE: &str = "com.braintrust.bt.cli"; @@ -338,6 +336,21 @@ enum ApiKeyOrgMismatchAction { UseOauth, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OauthOrgMismatchAction { + SelectAvailableOrg, + SaveWithoutOrg, + Cancel, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OauthRequestedOrgResolution { + NoRequest, + UseRequested, + SelectAvailable, + SaveWithoutSelection, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RequestedOrgResolution { NoRequestedOrg, @@ -389,11 +402,6 @@ pub struct LogoutArgs { force: bool, } -struct PostLoginContextUpdate { - display: String, - path: PathBuf, -} - pub async fn run_login_command(base: BaseArgs, args: LoginArgs) -> Result<()> { if args.refresh { run_login_refresh(&base).await @@ -1301,16 +1309,6 @@ async fn run_login_set(base: &BaseArgs, args: LoginArgs) -> Result<()> { Some(login_app_url.clone()), org_constraint.as_ref().map(|org| 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(), - ) - .await - .context("login succeeded, but failed to update active context")?; let human = format_login_success(&selected_org, &profile_name, &selected_api_url); emit_result( @@ -1325,13 +1323,6 @@ async fn run_login_set(base: &BaseArgs, args: LoginArgs) -> Result<()> { }), || { ui::print_command_status(ui::CommandStatus::Success, &human); - ui::print_command_status( - ui::CommandStatus::Success, - &format!("Switched to {}", context_update.display), - ); - if base.verbose { - eprintln!("Wrote to {}", context_update.path.display()); - } }, ) } @@ -1406,7 +1397,30 @@ async fn run_login_oauth(base: &BaseArgs, args: LoginArgs) -> Result<()> { .await?; let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; let store = load_auth_store()?; - let selected_org = select_explicit_oauth_login_org(&login_orgs, base.org_name.as_deref())?; + let requested_org_resolution = resolve_requested_org_for_oauth_login( + &login_orgs, + base.org_name.as_deref(), + ui::can_prompt(), + prompt_for_oauth_org_mismatch, + )?; + let selected_org = match requested_org_resolution { + OauthRequestedOrgResolution::NoRequest + | OauthRequestedOrgResolution::SaveWithoutSelection => None, + OauthRequestedOrgResolution::UseRequested => base + .org_name + .as_deref() + .and_then(|name| find_login_org(&login_orgs, name)) + .cloned(), + OauthRequestedOrgResolution::SelectAvailable => select_login_org( + login_orgs.clone(), + None, + None, + true, + base.verbose, + false, + explicitly_quiet(base), + )?, + }; let selected_api_url = if selected_org.is_some() { resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)? } else { @@ -1426,21 +1440,6 @@ async fn run_login_oauth(base: &BaseArgs, args: LoginArgs) -> Result<()> { app_url.clone(), client_id.clone(), )?; - let context_update = match selected_org.as_ref() { - Some(org) => { - persist_post_login_context( - base, - &profile_name, - &oauth_tokens.access_token, - &selected_api_url, - &app_url, - Some(org), - ) - .await - } - None => persist_identity_login_context(&profile_name), - } - .context("login succeeded, but failed to update active context")?; let human = format_login_success(&selected_org, &profile_name, &selected_api_url); emit_result( @@ -1455,15 +1454,6 @@ async fn run_login_oauth(base: &BaseArgs, args: LoginArgs) -> Result<()> { }), || { ui::print_command_status(ui::CommandStatus::Success, &human); - let context_status = if selected_org.is_some() { - format!("Switched to {}", context_update.display) - } else { - format!("Using profile '{profile_name}' (organization unchanged)") - }; - ui::print_command_status(ui::CommandStatus::Success, &context_status); - if base.verbose { - eprintln!("Wrote to {}", context_update.path.display()); - } }, ) } @@ -1857,113 +1847,6 @@ fn format_login_success( } } -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(), - profile: None, - } -} - -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 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( - base: &BaseArgs, - profile_name: &str, - credential: &str, - api_url: &str, - app_url: &str, - selected_org: Option<&LoginOrgInfo>, -) -> Result { - 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 mut cfg = config::load_file(&path); - switch::apply_switch_config( - &mut cfg, - Some(profile_name), - selected_org.map(|org| org.name.as_str()), - project.as_ref(), - ); - config::save_file(&path, &cfg) - .context(format!("Could not save config to {}", path.display()))?; - - Ok(PostLoginContextUpdate { - display: format_post_login_context(selected_org, project.as_ref()), - path, - }) -} - -fn persist_identity_login_context(profile_name: &str) -> Result { - let path = if ui::can_prompt() && config::local_path().is_some() { - switch::select_scope()?.0 - } else { - config::global_path()? - }; - - let mut cfg = config::load_file(&path); - cfg.profile = Some(profile_name.to_string()); - config::save_file(&path, &cfg) - .context(format!("Could not save config to {}", path.display()))?; - - Ok(PostLoginContextUpdate { - display: profile_name.to_string(), - path, - }) -} - /// 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 { @@ -2301,17 +2184,58 @@ fn single_org_api_key_constraint<'a>( (credential.trim().starts_with("sk-") && orgs.len() == 1).then(|| &orgs[0]) } -fn select_explicit_oauth_login_org( +fn resolve_requested_org_for_oauth_login( orgs: &[LoginOrgInfo], requested_org_name: Option<&str>, -) -> Result> { - requested_org_name - .map(|org_name| { - find_login_org(orgs, org_name) - .cloned() - .ok_or_else(|| missing_requested_org_error(orgs, org_name)) - }) - .transpose() + can_prompt: bool, + choose_action: F, +) -> Result +where + F: FnOnce(&str, &[LoginOrgInfo]) -> Result, +{ + let Some(requested_org_name) = requested_org_name else { + return Ok(OauthRequestedOrgResolution::NoRequest); + }; + + if find_login_org(orgs, requested_org_name).is_some() { + return Ok(OauthRequestedOrgResolution::UseRequested); + } + + if !can_prompt { + return Err(missing_requested_org_error(orgs, requested_org_name)); + } + + match choose_action(requested_org_name, orgs)? { + OauthOrgMismatchAction::SelectAvailableOrg => { + Ok(OauthRequestedOrgResolution::SelectAvailable) + } + OauthOrgMismatchAction::SaveWithoutOrg => { + Ok(OauthRequestedOrgResolution::SaveWithoutSelection) + } + OauthOrgMismatchAction::Cancel => bail!("login cancelled"), + } +} + +fn prompt_for_oauth_org_mismatch( + requested_org_name: &str, + _orgs: &[LoginOrgInfo], +) -> Result { + let actions = [ + "Select a valid organization", + "Save login without selecting an organization", + "Cancel", + ]; + let selection = ui::fuzzy_select( + &format!("Org '{requested_org_name}' is not available. Continue with"), + &actions, + 0, + )?; + Ok(match selection { + 0 => OauthOrgMismatchAction::SelectAvailableOrg, + 1 => OauthOrgMismatchAction::SaveWithoutOrg, + 2 => OauthOrgMismatchAction::Cancel, + _ => unreachable!("fuzzy_select returned out-of-range index"), + }) } fn select_login_org( @@ -5047,93 +4971,54 @@ mod tests { } #[test] - fn oauth_login_only_selects_an_explicit_org() { + fn oauth_login_org_resolution_handles_valid_missing_and_cancelled_requests() { let orgs = vec![ login_org("org_1", "test-org"), login_org("org_2", "other-org"), ]; - assert!(select_explicit_oauth_login_org(&orgs, None) - .expect("no org selection") - .is_none()); assert_eq!( - select_explicit_oauth_login_org(&orgs, Some("other-org")) - .expect("explicit org selection") - .map(|org| org.id), - Some("org_2".to_string()) + resolve_requested_org_for_oauth_login(&orgs, None, false, |_, _| { + panic!("prompt should not be called") + }) + .expect("no org selection"), + OauthRequestedOrgResolution::NoRequest ); - } - - #[tokio::test] - async fn identity_login_preserves_existing_org_and_project() { - let _env = TestEnv::new(None, None).await; - crate::config::save_global(&crate::config::Config { - profile: Some("old-profile".to_string()), - org: Some("test-org".to_string()), - project: Some("test-project".to_string()), - project_id: Some("proj_test".to_string()), - ..Default::default() - }) - .expect("save initial config"); - - persist_identity_login_context("new-profile").expect("persist identity"); - let cfg = crate::config::load_global().expect("load global config"); - - assert_eq!(cfg.profile.as_deref(), Some("new-profile")); - 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")); - } - - #[tokio::test] - 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()), - ..Default::default() - }) - .expect("save initial config"); - - let update = persist_post_login_context( - &make_base(), - "work", - "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"); - - assert_eq!(update.display, "acme"); - assert_eq!(cfg.profile.as_deref(), Some("work")); - assert_eq!(cfg.org.as_deref(), Some("acme")); - assert_eq!(cfg.project, None); - assert_eq!(cfg.project_id, None); - } - - #[tokio::test] - async fn resolve_post_login_project_rejects_cross_org_default_project() { - let mut base = make_base(); - base.project = Some("demo-project".to_string()); - - let err = resolve_post_login_project( - &base, - "test-api-key", - "https://api.example.test", - "https://www.example.test", - None, - ) - .await - .expect_err("cross-org project selection should fail"); - + assert_eq!( + resolve_requested_org_for_oauth_login(&orgs, Some("other-org"), false, |_, _| { + panic!("prompt should not be called") + }) + .expect("matching org"), + OauthRequestedOrgResolution::UseRequested + ); + let err = + resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), false, |_, _| { + panic!("prompt should not be called") + }) + .expect_err("non-interactive mismatch should fail"); assert!(err .to_string() - .contains("cannot set a default project in cross-org mode")); + .contains("org 'missing-org' not found. Available: test-org, other-org")); + assert_eq!( + resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), true, |_, _| { + Ok(OauthOrgMismatchAction::SelectAvailableOrg) + }) + .expect("select another org"), + OauthRequestedOrgResolution::SelectAvailable + ); + assert_eq!( + resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), true, |_, _| { + Ok(OauthOrgMismatchAction::SaveWithoutOrg) + }) + .expect("save without org"), + OauthRequestedOrgResolution::SaveWithoutSelection + ); + let err = + resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), true, |_, _| { + Ok(OauthOrgMismatchAction::Cancel) + }) + .expect_err("cancel login"); + assert_eq!(err.to_string(), "login cancelled"); } #[test] From 9e1bf9fffd824d0e5439659ce627a9f0775f17db Mon Sep 17 00:00:00 2001 From: Ankur Goyal Date: Sat, 1 Aug 2026 14:00:17 -0700 Subject: [PATCH 7/7] reject org and project in login --- src/args.rs | 68 +++++-- src/auth.rs | 388 ++------------------------------------- src/config/mod.rs | 24 +-- src/datasets/pipeline.rs | 21 +-- src/functions/push.rs | 21 +-- src/main.rs | 27 ++- src/setup/mod.rs | 21 +-- src/switch.rs | 17 +- src/traces.rs | 21 +-- 9 files changed, 91 insertions(+), 517 deletions(-) diff --git a/src/args.rs b/src/args.rs index c15c6478..313341bf 100644 --- a/src/args.rs +++ b/src/args.rs @@ -11,8 +11,8 @@ pub enum ArgValueSource { EnvVariable, } -#[derive(Debug, Clone, Args)] -pub struct BaseArgs { +#[derive(Debug, Clone, Args, Default)] +pub struct LoginBaseArgs { /// Output as JSON #[arg(long, global = true)] pub json: bool, @@ -46,20 +46,6 @@ pub struct BaseArgs { #[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, - - /// Override active project - #[arg( - short = 'p', - long, - env = "BRAINTRUST_DEFAULT_PROJECT", - hide_env_values = true, - global = true - )] - pub project: Option, - /// Override stored API key (or via BRAINTRUST_API_KEY) #[arg(long, env = "BRAINTRUST_API_KEY", global = true, hide = true)] pub api_key: Option, @@ -108,16 +94,36 @@ pub struct BaseArgs { pub env_file: Option, } +#[derive(Debug, Clone, Args, Default)] +pub struct BaseArgs { + #[command(flatten)] + pub login: LoginBaseArgs, + + /// Override active org (or via BRAINTRUST_ORG_NAME) + #[arg(short = 'o', long = "org", env = "BRAINTRUST_ORG_NAME", global = true)] + pub org_name: Option, + + /// Override active project + #[arg( + short = 'p', + long, + env = "BRAINTRUST_DEFAULT_PROJECT", + hide_env_values = true, + global = true + )] + pub project: Option, +} + #[derive(Debug, Clone, Args)] -pub struct CLIArgs { +pub struct CLIArgs { #[command(flatten)] pub args: T, #[command(flatten, next_help_heading = "Global options")] - pub base: BaseArgs, + pub base: B, } -impl BaseArgs { +impl LoginBaseArgs { pub fn ca_cert(&self) -> Option<&Path> { self.ca_cert.as_deref() } @@ -127,6 +133,30 @@ impl BaseArgs { } } +impl std::ops::Deref for BaseArgs { + type Target = LoginBaseArgs; + + fn deref(&self) -> &Self::Target { + &self.login + } +} + +impl std::ops::DerefMut for BaseArgs { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.login + } +} + +impl From for BaseArgs { + fn from(login: LoginBaseArgs) -> Self { + Self { + login, + org_name: None, + project: None, + } + } +} + pub fn has_explicit_profile_arg(args: &[OsString]) -> bool { let mut idx = 1usize; while idx < args.len() { diff --git a/src/auth.rs b/src/auth.rs index 30e01df4..032271c5 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -330,35 +330,6 @@ struct LoginOrgInfo { api_url: Option, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ApiKeyOrgMismatchAction { - UseApiKey, - UseOauth, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum OauthOrgMismatchAction { - SelectAvailableOrg, - SaveWithoutOrg, - Cancel, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum OauthRequestedOrgResolution { - NoRequest, - UseRequested, - SelectAvailable, - SaveWithoutSelection, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum RequestedOrgResolution { - NoRequestedOrg, - UseRequestedOrg, - IgnoreRequestedOrg, - SwitchToOauth, -} - #[derive(Debug, Clone, Deserialize)] struct OAuthTokenResponse { access_token: String, @@ -1265,26 +1236,9 @@ async fn run_login_set(base: &BaseArgs, args: LoginArgs) -> Result<()> { let login_orgs = fetch_login_orgs(&api_key, &login_app_url).await?; let org_constraint = single_org_api_key_constraint(&api_key, &login_orgs).cloned(); let store = load_auth_store()?; - let requested_org_resolution = resolve_requested_org_for_api_key_login( - &login_orgs, - base.org_name.as_deref(), - ui::can_prompt(), - prompt_for_auth_method_for_missing_requested_org, - )?; - 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 default_org_name = default_login_org_name(&store, base.profile.as_deref()); let selected_org = select_login_org( login_orgs.clone(), - match requested_org_resolution { - RequestedOrgResolution::UseRequestedOrg => base.org_name.as_deref(), - RequestedOrgResolution::NoRequestedOrg | RequestedOrgResolution::IgnoreRequestedOrg => { - None - } - RequestedOrgResolution::SwitchToOauth => unreachable!("handled above"), - }, default_org_name.as_deref(), interactive, base.verbose, @@ -1395,37 +1349,10 @@ async fn run_login_oauth(base: &BaseArgs, args: LoginArgs) -> Result<()> { pkce_verifier, ) .await?; - let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; + fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; let store = load_auth_store()?; - let requested_org_resolution = resolve_requested_org_for_oauth_login( - &login_orgs, - base.org_name.as_deref(), - ui::can_prompt(), - prompt_for_oauth_org_mismatch, - )?; - let selected_org = match requested_org_resolution { - OauthRequestedOrgResolution::NoRequest - | OauthRequestedOrgResolution::SaveWithoutSelection => None, - OauthRequestedOrgResolution::UseRequested => base - .org_name - .as_deref() - .and_then(|name| find_login_org(&login_orgs, name)) - .cloned(), - OauthRequestedOrgResolution::SelectAvailable => select_login_org( - login_orgs.clone(), - None, - None, - true, - base.verbose, - false, - explicitly_quiet(base), - )?, - }; - let selected_api_url = if selected_org.is_some() { - resolve_profile_api_url(base.api_url.clone(), selected_org.as_ref(), &login_orgs)? - } else { - api_url.clone() - }; + let selected_org: Option = None; + let selected_api_url = api_url.clone(); 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(), &app_url, &jwt_id, &store)?; @@ -1665,18 +1592,7 @@ fn resolve_profile_name( .to_string()) } -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 default_login_org_name(store: &AuthStore, profile_name: Option<&str>) -> Option { let profile_name = profile_name .map(str::trim) .filter(|name| !name.is_empty())?; @@ -1924,14 +1840,15 @@ fn run_login_delete(profile_name: &str, force: bool, base_json: bool) -> Result< } fn run_login_logout(base: BaseArgs, args: LogoutArgs) -> Result<()> { + let base_json = base.json; let store = load_auth_store()?; if store.profiles.is_empty() { - return emit_result(base.json, serde_json::json!({ "status": "empty" }), || { + return emit_result(base_json, serde_json::json!({ "status": "empty" }), || { println!("No saved profiles.") }); } - let profile_name = if let Some(p) = base.profile { + let profile_name = if let Some(p) = base.login.profile { let p = p.trim().to_string(); if !store.profiles.contains_key(&p) { return Err(profile_not_found_err(&p, &store)); @@ -1947,7 +1864,7 @@ fn run_login_logout(base: BaseArgs, args: LogoutArgs) -> Result<()> { bail!("multiple profiles exist. Use --profile to specify which one."); }; - run_login_delete(&profile_name, args.force, base.json) + run_login_delete(&profile_name, args.force, base_json) } enum ProfileStatus { @@ -2184,63 +2101,8 @@ fn single_org_api_key_constraint<'a>( (credential.trim().starts_with("sk-") && orgs.len() == 1).then(|| &orgs[0]) } -fn resolve_requested_org_for_oauth_login( - orgs: &[LoginOrgInfo], - requested_org_name: Option<&str>, - can_prompt: bool, - choose_action: F, -) -> Result -where - F: FnOnce(&str, &[LoginOrgInfo]) -> Result, -{ - let Some(requested_org_name) = requested_org_name else { - return Ok(OauthRequestedOrgResolution::NoRequest); - }; - - if find_login_org(orgs, requested_org_name).is_some() { - return Ok(OauthRequestedOrgResolution::UseRequested); - } - - if !can_prompt { - return Err(missing_requested_org_error(orgs, requested_org_name)); - } - - match choose_action(requested_org_name, orgs)? { - OauthOrgMismatchAction::SelectAvailableOrg => { - Ok(OauthRequestedOrgResolution::SelectAvailable) - } - OauthOrgMismatchAction::SaveWithoutOrg => { - Ok(OauthRequestedOrgResolution::SaveWithoutSelection) - } - OauthOrgMismatchAction::Cancel => bail!("login cancelled"), - } -} - -fn prompt_for_oauth_org_mismatch( - requested_org_name: &str, - _orgs: &[LoginOrgInfo], -) -> Result { - let actions = [ - "Select a valid organization", - "Save login without selecting an organization", - "Cancel", - ]; - let selection = ui::fuzzy_select( - &format!("Org '{requested_org_name}' is not available. Continue with"), - &actions, - 0, - )?; - Ok(match selection { - 0 => OauthOrgMismatchAction::SelectAvailableOrg, - 1 => OauthOrgMismatchAction::SaveWithoutOrg, - 2 => OauthOrgMismatchAction::Cancel, - _ => unreachable!("fuzzy_select returned out-of-range index"), - }) -} - fn select_login_org( mut orgs: Vec, - requested_org_name: Option<&str>, default_org_name: Option<&str>, interactive: bool, verbose: bool, @@ -2257,13 +2119,6 @@ fn select_login_org( .then_with(|| a.name.cmp(&b.name)) }); - if let Some(name) = requested_org_name { - return find_login_org(&orgs, name) - .cloned() - .map(Some) - .ok_or_else(|| missing_requested_org_error(&orgs, name)); - } - if orgs.len() == 1 { return Ok(Some(orgs.into_iter().next().expect("org exists"))); } @@ -2320,13 +2175,6 @@ fn move_default_login_org_first( true } -fn find_login_org<'a>( - orgs: &'a [LoginOrgInfo], - requested_org_name: &str, -) -> Option<&'a LoginOrgInfo> { - find_login_org_index(orgs, requested_org_name).map(|idx| &orgs[idx]) -} - fn find_login_org_index(orgs: &[LoginOrgInfo], requested_org_name: &str) -> Option { orgs.iter() .position(|org| org.name == requested_org_name) @@ -2337,65 +2185,6 @@ 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() - .map(|org| org.name.as_str()) - .collect::>() - .join(", "); - anyhow::anyhow!("org '{requested_org_name}' not found. Available: {available}") -} - -fn resolve_requested_org_for_api_key_login( - orgs: &[LoginOrgInfo], - requested_org_name: Option<&str>, - can_prompt: bool, - choose_auth_method: F, -) -> Result -where - F: FnOnce(&str, &[LoginOrgInfo]) -> Result, -{ - let Some(requested_org_name) = requested_org_name else { - return Ok(RequestedOrgResolution::NoRequestedOrg); - }; - - if find_login_org(orgs, requested_org_name).is_some() { - return Ok(RequestedOrgResolution::UseRequestedOrg); - } - - if !can_prompt { - return Err(missing_requested_org_error(orgs, requested_org_name)); - } - - match choose_auth_method(requested_org_name, orgs)? { - ApiKeyOrgMismatchAction::UseApiKey => Ok(RequestedOrgResolution::IgnoreRequestedOrg), - ApiKeyOrgMismatchAction::UseOauth => Ok(RequestedOrgResolution::SwitchToOauth), - } -} - -fn prompt_for_auth_method_for_missing_requested_org( - requested_org_name: &str, - orgs: &[LoginOrgInfo], -) -> Result { - let api_key_label = if orgs.len() == 1 { - format!("API key ({})", orgs[0].name) - } else { - "API key (use available org)".to_string() - }; - let methods = ["OAuth (browser)".to_string(), api_key_label]; - let method_refs: Vec<&str> = methods.iter().map(String::as_str).collect(); - let selection = ui::fuzzy_select( - &format!("Org '{requested_org_name}' is not available for this API key. Continue with"), - &method_refs, - 0, - )?; - Ok(match selection { - 0 => ApiKeyOrgMismatchAction::UseOauth, - 1 => ApiKeyOrgMismatchAction::UseApiKey, - _ => unreachable!("fuzzy_select returned out-of-range index"), - }) -} - fn resolve_profile_api_url( explicit_api_url: Option, selected_org: Option<&LoginOrgInfo>, @@ -3579,26 +3368,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, - profile: None, - profile_explicit: false, - project: None, - org_name: None, - api_key: None, - api_key_source: None, - prefer_profile: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn auth_config(profile: Option<&str>, org: Option<&str>) -> crate::config::Config { @@ -4741,7 +4511,7 @@ mod tests { ); assert_eq!( - default_login_org_name(&store, Some(" work "), None).as_deref(), + default_login_org_name(&store, Some(" work ")).as_deref(), Some("acme") ); } @@ -4750,24 +4520,7 @@ mod tests { fn default_login_org_name_does_not_treat_profile_name_as_org() { let store = AuthStore::default(); - assert_eq!(default_login_org_name(&store, Some(" acme "), None), None); - } - - #[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 - ); + assert_eq!(default_login_org_name(&store, Some(" acme ")), None); } #[test] @@ -4970,123 +4723,6 @@ mod tests { assert!(single_org_api_key_constraint("sk-test-key", &multiple_orgs).is_none()); } - #[test] - fn oauth_login_org_resolution_handles_valid_missing_and_cancelled_requests() { - let orgs = vec![ - login_org("org_1", "test-org"), - login_org("org_2", "other-org"), - ]; - - assert_eq!( - resolve_requested_org_for_oauth_login(&orgs, None, false, |_, _| { - panic!("prompt should not be called") - }) - .expect("no org selection"), - OauthRequestedOrgResolution::NoRequest - ); - assert_eq!( - resolve_requested_org_for_oauth_login(&orgs, Some("other-org"), false, |_, _| { - panic!("prompt should not be called") - }) - .expect("matching org"), - OauthRequestedOrgResolution::UseRequested - ); - let err = - resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), false, |_, _| { - panic!("prompt should not be called") - }) - .expect_err("non-interactive mismatch should fail"); - assert!(err - .to_string() - .contains("org 'missing-org' not found. Available: test-org, other-org")); - assert_eq!( - resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), true, |_, _| { - Ok(OauthOrgMismatchAction::SelectAvailableOrg) - }) - .expect("select another org"), - OauthRequestedOrgResolution::SelectAvailable - ); - assert_eq!( - resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), true, |_, _| { - Ok(OauthOrgMismatchAction::SaveWithoutOrg) - }) - .expect("save without org"), - OauthRequestedOrgResolution::SaveWithoutSelection - ); - let err = - resolve_requested_org_for_oauth_login(&orgs, Some("missing-org"), true, |_, _| { - Ok(OauthOrgMismatchAction::Cancel) - }) - .expect_err("cancel login"); - assert_eq!(err.to_string(), "login cancelled"); - } - - #[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( - &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) - }, - ) - .expect("resolve"); - - assert_eq!(resolution, RequestedOrgResolution::IgnoreRequestedOrg); - } - #[test] fn obscure_api_key_standard() { assert_eq!(obscure_api_key("sk-LumEdp0BbLRzhJwO"), "sk-****zhJwO"); diff --git a/src/config/mod.rs b/src/config/mod.rs index e1e4bf9b..5639e90a 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -427,24 +427,12 @@ mod tests { fn base_with_profile(profile: Option<&str>) -> BaseArgs { BaseArgs { - json: false, - verbose: false, - verbose_source: None, - quiet: false, - quiet_source: None, - 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, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, + login: crate::args::LoginBaseArgs { + profile: profile.map(str::to_string), + profile_explicit: profile.is_some(), + ..Default::default() + }, + ..Default::default() } } diff --git a/src/datasets/pipeline.rs b/src/datasets/pipeline.rs index f1c43e49..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, - profile: None, - profile_explicit: false, - org_name: None, - project: None, - api_key: None, - api_key_source: None, - prefer_profile: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn test_source() -> PipelineSourceInspect { diff --git a/src/functions/push.rs b/src/functions/push.rs index 4d986b2f..a4b21fda 100644 --- a/src/functions/push.rs +++ b/src/functions/push.rs @@ -4055,25 +4055,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, - profile: None, - profile_explicit: false, - org_name: None, - project: None, - api_key: None, - api_key_source: None, - prefer_profile: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } } diff --git a/src/main.rs b/src/main.rs index 24341fc6..ee2c405c 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::{has_explicit_profile_arg, ArgValueSource, CLIArgs, LoginBaseArgs}; const DEFAULT_CANARY_VERSION: &str = concat!(env!("CARGO_PKG_VERSION"), "-canary.dev"); pub(crate) const CLI_VERSION: &str = match option_env!("BT_VERSION_STRING") { @@ -131,7 +131,7 @@ enum Commands { /// Run SQL queries against Braintrust Sql(CLIArgs), /// Log in to Braintrust - Login(CLIArgs), + Login(CLIArgs), /// Remove a saved Braintrust login Logout(CLIArgs), /// View logs, traces, and spans @@ -173,7 +173,7 @@ enum Commands { } impl Commands { - fn base(&self) -> &BaseArgs { + fn base(&self) -> &LoginBaseArgs { match self { Commands::Init(cmd) => &cmd.base, Commands::Setup(cmd) => &cmd.base, @@ -201,7 +201,7 @@ impl Commands { } } - fn base_mut(&mut self) -> &mut BaseArgs { + fn base_mut(&mut self) -> &mut LoginBaseArgs { match self { Commands::Init(cmd) => &mut cmd.base, Commands::Setup(cmd) => &mut cmd.base, @@ -281,7 +281,7 @@ fn handle_version_json(argv: &[OsString]) -> Result { Ok(true) } -fn apply_runtime_env_overrides(base: &BaseArgs) { +fn apply_runtime_env_overrides(base: &LoginBaseArgs) { // Apply the CLI-owned override once so reqwest and inherited child // commands consistently observe BRAINTRUST_CA_CERT/--ca-cert precedence // over any ambient SSL_CERT_FILE. @@ -312,7 +312,7 @@ fn try_main() -> Result<()> { let command_result: Result<()> = runtime.block_on(async move { match cli.command { - Commands::Login(cmd) => auth::run_login_command(cmd.base, cmd.args).await?, + Commands::Login(cmd) => auth::run_login_command(cmd.base.into(), cmd.args).await?, Commands::Logout(cmd) => auth::run_logout_command(cmd.base, cmd.args)?, Commands::View(cmd) => traces::run(cmd.base, cmd.args).await?, Commands::Init(cmd) => init::run(cmd.base, cmd.args).await?, @@ -350,7 +350,7 @@ fn try_main() -> Result<()> { command_result } -fn apply_base_arg_sources(matches: &ArgMatches, base: &mut BaseArgs) { +fn apply_base_arg_sources(matches: &ArgMatches, base: &mut LoginBaseArgs) { 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.api_key_source = find_value_source(matches, "api_key").and_then(map_value_source); @@ -386,7 +386,7 @@ fn map_value_source(source: ValueSource) -> Option { } } -fn configure_output(base: &BaseArgs) { +fn configure_output(base: &LoginBaseArgs) { let mut disable_color = base.no_color || std::env::var_os("NO_COLOR").is_some(); // TERM is a terminal capability signal; it isn't a user-facing config knob. @@ -587,6 +587,17 @@ mod tests { assert!(cli.command.base().verbose_explicit()); } + #[test] + fn login_rejects_context_selection_flags() { + for args in [ + ["bt", "login", "--org", "test-org"], + ["bt", "login", "--project", "test-project"], + ] { + let err = Cli::try_parse_from(args).expect_err("context flag should be rejected"); + assert_eq!(err.kind(), clap::error::ErrorKind::UnknownArgument); + } + } + #[test] fn default_verbose_output_is_not_explicit_verbose() { let matches = Cli::command() diff --git a/src/setup/mod.rs b/src/setup/mod.rs index 075bf35e..e77b2228 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -5395,26 +5395,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, - profile: None, - profile_explicit: false, - org_name: None, - project: None, - api_key: None, - api_key_source: None, - prefer_profile: 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/switch.rs b/src/switch.rs index 3916b878..865f4f12 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -317,24 +317,9 @@ mod tests { 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, - 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, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, + ..Default::default() } } diff --git a/src/traces.rs b/src/traces.rs index 1ec1c80d..99184356 100644 --- a/src/traces.rs +++ b/src/traces.rs @@ -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, - profile: None, - profile_explicit: false, - org_name: None, - project: None, - api_key: None, - api_key_source: None, - prefer_profile: false, - api_url: None, - app_url: None, - ca_cert: None, - env_file: None, - } + BaseArgs::default() } fn parsed_url_with_org(org: &str) -> ParsedTraceUrl {