From 080838387a76b3ea7fc841c8cabfd1ac00a7e444 Mon Sep 17 00:00:00 2001 From: Neel Date: Thu, 18 Jun 2026 19:27:21 +0100 Subject: [PATCH] agent: Preserve saved model selection until provider loads (#59417) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR fixes a bug where a thread could fall back to a BYOK model when a remote provider hasn't yet returned its list of available models. On a cold start, a thread previously using a Cloud model would fall over to a BYOK model (if a key is defined) because Cloud models aren't loaded yet. Note: there is a small behaviour change here, where previously the priority was `default_model → profile_model → default_model`. The new order is `profile_model → default_model`. Release Notes: - Improved saved thread model selection for cloud-based providers --- crates/agent/src/agent.rs | 214 +++++++++++++++++++++++++++++++++++-- crates/agent/src/thread.rs | 143 +++++++++++++++++-------- 2 files changed, 306 insertions(+), 51 deletions(-) diff --git a/crates/agent/src/agent.rs b/crates/agent/src/agent.rs index eed1f118db..2b0d852bc8 100644 --- a/crates/agent/src/agent.rs +++ b/crates/agent/src/agent.rs @@ -1369,12 +1369,8 @@ impl NativeAgent { for session in self.sessions.values_mut() { session.thread.update(cx, |thread, cx| { - if thread.model().is_none() - && let Some(model) = default_model.clone() - { - thread.set_model(model, cx); - cx.notify(); - } + thread.ensure_model(default_model.as_ref(), cx); + if let Some(model) = summarization_model.clone() { if thread.summarization_model().is_none() || matches!(event, language_model::Event::ThreadSummaryModelChanged) @@ -6023,6 +6019,212 @@ mod internal_tests { drop(reloaded_acp_thread); } + async fn persist_thread_with_fake_corp_model( + cx: &mut TestAppContext, + ) -> ( + Entity, + Rc, + Entity, + acp::SessionId, + Arc, + ) { + let fs = FakeFs::new(cx.executor()); + fs.insert_tree("/", json!({ "a": {} })).await; + let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await; + let thread_store = cx.new(|cx| ThreadStore::new(cx)); + let agent = cx + .update(|cx| NativeAgent::new(thread_store.clone(), Templates::new(), fs.clone(), cx)); + let connection = Rc::new(NativeAgentConnection(agent.clone())); + + let model = Arc::new(FakeLanguageModel::with_id_and_thinking( + "fake-corp", + "custom-model-id", + "Custom Model Display Name", + false, + )); + let provider = Arc::new( + FakeLanguageModelProvider::new( + LanguageModelProviderId::from("fake-corp".to_string()), + LanguageModelProviderName::from("Fake Corp".to_string()), + ) + .with_models(vec![model.clone()]), + ); + cx.update(|cx| { + LanguageModelRegistry::global(cx).update(cx, |registry, cx| { + registry.register_provider(provider.clone(), cx); + }); + }); + agent.update(cx, |agent, cx| agent.models.refresh_list(cx)); + + let acp_thread = cx + .update(|cx| { + connection.clone().new_session( + project.clone(), + PathList::new(&[Path::new("/a")]), + cx, + ) + }) + .await + .unwrap(); + let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone()); + + let selector = connection.model_selector(&session_id).unwrap(); + cx.update(|cx| selector.select_model(AgentModelId::new("fake-corp/custom-model-id"), cx)) + .await + .unwrap(); + + let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx)); + let send = cx.foreground_executor().spawn(send); + cx.run_until_parked(); + model.send_last_completion_stream_text_chunk("Response."); + model.end_last_completion_stream(); + send.await.unwrap(); + cx.run_until_parked(); + + cx.update(|cx| connection.clone().close_session(&session_id, cx)) + .await + .unwrap(); + drop(acp_thread); + + (agent, connection, project, session_id, provider) + } + + fn unregister_fake_corp(cx: &mut TestAppContext) { + cx.update(|cx| { + LanguageModelRegistry::global(cx).update(cx, |registry, cx| { + registry.unregister_provider( + LanguageModelProviderId::from("fake-corp".to_string()), + cx, + ); + }); + }); + } + + #[gpui::test] + async fn test_loaded_thread_resolves_model_when_provider_loads_late(cx: &mut TestAppContext) { + init_test(cx); + let (agent, _connection, project, session_id, provider) = + persist_thread_with_fake_corp_model(cx).await; + + // Simulate a restart where the provider hasn't fetched its model list + // yet, so the saved selection can't be resolved at load time. + unregister_fake_corp(cx); + + let reloaded_acp_thread = agent + .update(cx, |agent, cx| { + agent.open_thread(session_id.clone(), project.clone(), cx) + }) + .await + .unwrap(); + let thread = agent.read_with(cx, |agent, _| { + agent.sessions.get(&session_id).unwrap().thread.clone() + }); + thread.read_with(cx, |thread, _| { + assert!( + thread.model().is_none(), + "should not fall back to an unrelated model" + ); + }); + + // The original selection is persisted even while unresolved, so a save + // during the window can't overwrite the user's choice with a fallback. + let db_thread = thread.read_with(cx, |thread, cx| thread.to_db(cx)).await; + let saved = db_thread.model.expect("selection should be persisted"); + assert_eq!(saved.provider, "fake-corp"); + assert_eq!(saved.model, "custom-model-id"); + + cx.update(|cx| { + LanguageModelRegistry::global(cx).update(cx, |registry, cx| { + registry.register_provider(provider.clone(), cx); + }); + }); + cx.run_until_parked(); + + thread.read_with(cx, |thread, _| { + assert_eq!( + thread + .model() + .expect("model should resolve once provider loads") + .id() + .0 + .as_ref(), + "custom-model-id" + ); + }); + + drop(reloaded_acp_thread); + } + + #[gpui::test] + async fn test_explicit_model_selection_cancels_pending(cx: &mut TestAppContext) { + init_test(cx); + let (agent, connection, project, session_id, provider) = + persist_thread_with_fake_corp_model(cx).await; + + unregister_fake_corp(cx); + + let reloaded_acp_thread = agent + .update(cx, |agent, cx| { + agent.open_thread(session_id.clone(), project.clone(), cx) + }) + .await + .unwrap(); + let thread = agent.read_with(cx, |agent, _| { + agent.sessions.get(&session_id).unwrap().thread.clone() + }); + thread.read_with(cx, |thread, _| { + assert!(thread.model().is_none()); + }); + + // The user explicitly picks a different, available model. + let other_model = Arc::new(FakeLanguageModel::with_id_and_thinking( + "other-corp", + "other-model-id", + "Other Model", + false, + )); + let other_provider = Arc::new( + FakeLanguageModelProvider::new( + LanguageModelProviderId::from("other-corp".to_string()), + LanguageModelProviderName::from("Other Corp".to_string()), + ) + .with_models(vec![other_model.clone()]), + ); + cx.update(|cx| { + LanguageModelRegistry::global(cx).update(cx, |registry, cx| { + registry.register_provider(other_provider, cx); + }); + }); + cx.run_until_parked(); + + let selector = connection.model_selector(&session_id).unwrap(); + cx.update(|cx| selector.select_model(AgentModelId::new("other-corp/other-model-id"), cx)) + .await + .unwrap(); + + thread.read_with(cx, |thread, _| { + assert_eq!(thread.model().unwrap().id().0.as_ref(), "other-model-id"); + }); + + // The original provider returning must not clobber the explicit choice. + cx.update(|cx| { + LanguageModelRegistry::global(cx).update(cx, |registry, cx| { + registry.register_provider(provider.clone(), cx); + }); + }); + cx.run_until_parked(); + + thread.read_with(cx, |thread, _| { + assert_eq!( + thread.model().unwrap().id().0.as_ref(), + "other-model-id", + "a late provider load must not override the explicit selection" + ); + }); + + drop(reloaded_acp_thread); + } + #[gpui::test] async fn test_save_load_thread(cx: &mut TestAppContext) { init_test(cx); diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 0bae32da67..22326ee8d4 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1136,6 +1136,37 @@ enum CompletionError { Other(#[from] anyhow::Error), } +pub(crate) enum ThreadModel { + Ready(Arc), + Unresolved(SelectedModel), + Unset, +} + +impl ThreadModel { + fn as_model(&self) -> Option<&Arc> { + match self { + Self::Ready(model) => Some(model), + Self::Unresolved(_) | Self::Unset => None, + } + } +} + +impl From<&ThreadModel> for Option { + fn from(model: &ThreadModel) -> Self { + match model { + ThreadModel::Ready(model) => Some(DbLanguageModel { + provider: model.provider_id().to_string(), + model: model.id().0.to_string(), + }), + ThreadModel::Unresolved(selection) => Some(DbLanguageModel { + provider: selection.provider.0.to_string(), + model: selection.model.0.to_string(), + }), + ThreadModel::Unset => None, + } + } +} + pub struct Thread { id: acp::SessionId, prompt_id: PromptId, @@ -1169,7 +1200,7 @@ pub struct Thread { profile_id: AgentProfileId, project_context: Entity, pub(crate) templates: Arc, - model: Option>, + model: ThreadModel, summarization_model: Option>, thinking_enabled: bool, thinking_effort: Option, @@ -1278,6 +1309,7 @@ impl Thread { .and_then(|model| model.speed); let (prompt_capabilities_tx, prompt_capabilities_rx) = watch::channel(Self::prompt_capabilities(model.as_deref())); + let model = model.map_or(ThreadModel::Unset, ThreadModel::Ready); Self { id: acp::SessionId::new(uuid::Uuid::new_v4().to_string()), prompt_id: PromptId::new(), @@ -1354,13 +1386,13 @@ impl Thread { return; }; - self.model = Some(model.clone()); self.thinking_enabled = selection.enable_thinking && model.supports_thinking(); self.thinking_effort = selection.effort.clone(); self.speed = selection.speed.filter(|_| model.supports_fast_mode()); self.prompt_capabilities_tx - .send(Self::prompt_capabilities(self.model.as_deref())) + .send(Self::prompt_capabilities(Some(model.as_ref()))) .log_err(); + self.model = ThreadModel::Ready(model); } pub fn id(&self) -> &acp::SessionId { @@ -1630,31 +1662,33 @@ impl Thread { .profile .unwrap_or_else(|| settings.default_profile.clone()); - let mut model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| { - db_thread - .model - .and_then(|model| { - let model = SelectedModel { - provider: model.provider.clone().into(), - model: model.model.into(), - }; - registry.select_model(&model, cx) - }) - .or_else(|| registry.default_model()) - .map(|model| model.model) + let saved_selection = db_thread.model.map(|model| SelectedModel { + provider: model.provider.into(), + model: model.model.into(), }); - if model.is_none() { - model = Self::resolve_profile_model(&profile_id, cx); - } - if model.is_none() { - model = LanguageModelRegistry::global(cx).update(cx, |registry, _cx| { - registry.default_model().map(|model| model.model) - }); - } + let resolved_saved_model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| { + saved_selection + .as_ref() + .and_then(|selection| registry.select_model(selection, cx)) + .map(|configured| configured.model) + }); - let (prompt_capabilities_tx, prompt_capabilities_rx) = - watch::channel(Self::prompt_capabilities(model.as_deref())); + let model = match (resolved_saved_model, saved_selection) { + (Some(model), _) => ThreadModel::Ready(model), + (None, Some(selection)) => ThreadModel::Unresolved(selection), + (None, None) => Self::resolve_profile_model(&profile_id, cx) + .or_else(|| { + LanguageModelRegistry::global(cx).update(cx, |registry, _cx| { + registry.default_model().map(|model| model.model) + }) + }) + .map_or(ThreadModel::Unset, ThreadModel::Ready), + }; + + let (prompt_capabilities_tx, prompt_capabilities_rx) = watch::channel( + Self::prompt_capabilities(model.as_model().map(|model| model.as_ref())), + ); let action_log = cx.new(|_| ActionLog::new(project.clone())); @@ -1719,10 +1753,7 @@ impl Thread { initial_project_snapshot: None, cumulative_token_usage: self.cumulative_token_usage, request_token_usage: self.request_token_usage.clone(), - model: self.model.as_ref().map(|model| DbLanguageModel { - provider: model.provider_id().to_string(), - model: model.id().0.to_string(), - }), + model: (&self.model).into(), profile: Some(self.profile_id.clone()), imported: self.imported, subagent_context: self.subagent_context.clone(), @@ -1795,13 +1826,35 @@ impl Thread { } pub fn model(&self) -> Option<&Arc> { - self.model.as_ref() + self.model.as_model() + } + + pub(crate) fn ensure_model( + &mut self, + default_model: Option<&Arc>, + cx: &mut Context, + ) { + let resolved = match &self.model { + ThreadModel::Ready(_) => return, + ThreadModel::Unresolved(selection) => { + LanguageModelRegistry::global(cx).update(cx, |registry, cx| { + registry + .select_model(selection, cx) + .map(|configured| configured.model) + }) + } + ThreadModel::Unset => default_model.cloned(), + }; + + if let Some(model) = resolved { + self.set_model(model, cx); + } } pub fn set_model(&mut self, model: Arc, cx: &mut Context) { let old_usage = self.latest_token_usage(); - self.model = Some(model.clone()); - let new_caps = Self::prompt_capabilities(self.model.as_deref()); + self.model = ThreadModel::Ready(model.clone()); + let new_caps = Self::prompt_capabilities(self.model.as_model().map(|model| model.as_ref())); let new_usage = self.latest_token_usage(); if old_usage != new_usage { cx.emit(TokenUsageUpdated(new_usage)); @@ -2137,7 +2190,7 @@ impl Thread { pub fn latest_token_usage(&self) -> Option { let usage = self.latest_request_token_usage()?; - let model = self.model.clone()?; + let model = self.model()?; let input_tokens = total_input_tokens(usage); Some(acp_thread::TokenUsage { @@ -2256,8 +2309,8 @@ impl Thread { cx: &mut Context, ) -> Result>> { let model = self - .model - .clone() + .model() + .cloned() .ok_or_else(|| anyhow!(NoModelConfiguredError))?; // Flush any pending message and cancel an in-flight turn before we @@ -2536,7 +2589,7 @@ impl Thread { let (model, request) = this.update(cx, |this, cx| { let model = refusal_fallback_model .clone() - .or_else(|| this.model.clone()) + .or_else(|| this.model().cloned()) .ok_or_else(|| anyhow!(NoModelConfiguredError))?; this.refresh_turn_tools(cx); let request = this.build_completion_request(intent, cx)?; @@ -2683,7 +2736,7 @@ impl Thread { if had_refusal { let maybe_fallback = this.update(cx, |this, cx| -> Option> { - let current_model = refusal_fallback_model.as_ref().or(this.model.as_ref())?; + let current_model = refusal_fallback_model.as_ref().or(this.model())?; let fallback_id = match current_model.refusal_fallback_model_id() { Some(id) => id, None => { @@ -2830,7 +2883,7 @@ impl Thread { ) -> Result> { let Some((model, request, insertion_ix)) = this.update(cx, |this, cx| { let insertion_ix = this.compaction_message_target_ix(cx)?; - let model = this.model.clone()?; + let model = this.model().cloned()?; let request = this.build_compaction_request(insertion_ix, &model, cx); this.current_request_token_usage = TokenUsage::default(); // Preserve telemetry across retries so the retry count keeps @@ -2997,7 +3050,7 @@ impl Thread { attempt: u8, plan: Option, ) -> Result { - let Some(model) = self.model.as_ref() else { + let Some(model) = self.model() else { return Err(anyhow!(error)); }; @@ -3103,8 +3156,8 @@ impl Thread { thread_id = self.id.to_string(), parent_thread_id = self.parent_thread_id().map(|id| id.to_string()), prompt_id = self.prompt_id.to_string(), - model = self.model.as_ref().map(|m| m.telemetry_id()), - model_provider = self.model.as_ref().map(|m| m.provider_id().to_string()), + model = self.model().map(|m| m.telemetry_id()), + model_provider = self.model().map(|m| m.provider_id().to_string()), input_tokens = usage.input_tokens, output_tokens = usage.output_tokens, cache_creation_input_tokens = usage.cache_creation_input_tokens, @@ -3720,7 +3773,7 @@ impl Thread { } fn enabled_tools(&self, cx: &App) -> BTreeMap> { - let Some(model) = self.model.as_ref() else { + let Some(model) = self.model() else { return BTreeMap::new(); }; let Some(profile) = AgentSettings::get_global(cx).profiles.get(&self.profile_id) else { @@ -3907,7 +3960,7 @@ impl Thread { let system_prompt = SystemPromptTemplate { project: self.project_context.read(cx), available_tools, - model_name: self.model.as_ref().map(|m| m.name().0.to_string()), + model_name: self.model().map(|m| m.name().0.to_string()), date: Local::now().format("%Y-%m-%d").to_string(), user_agents_md, sandboxing: crate::sandboxing::sandboxing_enabled_for_project( @@ -3972,7 +4025,7 @@ impl Thread { trigger: &'static str, cx: &App, ) -> Option { - let model = self.model.as_ref()?; + let model = self.model()?; let auto_compact = AgentSettings::get_global(cx).auto_compact; let max_tokens = model.max_token_count(); let max_input_tokens = max_tokens.saturating_sub(model.max_output_tokens().unwrap_or(0)); @@ -4014,7 +4067,7 @@ impl Thread { return None; } - let model = self.model.as_ref()?; + let model = self.model()?; let max_token_count = model.max_token_count(); let max_input_tokens = max_token_count.saturating_sub(model.max_output_tokens().unwrap_or(0));