ai_sdk_core/tool/

mod.rs

mod tool_output;

pub use tool_output::ToolOutput;

use crate::error::ToolError;
use ai_sdk_provider::language_model::{
    FunctionTool, Message, ToolCallPart, ToolResultOutput, ToolResultPart,
};
use ai_sdk_provider::JsonValue;
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;

/// Context information provided to tools during execution.
///
/// This structure contains metadata and conversation history that tools may need
/// during execution. It provides tools with access to the tool call identifier and
/// the full conversation context that led to the tool invocation.
pub struct ToolContext {
    /// Unique identifier for the tool call being executed.
    /// This ID is used to correlate tool results back to the original invocation.
    pub tool_call_id: String,
    /// Conversation messages leading up to this tool call.
    /// Includes all previous user, assistant, and tool result messages
    /// that provide context for the current execution.
    pub messages: Vec<Message>,
}

/// Trait that tools must implement to be available for language model invocation.
///
/// Tools are callable units of functionality that language models can invoke to
/// perform actions, retrieve information, or integrate with external systems.
/// Implementations must provide metadata (name, description, schema) and an
/// execution method that processes the tool input and returns results.
#[async_trait]
pub trait Tool: Send + Sync {
    /// Returns the name of the tool as recognized by the language model.
    ///
    /// This name must be unique among all available tools and is used by the
    /// language model to reference this specific tool when making invocations.
    /// Names should be alphanumeric with underscores and clearly describe the tool's purpose.
    fn name(&self) -> &str;

    /// Returns a human-readable description of what the tool does.
    ///
    /// This description is provided to the language model to help it understand
    /// the tool's purpose and when to use it. It should be clear, concise, and
    /// explain the primary function and use cases.
    fn description(&self) -> &str;

    /// Returns the JSON Schema that defines the structure of the tool's input parameters.
    ///
    /// The schema should be a valid JSON Schema document that describes all required
    /// and optional input parameters. This schema is used by the language model to
    /// understand what inputs the tool expects and by the runtime to validate inputs.
    fn input_schema(&self) -> Value;

    /// Executes the tool with the given input and context.
    ///
    /// This method is responsible for performing the actual work of the tool.
    /// It may return either a single value or a stream of preliminary results
    /// (useful for long-running operations that can produce incremental output).
    ///
    /// # Arguments
    ///
    /// * `input` - The parsed JSON input to the tool, must conform to the schema
    ///   returned by `input_schema()`.
    /// * `context` - Additional execution context including the tool call ID and
    ///   conversation history.
    ///
    /// # Returns
    ///
    /// * `Ok(ToolOutput)` - Either a single result value or a stream of preliminary results.
    /// * `Err(ToolError)` - An error that occurred during execution.
    async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolOutput, ToolError>;

    /// Determines whether this tool requires explicit user approval before execution.
    ///
    /// Implement this method to require approval for potentially sensitive operations.
    /// The runtime will check this method before invoking the tool and will deny
    /// execution if approval is required.
    ///
    /// # Arguments
    ///
    /// * `_input` - The input parameters that would be passed to the tool. Can be
    ///   examined to make approval decisions based on the specific operation.
    ///
    /// # Returns
    ///
    /// * `true` if approval is required before execution.
    /// * `false` if the tool can be executed automatically (default).
    fn needs_approval(&self, _input: &Value) -> bool {
        false
    }

    /// Converts the tool's raw output into a structured format for the language model.
    ///
    /// This method allows customization of how tool output is formatted and presented
    /// back to the language model. The default implementation converts strings to text
    /// output and other values to JSON output.
    ///
    /// # Arguments
    ///
    /// * `output` - The raw output value from `execute()`.
    ///
    /// # Returns
    ///
    /// A `ToolResultOutput` variant representing the formatted output that will be
    /// returned to the language model.
    fn to_model_output(&self, output: JsonValue) -> ToolResultOutput {
        // Default implementation
        match output {
            JsonValue::String(s) => ToolResultOutput::Text {
                value: s,
                provider_metadata: None,
            },
            other => ToolResultOutput::Json {
                value: other,
                provider_metadata: None,
            },
        }
    }
}

/// Manages the execution of tools invoked by language models.
///
/// The `ToolExecutor` is responsible for:
/// - Maintaining a registry of available tools
/// - Executing tool calls invoked by the language model
/// - Handling approval checks for sensitive operations
/// - Managing both single and streaming tool execution
/// - Converting tool results to structured output formats
///
/// It supports both parallel execution of multiple tool calls and sequential
/// streaming execution with preliminary result callbacks.
pub struct ToolExecutor {
    tools: Vec<Arc<dyn Tool>>,
}

