Decouple agent reregistration from settings changes (#39528)

Richard Feldman created 1 month ago

Fixes a `--release`-only bug in feature-flagged agents where the feature
flag isn't picked up in some situations (unless there was a settings
change to go with it - due to an early return when settings didn't
change).

Release Notes:

- N/A

Change summary

crates/agent_ui/src/acp/thread_view.rs   | 40 ++++++++++++++++++++++++-
crates/project/src/agent_server_store.rs | 29 +++++++++++++++---
crates/project/src/project.rs            |  3 +
3 files changed, 63 insertions(+), 9 deletions(-)

Detailed changes

crates/agent_ui/src/acp/thread_view.rs

@@ -289,8 +289,9 @@ pub struct AcpThreadView {
     available_commands: Rc<RefCell<Vec<acp::AvailableCommand>>>,
     is_loading_contents: bool,
     new_server_version_available: Option<SharedString>,
+    resume_thread_metadata: Option<DbThreadMetadata>,
     _cancel_task: Option<Task<()>>,
-    _subscriptions: [Subscription; 4],
+    _subscriptions: [Subscription; 5],
 }
 
 enum ThreadState {
@@ -380,11 +381,17 @@ impl AcpThreadView {
             )
         });
 
+        let agent_server_store = project.read(cx).agent_server_store().clone();
         let subscriptions = [
             cx.observe_global_in::<SettingsStore>(window, Self::agent_ui_font_size_changed),
             cx.observe_global_in::<AgentFontSize>(window, Self::agent_ui_font_size_changed),
             cx.subscribe_in(&message_editor, window, Self::handle_message_editor_event),
             cx.subscribe_in(&entry_view_state, window, Self::handle_entry_view_event),
+            cx.subscribe_in(
+                &agent_server_store,
+                window,
+                Self::handle_agent_servers_updated,
+            ),
         ];
 
         Self {
@@ -392,7 +399,14 @@ impl AcpThreadView {
             workspace: workspace.clone(),
             project: project.clone(),
             entry_view_state,
-            thread_state: Self::initial_state(agent, resume_thread, workspace, project, window, cx),
+            thread_state: Self::initial_state(
+                agent.clone(),
+                resume_thread.clone(),
+                workspace.clone(),
+                project.clone(),
+                window,
+                cx,
+            ),
             login: None,
             message_editor,
             model_selector: None,
@@ -421,13 +435,14 @@ impl AcpThreadView {
             _cancel_task: None,
             focus_handle: cx.focus_handle(),
             new_server_version_available: None,
+            resume_thread_metadata: resume_thread,
         }
     }
 
     fn reset(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         self.thread_state = Self::initial_state(
             self.agent.clone(),
-            None,
+            self.resume_thread_metadata.clone(),
             self.workspace.clone(),
             self.project.clone(),
             window,
@@ -775,6 +790,25 @@ impl AcpThreadView {
         cx.notify();
     }
 
+    fn handle_agent_servers_updated(
+        &mut self,
+        _agent_server_store: &Entity<project::AgentServerStore>,
+        _event: &project::AgentServersUpdated,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        // If we're in a LoadError state OR have a thread_error set (which can happen
+        // when agent.connect() fails during loading), retry loading the thread.
+        // This handles the case where a thread is restored before authentication completes.
+        let should_retry =
+            matches!(&self.thread_state, ThreadState::LoadError(_)) || self.thread_error.is_some();
+
+        if should_retry {
+            self.thread_error = None;
+            self.reset(window, cx);
+        }
+    }
+
     pub fn workspace(&self) -> &WeakEntity<Workspace> {
         &self.workspace
     }

crates/project/src/agent_server_store.rs

@@ -146,10 +146,6 @@ impl AgentServerStore {
 
     fn agent_servers_settings_changed(&mut self, cx: &mut Context<Self>) {
         let AgentServerStoreState::Local {
-            node_runtime,
-            fs,
-            project_environment,
-            downstream_client,
             settings: old_settings,
             ..
         } = &mut self.state
@@ -168,6 +164,29 @@ impl AgentServerStore {
             return;
         }
 
+        self.reregister_agents(cx);
+    }
+
+    fn reregister_agents(&mut self, cx: &mut Context<Self>) {
+        let AgentServerStoreState::Local {
+            node_runtime,
+            fs,
+            project_environment,
+            downstream_client,
+            settings: old_settings,
+            ..
+        } = &mut self.state
+        else {
+            debug_panic!("Non-local projects should never attempt to reregister. This is a bug!");
+
+            return;
+        };
+
+        let new_settings = cx
+            .global::<SettingsStore>()
+            .get::<AllAgentServersSettings>(None)
+            .clone();
+
         self.external_agents.clear();
         self.external_agents.insert(
             GEMINI_NAME.into(),
@@ -256,7 +275,7 @@ impl AgentServerStore {
         let feature_flags_subscription =
             cx.observe_flag::<feature_flags::CodexAcpFeatureFlag, _>(move |_enabled, cx| {
                 let _ = this_handle.update(cx, |this, cx| {
-                    this.agent_servers_settings_changed(cx);
+                    this.reregister_agents(cx);
                 });
             });
         let mut this = Self {

crates/project/src/project.rs

@@ -36,10 +36,11 @@ use dap::inline_value::{InlineValueLocation, VariableLookupKind, VariableScope};
 use task::Shell;
 
 use crate::{
-    agent_server_store::{AgentServerStore, AllAgentServersSettings},
+    agent_server_store::AllAgentServersSettings,
     git_store::GitStore,
     lsp_store::{SymbolLocation, log_store::LogKind},
 };
+pub use agent_server_store::{AgentServerStore, AgentServersUpdated};
 pub use git_store::{
     ConflictRegion, ConflictSet, ConflictSetSnapshot, ConflictSetUpdate,
     git_traversal::{ChildEntriesGitIter, GitEntry, GitEntryRef, GitTraversal},