use std::future::Future; use std::pin::Pin; use std::sync::Arc; use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; use crate::types::{Message, ToolCall, ToolDefinition, ToolResult}; /// Context passed to tool execute handlers (Section 5.2). #[derive(Clone)] pub struct ToolContext { pub tool_call_id: String, pub messages: Vec, pub abort_signal: Option, } /// An execute handler for a tool. pub type ExecuteHandler = Arc< dyn Fn( serde_json::Value, ToolContext, ) -> Pin> + Send>> + Send + Sync, >; /// A tool with an optional execute handler (Section 5.1, 5.5). /// "Active" tools have an execute handler and are automatically executed. /// "Passive" tools have no handler and are returned to the caller. pub struct Tool { pub definition: ToolDefinition, pub execute: Option, } impl Tool { /// Create a passive tool (no execute handler). /// /// # Panics /// /// Panics if the tool name is invalid (see [`validate_tool_name`]). #[must_use] pub fn passive(name: &str, description: &str, parameters: serde_json::Value) -> Self { if let Err(e) = validate_tool_name(name) { panic!("Invalid tool name: {e}"); } Self { definition: ToolDefinition { name: name.to_string(), description: description.to_string(), parameters, }, execute: None, } } /// Create an active tool with an execute handler. /// /// # Panics /// /// Panics if the tool name is invalid (see [`validate_tool_name`]). pub fn active( name: &str, description: &str, parameters: serde_json::Value, handler: F, ) -> Self where F: Fn(serde_json::Value, ToolContext) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { if let Err(e) = validate_tool_name(name) { panic!("Invalid tool name: {e}"); } Self { definition: ToolDefinition { name: name.to_string(), description: description.to_string(), parameters, }, execute: Some(Arc::new(move |args, ctx| Box::pin(handler(args, ctx)))), } } #[must_use] pub fn is_active(&self) -> bool { self.execute.is_some() } } /// Validate a tool name: [a-zA-Z][a-zA-Z0-9_]* max 64 chars (Section 5.1). /// /// # Errors /// /// Returns a description of the validation failure if the name is empty, /// too long, starts with a non-letter, or contains invalid characters. pub fn validate_tool_name(name: &str) -> Result<(), String> { if name.is_empty() { return Err("Tool name cannot be empty".to_string()); } if name.len() > 64 { return Err(format!("Tool name '{name}' exceeds 64 character limit")); } let mut chars = name.chars(); if let Some(first) = chars.next() { if !first.is_ascii_alphabetic() { return Err(format!("Tool name '{name}' must start with a letter")); } } for ch in chars { if !ch.is_ascii_alphanumeric() && ch != '_' { return Err(format!( "Tool name '{name}' contains invalid character '{ch}'" )); } } Ok(()) } /// A callback to repair invalid tool call arguments (Section 5.8). /// Receives the tool call and the validation error message, returns repaired /// arguments or an error if repair is not possible. pub type RepairToolCallFn = Arc< dyn Fn( ToolCall, String, ) -> Pin> + Send>> + Send + Sync, >; /// Validate tool call arguments against the tool's parameter schema. /// Performs a lightweight structural check: verifies that when the schema /// specifies `"type": "object"`, the arguments are a JSON object, and that /// required properties are present. fn validate_tool_args(args: &serde_json::Value, schema: &serde_json::Value) -> Result<(), String> { let schema_type = schema.get("type").and_then(serde_json::Value::as_str); if schema_type == Some("object") && !args.is_object() { return Err(format!( "Expected object arguments, got {}", args_type_name(args) )); } if let (Some(obj), Some(required)) = ( args.as_object(), schema.get("required").and_then(serde_json::Value::as_array), ) { let missing: Vec<&str> = required .iter() .filter_map(serde_json::Value::as_str) .filter(|key| !obj.contains_key(*key)) .collect(); if !missing.is_empty() { return Err(format!( "Missing required properties: {}", missing.join(", ") )); } } Ok(()) } const fn args_type_name(value: &serde_json::Value) -> &'static str { match value { serde_json::Value::Null => "null", serde_json::Value::Bool(_) => "boolean", serde_json::Value::Number(_) => "number", serde_json::Value::String(_) => "string", serde_json::Value::Array(_) => "array", serde_json::Value::Object(_) => "object", } } /// Execute all tool calls with optional schema validation and repair (Section /// 5.8). /// /// Before calling a tool's execute handler, validates the arguments against the /// tool's parameter schema. If validation fails and a `repair` callback is /// provided, calls it to attempt repair. If repair succeeds, uses the repaired /// arguments. If repair fails or is not configured, returns an error /// `ToolResult`. pub async fn execute_all_tools_with_repair( tools: &[&Tool], tool_calls: &[ToolCall], messages: &[Message], abort_signal: Option<&CancellationToken>, repair: Option<&RepairToolCallFn>, ) -> Vec { use futures::future::join_all; let futures: Vec<_> = tool_calls .iter() .map(|call| { let tool = tools.iter().find(|t| t.definition.name == call.name).copied(); let call_id = call.id.clone(); let call_name = call.name.clone(); let args = call.arguments.clone(); let call_clone = call.clone(); let ctx = ToolContext { tool_call_id: call_id.clone(), messages: messages.to_vec(), abort_signal: abort_signal.cloned(), }; async move { let Some(t) = tool else { return ToolResult::error(call_id, format!("Unknown tool: {call_name}")); }; let Some(handler) = &t.execute else { return ToolResult::error(call_id, format!("Unknown tool: {call_name}")); }; let validated_args = match validate_tool_args(&args, &t.definition.parameters) { Ok(()) => args, Err(validation_error) => { debug!(tool = %call_name, "Tool call validation failed"); if let Some(repair_fn) = repair { match repair_fn(call_clone, validation_error).await { Ok(repaired) => repaired, Err(repair_error) => { warn!(tool = %call_name, "Tool call repair failed"); return ToolResult::error( call_id, format!("Tool call validation failed and repair failed: {repair_error}"), ); } } } else { return ToolResult::error( call_id, format!("Tool call validation failed: {validation_error}"), ); } } }; match handler(validated_args, ctx).await { Ok(result) => ToolResult::success(call_id, result), Err(err_msg) => { warn!(tool = %call_name, "Tool execution returned error"); ToolResult::error(call_id, err_msg) } } } }) .collect(); join_all(futures).await } #[cfg(test)] mod tests { use super::*; #[test] fn validate_tool_name_valid() { assert!(validate_tool_name("get_weather").is_ok()); assert!(validate_tool_name("a").is_ok()); assert!(validate_tool_name("myTool123").is_ok()); assert!(validate_tool_name("A_B_C").is_ok()); } #[test] fn validate_tool_name_empty() { assert!(validate_tool_name("").is_err()); } #[test] fn validate_tool_name_starts_with_number() { assert!(validate_tool_name("1tool").is_err()); } #[test] fn validate_tool_name_starts_with_underscore() { assert!(validate_tool_name("_tool").is_err()); } #[test] fn validate_tool_name_contains_dash() { assert!(validate_tool_name("my-tool").is_err()); } #[test] fn validate_tool_name_too_long() { let name = "a".repeat(65); assert!(validate_tool_name(&name).is_err()); } #[test] fn validate_tool_name_max_length_ok() { let name = "a".repeat(64); assert!(validate_tool_name(&name).is_ok()); } #[test] fn passive_tool_is_not_active() { let tool = Tool::passive( "test", "test tool", serde_json::json!({"type": "object", "properties": {}}), ); assert!(!tool.is_active()); } #[test] fn active_tool_is_active() { let tool = Tool::active( "test", "test tool", serde_json::json!({"type": "object", "properties": {}}), |_args, _ctx| async { Ok(serde_json::json!("result")) }, ); assert!(tool.is_active()); } #[test] #[should_panic(expected = "Invalid tool name")] fn passive_tool_panics_on_invalid_name() { let _ = Tool::passive( "1invalid", "bad name", serde_json::json!({"type": "object"}), ); } #[test] #[should_panic(expected = "Invalid tool name")] fn active_tool_panics_on_invalid_name() { Tool::active( "my-tool", "bad name", serde_json::json!({"type": "object"}), |_args, _ctx| async { Ok(serde_json::json!("result")) }, ); } #[test] fn validate_tool_args_valid_object() { let schema = serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}}); let args = serde_json::json!({"name": "Alice"}); assert!(validate_tool_args(&args, &schema).is_ok()); } #[test] fn validate_tool_args_non_object_when_object_expected() { let schema = serde_json::json!({"type": "object", "properties": {}}); let args = serde_json::json!("not an object"); let result = validate_tool_args(&args, &schema); assert!(result.is_err()); assert!(result.unwrap_err().contains("Expected object")); } #[test] fn validate_tool_args_missing_required_properties() { let schema = serde_json::json!({ "type": "object", "properties": {"name": {"type": "string"}, "age": {"type": "number"}}, "required": ["name", "age"] }); let args = serde_json::json!({"name": "Alice"}); let result = validate_tool_args(&args, &schema); assert!(result.is_err()); assert!(result.unwrap_err().contains("age")); } #[test] fn validate_tool_args_no_schema_type_passes() { let schema = serde_json::json!({}); let args = serde_json::json!("anything"); assert!(validate_tool_args(&args, &schema).is_ok()); } #[tokio::test] async fn execute_with_repair_valid_args_no_repair_needed() { let tools = [Tool::active( "greet", "Greet someone", serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}}), |args, _ctx| async move { let name = args["name"].as_str().unwrap_or("world"); Ok(serde_json::json!(format!("Hello, {}!", name))) }, )]; let calls = vec![ToolCall::new( "call_1", "greet", serde_json::json!({"name": "Alice"}), )]; let tool_refs: Vec<&Tool> = tools.iter().collect(); let results = execute_all_tools_with_repair(&tool_refs, &calls, &[], None, None).await; assert_eq!(results.len(), 1); assert!(!results[0].is_error); assert_eq!(results[0].content, serde_json::json!("Hello, Alice!")); } #[tokio::test] async fn execute_with_repair_invalid_args_no_repair_fn() { let tools = [Tool::active( "greet", "Greet someone", serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}), |args, _ctx| async move { let name = args["name"].as_str().unwrap_or("world"); Ok(serde_json::json!(format!("Hello, {}!", name))) }, )]; let calls = vec![ToolCall::new("call_1", "greet", serde_json::json!({}))]; let tool_refs: Vec<&Tool> = tools.iter().collect(); let results = execute_all_tools_with_repair(&tool_refs, &calls, &[], None, None).await; assert_eq!(results.len(), 1); assert!(results[0].is_error); assert!( results[0] .content .as_str() .unwrap() .contains("validation failed") ); } #[tokio::test] async fn execute_with_repair_invalid_args_repair_succeeds() { let tools = [Tool::active( "greet", "Greet someone", serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}), |args, _ctx| async move { let name = args["name"].as_str().unwrap_or("world"); Ok(serde_json::json!(format!("Hello, {}!", name))) }, )]; let calls = vec![ToolCall::new("call_1", "greet", serde_json::json!({}))]; let tool_refs: Vec<&Tool> = tools.iter().collect(); let repair: RepairToolCallFn = Arc::new(|_call, _error| { Box::pin(async { Ok(serde_json::json!({"name": "Repaired"})) }) }); let results = execute_all_tools_with_repair(&tool_refs, &calls, &[], None, Some(&repair)).await; assert_eq!(results.len(), 1); assert!(!results[0].is_error); assert_eq!(results[0].content, serde_json::json!("Hello, Repaired!")); } #[tokio::test] async fn execute_with_repair_invalid_args_repair_fails() { let tools = [Tool::active( "greet", "Greet someone", serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}), |args, _ctx| async move { let name = args["name"].as_str().unwrap_or("world"); Ok(serde_json::json!(format!("Hello, {}!", name))) }, )]; let calls = vec![ToolCall::new("call_1", "greet", serde_json::json!({}))]; let tool_refs: Vec<&Tool> = tools.iter().collect(); let repair: RepairToolCallFn = Arc::new(|_call, _error| Box::pin(async { Err("cannot repair".to_string()) })); let results = execute_all_tools_with_repair(&tool_refs, &calls, &[], None, Some(&repair)).await; assert_eq!(results.len(), 1); assert!(results[0].is_error); assert!( results[0] .content .as_str() .unwrap() .contains("repair failed") ); } // --- args_type_name --- #[test] fn args_type_name_null() { assert_eq!(args_type_name(&serde_json::Value::Null), "null"); } #[test] fn args_type_name_bool() { assert_eq!(args_type_name(&serde_json::json!(true)), "boolean"); } #[test] fn args_type_name_number() { assert_eq!(args_type_name(&serde_json::json!(42)), "number"); } #[test] fn args_type_name_string() { assert_eq!(args_type_name(&serde_json::json!("hello")), "string"); } #[test] fn args_type_name_array() { assert_eq!(args_type_name(&serde_json::json!([1, 2])), "array"); } #[test] fn args_type_name_object() { assert_eq!(args_type_name(&serde_json::json!({})), "object"); } }