owlen/crates/owlen-core/src/tools/registry.rs

use std::collections::HashMap;
use std::sync::Arc;

use crate::{Error, Result};
use anyhow::Context;
use serde_json::Value;

use super::{
    Tool, ToolResult, WEB_SEARCH_TOOL_NAME, canonical_tool_name, tool_identifier_violation,
};
use crate::config::Config;
use crate::mode::Mode;
use crate::ui::UiController;

pub struct ToolRegistry {
    tools: HashMap<String, Arc<dyn Tool>>,
    config: Arc<tokio::sync::Mutex<Config>>,
    ui: Arc<dyn UiController>,
}

impl ToolRegistry {
    pub fn new(config: Arc<tokio::sync::Mutex<Config>>, ui: Arc<dyn UiController>) -> Self {
        Self {
            tools: HashMap::new(),
            config,
            ui,
        }
    }

    pub fn register<T>(&mut self, tool: T) -> Result<()>
    where
        T: Tool + 'static,
    {
        let tool: Arc<dyn Tool> = Arc::new(tool);
        let name = tool.name();

        if let Some(reason) = tool_identifier_violation(name) {
            log::error!("Tool '{}' failed validation: {}", name, reason);
            return Err(Error::InvalidInput(format!(
                "Tool '{name}' is not a valid MCP identifier: {reason}"
            )));
        }

        if self
            .tools
            .insert(name.to_string(), Arc::clone(&tool))
            .is_some()
        {
            log::warn!(
                "Tool '{}' was already registered; overwriting previous entry.",
                name
            );
        }

        Ok(())
    }

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

    pub fn all(&self) -> Vec<Arc<dyn Tool>> {
        self.tools.values().cloned().collect()
    }

    pub async fn execute(&self, name: &str, args: Value, mode: Mode) -> Result<ToolResult> {
        let canonical = canonical_tool_name(name);
        let tool = self
            .get(canonical)
            .with_context(|| format!("Tool not registered: {}", name))?;

        let mut config = self.config.lock().await;

        // Check mode-based tool availability first
        if !(config.modes.is_tool_allowed(mode, canonical)
            || config.modes.is_tool_allowed(mode, name))
        {
            let alternate_mode = match mode {
                Mode::Chat => Mode::Code,
                Mode::Code => Mode::Chat,
            };

            if config.modes.is_tool_allowed(alternate_mode, canonical)
                || config.modes.is_tool_allowed(alternate_mode, name)
            {
                return Ok(ToolResult::error(&format!(
                    "Tool '{}' is not available in {} mode. Switch to {} mode to use this tool (use :mode {} command).",
                    name, mode, alternate_mode, alternate_mode
                )));
            } else {
                return Ok(ToolResult::error(&format!(
                    "Tool '{}' is not available in any mode. Check your configuration.",
                    name
                )));
            }
        }

        let is_enabled = match canonical {
            WEB_SEARCH_TOOL_NAME => config.tools.web_search.enabled,
            "code_exec" => config.tools.code_exec.enabled,
            _ => true, // All other tools are considered enabled by default
        };

        if !is_enabled {
            let prompt = format!(
                "Tool '{}' is disabled. Would you like to enable it for this session?",
                name
            );
            if self.ui.confirm(&prompt).await {
                // Enable the tool in the in-memory config for the current session
                match canonical {
                    WEB_SEARCH_TOOL_NAME => config.tools.web_search.enabled = true,
                    "code_exec" => config.tools.code_exec.enabled = true,
                    _ => {}
                }
            } else {
                return Ok(ToolResult::cancelled(&format!(
                    "Tool '{}' execution was cancelled by the user.",
                    name
                )));
            }
        }

        tool.execute(args).await
    }

    /// Get all tools available in the given mode
    pub async fn available_tools(&self, mode: Mode) -> Vec<String> {
        let config = self.config.lock().await;
        self.tools
            .keys()
            .filter(|name| config.modes.is_tool_allowed(mode, name))
            .cloned()
            .collect()
    }

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Config;
    use crate::tools::{Tool, ToolResult, WEB_SEARCH_TOOL_NAME};
    use crate::ui::NoOpUiController;
    use async_trait::async_trait;
    use serde_json::{Value, json};
    use std::sync::Arc;

    struct DummyTool {
        name: &'static str,
    }

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

        fn description(&self) -> &'static str {
            "dummy tool"
        }

        fn schema(&self) -> Value {
            json!({ "type": "object" })
        }

        fn aliases(&self) -> &'static [&'static str] {
            &[]
        }

        async fn execute(&self, _args: Value) -> Result<ToolResult> {
            Ok(ToolResult::success(json!({ "echo": true })))
        }
    }

    fn registry() -> ToolRegistry {
        let config = Arc::new(tokio::sync::Mutex::new(Config::default()));
        let ui = Arc::new(NoOpUiController);
        ToolRegistry::new(config, ui)
    }

    #[test]
    fn rejects_invalid_tool_identifier() {
        let mut registry = registry();
        let tool = DummyTool {
            name: "invalid.tool",
        };

        let err = registry.register(tool).unwrap_err();
        assert!(matches!(err, Error::InvalidInput(_)));
    }

    #[test]
    fn registers_spec_compliant_tool() {
        let mut registry = registry();
        let tool = DummyTool {
            name: WEB_SEARCH_TOOL_NAME,
        };

        registry.register(tool).unwrap();
        assert!(registry.get(WEB_SEARCH_TOOL_NAME).is_some());
    }
}