From 48e78f5da6c7b3c78bb9ac70061f2422d941185d Mon Sep 17 00:00:00 2001 From: KelvinJRosado Date: Mon, 25 May 2026 09:06:54 -0400 Subject: [PATCH] fix: Preserve custom models during LLM preference reconciliation When reconciling LLM preferences, fall back to custom LLM info if the model is not found in the built-in models list. This ensures custom models saved on execution profiles are not cleared during preference updates. --- app/src/ai/llms.rs | 11 +++++- app/src/ai/llms_tests.rs | 77 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 88cadfacd9..fa7d9b8b78 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -1180,7 +1180,10 @@ impl LLMPreferences { let effective_base_model_usable = self .models_by_feature .agent_mode - .usable_info_for_id(effective_base_model_id, ctx); + .usable_info_for_id(effective_base_model_id, ctx) + .or_else(|| { + self.custom_llm_info_for_id_if_enabled(effective_base_model_id, ctx) + }); let effective_base_model_unusable = effective_base_model_usable.is_none(); let effective_base_model_is_configurable = effective_base_model_usable .is_some_and(|info| info.context_window.is_configurable); @@ -1199,6 +1202,9 @@ impl LLMPreferences { .models_by_feature .coding .usable_info_for_id(preferred_llm_id, ctx) + .or_else(|| { + self.custom_llm_info_for_id_if_enabled(preferred_llm_id, ctx) + }) .is_none() { profiles.set_coding_model(profile_id, None, ctx); @@ -1208,6 +1214,9 @@ impl LLMPreferences { if self .get_cli_agent_available() .usable_info_for_id(preferred_llm_id, ctx) + .or_else(|| { + self.custom_llm_info_for_id_if_enabled(preferred_llm_id, ctx) + }) .is_none() { profiles.set_cli_agent_model(profile_id, None, ctx); diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index def397929f..386019ed49 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -1,5 +1,21 @@ use super::*; +use warpui::App; + +use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; +use crate::ai::mcp::TemplatableMCPServerManager; +use crate::auth::auth_manager::AuthManager; +use crate::auth::AuthStateProvider; +use crate::cloud_object::model::persistence::CloudModel; +use crate::network::NetworkStatus; +use crate::server::cloud_objects::update_manager::UpdateManager; +use crate::server::server_api::ServerApiProvider; +use crate::server::sync_queue::SyncQueue; +use crate::test_util::settings::initialize_settings_for_tests; +use crate::workspaces::team_tester::TeamTesterStatus; +use crate::workspaces::user_workspaces::UserWorkspaces; +use crate::LaunchMode; + // -- DisableReason::should_clear_preference tests -- #[test] @@ -303,3 +319,64 @@ fn removing_endpoint_purges_all_its_models_from_custom_llms() { assert_eq!(infos.len(), 1); assert_eq!(infos[0].id.as_str(), "uuid-k1"); } + +#[test] +fn reconcile_preserves_custom_models_saved_on_execution_profile() { + App::test((), |mut app| async move { + let _custom_inference_flag = FeatureFlag::CustomInferenceEndpoints.override_enabled(true); + + initialize_settings_for_tests(&mut app); + app.add_singleton_model(|_| ServerApiProvider::new_for_test()); + app.add_singleton_model(|_| AuthStateProvider::new_for_test()); + app.add_singleton_model(AuthManager::new_for_test); + app.add_singleton_model(|_| NetworkStatus::new()); + app.add_singleton_model(UserWorkspaces::default_mock); + app.add_singleton_model(CloudModel::mock); + app.add_singleton_model(TeamTesterStatus::mock); + app.add_singleton_model(SyncQueue::mock); + app.add_singleton_model(UpdateManager::mock); + app.add_singleton_model(|_| TemplatableMCPServerManager::default()); + + let profiles_model = app.add_singleton_model(|ctx| { + AIExecutionProfilesModel::new(&LaunchMode::new_for_unit_test(), ctx) + }); + let llm_preferences = app.add_singleton_model(LLMPreferences::new); + + let custom_model_id = LLMId::from("custom-model-config-key"); + ApiKeyManager::handle(&app).update(&mut app, |api_key_manager, ctx| { + api_key_manager.add_custom_endpoint( + "local".to_string(), + "https://example.com/v1".to_string(), + "test-key".to_string(), + vec![( + "custom-model".to_string(), + Some("Custom Model".to_string()), + Some(custom_model_id.to_string()), + )], + ctx, + ); + }); + + let default_profile_id = + profiles_model.read(&app, |profiles, _| profiles.default_profile_id()); + profiles_model.update(&mut app, |profiles, ctx| { + profiles.set_base_model(default_profile_id, Some(custom_model_id.clone()), ctx); + profiles.set_coding_model(default_profile_id, Some(custom_model_id.clone()), ctx); + profiles.set_cli_agent_model(default_profile_id, Some(custom_model_id.clone()), ctx); + }); + + llm_preferences.update(&mut app, |preferences, ctx| { + preferences.update_feature_model_choices(Ok(ModelsByFeature::default()), ctx); + }); + + profiles_model.read(&app, |profiles, ctx| { + let profile = profiles.default_profile(ctx); + assert_eq!(profile.data().base_model.as_ref(), Some(&custom_model_id)); + assert_eq!(profile.data().coding_model.as_ref(), Some(&custom_model_id)); + assert_eq!( + profile.data().cli_agent_model.as_ref(), + Some(&custom_model_id) + ); + }); + }); +}