mcp_hub/

aggregator.rs

use anyhow::{Context, Result};
use rmcp::{
    model::{
        CallToolRequestParam, CallToolResult, ErrorCode, ErrorData, Implementation,
        InitializeRequestParam, InitializeResult, ListToolsResult, PaginatedRequestParam,
        ProtocolVersion, ServerCapabilities, Tool, ToolsCapability,
    },
    service::{Peer, RequestContext, RoleServer},
    Error as McpError, ServerHandler,
};
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot, Mutex};
use tracing::{debug, error, info, warn};

use crate::{
    config::Config,
    hub_actor::{spawn_hub_actor, HubMessage},
    metrics::{HubMetrics, MetricsTimer},
    supervisor::SupervisorConfig,
};

/// Main MCP aggregator that coordinates with the hub actor
#[derive(Debug, Clone)]
pub struct McpAggregator {
    /// Channel to communicate with the hub actor
    hub_sender: mpsc::Sender<HubMessage>,
    /// Server information
    server_info: Implementation,
    /// Instructions for the aggregator
    pub instructions: Option<String>,
    /// Peer connection (set during initialization)
    peer: Option<Peer<RoleServer>>,
    /// Shutdown coordination
    shutdown_tx: Arc<Mutex<Option<oneshot::Sender<()>>>>,
    /// Metrics collector
    metrics: HubMetrics,
}

impl McpAggregator {
    /// Create a new MCP aggregator
    pub fn new(config: Config, supervisor_config: SupervisorConfig) -> Self {
        let (server_info, instructions) = Self::create_server_info(&config);

        // Create shutdown coordination
        let (shutdown_tx, shutdown_rx) = oneshot::channel();
        let shutdown_tx = Arc::new(Mutex::new(Some(shutdown_tx)));

        // Spawn the hub actor
        let hub_sender = spawn_hub_actor(config, supervisor_config, shutdown_rx);

        Self {
            hub_sender,
            server_info,
            instructions,
            peer: None,
            shutdown_tx,
            metrics: HubMetrics::new(),
        }
    }

    /// Create server info from configuration
    fn create_server_info(config: &Config) -> (Implementation, Option<String>) {
        let server_names = config
            .servers()
            .iter()
            .map(|s| s.name.as_str())
            .collect::<Vec<_>>()
            .join(", ");

        let server_info = Implementation {
            name: "MCP Hub Aggregator".into(),
            version: env!("CARGO_PKG_VERSION").into(),
        };

        let instructions = Some(format!(
            "MCP Hub aggregating {} servers: {}",
            config.servers().len(),
            server_names
        ));

        (server_info, instructions)
    }

    /// Initialize the hub actor
    async fn initialize_hub(&self) -> Result<Vec<Tool>> {
        debug!("Initializing hub actor - entering method");
        debug!(
            "Hub sender channel capacity: {}",
            self.hub_sender.capacity()
        );

        let (response_tx, response_rx) = oneshot::channel();

        debug!("Sending Initialize message to hub actor");
        self.hub_sender
            .send(HubMessage::Initialize {
                response: response_tx,
            })
            .await
            .context("Failed to send initialize message to hub actor")?;

        debug!("Initialize message sent, waiting for response");

        let result = response_rx
            .await
            .context("Hub actor did not respond to initialize")?;

        debug!(
            "Received response from hub actor with {} tools",
            result.as_ref().map(|t| t.len()).unwrap_or(0)
        );
        result
    }

    /// Get aggregated tools from the hub
    async fn get_aggregated_tools(&self) -> Vec<Tool> {
        let (response_tx, response_rx) = oneshot::channel();

        if self
            .hub_sender
            .send(HubMessage::ListTools {
                response: response_tx,
            })
            .await
            .is_err()
        {
            warn!("Failed to send list_tools message to hub actor");
            return Vec::new();
        }

        response_rx.await.unwrap_or_default()
    }

