funera-core 0.3.0

#![cfg(feature = "tool")]

use std::{collections::HashMap, fmt::Display, sync::Arc};

use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;
use thiserror::Error;

/// The type of a tool, as communicated to the LLM.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ToolType {
    /// Standard OpenAI-compatible function tool.
    Function,
}

impl Display for ToolType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "function")
    }
}

/// A callable tool exposed to the LLM agent.
///
/// Implement this trait to define custom tools. The framework will expose
/// the tool's [`schema`](Tool::schema) to the LLM and invoke
/// [`execute`](Tool::execute) when the LLM requests it.
#[async_trait]
pub trait Tool: Send + Sync {
    /// Unique name for this tool (e.g. `"read"`, `"shell"`).
    fn name(&self) -> &str;

    /// Human-readable description sent to the LLM.
    fn description(&self) -> &str;

    /// Execute the tool with the given JSON arguments.
    ///
    /// Returns a string result on success, or a [`ToolCallError`] on failure.
    async fn execute(&self, args: JsonValue) -> Result<String, ToolCallError>;

    /// Returns the tool type (defaults to [`ToolType::Function`]).
    fn get_type(&self) -> ToolType {
        ToolType::Function
    }

    /// Returns the JSON schema describing this tool's parameters.
    ///
    /// This is sent to the LLM so it can generate well-formed invocations.
    fn schema(&self) -> JsonValue;
}

/// Errors that can occur during tool execution.
#[derive(Debug, Error)]
pub enum ToolCallError {
    /// The arguments did not match the expected schema.
    #[error("parameter mismatch: {0}")]
    ParameterMismatch(JsonValue),

    /// The tool encountered a runtime error during execution.
    #[error("tool execution error: {0}")]
    ToolExecutionError(#[from] anyhow::Error),

    /// The tool exists but is currently unavailable (e.g. disabled by policy).
    #[error("tool unavailable: {0}")]
    ToolUnavailable(String),

    /// No tool with the given name is registered.
    #[error("tool not found: {0}")]
    ToolNotFound(String),

    /// The tool call requires user approval before proceeding.
    #[error("approval required for {tool_name}: {reason}")]
    ApprovalRequired {
        call_id: String,
        tool_name: String,
        reason: String,
    },

    /// The tool call was rejected after an approval request was denied.
    #[error("tool call rejected: {reason}")]
    Rejected { reason: String },
}

/// An entry in the tool registry, pairing a tool with its availability status.
#[derive(Clone)]
pub struct ToolRegistryEntry {
    pub tool: Arc<dyn Tool>,
    pub available: bool,
}
impl ToolRegistryEntry {
    /// Create a new registry entry with explicit availability.
    pub fn new(tool: Arc<dyn Tool>, available: bool) -> Self {
        Self { tool, available }
    }

    /// Whether the tool is currently available for execution.
    pub fn is_available(&self) -> bool {
        self.available
    }

    /// Create a new registry entry with the tool available.
    pub fn new_available(tool: Arc<dyn Tool>) -> Self {
        Self::new(tool, true)
    }

    /// Create a new registry entry with the tool unavailable.
    pub fn new_unavailable(tool: Arc<dyn Tool>) -> Self {
        Self::new(tool, false)
    }
}

/// Raw tool registry (no security checks).
///
/// When the `security` feature is enabled, [`ToolRegistry`] aliases to
/// [`GuardedToolRegistry`](crate::security::registry::GuardedToolRegistry)
/// instead, which wraps this registry with policy checks and audit logging.
#[doc(hidden)]
#[derive(Clone)]
pub struct RawToolRegistry {
    tools: HashMap<String, ToolRegistryEntry>,
}

impl Default for RawToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl RawToolRegistry {
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    pub fn add_tool(&mut self, tool: Arc<dyn Tool>) {
        self.tools.insert(
            tool.name().to_string(),
            ToolRegistryEntry::new_available(tool),
        );
    }
    pub fn get_tool(&self, name: &str) -> Option<&ToolRegistryEntry> {
        self.tools.get(name)
    }

    /// Clone the tool's `Arc` if it exists and is available.
    ///
    /// Used by the executor to run a tool outside the registry lock.
    pub fn get_tool_arc(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.get_tool(name)
            .filter(|entry| entry.is_available())
            .map(|entry| entry.tool.clone())
    }
    pub fn remove_tool(&mut self, name: &str) {
        self.tools.remove(name);
    }

