diff --git a/crates/goat-code/src/auth.rs b/crates/goat-code/src/auth.rs index d07ed8d..1bfae8b 100644 --- a/crates/goat-code/src/auth.rs +++ b/crates/goat-code/src/auth.rs @@ -6,7 +6,7 @@ use goat_auth::{ Credential, CredentialKey, CredentialKind, CredentialService, CredentialStore, SecretString, }; use goat_provider::{AuthMethod, LoginEndpointMetadata, ProviderId, ProviderMetadata}; -use goat_providers::{DEFAULT_ACCOUNT, Registry}; +use goat_providers::Registry; use tokio::sync::mpsc; use crate::{ @@ -31,8 +31,26 @@ pub async fn run_provider(command: ProviderCommand) -> color_eyre::Result<()> { Some(provider) => provider, None => pick_login_provider(&store)?, }; - let account = account.as_deref().unwrap_or(DEFAULT_ACCOUNT); - login(&store, &provider, account, key, endpoint).await + let existing = provider_accounts(&store.entries(), &provider) + .into_iter() + .map(|(account, _)| account) + .collect::>(); + let resolution = cli_ui::resolve_account( + "provider", + &provider, + account.as_deref(), + &existing, + goat_providers::DEFAULT_ACCOUNT, + )?; + login( + &store, + &provider, + &resolution.name, + key, + endpoint, + resolution.replacing, + ) + .await } ProviderCommand::List => { list_providers(&store); @@ -94,6 +112,7 @@ async fn login( account: &str, key: Option, endpoint: Option, + replacing: bool, ) -> color_eyre::Result<()> { let registry = Registry::new(store); let provider_handle = registry @@ -169,7 +188,8 @@ async fn login( } } - cli_ui::success(&format!("stored credential for {provider} ({account})")); + let verb = if replacing { "updated" } else { "stored" }; + cli_ui::success(&format!("{verb} credential for {provider} ({account})")); verify(store, provider, account).await; Ok(()) } diff --git a/crates/goat-code/src/cli_ui.rs b/crates/goat-code/src/cli_ui.rs index 5c7b969..9bf32bf 100644 --- a/crates/goat-code/src/cli_ui.rs +++ b/crates/goat-code/src/cli_ui.rs @@ -1,7 +1,7 @@ use std::io::IsTerminal; use color_eyre::eyre::{Report, Result, eyre}; -use dialoguer::{Input, Password, Select}; +use dialoguer::{Confirm, Input, Password, Select}; use goat_provider::AuthMethod; use crate::{ @@ -21,6 +21,130 @@ pub enum AuthPick { ApiKey, } +#[derive(Debug)] +pub struct AccountResolution { + pub name: String, + pub replacing: bool, +} + +pub fn resolve_account( + service: &str, + provider: &str, + requested: Option<&str>, + existing: &[String], + default_name: &str, +) -> Result { + let taken = |name: &str| existing.iter().any(|account| account == name); + let interactive = std::io::stdin().is_terminal() && std::io::stderr().is_terminal(); + + if let Some(requested) = requested { + let name = requested.trim(); + if name.is_empty() { + return Err(report("account name must not be empty")); + } + return Ok(AccountResolution { + replacing: taken(name), + name: name.to_owned(), + }); + } + + if existing.is_empty() { + if !interactive { + return Ok(AccountResolution { + name: default_name.to_owned(), + replacing: false, + }); + } + let name = prompt_account_name(provider, Some(default_name))?; + return Ok(AccountResolution { + replacing: taken(&name), + name, + }); + } + + if !interactive { + return Err(report_hint( + format!( + "{service} {provider} already has an account: {}", + existing.join(", ") + ), + format!("pass `--account ` to choose or add an account for {provider}"), + )); + } + + match pick_account_slot(provider, existing)? { + AccountSlot::Existing(name) => Ok(AccountResolution { + replacing: true, + name, + }), + AccountSlot::New => { + let name = prompt_account_name(provider, None)?; + if taken(&name) && !confirm_replace(provider, &name)? { + return Err(report("login cancelled")); + } + Ok(AccountResolution { + replacing: taken(&name), + name, + }) + } + } +} + +enum AccountSlot { + Existing(String), + New, +} + +fn pick_account_slot(provider: &str, existing: &[String]) -> Result { + terminal_required()?; + let color = ColorMode::detect_stderr(); + let mut sorted = existing.to_vec(); + sorted.sort(); + let mut labels: Vec = sorted + .iter() + .map(|account| provider_table::option_label(color, account, Palette::Value, "update")) + .collect(); + labels.push(provider_table::option_label( + color, + "+ new account", + Palette::Muted, + "add", + )); + let index = Select::with_theme(goat_theme()) + .with_prompt(select_prompt(color, "account", Some(provider))) + .items(&labels) + .default(0) + .report(false) + .interact_opt() + .map_err(dialoguer_error)? + .ok_or_else(|| report("login cancelled"))?; + Ok(sorted + .get(index) + .cloned() + .map_or(AccountSlot::New, AccountSlot::Existing)) +} + +fn prompt_account_name(provider: &str, default: Option<&str>) -> Result { + let color = ColorMode::detect_stderr(); + let name = prompt_text(&input_prompt(color, "account", Some(provider)), default)?; + if name.is_empty() { + return Err(report("account name must not be empty")); + } + Ok(name) +} + +fn confirm_replace(provider: &str, account: &str) -> Result { + terminal_required()?; + Confirm::with_theme(goat_theme()) + .with_prompt(format!( + "account `{account}` already exists for {provider}; update it?" + )) + .default(false) + .report(false) + .interact() + .map_err(dialoguer_error) +} + pub fn terminal_required() -> Result<()> { if !std::io::stdin().is_terminal() || !std::io::stderr().is_terminal() { return fail("interactive commands require a terminal"); @@ -215,9 +339,58 @@ pub fn format_failure(message: &str, hint: Option) -> String { mod tests { use super::{ format_failure, input_prompt, oauth_status, parse_browser_url, parse_device_code_message, + resolve_account, }; use crate::style::ColorMode; + #[test] + fn resolve_account_uses_requested_name() { + let resolved = + resolve_account("provider", "anthropic", Some("work"), &[], "default").unwrap(); + assert_eq!(resolved.name, "work"); + assert!(!resolved.replacing); + } + + #[test] + fn resolve_account_flags_requested_name_as_replacing() { + let existing = vec!["work".to_owned()]; + let resolved = + resolve_account("provider", "anthropic", Some("work"), &existing, "default").unwrap(); + assert_eq!(resolved.name, "work"); + assert!(resolved.replacing); + } + + #[test] + fn resolve_account_rejects_blank_requested_name() { + let error = resolve_account("provider", "anthropic", Some(" "), &[], "default") + .unwrap_err() + .to_string(); + assert!(error.contains("must not be empty")); + } + + #[test] + fn resolve_account_defaults_when_no_accounts_noninteractive() { + let resolved = resolve_account("provider", "anthropic", None, &[], "default").unwrap(); + assert_eq!(resolved.name, "default"); + assert!(!resolved.replacing); + } + + #[test] + fn resolve_account_honors_provider_default_name() { + let resolved = resolve_account("search provider", "duckduckgo", None, &[], "html").unwrap(); + assert_eq!(resolved.name, "html"); + } + + #[test] + fn resolve_account_errors_noninteractive_when_accounts_exist() { + let existing = vec!["default".to_owned()]; + let error = resolve_account("provider", "anthropic", None, &existing, "default") + .unwrap_err() + .to_string(); + assert!(error.contains("already has an account")); + assert!(error.contains("--account")); + } + #[test] fn input_prompt_names_provider() { let prompt = input_prompt(ColorMode::Plain, "api key", Some("openrouter")); diff --git a/crates/goat-code/src/search.rs b/crates/goat-code/src/search.rs index 9b6c51e..cee5777 100644 --- a/crates/goat-code/src/search.rs +++ b/crates/goat-code/src/search.rs @@ -57,7 +57,15 @@ fn login( "run `goat search list` to see available providers", ) })?; - let account = account.unwrap_or(metadata.default_account); + let existing = existing_search_accounts(provider); + let resolution = cli_ui::resolve_account( + "search provider", + provider, + account, + &existing, + metadata.default_account, + )?; + let account = resolution.name.as_str(); let entry = build_search_account_config(provider, account, endpoint, engine).map_err(cli_ui::report)?; store_search_key(provider, account, key, metadata.credential)?; @@ -72,14 +80,35 @@ fn login( config.search.default_target = Some(target.clone()); } config.save()?; + let verb = if resolution.replacing { + "updated" + } else { + "connected" + }; if make_default { - cli_ui::success(&format!("connected {target} (default)")); + cli_ui::success(&format!("{verb} {target} (default)")); } else { - cli_ui::success(&format!("connected {target}")); + cli_ui::success(&format!("{verb} {target}")); } Ok(()) } +fn existing_search_accounts(provider: &str) -> Vec { + let mut accounts = Config::load() + .search + .accounts + .iter() + .filter(|account| configured_search_provider(account) == provider) + .map(|account| configured_search_target(account).account.to_owned()) + .collect::>(); + for (key, _) in search_credentials() { + if key.provider == provider && !accounts.contains(&key.account) { + accounts.push(key.account); + } + } + accounts +} + fn store_search_key( provider: &str, account: &str,