    /// Route a tool call to the hub
    async fn route_tool_call(&self, params: CallToolRequestParam) -> Result<CallToolResult> {
        debug!("Routing tool call: {}", params.name);

        // Create metrics timer for this tool call
        let timer = MetricsTimer::new(
            self.metrics.clone(),
            "hub".to_string(), // We'll use "hub" as server name for aggregated metrics
            params.name.to_string(),
        );

        let (response_tx, response_rx) = oneshot::channel();

        let send_result = self
            .hub_sender
            .send(HubMessage::CallTool {
                params,
                response: response_tx,
            })
            .await;

        match send_result {
            Ok(_) => {
                let response_result = response_rx
                    .await
                    .context("Hub actor did not respond to call_tool");

                match response_result {
                    Ok(result) => {
                        timer.finish_with_status("success");
                        result
                    }
                    Err(e) => {
                        timer.finish_with_error("hub_timeout");
                        Err(e)
                    }
                }
            }
            Err(e) => {
                timer.finish_with_error("send_failed");
                Err(anyhow::anyhow!(
                    "Failed to send call_tool message to hub actor: {}",
                    e
                ))
            }
        }
    }

    /// Reload configuration in the hub
    #[allow(dead_code)]
    pub async fn reload_config(&self, config: Config) -> Result<()> {
        let (response_tx, response_rx) = oneshot::channel();

        self.hub_sender
            .send(HubMessage::ReloadConfig {
                config,
                response: response_tx,
            })
            .await
            .context("Failed to send reload_config message to hub actor")?;

        response_rx
            .await
            .context("Hub actor did not respond to reload_config")?
    }

    /// Get status of all supervisors
    #[allow(dead_code)]
    pub async fn get_status(
        &self,
    ) -> std::collections::HashMap<String, crate::supervisor::ProcessState> {
        let (response_tx, response_rx) = oneshot::channel();

        if self
            .hub_sender
            .send(HubMessage::GetStatus {
                response: response_tx,
            })
            .await
            .is_err()
        {
            warn!("Failed to send get_status message to hub actor");
            return std::collections::HashMap::new();
        }

        response_rx.await.unwrap_or_default()
    }

    /// Get a handle to communicate with the hub actor
    pub fn hub_handle(&self) -> mpsc::Sender<HubMessage> {
        self.hub_sender.clone()
    }
}

impl ServerHandler for McpAggregator {
    fn get_info(&self) -> InitializeResult {
        debug!("get_info() called - returning server capabilities");
        InitializeResult {
            protocol_version: ProtocolVersion::default(),
            capabilities: ServerCapabilities {
                experimental: None,
                logging: None,
                prompts: None,
                resources: None,
                tools: Some(ToolsCapability {
                    list_changed: Some(true),
                }),
            },
            server_info: self.server_info.clone(),
            instructions: self.instructions.clone(),
        }
    }

    fn get_peer(&self) -> Option<Peer<RoleServer>> {
        self.peer.clone()
    }

    fn set_peer(&mut self, peer: Peer<RoleServer>) {
        self.peer = Some(peer);
    }

    fn initialize(
        &self,
        _request: InitializeRequestParam,
        _context: RequestContext<RoleServer>,
    ) -> impl std::future::Future<Output = Result<InitializeResult, McpError>> + Send + '_ {
        async move {
            info!("Initializing MCP aggregator");
            debug!("About to call initialize_hub()");

            // Initialize the hub actor
            let _tools = self.initialize_hub().await.map_err(|e| {
                error!("Failed to initialize hub: {}", e);
                ErrorData {
                    code: ErrorCode::INTERNAL_ERROR,
                    message: format!("Failed to initialize hub: {}", e).into(),
                    data: None,
                }
            })?;

            debug!("Hub initialization completed, got {} tools", _tools.len());
            info!("Aggregator initialized successfully");

            let server_capabilities = ServerCapabilities {
                experimental: None,
                logging: None,
                prompts: None,
                resources: None,
                tools: Some(ToolsCapability {
                    list_changed: Some(true),
                }),
            };

            Ok(InitializeResult {
                protocol_version: ProtocolVersion::default(),
                capabilities: server_capabilities,
                server_info: self.server_info.clone(),
                instructions: self.instructions.clone(),
            })
        }
    }

    async fn list_tools(
        &self,
        _request: PaginatedRequestParam,
        _context: RequestContext<RoleServer>,
    ) -> Result<ListToolsResult, McpError> {
        debug!("Listing all aggregated tools");

        let tools = self.get_aggregated_tools().await;

        debug!("Returning {} tools", tools.len());

        Ok(ListToolsResult {
            tools,
            next_cursor: None,
        })
    }

    async fn call_tool(
        &self,
        request: CallToolRequestParam,
        _context: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, McpError> {
        info!("Calling tool: {}", request.name);

        self.route_tool_call(request).await.map_err(|e| {
            error!("Tool call failed: {}", e);
            ErrorData {
                code: ErrorCode::METHOD_NOT_FOUND,
                message: format!("Tool call failed: {}", e).into(),
                data: None,
            }
        })
    }

    async fn ping(&self, _context: RequestContext<RoleServer>) -> Result<(), McpError> {
        debug!("Ping received");
        Ok(())
    }

    async fn on_initialized(&self) {
        info!("MCP aggregator initialized and ready");
    }
}