impl ToolExecutor {
    /// Creates a new `ToolExecutor` with the provided set of tools.
    ///
    /// # Arguments
    ///
    /// * `tools` - A vector of tool implementations to be managed by this executor.
    ///   Tools are wrapped in `Arc` for thread-safe shared ownership.
    pub fn new(tools: Vec<Arc<dyn Tool>>) -> Self {
        Self { tools }
    }

    /// Retrieves the tool definitions in a format suitable for language model APIs.
    ///
    /// This converts the internal tool representations into `FunctionTool` definitions
    /// that can be provided to language models. These definitions include the tool's
    /// name, description, and input schema.
    ///
    /// # Returns
    ///
    /// A vector of `FunctionTool` definitions, one for each registered tool.
    pub fn tool_definitions(&self) -> Vec<FunctionTool> {
        self.tools
            .iter()
            .map(|tool| FunctionTool {
                name: tool.name().to_string(),
                description: Some(tool.description().to_string()),
                input_schema: tool.input_schema(),
                provider_options: None,
            })
            .collect()
    }

    /// Executes multiple tool calls in parallel and returns their results.
    ///
    /// This method processes all provided tool calls concurrently, allowing for
    /// efficient execution of multiple independent tools. Each tool call is handled
    /// independently, with its own error handling and result conversion.
    ///
    /// The execution includes:
    /// - Tool lookup by name
    /// - Input validation and parsing
    /// - Approval checks (tools may require user approval)
    /// - Actual tool execution
    /// - Result conversion to structured format
    ///
    /// # Arguments
    ///
    /// * `tool_calls` - A vector of tool calls to execute, each specifying the tool
    ///   name, call ID, and JSON-encoded input parameters.
    ///
    /// # Returns
    ///
    /// A vector of `ToolResultPart` structures, one for each input tool call.
    /// Results include either successful output or error information for each tool.
    pub async fn execute_tools(&self, tool_calls: Vec<ToolCallPart>) -> Vec<ToolResultPart> {
        let mut futures = Vec::new();

        for tool_call in tool_calls {
            let tool_opt = self.find_tool(&tool_call.tool_name);
            let tool_call_id = tool_call.tool_call_id.clone();
            let tool_name = tool_call.tool_name.clone();
            let input_str = tool_call.input.clone();

            let future = async move {
                // Handle tool not found
                let tool = match tool_opt {
                    Some(t) => t,
                    None => {
                        return ToolResultPart {
                            tool_call_id,
                            tool_name: tool_name.clone(),
                            output: ToolResultOutput::ErrorText {
                                value: format!("Tool '{}' not found", tool_name),
                                provider_metadata: None,
                            },
                            preliminary: None,
                            provider_metadata: None,
                        };
                    }
                };

                let context = ToolContext {
                    tool_call_id: tool_call_id.clone(),
                    messages: vec![], // TODO: pass actual messages
                };

                // Parse input
                let input: Value = match serde_json::from_str(&input_str) {
                    Ok(v) => v,
                    Err(e) => {
                        // Return error result instead of propagating
                        return ToolResultPart {
                            tool_call_id,
                            tool_name,
                            output: ToolResultOutput::ErrorText {
                                value: format!("Invalid input: {}", e),
                                provider_metadata: None,
                            },
                            preliminary: None,
                            provider_metadata: None,
                        };
                    }
                };

                // Check approval
                if tool.needs_approval(&input) {
                    // Return denial result
                    return ToolResultPart {
                        tool_call_id,
                        tool_name,
                        output: ToolResultOutput::ExecutionDenied {
                            reason: Some("Execution denied by user".to_string()),
                            provider_metadata: None,
                        },
                        preliminary: None,
                        provider_metadata: None,
                    };
                }

                // Execute tool and convert to structured output
                let output = match tool.execute(input, &context).await {
                    Ok(tool_output) => {
                        // Handle both Value and Stream outputs
                        match tool_output {
                            ToolOutput::Value(raw_output) => {
                                // Success - convert to structured output
                                tool.to_model_output(raw_output)
                            }
                            ToolOutput::Stream(mut stream) => {
                                // For non-streaming execute_tools, just get the last value
                                use futures::stream::StreamExt;
                                let mut last_output = None;
                                while let Some(item) = stream.next().await {
                                    match item {
                                        Ok(output) => last_output = Some(output),
                                        Err(e) => {
                                            return ToolResultPart {
                                                tool_call_id,
                                                tool_name,
                                                output: ToolResultOutput::ErrorText {
                                                    value: e.to_string(),
                                                    provider_metadata: None,
                                                },
                                                preliminary: None,
                                                provider_metadata: None,
                                            };
                                        }
                                    }
                                }
                                // Convert final output
                                let final_value = last_output.unwrap_or(JsonValue::Null);
                                tool.to_model_output(final_value)
                            }
                        }
                    }
                    Err(error) => {
                        // Execution error - return error output
                        ToolResultOutput::ErrorText {
                            value: error.to_string(),
                            provider_metadata: None,
                        }
                    }
                };

                ToolResultPart {
                    tool_call_id,
                    tool_name,
                    output,
                    preliminary: None,
                    provider_metadata: None,
                }
            };

            futures.push(future);
        }

        // Execute all tools in parallel
        futures::future::join_all(futures).await
    }

