From f0fad7aa2ebdb976757ad0d24047c51c1031cf7a Mon Sep 17 00:00:00 2001 From: Hunter Bown Date: Fri, 1 May 2026 06:07:59 -0500 Subject: [PATCH] refactor(engine): modularize turn tool setup --- crates/tui/src/core/engine.rs | 77 +++------------------- crates/tui/src/core/engine/tests.rs | 70 ++++++++++++++++++++ crates/tui/src/core/engine/tool_catalog.rs | 35 ++++++++++ crates/tui/src/core/engine/tool_setup.rs | 52 +++++++++++++++ 4 files changed, 167 insertions(+), 67 deletions(-) create mode 100644 crates/tui/src/core/engine/tool_setup.rs diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index 04bb656a..af3998c6 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -1302,43 +1302,7 @@ impl Engine { let plan_state = self.config.plan_state.clone(); let tool_context = self.build_tool_context(mode, auto_approve); - let mut builder = if mode == AppMode::Plan { - ToolRegistryBuilder::new() - .with_read_only_file_tools() - .with_search_tools() - .with_git_tools() - .with_git_history_tools() - .with_diagnostics_tool() - .with_validation_tools() - .with_runtime_task_tools() - .with_todo_tool(todo_list.clone()) - .with_plan_tool(plan_state.clone()) - } else { - ToolRegistryBuilder::new() - .with_agent_tools(self.session.allow_shell) - .with_todo_tool(todo_list.clone()) - .with_plan_tool(plan_state.clone()) - }; - - builder = builder - .with_review_tool(self.deepseek_client.clone(), self.session.model.clone()) - .with_rlm_tool(self.deepseek_client.clone(), self.session.model.clone()) - .with_user_input_tool() - .with_parallel_tool(); - - if self.config.features.enabled(Feature::ApplyPatch) && mode != AppMode::Plan { - builder = builder.with_patch_tools(); - } - if self.config.features.enabled(Feature::WebSearch) { - builder = builder.with_web_tools(); - } - // Plan mode now keeps shell available — the existing approval flow - // and command-safety classifier gate destructive commands. Writes - // and patches stay blocked above; that's the only "destructive" - // boundary plan mode enforces by tool registration. - if self.config.features.enabled(Feature::ShellTool) && self.session.allow_shell { - builder = builder.with_shell_tools(); - } + let builder = self.build_turn_tool_registry_builder(mode, todo_list, plan_state); // Mailbox for structured sub-agent envelopes (#128/#130). One per // turn: the receiver is drained by a short-lived task that converts @@ -1411,31 +1375,9 @@ impl Engine { } else { Vec::new() }; - let tools = tool_registry.as_ref().map(|registry| { - let mut tools = registry.to_api_tools(); - for tool in &mut tools { - tool.defer_loading = Some(should_default_defer_tool(&tool.name, mode)); - } - let mut mcp_tools = mcp_tools; - for tool in &mut mcp_tools { - if mode == AppMode::Yolo { - tool.defer_loading = Some(false); - continue; - } - - let keep_loaded = matches!( - tool.name.as_str(), - "list_mcp_resources" - | "list_mcp_resource_templates" - | "mcp_read_resource" - | "read_mcp_resource" - | "mcp_get_prompt" - ); - tool.defer_loading = Some(!keep_loaded); - } - tools.extend(mcp_tools); - tools - }); + let tools = tool_registry + .as_ref() + .map(|registry| build_model_tool_catalog(registry.to_api_tools(), mcp_tools, mode)); // Main turn loop let (status, error) = self @@ -2464,6 +2406,7 @@ mod approval; mod capacity_flow; mod dispatch; mod tool_catalog; +mod tool_setup; mod turn_loop; use self::approval::{ApprovalDecision, ApprovalResult, UserInputDecision}; @@ -2473,14 +2416,14 @@ use self::dispatch::{ mcp_tool_is_read_only, parse_parallel_tool_calls, parse_tool_input, should_force_update_plan_first, should_parallelize_tool_batch, should_stop_after_plan_tool, }; -#[cfg(test)] -use self::tool_catalog::TOOL_SEARCH_BM25_NAME; use self::tool_catalog::{ CODE_EXECUTION_TOOL_NAME, MULTI_TOOL_PARALLEL_NAME, REQUEST_USER_INPUT_NAME, - active_tools_for_step, ensure_advanced_tooling, execute_code_execution_tool, - execute_tool_search, initial_active_tools, is_tool_search_tool, - maybe_activate_requested_deferred_tool, missing_tool_error_message, should_default_defer_tool, + active_tools_for_step, build_model_tool_catalog, ensure_advanced_tooling, + execute_code_execution_tool, execute_tool_search, initial_active_tools, is_tool_search_tool, + maybe_activate_requested_deferred_tool, missing_tool_error_message, }; +#[cfg(test)] +use self::tool_catalog::{TOOL_SEARCH_BM25_NAME, should_default_defer_tool}; #[cfg(test)] mod tests; diff --git a/crates/tui/src/core/engine/tests.rs b/crates/tui/src/core/engine/tests.rs index d4194473..8dfc291c 100644 --- a/crates/tui/src/core/engine/tests.rs +++ b/crates/tui/src/core/engine/tests.rs @@ -36,6 +36,20 @@ fn make_plan( } } +fn api_tool(name: &str) -> Tool { + Tool { + tool_type: Some("function".to_string()), + name: name.to_string(), + description: format!("Test tool {name}"), + input_schema: json!({"type": "object"}), + allowed_callers: Some(vec!["direct".to_string()]), + defer_loading: None, + input_examples: None, + strict: None, + cache_control: None, + } +} + #[test] fn engine_handle_cancel_tracks_latest_turn_token() { let (mut engine, handle) = Engine::new(EngineConfig::default(), &Config::default()); @@ -205,6 +219,62 @@ fn non_yolo_mode_retains_default_defer_policy() { )); } +#[test] +fn model_tool_catalog_applies_native_and_mcp_deferral() { + let catalog = build_model_tool_catalog( + vec![ + api_tool("read_file"), + api_tool("exec_shell"), + api_tool("project_map"), + ], + vec![api_tool("list_mcp_resources"), api_tool("mcp_server_write")], + AppMode::Agent, + ); + + let defer_loading = |name: &str| { + catalog + .iter() + .find(|tool| tool.name == name) + .and_then(|tool| tool.defer_loading) + }; + + assert_eq!(defer_loading("read_file"), Some(false)); + assert_eq!(defer_loading("exec_shell"), Some(false)); + assert_eq!(defer_loading("project_map"), Some(true)); + assert_eq!(defer_loading("list_mcp_resources"), Some(false)); + assert_eq!(defer_loading("mcp_server_write"), Some(true)); +} + +#[test] +fn model_tool_catalog_keeps_everything_loaded_in_yolo_mode() { + let catalog = build_model_tool_catalog( + vec![api_tool("project_map")], + vec![api_tool("mcp_server_write")], + AppMode::Yolo, + ); + + assert!(catalog.iter().all(|tool| tool.defer_loading == Some(false))); +} + +#[test] +fn turn_tool_registry_builder_keeps_plan_mode_read_only_for_files() { + let (engine, _handle) = Engine::new(EngineConfig::default(), &Config::default()); + let registry = engine + .build_turn_tool_registry_builder( + AppMode::Plan, + engine.config.todos.clone(), + engine.config.plan_state.clone(), + ) + .build(engine.build_tool_context(AppMode::Plan, false)); + + assert!(registry.contains("read_file")); + assert!(registry.contains("list_dir")); + assert!(!registry.contains("write_file")); + assert!(!registry.contains("edit_file")); + assert!(registry.contains("update_plan")); + assert!(registry.contains("task_create")); +} + #[test] fn agent_mode_can_build_auto_approved_tool_context() { let (engine, _handle) = Engine::new(EngineConfig::default(), &Config::default()); diff --git a/crates/tui/src/core/engine/tool_catalog.rs b/crates/tui/src/core/engine/tool_catalog.rs index 011d632e..2581fd4f 100644 --- a/crates/tui/src/core/engine/tool_catalog.rs +++ b/crates/tui/src/core/engine/tool_catalog.rs @@ -74,6 +74,41 @@ pub(super) fn should_default_defer_tool(name: &str, mode: AppMode) -> bool { ) } +pub(super) fn apply_native_tool_deferral(catalog: &mut [Tool], mode: AppMode) { + for tool in catalog { + tool.defer_loading = Some(should_default_defer_tool(&tool.name, mode)); + } +} + +fn should_keep_mcp_tool_loaded(name: &str) -> bool { + matches!( + name, + "list_mcp_resources" + | "list_mcp_resource_templates" + | "mcp_read_resource" + | "read_mcp_resource" + | "mcp_get_prompt" + ) +} + +pub(super) fn apply_mcp_tool_deferral(catalog: &mut [Tool], mode: AppMode) { + for tool in catalog { + tool.defer_loading = + Some(mode != AppMode::Yolo && !should_keep_mcp_tool_loaded(&tool.name)); + } +} + +pub(super) fn build_model_tool_catalog( + mut native_tools: Vec, + mut mcp_tools: Vec, + mode: AppMode, +) -> Vec { + apply_native_tool_deferral(&mut native_tools, mode); + apply_mcp_tool_deferral(&mut mcp_tools, mode); + native_tools.extend(mcp_tools); + native_tools +} + pub(super) fn ensure_advanced_tooling(catalog: &mut Vec) { if !catalog.iter().any(|t| t.name == CODE_EXECUTION_TOOL_NAME) { catalog.push(Tool { diff --git a/crates/tui/src/core/engine/tool_setup.rs b/crates/tui/src/core/engine/tool_setup.rs new file mode 100644 index 00000000..220d4706 --- /dev/null +++ b/crates/tui/src/core/engine/tool_setup.rs @@ -0,0 +1,52 @@ +//! Per-turn tool registry setup. +//! +//! This keeps mode/feature-specific registry construction out of the send path. + +use super::*; + +impl Engine { + pub(super) fn build_turn_tool_registry_builder( + &self, + mode: AppMode, + todo_list: SharedTodoList, + plan_state: SharedPlanState, + ) -> ToolRegistryBuilder { + let mut builder = if mode == AppMode::Plan { + ToolRegistryBuilder::new() + .with_read_only_file_tools() + .with_search_tools() + .with_git_tools() + .with_git_history_tools() + .with_diagnostics_tool() + .with_validation_tools() + .with_runtime_task_tools() + .with_todo_tool(todo_list) + .with_plan_tool(plan_state) + } else { + ToolRegistryBuilder::new() + .with_agent_tools(self.session.allow_shell) + .with_todo_tool(todo_list) + .with_plan_tool(plan_state) + }; + + builder = builder + .with_review_tool(self.deepseek_client.clone(), self.session.model.clone()) + .with_rlm_tool(self.deepseek_client.clone(), self.session.model.clone()) + .with_user_input_tool() + .with_parallel_tool(); + + if self.config.features.enabled(Feature::ApplyPatch) && mode != AppMode::Plan { + builder = builder.with_patch_tools(); + } + if self.config.features.enabled(Feature::WebSearch) { + builder = builder.with_web_tools(); + } + // Plan mode keeps shell available when the session allows it; command + // safety and approval checks still gate risky commands. + if self.config.features.enabled(Feature::ShellTool) && self.session.allow_shell { + builder = builder.with_shell_tools(); + } + + builder + } +}