From 39c016a322582410e8a32754c008e4fd03dc3e8a Mon Sep 17 00:00:00 2001 From: Bennet Bo Fenner Date: Thu, 7 May 2026 14:13:39 +0200 Subject: [PATCH 1/2] agent: Fix race-condition for LSP tool registration --- crates/agent/src/tests/mod.rs | 100 ++++++++++++++++++++++++++++++++++ crates/agent/src/thread.rs | 43 +++++++++------ 2 files changed, 125 insertions(+), 18 deletions(-) diff --git a/crates/agent/src/tests/mod.rs b/crates/agent/src/tests/mod.rs index 57cec0bc5d07a9..e4a1be9d008ae3 100644 --- a/crates/agent/src/tests/mod.rs +++ b/crates/agent/src/tests/mod.rs @@ -5545,6 +5545,106 @@ 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), + cx, + ); + thread.add_default_tools(environment, cx); + thread + }); + + let lsp_tool_names = [ + DiagnosticsTool::NAME, + 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); diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 78a4b2fd488918..8ce24e81875865 100644 --- a/crates/agent/src/thread.rs +++ b/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::() { - 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,18 @@ impl Thread { None } }) + .filter(|(tool_name, _)| { + cx.has_flag::() + || !matches!( + tool_name.as_ref(), + DiagnosticsTool::NAME + | FindReferencesTool::NAME + | GetCodeActionsTool::NAME + | ApplyCodeActionTool::NAME + | GoToDefinitionTool::NAME + | RenameTool::NAME + ) + }) .collect::>(); let mut context_server_tools = Vec::new(); @@ -2957,10 +2968,6 @@ impl Thread { self.tools.contains_key(name) } - pub fn registered_tool_names(&self) -> Vec { - self.tools.keys().cloned().collect() - } - pub(crate) fn register_running_subagent(&mut self, subagent: WeakEntity) { self.running_subagents.push(subagent); } From 36d957a546917472438226247676e36e46252e61 Mon Sep 17 00:00:00 2001 From: Bennet Bo Fenner Date: Thu, 7 May 2026 14:50:46 +0200 Subject: [PATCH 2/2] Update filter --- crates/agent/src/tests/mod.rs | 1 - crates/agent/src/thread.rs | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/crates/agent/src/tests/mod.rs b/crates/agent/src/tests/mod.rs index e4a1be9d008ae3..e4dd5a24257501 100644 --- a/crates/agent/src/tests/mod.rs +++ b/crates/agent/src/tests/mod.rs @@ -5575,7 +5575,6 @@ async fn test_lsp_tools_gated_by_feature_flag(cx: &mut TestAppContext) { }); let lsp_tool_names = [ - DiagnosticsTool::NAME, FindReferencesTool::NAME, GetCodeActionsTool::NAME, ApplyCodeActionTool::NAME, diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 8ce24e81875865..ef03f47a8d3703 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -2897,8 +2897,7 @@ impl Thread { cx.has_flag::() || !matches!( tool_name.as_ref(), - DiagnosticsTool::NAME - | FindReferencesTool::NAME + FindReferencesTool::NAME | GetCodeActionsTool::NAME | ApplyCodeActionTool::NAME | GoToDefinitionTool::NAME