From c37b64124c3f2342957c02255303124bb2cc6c35 Mon Sep 17 00:00:00 2001 From: Zaki Manian Date: Wed, 11 Mar 2026 18:14:25 -0700 Subject: [PATCH] fix(setup): preserve model selection on provider re-run (#679) (#987) --- src/setup/README.md | 2 ++ src/setup/wizard.rs | 78 ++++++++++++++++++++++++++++++--------------- 2 files changed, 54 insertions(+), 26 deletions(-) diff --git a/src/setup/README.md b/src/setup/README.md index 4b734e36..a1a1d3aa 100644 --- a/src/setup/README.md +++ b/src/setup/README.md @@ -227,6 +227,8 @@ with its own secret name and env var. It is **not** stored as `openai_compatible 2. Otherwise prompt for key entry via `secret_input()` 3. Store encrypted in secrets via `init_secrets_context()` 4. **Cache key in `self.llm_api_key`** for model fetching in Step 4 +5. Preserve `selected_model` on a same-backend re-run; clear it only when + switching to a different backend **NEAR AI** (`setup_nearai`): - Calls `session_manager.ensure_authenticated()` which shows the auth menu: diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 89d1e5be..f8c695f1 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1082,7 +1082,7 @@ impl SetupWizard { "Provider '{}' has no setup wizard. Configure via environment variables.", provider_id )); - self.settings.llm_backend = Some(provider_id.to_string()); + self.set_llm_backend_preserving_model(provider_id); return Ok(()); }; @@ -1137,9 +1137,19 @@ impl SetupWizard { Ok(()) } + /// Update the selected LLM backend while preserving the current model when + /// the backend did not actually change. + fn set_llm_backend_preserving_model(&mut self, backend: &str) { + let backend_changed = self.settings.llm_backend.as_deref() != Some(backend); + self.settings.llm_backend = Some(backend.to_string()); + if backend_changed { + self.settings.selected_model = None; + } + } + /// NEAR AI provider setup (extracted from the old step_authentication). async fn setup_nearai(&mut self) -> Result<(), SetupError> { - self.settings.llm_backend = Some("nearai".to_string()); + self.set_llm_backend_preserving_model("nearai"); // Check if we already have a session if let Some(ref session) = self.session_manager @@ -1224,11 +1234,7 @@ impl SetupWizard { /// Anthropic OAuth setup: extract token from `claude login` credentials. async fn setup_anthropic_oauth(&mut self) -> Result<(), SetupError> { - // Clear model only when switching providers (old model may be invalid) - if self.settings.llm_backend.as_deref() != Some("anthropic") { - self.settings.selected_model = None; - } - self.settings.llm_backend = Some("anthropic".to_string()); + self.set_llm_backend_preserving_model("anthropic"); // Try to extract existing OAuth token from Claude Code credentials if let Some(token) = crate::config::ClaudeCodeConfig::extract_oauth_token() { @@ -1322,11 +1328,7 @@ impl SetupWizard { other => other, }); - // Clear model only when switching providers (old model may be invalid) - if self.settings.llm_backend.as_deref() != Some(backend) { - self.settings.selected_model = None; - } - self.settings.llm_backend = Some(backend.to_string()); + self.set_llm_backend_preserving_model(backend); // Check env var first if let Ok(existing) = std::env::var(env_var) { @@ -1385,11 +1387,7 @@ impl SetupWizard { &mut self, def: &crate::llm::ProviderDefinition, ) -> Result<(), SetupError> { - // Clear model only when switching providers (old model may be invalid) - if self.settings.llm_backend.as_deref() != Some(&def.id) { - self.settings.selected_model = None; - } - self.settings.llm_backend = Some(def.id.clone()); + self.set_llm_backend_preserving_model(&def.id); let default_url = self .settings @@ -1419,10 +1417,7 @@ impl SetupWizard { /// AWS Bedrock provider setup: region, auth, and cross-region config. async fn setup_bedrock(&mut self) -> Result<(), SetupError> { - if self.settings.llm_backend.as_deref() != Some("bedrock") { - self.settings.selected_model = None; - } - self.settings.llm_backend = Some("bedrock".to_string()); + self.set_llm_backend_preserving_model("bedrock"); // Region let default_region = self @@ -1513,11 +1508,7 @@ impl SetupWizard { secret_name: &str, display_name: &str, ) -> Result<(), SetupError> { - // Clear model only when switching providers (old model may be invalid) - if self.settings.llm_backend.as_deref() != Some(backend_id) { - self.settings.selected_model = None; - } - self.settings.llm_backend = Some(backend_id.to_string()); + self.set_llm_backend_preserving_model(backend_id); let existing_url = self .settings @@ -3872,6 +3863,41 @@ mod tests { } } + #[test] + fn test_set_llm_backend_preserves_model_when_backend_unchanged() { + let mut wizard = SetupWizard::new(); + wizard.settings.llm_backend = Some("openai".to_string()); + wizard.settings.selected_model = Some("gpt-4o".to_string()); + + wizard.set_llm_backend_preserving_model("openai"); + + assert_eq!(wizard.settings.llm_backend.as_deref(), Some("openai")); + assert_eq!(wizard.settings.selected_model.as_deref(), Some("gpt-4o")); + } + + #[test] + fn test_set_llm_backend_clears_model_when_backend_was_unset() { + let mut wizard = SetupWizard::new(); + wizard.settings.selected_model = Some("gpt-4o".to_string()); + + wizard.set_llm_backend_preserving_model("openai"); + + assert_eq!(wizard.settings.llm_backend.as_deref(), Some("openai")); + assert_eq!(wizard.settings.selected_model, None); + } + + #[test] + fn test_set_llm_backend_clears_model_when_backend_changes() { + let mut wizard = SetupWizard::new(); + wizard.settings.llm_backend = Some("openai".to_string()); + wizard.settings.selected_model = Some("gpt-4o".to_string()); + + wizard.set_llm_backend_preserving_model("anthropic"); + + assert_eq!(wizard.settings.llm_backend.as_deref(), Some("anthropic")); + assert_eq!(wizard.settings.selected_model, None); + } + /// Regression test for #600: re-running provider setup for the same backend /// must NOT clear selected_model. Only switching to a different backend should. #[test]