diff --git a/src/app.rs b/src/app.rs index f7d5aabb..738d659c 100644 --- a/src/app.rs +++ b/src/app.rs @@ -21,7 +21,7 @@ use crate::secrets::SecretsStore; use crate::skills::SkillRegistry; use crate::skills::catalog::SkillCatalog; use crate::tools::ToolRegistry; -use crate::tools::mcp::McpSessionManager; +use crate::tools::mcp::{McpProcessManager, McpSessionManager}; use crate::tools::wasm::SharedCredentialRegistry; use crate::tools::wasm::WasmToolRuntime; use crate::workspace::{EmbeddingProvider, Workspace}; @@ -41,6 +41,7 @@ pub struct AppComponents { pub workspace: Option>, pub extension_manager: Option>, pub mcp_session_manager: Arc, + pub mcp_process_manager: Arc, pub wasm_tool_runtime: Option>, pub log_broadcaster: Arc, pub context_manager: Arc, @@ -420,6 +421,7 @@ impl AppBuilder { ) -> Result< ( Arc, + Arc, Option>, Option>, Vec, @@ -427,10 +429,13 @@ impl AppBuilder { ), anyhow::Error, > { - use crate::tools::mcp::{McpClient, config::load_mcp_servers_from_db, is_authenticated}; + use crate::tools::mcp::{ + McpClient, McpTransport, config::load_mcp_servers_from_db, is_authenticated, + }; use crate::tools::wasm::{WasmToolLoader, load_dev_tools}; let mcp_session_manager = Arc::new(McpSessionManager::new()); + let mcp_process_manager = Arc::new(McpProcessManager::new()); // Create WASM tool runtime eagerly so extensions installed after startup // (e.g. via the web UI) can still be activated. The tools directory is only @@ -506,97 +511,175 @@ impl AppBuilder { let db = self.db.clone(); let tools = Arc::clone(tools); let mcp_sm = Arc::clone(&mcp_session_manager); + let pm = Arc::clone(&mcp_process_manager); async move { - if let Some(ref secrets) = secrets_store { - let servers_result = if let Some(ref d) = db { - load_mcp_servers_from_db(d.as_ref(), "default").await - } else { - crate::tools::mcp::config::load_mcp_servers().await - }; - match servers_result { - Ok(servers) => { - let enabled: Vec<_> = servers.enabled_servers().cloned().collect(); - if !enabled.is_empty() { - tracing::info!( - "Loading {} configured MCP server(s)...", - enabled.len() - ); - } + let servers_result = if let Some(ref d) = db { + load_mcp_servers_from_db(d.as_ref(), "default").await + } else { + crate::tools::mcp::config::load_mcp_servers().await + }; + match servers_result { + Ok(servers) => { + let enabled: Vec<_> = servers.enabled_servers().cloned().collect(); + if !enabled.is_empty() { + tracing::info!("Loading {} configured MCP server(s)...", enabled.len()); + } - let mut join_set = tokio::task::JoinSet::new(); - for server in enabled { - let mcp_sm = Arc::clone(&mcp_sm); - let secrets = Arc::clone(secrets); - let tools = Arc::clone(&tools); + let mut join_set = tokio::task::JoinSet::new(); + for server in enabled { + let mcp_sm = Arc::clone(&mcp_sm); + let secrets = secrets_store.clone(); + let tools = Arc::clone(&tools); + let pm = Arc::clone(&pm); - join_set.spawn(async move { - let server_name = server.name.clone(); - let has_tokens = - is_authenticated(&server, &secrets, "default").await; + join_set.spawn(async move { + let server_name = server.name.clone(); - let client = if has_tokens || server.requires_auth() { - McpClient::new_authenticated( - server, mcp_sm, secrets, "default", - ) - } else { - McpClient::new_with_name(&server_name, &server.url) - }; - - match client.list_tools().await { - Ok(mcp_tools) => { - let tool_count = mcp_tools.len(); - match client.create_tools().await { - Ok(tool_impls) => { - for tool in tool_impls { - tools.register(tool).await; - } - tracing::info!( - "Loaded {} tools from MCP server '{}'", - tool_count, - server_name - ); - } - Err(e) => { - tracing::warn!( - "Failed to create tools from MCP server '{}': {}", - server_name, - e - ); - } + let client: McpClient = match server.effective_transport() { + crate::tools::mcp::config::EffectiveTransport::Stdio { + command, + args, + env, + } => { + match pm + .spawn_stdio( + &server_name, + command, + args.to_vec(), + env.clone(), + ) + .await + { + Ok(transport) => McpClient::new_with_transport( + &server_name, + transport as Arc, + None, + secrets, + "default", + Some(server), + ), + Err(e) => { + tracing::warn!( + "Failed to spawn stdio MCP server '{}': {}", + server_name, + e + ); + return; } } - Err(e) => { - let err_str = e.to_string(); - if err_str.contains("401") - || err_str.contains("authentication") - { + } + #[cfg(unix)] + crate::tools::mcp::config::EffectiveTransport::Unix { + socket_path, + } => { + match crate::tools::mcp::unix_transport::UnixMcpTransport::connect( + &server_name, + socket_path, + ) + .await + { + Ok(transport) => McpClient::new_with_transport( + &server_name, + Arc::new(transport) as Arc, + None, + secrets, + "default", + Some(server), + ), + Err(e) => { tracing::warn!( - "MCP server '{}' requires authentication. \ - Run: ironclaw mcp auth {}", + "Failed to connect to Unix MCP server '{}': {}", server_name, + e + ); + return; + } + } + } + #[cfg(not(unix))] + crate::tools::mcp::config::EffectiveTransport::Unix { .. } => { + tracing::warn!( + "Unix socket transport is not supported on this platform (server '{}')", + server_name + ); + return; + } + crate::tools::mcp::config::EffectiveTransport::Http => { + if let Some(ref secrets) = secrets { + let has_tokens = + is_authenticated(&server, secrets, "default") + .await; + + if has_tokens || server.requires_auth() { + McpClient::new_authenticated( + server, + Arc::clone(&mcp_sm), + Arc::clone(secrets), + "default", + ) + } else { + McpClient::new_with_config(server) + } + } else { + McpClient::new_with_config(server) + } + } + }; + + match client.list_tools().await { + Ok(mcp_tools) => { + let tool_count = mcp_tools.len(); + match client.create_tools().await { + Ok(tool_impls) => { + for tool in tool_impls { + tools.register(tool).await; + } + tracing::info!( + "Loaded {} tools from MCP server '{}'", + tool_count, server_name ); - } else { + } + Err(e) => { tracing::warn!( - "Failed to connect to MCP server '{}': {}", + "Failed to create tools from MCP server '{}': {}", server_name, e ); } } } - }); - } - - while let Some(result) = join_set.join_next().await { - if let Err(e) = result { - tracing::warn!("MCP server loading task panicked: {}", e); + Err(e) => { + let err_str = e.to_string(); + if err_str.contains("401") + || err_str.contains("authentication") + { + tracing::warn!( + "MCP server '{}' requires authentication. \ + Run: ironclaw mcp auth {}", + server_name, + server_name + ); + } else { + tracing::warn!( + "Failed to connect to MCP server '{}': {}", + server_name, + e + ); + } + } } + }); + } + + while let Some(result) = join_set.join_next().await { + if let Err(e) = result { + tracing::warn!("MCP server loading task panicked: {}", e); } } - Err(e) => { - tracing::debug!("No MCP servers configured ({})", e); - } + } + Err(e) => { + tracing::debug!("No MCP servers configured ({})", e); } } } @@ -667,6 +750,7 @@ impl AppBuilder { Ok(( mcp_session_manager, + mcp_process_manager, wasm_tool_runtime, extension_manager, catalog_entries, @@ -702,6 +786,7 @@ impl AppBuilder { let ( mcp_session_manager, + mcp_process_manager, wasm_tool_runtime, extension_manager, catalog_entries, @@ -799,6 +884,7 @@ impl AppBuilder { workspace, extension_manager, mcp_session_manager, + mcp_process_manager, wasm_tool_runtime, log_broadcaster: self.log_broadcaster, context_manager, diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index 5e2f4dea..f9d3acf0 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -2,10 +2,11 @@ //! //! Commands for adding, removing, authenticating, and testing MCP servers. +use std::collections::HashMap; use std::io::Write; use std::sync::Arc; -use clap::Subcommand; +use clap::{Args, Subcommand}; use crate::config::Config; use crate::db::Database; @@ -15,39 +16,67 @@ use crate::secrets::{SecretsCrypto, SecretsStore}; use crate::tools::mcp::{ McpClient, McpServerConfig, McpSessionManager, OAuthConfig, auth::{authorize_mcp_server, is_authenticated}, - config::{self, McpServersFile}, + config::{self, EffectiveTransport, McpServersFile}, }; +/// Arguments for the `mcp add` subcommand. +#[derive(Args, Debug, Clone)] +pub struct McpAddArgs { + /// Server name (e.g., "notion", "github") + pub name: String, + + /// Server URL (e.g., "https://mcp.notion.com") -- required for http transport + pub url: Option, + + /// Transport type: http (default), stdio, unix + #[arg(long, default_value = "http")] + pub transport: String, + + /// Command to run (stdio transport) + #[arg(long)] + pub command: Option, + + /// Command arguments (stdio transport, can be repeated) + #[arg(long = "arg", num_args = 1..)] + pub cmd_args: Vec, + + /// Environment variables (stdio transport, KEY=VALUE format, can be repeated) + #[arg(long = "env", value_parser = parse_env_var)] + pub env: Vec<(String, String)>, + + /// Unix socket path (unix transport) + #[arg(long)] + pub socket: Option, + + /// Custom HTTP headers (KEY:VALUE format, can be repeated) + #[arg(long = "header", value_parser = parse_header)] + pub headers: Vec<(String, String)>, + + /// OAuth client ID (if authentication is required) + #[arg(long)] + pub client_id: Option, + + /// OAuth authorization URL (optional, can be discovered) + #[arg(long)] + pub auth_url: Option, + + /// OAuth token URL (optional, can be discovered) + #[arg(long)] + pub token_url: Option, + + /// Scopes to request (comma-separated) + #[arg(long)] + pub scopes: Option, + + /// Server description + #[arg(long)] + pub description: Option, +} + #[derive(Subcommand, Debug, Clone)] pub enum McpCommand { /// Add an MCP server - Add { - /// Server name (e.g., "notion", "github") - name: String, - - /// Server URL (e.g., "https://mcp.notion.com") - url: String, - - /// OAuth client ID (if authentication is required) - #[arg(long)] - client_id: Option, - - /// OAuth authorization URL (optional, can be discovered) - #[arg(long)] - auth_url: Option, - - /// OAuth token URL (optional, can be discovered) - #[arg(long)] - token_url: Option, - - /// Scopes to request (comma-separated) - #[arg(long)] - scopes: Option, - - /// Server description - #[arg(long)] - description: Option, - }, + Add(Box), /// Remove an MCP server Remove { @@ -97,29 +126,24 @@ pub enum McpCommand { }, } +fn parse_header(s: &str) -> Result<(String, String), String> { + let pos = s + .find(':') + .ok_or_else(|| format!("invalid header format '{}', expected KEY:VALUE", s))?; + Ok((s[..pos].trim().to_string(), s[pos + 1..].trim().to_string())) +} + +fn parse_env_var(s: &str) -> Result<(String, String), String> { + let pos = s + .find('=') + .ok_or_else(|| format!("invalid env var format '{}', expected KEY=VALUE", s))?; + Ok((s[..pos].to_string(), s[pos + 1..].to_string())) +} + /// Run an MCP command. pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { match cmd { - McpCommand::Add { - name, - url, - client_id, - auth_url, - token_url, - scopes, - description, - } => { - add_server( - name, - url, - client_id, - auth_url, - token_url, - scopes, - description, - ) - .await - } + McpCommand::Add(args) => add_server(*args).await, McpCommand::Remove { name } => remove_server(name).await, McpCommand::List { verbose } => list_servers(verbose).await, McpCommand::Auth { name, user } => auth_server(name, user).await, @@ -133,16 +157,58 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { } /// Add a new MCP server. -async fn add_server( - name: String, - url: String, - client_id: Option, - auth_url: Option, - token_url: Option, - scopes: Option, - description: Option, -) -> anyhow::Result<()> { - let mut config = McpServerConfig::new(&name, &url); +async fn add_server(args: McpAddArgs) -> anyhow::Result<()> { + let McpAddArgs { + name, + url, + transport, + command, + cmd_args, + env, + socket, + headers, + client_id, + auth_url, + token_url, + scopes, + description, + } = args; + + let transport_lower = transport.to_lowercase(); + + let mut config = match transport_lower.as_str() { + "stdio" => { + let cmd = command + .clone() + .ok_or_else(|| anyhow::anyhow!("--command is required for stdio transport"))?; + let env_map: HashMap = env.into_iter().collect(); + McpServerConfig::new_stdio(&name, &cmd, cmd_args.clone(), env_map) + } + "unix" => { + let socket_path = socket + .clone() + .ok_or_else(|| anyhow::anyhow!("--socket is required for unix transport"))?; + McpServerConfig::new_unix(&name, &socket_path) + } + "http" => { + let url_val = url + .as_deref() + .ok_or_else(|| anyhow::anyhow!("URL is required for http transport"))?; + McpServerConfig::new(&name, url_val) + } + other => { + anyhow::bail!( + "Unknown transport type '{}'. Supported: http, stdio, unix", + other + ); + } + }; + + // Apply headers if any + if !headers.is_empty() { + let headers_map: HashMap = headers.into_iter().collect(); + config = config.with_headers(headers_map); + } if let Some(desc) = description { config = config.with_description(desc); @@ -151,8 +217,12 @@ async fn add_server( // Track if auth is required let requires_auth = client_id.is_some(); - // Set up OAuth if client_id is provided + // Set up OAuth if client_id is provided (HTTP transport only) if let Some(client_id) = client_id { + if transport_lower != "http" { + anyhow::bail!("OAuth authentication is only supported with http transport"); + } + let mut oauth = OAuthConfig::new(client_id); if let (Some(auth), Some(token)) = (auth_url, token_url) { @@ -181,7 +251,24 @@ async fn add_server( println!(); println!(" ✓ Added MCP server '{}'", name); - println!(" URL: {}", url); + + match transport_lower.as_str() { + "stdio" => { + println!( + " Transport: stdio (command: {})", + command.as_deref().unwrap_or("") + ); + } + "unix" => { + println!( + " Transport: unix (socket: {})", + socket.as_deref().unwrap_or("") + ); + } + _ => { + println!(" URL: {}", url.as_deref().unwrap_or("")); + } + } if requires_auth { println!(); @@ -236,9 +323,40 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> { "" }; + let effective = server.effective_transport(); + + let transport_label = match &effective { + EffectiveTransport::Http => "http".to_string(), + EffectiveTransport::Stdio { command, .. } => { + format!("stdio ({})", command) + } + EffectiveTransport::Unix { socket_path } => { + format!("unix ({})", socket_path) + } + }; + if verbose { println!(" {} {}{}", status, server.name, auth_status); - println!(" URL: {}", server.url); + println!(" Transport: {}", transport_label); + match &effective { + EffectiveTransport::Http => { + println!(" URL: {}", server.url); + } + EffectiveTransport::Stdio { command, args, env } => { + println!(" Command: {}", command); + if !args.is_empty() { + println!(" Args: {}", args.join(", ")); + } + if !env.is_empty() { + // Only print env var names, not values (may contain secrets). + let env_keys: Vec<&str> = env.keys().map(|k| k.as_str()).collect(); + println!(" Env: {}", env_keys.join(", ")); + } + } + EffectiveTransport::Unix { socket_path } => { + println!(" Socket: {}", socket_path); + } + } if let Some(ref desc) = server.description { println!(" Description: {}", desc); } @@ -248,11 +366,27 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> { println!(" Scopes: {}", oauth.scopes.join(", ")); } } + if !server.headers.is_empty() { + let header_keys: Vec<&String> = server.headers.keys().collect(); + println!( + " Headers: {}", + header_keys + .iter() + .map(|k| k.as_str()) + .collect::>() + .join(", ") + ); + } println!(); } else { + let display = match &effective { + EffectiveTransport::Http => server.url.clone(), + EffectiveTransport::Stdio { command, .. } => command.to_string(), + EffectiveTransport::Unix { socket_path } => socket_path.to_string(), + }; println!( - " {} {} - {}{}", - status, server.name, server.url, auth_status + " {} {} - {} [{}]{}", + status, server.name, display, transport_label, auth_status ); } } @@ -374,7 +508,7 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { return Ok(()); } else { // No OAuth and no tokens - try unauthenticated - McpClient::new_with_name(&server.name, &server.url) + McpClient::new_with_config(server.clone()) }; // Test connection @@ -579,4 +713,46 @@ mod tests { TestCli::command().debug_assert(); } + + #[test] + fn test_parse_header_valid() { + let result = parse_header("Authorization: Bearer token123").unwrap(); + assert_eq!(result.0, "Authorization"); + assert_eq!(result.1, "Bearer token123"); + } + + #[test] + fn test_parse_header_no_spaces() { + let result = parse_header("X-Api-Key:abc123").unwrap(); + assert_eq!(result.0, "X-Api-Key"); + assert_eq!(result.1, "abc123"); + } + + #[test] + fn test_parse_header_invalid() { + let result = parse_header("no-colon-here"); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("invalid header format")); + } + + #[test] + fn test_parse_env_var_valid() { + let result = parse_env_var("NODE_ENV=production").unwrap(); + assert_eq!(result.0, "NODE_ENV"); + assert_eq!(result.1, "production"); + } + + #[test] + fn test_parse_env_var_with_equals_in_value() { + let result = parse_env_var("KEY=value=with=equals").unwrap(); + assert_eq!(result.0, "KEY"); + assert_eq!(result.1, "value=with=equals"); + } + + #[test] + fn test_parse_env_var_invalid() { + let result = parse_env_var("no-equals-here"); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("invalid env var format")); + } } diff --git a/src/cli/mod.rs b/src/cli/mod.rs index f266b9b6..1e47ccfd 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -132,7 +132,7 @@ pub enum Command { about = "Manage MCP servers", long_about = "Add, auth, list, or test MCP servers.\nExample: ironclaw mcp add notion https://mcp.notion.com" )] - Mcp(McpCommand), + Mcp(Box), /// Query and manage workspace memory #[command( diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 193ec56e..3f51511e 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -1596,6 +1596,7 @@ impl ExtensionManager { &metadata.scopes_supported, Some(&pkce), &std::collections::HashMap::new(), + None, ); // Store pending auth for later callback handling @@ -2476,7 +2477,7 @@ impl ExtensionManager { &self.user_id, ) } else { - McpClient::new_with_name(&server.name, &server.url) + McpClient::new_with_config(server.clone()) }; // Try to list and create tools diff --git a/src/main.rs b/src/main.rs index 9e79d378..5ac7d315 100644 --- a/src/main.rs +++ b/src/main.rs @@ -75,7 +75,7 @@ async fn async_main() -> anyhow::Result<()> { } Some(Command::Mcp(mcp_cmd)) => { init_cli_tracing(); - return run_mcp_command(mcp_cmd.clone()).await; + return run_mcp_command(*mcp_cmd.clone()).await; } Some(Command::Memory(mem_cmd)) => { init_cli_tracing(); @@ -723,6 +723,9 @@ async fn async_main() -> anyhow::Result<()> { // ── Shutdown ──────────────────────────────────────────────────────── + // Shut down all stdio MCP server child processes. + components.mcp_process_manager.shutdown_all().await; + // Flush LLM trace recording if enabled if let Some(ref recorder) = components.recording_handle && let Err(e) = recorder.flush().await diff --git a/src/tools/mcp/auth.rs b/src/tools/mcp/auth.rs index 0f7cd3a5..2e483b60 100644 --- a/src/tools/mcp/auth.rs +++ b/src/tools/mcp/auth.rs @@ -4,6 +4,7 @@ //! See: https://spec.modelcontextprotocol.io/specification/2025-03-26/basic/authorization/ use std::collections::HashMap; +use std::net::IpAddr; use std::sync::Arc; use std::time::Duration; @@ -199,23 +200,285 @@ impl PkceChallenge { } } +// --------------------------------------------------------------------------- +// Well-known URI construction (RFC 8414 / RFC 9728) +// --------------------------------------------------------------------------- + +/// Build a well-known URI according to RFC 8414 / RFC 9728. +/// +/// The path component of the base URL is placed *after* the well-known suffix: +/// ```text +/// https://example.com/path + oauth-authorization-server +/// -> https://example.com/.well-known/oauth-authorization-server/path +/// ``` +pub fn build_well_known_uri(base_url: &str, suffix: &str) -> Result { + let parsed = reqwest::Url::parse(base_url) + .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid URL: {}", e)))?; + let origin = parsed.origin().ascii_serialization(); + let path = parsed.path().trim_end_matches('/'); + Ok(format!("{}/.well-known/{}{}", origin, suffix, path)) +} + +// --------------------------------------------------------------------------- +// RFC 8707 resource parameter +// --------------------------------------------------------------------------- + +/// Compute the canonical resource URI for RFC 8707. +/// +/// Strips fragments and trailing slashes from the server URL. +pub fn canonical_resource_uri(server_url: &str) -> String { + match reqwest::Url::parse(server_url) { + Ok(mut parsed) => { + parsed.set_fragment(None); + let s = parsed.to_string(); + s.trim_end_matches('/').to_string() + } + Err(_) => server_url.trim_end_matches('/').to_string(), + } +} + +// --------------------------------------------------------------------------- +// SSRF protection +// --------------------------------------------------------------------------- + +/// Check if an IP address is dangerous (loopback, link-local, private, etc.) +fn is_dangerous_ip(ip: IpAddr) -> bool { + match ip { + IpAddr::V4(v4) => { + v4.is_loopback() + || v4.is_private() + || v4.is_link_local() + || v4.is_broadcast() + || v4.is_unspecified() + || (v4.octets()[0] == 169 && v4.octets()[1] == 254) // link-local + || (v4.octets()[0] == 100 && (v4.octets()[1] & 0xC0) == 64) // CGNAT 100.64/10 + } + IpAddr::V6(v6) => { + let segs = v6.segments(); + v6.is_loopback() + || v6.is_unspecified() + // Link-local (fe80::/10) + || (segs[0] & 0xffc0) == 0xfe80 + // Site-local / deprecated (fec0::/10) + || (segs[0] & 0xffc0) == 0xfec0 + // Unique local (fc00::/7) + || (segs[0] & 0xfe00) == 0xfc00 + // Documentation (2001:db8::/32) + || (segs[0] == 0x2001 && segs[1] == 0x0db8) + // Check for IPv4-mapped IPv6 (::ffff:x.x.x.x) + || v6 + .to_ipv4_mapped() + .is_some_and(|v4| is_dangerous_ip(IpAddr::V4(v4))) + } + } +} + +/// Validate that a URL is safe for server-side requests (SSRF protection). +async fn validate_url_safe(url: &str) -> Result<(), AuthError> { + let parsed = reqwest::Url::parse(url) + .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid URL: {}", e)))?; + + // Must be HTTPS. HTTP is only allowed for localhost/loopback (dev scenarios). + let scheme = parsed.scheme(); + if scheme != "https" && scheme != "http" { + return Err(AuthError::DiscoveryFailed(format!( + "Unsupported scheme: {}", + scheme + ))); + } + if scheme == "http" { + let host = parsed.host_str().unwrap_or(""); + let is_localhost = + host == "localhost" || host == "127.0.0.1" || host == "::1" || host == "[::1]"; + if !is_localhost { + return Err(AuthError::DiscoveryFailed(format!( + "HTTP is only allowed for localhost; use HTTPS for '{}'", + host + ))); + } + // Localhost HTTP is allowed for dev — skip SSRF checks since we've + // already validated the host is localhost/loopback. + return Ok(()); + } + + let host = parsed + .host_str() + .ok_or_else(|| AuthError::DiscoveryFailed("URL has no host".to_string()))?; + + // For IP literals, parse directly and check. + if let Ok(ip) = host.parse::() + && is_dangerous_ip(ip) + { + return Err(AuthError::DiscoveryFailed(format!( + "URL points to a restricted IP address: {}", + host + ))); + } + + // For hostnames, resolve DNS and check each resolved address. + // This prevents DNS-based SSRF where a hostname resolves to an internal IP + // (e.g., 169.254.169.254 for cloud metadata endpoints). + if host.parse::().is_err() { + let addr = format!("{}:{}", host, parsed.port_or_known_default().unwrap_or(443)); + match tokio::net::lookup_host(&addr).await { + Ok(addrs) => { + for socket_addr in addrs { + if is_dangerous_ip(socket_addr.ip()) { + return Err(AuthError::DiscoveryFailed(format!( + "URL hostname '{}' resolves to restricted IP address: {}", + host, + socket_addr.ip() + ))); + } + } + } + Err(e) => { + // DNS failure = fail closed (do not allow the request) + return Err(AuthError::DiscoveryFailed(format!( + "DNS resolution failed for '{}': {}", + host, e + ))); + } + } + } + + Ok(()) +} + +// --------------------------------------------------------------------------- +// Multi-strategy OAuth discovery helpers +// --------------------------------------------------------------------------- + +/// Parse the resource_metadata URL from a WWW-Authenticate header value. +fn parse_resource_metadata_url(www_authenticate: &str) -> Option { + // Try comma-separated parameters first + for part in www_authenticate.split(',') { + let part = part.trim(); + if let Some(rest) = part.strip_prefix("resource_metadata=\"") { + return rest.strip_suffix('"').map(|s| s.to_string()); + } + if let Some(rest) = part.strip_prefix("resource_metadata=") { + let val = rest.trim_matches('"'); + return Some(val.to_string()); + } + } + // Also try whitespace-separated tokens (e.g. Bearer resource_metadata="url") + for part in www_authenticate.split_whitespace() { + if let Some(rest) = part.strip_prefix("resource_metadata=\"") { + return rest + .trim_end_matches(',') + .strip_suffix('"') + .map(|s| s.to_string()); + } + if let Some(rest) = part.strip_prefix("resource_metadata=") { + let val = rest.trim_matches('"').trim_end_matches(','); + return Some(val.to_string()); + } + } + None +} + +/// Fetch protected resource metadata from a URL. +async fn fetch_resource_metadata(url: &str) -> Result { + validate_url_safe(url).await?; + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|e| AuthError::Http(e.to_string()))?; + + let response = client + .get(url) + .send() + .await + .map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?; + + if !response.status().is_success() { + return Err(AuthError::DiscoveryFailed(format!( + "HTTP {}", + response.status() + ))); + } + + response + .json() + .await + .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid metadata: {}", e))) +} + +/// Try to discover OAuth metadata via 401 challenge response. +async fn discover_via_401(server_url: &str) -> Result { + validate_url_safe(server_url).await?; + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|e| AuthError::Http(e.to_string()))?; + + let response = client + .post(server_url) + .header("Content-Type", "application/json") + .body("{}") + .send() + .await + .map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?; + + if response.status().as_u16() != 401 { + return Err(AuthError::DiscoveryFailed(format!( + "Expected 401, got {}", + response.status() + ))); + } + + let www_auth = response + .headers() + .get("WWW-Authenticate") + .and_then(|v| v.to_str().ok()) + .ok_or_else(|| { + AuthError::DiscoveryFailed("No WWW-Authenticate header in 401 response".to_string()) + })?; + + let resource_metadata_url = parse_resource_metadata_url(www_auth).ok_or_else(|| { + AuthError::DiscoveryFailed( + "No resource_metadata URL in WWW-Authenticate header".to_string(), + ) + })?; + + let resource_meta = fetch_resource_metadata(&resource_metadata_url).await?; + try_discover_from_auth_servers(&resource_meta).await +} + +/// Try to discover auth server metadata from resource metadata's authorization_servers list. +async fn try_discover_from_auth_servers( + resource_meta: &ProtectedResourceMetadata, +) -> Result { + let auth_server_url = resource_meta + .authorization_servers + .first() + .ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?; + + discover_authorization_server(auth_server_url).await +} + +// --------------------------------------------------------------------------- +// Discovery functions +// --------------------------------------------------------------------------- + /// Discover protected resource metadata from an MCP server. pub async fn discover_protected_resource( server_url: &str, ) -> Result { + validate_url_safe(server_url).await?; + let client = reqwest::Client::builder() .timeout(Duration::from_secs(10)) + .redirect(reqwest::redirect::Policy::none()) .build() .map_err(|e| AuthError::Http(e.to_string()))?; - // Parse the server URL to extract the origin (scheme + host + port) - // The .well-known endpoints are always at the root of the origin, not under any path - let parsed = reqwest::Url::parse(server_url) - .map_err(|e| AuthError::DiscoveryFailed(format!("Invalid server URL: {}", e)))?; - let origin = parsed.origin().ascii_serialization(); - - // Try the well-known endpoint at the origin root - let well_known_url = format!("{}/.well-known/oauth-protected-resource", origin); + let well_known_url = build_well_known_uri(server_url, "oauth-protected-resource")?; let response = client .get(&well_known_url) @@ -237,13 +500,15 @@ pub async fn discover_protected_resource( pub async fn discover_authorization_server( auth_server_url: &str, ) -> Result { + validate_url_safe(auth_server_url).await?; + let client = reqwest::Client::builder() .timeout(Duration::from_secs(10)) + .redirect(reqwest::redirect::Policy::none()) .build() .map_err(|e| AuthError::Http(e.to_string()))?; - let base_url = auth_server_url.trim_end_matches('/'); - let well_known_url = format!("{}/.well-known/oauth-authorization-server", base_url); + let well_known_url = build_well_known_uri(auth_server_url, "oauth-authorization-server")?; let response = client .get(&well_known_url) @@ -298,20 +563,27 @@ pub async fn discover_oauth_endpoints( /// Discover full OAuth metadata including DCR support. /// /// Returns authorization server metadata which includes registration_endpoint if DCR is supported. +/// Uses a 3-strategy discovery chain: +/// 1. **401-based**: POST to MCP server, parse WWW-Authenticate header for resource_metadata URL +/// 2. **RFC 9728**: Discover protected resource metadata, then authorization server from it +/// 3. **Direct**: Treat MCP server as its own auth server pub async fn discover_full_oauth_metadata( server_url: &str, ) -> Result { - // Try to discover from the server - let resource_meta = discover_protected_resource(server_url).await?; + // Strategy 1: 401-based discovery + if let Ok(meta) = discover_via_401(server_url).await { + return Ok(meta); + } - // Get the first authorization server - let auth_server_url = resource_meta - .authorization_servers - .first() - .ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?; + // Strategy 2: RFC 9728 protected resource discovery + if let Ok(resource_meta) = discover_protected_resource(server_url).await + && let Ok(meta) = try_discover_from_auth_servers(&resource_meta).await + { + return Ok(meta); + } - // Discover the authorization server metadata - discover_authorization_server(auth_server_url).await + // Strategy 3: Direct - treat MCP server as its own auth server + discover_authorization_server(server_url).await } /// Perform Dynamic Client Registration with an authorization server. @@ -321,8 +593,11 @@ pub async fn register_client( registration_endpoint: &str, redirect_uri: &str, ) -> Result { + validate_url_safe(registration_endpoint).await?; + let client = reqwest::Client::builder() .timeout(Duration::from_secs(30)) + .redirect(reqwest::redirect::Policy::none()) .build() .map_err(|e| AuthError::Http(e.to_string()))?; @@ -417,7 +692,7 @@ pub async fn authorize_mcp_server( println!(" Registering client dynamically..."); let registration = register_client(®istration_endpoint, &redirect_uri).await?; - println!(" ✓ Client registered: {}", registration.client_id); + println!(" Client registered: {}", registration.client_id); ( registration.client_id, @@ -436,6 +711,15 @@ pub async fn authorize_mcp_server( None }; + // Compute canonical resource URI for RFC 8707 + let resource = canonical_resource_uri(&server_config.url); + + // Validate the discovered authorization URL to prevent a malicious MCP server + // from redirecting the user to a phishing page or non-HTTPS endpoint. + validate_url_safe(&authorization_url) + .await + .map_err(|e| AuthError::DiscoveryFailed(format!("Unsafe authorization endpoint: {}", e)))?; + // Build authorization URL let auth_url = build_authorization_url( &authorization_url, @@ -444,6 +728,7 @@ pub async fn authorize_mcp_server( &scopes, pkce.as_ref(), &extra_params, + Some(&resource), ); // Open browser @@ -462,9 +747,15 @@ pub async fn authorize_mcp_server( println!(" Exchanging code for token..."); // Exchange code for token - let token = - exchange_code_for_token(&token_url, &client_id, &code, &redirect_uri, pkce.as_ref()) - .await?; + let token = exchange_code_for_token( + &token_url, + &client_id, + &code, + &redirect_uri, + pkce.as_ref(), + Some(&resource), + ) + .await?; // Store the tokens store_tokens(secrets, user_id, server_config, &token).await?; @@ -493,6 +784,7 @@ pub fn build_authorization_url( scopes: &[String], pkce: Option<&PkceChallenge>, extra_params: &HashMap, + resource: Option<&str>, ) -> String { let mut url = format!( "{}?client_id={}&response_type=code&redirect_uri={}", @@ -523,6 +815,10 @@ pub fn build_authorization_url( )); } + if let Some(resource) = resource { + url.push_str(&format!("&resource={}", urlencoding::encode(resource))); + } + url } @@ -553,9 +849,13 @@ pub async fn exchange_code_for_token( code: &str, redirect_uri: &str, pkce: Option<&PkceChallenge>, + resource: Option<&str>, ) -> Result { + validate_url_safe(token_url).await?; + let client = reqwest::Client::builder() .timeout(Duration::from_secs(30)) + .redirect(reqwest::redirect::Policy::none()) .build() .map_err(|e| AuthError::Http(e.to_string()))?; @@ -570,6 +870,10 @@ pub async fn exchange_code_for_token( params.push(("code_verifier", pkce.verifier.clone())); } + if let Some(resource) = resource { + params.push(("resource", resource.to_string())); + } + let response = client .post(token_url) .form(¶ms) @@ -738,15 +1042,22 @@ pub async fn refresh_access_token( auth_meta.token_endpoint }; + validate_url_safe(&token_url).await?; + let client = reqwest::Client::builder() .timeout(Duration::from_secs(30)) + .redirect(reqwest::redirect::Policy::none()) .build() .map_err(|e| AuthError::Http(e.to_string()))?; + // Compute canonical resource URI for RFC 8707 + let resource = canonical_resource_uri(&server_config.url); + let params = vec![ ("grant_type", "refresh_token".to_string()), ("refresh_token", refresh_token.expose().to_string()), ("client_id", client_id), + ("resource", resource), ]; let response = client @@ -815,6 +1126,7 @@ mod tests { &["read".to_string(), "write".to_string()], None, &HashMap::new(), + None, ); assert!(url.starts_with("https://auth.example.com/authorize?")); @@ -834,6 +1146,7 @@ mod tests { &[], Some(&pkce), &HashMap::new(), + None, ); assert!(url.contains(&format!("code_challenge={}", pkce.challenge))); @@ -853,6 +1166,7 @@ mod tests { &[], None, &extra, + None, ); assert!(url.contains("owner=user")); @@ -880,6 +1194,7 @@ mod tests { &[], None, &HashMap::new(), + None, ); // With no scopes, the URL must not contain a scope parameter at all. @@ -895,6 +1210,7 @@ mod tests { &[], None, &HashMap::new(), + None, ); // Spaces and ampersands in client_id must be percent-encoded. @@ -1164,4 +1480,235 @@ mod tests { ); } } + + // --- New tests for well-known URI construction --- + + #[test] + fn test_build_well_known_uri_no_path() { + let uri = + build_well_known_uri("https://example.com", "oauth-authorization-server").unwrap(); + assert_eq!( + uri, + "https://example.com/.well-known/oauth-authorization-server" + ); + } + + #[test] + fn test_build_well_known_uri_with_path() { + let uri = + build_well_known_uri("https://example.com/path", "oauth-authorization-server").unwrap(); + assert_eq!( + uri, + "https://example.com/.well-known/oauth-authorization-server/path" + ); + } + + #[test] + fn test_build_well_known_uri_with_trailing_slash() { + let uri = + build_well_known_uri("https://example.com/path/", "oauth-protected-resource").unwrap(); + assert_eq!( + uri, + "https://example.com/.well-known/oauth-protected-resource/path" + ); + } + + #[test] + fn test_build_well_known_uri_root_trailing_slash() { + let uri = + build_well_known_uri("https://example.com/", "oauth-authorization-server").unwrap(); + assert_eq!( + uri, + "https://example.com/.well-known/oauth-authorization-server" + ); + } + + // --- New tests for canonical_resource_uri --- + + #[test] + fn test_canonical_resource_uri_strips_fragment() { + assert_eq!( + canonical_resource_uri("https://mcp.example.com/v1#section"), + "https://mcp.example.com/v1" + ); + } + + #[test] + fn test_canonical_resource_uri_strips_trailing_slash() { + assert_eq!( + canonical_resource_uri("https://mcp.example.com/v1/"), + "https://mcp.example.com/v1" + ); + } + + #[test] + fn test_canonical_resource_uri_no_changes_needed() { + assert_eq!( + canonical_resource_uri("https://mcp.example.com/v1"), + "https://mcp.example.com/v1" + ); + } + + // --- New tests for SSRF protection --- + + #[test] + fn test_is_dangerous_ip_loopback_v4() { + assert!(is_dangerous_ip("127.0.0.1".parse().unwrap())); + assert!(is_dangerous_ip("127.0.0.2".parse().unwrap())); + } + + #[test] + fn test_is_dangerous_ip_private_v4() { + assert!(is_dangerous_ip("10.0.0.1".parse().unwrap())); + assert!(is_dangerous_ip("172.16.0.1".parse().unwrap())); + assert!(is_dangerous_ip("192.168.1.1".parse().unwrap())); + } + + #[test] + fn test_is_dangerous_ip_link_local_v4() { + assert!(is_dangerous_ip("169.254.169.254".parse().unwrap())); + } + + #[test] + fn test_is_dangerous_ip_cgnat() { + assert!(is_dangerous_ip("100.64.0.1".parse().unwrap())); + assert!(is_dangerous_ip("100.127.255.254".parse().unwrap())); + } + + #[test] + fn test_is_dangerous_ip_safe_v4() { + assert!(!is_dangerous_ip("8.8.8.8".parse().unwrap())); + assert!(!is_dangerous_ip("1.1.1.1".parse().unwrap())); + } + + #[test] + fn test_is_dangerous_ip_ipv4_mapped_v6_loopback() { + // ::ffff:127.0.0.1 must be blocked + let ip: IpAddr = "::ffff:127.0.0.1".parse().unwrap(); + assert!(is_dangerous_ip(ip)); + } + + #[test] + fn test_is_dangerous_ip_ipv4_mapped_v6_link_local() { + // ::ffff:169.254.169.254 must be blocked + let ip: IpAddr = "::ffff:169.254.169.254".parse().unwrap(); + assert!(is_dangerous_ip(ip)); + } + + #[test] + fn test_is_dangerous_ip_unspecified() { + assert!(is_dangerous_ip("0.0.0.0".parse().unwrap())); + assert!(is_dangerous_ip("::".parse().unwrap())); + } + + #[test] + fn test_is_dangerous_ip_v6_loopback() { + assert!(is_dangerous_ip("::1".parse().unwrap())); + } + + #[tokio::test] + async fn test_validate_url_safe_https() { + assert!(validate_url_safe("https://example.com/path").await.is_ok()); + } + + #[tokio::test] + async fn test_validate_url_safe_http_localhost_allowed() { + // HTTP is only allowed for localhost dev scenarios + assert!(validate_url_safe("http://localhost/path").await.is_ok()); + assert!( + validate_url_safe("http://localhost:8080/path") + .await + .is_ok() + ); + } + + #[tokio::test] + async fn test_validate_url_safe_http_non_localhost_rejected() { + // HTTP to non-localhost hosts must be rejected (plaintext credential risk) + assert!(validate_url_safe("http://example.com/path").await.is_err()); + } + + #[tokio::test] + async fn test_validate_url_safe_bad_scheme() { + assert!(validate_url_safe("ftp://example.com/path").await.is_err()); + assert!(validate_url_safe("file:///etc/passwd").await.is_err()); + } + + #[tokio::test] + async fn test_validate_url_safe_private_ip() { + // 127.0.0.1 over HTTP is allowed (localhost dev scenario) + assert!(validate_url_safe("http://127.0.0.1/path").await.is_ok()); + // Private/link-local IPs over HTTPS are blocked (SSRF protection) + assert!(validate_url_safe("https://10.0.0.1/path").await.is_err()); + assert!( + validate_url_safe("https://169.254.169.254/latest/meta-data") + .await + .is_err() + ); + // Private IPs over HTTP (non-localhost) are blocked + assert!(validate_url_safe("http://10.0.0.1/path").await.is_err()); + } + + #[tokio::test] + async fn test_validate_url_safe_public_ip() { + assert!(validate_url_safe("https://8.8.8.8/dns").await.is_ok()); + } + + // --- New tests for parse_resource_metadata_url --- + + #[test] + fn test_parse_resource_metadata_url_bearer() { + let header = r#"Bearer resource_metadata="https://res.example.com/.well-known/oauth-protected-resource""#; + let url = parse_resource_metadata_url(header); + assert_eq!( + url.as_deref(), + Some("https://res.example.com/.well-known/oauth-protected-resource") + ); + } + + #[test] + fn test_parse_resource_metadata_url_with_other_params() { + let header = r#"Bearer realm="example", resource_metadata="https://res.example.com/meta""#; + let url = parse_resource_metadata_url(header); + assert_eq!(url.as_deref(), Some("https://res.example.com/meta")); + } + + #[test] + fn test_parse_resource_metadata_url_missing() { + let header = r#"Bearer realm="example""#; + let url = parse_resource_metadata_url(header); + assert!(url.is_none()); + } + + // --- New tests for resource parameter in authorization URL --- + + #[test] + fn test_build_authorization_url_with_resource() { + let url = build_authorization_url( + "https://auth.example.com/authorize", + "client-123", + "http://localhost:9876/callback", + &[], + None, + &HashMap::new(), + Some("https://mcp.example.com/v1"), + ); + + assert!(url.contains("resource=https%3A%2F%2Fmcp.example.com%2Fv1")); + } + + #[test] + fn test_build_authorization_url_without_resource() { + let url = build_authorization_url( + "https://auth.example.com/authorize", + "client-123", + "http://localhost:9876/callback", + &[], + None, + &HashMap::new(), + None, + ); + + assert!(!url.contains("resource=")); + } } diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index e6af2d1c..aa14189b 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -1,11 +1,11 @@ //! MCP client for connecting to MCP servers. //! //! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers. -//! Uses the Streamable HTTP transport with session management. +//! Uses pluggable transports (HTTP, stdio, Unix) via the `McpTransport` trait. +use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; use async_trait::async_trait; use tokio::sync::RwLock; @@ -14,27 +14,29 @@ use crate::context::JobContext; use crate::secrets::SecretsStore; use crate::tools::mcp::auth::refresh_access_token; use crate::tools::mcp::config::McpServerConfig; +use crate::tools::mcp::http_transport::HttpMcpTransport; use crate::tools::mcp::protocol::{ CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool, }; use crate::tools::mcp::session::McpSessionManager; +use crate::tools::mcp::transport::McpTransport; use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput}; /// MCP client for communicating with MCP servers. /// -/// Supports two modes: -/// - Simple: Just a URL, no auth or session management (for local/test servers) -/// - Authenticated: Full OAuth support with session management (for hosted servers) +/// Supports multiple transport types: +/// - HTTP: For remote MCP servers (created via `new`, `new_with_name`, `new_authenticated`) +/// - Stdio/Unix: Via `new_with_transport` with a custom `McpTransport` implementation pub struct McpClient { - /// Server URL (for HTTP transport). + /// Transport for sending requests. + transport: Arc, + + /// Server URL (kept for accessor compatibility). server_url: String, /// Server name (for logging and session management). server_name: String, - /// HTTP client. - http_client: reqwest::Client, - /// Request ID counter. next_id: AtomicU64, @@ -52,6 +54,9 @@ pub struct McpClient { /// Server configuration (for token secret name lookup). server_config: Option, + + /// Custom headers to include in every request. + custom_headers: HashMap, } impl McpClient { @@ -59,22 +64,21 @@ impl McpClient { /// /// Use this for local development servers or servers that don't require auth. pub fn new(server_url: impl Into) -> Self { - let url = server_url.into(); + let url: String = server_url.into(); let name = extract_server_name(&url); + let transport = Arc::new(HttpMcpTransport::new(url.clone(), name.clone())); Self { + transport, server_url: url, server_name: name, - http_client: reqwest::Client::builder() - .timeout(Duration::from_secs(30)) - .build() - .expect("Failed to create HTTP client"), next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager: None, secrets: None, user_id: "default".to_string(), server_config: None, + custom_headers: HashMap::new(), } } @@ -82,19 +86,52 @@ impl McpClient { /// /// Use this when you have a configured server name but no authentication. pub fn new_with_name(server_name: impl Into, server_url: impl Into) -> Self { + let name: String = server_name.into(); + let url: String = server_url.into(); + let transport = Arc::new(HttpMcpTransport::new(url.clone(), name.clone())); + Self { - server_url: server_url.into(), - server_name: server_name.into(), - http_client: reqwest::Client::builder() - .timeout(Duration::from_secs(30)) - .build() - .expect("Failed to create HTTP client"), + transport, + server_url: url, + server_name: name, next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager: None, secrets: None, user_id: "default".to_string(), server_config: None, + custom_headers: HashMap::new(), + } + } + + /// Create a new simple MCP client from an HTTP server configuration (no authentication). + /// + /// Use this when you have an `McpServerConfig` with custom headers but no OAuth. + /// The config must use HTTP transport (the default); for stdio/UDS use `new_with_transport`. + pub fn new_with_config(config: McpServerConfig) -> Self { + assert!( + matches!( + config.effective_transport(), + crate::tools::mcp::config::EffectiveTransport::Http + ), + "new_with_config only supports HTTP transport; use new_with_transport for stdio/UDS" + ); + let transport = Arc::new(HttpMcpTransport::new( + config.url.clone(), + config.name.clone(), + )); + + Self { + transport, + server_url: config.url.clone(), + server_name: config.name.clone(), + next_id: AtomicU64::new(1), + tools_cache: RwLock::new(None), + session_manager: None, + secrets: None, + user_id: "default".to_string(), + custom_headers: config.headers.clone(), + server_config: Some(config), } } @@ -107,19 +144,59 @@ impl McpClient { secrets: Arc, user_id: impl Into, ) -> Self { + let transport = Arc::new( + HttpMcpTransport::new(config.url.clone(), config.name.clone()) + .with_session_manager(session_manager.clone()), + ); + + let custom_headers = config.headers.clone(); + Self { + transport, server_url: config.url.clone(), server_name: config.name.clone(), - http_client: reqwest::Client::builder() - .timeout(Duration::from_secs(30)) - .build() - .expect("Failed to create HTTP client"), next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager: Some(session_manager), secrets: Some(secrets), user_id: user_id.into(), server_config: Some(config), + custom_headers, + } + } + + /// Create a new MCP client with a custom transport. + /// + /// Use this for stdio, UDS, or other non-HTTP transports. + pub fn new_with_transport( + server_name: impl Into, + transport: Arc, + session_manager: Option>, + secrets: Option>, + user_id: impl Into, + server_config: Option, + ) -> Self { + let name: String = server_name.into(); + let url = server_config + .as_ref() + .map(|c| c.url.clone()) + .unwrap_or_default(); + let custom_headers = server_config + .as_ref() + .map(|c| c.headers.clone()) + .unwrap_or_default(); + + Self { + transport, + server_url: url, + server_name: name, + next_id: AtomicU64::new(1), + tools_cache: RwLock::new(None), + session_manager, + secrets, + user_id: user_id.into(), + server_config, + custom_headers, } } @@ -139,19 +216,13 @@ impl McpClient { } /// Get the access token for this server (if authenticated). - /// - /// Returns the stored token regardless of whether OAuth was pre-configured - /// or obtained via Dynamic Client Registration. async fn get_access_token(&self) -> Result, ToolError> { let Some(ref secrets) = self.secrets else { return Ok(None); }; - let Some(ref config) = self.server_config else { return Ok(None); }; - - // Try to get stored token (from either pre-configured OAuth or DCR) match secrets .get_decrypted(&self.user_id, &config.token_secret_name()) .await @@ -165,46 +236,41 @@ impl McpClient { } } + /// Build the headers map for a request (auth, session-id, custom headers). + async fn build_request_headers(&self) -> Result, ToolError> { + let mut headers = self.custom_headers.clone(); + if let Some(token) = self.get_access_token().await? { + headers.insert("Authorization".to_string(), format!("Bearer {}", token)); + } + if let Some(ref session_manager) = self.session_manager + && let Some(session_id) = session_manager.get_session_id(&self.server_name).await + { + headers.insert("Mcp-Session-Id".to_string(), session_id); + } + Ok(headers) + } + /// Send a request to the MCP server with auth and session headers. - /// Automatically attempts token refresh on 401 errors. + /// Automatically attempts token refresh on 401 errors (HTTP transports only). async fn send_request(&self, request: McpRequest) -> Result { - // Try up to 2 times: first attempt, then retry after token refresh + // For non-HTTP transports, just send directly without retry logic + if !self.transport.supports_http_features() { + let headers = self.build_request_headers().await?; + return self.transport.send(&request, &headers).await; + } + + // HTTP transport: try up to 2 times (first attempt, then retry after token refresh) for attempt in 0..2 { - // Request both JSON and SSE as per MCP spec - let mut req_builder = self - .http_client - .post(&self.server_url) - .header("Accept", "application/json, text/event-stream") - .header("Content-Type", "application/json") - .json(&request); + let headers = self.build_request_headers().await?; + let result = self.transport.send(&request, &headers).await; - // Add Authorization header if we have a token - if let Some(token) = self.get_access_token().await? { - req_builder = req_builder.header("Authorization", format!("Bearer {}", token)); - } - - // Add Mcp-Session-Id header if we have a session - if let Some(ref session_manager) = self.session_manager - && let Some(session_id) = session_manager.get_session_id(&self.server_name).await - { - req_builder = req_builder.header("Mcp-Session-Id", session_id); - } - - let response = req_builder.send().await.map_err(|e| { - let mut chain = format!("MCP request failed: {}", e); - let mut source = std::error::Error::source(&e); - while let Some(cause) = source { - chain.push_str(&format!(" -> {}", cause)); - source = cause.source(); - } - ToolError::ExternalService(chain) - })?; - - // Check for 401 Unauthorized - try to refresh token on first attempt - if response.status() == reqwest::StatusCode::UNAUTHORIZED { - if attempt == 0 { - // Try to refresh the token - if let Some(ref secrets) = self.secrets + match result { + Ok(response) => return Ok(response), + Err(ToolError::ExternalService(ref msg)) + if msg.contains("401") || msg.contains("Unauthorized") => + { + if attempt == 0 + && let Some(ref secrets) = self.secrets && let Some(ref config) = self.server_config { tracing::debug!( @@ -214,7 +280,6 @@ impl McpClient { match refresh_access_token(config, secrets, &self.user_id).await { Ok(_) => { tracing::info!("MCP token refreshed for '{}'", self.server_name); - // Continue to next iteration to retry with new token continue; } Err(e) => { @@ -223,108 +288,30 @@ impl McpClient { self.server_name, e ); - // Fall through to return auth error } } } + return Err(ToolError::ExternalService(format!( + "MCP server '{}' requires authentication. Run: ironclaw mcp auth {}", + self.server_name, self.server_name + ))); } - return Err(ToolError::ExternalService(format!( - "MCP server '{}' requires authentication. Run: ironclaw mcp auth {}", - self.server_name, self.server_name - ))); + Err(e) => return Err(e), } - - // Success path - return the parsed response - return self.parse_response(response).await; } - // Should not reach here, but just in case Err(ToolError::ExternalService( "MCP request failed after retry".to_string(), )) } - /// Parse the HTTP response into an MCP response. - async fn parse_response(&self, response: reqwest::Response) -> Result { - // Extract session ID from response header - if let Some(ref session_manager) = self.session_manager - && let Some(session_id) = response - .headers() - .get("Mcp-Session-Id") - .and_then(|v| v.to_str().ok()) - { - session_manager - .update_session_id(&self.server_name, Some(session_id.to_string())) - .await; - } - - if !response.status().is_success() { - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - let preview = sanitize_error_body(&body); - return Err(ToolError::ExternalService(format!( - "MCP server returned status: {status} - {preview}", - ))); - } - - // Check content type to handle SSE vs JSON responses - let content_type = response - .headers() - .get("content-type") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); - - if content_type.contains("text/event-stream") { - // SSE response - read chunks until we get a complete JSON message - use futures::StreamExt; - - let mut stream = response.bytes_stream(); - let mut buffer = String::new(); - - while let Some(chunk) = stream.next().await { - let chunk = chunk.map_err(|e| { - ToolError::ExternalService(format!("Failed to read SSE chunk: {}", e)) - })?; - - buffer.push_str(&String::from_utf8_lossy(&chunk)); - - // Look for complete SSE data lines - for line in buffer.lines() { - if let Some(json_str) = line.strip_prefix("data: ") { - // Try to parse - if valid JSON, we're done - if let Ok(response) = serde_json::from_str::(json_str) { - return Ok(response); - } - } - } - } - - Err(ToolError::ExternalService(format!( - "No valid data in SSE response: {}", - buffer - ))) - } else { - // JSON response - response.json().await.map_err(|e| { - ToolError::ExternalService(format!("Failed to parse MCP response: {}", e)) - }) - } - } - /// Initialize the connection to the MCP server. - /// - /// This should be called once per session to establish capabilities. pub async fn initialize(&self) -> Result { - // Check if already initialized if let Some(ref session_manager) = self.session_manager && session_manager.is_initialized(&self.server_name).await { - // Return cached/default capabilities return Ok(InitializeResult::default()); } - - // Ensure we have a session if let Some(ref session_manager) = self.session_manager { session_manager .get_or_create(&self.server_name, &self.server_url) @@ -352,14 +339,11 @@ impl McpClient { }) })?; - // Mark session as initialized if let Some(ref session_manager) = self.session_manager { session_manager.mark_initialized(&self.server_name).await; } - // Send initialized notification let notification = McpRequest::initialized_notification(); - // Fire and forget - notifications don't have responses let _ = self.send_request(notification).await; Ok(result) @@ -367,12 +351,9 @@ impl McpClient { /// List available tools from the MCP server. pub async fn list_tools(&self) -> Result, ToolError> { - // Check cache first if let Some(tools) = self.tools_cache.read().await.as_ref() { return Ok(tools.clone()); } - - // Ensure initialized for authenticated sessions if self.session_manager.is_some() { self.initialize().await?; } @@ -395,9 +376,7 @@ impl McpClient { .map_err(|e| ToolError::ExternalService(format!("Invalid tools list: {}", e))) })?; - // Cache the tools *self.tools_cache.write().await = Some(result.tools.clone()); - Ok(result.tools) } @@ -407,7 +386,6 @@ impl McpClient { name: &str, arguments: serde_json::Value, ) -> Result { - // Ensure initialized for authenticated sessions if self.session_manager.is_some() { self.initialize().await?; } @@ -440,7 +418,6 @@ impl McpClient { pub async fn create_tools(&self) -> Result>, ToolError> { let mcp_tools = self.list_tools().await?; let client = Arc::new(self.clone()); - Ok(mcp_tools .into_iter() .map(|t| { @@ -465,15 +442,16 @@ impl McpClient { impl Clone for McpClient { fn clone(&self) -> Self { Self { + transport: self.transport.clone(), server_url: self.server_url.clone(), server_name: self.server_name.clone(), - http_client: self.http_client.clone(), next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)), tools_cache: RwLock::new(None), session_manager: self.session_manager.clone(), secrets: self.secrets.clone(), user_id: self.user_id.clone(), server_config: self.server_config.clone(), + custom_headers: self.custom_headers.clone(), } } } @@ -490,7 +468,6 @@ fn extract_server_name(url: &str) -> String { /// Wrapper that implements Tool for an MCP tool. struct McpToolWrapper { tool: McpTool, - /// Prefixed name (server_name_tool_name) for unique identification. prefixed_name: String, client: Arc, } @@ -500,11 +477,9 @@ impl Tool for McpToolWrapper { fn name(&self) -> &str { &self.prefixed_name } - fn description(&self) -> &str { &self.tool.description } - fn parameters_schema(&self) -> serde_json::Value { self.tool.input_schema.clone() } @@ -515,31 +490,24 @@ impl Tool for McpToolWrapper { _ctx: &JobContext, ) -> Result { let start = std::time::Instant::now(); - - // Use the original tool name (without prefix) for the actual call let result = self.client.call_tool(&self.tool.name, params).await?; - - // Convert content blocks to a single result let content: String = result .content .iter() - .filter_map(|block| block.as_text()) + .filter_map(|b| b.as_text()) .collect::>() .join("\n"); - if result.is_error { return Err(ToolError::ExecutionFailed(content)); } - Ok(ToolOutput::text(content, start.elapsed())) } fn requires_sanitization(&self) -> bool { - true // MCP tools are external, always sanitize + true } fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - // Delegate to the MCP protocol type's own requires_approval() bool method if self.tool.requires_approval() { ApprovalRequirement::UnlessAutoApproved } else { @@ -551,55 +519,6 @@ impl Tool for McpToolWrapper { /// Sanitize an HTTP error response body for safe display. /// /// Detects full HTML error pages (containing ` String { - const MAX_CHARS: usize = 200; - - // Only strip tags when the body looks like a full HTML document. - // Plain text that happens to contain `<` / `>` (e.g. log lines, - // comparison expressions) is left untouched. - let lower = body.to_ascii_lowercase(); - let is_html_document = lower.contains("' { - (out, false) - } else if !in_tag { - out.push(c); - (out, false) - } else { - (out, true) - } - }) - .0; - stripped.split_whitespace().collect::>().join(" ") - } else { - body.to_string() - }; - - // Truncate at a char boundary (safe for multi-byte UTF-8). - if text.chars().count() > MAX_CHARS { - let byte_offset = text - .char_indices() - .nth(MAX_CHARS) - .map(|(i, _)| i) - .unwrap_or(text.len()); - format!("{}... ({} bytes total)", &text[..byte_offset], body.len()) - } else { - text - } -} - #[cfg(test)] mod tests { use super::*; @@ -712,42 +631,61 @@ mod tests { #[test] fn test_clone_preserves_fields() { let client = McpClient::new_with_name("cloned-server", "http://localhost:5555"); - // Bump the request ID a few times client.next_request_id(); client.next_request_id(); - let cloned = client.clone(); assert_eq!(cloned.server_url(), "http://localhost:5555"); assert_eq!(cloned.server_name(), "cloned-server"); assert_eq!(cloned.user_id, "default"); - // The atomic counter value is copied assert_eq!(cloned.next_id.load(Ordering::SeqCst), 3); } #[tokio::test] async fn test_clone_resets_tools_cache() { let client = McpClient::new("http://localhost:5555"); - // The clone implementation resets tools_cache to None let cloned = client.clone(); let cache = cloned.tools_cache.read().await; assert!(cache.is_none()); } + #[test] + fn test_new_with_config_carries_custom_headers() { + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "secret".to_string()); + headers.insert("X-Custom".to_string(), "value".to_string()); + + let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers); + let client = McpClient::new_with_config(config.clone()); + + assert_eq!(client.server_name(), "test"); + assert_eq!(client.server_url(), "http://localhost:8080"); + assert_eq!(client.custom_headers.len(), 2); + assert_eq!(client.custom_headers.get("X-API-Key").unwrap(), "secret"); + assert!(client.server_config.is_some()); + } + + #[test] + fn test_new_with_config_no_headers() { + let config = McpServerConfig::new("bare", "http://localhost:9090"); + let client = McpClient::new_with_config(config); + + assert_eq!(client.server_name(), "bare"); + assert!(client.custom_headers.is_empty()); + assert!(client.secrets.is_none()); + assert!(client.session_manager.is_none()); + } + #[test] fn test_next_request_id_monotonically_increasing() { let client = McpClient::new("http://localhost:1234"); - let id1 = client.next_request_id(); - let id2 = client.next_request_id(); - let id3 = client.next_request_id(); - assert_eq!(id1, 1); - assert_eq!(id2, 2); - assert_eq!(id3, 3); + assert_eq!(client.next_request_id(), 1); + assert_eq!(client.next_request_id(), 2); + assert_eq!(client.next_request_id(), 3); } #[test] fn test_mcp_tool_requires_approval_destructive() { use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations}; - let tool = McpTool { name: "delete_all".to_string(), description: "Deletes everything".to_string(), @@ -765,7 +703,6 @@ mod tests { #[test] fn test_mcp_tool_no_approval_when_not_destructive() { use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations}; - let tool = McpTool { name: "read_data".to_string(), description: "Reads data".to_string(), @@ -783,7 +720,6 @@ mod tests { #[test] fn test_mcp_tool_no_approval_when_no_annotations() { use crate::tools::mcp::protocol::McpTool; - let tool = McpTool { name: "simple_tool".to_string(), description: "A simple tool".to_string(), @@ -793,72 +729,81 @@ mod tests { assert!(!tool.requires_approval()); } - // Regression tests for #263: HTML error bodies must not propagate raw - // markup through the error chain into the web UI. - - #[test] - fn test_sanitize_error_body_strips_html_tags() { - let html = - r#"

422 Error

Invalid token

"#; - let result = sanitize_error_body(html); - assert!(!result.contains('<'), "HTML tags must be stripped"); - assert!(!result.contains('>'), "HTML tags must be stripped"); - assert!(result.contains("422 Error")); - assert!(result.contains("Invalid token")); + /// Mock transport for testing transport abstraction behavior. + struct MockTransport { + supports_http: bool, + responses: std::sync::Mutex>, + recorded_headers: std::sync::Mutex>>, } - #[test] - fn test_sanitize_error_body_truncates_large_html_page() { - let html = format!( - "

{}

", - "error detail ".repeat(50) + impl MockTransport { + fn new(supports_http: bool, responses: Vec) -> Self { + Self { + supports_http, + responses: std::sync::Mutex::new(responses), + recorded_headers: std::sync::Mutex::new(Vec::new()), + } + } + fn recorded_headers(&self) -> Vec> { + self.recorded_headers.lock().unwrap().clone() + } + } + + #[async_trait] + impl McpTransport for MockTransport { + async fn send( + &self, + _request: &McpRequest, + headers: &HashMap, + ) -> Result { + self.recorded_headers.lock().unwrap().push(headers.clone()); + let mut responses = self.responses.lock().unwrap(); + if responses.is_empty() { + return Err(ToolError::ExternalService( + "No more mock responses".to_string(), + )); + } + Ok(responses.remove(0)) + } + async fn shutdown(&self) -> Result<(), ToolError> { + Ok(()) + } + fn supports_http_features(&self) -> bool { + self.supports_http + } + } + + #[tokio::test] + async fn test_non_http_transport_skips_401_retry() { + let response = McpResponse { + jsonrpc: "2.0".to_string(), + id: 1, + result: Some(serde_json::json!({"tools": []})), + error: None, + }; + let transport = Arc::new(MockTransport::new(false, vec![response])); + let client = McpClient::new_with_transport( + "test-stdio", + transport.clone(), + None, + None, + "default", + None, ); - let result = sanitize_error_body(&html); - assert!(result.contains("...")); - assert!(result.contains("bytes total)")); - assert!(!result.contains('<')); + let result = client.list_tools().await; + assert!(result.is_ok()); + assert_eq!(result.unwrap().len(), 0); + let headers = transport.recorded_headers(); + assert_eq!(headers.len(), 1); + assert!(!headers[0].contains_key("Authorization")); + assert!(!headers[0].contains_key("Mcp-Session-Id")); } - #[test] - fn test_sanitize_error_body_passes_short_plain_text() { - assert_eq!(sanitize_error_body("Not Found"), "Not Found"); - } - - #[test] - fn test_sanitize_error_body_truncates_long_plain_text() { - let long = "x".repeat(300); - let result = sanitize_error_body(&long); - assert!(result.contains("...")); - assert!(result.contains("300 bytes total)")); - } - - #[test] - fn test_sanitize_error_body_multibyte_no_panic() { - // 300 CJK characters = 900 bytes; truncation must land on a - // char boundary, not in the middle of a multi-byte sequence. - let cjk = "错误".repeat(150); - let result = sanitize_error_body(&cjk); - assert!(result.contains("...")); - // Must be valid UTF-8 (would have panicked otherwise). - assert!(result.is_char_boundary(result.len())); - } - - #[test] - fn test_sanitize_error_body_strips_uppercase_html() { - let html = "

500 Internal Server Error

"; - let result = sanitize_error_body(html); - assert!( - !result.contains('<'), - "uppercase HTML tags must be stripped" - ); - assert!(result.contains("500 Internal Server Error")); - } - - #[test] - fn test_sanitize_error_body_preserves_angle_brackets_in_non_html() { - // Text with < and > that is NOT an HTML document should be - // left untouched (e.g. log lines, comparison expressions). - let text = "value < 10 and value > 0"; - assert_eq!(sanitize_error_body(text), text); + #[tokio::test] + async fn test_transport_supports_http_features_accessor() { + let http_transport = HttpMcpTransport::new("http://localhost:8080", "test"); + assert!(http_transport.supports_http_features()); + let mock_non_http = MockTransport::new(false, vec![]); + assert!(!mock_non_http.supports_http_features()); } } diff --git a/src/tools/mcp/config.rs b/src/tools/mcp/config.rs index 784f0aa2..7dd4be57 100644 --- a/src/tools/mcp/config.rs +++ b/src/tools/mcp/config.rs @@ -12,6 +12,24 @@ use tokio::fs; use crate::bootstrap::ironclaw_base_dir; use crate::tools::tool::ToolError; +/// Transport configuration for an MCP server. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "transport", rename_all = "lowercase")] +pub enum McpTransportConfig { + /// HTTP/HTTPS transport (uses the `url` field on McpServerConfig). + Http, + /// Stdio transport — spawns a child process. + Stdio { + command: String, + #[serde(default)] + args: Vec, + #[serde(default)] + env: HashMap, + }, + /// Unix domain socket transport. + Unix { socket_path: String }, +} + /// Configuration for connecting to a remote MCP server. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct McpServerConfig { @@ -21,6 +39,14 @@ pub struct McpServerConfig { /// Server URL (must be HTTPS for remote servers). pub url: String, + /// Transport configuration. If `None`, defaults to Http using `url`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub transport: Option, + + /// Custom headers to include in every HTTP request. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub headers: HashMap, + /// OAuth configuration (if server requires authentication). #[serde(skip_serializing_if = "Option::is_none")] pub oauth: Option, @@ -44,6 +70,45 @@ impl McpServerConfig { Self { name: name.into(), url: url.into(), + transport: None, + headers: HashMap::new(), + oauth: None, + enabled: true, + description: None, + } + } + + /// Create a new stdio transport MCP server configuration. + pub fn new_stdio( + name: impl Into, + command: impl Into, + args: Vec, + env: HashMap, + ) -> Self { + Self { + name: name.into(), + url: String::new(), + transport: Some(McpTransportConfig::Stdio { + command: command.into(), + args, + env, + }), + headers: HashMap::new(), + oauth: None, + enabled: true, + description: None, + } + } + + /// Create a new Unix socket transport MCP server configuration. + pub fn new_unix(name: impl Into, socket_path: impl Into) -> Self { + Self { + name: name.into(), + url: String::new(), + transport: Some(McpTransportConfig::Unix { + socket_path: socket_path.into(), + }), + headers: HashMap::new(), oauth: None, enabled: true, description: None, @@ -62,6 +127,25 @@ impl McpServerConfig { self } + /// Set custom headers. + pub fn with_headers(mut self, headers: HashMap) -> Self { + self.headers = headers; + self + } + + /// Get the effective transport type. + pub fn effective_transport(&self) -> EffectiveTransport<'_> { + match &self.transport { + Some(McpTransportConfig::Http) | None => EffectiveTransport::Http, + Some(McpTransportConfig::Stdio { command, args, env }) => { + EffectiveTransport::Stdio { command, args, env } + } + Some(McpTransportConfig::Unix { socket_path }) => { + EffectiveTransport::Unix { socket_path } + } + } + } + /// Validate the server configuration. pub fn validate(&self) -> Result<(), ConfigError> { if self.name.is_empty() { @@ -70,19 +154,38 @@ impl McpServerConfig { }); } - if self.url.is_empty() { - return Err(ConfigError::InvalidConfig { - reason: "Server URL cannot be empty".to_string(), - }); - } + match self.effective_transport() { + EffectiveTransport::Http => { + if self.url.is_empty() { + return Err(ConfigError::InvalidConfig { + reason: "Server URL cannot be empty".to_string(), + }); + } - // Remote servers must use HTTPS (localhost is allowed for development) - let url_lower = self.url.to_lowercase(); - let is_localhost = url_lower.contains("localhost") || url_lower.contains("127.0.0.1"); - if !is_localhost && !url_lower.starts_with("https://") { - return Err(ConfigError::InvalidConfig { - reason: "Remote MCP servers must use HTTPS".to_string(), - }); + // Remote servers must use HTTPS (localhost is allowed for development) + let url_lower = self.url.to_lowercase(); + let is_localhost = + url_lower.contains("localhost") || url_lower.contains("127.0.0.1"); + if !is_localhost && !url_lower.starts_with("https://") { + return Err(ConfigError::InvalidConfig { + reason: "Remote MCP servers must use HTTPS".to_string(), + }); + } + } + EffectiveTransport::Stdio { command, .. } => { + if command.is_empty() { + return Err(ConfigError::InvalidConfig { + reason: "Stdio transport command cannot be empty".to_string(), + }); + } + } + EffectiveTransport::Unix { socket_path } => { + if socket_path.is_empty() { + return Err(ConfigError::InvalidConfig { + reason: "Unix socket path cannot be empty".to_string(), + }); + } + } } Ok(()) @@ -92,7 +195,14 @@ impl McpServerConfig { /// /// Returns true if OAuth is pre-configured OR if this is a remote HTTPS server /// (which likely supports Dynamic Client Registration even without pre-configured OAuth). + /// + /// Non-HTTP transports (stdio, unix) never require auth. pub fn requires_auth(&self) -> bool { + // Non-HTTP transports don't use HTTP auth + if !matches!(self.effective_transport(), EffectiveTransport::Http) { + return false; + } + if self.oauth.is_some() { return true; } @@ -426,6 +536,20 @@ fn is_localhost_url(url: &str) -> bool { } } +/// Resolved transport type (borrows from config). +#[derive(Debug)] +pub enum EffectiveTransport<'a> { + Http, + Stdio { + command: &'a str, + args: &'a [String], + env: &'a HashMap, + }, + Unix { + socket_path: &'a str, + }, +} + #[cfg(test)] mod tests { use super::*; @@ -593,4 +717,250 @@ mod tests { let config = McpServerConfig::new("bad", "http://mcp.example.com"); assert!(!config.requires_auth()); } + + #[test] + fn test_stdio_config_creation() { + let env = HashMap::from([("PATH".to_string(), "/usr/bin".to_string())]); + let config = McpServerConfig::new_stdio( + "my-server", + "npx", + vec!["-y".to_string(), "@modelcontextprotocol/server".to_string()], + env.clone(), + ); + + assert_eq!(config.name, "my-server"); + assert!(config.url.is_empty()); + assert!(config.enabled); + assert!(config.oauth.is_none()); + assert!(config.headers.is_empty()); + + match &config.transport { + Some(McpTransportConfig::Stdio { + command, + args, + env: e, + }) => { + assert_eq!(command, "npx"); + assert_eq!( + args, + &["-y".to_string(), "@modelcontextprotocol/server".to_string()] + ); + assert_eq!(e, &env); + } + other => panic!("Expected Stdio transport, got {:?}", other), + } + } + + #[test] + fn test_unix_config_creation() { + let config = McpServerConfig::new_unix("local-server", "/tmp/mcp.sock"); + + assert_eq!(config.name, "local-server"); + assert!(config.url.is_empty()); + assert!(config.enabled); + + match &config.transport { + Some(McpTransportConfig::Unix { socket_path }) => { + assert_eq!(socket_path, "/tmp/mcp.sock"); + } + other => panic!("Expected Unix transport, got {:?}", other), + } + } + + #[test] + fn test_stdio_validation() { + // Valid stdio config + let config = McpServerConfig::new_stdio("server", "npx", vec![], HashMap::new()); + assert!(config.validate().is_ok()); + + // Invalid: empty command + let config = McpServerConfig::new_stdio("server", "", vec![], HashMap::new()); + assert!(config.validate().is_err()); + let err = config.validate().unwrap_err().to_string(); + assert!( + err.contains("command"), + "Error should mention command: {}", + err + ); + + // Invalid: empty name + let config = McpServerConfig::new_stdio("", "npx", vec![], HashMap::new()); + assert!(config.validate().is_err()); + } + + #[test] + fn test_unix_validation() { + // Valid unix config + let config = McpServerConfig::new_unix("server", "/tmp/mcp.sock"); + assert!(config.validate().is_ok()); + + // Invalid: empty socket path + let config = McpServerConfig::new_unix("server", ""); + assert!(config.validate().is_err()); + let err = config.validate().unwrap_err().to_string(); + assert!( + err.contains("socket"), + "Error should mention socket: {}", + err + ); + + // Invalid: empty name + let config = McpServerConfig::new_unix("", "/tmp/mcp.sock"); + assert!(config.validate().is_err()); + } + + #[test] + fn test_requires_auth_stdio_never() { + // Stdio transport should never require auth, even with OAuth configured + let mut config = McpServerConfig::new_stdio("server", "npx", vec![], HashMap::new()); + assert!(!config.requires_auth()); + + // Even if OAuth is set, stdio doesn't use HTTP auth + config.oauth = Some(OAuthConfig::new("client-123")); + assert!(!config.requires_auth()); + } + + #[test] + fn test_requires_auth_unix_never() { + // Unix transport should never require auth + let mut config = McpServerConfig::new_unix("server", "/tmp/mcp.sock"); + assert!(!config.requires_auth()); + + config.oauth = Some(OAuthConfig::new("client-123")); + assert!(!config.requires_auth()); + } + + #[test] + fn test_custom_headers() { + let headers = HashMap::from([ + ("X-Api-Key".to_string(), "secret".to_string()), + ("Authorization".to_string(), "Bearer token".to_string()), + ]); + let config = + McpServerConfig::new("server", "https://mcp.example.com").with_headers(headers.clone()); + + assert_eq!(config.headers, headers); + assert_eq!(config.headers.get("X-Api-Key").unwrap(), "secret"); + } + + #[test] + fn test_transport_config_serde_http() { + let transport = McpTransportConfig::Http; + let json = serde_json::to_string(&transport).unwrap(); + assert!(json.contains("\"transport\":\"http\"")); + + let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap(); + assert!(matches!(parsed, McpTransportConfig::Http)); + } + + #[test] + fn test_transport_config_serde_stdio() { + let transport = McpTransportConfig::Stdio { + command: "npx".to_string(), + args: vec!["-y".to_string(), "server".to_string()], + env: HashMap::from([("KEY".to_string(), "val".to_string())]), + }; + let json = serde_json::to_string(&transport).unwrap(); + assert!(json.contains("\"transport\":\"stdio\"")); + assert!(json.contains("\"command\":\"npx\"")); + + let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap(); + match parsed { + McpTransportConfig::Stdio { command, args, env } => { + assert_eq!(command, "npx"); + assert_eq!(args, vec!["-y".to_string(), "server".to_string()]); + assert_eq!(env.get("KEY").unwrap(), "val"); + } + other => panic!("Expected Stdio, got {:?}", other), + } + } + + #[test] + fn test_transport_config_serde_unix() { + let transport = McpTransportConfig::Unix { + socket_path: "/tmp/mcp.sock".to_string(), + }; + let json = serde_json::to_string(&transport).unwrap(); + assert!(json.contains("\"transport\":\"unix\"")); + assert!(json.contains("\"socket_path\":\"/tmp/mcp.sock\"")); + + let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap(); + match parsed { + McpTransportConfig::Unix { socket_path } => { + assert_eq!(socket_path, "/tmp/mcp.sock"); + } + other => panic!("Expected Unix, got {:?}", other), + } + } + + #[test] + fn test_backward_compat_no_transport_field() { + // Existing configs without transport field should still deserialize + let json = r#"{ + "name": "notion", + "url": "https://mcp.notion.com", + "enabled": true + }"#; + let config: McpServerConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.name, "notion"); + assert_eq!(config.url, "https://mcp.notion.com"); + assert!(config.transport.is_none()); + assert!(config.headers.is_empty()); + assert!(matches!( + config.effective_transport(), + EffectiveTransport::Http + )); + } + + #[test] + fn test_config_roundtrip_with_transport() { + // Test full roundtrip with stdio transport + let config = McpServerConfig::new_stdio( + "test-server", + "node", + vec!["server.js".to_string()], + HashMap::from([("NODE_ENV".to_string(), "production".to_string())]), + ) + .with_description("A test server"); + + let json = serde_json::to_string_pretty(&config).unwrap(); + let parsed: McpServerConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(parsed.name, "test-server"); + assert!(parsed.url.is_empty()); + assert_eq!(parsed.description.as_deref(), Some("A test server")); + + match &parsed.transport { + Some(McpTransportConfig::Stdio { command, args, env }) => { + assert_eq!(command, "node"); + assert_eq!(args, &["server.js".to_string()]); + assert_eq!(env.get("NODE_ENV").unwrap(), "production"); + } + other => panic!("Expected Stdio transport, got {:?}", other), + } + + // Test full roundtrip with unix transport + let config = McpServerConfig::new_unix("unix-server", "/var/run/mcp.sock"); + let json = serde_json::to_string_pretty(&config).unwrap(); + let parsed: McpServerConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(parsed.name, "unix-server"); + match &parsed.transport { + Some(McpTransportConfig::Unix { socket_path }) => { + assert_eq!(socket_path, "/var/run/mcp.sock"); + } + other => panic!("Expected Unix transport, got {:?}", other), + } + + // Test roundtrip with HTTP + headers + let headers = HashMap::from([("X-Custom".to_string(), "value".to_string())]); + let config = + McpServerConfig::new("http-server", "https://mcp.example.com").with_headers(headers); + let json = serde_json::to_string_pretty(&config).unwrap(); + let parsed: McpServerConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(parsed.name, "http-server"); + assert!(parsed.transport.is_none()); + assert_eq!(parsed.headers.get("X-Custom").unwrap(), "value"); + } } diff --git a/src/tools/mcp/http_transport.rs b/src/tools/mcp/http_transport.rs new file mode 100644 index 00000000..2a51ae63 --- /dev/null +++ b/src/tools/mcp/http_transport.rs @@ -0,0 +1,386 @@ +//! HTTP transport for MCP servers. +//! +//! Implements the Streamable HTTP transport, communicating with MCP servers +//! over HTTP POST with JSON and SSE response support. + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::tools::mcp::protocol::{McpRequest, McpResponse}; +use crate::tools::mcp::session::McpSessionManager; +use crate::tools::mcp::transport::McpTransport; +use crate::tools::tool::ToolError; + +/// MCP transport that communicates with a server over HTTP. +/// +/// Sends JSON-RPC requests as HTTP POST with `Content-Type: application/json` +/// and accepts either JSON or SSE (`text/event-stream`) responses. Optionally +/// manages session IDs via [`McpSessionManager`] and supports custom headers. +pub struct HttpMcpTransport { + server_url: String, + server_name: String, + http_client: reqwest::Client, + session_manager: Option>, + custom_headers: HashMap, +} + +impl HttpMcpTransport { + /// Create a new HTTP transport for the given server URL. + pub fn new(server_url: impl Into, server_name: impl Into) -> Self { + Self { + server_url: server_url.into(), + server_name: server_name.into(), + // reqwest::Client::builder().build() only fails if the TLS backend + // cannot initialize, which does not happen with the default rustls + // feature set. Panic is acceptable here (same as reqwest's own + // `Client::new()`). + http_client: reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .expect("Failed to create HTTP client"), + session_manager: None, + custom_headers: HashMap::new(), + } + } + + /// Attach a session manager for Mcp-Session-Id tracking. + pub fn with_session_manager(mut self, session_manager: Arc) -> Self { + self.session_manager = Some(session_manager); + self + } + + /// Set custom headers that will be sent with every request. + #[cfg(test)] + pub fn with_custom_headers(mut self, headers: HashMap) -> Self { + self.custom_headers = headers; + self + } + + /// Get the server URL. + #[cfg(test)] + pub(crate) fn server_url(&self) -> &str { + &self.server_url + } + + /// Get the session manager, if one is configured. + #[cfg(test)] + pub(crate) fn session_manager(&self) -> Option<&Arc> { + self.session_manager.as_ref() + } +} + +#[async_trait] +impl McpTransport for HttpMcpTransport { + async fn send( + &self, + request: &McpRequest, + headers: &HashMap, + ) -> Result { + // Build the HTTP request. + let mut req_builder = self + .http_client + .post(&self.server_url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .json(request); + + // Apply custom headers configured on the transport. + for (key, value) in &self.custom_headers { + req_builder = req_builder.header(key.as_str(), value.as_str()); + } + + // Apply per-request headers (e.g. Authorization, Mcp-Session-Id). + for (key, value) in headers { + req_builder = req_builder.header(key.as_str(), value.as_str()); + } + + // Send the request. + let response = req_builder.send().await.map_err(|e| { + let mut chain = format!("[{}] MCP HTTP request failed: {}", self.server_name, e); + let mut source = std::error::Error::source(&e); + while let Some(cause) = source { + chain.push_str(&format!(" -> {}", cause)); + source = cause.source(); + } + ToolError::ExternalService(chain) + })?; + + // Extract session ID from response headers before consuming the body. + if let Some(ref session_manager) = self.session_manager + && let Some(session_id) = response + .headers() + .get("Mcp-Session-Id") + .and_then(|v| v.to_str().ok()) + { + session_manager + .update_session_id(&self.server_name, Some(session_id.to_string())) + .await; + } + + // Handle error status codes. + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + let sanitized = sanitize_error_body(&body); + return Err(ToolError::ExternalService(format!( + "[{}] MCP server returned status: {} - {}", + self.server_name, status, sanitized + ))); + } + + // Determine response format from Content-Type. + let content_type = response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); + + if content_type.contains("text/event-stream") { + self.parse_sse_response(response).await + } else { + response.json().await.map_err(|e| { + ToolError::ExternalService(format!( + "[{}] Failed to parse MCP response: {}", + self.server_name, e + )) + }) + } + } + + async fn shutdown(&self) -> Result<(), ToolError> { + // HTTP transport is stateless; nothing to shut down. + Ok(()) + } + + fn supports_http_features(&self) -> bool { + true + } +} + +impl HttpMcpTransport { + /// Parse a Server-Sent Events response, returning the first valid JSON-RPC + /// `data:` line as an [`McpResponse`]. + async fn parse_sse_response( + &self, + response: reqwest::Response, + ) -> Result { + use futures::StreamExt; + + const MAX_SSE_BUFFER: usize = 10 * 1024 * 1024; // 10 MB + + let mut stream = response.bytes_stream(); + let mut buffer = String::new(); + + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|e| { + ToolError::ExternalService(format!( + "[{}] Failed to read SSE chunk: {}", + self.server_name, e + )) + })?; + + buffer.push_str(&String::from_utf8_lossy(&chunk)); + + if buffer.len() > MAX_SSE_BUFFER { + return Err(ToolError::ExternalService(format!( + "[{}] SSE response exceeded {} byte limit", + self.server_name, MAX_SSE_BUFFER + ))); + } + + // Process only complete lines (terminated by \n). The last + // element of split('\n') may be an incomplete line; keep it + // in the buffer for the next chunk. + let mut remaining_start = 0; + let bytes = buffer.as_bytes(); + for (i, &b) in bytes.iter().enumerate() { + if b == b'\n' { + let line = &buffer[remaining_start..i]; + remaining_start = i + 1; + + if let Some(json_str) = line.strip_prefix("data: ") + && let Ok(response) = serde_json::from_str::(json_str) + { + return Ok(response); + } + } + } + // Keep only the unprocessed trailing fragment. + if remaining_start > 0 { + buffer = buffer[remaining_start..].to_string(); + } + } + + // Process any remaining data without a trailing newline. + if let Some(json_str) = buffer.strip_prefix("data: ") + && let Ok(response) = serde_json::from_str::(json_str.trim()) + { + return Ok(response); + } + + Err(ToolError::ExternalService(format!( + "[{}] No valid data in SSE response: {}", + self.server_name, buffer + ))) + } +} + +/// Sanitize an HTTP error body for safe inclusion in error messages. +/// +/// When the body looks like a full HTML document (` String { + const MAX_CHARS: usize = 200; + + // Only strip tags when the body looks like a full HTML document. + // Plain text that happens to contain `<` / `>` (e.g. log lines, + // comparison expressions) is left untouched. + let lower = body.to_ascii_lowercase(); + let is_html_document = lower.contains("' { + (out, false) + } else if !in_tag { + out.push(c); + (out, false) + } else { + (out, true) + } + }) + .0; + stripped.split_whitespace().collect::>().join(" ") + } else { + body.to_string() + }; + + // Truncate at a char boundary (safe for multi-byte UTF-8). + if text.chars().count() > MAX_CHARS { + let byte_offset = text + .char_indices() + .nth(MAX_CHARS) + .map(|(i, _)| i) + .unwrap_or(text.len()); + format!("{}... ({} bytes total)", &text[..byte_offset], body.len()) + } else { + text + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sanitize_error_body_strips_html_tags() { + let html = + r#"

422 Error

Invalid token

"#; + let result = sanitize_error_body(html); + assert!(!result.contains('<'), "HTML tags must be stripped"); + assert!(!result.contains('>'), "HTML tags must be stripped"); + assert!(result.contains("422 Error")); + assert!(result.contains("Invalid token")); + } + + #[test] + fn test_sanitize_error_body_truncates_large_html_page() { + let html = format!( + "

{}

", + "error detail ".repeat(50) + ); + let result = sanitize_error_body(&html); + assert!(result.contains("...")); + assert!(result.contains("bytes total)")); + assert!(!result.contains('<')); + } + + #[test] + fn test_sanitize_error_body_passes_short_plain_text() { + assert_eq!(sanitize_error_body("Not Found"), "Not Found"); + } + + #[test] + fn test_sanitize_error_body_truncates_long_plain_text() { + let long = "x".repeat(300); + let result = sanitize_error_body(&long); + assert!(result.contains("...")); + assert!(result.contains("300 bytes total)")); + } + + #[test] + fn test_sanitize_error_body_multibyte_no_panic() { + // 300 CJK characters = 900 bytes; truncation must land on a + // char boundary, not in the middle of a multi-byte sequence. + let cjk = "错误".repeat(150); + let result = sanitize_error_body(&cjk); + assert!(result.contains("...")); + // Must be valid UTF-8 (would have panicked otherwise). + assert!(result.is_char_boundary(result.len())); + } + + #[test] + fn test_sanitize_error_body_strips_uppercase_html() { + let html = "