impl Drop for McpAggregator {
    fn drop(&mut self) {
        // Trigger shutdown when aggregator is dropped
        if let Some(shutdown_tx) = self
            .shutdown_tx
            .try_lock()
            .ok()
            .and_then(|mut guard| guard.take())
        {
            let _ = shutdown_tx.send(());
        }
    }
}

/// Builder for creating an MCP aggregator
pub struct AggregatorBuilder {
    config: Option<Config>,
    supervisor_config: Option<SupervisorConfig>,
}

impl AggregatorBuilder {
    /// Create a new builder
    pub fn new() -> Self {
        Self {
            config: None,
            supervisor_config: None,
        }
    }

    /// Set configuration
    pub fn with_config(mut self, config: Config) -> Self {
        self.config = Some(config);
        self
    }

    /// Set supervisor configuration
    pub fn with_supervisor_config(mut self, supervisor_config: SupervisorConfig) -> Self {
        self.supervisor_config = Some(supervisor_config);
        self
    }

    /// Build the aggregator
    pub fn build(self) -> Result<McpAggregator> {
        let config = self
            .config
            .ok_or_else(|| anyhow::anyhow!("Configuration is required"))?;

        let supervisor_config = self.supervisor_config.unwrap_or_default();

        Ok(McpAggregator::new(config, supervisor_config))
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::ServerConfig;
    use std::collections::HashMap;

    fn create_test_config() -> Config {
        let server_configs = vec![
            ServerConfig {
                name: "test1".to_string(),
                cmd: "echo".to_string(),
                args: vec!["hello".to_string()],
                env: HashMap::new(),
                whitelist: None,
                blacklist: None,
                prefix: Some("test1_".to_string()),
                description_suffix: Some(" (test server 1)".to_string()),
                description_prefix: None,
            },
            ServerConfig {
                name: "test2".to_string(),
                cmd: "echo".to_string(),
                args: vec!["world".to_string()],
                env: HashMap::new(),
                whitelist: None,
                blacklist: None,
                prefix: Some("test2_".to_string()),
                description_suffix: Some(" (test server 2)".to_string()),
                description_prefix: None,
            },
        ];

        Config {
            server: server_configs,
        }
    }

    #[tokio::test]
    async fn test_aggregator_creation() {
        let config = create_test_config();
        let aggregator = McpAggregator::new(config, SupervisorConfig::default());

        let info = aggregator.get_info();
        assert_eq!(info.server_info.name, "MCP Hub Aggregator");
        assert!(aggregator
            .instructions
            .as_ref()
            .unwrap()
            .contains("test1, test2"));
    }

    #[tokio::test]
    async fn test_builder_pattern() -> Result<()> {
        let config = create_test_config();
        let aggregator = AggregatorBuilder::new()
            .with_config(config)
            .with_supervisor_config(SupervisorConfig::default())
            .build()?;

        assert_eq!(aggregator.get_info().server_info.name, "MCP Hub Aggregator");
        Ok(())
    }

    #[tokio::test]
    async fn test_builder_requires_config() {
        let result = AggregatorBuilder::new().build();
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Configuration is required"));
    }
}