PicoBot/src/tools/registry.rs

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use crate::providers::{Tool, ToolFunction};

use super::traits::DelegationPolicy;
use super::traits::Tool as ToolTrait;

pub struct ToolRegistry {
    tools: Mutex<HashMap<String, Arc<dyn ToolTrait>>>,
}

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

    pub fn register<T: ToolTrait + 'static>(&self, tool: T) {
        self.tools
            .lock()
            .unwrap()
            .insert(tool.name().to_string(), Arc::new(tool));
    }

    /// Register an existing Arc-wrapped tool by name
    pub fn register_raw(&self, name: String, tool: Arc<dyn ToolTrait>) {
        self.tools.lock().unwrap().insert(name, tool);
    }

    pub fn get(&self, name: &str) -> Option<Arc<dyn ToolTrait>> {
        self.tools.lock().unwrap().get(name).cloned()
    }

    /// Get all registered tools.
    /// Used for concurrent tool execution when we need to look up tools by name.
    pub fn get_all(&self) -> Vec<Arc<dyn ToolTrait>> {
        self.tools.lock().unwrap().values().cloned().collect()
    }

    pub fn get_definitions(&self) -> Vec<Tool> {
        let mut defs: Vec<Tool> = self
            .tools
            .lock()
            .unwrap()
            .values()
            .map(|tool| Tool {
                tool_type: "function".to_string(),
                function: ToolFunction {
                    name: tool.name().to_string(),
                    description: tool.description().to_string(),
                    parameters: tool.parameters_schema(),
                },
            })
            .collect();

        defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));
        defs
    }

    pub fn has_tools(&self) -> bool {
        !self.tools.lock().unwrap().is_empty()
    }

    pub fn tool_names(&self) -> Vec<String> {
        self.tools.lock().unwrap().keys().cloned().collect()
    }

    pub fn iter(&self) -> Vec<(String, Arc<dyn ToolTrait>)> {
        self.tools
            .lock()
            .unwrap()
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect()
    }

    /// Clone this registry while excluding tools that a constrained execution
    /// path must not call (for example scheduler-managed delivery).
    pub fn without(&self, excluded: &[&str]) -> Arc<Self> {
        let filtered = Self::new();
        for (name, tool) in self.iter() {
            if !excluded.contains(&name.as_str()) {
                filtered.register_raw(name, tool);
            }
        }
        Arc::new(filtered)
    }

    /// Build an immutable execution view for a named Agent. Definition tools
    /// must be explicitly delegatable; runtime tools are passed separately and
    /// can only carry the RuntimeInjected policy.
    pub fn scoped_for_agent(
        &self,
        tool_names: &[String],
        runtime_tools: Vec<Arc<dyn ToolTrait>>,
    ) -> Result<Arc<Self>, String> {
        let scoped = Self::new();
        for name in tool_names {
            let tool = self
                .get(name)
                .ok_or_else(|| format!("tool '{name}' is not registered"))?;
            if tool.delegation_policy() != DelegationPolicy::Delegatable {
                return Err(format!("tool '{name}' is not delegatable"));
            }
            scoped.register_raw(name.clone(), tool);
        }
        for tool in runtime_tools {
            if tool.delegation_policy() != DelegationPolicy::RuntimeInjected {
                return Err(format!(
                    "runtime tool '{}' is missing RuntimeInjected policy",
                    tool.name()
                ));
            }
            scoped.register_raw(tool.name().to_string(), tool);
        }
        Ok(Arc::new(scoped))
    }

    /// 生成工具列表描述，用于子 Agent 系统提示词
    pub fn describe_for_prompt(&self) -> String {
        let mut entries: Vec<String> = self
            .iter()
            .into_iter()
            .map(|(name, tool)| format!("- {}: {}", name, tool.description()))
            .collect();
        entries.sort();
        entries.join("\n")
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tools::traits::ToolResult;
    use async_trait::async_trait;
    use serde_json::json;

    struct TestTool(&'static str);

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

        fn description(&self) -> &str {
            self.0
        }

        fn parameters_schema(&self) -> serde_json::Value {
            json!({})
        }

        async fn execute(&self, _args: serde_json::Value) -> anyhow::Result<ToolResult> {
            Ok(ToolResult {
                success: true,
                output: "ok".to_string(),
                error: None,
            })
        }
    }

    #[test]
    fn test_get_definitions_sorted_by_name() {
        let registry = ToolRegistry::new();
        registry.register(TestTool("zeta"));
        registry.register(TestTool("alpha"));
        registry.register(TestTool("beta"));

        let defs = registry.get_definitions();
        let names: Vec<_> = defs.into_iter().map(|tool| tool.function.name).collect();

        assert_eq!(names, vec!["alpha", "beta", "zeta"]);
    }
}