diff --git a/src/auth.rs b/src/auth.rs index 1d5af91..b28e967 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1161,6 +1161,11 @@ where })?) }; + // Once a profile is selected, the profile's own org_name wins over the config-file + // `org` (cfg_org). cfg_org may be a stale/legacy value authored for a *different* + // profile (e.g. global config still has the previous profile's org), and letting it + // override the selected profile's org produces a JWT/org-name mismatch and makes + // `bt switch --org ` switch to the wrong org. An explicit --org flag still wins. return Ok(ResolvedAuth { api_key, api_url: base.api_url.clone().or_else(|| profile.api_url.clone()), @@ -1168,8 +1173,8 @@ where org_name: base .org_name .clone() - .or_else(|| cfg_org.clone()) - .or_else(|| profile.org_name.clone()), + .or_else(|| profile.org_name.clone()) + .or_else(|| cfg_org.clone()), is_oauth, }); } @@ -1262,9 +1267,15 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { } async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { - let api_url = base + // Resolve the deployment to authenticate against using exactly one OAuth browser flow. + // Priority: explicit --api-url, then the existing profile's api_url (so re-logging into a + // staging profile re-authenticates against staging and gets a staging-issued JWT), then + // the default (prod). This must match the deployment the selected org lives on; a JWT is + // only valid for the deployment that issued it. + let exchange_api_url = base .api_url .clone() + .or_else(|| existing_profile_api_url(base.profile.as_deref())) .unwrap_or_else(|| DEFAULT_API_URL.to_string()); let app_url = base .app_url @@ -1281,6 +1292,105 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { .clone() .unwrap_or_else(|| default_oauth_client_id(provisional_profile)); + // Run the browser OAuth authorize + token exchange against `exchange_api_url`, returning + // the issued tokens. This is the only browser prompt. The deployment that issues the JWT + // is the only deployment the JWT is valid against, so the profile's stored `api_url` must + // match. + let oauth_tokens = + oauth_authorize_and_exchange(&exchange_api_url, &client_id, base, &args).await?; + + // List the orgs available to this token and let the user pick one. + let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; + let selected_org = select_login_org( + login_orgs.clone(), + base.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)?; + + // A JWT is only valid for the deployment that issued it. If the user picked an org that + // lives on a *different* deployment than the one we just authenticated against, the + // token cannot serve that org — we cannot fix this with another browser flow (that would + // be a second login and still wouldn't change which deployment issued the JWT we keep). + // Surface an actionable error telling the user how to authenticate against that org's + // deployment instead. Never re-open the browser here. + if let Some(org) = selected_org.as_ref() { + let org_api_url = org + .api_url + .as_deref() + .unwrap_or(&selected_api_url) + .trim_end_matches('/'); + if !same_deployment(org_api_url, &exchange_api_url) { + bail!( + "org '{}' lives on a different deployment ({}) than the one used for login ({}).\n\ + Re-run with `--api-url {}` to authenticate against that deployment.", + org.name, + org_api_url, + exchange_api_url, + org_api_url, + ); + } + } + + let store = load_auth_store()?; + let jwt_id = decode_jwt_identity(&oauth_tokens.access_token); + let (profile_name, should_confirm_overwrite) = resolve_oauth_login_profile_name( + base.profile.as_deref(), + selected_org.as_ref().map(|org| org.name.as_str()), + &selected_api_url, + &app_url, + &jwt_id, + &store, + )?; + if should_confirm_overwrite { + confirm_profile_overwrite(&profile_name)?; + } + + commit_oauth_profile( + &profile_name, + &oauth_tokens, + selected_api_url.clone(), + app_url.clone(), + client_id.clone(), + selected_org.as_ref().map(|org| org.name.clone()), + )?; + + ui::print_command_status( + ui::CommandStatus::Success, + &format_login_success(&selected_org, &profile_name, &selected_api_url), + ); + + Ok(()) +} + +/// Look up the stored `api_url` for an existing profile, so re-logging into a profile (e.g. +/// `bt auth login --profile "BT Staging"` where that profile already points at staging) +/// authenticates against the same deployment and gets a JWT issued by that deployment. +fn existing_profile_api_url(profile: Option<&str>) -> Option { + let name = profile?.trim(); + if name.is_empty() { + return None; + } + let store = load_auth_store().ok()?; + store + .profiles + .get(name) + .and_then(|profile| profile.api_url.clone()) +} + +/// Run the browser-based OAuth authorize + PKCE token exchange against `api_url`. +/// Returns the issued `OAuthTokenResponse`. This is the sole browser prompt during +/// `bt auth login --oauth`; it must never be called twice in one invocation. +async fn oauth_authorize_and_exchange( + api_url: &str, + client_id: &str, + base: &BaseArgs, + args: &AuthLoginArgs, +) -> Result { let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256(); let state = generate_random_token(32)?; @@ -1292,7 +1402,7 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { .context("failed to read callback listener address")? .port(); let redirect_uri = format!("http://127.0.0.1:{callback_port}/callback"); - let oauth_client = build_oauth_client(&api_url, &client_id, Some(&redirect_uri))?; + let oauth_client = build_oauth_client(api_url, client_id, Some(&redirect_uri))?; let (authorize_url, _) = oauth_client .authorize_url(|| CsrfToken::new(state.clone())) .add_scope(Scope::new(OAUTH_SCOPE.to_string())) @@ -1327,54 +1437,24 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { bail!("oauth state mismatch; please try again"); } - let oauth_tokens = exchange_oauth_authorization_code( - &api_url, - &client_id, - &redirect_uri, - &auth_code, - pkce_verifier, - ) - .await?; - let login_orgs = fetch_login_orgs(&oauth_tokens.access_token, &app_url).await?; - let selected_org = select_login_org( - login_orgs.clone(), - base.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 store = load_auth_store()?; - let jwt_id = decode_jwt_identity(&oauth_tokens.access_token); - let (profile_name, should_confirm_overwrite) = resolve_oauth_login_profile_name( - base.profile.as_deref(), - selected_org.as_ref().map(|org| org.name.as_str()), - &selected_api_url, - &app_url, - &jwt_id, - &store, - )?; - if should_confirm_overwrite { - confirm_profile_overwrite(&profile_name)?; - } - - commit_oauth_profile( - &profile_name, - &oauth_tokens, - selected_api_url.clone(), - app_url.clone(), - client_id.clone(), - selected_org.as_ref().map(|org| org.name.clone()), - )?; - - ui::print_command_status( - ui::CommandStatus::Success, - &format_login_success(&selected_org, &profile_name, &selected_api_url), - ); + exchange_oauth_authorization_code(api_url, client_id, &redirect_uri, &auth_code, pkce_verifier) + .await +} - Ok(()) +/// Compare two API URLs by host (ignoring trailing slashes), used to decide whether a +/// selected org's `api_url` matches the deployment that issued the current JWT. +fn same_deployment(a: &str, b: &str) -> bool { + let host_of = |url: &str| { + url.trim_end_matches('/') + .split("//") + .nth(1) + .unwrap_or("") + .split('/') + .next() + .unwrap_or("") + .to_ascii_lowercase() + }; + !a.is_empty() && !b.is_empty() && host_of(a) == host_of(b) } pub(crate) fn commit_api_key_profile( @@ -4223,7 +4303,11 @@ mod tests { } #[test] - fn resolve_auth_config_org_overrides_profile_org() { + fn resolve_auth_selected_profile_org_wins_over_stale_config_org() { + // A selected profile owns its org. The config-file `org` (cfg_org) may be a stale + // value left over from a different profile (e.g. global config still has the previous + // profile's org). Letting it override the selected profile's org produces a JWT/org + // mismatch, so the profile's org_name must win once a profile is selected. let mut base = make_base(); base.profile = Some("default-profile".to_string()); @@ -4245,6 +4329,60 @@ mod tests { ) .expect("resolve"); assert_eq!(resolved.api_key.as_deref(), Some("profile-key")); + assert_eq!(resolved.org_name.as_deref(), Some("profile-org")); + } + + #[test] + fn resolve_auth_explicit_org_flag_overrides_profile_org() { + // An explicit --org flag must still win over the selected profile's org_name. + let mut base = make_base(); + base.profile = Some("default-profile".to_string()); + base.org_name = Some("flag-org".to_string()); + + let mut store = AuthStore::default(); + store.profiles.insert( + "default-profile".into(), + AuthProfile { + org_name: Some("profile-org".into()), + ..Default::default() + }, + ); + let cfg_org = Some("local-org".to_string()); + + let resolved = resolve_auth_from_store_with_secret_lookup( + &base, + &store, + |_| Ok(Some("profile-key".into())), + &cfg_org, + ) + .expect("resolve"); + assert_eq!(resolved.org_name.as_deref(), Some("flag-org")); + } + + #[test] + fn resolve_auth_profile_without_org_falls_back_to_config_org() { + // If the selected profile has no org_name, keep using cfg_org so behavior is + // unchanged for legacy/api-key profiles that rely on the config org. + let mut base = make_base(); + base.profile = Some("default-profile".to_string()); + + let mut store = AuthStore::default(); + store.profiles.insert( + "default-profile".into(), + AuthProfile { + org_name: None, + ..Default::default() + }, + ); + let cfg_org = Some("local-org".to_string()); + + let resolved = resolve_auth_from_store_with_secret_lookup( + &base, + &store, + |_| Ok(Some("profile-key".into())), + &cfg_org, + ) + .expect("resolve"); assert_eq!(resolved.org_name.as_deref(), Some("local-org")); } @@ -4681,4 +4819,36 @@ mod tests { let result = env.login_read_only_with_base(base).await; assert_err_contains(result, "no login credentials found"); } + + #[test] + fn same_deployment_compares_hosts_ignoring_trailing_slash() { + assert!(same_deployment( + "https://staging-api.braintrust.dev", + "https://staging-api.braintrust.dev/", + )); + assert!(same_deployment( + "https://api.braintrust.dev", + "https://api.braintrust.dev", + )); + } + + #[test] + fn same_deployment_detects_cross_deployment_orgs() { + // A prod-issued JWT vs a staging org's api_url must be detected as a mismatch so + // `run_login_oauth` re-runs OAuth against the staging deployment. + assert!(!same_deployment( + "https://staging-api.braintrust.dev", + "https://api.braintrust.dev", + )); + } + + #[test] + fn same_deployment_is_case_insensitive_and_handles_empty() { + assert!(same_deployment( + "https://API.Braintrust.dev", + "https://api.braintrust.dev", + )); + assert!(!same_deployment("", "https://api.braintrust.dev")); + assert!(!same_deployment("https://api.braintrust.dev", "")); + } } diff --git a/src/status.rs b/src/status.rs index 93b555e..25c603d 100644 --- a/src/status.rs +++ b/src/status.rs @@ -76,8 +76,18 @@ pub async fn run(base: BaseArgs, _args: StatusArgs) -> Result<()> { }); let profile_info = resolve_profile_info(selected_profile.as_deref(), org.as_deref()); - if selected_profile.is_some() && org.as_deref().map(str::trim).is_none_or(str::is_empty) { - if let Some(profile_org) = profile_info.as_ref().and_then(|p| p.org_name.clone()) { + // A selected profile owns its org: if the profile has an org_name, it wins over a stale + // config-file `org` that may have been authored for a different profile. An explicit + // --org flag already won during `resolve_config` and is preserved here. + if let Some(profile_org) = profile_info.as_ref().and_then(|p| p.org_name.clone()) { + let cli_or_profile_explicit = base + .org_name + .as_deref() + .map(str::trim) + .is_some_and(|value| !value.is_empty()); + if selected_profile.is_some() + && (org.as_deref().map(str::trim).is_none_or(str::is_empty) || !cli_or_profile_explicit) + { org = Some(profile_org); } } diff --git a/src/switch.rs b/src/switch.rs index 2b34046..045d064 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -110,7 +110,25 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { let ctx = login(&login_base).await?; let client = ApiClient::new(&ctx)?; - let org_name = client.org_name().to_string(); + // Source the org we write from the resolved profile so the written config always reflects + // the active profile's org. The API client's org_name can be empty when the login state + // doesn't echo an org name, so fall back to the profile's stored org_name. + let mut org_name = client.org_name().to_string(); + if org_name.trim().is_empty() { + let effective_profile = + config::trimmed_option(profile_name.as_deref().or(base.profile.as_deref())) + .map(str::to_string); + if let Some(profile_name) = effective_profile { + if let Some(info) = auth::list_profiles()? + .into_iter() + .find(|p| p.name == profile_name) + { + if let Some(profile_org) = info.org_name { + org_name = profile_org; + } + } + } + } let project = match resolved_project { Some(p) => validate_or_create_project(&client, &p).await?, diff --git a/src/traces.rs b/src/traces.rs index 835ef40..9b6845c 100644 --- a/src/traces.rs +++ b/src/traces.rs @@ -30,7 +30,7 @@ use reqwest::Url; use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; use unicode_width::UnicodeWidthStr; -use urlencoding::encode; +use urlencoding::{decode, encode}; use crate::args::BaseArgs; use crate::auth::{self, login}; @@ -931,7 +931,12 @@ impl TraceViewerApp { pub async fn run(base: BaseArgs, args: ViewArgs) -> Result<()> { let parsed_startup_url = parse_startup_trace_url_from_view_args(&args)?; - let base = apply_url_hints_to_base(base, parsed_startup_url.as_ref()); + let (base, hints) = + apply_url_hints_with_profile_resolver(base, parsed_startup_url.as_ref(), |org| { + let profiles = auth::list_profiles().ok()?; + auth::resolve_org_to_profile(org, &profiles).ok() + }); + maybe_surface_url_profile_override(&base, &hints); let ctx = login(&base).await?; let client = ApiClient::new(&ctx)?; @@ -5567,23 +5572,26 @@ 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() - }) +/// Result of applying a trace URL's org/profile hints to the base args. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +struct UrlHints { + /// Profile name the URL resolved to (if any), regardless of whether `base.profile` was + /// already set. `None` means the URL did not resolve a profile. + resolved_profile: Option, + /// Org name the URL contributed (if any). + resolved_org: Option, } fn apply_url_hints_with_profile_resolver( mut base: BaseArgs, parsed_url: Option<&ParsedTraceUrl>, resolve_profile_for_org: F, -) -> BaseArgs +) -> (BaseArgs, UrlHints) where F: Fn(&str) -> Option, { let Some(parsed) = parsed_url else { - return base; + return (base, UrlHints::default()); }; let Some(url_org) = parsed .org @@ -5591,7 +5599,7 @@ where .map(str::trim) .filter(|v| !v.is_empty()) else { - return base; + return (base, UrlHints::default()); }; let has_profile_override = base @@ -5605,15 +5613,65 @@ where .map(str::trim) .is_some_and(|v| !v.is_empty()); + let mut hints = UrlHints { + resolved_org: Some(url_org.to_string()), + ..Default::default() + }; + 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.profile = Some(profile_name.clone()); + hints.resolved_profile = Some(profile_name); } } - base + + (base, hints) +} + +/// Whether a URL-driven profile/org overrides the currently active config profile/org. +/// Used to surface a non-surprising info note so users know which credentials a trace URL +/// actually used (e.g. viewing a staging trace while switched to a prod org). +fn url_hints_override_active( + hints: &UrlHints, + active_profile: Option<&str>, + active_org: Option<&str>, +) -> bool { + hints.resolved_profile.is_some() && hints.resolved_profile.as_deref() != active_profile + || hints.resolved_org.is_some() && hints.resolved_org.as_deref() != active_org +} + +/// Surface a non-surprising note when a trace URL causes `bt view` to use a profile/org +/// different from the currently active config. A trace URL is self-describing (it encodes +/// the org/project), so `bt view` intentionally follows the URL; this just makes that +/// override visible instead of silently crossing deployments. +fn maybe_surface_url_profile_override(base: &BaseArgs, hints: &UrlHints) { + if base.json || crate::ui::is_quiet() { + return; + } + if hints.resolved_profile.is_none() && hints.resolved_org.is_none() { + return; + } + let active = crate::config::load().unwrap_or_default(); + let active_profile = crate::config::trimmed_option(active.profile.as_deref()); + let active_org = crate::config::trimmed_option(active.org.as_deref()); + if !url_hints_override_active(hints, active_profile, active_org) { + return; + } + + let resolved_profile = hints.resolved_profile.as_deref().unwrap_or("(none)"); + let resolved_org = hints.resolved_org.as_deref().unwrap_or("(none)"); + let current = match (active_profile, active_org) { + (Some(p), Some(o)) => format!("{p} (org: {o})"), + (Some(p), None) => p.to_string(), + (None, Some(o)) => format!("(profile: none, org: {o})"), + (None, None) => "(none)".to_string(), + }; + eprintln!( + "info: trace URL resolves to profile '{resolved_profile}' (org: {resolved_org}); using it instead of the active profile '{current}'." + ); } fn select_startup_url( @@ -5659,13 +5717,23 @@ fn parse_trace_url(input: &str) -> Result { }; if let Some(segments) = parsed_url.path_segments() { - let parts: Vec<&str> = segments.filter(|part| !part.is_empty()).collect(); + // `Url::path_segments` returns percent-encoded segments (e.g. a space becomes + // `%20`), so decode the org/project/page path segments before using them as + // identifiers. Query-pair values are already decoded by `query_pairs`. + let parts: Vec = segments + .filter(|part| !part.is_empty()) + .map(|part| { + decode(part) + .map(|cow| cow.into_owned()) + .unwrap_or_else(|_| part.to_string()) + }) + .collect(); if parts.len() >= 2 && parts[0] == "app" { - parsed.org = Some(parts[1].to_string()); + parsed.org = Some(parts[1].clone()); if parts.len() >= 4 && parts[2] == "p" { - parsed.project = Some(parts[3].to_string()); + parsed.project = Some(parts[3].clone()); if parts.len() >= 5 { - parsed.page = Some(parts[4].to_string()); + parsed.page = Some(parts[4].clone()); } } } @@ -6799,12 +6867,14 @@ mod tests { let base = base_args(); let parsed = parsed_url_with_org("Lovable"); - let updated = apply_url_hints_with_profile_resolver(base, Some(&parsed), |org| { + let (updated, hints) = apply_url_hints_with_profile_resolver(base, Some(&parsed), |org| { (org == "Lovable").then(|| "lovable-profile".to_string()) }); assert_eq!(updated.org_name.as_deref(), Some("Lovable")); assert_eq!(updated.profile.as_deref(), Some("lovable-profile")); + assert_eq!(hints.resolved_org.as_deref(), Some("Lovable")); + assert_eq!(hints.resolved_profile.as_deref(), Some("lovable-profile")); } #[test] @@ -6813,12 +6883,15 @@ 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), |_| { + let (updated, hints) = apply_url_hints_with_profile_resolver(base, Some(&parsed), |_| { Some("other".to_string()) }); assert_eq!(updated.profile.as_deref(), Some("explicit-profile")); assert!(updated.org_name.is_none()); + // URL org is still reported even when it didn't override an explicit profile. + assert_eq!(hints.resolved_org.as_deref(), Some("Lovable")); + assert!(hints.resolved_profile.is_none()); } #[test] @@ -6916,4 +6989,90 @@ mod tests { assert_eq!(parsed.span_id.as_deref(), Some("span")); assert_eq!(parsed.trace_view_type.as_deref(), Some("timeline")); } + + #[test] + fn parse_trace_url_decodes_percent_encoded_org_and_project() { + // Org/project path segments arrive percent-encoded from `Url::path_segments` (spaces + // become `%20`); they must be decoded so downstream org resolution gets "BT Staging", + // not "BT%20Staging". + let parsed = + parse_trace_url("https://www.braintrust.dev/app/BT%20Staging/p/cedric-traces/logs?r=aa1a6fd0-e1b1-46e4-a176-082350262954") + .expect("parse"); + assert_eq!(parsed.org.as_deref(), Some("BT Staging")); + assert_eq!(parsed.project.as_deref(), Some("cedric-traces")); + assert_eq!( + parsed.row_ref.as_deref(), + Some("aa1a6fd0-e1b1-46e4-a176-082350262954") + ); + } + + #[test] + fn parse_trace_url_decodes_other_percent_encoded_segments() { + let parsed = parse_trace_url( + "https://www.braintrust.dev/app/A%26B%20Corp/p/my%20project/logs?r=root", + ) + .expect("parse"); + assert_eq!(parsed.org.as_deref(), Some("A&B Corp")); + assert_eq!(parsed.project.as_deref(), Some("my project")); + } + + #[test] + fn parse_trace_url_decodes_project_with_no_page() { + let parsed = + parse_trace_url("https://www.braintrust.dev/app/BT%20Staging/p/cedric-traces?r=root") + .expect("parse"); + assert_eq!(parsed.org.as_deref(), Some("BT Staging")); + assert_eq!(parsed.project.as_deref(), Some("cedric-traces")); + assert_eq!(parsed.page.as_deref(), None); + } + + #[test] + fn url_hints_override_active_detects_profile_change() { + let hints = UrlHints { + resolved_profile: Some("BT Staging".to_string()), + resolved_org: Some("BT Staging".to_string()), + }; + // Active config is ced-test-1; URL resolved to BT Staging => override. + assert!(url_hints_override_active( + &hints, + Some("ced-test-1"), + Some("ced-test-1") + )); + } + + #[test] + fn url_hints_override_active_no_change_when_same_profile() { + let hints = UrlHints { + resolved_profile: Some("ced-test-1".to_string()), + resolved_org: Some("ced-test-1".to_string()), + }; + assert!(!url_hints_override_active( + &hints, + Some("ced-test-1"), + Some("ced-test-1") + )); + } + + #[test] + fn url_hints_override_active_no_change_when_no_hints() { + let hints = UrlHints::default(); + assert!(!url_hints_override_active( + &hints, + Some("ced-test-1"), + Some("ced-test-1") + )); + } + + #[test] + fn url_hints_override_active_detects_org_only_change() { + let hints = UrlHints { + resolved_profile: None, + resolved_org: Some("BT Staging".to_string()), + }; + assert!(url_hints_override_active( + &hints, + Some("ced-test-1"), + Some("ced-test-1") + )); + } } diff --git a/src/ui/ratatui_table.rs b/src/ui/ratatui_table.rs index 37ca9e0..9082176 100644 --- a/src/ui/ratatui_table.rs +++ b/src/ui/ratatui_table.rs @@ -271,10 +271,7 @@ fn common_prefix_len(tokenized: &[Vec]) -> usize { return 0; }; let mut len = 0; - 'outer: loop { - let Some(candidate) = first.get(len) else { - break; - }; + 'outer: while let Some(candidate) = first.get(len) { for tokens in tokenized.iter().skip(1) { if tokens.get(len) != Some(candidate) { break 'outer;