cortexai-mcp 0.1.0

//! Cortex MCP Server
//!
//! High-level builder that wires AgentEngine, Crew, and ToolRegistry
//! into a single MCP server with both HTTP and SSE transports.

#[cfg(feature = "engine")]
use std::sync::Arc;

#[cfg(feature = "engine")]
use cortexai_agents::AgentEngine;
#[cfg(feature = "engine")]
use cortexai_core::tool::ToolRegistry;
#[cfg(feature = "engine")]
use cortexai_providers::LLMBackend;

#[cfg(feature = "engine")]
use crate::crew_engine_handler::CrewEngineHandler;
#[cfg(feature = "engine")]
use crate::engine_handler::EngineHandler;
#[cfg(feature = "engine")]
use crate::server::McpServer;
#[cfg(feature = "engine")]
use crate::tool_proxy_handler::ToolProxyHandler;

/// A Cortex-aware MCP server that exposes AgentEngine, Crew, and
/// ToolRegistry as MCP tools.
///
/// Use the builder pattern to configure which capabilities are exposed:
///
/// ```rust,ignore
/// let server = CortexMcpServer::builder()
///     .engine(engine)
///     .backend(backend)
///     .tool_registry(registry)
///     .enable_agent_tool(true)
///     .enable_crew_tool(true)
///     .enable_tool_proxy(true)
///     .build();
/// ```
#[cfg(feature = "engine")]
pub struct CortexMcpServer {
    inner: Arc<McpServer>,
}

#[cfg(feature = "engine")]
impl CortexMcpServer {
    /// Create a new builder.
    pub fn builder() -> CortexMcpServerBuilder {
        CortexMcpServerBuilder::new()
    }

    /// Get the underlying MCP server.
    pub fn inner(&self) -> &Arc<McpServer> {
        &self.inner
    }

    /// Get the number of registered tools.
    pub fn tool_count(&self) -> usize {
        self.inner.tool_count()
    }
}

/// Builder for `CortexMcpServer`.
#[cfg(feature = "engine")]
pub struct CortexMcpServerBuilder {
    name: String,
    version: String,
    engine: Option<Arc<AgentEngine>>,
    backend: Option<Arc<dyn LLMBackend>>,
    tool_registry: Option<Arc<ToolRegistry>>,
    enable_agent: bool,
    enable_crew: bool,
    enable_proxy: bool,
}

#[cfg(feature = "engine")]
impl CortexMcpServerBuilder {
    fn new() -> Self {
        Self {
            name: "cortex-mcp-server".to_string(),
            version: env!("CARGO_PKG_VERSION").to_string(),
            engine: None,
            backend: None,
            tool_registry: None,
            enable_agent: false,
            enable_crew: false,
            enable_proxy: false,
        }
    }

    /// Set the server name.
    pub fn name(mut self, name: impl Into<String>) -> Self {
        self.name = name.into();
        self
    }

    /// Set the server version.
    pub fn version(mut self, version: impl Into<String>) -> Self {
        self.version = version.into();
        self
    }

    /// Set the AgentEngine.
    pub fn engine(mut self, engine: Arc<AgentEngine>) -> Self {
        self.engine = Some(engine);
        self
    }

    /// Set the LLM backend.
    pub fn backend(mut self, backend: Arc<dyn LLMBackend>) -> Self {
        self.backend = Some(backend);
        self
    }

    /// Set the ToolRegistry.
    pub fn tool_registry(mut self, registry: Arc<ToolRegistry>) -> Self {
        self.tool_registry = Some(registry);
        self
    }

    /// Enable/disable the `run_agent` tool.
    pub fn enable_agent_tool(mut self, enable: bool) -> Self {
        self.enable_agent = enable;
        self
    }

    /// Enable/disable the `run_crew` tool.
    pub fn enable_crew_tool(mut self, enable: bool) -> Self {
        self.enable_crew = enable;
        self
    }

    /// Enable/disable proxying all registry tools as MCP tools.
    pub fn enable_tool_proxy(mut self, enable: bool) -> Self {
        self.enable_proxy = enable;
        self
    }

