agent: Fix race-condition for LSP tool registration (#56044)

Bennet Bo Fenner created

This fixes a race condition where the thread would not get the LSP tools
at startup if the feature flag was not resolved yet.
We now always add the tools, but filter them out when we start a new
turn if the feature flag is not set.

Release Notes:

- N/A

Change summary

crates/agent/src/tests/mod.rs | 99 +++++++++++++++++++++++++++++++++++++
crates/agent/src/thread.rs    | 42 ++++++++------
2 files changed, 123 insertions(+), 18 deletions(-)

Detailed changes

crates/agent/src/tests/mod.rs 🔗

@@ -5545,6 +5545,105 @@ async fn test_max_subagent_depth_prevents_tool_registration(cx: &mut TestAppCont
     });
 }
 
+#[gpui::test]
+async fn test_lsp_tools_gated_by_feature_flag(cx: &mut TestAppContext) {
+    init_test(cx);
+
+    let fs = FakeFs::new(cx.executor());
+    fs.insert_tree(path!("/test"), json!({})).await;
+    let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
+    let project_context = cx.new(|_cx| ProjectContext::default());
+    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());
+    let context_server_registry =
+        cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
+    let model = Arc::new(FakeLanguageModel::default());
+    let environment = Rc::new(cx.update(|cx| {
+        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
+    }));
+
+    let thread = cx.new(|cx| {
+        let mut thread = Thread::new(
+            project,
+            project_context,
+            context_server_registry,
+            Templates::new(),
+            Some(model.clone() as Arc<dyn LanguageModel>),
+            cx,
+        );
+        thread.add_default_tools(environment, cx);
+        thread
+    });
+
+    let lsp_tool_names = [
+        FindReferencesTool::NAME,
+        GetCodeActionsTool::NAME,
+        ApplyCodeActionTool::NAME,
+        GoToDefinitionTool::NAME,
+        RenameTool::NAME,
+    ];
+
+    // All LSP tools should be registered on the thread regardless of the flag,
+    // since the feature flag now only controls exposure to the model rather
+    // than registration.
+    thread.read_with(cx, |thread, _| {
+        for name in &lsp_tool_names {
+            assert!(
+                thread.has_registered_tool(name),
+                "expected LSP tool {name} to be registered"
+            );
+        }
+    });
+
+    // Without the `lsp-tool` flag, sending a message should produce a
+    // completion request whose tool list excludes the LSP tools.
+    thread
+        .update(cx, |thread, cx| {
+            thread.send(UserMessageId::new(), ["hello"], cx)
+        })
+        .unwrap();
+    cx.run_until_parked();
+
+    let completion = model.pending_completions().pop().unwrap();
+    let tool_names = tool_names_for_completion(&completion);
+    for name in &lsp_tool_names {
+        assert!(
+            !tool_names.iter().any(|t| t == name),
+            "expected LSP tool {name} to be hidden without the lsp-tool flag, \
+             but completion tools were: {tool_names:?}"
+        );
+    }
+    // Sanity check: a non-LSP default tool should still be exposed.
+    assert!(
+        tool_names.iter().any(|t| t == ReadFileTool::NAME),
+        "expected non-LSP tools to still be exposed, got: {tool_names:?}"
+    );
+    model.end_last_completion_stream();
+    cx.run_until_parked();
+
+    // Enable the `lsp-tool` flag and send another message; the LSP tools
+    // should now appear in the completion request.
+    cx.update(|cx| {
+        cx.update_flags(false, vec!["lsp-tool".to_string()]);
+    });
+
+    thread
+        .update(cx, |thread, cx| {
+            thread.send(UserMessageId::new(), ["hello again"], cx)
+        })
+        .unwrap();
+    cx.run_until_parked();
+
+    let completion = model.pending_completions().pop().unwrap();
+    let tool_names = tool_names_for_completion(&completion);
+    for name in &lsp_tool_names {
+        assert!(
+            tool_names.iter().any(|t| t == name),
+            "expected LSP tool {name} to be exposed when lsp-tool flag is on, \
+             but completion tools were: {tool_names:?}"
+        );
+    }
+}
+
 #[gpui::test]
 async fn test_parent_cancel_stops_subagent(cx: &mut TestAppContext) {
     init_test(cx);

crates/agent/src/thread.rs 🔗

@@ -1577,20 +1577,19 @@ impl Thread {
         self.add_tool(WebSearchTool);
 
         self.add_tool(DiagnosticsTool::new(self.project.clone()));
-        if cx.has_flag::<LspToolFeatureFlag>() {
-            let code_action_store: CodeActionStore = cx.new(|_cx| None);
-            self.add_tool(FindReferencesTool::new(self.project.clone()));
-            self.add_tool(GetCodeActionsTool::new(
-                self.project.clone(),
-                code_action_store.clone(),
-            ));
-            self.add_tool(ApplyCodeActionTool::new(
-                self.project.clone(),
-                code_action_store,
-            ));
-            self.add_tool(GoToDefinitionTool::new(self.project.clone()));
-            self.add_tool(RenameTool::new(self.project.clone()));
-        }
+
+        let code_action_store: CodeActionStore = cx.new(|_cx| None);
+        self.add_tool(FindReferencesTool::new(self.project.clone()));
+        self.add_tool(GetCodeActionsTool::new(
+            self.project.clone(),
+            code_action_store.clone(),
+        ));
+        self.add_tool(ApplyCodeActionTool::new(
+            self.project.clone(),
+            code_action_store,
+        ));
+        self.add_tool(GoToDefinitionTool::new(self.project.clone()));
+        self.add_tool(RenameTool::new(self.project.clone()));
 
         if self.depth() < MAX_SUBAGENT_DEPTH {
             self.add_tool(SpawnAgentTool::new(environment));
@@ -2894,6 +2893,17 @@ impl Thread {
                     None
                 }
             })
+            .filter(|(tool_name, _)| {
+                cx.has_flag::<LspToolFeatureFlag>()
+                    || !matches!(
+                        tool_name.as_ref(),
+                        FindReferencesTool::NAME
+                            | GetCodeActionsTool::NAME
+                            | ApplyCodeActionTool::NAME
+                            | GoToDefinitionTool::NAME
+                            | RenameTool::NAME
+                    )
+            })
             .collect::<BTreeMap<_, _>>();
 
         let mut context_server_tools = Vec::new();
@@ -2957,10 +2967,6 @@ impl Thread {
         self.tools.contains_key(name)
     }
 
-    pub fn registered_tool_names(&self) -> Vec<SharedString> {
-        self.tools.keys().cloned().collect()
-    }
-
     pub(crate) fn register_running_subagent(&mut self, subagent: WeakEntity<Thread>) {
         self.running_subagents.push(subagent);
     }