    /// Executes a single tool call with streaming support and preliminary result callbacks.
    ///
    /// This method is specialized for tools that produce streaming output. It handles
    /// the execution of a single tool call and invokes a callback for each preliminary
    /// result as they arrive. This is useful for long-running operations that can
    /// produce incremental output.
    ///
    /// The method:
    /// - Finds and validates the tool
    /// - Parses and validates input
    /// - Checks approval requirements
    /// - Executes the tool
    /// - For streaming results, invokes the callback for each intermediate value
    /// - Returns the final result after all streaming is complete
    ///
    /// # Arguments
    ///
    /// * `tool_call` - The tool call to execute, specifying the tool name, call ID,
    ///   and JSON-encoded input parameters.
    /// * `on_preliminary` - A callback function that is invoked for each preliminary
    ///   result as they stream in. Useful for real-time processing or UI updates.
    ///   The callback is not invoked for the final result.
    ///
    /// # Returns
    ///
    /// A `ToolResultPart` containing the final tool result, or an error if the
    /// tool execution or streaming fails. The final result does not have the
    /// `preliminary` flag set.
    pub async fn execute_tool_with_stream<F>(
        &self,
        tool_call: ToolCallPart,
        on_preliminary: F,
    ) -> ToolResultPart
    where
        F: Fn(ToolResultPart) + Send,
    {
        let tool_call_id = tool_call.tool_call_id.clone();
        let tool_name = tool_call.tool_name.clone();

        // Find tool
        let tool = match self.find_tool(&tool_call.tool_name) {
            Some(t) => t,
            None => {
                return ToolResultPart {
                    tool_call_id,
                    tool_name: tool_name.clone(),
                    output: ToolResultOutput::ErrorText {
                        value: format!("Tool '{}' not found", tool_name),
                        provider_metadata: None,
                    },
                    preliminary: None,
                    provider_metadata: None,
                };
            }
        };

        let context = ToolContext {
            tool_call_id: tool_call_id.clone(),
            messages: vec![], // TODO: pass actual messages
        };

        // Parse input
        let input: Value = match serde_json::from_str(&tool_call.input) {
            Ok(v) => v,
            Err(e) => {
                return ToolResultPart {
                    tool_call_id,
                    tool_name,
                    output: ToolResultOutput::ErrorText {
                        value: format!("Invalid input: {}", e),
                        provider_metadata: None,
                    },
                    preliminary: None,
                    provider_metadata: None,
                };
            }
        };

        // Check approval
        if tool.needs_approval(&input) {
            return ToolResultPart {
                tool_call_id,
                tool_name,
                output: ToolResultOutput::ExecutionDenied {
                    reason: Some("Execution denied by user".to_string()),
                    provider_metadata: None,
                },
                preliminary: None,
                provider_metadata: None,
            };
        }

        // Execute tool
        match tool.execute(input, &context).await {
            Ok(ToolOutput::Value(value)) => {
                // Simple case: single result
                let output = tool.to_model_output(value);
                ToolResultPart {
                    tool_call_id,
                    tool_name,
                    output,
                    preliminary: None,
                    provider_metadata: None,
                }
            }
            Ok(ToolOutput::Stream(mut stream)) => {
                use futures::stream::StreamExt;

                let mut last_output = None;

                // Process all stream items
                while let Some(item) = stream.next().await {
                    match item {
                        Ok(output) => {
                            // Emit preliminary result
                            let structured = tool.to_model_output(output.clone());
                            let preliminary_result = ToolResultPart {
                                tool_call_id: tool_call_id.clone(),
                                tool_name: tool_name.clone(),
                                output: structured,
                                preliminary: Some(true),
                                provider_metadata: None,
                            };

                            on_preliminary(preliminary_result);
                            last_output = Some(output);
                        }
                        Err(e) => {
                            return ToolResultPart {
                                tool_call_id,
                                tool_name,
                                output: ToolResultOutput::ErrorText {
                                    value: e.to_string(),
                                    provider_metadata: None,
                                },
                                preliminary: None,
                                provider_metadata: None,
                            };
                        }
                    }
                }

                // Return final result (last output without preliminary flag)
                let final_value = last_output.unwrap_or(JsonValue::Null);
                let final_output = tool.to_model_output(final_value);

                ToolResultPart {
                    tool_call_id,
                    tool_name,
                    output: final_output,
                    preliminary: None, // Final result
                    provider_metadata: None,
                }
            }
            Err(error) => ToolResultPart {
                tool_call_id,
                tool_name,
                output: ToolResultOutput::ErrorText {
                    value: error.to_string(),
                    provider_metadata: None,
                },
                preliminary: None,
                provider_metadata: None,
            },
        }
    }

