owlen/crates/owlen-core/src/mcp/factory.rs

/// MCP Client Factory
///
/// Provides a unified interface for creating MCP clients based on configuration.
/// Supports switching between local (in-process) and remote (STDIO) execution modes.
use super::client::McpClient;
use super::{
    LocalMcpClient,
    remote_client::{McpRuntimeSecrets, RemoteMcpClient},
};
use crate::config::{Config, McpMode};
use crate::tools::registry::ToolRegistry;
use crate::validation::SchemaValidator;
use crate::{Error, Result};
use log::{info, warn};
use std::sync::Arc;

/// Factory for creating MCP clients based on configuration
pub struct McpClientFactory {
    config: Arc<Config>,
    registry: Arc<ToolRegistry>,
    validator: Arc<SchemaValidator>,
}

impl McpClientFactory {
    pub fn new(
        config: Arc<Config>,
        registry: Arc<ToolRegistry>,
        validator: Arc<SchemaValidator>,
    ) -> Self {
        Self {
            config,
            registry,
            validator,
        }
    }

    /// Create an MCP client based on the current configuration.
    pub fn create(&self) -> Result<Box<dyn McpClient>> {
        self.create_with_secrets(None)
    }

    /// Create an MCP client using optional runtime secrets (OAuth tokens, env overrides).
    pub fn create_with_secrets(
        &self,
        runtime: Option<McpRuntimeSecrets>,
    ) -> Result<Box<dyn McpClient>> {
        match self.config.mcp.mode {
            McpMode::Disabled => Err(Error::Config(
                "MCP mode is set to 'disabled'; tooling cannot function in this configuration."
                    .to_string(),
            )),
            McpMode::LocalOnly | McpMode::Legacy => {
                if matches!(self.config.mcp.mode, McpMode::Legacy) {
                    warn!("Using deprecated MCP legacy mode; consider switching to 'local_only'.");
                }
                Ok(Box::new(LocalMcpClient::new(
                    self.registry.clone(),
                    self.validator.clone(),
                )))
            }
            McpMode::RemoteOnly => {
                let server_cfg = self.config.effective_mcp_servers().first().ok_or_else(|| {
                    Error::Config(
                        "MCP mode 'remote_only' requires at least one entry in [[mcp_servers]]"
                            .to_string(),
                    )
                })?;

                RemoteMcpClient::new_with_runtime(server_cfg, runtime)
                    .map(|client| Box::new(client) as Box<dyn McpClient>)
                    .map_err(|e| {
                        Error::Config(format!(
                            "Failed to start remote MCP client '{}': {e}",
                            server_cfg.name
                        ))
                    })
            }
            McpMode::RemotePreferred => {
                if let Some(server_cfg) = self.config.effective_mcp_servers().first() {
                    match RemoteMcpClient::new_with_runtime(server_cfg, runtime.clone()) {
                        Ok(client) => {
                            info!(
                                "Connected to remote MCP server '{}' via {} transport.",
                                server_cfg.name, server_cfg.transport
                            );
                            Ok(Box::new(client) as Box<dyn McpClient>)
                        }
                        Err(e) if self.config.mcp.allow_fallback => {
                            warn!(
                                "Failed to start remote MCP client '{}': {}. Falling back to local tooling.",
                                server_cfg.name, e
                            );
                            Ok(Box::new(LocalMcpClient::new(
                                self.registry.clone(),
                                self.validator.clone(),
                            )))
                        }
                        Err(e) => Err(Error::Config(format!(
                            "Failed to start remote MCP client '{}': {e}. To allow fallback, set [mcp].allow_fallback = true.",
                            server_cfg.name
                        ))),
                    }
                } else {
                    warn!("No MCP servers configured; using local MCP tooling.");
                    Ok(Box::new(LocalMcpClient::new(
                        self.registry.clone(),
                        self.validator.clone(),
                    )))
                }
            }
        }
    }

    /// Check if remote MCP mode is available
    pub fn is_remote_available() -> bool {
        RemoteMcpClient::new().is_ok()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Error;
    use crate::config::McpServerConfig;

    fn build_factory(config: Config) -> McpClientFactory {
        let ui = Arc::new(crate::ui::NoOpUiController);
        let registry = Arc::new(ToolRegistry::new(
            Arc::new(tokio::sync::Mutex::new(config.clone())),
            ui,
        ));
        let validator = Arc::new(SchemaValidator::new());

        McpClientFactory::new(Arc::new(config), registry, validator)
    }

    #[test]
    fn test_factory_creates_local_client_when_no_servers_configured() {
        let mut config = Config::default();
        config.refresh_mcp_servers(None).unwrap();

        let factory = build_factory(config);

        // Should create without error and fall back to local client
        let result = factory.create();
        assert!(result.is_ok());
    }

    #[test]
    fn test_remote_only_without_servers_errors() {
        let mut config = Config::default();
        config.mcp.mode = McpMode::RemoteOnly;
        config.mcp_servers.clear();
        config.refresh_mcp_servers(None).unwrap();

        let factory = build_factory(config);
        let result = factory.create();
        assert!(matches!(result, Err(Error::Config(_))));
    }

    #[test]
    fn test_remote_preferred_without_fallback_propagates_remote_error() {
        let mut config = Config::default();
        config.mcp.mode = McpMode::RemotePreferred;
        config.mcp.allow_fallback = false;
        config.mcp_servers = vec![McpServerConfig {
            name: "invalid".to_string(),
            command: "nonexistent-mcp-server-binary".to_string(),
            args: Vec::new(),
            transport: "stdio".to_string(),
            env: std::collections::HashMap::new(),
            oauth: None,
        }];
        config.refresh_mcp_servers(None).unwrap();

        let factory = build_factory(config);
        let result = factory.create();
        assert!(
            matches!(result, Err(Error::Config(message)) if message.contains("Failed to start remote MCP client"))
        );
    }

    #[test]
    fn test_legacy_mode_uses_local_client() {
        let mut config = Config::default();
        config.mcp.mode = McpMode::Legacy;

        let factory = build_factory(config);
        let result = factory.create();
        assert!(result.is_ok());
    }
}