    /// Remove the tool only if the registered entry is the same `Arc` value.
    /// This prevents a stale disposer from deleting a replacement tool that
    /// reuses the same name (e.g. HMR replacement).
    pub fn remove_tool_if_same(&mut self, name: &str, tool: &Arc<dyn Tool>) -> bool {
        match self.tools.get(name) {
            Some(entry) if Arc::ptr_eq(&entry.tool, tool) => {
                self.tools.remove(name);
                true
            }
            _ => false,
        }
    }
    pub fn tool_exists(&self, name: &str) -> bool {
        self.tools.contains_key(name)
    }
    pub fn tool_count(&self) -> usize {
        self.tools.len()
    }
    pub fn get_all_tools(&self) -> &HashMap<String, ToolRegistryEntry> {
        &self.tools
    }
    pub fn available_tools_json(&self) -> JsonValue {
        self.tools
            .values()
            .filter_map(|tool| {
                if tool.is_available() {
                    Some(tool.tool.schema())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
            .into()
    }
    pub async fn call_tool(&self, name: &str, args: JsonValue) -> Result<String, ToolCallError> {
        if let Some(tool) = self.get_tool(name) {
            if tool.is_available() {
                tool.tool.execute(args).await
            } else {
                Err(ToolCallError::ToolUnavailable(name.to_string()))
            }
        } else {
            Err(ToolCallError::ToolNotFound(name.to_string()))
        }
    }
}

/// The active tool registry type.
///
/// When the `security` feature is enabled, this aliases to
/// [`GuardedToolRegistry`](crate::security::registry::GuardedToolRegistry)
/// which enforces tool policies and logs audit events on every tool call.
/// Without `security`, it is the raw registry with no policy checks.
#[cfg(feature = "security")]
pub use crate::security::registry::GuardedToolRegistry as ToolRegistry;

#[cfg(not(feature = "security"))]
pub use RawToolRegistry as ToolRegistry;

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

    #[test]
    fn rejected_error_display() {
        let e = ToolCallError::Rejected {
            reason: "access denied".into(),
        };
        let msg = format!("{e}");
        assert!(msg.contains("access denied"), "msg: {msg}");
    }

    #[test]
    fn approval_required_error_display() {
        let e = ToolCallError::ApprovalRequired {
            call_id: "c1".into(),
            tool_name: "shell".into(),
            reason: "needs approval".into(),
        };
        let msg = format!("{e}");
        assert!(msg.contains("shell"), "msg: {msg}");
        assert!(msg.contains("approval"), "msg: {msg}");
    }

    struct MockTool;
    #[async_trait]
    impl Tool for MockTool {
        fn name(&self) -> &str {
            "mock"
        }
        fn description(&self) -> &str {
            "mock tool"
        }
        fn schema(&self) -> JsonValue {
            json!({"type": "function", "function": {"name": "mock"}})
        }
        async fn execute(&self, _args: JsonValue) -> Result<String, ToolCallError> {
            Ok("done".into())
        }
    }

    #[test]
    fn get_tool_arc_returns_available_tool_or_none() {
        let mut reg = RawToolRegistry::new();
        assert!(reg.get_tool_arc("mock").is_none());

        reg.add_tool(Arc::new(MockTool));
        let tool = reg.get_tool_arc("mock");
        assert!(
            tool.is_some(),
            "registered tool must be clonable via get_tool_arc"
        );
        assert_eq!(tool.unwrap().name(), "mock");

        assert!(reg.get_tool_arc("missing").is_none());
    }

    #[test]
    fn remove_tool_if_same_only_removes_matching_arc() {
        let mut reg = RawToolRegistry::new();
        let original: Arc<dyn Tool> = Arc::new(MockTool);
        reg.add_tool(Arc::clone(&original));

        // A different Arc (e.g. a replacement with the same name) is kept.
        let other: Arc<dyn Tool> = Arc::new(MockTool);
        assert!(!reg.remove_tool_if_same("mock", &other));
        assert!(reg.tool_exists("mock"));

        // Removing the registered Arc succeeds.
        assert!(reg.remove_tool_if_same("mock", &original));
        assert!(!reg.tool_exists("mock"));
    }
}