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); }