From 8b18be4182821b98f29a626103cc51c5c0c44d31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Fri, 12 Jun 2026 18:23:37 -0700 Subject: [PATCH 1/2] chore: switch profile/org after bt auth login --- README.md | 1 + src/auth.rs | 184 +++++++++++++++++++++++++++++++++++++++++++++++++- src/switch.rs | 56 ++++++++++++--- 3 files changed, 228 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 932c72b..76e60d2 100644 --- a/README.md +++ b/README.md @@ -289,6 +289,7 @@ Local version and pagination-key conversion helpers: - `bt auth login` - First prompt chooses: `OAuth (browser)` (default) or `API key`. - If your API key can access multiple orgs, `bt` uses a searchable picker (alphabetized) and lets you choose a specific org or no default org (cross-org mode). + - After login, `bt` updates the active profile/org context immediately. If `--project` is set, it also switches that project; otherwise it clears any stale default project for the new login. - `bt` confirms the resolved API URL before saving. - Login with OAuth (browser-based, stores refresh token in secure credential store): - `bt auth login --oauth --profile work` diff --git a/src/auth.rs b/src/auth.rs index 4cfe9c4..4d1f430 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -25,8 +25,10 @@ use tokio::net::TcpListener; use crate::{ args::{BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL}, - http::{build_http_client, build_http_client_from_builder}, - ui, + config, + http::{build_http_client, build_http_client_from_builder, ApiClient}, + projects::api, + switch, ui, }; const KEYCHAIN_SERVICE: &str = "com.braintrust.bt.cli"; @@ -420,6 +422,11 @@ struct AuthLogoutArgs { force: bool, } +struct PostLoginContextUpdate { + display: String, + 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, @@ -1231,11 +1238,28 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { base.app_url.clone(), selected_org.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")?; ui::print_command_status( ui::CommandStatus::Success, &format_login_success(&selected_org, &profile_name, &selected_api_url), ); + ui::print_command_status( + ui::CommandStatus::Success, + &format!("Switched to {}", context_update.display), + ); + if base.verbose { + eprintln!("Wrote to {}", context_update.path.display()); + } Ok(()) } @@ -1346,11 +1370,28 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { client_id.clone(), selected_org.as_ref().map(|org| org.name.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 + .context("login succeeded, but failed to update active context")?; ui::print_command_status( ui::CommandStatus::Success, &format_login_success(&selected_org, &profile_name, &selected_api_url), ); + ui::print_command_status( + ui::CommandStatus::Success, + &format!("Switched to {}", context_update.display), + ); + if base.verbose { + eprintln!("Wrote to {}", context_update.path.display()); + } Ok(()) } @@ -1680,6 +1721,94 @@ 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(), + } +} + +fn format_post_login_context( + selected_org: Option<&LoginOrgInfo>, + project: Option<&api::Project>, +) -> String { + match (selected_org, project) { + (Some(org), Some(project)) => format!("{}/{}", org.name, project.name), + (Some(org), None) => org.name.clone(), + (None, _) => "cross-org mode".to_string(), + } +} + +async fn resolve_post_login_project( + base: &BaseArgs, + credential: &str, + api_url: &str, + app_url: &str, + selected_org: Option<&LoginOrgInfo>, +) -> Result> { + let Some(project_name) = config::trimmed_option(base.project.as_deref()) else { + return Ok(None); + }; + + let selected_org = selected_org.ok_or_else(|| { + anyhow::anyhow!( + "cannot set a default project in cross-org mode; rerun `bt auth login --org --project `" + ) + })?; + let ctx = + build_login_context_for_selected_org(credential, api_url, app_url, Some(selected_org)); + let client = ApiClient::new(&ctx)?; + switch::validate_or_create_project(&client, project_name) + .await + .map(Some) +} + +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, + }) +} + async fn run_profiles(base: &BaseArgs, _args: AuthProfilesArgs) -> Result<()> { let store = load_auth_store()?; if store.profiles.is_empty() { @@ -4258,6 +4387,57 @@ mod tests { } } + #[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!(err + .to_string() + .contains("cannot set a default project in cross-org mode")); + } + #[test] fn resolve_requested_org_for_api_key_login_keeps_matching_requested_org() { let orgs = vec![login_org("org_1", "acme")]; diff --git a/src/switch.rs b/src/switch.rs index 2b34046..35495bb 100644 --- a/src/switch.rs +++ b/src/switch.rs @@ -150,7 +150,12 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { let config_profile = config::trimmed_option(profile_name.as_deref().or(base.profile.as_deref())) .map(str::to_string); - apply_switch_config(&mut cfg, config_profile.as_deref(), &org_name, &project); + apply_switch_config( + &mut cfg, + config_profile.as_deref(), + Some(&org_name), + Some(&project), + ); config::save_file(&path, &cfg) .context(format!("Could not save config to {}", path.display()))?; @@ -176,7 +181,7 @@ pub async fn run(base: BaseArgs, args: SwitchArgs) -> Result<()> { Ok(()) } -fn select_scope() -> Result<(std::path::PathBuf, &'static str)> { +pub(crate) fn select_scope() -> Result<(std::path::PathBuf, &'static str)> { let global = config::global_path()?; let local = config::local_path().unwrap(); let options = [ @@ -225,7 +230,10 @@ fn select_scope() -> Result<(std::path::PathBuf, &'static str)> { } } -async fn validate_or_create_project(client: &ApiClient, name: &str) -> Result { +pub(crate) async fn validate_or_create_project( + client: &ApiClient, + name: &str, +) -> Result { let exists = with_spinner("Loading project...", api::get_project_by_name(client, name)).await?; if let Some(project) = exists { @@ -248,18 +256,26 @@ async fn validate_or_create_project(client: &ApiClient, name: &str) -> Result, - org_name: &str, - project: &api::Project, + org_name: Option<&str>, + project: Option<&api::Project>, ) { if let Some(profile_name) = config::trimmed_option(profile_name) { cfg.profile = Some(profile_name.to_string()); } - cfg.org = Some(org_name.to_string()); - cfg.project = Some(project.name.clone()); - cfg.project_id = Some(project.id.clone()); + cfg.org = config::trimmed_option(org_name).map(str::to_string); + match project { + Some(project) => { + cfg.project = Some(project.name.clone()); + cfg.project_id = Some(project.id.clone()); + } + None => { + cfg.project = None; + cfg.project_id = None; + } + } } fn resolve_profile_for_switch( @@ -530,7 +546,7 @@ mod tests { description: None, }; - apply_switch_config(&mut cfg, Some("work"), "acme-org", &project); + apply_switch_config(&mut cfg, Some("work"), Some("acme-org"), Some(&project)); assert_eq!(cfg.profile.as_deref(), Some("work")); assert_eq!(cfg.org.as_deref(), Some("acme-org")); @@ -551,12 +567,30 @@ mod tests { description: None, }; - apply_switch_config(&mut cfg, None, "acme-org", &project); + apply_switch_config(&mut cfg, None, Some("acme-org"), Some(&project)); assert_eq!(cfg.profile.as_deref(), Some("work")); assert_eq!(cfg.project.as_deref(), Some("next-project")); } + #[test] + fn apply_switch_config_clears_project_and_org_when_context_is_org_only() { + let mut cfg = config::Config { + profile: Some("work".to_string()), + org: Some("old-org".to_string()), + project: Some("stale-project".to_string()), + project_id: Some("proj_stale".to_string()), + ..Default::default() + }; + + apply_switch_config(&mut cfg, Some("next"), None, None); + + assert_eq!(cfg.profile.as_deref(), Some("next")); + assert_eq!(cfg.org, None); + assert_eq!(cfg.project, None); + assert_eq!(cfg.project_id, None); + } + #[test] fn resolve_profile_for_switch_skips_org_prompt_when_api_key_infers_profile() { let mut interactive = false; From 660e766c5f2f9a2bfddc173123e95009f6ba0a15 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Halber?= Date: Fri, 12 Jun 2026 18:36:02 -0700 Subject: [PATCH 2/2] add: flag to keep previous profile after bt auth login --- README.md | 1 + src/auth.rs | 104 +++++++++++++++++++++++++++++++++++++++++++++------- 2 files changed, 91 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 76e60d2..6d7bbb3 100644 --- a/README.md +++ b/README.md @@ -290,6 +290,7 @@ Local version and pagination-key conversion helpers: - 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. + - Pass `--no-profile-switch` (or set `BRAINTRUST_NO_PROFILE_SWITCH=1`) to keep the old behavior and avoid changing active profile/org/project context. - `bt` confirms the resolved API URL before saving. - Login with OAuth (browser-based, stores refresh token in secure credential store): - `bt auth login --oauth --profile work` diff --git a/src/auth.rs b/src/auth.rs index 4d1f430..be77f17 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -409,6 +409,15 @@ struct AuthLoginArgs { /// Do not try to open a browser automatically #[arg(long)] no_browser: bool, + + /// Log in without updating active profile/org/project context + #[arg( + long, + env = "BRAINTRUST_NO_PROFILE_SWITCH", + value_parser = clap::builder::BoolishValueParser::new(), + default_value_t = false + )] + no_profile_switch: bool, } #[derive(Debug, Clone, Args)] @@ -1238,8 +1247,9 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { base.app_url.clone(), selected_org.as_ref().map(|org| org.name.clone()), )?; - let context_update = persist_post_login_context( + let context_update = maybe_persist_post_login_context( base, + &args, &profile_name, &api_key, &selected_api_url, @@ -1253,12 +1263,14 @@ async fn run_login_set(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { ui::CommandStatus::Success, &format_login_success(&selected_org, &profile_name, &selected_api_url), ); - ui::print_command_status( - ui::CommandStatus::Success, - &format!("Switched to {}", context_update.display), - ); - if base.verbose { - eprintln!("Wrote to {}", context_update.path.display()); + if let Some(context_update) = context_update { + ui::print_command_status( + ui::CommandStatus::Success, + &format!("Switched to {}", context_update.display), + ); + if base.verbose { + eprintln!("Wrote to {}", context_update.path.display()); + } } Ok(()) } @@ -1370,8 +1382,9 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { client_id.clone(), selected_org.as_ref().map(|org| org.name.clone()), )?; - let context_update = persist_post_login_context( + let context_update = maybe_persist_post_login_context( base, + &args, &profile_name, &oauth_tokens.access_token, &selected_api_url, @@ -1385,12 +1398,14 @@ async fn run_login_oauth(base: &BaseArgs, args: AuthLoginArgs) -> Result<()> { ui::CommandStatus::Success, &format_login_success(&selected_org, &profile_name, &selected_api_url), ); - ui::print_command_status( - ui::CommandStatus::Success, - &format!("Switched to {}", context_update.display), - ); - if base.verbose { - eprintln!("Wrote to {}", context_update.path.display()); + if let Some(context_update) = context_update { + ui::print_command_status( + ui::CommandStatus::Success, + &format!("Switched to {}", context_update.display), + ); + if base.verbose { + eprintln!("Wrote to {}", context_update.path.display()); + } } Ok(()) @@ -1809,6 +1824,31 @@ async fn persist_post_login_context( }) } +async fn maybe_persist_post_login_context( + base: &BaseArgs, + args: &AuthLoginArgs, + profile_name: &str, + credential: &str, + api_url: &str, + app_url: &str, + selected_org: Option<&LoginOrgInfo>, +) -> Result> { + if args.no_profile_switch { + return Ok(None); + } + + persist_post_login_context( + base, + profile_name, + credential, + api_url, + app_url, + selected_org, + ) + .await + .map(Some) +} + async fn run_profiles(base: &BaseArgs, _args: AuthProfilesArgs) -> Result<()> { let store = load_auth_store()?; if store.profiles.is_empty() { @@ -4418,6 +4458,42 @@ mod tests { assert_eq!(cfg.project_id, None); } + #[tokio::test] + async fn maybe_persist_post_login_context_skips_config_write_when_no_profile_switch_is_set() { + let _env = TestEnv::new(None, None).await; + let original = 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() + }; + crate::config::save_global(&original).expect("save initial config"); + + let args = AuthLoginArgs { + oauth: false, + client_id: None, + no_browser: false, + no_profile_switch: true, + }; + + let update = maybe_persist_post_login_context( + &make_base(), + &args, + "work", + "test-api-key", + "https://api.example.test", + "https://www.example.test", + Some(&login_org("org_123", "acme")), + ) + .await + .expect("skip context update"); + let cfg = crate::config::load_global().expect("load global config"); + + assert!(update.is_none()); + assert_eq!(cfg, original); + } + #[tokio::test] async fn resolve_post_login_project_rejects_cross_org_default_project() { let mut base = make_base();