diff --git a/src/app.rs b/src/app.rs index 766aff30..53429057 100644 --- a/src/app.rs +++ b/src/app.rs @@ -542,7 +542,7 @@ impl AppBuilder { server, mcp_sm, secrets, "default", ) } else { - McpClient::new_with_name(&server_name, &server.url) + McpClient::new_with_config(server.clone()) }; match client.list_tools().await { diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index 5e2f4dea..bc3ddf45 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -47,6 +47,10 @@ pub enum McpCommand { /// Server description #[arg(long)] description: Option, + + /// Custom HTTP headers (format: "Key:Value", can be repeated) + #[arg(long = "header", short = 'H')] + headers: Vec, }, /// Remove an MCP server @@ -108,6 +112,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { token_url, scopes, description, + headers, } => { add_server( name, @@ -117,6 +122,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { token_url, scopes, description, + headers, ) .await } @@ -133,6 +139,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { } /// Add a new MCP server. +#[allow(clippy::too_many_arguments)] async fn add_server( name: String, url: String, @@ -141,6 +148,7 @@ async fn add_server( token_url: Option, scopes: Option, description: Option, + headers: Vec, ) -> anyhow::Result<()> { let mut config = McpServerConfig::new(&name, &url); @@ -148,6 +156,18 @@ async fn add_server( config = config.with_description(desc); } + // Parse custom headers (format: "Key:Value") + if !headers.is_empty() { + let mut header_map = std::collections::HashMap::new(); + for h in &headers { + let (key, value) = h.split_once(':').ok_or_else(|| { + anyhow::anyhow!("Invalid header format '{}'. Expected 'Key:Value'.", h) + })?; + header_map.insert(key.trim().to_string(), value.trim().to_string()); + } + config = config.with_headers(header_map); + } + // Track if auth is required let requires_auth = client_id.is_some(); @@ -242,6 +262,17 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> { if let Some(ref desc) = server.description { println!(" Description: {}", desc); } + if !server.headers.is_empty() { + println!( + " Custom headers: {}", + server + .headers + .keys() + .cloned() + .collect::>() + .join(", ") + ); + } if let Some(ref oauth) = server.oauth { println!(" OAuth Client ID: {}", oauth.client_id); if !oauth.scopes.is_empty() { @@ -374,7 +405,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 diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 6ef47d22..e804f8c8 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -2508,7 +2508,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/tools/mcp/client.rs b/src/tools/mcp/client.rs index 316851fc..30acb474 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -3,6 +3,7 @@ //! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers. //! Uses the Streamable HTTP transport with session management. +use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; @@ -52,6 +53,9 @@ pub struct McpClient { /// Server configuration (for token secret name lookup). server_config: Option, + + /// Custom HTTP headers injected into every request. + custom_headers: HashMap, } impl McpClient { @@ -75,6 +79,7 @@ impl McpClient { secrets: None, user_id: "default".to_string(), server_config: None, + custom_headers: HashMap::new(), } } @@ -95,6 +100,28 @@ impl McpClient { secrets: None, user_id: "default".to_string(), server_config: None, + custom_headers: HashMap::new(), + } + } + + /// Create a new simple MCP client from a server configuration (no authentication). + /// + /// Use this when you have an `McpServerConfig` with custom headers but no OAuth. + pub fn new_with_config(config: McpServerConfig) -> Self { + Self { + server_name: config.name.clone(), + server_url: config.url.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: None, + secrets: None, + user_id: "default".to_string(), + custom_headers: config.headers.clone(), + server_config: Some(config), } } @@ -119,6 +146,7 @@ impl McpClient { session_manager: Some(session_manager), secrets: Some(secrets), user_id: user_id.into(), + custom_headers: config.headers.clone(), server_config: Some(config), } } @@ -178,7 +206,12 @@ impl McpClient { .header("Content-Type", "application/json") .json(&request); - // Add Authorization header if we have a token + // Add custom headers from config + for (key, value) in &self.custom_headers { + req_builder = req_builder.header(key, value); + } + + // Add Authorization header if we have a token (overrides custom Authorization) if let Some(token) = self.get_access_token().await? { req_builder = req_builder.header("Authorization", format!("Bearer {}", token)); } @@ -474,6 +507,7 @@ impl Clone for McpClient { secrets: self.secrets.clone(), user_id: self.user_id.clone(), server_config: self.server_config.clone(), + custom_headers: self.custom_headers.clone(), } } } @@ -692,6 +726,26 @@ mod tests { assert_eq!(id3, 3); } + #[test] + fn test_custom_headers_from_config() { + use std::collections::HashMap; + 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); + assert_eq!(client.custom_headers.len(), 2); + assert_eq!(client.custom_headers.get("X-API-Key").unwrap(), "secret"); + } + + #[test] + fn test_new_has_no_custom_headers() { + let client = McpClient::new("http://localhost:8080"); + assert!(client.custom_headers.is_empty()); + } + #[test] fn test_mcp_tool_requires_approval_destructive() { use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations}; diff --git a/src/tools/mcp/config.rs b/src/tools/mcp/config.rs index 784f0aa2..1beea242 100644 --- a/src/tools/mcp/config.rs +++ b/src/tools/mcp/config.rs @@ -32,6 +32,13 @@ pub struct McpServerConfig { /// Optional description for the server. #[serde(skip_serializing_if = "Option::is_none")] pub description: Option, + + /// Custom HTTP headers to send with every request to this server. + /// + /// Useful for MCP servers that require non-OAuth authentication + /// (e.g., API keys via `X-API-Key` header). + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub headers: HashMap, } fn default_true() -> bool { @@ -47,9 +54,16 @@ impl McpServerConfig { oauth: None, enabled: true, description: None, + headers: HashMap::new(), } } + /// Set custom HTTP headers for this server. + pub fn with_headers(mut self, headers: HashMap) -> Self { + self.headers = headers; + self + } + /// Set OAuth configuration. pub fn with_oauth(mut self, oauth: OAuthConfig) -> Self { self.oauth = Some(oauth); @@ -593,4 +607,79 @@ mod tests { let config = McpServerConfig::new("bad", "http://mcp.example.com"); assert!(!config.requires_auth()); } + + #[test] + fn test_custom_headers_default_empty() { + let config = McpServerConfig::new("test", "http://localhost:8080"); + assert!(config.headers.is_empty()); + } + + #[test] + fn test_custom_headers_with_builder() { + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "secret123".to_string()); + headers.insert("X-Custom".to_string(), "value".to_string()); + + let config = McpServerConfig::new("browser-use", "https://mcp.browser-use.com") + .with_headers(headers.clone()); + + assert_eq!(config.headers.len(), 2); + assert_eq!(config.headers.get("X-API-Key").unwrap(), "secret123"); + } + + #[test] + fn test_custom_headers_serde_roundtrip() { + let mut headers = HashMap::new(); + headers.insert("Authorization".to_string(), "Bearer tok_123".to_string()); + + let config = + McpServerConfig::new("test-serde", "http://localhost:3000").with_headers(headers); + + let json = serde_json::to_string(&config).unwrap(); + let deserialized: McpServerConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(deserialized.headers.len(), 1); + assert_eq!( + deserialized.headers.get("Authorization").unwrap(), + "Bearer tok_123" + ); + } + + #[test] + fn test_custom_headers_absent_in_json_defaults_empty() { + let json = serde_json::json!({ + "name": "legacy", + "url": "http://localhost:8080" + }); + let config: McpServerConfig = serde_json::from_value(json).unwrap(); + assert!(config.headers.is_empty()); + } + + #[test] + fn test_custom_headers_skipped_when_empty_in_serialization() { + let config = McpServerConfig::new("minimal", "http://localhost:8080"); + let json = serde_json::to_value(&config).unwrap(); + // Empty headers map should not appear in serialized output + assert!(json.get("headers").is_none()); + } + + #[tokio::test] + async fn test_custom_headers_persist_to_disk() { + let dir = tempdir().unwrap(); + let path = dir.path().join("mcp-headers-test.json"); + + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "key123".to_string()); + + let mut config = McpServersFile::default(); + config.upsert( + McpServerConfig::new("headered", "http://localhost:9090").with_headers(headers), + ); + + save_mcp_servers_to(&config, &path).await.unwrap(); + let loaded = load_mcp_servers_from(&path).await.unwrap(); + + let server = loaded.get("headered").unwrap(); + assert_eq!(server.headers.get("X-API-Key").unwrap(), "key123"); + } }