diff --git a/apps/desktop/src/features/servers/ServersPage.tsx b/apps/desktop/src/features/servers/ServersPage.tsx index 07ad188e..599e2ff0 100644 --- a/apps/desktop/src/features/servers/ServersPage.tsx +++ b/apps/desktop/src/features/servers/ServersPage.tsx @@ -709,15 +709,9 @@ export function ServersPage() { }; const handleUninstall = async (server: ServerViewModel) => { - // Import source-aware helpers - const { getUninstallConfirmMessage, getUninstallLabel } = await import('@/components/SourceBadge'); - const confirmMsg = getUninstallConfirmMessage(server.name, server.installation_source); + const { getUninstallLabel } = await import('@/components/SourceBadge'); const actionLabel = getUninstallLabel(server.installation_source); - - if (!confirm(confirmMsg)) { - return; - } - + setActionLoading(`uninstall-${server.id}`); try { const { uninstallServer } = await import('@/lib/api/registry'); diff --git a/crates/mcpmux-gateway/src/pool/transport/http.rs b/crates/mcpmux-gateway/src/pool/transport/http.rs index d1e99ba8..250f0046 100644 --- a/crates/mcpmux-gateway/src/pool/transport/http.rs +++ b/crates/mcpmux-gateway/src/pool/transport/http.rs @@ -29,7 +29,6 @@ use crate::pool::credential_store::DatabaseCredentialStore; /// automatically refreshed by RMCP on every request when needed. pub struct HttpTransport { url: String, - #[allow(dead_code)] // Reserved for future custom headers headers: HashMap, space_id: Uuid, server_id: String, @@ -99,7 +98,7 @@ impl HttpTransport { oauth_indicators.iter().any(|s| error_lower.contains(s)) } - /// Connect with OAuth using DatabaseCredentialStore. + /// Connect with OAuth using DatabaseCredentialStore (with definition headers if any). /// /// RMCP's AuthClient will automatically: /// - Load tokens from the credential store @@ -107,11 +106,18 @@ impl HttpTransport { /// - Save refreshed tokens back to the store /// - Add auth header to every request /// + /// Definition headers are applied as default_headers on the underlying reqwest::Client, + /// so they're sent alongside OAuth tokens on every request. + /// /// If RMCP's metadata discovery fails (non-spec-compliant servers), we use /// stored metadata from the initial OAuth flow. - async fn connect_with_auth(&self) -> TransportConnectResult { + async fn connect_with_auth( + &self, + header_map: reqwest::header::HeaderMap, + ) -> TransportConnectResult { debug!( server_id = %self.server_id, + header_count = header_map.len(), "Connecting with OAuth via CredentialStore" ); @@ -119,8 +125,9 @@ impl HttpTransport { LogLevel::Info, LogSource::HttpRequest, format!( - "Connecting to {} with OAuth (auto-refresh enabled)", - self.url + "Connecting to {} with OAuth (auto-refresh enabled, {} custom header(s))", + self.url, + header_map.len() ), ) .await; @@ -234,12 +241,17 @@ impl HttpTransport { ) .await; - return self.connect_with_manual_token().await; + return self.connect_with_manual_token(header_map).await; } } - // Create AuthClient - this wraps reqwest::Client with automatic token injection & refresh - let auth_client = AuthClient::new(reqwest::Client::default(), auth_manager); + // Create AuthClient - wraps reqwest::Client with automatic token injection & refresh. + // Definition headers are baked into the client so they're sent on every request. + let base_client = match self.build_http_client(header_map) { + Ok(c) => c, + Err(err) => return TransportConnectResult::Failed(err), + }; + let auth_client = AuthClient::new(base_client, auth_manager); let transport_config = StreamableHttpClientTransportConfig::with_uri(self.url.as_str()); let transport = StreamableHttpClientTransport::with_client(auth_client, transport_config); @@ -303,8 +315,12 @@ impl HttpTransport { /// /// Some servers (like Cloudflare) don't serve OAuth metadata at the standard location /// that RMCP expects. In this case, we manually inject the stored token into requests. + /// Definition headers are merged in (token Authorization header takes precedence). /// NOTE: Auto-refresh won't work in this mode - tokens must be refreshed manually. - async fn connect_with_manual_token(&self) -> TransportConnectResult { + async fn connect_with_manual_token( + &self, + mut header_map: reqwest::header::HeaderMap, + ) -> TransportConnectResult { debug!( server_id = %self.server_id, "Connecting with manual token injection (RMCP metadata failed)" @@ -341,12 +357,11 @@ impl HttpTransport { ) .await; - // Build HTTP client with Authorization header - let mut headers = reqwest::header::HeaderMap::new(); + // Add Authorization header to the definition headers (overrides if already present) let auth_value = format!("Bearer {}", access_token); match reqwest::header::HeaderValue::from_str(&auth_value) { Ok(val) => { - headers.insert(reqwest::header::AUTHORIZATION, val); + header_map.insert(reqwest::header::AUTHORIZATION, val); } Err(e) => { let err = format!("Invalid token format: {}", e); @@ -355,13 +370,9 @@ impl HttpTransport { } } - let client = match reqwest::Client::builder().default_headers(headers).build() { + let client = match self.build_http_client(header_map) { Ok(c) => c, - Err(e) => { - let err = format!("Failed to build HTTP client: {}", e); - error!(server_id = %self.server_id, "{}", err); - return TransportConnectResult::Failed(err); - } + Err(err) => return TransportConnectResult::Failed(err), }; let transport_config = StreamableHttpClientTransportConfig::with_uri(self.url.as_str()); @@ -423,21 +434,74 @@ impl HttpTransport { } } - /// Try connecting without authentication - async fn connect_without_auth(&self) -> TransportConnectResult { + /// Build a reqwest HeaderMap from definition-provided headers. + /// + /// These headers (resolved from `${input:ID}` placeholders) are always applied + /// to the HTTP client regardless of auth strategy. Returns an empty map if no + /// definition headers are configured. + fn build_default_headers(&self) -> Result { + let mut header_map = reqwest::header::HeaderMap::new(); + for (key, value) in &self.headers { + let header_name = + reqwest::header::HeaderName::from_bytes(key.as_bytes()).map_err(|e| { + let err = format!("Invalid header name '{}': {}", key, e); + error!(server_id = %self.server_id, "{}", err); + err + })?; + let header_value = reqwest::header::HeaderValue::from_str(value).map_err(|e| { + let err = format!("Invalid header value for '{}': {}", key, e); + error!(server_id = %self.server_id, "{}", err); + err + })?; + header_map.insert(header_name, header_value); + } + Ok(header_map) + } + + /// Build a reqwest::Client with definition headers as default_headers. + fn build_http_client( + &self, + header_map: reqwest::header::HeaderMap, + ) -> Result { + reqwest::Client::builder() + .default_headers(header_map) + .build() + .map_err(|e| { + let err = format!("Failed to build HTTP client: {}", e); + error!(server_id = %self.server_id, "{}", err); + err + }) + } + + /// Try connecting without authentication (but with definition headers if any) + async fn connect_without_auth( + &self, + header_map: reqwest::header::HeaderMap, + ) -> TransportConnectResult { debug!( server_id = %self.server_id, + header_count = header_map.len(), "Trying connection without auth" ); self.log( LogLevel::Info, LogSource::HttpRequest, - format!("Connecting to {} without auth", self.url), + format!( + "Connecting to {} without auth ({} custom header(s))", + self.url, + header_map.len() + ), ) .await; - let transport = StreamableHttpClientTransport::from_uri(self.url.as_str()); + let client = match self.build_http_client(header_map) { + Ok(c) => c, + Err(err) => return TransportConnectResult::Failed(err), + }; + + let transport_config = StreamableHttpClientTransportConfig::with_uri(self.url.as_str()); + let transport = StreamableHttpClientTransport::with_client(client, transport_config); let client_handler = create_client_handler( &self.server_id, self.space_id, @@ -519,7 +583,33 @@ impl Transport for HttpTransport { return TransportConnectResult::Failed(err); } - // Check if we have stored credentials for this server + // Build definition headers (always applied regardless of auth strategy) + let header_map = match self.build_default_headers() { + Ok(h) => h, + Err(err) => return TransportConnectResult::Failed(err), + }; + + if !header_map.is_empty() { + info!( + server_id = %self.server_id, + header_count = header_map.len(), + "Applying definition-provided headers to connection" + ); + } + + // Check if definition headers already include an Authorization header. + // If so, skip OAuth — the user explicitly provided auth via the definition (e.g., PAT). + let has_explicit_auth = header_map.contains_key(reqwest::header::AUTHORIZATION); + + if has_explicit_auth { + info!( + server_id = %self.server_id, + "Definition includes Authorization header, skipping OAuth" + ); + return self.connect_without_auth(header_map).await; + } + + // No explicit auth in headers — check for stored OAuth credentials let has_credentials = self .credential_repo .get( @@ -537,13 +627,13 @@ impl Transport for HttpTransport { server_id = %self.server_id, "Found stored credentials, connecting with OAuth (auto-refresh enabled)" ); - self.connect_with_auth().await + self.connect_with_auth(header_map).await } else { debug!( server_id = %self.server_id, "No stored credentials, trying without auth" ); - self.connect_without_auth().await + self.connect_without_auth(header_map).await } } @@ -555,3 +645,445 @@ impl Transport for HttpTransport { format!("http:{}", self.url) } } + +#[cfg(test)] +mod tests { + use super::*; + use mcpmux_core::{Credential, CredentialType, OutboundOAuthRegistration}; + + // ── Mock repos (minimal, sufficient for HttpTransport unit tests) ── + + #[derive(Clone)] + struct MockCredentialRepo { + credentials: Arc>>, + } + + impl MockCredentialRepo { + fn new() -> Self { + Self { + credentials: Arc::new(tokio::sync::RwLock::new(Vec::new())), + } + } + + fn with_credential(cred: Credential) -> Self { + Self { + credentials: Arc::new(tokio::sync::RwLock::new(vec![cred])), + } + } + } + + #[async_trait] + impl CredentialRepository for MockCredentialRepo { + async fn get( + &self, + space_id: &Uuid, + server_id: &str, + credential_type: &CredentialType, + ) -> anyhow::Result> { + let creds = self.credentials.read().await; + Ok(creds + .iter() + .find(|c| { + c.space_id == *space_id + && c.server_id == server_id + && c.credential_type == *credential_type + }) + .cloned()) + } + + async fn get_all( + &self, + space_id: &Uuid, + server_id: &str, + ) -> anyhow::Result> { + let creds = self.credentials.read().await; + Ok(creds + .iter() + .filter(|c| c.space_id == *space_id && c.server_id == server_id) + .cloned() + .collect()) + } + + async fn save(&self, credential: &Credential) -> anyhow::Result<()> { + let mut creds = self.credentials.write().await; + creds.retain(|c| { + !(c.space_id == credential.space_id + && c.server_id == credential.server_id + && c.credential_type == credential.credential_type) + }); + creds.push(credential.clone()); + Ok(()) + } + + async fn delete( + &self, + space_id: &Uuid, + server_id: &str, + credential_type: &CredentialType, + ) -> anyhow::Result<()> { + let mut creds = self.credentials.write().await; + creds.retain(|c| { + !(c.space_id == *space_id + && c.server_id == server_id + && c.credential_type == *credential_type) + }); + Ok(()) + } + + async fn delete_all(&self, space_id: &Uuid, server_id: &str) -> anyhow::Result<()> { + let mut creds = self.credentials.write().await; + creds.retain(|c| !(c.space_id == *space_id && c.server_id == server_id)); + Ok(()) + } + + async fn clear_tokens(&self, space_id: &Uuid, server_id: &str) -> anyhow::Result { + let mut creds = self.credentials.write().await; + let before = creds.len(); + creds.retain(|c| { + !(c.space_id == *space_id + && c.server_id == server_id + && c.credential_type.is_oauth()) + }); + Ok(creds.len() < before) + } + + async fn list_for_space(&self, space_id: &Uuid) -> anyhow::Result> { + let creds = self.credentials.read().await; + Ok(creds + .iter() + .filter(|c| c.space_id == *space_id) + .cloned() + .collect()) + } + } + + #[derive(Clone)] + struct MockOAuthRepo; + + #[async_trait] + impl OutboundOAuthRepository for MockOAuthRepo { + async fn get( + &self, + _space_id: &Uuid, + _server_id: &str, + ) -> anyhow::Result> { + Ok(None) + } + + async fn save(&self, _registration: &OutboundOAuthRegistration) -> anyhow::Result<()> { + Ok(()) + } + + async fn delete(&self, _space_id: &Uuid, _server_id: &str) -> anyhow::Result<()> { + Ok(()) + } + + async fn list_for_space( + &self, + _space_id: &Uuid, + ) -> anyhow::Result> { + Ok(vec![]) + } + } + + /// Helper to create an HttpTransport with given headers and credential repo. + fn make_transport( + headers: HashMap, + credential_repo: Arc, + ) -> HttpTransport { + HttpTransport::new( + "https://example.com/mcp".to_string(), + headers, + Uuid::new_v4(), + "test-server".to_string(), + credential_repo, + Arc::new(MockOAuthRepo), + None, + Duration::from_secs(10), + None, + ) + } + + fn make_transport_with_space( + headers: HashMap, + credential_repo: Arc, + space_id: Uuid, + server_id: &str, + ) -> HttpTransport { + HttpTransport::new( + "https://example.com/mcp".to_string(), + headers, + space_id, + server_id.to_string(), + credential_repo, + Arc::new(MockOAuthRepo), + None, + Duration::from_secs(10), + None, + ) + } + + // ── requires_oauth tests ── + + #[test] + fn test_requires_oauth_401() { + assert!(HttpTransport::requires_oauth("HTTP 401 Unauthorized")); + } + + #[test] + fn test_requires_oauth_bearer() { + assert!(HttpTransport::requires_oauth("Missing Bearer token")); + } + + #[test] + fn test_requires_oauth_www_authenticate() { + assert!(HttpTransport::requires_oauth("WWW-Authenticate: Bearer")); + } + + #[test] + fn test_requires_oauth_channel_closed() { + assert!(HttpTransport::requires_oauth("transport channel closed")); + } + + #[test] + fn test_requires_oauth_false_for_unrelated() { + assert!(!HttpTransport::requires_oauth("connection refused")); + assert!(!HttpTransport::requires_oauth("DNS lookup failed")); + assert!(!HttpTransport::requires_oauth("timeout")); + } + + // ── build_default_headers tests ── + + #[test] + fn test_build_default_headers_empty() { + let transport = make_transport(HashMap::new(), Arc::new(MockCredentialRepo::new())); + let headers = transport.build_default_headers().unwrap(); + assert!(headers.is_empty()); + } + + #[test] + fn test_build_default_headers_single() { + let mut h = HashMap::new(); + h.insert("Authorization".to_string(), "Bearer token123".to_string()); + let transport = make_transport(h, Arc::new(MockCredentialRepo::new())); + let headers = transport.build_default_headers().unwrap(); + + assert_eq!(headers.len(), 1); + assert_eq!( + headers.get(reqwest::header::AUTHORIZATION).unwrap(), + "Bearer token123" + ); + } + + #[test] + fn test_build_default_headers_multiple() { + let mut h = HashMap::new(); + h.insert("Authorization".to_string(), "Bearer pat_xxx".to_string()); + h.insert("X-Custom-Header".to_string(), "custom-value".to_string()); + let transport = make_transport(h, Arc::new(MockCredentialRepo::new())); + let headers = transport.build_default_headers().unwrap(); + + assert_eq!(headers.len(), 2); + assert_eq!( + headers.get(reqwest::header::AUTHORIZATION).unwrap(), + "Bearer pat_xxx" + ); + assert_eq!(headers.get("x-custom-header").unwrap(), "custom-value"); + } + + #[test] + fn test_build_default_headers_invalid_name() { + let mut h = HashMap::new(); + h.insert("Invalid Header\n".to_string(), "value".to_string()); + let transport = make_transport(h, Arc::new(MockCredentialRepo::new())); + let result = transport.build_default_headers(); + assert!(result.is_err()); + } + + #[test] + fn test_build_default_headers_invalid_value() { + let mut h = HashMap::new(); + h.insert("X-Header".to_string(), "bad\nvalue".to_string()); + let transport = make_transport(h, Arc::new(MockCredentialRepo::new())); + let result = transport.build_default_headers(); + assert!(result.is_err()); + } + + // ── build_http_client tests ── + + #[test] + fn test_build_http_client_empty_headers() { + let transport = make_transport(HashMap::new(), Arc::new(MockCredentialRepo::new())); + let client = transport.build_http_client(reqwest::header::HeaderMap::new()); + assert!(client.is_ok()); + } + + #[test] + fn test_build_http_client_with_headers() { + let transport = make_transport(HashMap::new(), Arc::new(MockCredentialRepo::new())); + let mut header_map = reqwest::header::HeaderMap::new(); + header_map.insert( + reqwest::header::AUTHORIZATION, + reqwest::header::HeaderValue::from_static("Bearer token"), + ); + let client = transport.build_http_client(header_map); + assert!(client.is_ok()); + } + + // ── connect() routing logic tests ── + + #[tokio::test] + async fn test_connect_invalid_url_fails() { + let transport = HttpTransport::new( + "not a valid url".to_string(), + HashMap::new(), + Uuid::new_v4(), + "test-server".to_string(), + Arc::new(MockCredentialRepo::new()), + Arc::new(MockOAuthRepo), + None, + Duration::from_secs(5), + None, + ); + + let result = transport.connect().await; + match result { + TransportConnectResult::Failed(msg) => { + assert!(msg.contains("Invalid URL"), "Got: {}", msg); + } + _ => panic!("Expected Failed for invalid URL"), + } + } + + #[tokio::test] + async fn test_connect_with_explicit_auth_header_skips_oauth_check() { + // When headers include Authorization, connect should NOT check credential_repo + // for OAuth tokens — it should go straight to connect_without_auth with headers. + // We verify this by giving it an unreachable server URL (connection will fail) + // and checking that the error is a connection failure, NOT OAuthRequired. + let mut h = HashMap::new(); + h.insert( + "Authorization".to_string(), + "Bearer ghp_testtoken123".to_string(), + ); + + let transport = HttpTransport::new( + "https://127.0.0.1:1/mcp".to_string(), // unreachable + h, + Uuid::new_v4(), + "test-server".to_string(), + Arc::new(MockCredentialRepo::new()), + Arc::new(MockOAuthRepo), + None, + Duration::from_secs(2), + None, + ); + + let result = transport.connect().await; + // Should be Failed (connection error) or timeout — NOT OAuthRequired + match result { + TransportConnectResult::Failed(_) => {} // expected + TransportConnectResult::OAuthRequired { .. } => { + panic!("Should not trigger OAuth when Authorization header is present") + } + _ => {} // Connected would be surprising but not wrong + } + } + + #[tokio::test] + async fn test_connect_no_headers_no_credentials_tries_no_auth() { + // No headers, no stored credentials → connect_without_auth + // Will fail to connect to unreachable server + let transport = HttpTransport::new( + "https://127.0.0.1:1/mcp".to_string(), + HashMap::new(), + Uuid::new_v4(), + "test-server".to_string(), + Arc::new(MockCredentialRepo::new()), + Arc::new(MockOAuthRepo), + None, + Duration::from_secs(2), + None, + ); + + let result = transport.connect().await; + // Should be Failed (connection error/timeout) or OAuthRequired (if 401 detected) + match result { + TransportConnectResult::Failed(_) | TransportConnectResult::OAuthRequired { .. } => {} + TransportConnectResult::Connected(_) => { + panic!("Should not connect to unreachable server") + } + } + } + + #[tokio::test] + async fn test_connect_with_stored_credentials_routes_to_oauth() { + // When stored credentials exist and no explicit Authorization header, + // connect should route to connect_with_auth (which will fail on unreachable server, + // but we verify it doesn't go to connect_without_auth by observing the error path). + let space_id = Uuid::new_v4(); + let server_id = "test-server"; + + let cred = Credential::access_token(space_id, server_id, "stored_token", None); + let cred_repo = Arc::new(MockCredentialRepo::with_credential(cred)); + + let transport = make_transport_with_space(HashMap::new(), cred_repo, space_id, server_id); + + let result = transport.connect().await; + // connect_with_auth will fail (can't reach server / AuthorizationManager::new fails) + // but it should NOT be OAuthRequired from the no-auth path's 401 detection + match result { + TransportConnectResult::Failed(_) | TransportConnectResult::OAuthRequired { .. } => {} + TransportConnectResult::Connected(_) => { + panic!("Should not connect to example.com MCP endpoint") + } + } + } + + #[tokio::test] + async fn test_connect_custom_headers_always_applied_with_credentials() { + // Even when we have stored OAuth credentials, custom (non-auth) headers + // from the definition should be applied. We can verify this indirectly: + // the transport should route to connect_with_auth (not connect_without_auth) + // because there's no Authorization in the custom headers. + let space_id = Uuid::new_v4(); + let server_id = "test-server"; + + let cred = Credential::access_token(space_id, server_id, "stored_token", None); + let cred_repo = Arc::new(MockCredentialRepo::with_credential(cred)); + + let mut h = HashMap::new(); + h.insert("X-MCP-Toolsets".to_string(), "tools-only".to_string()); + + let transport = make_transport_with_space(h, cred_repo, space_id, server_id); + + // Verify headers are built correctly (non-auth header present, no Authorization) + let headers = transport.build_default_headers().unwrap(); + assert_eq!(headers.len(), 1); + assert!(!headers.contains_key(reqwest::header::AUTHORIZATION)); + assert_eq!(headers.get("x-mcp-toolsets").unwrap(), "tools-only"); + + // connect() should route to connect_with_auth since no explicit Authorization header + let result = transport.connect().await; + match result { + TransportConnectResult::Failed(_) | TransportConnectResult::OAuthRequired { .. } => {} + TransportConnectResult::Connected(_) => { + panic!("Should not connect to example.com MCP endpoint") + } + } + } + + // ── transport_type / description tests ── + + #[test] + fn test_transport_type() { + let transport = make_transport(HashMap::new(), Arc::new(MockCredentialRepo::new())); + assert!(matches!(transport.transport_type(), TransportType::Http)); + } + + #[test] + fn test_description() { + let transport = make_transport(HashMap::new(), Arc::new(MockCredentialRepo::new())); + assert_eq!(transport.description(), "http:https://example.com/mcp"); + } +}