    /// Searches for a tool by name in the executor's tool registry.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the tool to find. Must match the value returned
    ///   by the tool's `name()` method exactly.
    ///
    /// # Returns
    ///
    /// * `Some(Arc<dyn Tool>)` if a tool with the specified name is found.
    /// * `None` if no tool with the specified name is available.
    fn find_tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.tools.iter().find(|t| t.name() == name).cloned()
    }

    /// Returns a reference to the list of all available tools in this executor.
    ///
    /// This provides access to the complete registry of tools that can be invoked
    /// by the language model. Useful for introspection and tool enumeration.
    ///
    /// # Returns
    ///
    /// A slice of all registered tools, allowing iteration over the available tools.
    pub fn tools(&self) -> &[Arc<dyn Tool>] {
        &self.tools
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct TestTool {
        name: String,
        result: String,
    }

    #[async_trait]
    impl Tool for TestTool {
        fn name(&self) -> &str {
            &self.name
        }

        fn description(&self) -> &str {
            "A test tool"
        }

        fn input_schema(&self) -> Value {
            serde_json::json!({
                "type": "object",
                "properties": {}
            })
        }

        async fn execute(
            &self,
            _input: Value,
            _context: &ToolContext,
        ) -> Result<ToolOutput, ToolError> {
            Ok(ToolOutput::Value(JsonValue::String(self.result.clone())))
        }
    }

    #[tokio::test]
    async fn test_tool_executor_find_tool() {
        let tool = Arc::new(TestTool {
            name: "test".to_string(),
            result: "success".to_string(),
        });

        let executor = ToolExecutor::new(vec![tool]);
        assert!(executor.find_tool("test").is_some());
        assert!(executor.find_tool("nonexistent").is_none());
    }

    #[tokio::test]
    async fn test_tool_executor_execute() {
        let tool = Arc::new(TestTool {
            name: "test".to_string(),
            result: "success".to_string(),
        });

        let executor = ToolExecutor::new(vec![tool]);

        let tool_call = ToolCallPart {
            tool_call_id: "call_123".to_string(),
            tool_name: "test".to_string(),
            input: "{}".to_string(),
            provider_executed: None,
            dynamic: None,
            provider_metadata: None,
        };

        let results = executor.execute_tools(vec![tool_call]).await;
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].tool_call_id, "call_123");
        assert_eq!(results[0].tool_name, "test");

        // Check that output is Text variant
        match &results[0].output {
            ToolResultOutput::Text { value, .. } => {
                assert_eq!(value, "success");
            }
            _ => panic!("Expected Text output variant"),
        }
    }

    #[tokio::test]
    async fn test_tool_executor_tool_not_found() {
        let executor = ToolExecutor::new(vec![]);

        let tool_call = ToolCallPart {
            tool_call_id: "call_123".to_string(),
            tool_name: "nonexistent".to_string(),
            input: "{}".to_string(),
            provider_executed: None,
            dynamic: None,
            provider_metadata: None,
        };

        let results = executor.execute_tools(vec![tool_call]).await;
        assert_eq!(results.len(), 1);

        // Check that output is ErrorText variant
        match &results[0].output {
            ToolResultOutput::ErrorText { value, .. } => {
                assert!(value.contains("not found"));
            }
            _ => panic!("Expected ErrorText output variant"),
        }
    }
}