From 7556cf8ced19cae275fb4d6678858050d4d51df4 Mon Sep 17 00:00:00 2001 From: Bennet Bo Fenner Date: Thu, 7 May 2026 15:47:13 +0200 Subject: [PATCH] agent: Fix race-condition for LSP tool registration (#56044) 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 --- crates/agent/src/tests/mod.rs | 99 +++++++++++++++++++++++++++++++++++ crates/agent/src/thread.rs | 42 ++++++++------- 2 files changed, 123 insertions(+), 18 deletions(-) diff --git a/crates/agent/src/tests/mod.rs b/crates/agent/src/tests/mod.rs index 57cec0bc5d07a99578d5a64d1ef442f47a5e04e9..e4dd5a2425750141104c7ddc496879079fa3643a 100644 --- a/crates/agent/src/tests/mod.rs +++ b/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), + 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); diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 78a4b2fd488918a0e2e773e7f07c55e013a52b48..ef03f47a8d37039f02bda12c6740f59f8dfbc1ff 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,17 @@ impl Thread { None } }) + .filter(|(tool_name, _)| { + cx.has_flag::() + || !matches!( + tool_name.as_ref(), + FindReferencesTool::NAME + | GetCodeActionsTool::NAME + | ApplyCodeActionTool::NAME + | GoToDefinitionTool::NAME + | RenameTool::NAME + ) + }) .collect::>(); 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 { - self.tools.keys().cloned().collect() - } - pub(crate) fn register_running_subagent(&mut self, subagent: WeakEntity) { self.running_subagents.push(subagent); }