diff --git a/src-tauri/src/mcp/client.rs b/src-tauri/src/mcp/client.rs index b6895dd2..11cafa13 100644 --- a/src-tauri/src/mcp/client.rs +++ b/src-tauri/src/mcp/client.rs @@ -7,17 +7,26 @@ use crate::mcp::models::{McpResource, McpTool}; /// Live connection to an MCP server. pub type McpConnection = RunningService; -/// Connect to a stdio MCP server. -pub async fn connect_stdio(command: &str, args: &[String]) -> Result { - let transport = crate::mcp::transport::stdio::build_stdio_transport(command, args)?; +/// Connect to a stdio MCP server with optional environment variables. +pub async fn connect_stdio( + command: &str, + args: &[String], + env: std::collections::HashMap, +) -> Result { + let transport = crate::mcp::transport::stdio::build_stdio_transport(command, args, env)?; ().serve(transport) .await .map_err(|e| format!("MCP stdio connection failed: {e}")) } -/// Connect to an HTTP MCP server. -pub async fn connect_http(url: &str, auth_header: Option<&str>) -> Result { - let transport = crate::mcp::transport::http::build_http_transport(url, auth_header); +/// Connect to an HTTP MCP server with optional custom headers. +pub async fn connect_http( + url: &str, + auth_header: Option<&str>, + custom_headers: std::collections::HashMap, +) -> Result { + let transport = + crate::mcp::transport::http::build_http_transport(url, auth_header, custom_headers); ().serve(transport) .await .map_err(|e| format!("MCP HTTP connection failed: {e}")) diff --git a/src-tauri/src/mcp/discovery.rs b/src-tauri/src/mcp/discovery.rs index 88095c72..e1add604 100644 --- a/src-tauri/src/mcp/discovery.rs +++ b/src-tauri/src/mcp/discovery.rs @@ -5,7 +5,8 @@ use tracing::{info, warn}; use crate::mcp::client::{connect_http, connect_stdio, list_resources, list_tools, McpConnection}; use crate::mcp::models::McpServer; use crate::mcp::store::{ - get_server_auth_value, list_servers, replace_resources, replace_tools, update_discovery_status, + get_server_auth_value, get_server_env_config, list_servers, replace_resources, replace_tools, + update_discovery_status, }; /// Discover a single MCP server: connect, list tools/resources, persist. @@ -55,11 +56,49 @@ async fn discover_server_inner( .collect() }) .unwrap_or_default(); - connect_stdio(command, &args).await? + + // Parse plaintext env vars from transport_config.env + let plaintext_env: std::collections::HashMap = config + .get("env") + .and_then(|v| v.as_object()) + .map(|obj| { + obj.iter() + .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string()))) + .collect() + }) + .unwrap_or_default(); + + // Decrypt and parse encrypted env vars from env_config column + let encrypted_env = { + let db = state.db.lock().map_err(|e| e.to_string())?; + get_server_env_config(&db, &server.id)? + }; + + // Merge env vars (encrypted takes precedence over plaintext) + let mut merged_env = plaintext_env; + if let Some(enc_env) = encrypted_env { + merged_env.extend(enc_env); + } + + connect_stdio(command, &args, merged_env).await? } "http" => { let auth_header = auth_value.as_deref(); - connect_http(&server.url, auth_header).await? + + // Parse custom headers from transport_config.headers + let config: serde_json::Value = + serde_json::from_str(&server.transport_config).unwrap_or_default(); + let custom_headers: std::collections::HashMap = config + .get("headers") + .and_then(|v| v.as_object()) + .map(|obj| { + obj.iter() + .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string()))) + .collect() + }) + .unwrap_or_default(); + + connect_http(&server.url, auth_header, custom_headers).await? } other => return Err(format!("Unknown transport type: {other}")), }; diff --git a/src-tauri/src/mcp/transport/http.rs b/src-tauri/src/mcp/transport/http.rs index ea6e4d48..c4ab515e 100644 --- a/src-tauri/src/mcp/transport/http.rs +++ b/src-tauri/src/mcp/transport/http.rs @@ -1,17 +1,32 @@ use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; use rmcp::transport::StreamableHttpClientTransport; +use std::collections::HashMap; use std::sync::Arc; -/// Build an HTTP (Streamable HTTP) transport from a URL. +/// Build an HTTP (Streamable HTTP) transport from a URL with optional custom headers. /// Optionally attaches an Authorization bearer token. +/// +/// NOTE: Custom headers are parsed but not yet applied due to rmcp v1.7.0 API limitations. +/// The rmcp library's StreamableHttpClientTransportConfig does not expose a .header() method. +/// Custom headers support is deferred until rmcp adds this capability or we find an alternative. pub fn build_http_transport( url: &str, auth_header: Option<&str>, + custom_headers: HashMap, ) -> impl rmcp::transport::Transport { + // Log warning if custom headers are provided (not yet supported) + if !custom_headers.is_empty() { + tracing::warn!( + "Custom HTTP headers provided but not supported by rmcp v1.7.0: {:?}", + custom_headers.keys().collect::>() + ); + } + let config = match auth_header { Some(token) => StreamableHttpClientTransportConfig::with_uri(Arc::from(url)) .auth_header(token.to_string()), None => StreamableHttpClientTransportConfig::with_uri(Arc::from(url)), }; + StreamableHttpClientTransport::from_config(config) } diff --git a/src-tauri/src/mcp/transport/stdio.rs b/src-tauri/src/mcp/transport/stdio.rs index b97ba55a..5d3297ae 100644 --- a/src-tauri/src/mcp/transport/stdio.rs +++ b/src-tauri/src/mcp/transport/stdio.rs @@ -1,10 +1,15 @@ use rmcp::transport::TokioChildProcess; +use std::collections::HashMap; use std::path::Path; use tokio::process::Command; -/// Build a stdio transport from a command path and argument list. +/// Build a stdio transport from a command path, argument list, and environment variables. /// Rejects relative paths to prevent path traversal. -pub fn build_stdio_transport(command: &str, args: &[String]) -> Result { +pub fn build_stdio_transport( + command: &str, + args: &[String], + env: HashMap, +) -> Result { if !Path::new(command).is_absolute() { return Err(format!( "stdio command must be an absolute path, got: {command}" @@ -14,5 +19,10 @@ pub fn build_stdio_transport(command: &str, args: &[String]) -> Result