500 Internal Server Error

"; + let result = sanitize_error_body(html); + assert!( + !result.contains('<'), + "uppercase HTML tags must be stripped" + ); + assert!(result.contains("500 Internal Server Error")); + } + + #[test] + fn test_sanitize_error_body_preserves_angle_brackets_in_non_html() { + let text = "value < 10 and value > 0"; + assert_eq!(sanitize_error_body(text), text); + } + + #[test] + fn test_sanitize_error_body_empty_string() { + assert_eq!(sanitize_error_body(""), ""); + } + + #[test] + fn test_new_creates_transport() { + let transport = HttpMcpTransport::new("http://localhost:8080", "test"); + assert_eq!(transport.server_url(), "http://localhost:8080"); + assert!(transport.session_manager().is_none()); + assert!(transport.custom_headers.is_empty()); + } + + #[test] + fn test_supports_http_features() { + let http_transport = HttpMcpTransport::new("http://localhost:8080", "test"); + assert!(http_transport.supports_http_features()); + } + + #[test] + fn test_with_session_manager() { + let session_manager = Arc::new(McpSessionManager::new()); + let transport = HttpMcpTransport::new("http://localhost:8080", "test") + .with_session_manager(session_manager.clone()); + assert!(transport.session_manager().is_some()); + } + + #[test] + fn test_with_custom_headers() { + let mut headers = HashMap::new(); + headers.insert("X-Custom".to_string(), "value".to_string()); + let transport = + HttpMcpTransport::new("http://localhost:8080", "test").with_custom_headers(headers); + assert_eq!(transport.custom_headers.get("X-Custom").unwrap(), "value"); + } +} diff --git a/src/tools/mcp/mod.rs b/src/tools/mcp/mod.rs index 8a8b92ea..8ab107c9 100644 --- a/src/tools/mcp/mod.rs +++ b/src/tools/mcp/mod.rs @@ -4,6 +4,8 @@ //! additional capabilities through a standardized protocol. //! //! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers. +//! Transport options include HTTP (Streamable HTTP / SSE), stdio (subprocess), +//! and Unix domain sockets. //! //! ## Usage //! @@ -29,11 +31,19 @@ pub mod auth; mod client; pub mod config; +pub(crate) mod http_transport; +pub(crate) mod process; mod protocol; pub mod session; +pub(crate) mod stdio_transport; +pub(crate) mod transport; +#[cfg(unix)] +pub(crate) mod unix_transport; pub use auth::{is_authenticated, refresh_access_token}; pub use client::McpClient; pub use config::{McpServerConfig, McpServersFile, OAuthConfig}; +pub use process::McpProcessManager; pub use protocol::{InitializeResult, McpRequest, McpResponse, McpTool}; pub use session::McpSessionManager; +pub use transport::McpTransport; diff --git a/src/tools/mcp/process.rs b/src/tools/mcp/process.rs new file mode 100644 index 00000000..85bc5715 --- /dev/null +++ b/src/tools/mcp/process.rs @@ -0,0 +1,206 @@ +//! MCP stdio process manager. +//! +//! Manages the lifecycle of MCP servers running as child processes. +//! Handles spawning, shutdown, and crash recovery with exponential backoff. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::RwLock; + +use crate::tools::mcp::stdio_transport::StdioMcpTransport; +use crate::tools::mcp::transport::McpTransport; +use crate::tools::tool::ToolError; + +/// Configuration for spawning a stdio MCP server. +#[derive(Debug, Clone)] +pub struct StdioSpawnConfig { + pub command: String, + pub args: Vec, + pub env: HashMap, +} + +/// Manages stdio MCP server processes. +/// +/// Handles spawning, tracking, and shutdown of child processes. +pub struct McpProcessManager { + transports: RwLock>>, + configs: RwLock>, +} + +impl McpProcessManager { + pub fn new() -> Self { + Self { + transports: RwLock::new(HashMap::new()), + configs: RwLock::new(HashMap::new()), + } + } + + /// Spawn a new stdio MCP server process. + pub async fn spawn_stdio( + &self, + name: impl Into, + command: impl Into, + args: Vec, + env: HashMap, + ) -> Result, ToolError> { + let name = name.into(); + let command = command.into(); + + // Store config for potential restart + self.configs.write().await.insert( + name.clone(), + StdioSpawnConfig { + command: command.clone(), + args: args.clone(), + env: env.clone(), + }, + ); + + let transport = Arc::new(StdioMcpTransport::spawn(&name, &command, args, env).await?); + + self.transports + .write() + .await + .insert(name, Arc::clone(&transport)); + + Ok(transport) + } + + /// Get a transport by server name. + pub async fn get(&self, name: &str) -> Option> { + self.transports.read().await.get(name).cloned() + } + + /// Shut down all managed transports. + pub async fn shutdown_all(&self) { + let transports: Vec<(String, Arc)> = { + let mut map = self.transports.write().await; + map.drain().collect() + }; + + for (name, transport) in transports { + if let Err(e) = transport.shutdown().await { + tracing::warn!("Failed to shut down MCP stdio server '{}': {}", name, e); + } + } + } + + /// Shut down a specific transport by name. + pub async fn shutdown(&self, name: &str) -> Result<(), ToolError> { + let transport = self.transports.write().await.remove(name); + + if let Some(transport) = transport { + transport.shutdown().await?; + } + + self.configs.write().await.remove(name); + Ok(()) + } + + /// Attempt to restart a crashed transport with exponential backoff. + /// + /// Tries up to 5 times with delays of 1s, 2s, 4s, 8s, 16s (total: 31s max wait). + pub async fn try_restart(&self, name: &str) -> Result, ToolError> { + let config = self + .configs + .read() + .await + .get(name) + .cloned() + .ok_or_else(|| { + ToolError::ExternalService(format!( + "No spawn config for MCP server '{}', cannot restart", + name + )) + })?; + + // Shut down and remove old transport to avoid orphaning a wedged process. + if let Some(old_transport) = self.transports.write().await.remove(name) { + let _ = old_transport.shutdown().await; + } + + let max_retries = 5; + let mut last_err = None; + + for attempt in 0..max_retries { + let delay = Duration::from_secs(1 << attempt); + tokio::time::sleep(delay).await; + + match StdioMcpTransport::spawn( + name, + &config.command, + config.args.clone(), + config.env.clone(), + ) + .await + { + Ok(transport) => { + let transport = Arc::new(transport); + self.transports + .write() + .await + .insert(name.to_string(), Arc::clone(&transport)); + tracing::info!( + "MCP stdio server '{}' restarted after {} attempt(s)", + name, + attempt + 1 + ); + return Ok(transport); + } + Err(e) => { + tracing::warn!( + "Restart attempt {}/{} for MCP server '{}' failed: {}", + attempt + 1, + max_retries, + name, + e + ); + last_err = Some(e); + } + } + } + + Err(last_err.unwrap_or_else(|| { + ToolError::ExternalService(format!( + "Failed to restart MCP server '{}' after {} attempts", + name, max_retries + )) + })) + } + + /// Get names of all managed transports. + pub async fn managed_servers(&self) -> Vec { + self.transports.read().await.keys().cloned().collect() + } +} + +impl Default for McpProcessManager { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_creates_empty_manager() { + let _manager = McpProcessManager::new(); + } + + #[tokio::test] + async fn test_managed_servers_returns_empty_list_initially() { + let manager = McpProcessManager::new(); + let servers = manager.managed_servers().await; + assert!(servers.is_empty()); + } + + #[tokio::test] + async fn test_shutdown_all_on_empty_manager_does_not_panic() { + let manager = McpProcessManager::new(); + manager.shutdown_all().await; + } +} diff --git a/src/tools/mcp/protocol.rs b/src/tools/mcp/protocol.rs index d5d9f052..cb55ef94 100644 --- a/src/tools/mcp/protocol.rs +++ b/src/tools/mcp/protocol.rs @@ -120,10 +120,15 @@ impl McpRequest { } /// Create an initialized notification (sent after initialize). + /// + /// Note: JSON-RPC 2.0 notifications should omit the `id` field entirely. + /// We set `id: 0` because `McpRequest` uses `u64` (not `Option`). + /// Most MCP servers tolerate this; a proper fix would use a separate + /// `McpNotification` type or make `id` optional with `skip_serializing_if`. pub fn initialized_notification() -> Self { Self { jsonrpc: "2.0".to_string(), - id: 0, // Notifications don't have IDs, but we need one for the struct + id: 0, method: "notifications/initialized".to_string(), params: None, } diff --git a/src/tools/mcp/stdio_transport.rs b/src/tools/mcp/stdio_transport.rs new file mode 100644 index 00000000..c1b79762 --- /dev/null +++ b/src/tools/mcp/stdio_transport.rs @@ -0,0 +1,263 @@ +//! Stdio transport for MCP servers. +//! +//! Spawns a child process and communicates via stdin/stdout using +//! newline-delimited JSON-RPC. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use tokio::io::BufReader; +use tokio::process::{Child, Command}; +use tokio::sync::{Mutex, oneshot}; +use tokio::task::JoinHandle; + +use crate::tools::mcp::protocol::{McpRequest, McpResponse}; +use crate::tools::mcp::transport::{McpTransport, spawn_jsonrpc_reader, write_jsonrpc_line}; +use crate::tools::tool::ToolError; + +/// MCP transport that communicates with a child process over stdin/stdout. +/// +/// The child process is spawned with piped stdin/stdout/stderr. Requests are +/// written as newline-delimited JSON to stdin, and responses are read from +/// stdout by a background reader task. Stderr is drained to tracing logs. +pub struct StdioMcpTransport { + server_name: String, + stdin: Arc>, + pending: Arc>>>, + reader_handle: Mutex>>, + stderr_handle: Mutex>>, + child: Arc>, +} + +impl StdioMcpTransport { + /// Spawn a child process and create a stdio transport. + /// + /// # Arguments + /// + /// * `name` - Human-readable server name for logging. + /// * `command` - The command to execute. + /// * `args` - Command-line arguments. + /// * `env` - Additional environment variables to set. + pub async fn spawn( + name: impl Into, + command: &str, + args: impl IntoIterator>, + env: impl IntoIterator, impl AsRef)>, + ) -> Result { + let server_name = name.into(); + + let mut cmd = Command::new(command); + cmd.args(args) + .envs(env) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()); + + let mut child = cmd.spawn().map_err(|e| { + ToolError::ExternalService(format!( + "[{}] Failed to spawn MCP server '{}': {}", + server_name, command, e + )) + })?; + + let stdin = child.stdin.take().ok_or_else(|| { + ToolError::ExternalService(format!( + "[{}] Failed to capture stdin of MCP server", + server_name + )) + })?; + + let stdout = child.stdout.take().ok_or_else(|| { + ToolError::ExternalService(format!( + "[{}] Failed to capture stdout of MCP server", + server_name + )) + })?; + + let stderr = child.stderr.take().ok_or_else(|| { + ToolError::ExternalService(format!( + "[{}] Failed to capture stderr of MCP server", + server_name + )) + })?; + + let pending: Arc>>> = + Arc::new(Mutex::new(HashMap::new())); + + let reader = BufReader::new(stdout); + let reader_handle = spawn_jsonrpc_reader(reader, pending.clone(), server_name.clone()); + + let stderr_name = server_name.clone(); + let stderr_handle = tokio::spawn(async move { + use tokio::io::{AsyncBufReadExt, BufReader as TokioBufReader}; + + let reader = TokioBufReader::new(stderr); + let mut lines = reader.lines(); + while let Ok(Some(line)) = lines.next_line().await { + tracing::debug!("[{}] stderr: {}", stderr_name, line); + } + }); + + Ok(Self { + server_name, + stdin: Arc::new(Mutex::new(stdin)), + pending, + reader_handle: Mutex::new(Some(reader_handle)), + stderr_handle: Mutex::new(Some(stderr_handle)), + child: Arc::new(Mutex::new(child)), + }) + } +} + +#[async_trait] +impl McpTransport for StdioMcpTransport { + async fn send( + &self, + request: &McpRequest, + _headers: &HashMap, + ) -> Result { + let (tx, rx) = oneshot::channel(); + + // Register the pending response handler before writing the request, + // so we don't miss a fast response from the child. + { + let mut pending = self.pending.lock().await; + pending.insert(request.id, tx); + } + + // Write the request to stdin. + { + let mut stdin = self.stdin.lock().await; + if let Err(e) = write_jsonrpc_line(&mut *stdin, request).await { + // Remove the pending entry on write failure. + let mut pending = self.pending.lock().await; + pending.remove(&request.id); + return Err(e); + } + } + + // Wait for the response with a timeout. + let timeout = Duration::from_secs(30); + match tokio::time::timeout(timeout, rx).await { + Ok(Ok(response)) => Ok(response), + Ok(Err(_)) => { + // Sender was dropped (reader task ended). Clean up pending entry. + let mut pending = self.pending.lock().await; + pending.remove(&request.id); + Err(ToolError::ExternalService(format!( + "[{}] MCP server closed connection before responding to request {}", + self.server_name, request.id + ))) + } + Err(_) => { + // Timeout: remove the pending entry. + let mut pending = self.pending.lock().await; + pending.remove(&request.id); + Err(ToolError::ExternalService(format!( + "[{}] Timeout waiting for response to request {} after {:?}", + self.server_name, request.id, timeout + ))) + } + } + } + + async fn shutdown(&self) -> Result<(), ToolError> { + // Kill the child process. + { + let mut child = self.child.lock().await; + let _ = child.kill().await; + } + + // Abort the reader tasks. + if let Some(handle) = self.reader_handle.lock().await.take() { + handle.abort(); + } + if let Some(handle) = self.stderr_handle.lock().await.take() { + handle.abort(); + } + + // Drain pending requests so waiters wake immediately instead of + // hanging until their 30s timeout. + { + let mut pending = self.pending.lock().await; + pending.clear(); // Dropping senders wakes receivers with Err + } + + tracing::debug!("[{}] Stdio transport shut down", self.server_name); + Ok(()) + } + + fn supports_http_features(&self) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_spawn_nonexistent_command_fails() { + let env: HashMap = HashMap::new(); + let result = StdioMcpTransport::spawn( + "test", + "this-command-does-not-exist-ironclaw-test", + std::iter::empty::<&str>(), + &env, + ) + .await; + + let err = result.err().expect("should be an error").to_string(); + assert!( + err.contains("Failed to spawn"), + "Error should mention spawn failure: {}", + err + ); + } + + #[tokio::test] + async fn test_spawn_and_shutdown() { + let env: HashMap = HashMap::new(); + let transport = + StdioMcpTransport::spawn("test-cat", "cat", std::iter::empty::<&str>(), &env) + .await + .expect("cat should be available"); + + // Verify shutdown completes without error. + transport.shutdown().await.expect("shutdown should succeed"); + } + + #[tokio::test] + async fn test_send_timeout_on_non_jsonrpc_server() { + // Spawn `cat` which echoes input back. Since the echoed input is the + // request (not a response with matching id), it will be ignored by the + // reader and we should hit the timeout. We use a short-lived test so + // we override the 30s timeout expectation by just checking the error type. + let env: HashMap = HashMap::new(); + let transport = + StdioMcpTransport::spawn("test-echo", "cat", std::iter::empty::<&str>(), &env) + .await + .expect("cat should be available"); + + let request = McpRequest::list_tools(999); + let headers = HashMap::new(); + + // The request will be echoed back by `cat`, but it won't parse as a + // valid McpResponse with matching id, so the reader will log a debug + // message and the send will eventually timeout. We don't want to wait + // 30 seconds in tests, so we just verify the transport was created and + // shut it down. + transport.shutdown().await.expect("shutdown should succeed"); + + // Verify that pending map is empty after shutdown. + let pending = transport.pending.lock().await; + assert!(pending.is_empty()); + drop(pending); + + // Verify send after shutdown fails (stdin is closed). + let result = transport.send(&request, &headers).await; + assert!(result.is_err()); + } +} diff --git a/src/tools/mcp/transport.rs b/src/tools/mcp/transport.rs new file mode 100644 index 00000000..98c4a478 --- /dev/null +++ b/src/tools/mcp/transport.rs @@ -0,0 +1,196 @@ +//! Shared MCP transport trait and JSON-RPC framing helpers. +//! +//! Provides the [`McpTransport`] trait that all MCP transports implement, +//! plus `write_jsonrpc_line` and `spawn_jsonrpc_reader` for newline-delimited +//! JSON-RPC over byte streams (used by stdio and unix socket transports). + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::sync::{Mutex, oneshot}; +use tokio::task::JoinHandle; + +use crate::tools::mcp::protocol::{McpRequest, McpResponse}; +use crate::tools::tool::ToolError; + +/// Trait for sending JSON-RPC requests to an MCP server and receiving responses. +/// +/// Implementations handle the underlying transport (HTTP, stdio, unix socket, etc.). +#[async_trait] +pub trait McpTransport: Send + Sync { + /// Send a request and wait for the corresponding response. + /// + /// `headers` are used by HTTP-based transports (e.g., `Mcp-Session-Id`); + /// stream-based transports may ignore them. + async fn send( + &self, + request: &McpRequest, + headers: &HashMap, + ) -> Result; + + /// Shut down the transport, releasing any resources (child processes, connections). + async fn shutdown(&self) -> Result<(), ToolError>; + + /// Whether this transport supports HTTP-specific features like session headers. + fn supports_http_features(&self) -> bool { + false + } +} + +/// Serialize an [`McpRequest`] as a single JSON line and write it to `writer`. +/// +/// The line is terminated with `\n` and the writer is flushed. +pub async fn write_jsonrpc_line( + writer: &mut (impl AsyncWrite + Unpin), + request: &McpRequest, +) -> Result<(), ToolError> { + let json = serde_json::to_string(request).map_err(|e| { + ToolError::ExternalService(format!("Failed to serialize JSON-RPC request: {e}")) + })?; + + writer.write_all(json.as_bytes()).await.map_err(|e| { + ToolError::ExternalService(format!("Failed to write JSON-RPC request: {e}")) + })?; + + writer + .write_all(b"\n") + .await + .map_err(|e| ToolError::ExternalService(format!("Failed to write newline: {e}")))?; + + writer + .flush() + .await + .map_err(|e| ToolError::ExternalService(format!("Failed to flush JSON-RPC writer: {e}")))?; + + Ok(()) +} + +/// Spawn a background task that reads newline-delimited JSON-RPC responses from +/// `reader` and dispatches them to the matching pending sender in `pending`. +/// +/// Each line is parsed as an [`McpResponse`]. If the response has an `id` that +/// matches a pending request, the corresponding [`oneshot::Sender`] is resolved. +/// Parse failures are logged at debug level and skipped. +pub fn spawn_jsonrpc_reader( + reader: R, + pending: Arc>>>, + server_name: String, +) -> JoinHandle<()> { + tokio::spawn(async move { + let mut lines = reader.lines(); + while let Ok(Some(line)) = lines.next_line().await { + let response = match serde_json::from_str::(&line) { + Ok(resp) => resp, + Err(e) => { + // Truncate logged line to avoid leaking sensitive data in large payloads. + let preview: String = line.chars().take(200).collect(); + tracing::debug!( + "[{}] Failed to parse JSON-RPC response: {} — line: {}{}", + server_name, + e, + preview, + if line.len() > 200 { "…" } else { "" } + ); + continue; + } + }; + + let id = response.id; + let mut map = pending.lock().await; + if let Some(tx) = map.remove(&id) { + // Ignore send error — the receiver may have been dropped (timeout). + let _ = tx.send(response); + } else { + tracing::debug!( + "[{}] Received response for unknown request id {}", + server_name, + id + ); + } + } + + tracing::debug!("[{}] JSON-RPC reader finished", server_name); + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_write_jsonrpc_line_serializes_and_flushes() { + let request = McpRequest { + jsonrpc: "2.0".into(), + id: 1, + method: "test/method".into(), + params: None, + }; + + let mut buf = Vec::new(); + write_jsonrpc_line(&mut buf, &request) + .await + .expect("write should succeed"); + + let written = String::from_utf8(buf).expect("should be valid UTF-8"); + assert!(written.ends_with('\n')); + + let parsed: serde_json::Value = + serde_json::from_str(written.trim()).expect("should be valid JSON"); + assert_eq!(parsed["id"], 1); + assert_eq!(parsed["method"], "test/method"); + } + + #[tokio::test] + async fn test_spawn_jsonrpc_reader_dispatches_response() { + let response = McpResponse { + jsonrpc: "2.0".into(), + id: 42, + result: Some(serde_json::json!({"tools": []})), + error: None, + }; + let line = format!("{}\n", serde_json::to_string(&response).unwrap()); + + let reader = std::io::Cursor::new(line.into_bytes()); + let pending: Arc>>> = + Arc::new(Mutex::new(HashMap::new())); + + let (tx, rx) = oneshot::channel(); + { + let mut map = pending.lock().await; + map.insert(42, tx); + } + + let handle = spawn_jsonrpc_reader(reader, pending.clone(), "test".into()); + + let resp = rx.await.expect("should receive response"); + assert_eq!(resp.id, 42); + assert!(resp.result.is_some()); + + handle.await.expect("reader task should finish"); + } + + #[tokio::test] + async fn test_spawn_jsonrpc_reader_skips_invalid_lines() { + let input = "this is not json\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":null}\n"; + let reader = std::io::Cursor::new(input.as_bytes().to_vec()); + let pending: Arc>>> = + Arc::new(Mutex::new(HashMap::new())); + + let (tx, rx) = oneshot::channel(); + { + let mut map = pending.lock().await; + map.insert(7, tx); + } + + let handle = spawn_jsonrpc_reader(reader, pending.clone(), "test".into()); + + let resp = rx + .await + .expect("should receive response despite earlier invalid line"); + assert_eq!(resp.id, 7); + + handle.await.expect("reader task should finish"); + } +} diff --git a/src/tools/mcp/unix_transport.rs b/src/tools/mcp/unix_transport.rs new file mode 100644 index 00000000..07ef7f17 --- /dev/null +++ b/src/tools/mcp/unix_transport.rs @@ -0,0 +1,269 @@ +//! Unix domain socket transport for MCP servers. +//! +//! Connects to an existing Unix socket and communicates using +//! newline-delimited JSON-RPC. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use tokio::io::BufReader; +use tokio::net::UnixStream; +use tokio::sync::{Mutex, oneshot}; +use tokio::task::JoinHandle; + +use crate::tools::mcp::protocol::{McpRequest, McpResponse}; +use crate::tools::mcp::transport::{McpTransport, spawn_jsonrpc_reader, write_jsonrpc_line}; +use crate::tools::tool::ToolError; + +/// MCP transport that communicates over a Unix domain socket. +/// +/// Connects to an existing Unix socket at the given path. Requests are +/// written as newline-delimited JSON to the write half, and responses are +/// read from the read half by a background reader task. +pub struct UnixMcpTransport { + socket_path: PathBuf, + server_name: String, + writer: Arc>>, + pending: Arc>>>, + reader_handle: Mutex>>, +} + +impl UnixMcpTransport { + /// Connect to an existing Unix domain socket and create a transport. + /// + /// # Arguments + /// + /// * `name` - Human-readable server name for logging. + /// * `socket_path` - Path to the Unix domain socket. + pub async fn connect( + name: impl Into, + socket_path: impl AsRef, + ) -> Result { + let server_name = name.into(); + let socket_path = socket_path.as_ref().to_path_buf(); + + let stream = UnixStream::connect(&socket_path).await.map_err(|e| { + ToolError::ExternalService(format!( + "[{}] Failed to connect to Unix socket '{}': {}", + server_name, + socket_path.display(), + e + )) + })?; + + let (read_half, write_half) = tokio::io::split(stream); + + let pending: Arc>>> = + Arc::new(Mutex::new(HashMap::new())); + + let reader = BufReader::new(read_half); + let reader_handle = spawn_jsonrpc_reader(reader, pending.clone(), server_name.clone()); + + Ok(Self { + socket_path, + server_name, + writer: Arc::new(Mutex::new(write_half)), + pending, + reader_handle: Mutex::new(Some(reader_handle)), + }) + } + + /// Get the path to the Unix domain socket. + #[cfg(test)] + pub(crate) fn socket_path(&self) -> &Path { + &self.socket_path + } + + /// Get the server name. + #[cfg(test)] + pub(crate) fn server_name(&self) -> &str { + &self.server_name + } +} + +#[async_trait] +impl McpTransport for UnixMcpTransport { + async fn send( + &self, + request: &McpRequest, + _headers: &HashMap, + ) -> Result { + let (tx, rx) = oneshot::channel(); + + // Register the pending response handler before writing the request, + // so we don't miss a fast response from the server. + { + let mut pending = self.pending.lock().await; + pending.insert(request.id, tx); + } + + // Write the request to the socket. + { + let mut writer = self.writer.lock().await; + if let Err(e) = write_jsonrpc_line(&mut *writer, request).await { + // Remove the pending entry on write failure. + let mut pending = self.pending.lock().await; + pending.remove(&request.id); + return Err(e); + } + } + + // Wait for the response with a timeout. + let timeout = Duration::from_secs(30); + match tokio::time::timeout(timeout, rx).await { + Ok(Ok(response)) => Ok(response), + Ok(Err(_)) => { + // Sender was dropped (reader task ended). Clean up pending entry. + let mut pending = self.pending.lock().await; + pending.remove(&request.id); + Err(ToolError::ExternalService(format!( + "[{}] MCP server closed connection before responding to request {}", + self.server_name, request.id + ))) + } + Err(_) => { + // Timeout: remove the pending entry. + let mut pending = self.pending.lock().await; + pending.remove(&request.id); + Err(ToolError::ExternalService(format!( + "[{}] Timeout waiting for response to request {} after {:?}", + self.server_name, request.id, timeout + ))) + } + } + } + + async fn shutdown(&self) -> Result<(), ToolError> { + // Abort the reader task. + if let Some(handle) = self.reader_handle.lock().await.take() { + handle.abort(); + } + + // Drain pending requests so waiters wake immediately instead of + // hanging until their 30s timeout. + { + let mut pending = self.pending.lock().await; + pending.clear(); // Dropping senders wakes receivers with Err + } + + tracing::debug!( + "[{}] Unix transport shut down (socket: {})", + self.server_name, + self.socket_path.display() + ); + Ok(()) + } + + fn supports_http_features(&self) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader as TokioBufReader}; + use tokio::net::UnixListener; + + #[tokio::test] + async fn test_connect_nonexistent_socket_fails() { + let tmp_dir = tempfile::tempdir().expect("create temp dir"); + let socket_path = tmp_dir.path().join("nonexistent.sock"); + + let result = UnixMcpTransport::connect("test", &socket_path).await; + + let err = result.err().expect("should be an error").to_string(); + assert!( + err.contains("Failed to connect"), + "Error should mention connection failure: {}", + err + ); + } + + #[tokio::test] + async fn test_round_trip_via_unix_socket() { + // Create a temporary directory for the socket. + let tmp_dir = tempfile::tempdir().expect("create temp dir"); + let socket_path = tmp_dir.path().join("test.sock"); + + // Bind a listener on the socket. + let listener = UnixListener::bind(&socket_path).expect("bind listener"); + + // Spawn an echo handler that reads one JSON-RPC request and writes + // back a valid McpResponse with the same id. + let handler = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept connection"); + let (read_half, mut write_half) = tokio::io::split(stream); + let mut reader = TokioBufReader::new(read_half); + let mut line = String::new(); + reader + .read_line(&mut line) + .await + .expect("read request line"); + + // Parse the request to extract the id. + let req: McpRequest = serde_json::from_str(&line).expect("parse request"); + + // Build a valid response. + let response = McpResponse { + jsonrpc: "2.0".to_string(), + id: req.id, + result: Some(serde_json::json!({"tools": []})), + error: None, + }; + + let mut resp_bytes = serde_json::to_vec(&response).expect("serialize response"); + resp_bytes.push(b'\n'); + write_half + .write_all(&resp_bytes) + .await + .expect("write response"); + write_half.flush().await.expect("flush"); + }); + + // Connect to the socket via our transport. + let transport = UnixMcpTransport::connect("test-uds", &socket_path) + .await + .expect("connect should succeed"); + + assert_eq!(transport.socket_path(), socket_path.as_path()); + assert_eq!(transport.server_name(), "test-uds"); + + // Send a list_tools request and verify the round-trip. + let request = McpRequest::list_tools(42); + let headers = HashMap::new(); + let response = transport.send(&request, &headers).await.expect("send"); + + assert_eq!(response.id, 42); + assert!(response.result.is_some()); + assert!(response.error.is_none()); + + // Clean up. + transport.shutdown().await.expect("shutdown"); + handler.await.expect("handler task"); + } + + #[tokio::test] + async fn test_shutdown_is_idempotent() { + let tmp_dir = tempfile::tempdir().expect("create temp dir"); + let socket_path = tmp_dir.path().join("idle.sock"); + + let listener = UnixListener::bind(&socket_path).expect("bind listener"); + + // Accept in the background so the connect succeeds. + let _handler = tokio::spawn(async move { + let _stream = listener.accept().await; + }); + + let transport = UnixMcpTransport::connect("test-idle", &socket_path) + .await + .expect("connect"); + + // Calling shutdown twice should not panic or error. + transport.shutdown().await.expect("first shutdown"); + transport.shutdown().await.expect("second shutdown"); + } +}