dgf1988/codewhale
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

refactor(engine): modularize turn tool setup

Browse Source
This commit is contained in:
Hunter Bown
2026-05-01 06:07:59 -05:00
parent 0887a88465
commit f0fad7aa2e
4 changed files with 167 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/tui/src/core/engine.rs
+10 -67
View File
@@ -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;
crates/tui/src/core/engine/tests.rs
+70
View File
@@ -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());
crates/tui/src/core/engine/tool_catalog.rs
+35
View File
@@ -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<Tool>,
mut mcp_tools: Vec<Tool>,
mode: AppMode,
) -> Vec<Tool> {
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<Tool>) {
if !catalog.iter().any(|t| t.name == CODE_EXECUTION_TOOL_NAME) {
catalog.push(Tool {
crates/tui/src/core/engine/tool_setup.rs
+52
View File
@@ -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
}
}