From 275c0aa5a9e3c35e7c8a13c929a247edf9d8d2d5 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Fri, 27 Feb 2026 23:47:46 -0500 Subject: [PATCH] Extract tool execution from session.rs into tool_execution.rs Move the tool execution subsystem (~250 lines) into a dedicated module: - execute_tool_calls: dispatches to parallel or sequential execution - execute_and_emit_one_tool: wraps execution with event emission - execute_one_tool: registry lookup, validation, and execution - validate_tool_args: JSON schema validation for tool arguments - truncate_tool_result: output truncation for history storage These functions have zero dependency on Session's history, state machine, or LLM interaction. Session now calls the extracted free functions, passing needed context as parameters. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/agent/src/tool_execution.rs | 276 +++++++++++++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 crates/agent/src/tool_execution.rs diff --git a/crates/agent/src/tool_execution.rs b/crates/agent/src/tool_execution.rs new file mode 100644 index 000000000..070d424ce --- /dev/null +++ b/crates/agent/src/tool_execution.rs @@ -0,0 +1,276 @@ +use crate::config::{SessionConfig, ToolApprovalFn}; +use crate::event::EventEmitter; +use crate::execution_env::ExecutionEnvironment; +use crate::tool_registry::ToolRegistry; +use crate::truncation::truncate_tool_output; +use crate::types::AgentEvent; +use llm::types::ToolResult; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; + +/// Execute tool calls, choosing parallel or sequential based on `parallel` flag. +#[allow(clippy::too_many_arguments)] +pub async fn execute_tool_calls( + tool_calls: &[llm::types::ToolCall], + parallel: bool, + registry: &ToolRegistry, + env: Arc, + tool_approval: Option<&ToolApprovalFn>, + cancel_token: &CancellationToken, + config: &SessionConfig, + emitter: &EventEmitter, + session_id: &str, +) -> Vec { + if parallel && tool_calls.len() > 1 { + execute_tool_calls_parallel( + tool_calls, registry, env, tool_approval, cancel_token, config, emitter, session_id, + ) + .await + } else { + execute_tool_calls_sequential( + tool_calls, registry, env, tool_approval, cancel_token, config, emitter, session_id, + ) + .await + } +} + +#[allow(clippy::too_many_arguments)] +async fn execute_tool_calls_sequential( + tool_calls: &[llm::types::ToolCall], + registry: &ToolRegistry, + env: Arc, + tool_approval: Option<&ToolApprovalFn>, + cancel_token: &CancellationToken, + config: &SessionConfig, + emitter: &EventEmitter, + session_id: &str, +) -> Vec { + let mut results = Vec::new(); + for tc in tool_calls { + if cancel_token.is_cancelled() { + results.push(ToolResult::error(tc.id.clone(), "Cancelled")); + continue; + } + + let result = execute_and_emit_one_tool( + tc, + registry, + env.clone(), + tool_approval, + cancel_token.child_token(), + config, + emitter, + session_id, + ) + .await; + results.push(result); + } + results +} + +#[allow(clippy::too_many_arguments)] +async fn execute_tool_calls_parallel( + tool_calls: &[llm::types::ToolCall], + registry: &ToolRegistry, + env: Arc, + tool_approval: Option<&ToolApprovalFn>, + cancel_token: &CancellationToken, + config: &SessionConfig, + emitter: &EventEmitter, + session_id: &str, +) -> Vec { + let futures: Vec<_> = tool_calls + .iter() + .map(|tc| { + let emitter = emitter.clone(); + let env = env.clone(); + let config = config.clone(); + let cancel_token = cancel_token.clone(); + let tc = tc.clone(); + let session_id = session_id.to_owned(); + let tool_approval = tool_approval.cloned(); + // Look up the tool before spawning since ToolRegistry is not Send. + let registered_tool = registry.get(&tc.name).cloned(); + async move { + execute_and_emit_one_tool_with_lookup( + &tc, + registered_tool.as_ref(), + env, + tool_approval.as_ref(), + cancel_token.child_token(), + &config, + &emitter, + &session_id, + ) + .await + } + }) + .collect(); + + futures::future::join_all(futures).await +} + +/// Execute a single tool call with event emission and output truncation. +#[allow(clippy::too_many_arguments)] +pub async fn execute_and_emit_one_tool( + tc: &llm::types::ToolCall, + registry: &ToolRegistry, + env: Arc, + tool_approval: Option<&ToolApprovalFn>, + cancel_token: CancellationToken, + config: &SessionConfig, + emitter: &EventEmitter, + session_id: &str, +) -> ToolResult { + execute_and_emit_one_tool_with_lookup( + tc, + registry.get(&tc.name), + env, + tool_approval, + cancel_token, + config, + emitter, + session_id, + ) + .await +} + +/// Execute a single tool call with event emission, using a pre-looked-up tool reference. +#[allow(clippy::too_many_arguments)] +async fn execute_and_emit_one_tool_with_lookup( + tc: &llm::types::ToolCall, + registered_tool: Option<&crate::tool_registry::RegisteredTool>, + env: Arc, + tool_approval: Option<&ToolApprovalFn>, + cancel_token: CancellationToken, + config: &SessionConfig, + emitter: &EventEmitter, + session_id: &str, +) -> ToolResult { + emitter.emit( + session_id.to_owned(), + AgentEvent::ToolCallStarted { + tool_name: tc.name.clone(), + tool_call_id: tc.id.clone(), + arguments: tc.arguments.clone(), + }, + ); + + let result = execute_one_tool( + &tc.id, + &tc.name, + &tc.arguments, + registered_tool, + env, + tool_approval, + cancel_token, + ) + .await; + + emitter.emit( + session_id.to_owned(), + AgentEvent::ToolCallOutputDelta { + delta: result.content.to_string(), + }, + ); + + emitter.emit( + session_id.to_owned(), + AgentEvent::ToolCallCompleted { + tool_name: tc.name.clone(), + tool_call_id: tc.id.clone(), + output: result.content.clone(), + is_error: result.is_error, + }, + ); + + truncate_tool_result(&result, &tc.name, config) +} + +/// Execute a single tool call: argument validation and execution. +async fn execute_one_tool( + tool_call_id: &str, + tool_name: &str, + arguments: &serde_json::Value, + registered_tool: Option<&crate::tool_registry::RegisteredTool>, + env: Arc, + tool_approval: Option<&ToolApprovalFn>, + cancel_token: CancellationToken, +) -> ToolResult { + if let Some(approval_fn) = tool_approval { + if let Err(denial_message) = approval_fn(tool_name, arguments) { + return ToolResult::error(tool_call_id, denial_message); + } + } + + match registered_tool { + Some(tool) => { + if let Err(validation_error) = + validate_tool_args(&tool.definition.parameters, arguments) + { + return ToolResult::error(tool_call_id, validation_error); + } + + let ctx = crate::tool_registry::ToolContext { + env, + cancel: cancel_token, + }; + match (tool.executor)(arguments.clone(), ctx).await { + Ok(output) => ToolResult::success(tool_call_id, serde_json::json!(output)), + Err(err) => ToolResult::error(tool_call_id, err), + } + } + None => ToolResult::error(tool_call_id, format!("Unknown tool: {tool_name}")), + } +} + +/// Truncate tool output for history storage while preserving identity fields. +fn truncate_tool_result( + result: &ToolResult, + tool_name: &str, + config: &SessionConfig, +) -> ToolResult { + let truncated_content = match &result.content { + serde_json::Value::String(s) => { + serde_json::json!(truncate_tool_output(s, tool_name, config)) + } + other => other.clone(), + }; + + ToolResult { + tool_call_id: result.tool_call_id.clone(), + content: truncated_content, + is_error: result.is_error, + image_data: result.image_data.clone(), + image_media_type: result.image_media_type.clone(), + } +} + +pub fn validate_tool_args( + schema: &serde_json::Value, + args: &serde_json::Value, +) -> Result<(), String> { + // Skip validation for empty/trivial schemas + if schema.is_null() { + return Ok(()); + } + if let Some(obj) = schema.as_object() { + if obj.is_empty() { + return Ok(()); + } + } + + let validator = + jsonschema::validator_for(schema).map_err(|e| format!("Invalid tool schema: {e}"))?; + + let errors: Vec = validator.iter_errors(args).map(|e| e.to_string()).collect(); + + if errors.is_empty() { + Ok(()) + } else { + Err(format!( + "Tool argument validation failed: {}", + errors.join("; ") + )) + } +}