    /// Build the server. Panics if required dependencies are missing.
    pub fn build(self) -> CortexMcpServer {
        let engine = self.engine.expect("engine is required");
        let backend = self.backend.expect("backend is required");
        let registry = self
            .tool_registry
            .unwrap_or_else(|| Arc::new(ToolRegistry::new()));

        let mut builder = McpServer::builder()
            .name(&self.name)
            .version(&self.version);

        if self.enable_agent {
            let handler =
                EngineHandler::new(engine.clone(), backend.clone(), registry.clone());
            builder = builder.add_tool(handler);
        }

        if self.enable_crew {
            let handler =
                CrewEngineHandler::new(engine.clone(), backend.clone(), registry.clone());
            builder = builder.add_tool(handler);
        }

        if self.enable_proxy {
            let proxy = ToolProxyHandler::new(registry.clone());
            for def in proxy.definitions() {
                // Create a thin wrapper that delegates to the proxy
                let proxy_clone = ToolProxyHandler::new(registry.clone());
                let tool_name = def.name.clone();
                let tool_desc = def.description.clone();
                let tool_schema = def.input_schema.clone();

                let wrapper = ProxiedTool {
                    name: tool_name,
                    description: tool_desc,
                    schema: tool_schema,
                    proxy: proxy_clone,
                };
                builder = builder.add_tool(wrapper);
            }
        }

        CortexMcpServer {
            inner: builder.build(),
        }
    }
}

/// A thin ToolHandler wrapper around a single proxied tool.
#[cfg(feature = "engine")]
struct ProxiedTool {
    name: String,
    description: Option<String>,
    schema: serde_json::Value,
    proxy: ToolProxyHandler,
}

#[cfg(feature = "engine")]
#[async_trait::async_trait]
impl crate::server::ToolHandler for ProxiedTool {
    fn definition(&self) -> crate::protocol::McpTool {
        crate::protocol::McpTool {
            name: self.name.clone(),
            description: self.description.clone(),
            input_schema: self.schema.clone(),
        }
    }

    async fn execute(
        &self,
        arguments: serde_json::Value,
    ) -> Result<crate::protocol::CallToolResult, crate::error::McpError> {
        self.proxy.call(&self.name, arguments).await
    }
}

#[cfg(all(test, feature = "engine"))]
mod tests {
    use std::sync::Arc;

    use cortexai_agents::AgentEngine;
    use cortexai_core::tool::ToolRegistry;
    use cortexai_providers::{LLMBackend, MockBackend, MockResponse};

    use super::*;

    #[tokio::test]
    async fn test_cortex_server_builder_all_features() {
        let engine = Arc::new(AgentEngine::new());
        let backend: Arc<dyn LLMBackend> =
            Arc::new(MockBackend::new().with_response(MockResponse::text("ok")));

        let mut registry = ToolRegistry::new();

        use async_trait::async_trait;
        use cortexai_core::errors::ToolError;
        use cortexai_core::tool::{ExecutionContext, Tool, ToolSchema};

        struct DummyTool;
        #[async_trait]
        impl Tool for DummyTool {
            fn schema(&self) -> ToolSchema {
                ToolSchema::new("dummy", "A dummy tool")
            }
            async fn execute(
                &self,
                _ctx: &ExecutionContext,
                _args: serde_json::Value,
            ) -> Result<serde_json::Value, ToolError> {
                Ok(serde_json::json!("ok"))
            }
        }
        registry.register(Arc::new(DummyTool));
        let registry = Arc::new(registry);

        let server = CortexMcpServer::builder()
            .engine(engine)
            .backend(backend)
            .tool_registry(registry)
            .enable_agent_tool(true)
            .enable_crew_tool(true)
            .enable_tool_proxy(true)
            .build();

        // run_agent + run_crew + dummy = 3 tools
        assert_eq!(server.tool_count(), 3);
    }

    #[tokio::test]
    async fn test_cortex_server_builder_agent_only() {
        let engine = Arc::new(AgentEngine::new());
        let backend: Arc<dyn LLMBackend> =
            Arc::new(MockBackend::new().with_response(MockResponse::text("ok")));
        let registry = Arc::new(ToolRegistry::new());

        let server = CortexMcpServer::builder()
            .engine(engine)
            .backend(backend)
            .tool_registry(registry)
            .enable_agent_tool(true)
            .enable_crew_tool(false)
            .enable_tool_proxy(false)
            .build();

        assert_eq!(server.tool_count(), 1);
    }
}