gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(mcp): transport abstraction, stdio/UDS transports, and OAuth fixes (#721)

Browse Source 
* feat(mcp): transport abstraction, stdio/UDS transports, and OAuth fixes

Extract McpTransport trait from HTTP-coupled McpClient, enabling pluggable
transport backends. Implements stdio and Unix domain socket transports for
local MCP server integration, fixes OAuth discovery per RFC 9728, and adds
SSRF protection.

Transport abstraction (Step 2):
- McpTransport trait with send(), shutdown(), supports_http_features()
- HttpMcpTransport extracted from McpClient with SSE parsing, session tracking
- Shared JSON-RPC framing helpers (write_jsonrpc_line, spawn_jsonrpc_reader)
- McpClient refactored to hold Arc<dyn McpTransport>

Stdio transport (#652, Step 4):
- StdioMcpTransport spawns child process, communicates via stdin/stdout
- McpProcessManager for lifecycle management with exponential backoff restart
- Background stderr drain task for debug logging

Unix domain socket transport (#134, Step 5):
- UnixMcpTransport connects to existing Unix sockets
- Reuses shared JSON-RPC framing from transport.rs

HTML error body sanitization (#263, Step 1):
- sanitize_error_body() detects HTML, strips control chars, truncates to 500

Custom headers (#639, Step 3):
- headers field on McpServerConfig, merged into every HTTP request
- --header CLI arg for `mcp add`

Config and CLI updates (Step 6):
- McpTransportConfig tagged enum (Http/Stdio/Unix) with serde support
- EffectiveTransport for zero-copy config dispatch
- CLI: --transport, --command, --arg, --env, --socket flags for `mcp add`
- `mcp list` shows transport type

OAuth fixes (#299, Step 8):
- Multi-strategy discovery (401-based, RFC 9728, direct)
- RFC 8707 resource parameter in auth and refresh flows
- SSRF protection with IPv4-mapped IPv6 bypass detection
- Well-known URI construction per RFC 8414

Closes #652, #134, #639, #263, #299

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(mcp): address audit findings from crate review

- Fix SSRF bypass: make validate_url_safe async with DNS resolution to
  block hostnames that resolve to private/link-local IPs
- Fix UTF-8 truncation: use char-based truncation in sanitize_error_body
  to avoid panicking on multi-byte characters
- Fix SSE parser: process only complete lines to handle chunks split
  across boundaries, add 10MB buffer size limit
- Add debug_assert for transport type mismatch in new_with_config
- Propagate custom headers in new_with_transport constructor
- Deduplicate effective_transport() calls in CLI list command
- Gate test-only accessors with #[cfg(test)] to eliminate dead_code warnings
- Document JSON-RPC notification id:0 limitation in protocol.rs
- Document total backoff wait time (31s) in process.rs
- Add regression test for multi-byte UTF-8 truncation

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(mcp): address PR review findings from Copilot, Gemini, and zmanian

Moderate/High fixes:
- Plumb custom headers through new_authenticated constructor
- Restrict HTTP to localhost only in validate_url_safe (prevent
  plaintext credential leaks over non-localhost HTTP)
- Add mcp_process_manager.shutdown_all() to app shutdown path to
  prevent orphaning stdio child processes
- Validate discovered authorization_url before opening browser
  (prevent malicious MCP server redirecting to phishing page)

Medium fixes:
- Upgrade debug_assert to assert in new_with_config (fires in release)
- Remove pending map entry on Ok(Err(_)) in stdio/unix send() to avoid
  stale entries and unnecessary 30s waits
- Shut down old transport in try_restart() before spawning replacement
- Redact env var values in mcp list --verbose (may contain secrets)
- Drain pending requests on shutdown to wake waiters immediately
- Add IPv6 link-local, site-local, unique-local, and documentation
  ranges to is_dangerous_ip SSRF protection

Low fixes:
- Truncate logged JSON parse error lines to 200 chars (prevent
  sensitive data in logs)
- Remove misleading shutdown comment in unix_transport
- Use tempfile::tempdir() instead of hardcoded /tmp/ path in test
- Adopt main's improved sanitize_error_body (HTML tag stripping,
  200-char truncation with char_indices)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(mcp): gate unix_transport with #[cfg(unix)] for Windows compat

- Add #[cfg(unix)] to unix_transport module declaration
- Add #[cfg(unix)]/#[cfg(not(unix))] branches in app.rs for Unix
  socket MCP server setup
- Remove unused sanitize_error_body import in client.rs tests

[skip-regression-check]

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-03-09 02:47:42 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 9401ab0d58
commit 02f85a8ad5
15 changed files with 2937 additions and 474 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/app.rs
+159 -73
View File
@@ -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<Arc<Workspace>>,
pub extension_manager: Option<Arc<ExtensionManager>>,
pub mcp_session_manager: Arc<McpSessionManager>,
pub mcp_process_manager: Arc<McpProcessManager>,
pub wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
pub log_broadcaster: Arc<LogBroadcaster>,
pub context_manager: Arc<ContextManager>,
@@ -420,6 +421,7 @@ impl AppBuilder {
) -> Result<
(
Arc<McpSessionManager>,
Arc<McpProcessManager>,
Option<Arc<WasmToolRuntime>>,
Option<Arc<ExtensionManager>>,
Vec<crate::extensions::RegistryEntry>,
@@ -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<dyn McpTransport>,
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<dyn McpTransport>,
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,
src/cli/mcp.rs
+241 -65
View File
@@ -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<String>,
/// Transport type: http (default), stdio, unix
#[arg(long, default_value = "http")]
pub transport: String,
/// Command to run (stdio transport)
#[arg(long)]
pub command: Option<String>,
/// Command arguments (stdio transport, can be repeated)
#[arg(long = "arg", num_args = 1..)]
pub cmd_args: Vec<String>,
/// 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<String>,
/// 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<String>,
/// OAuth authorization URL (optional, can be discovered)
#[arg(long)]
pub auth_url: Option<String>,
/// OAuth token URL (optional, can be discovered)
#[arg(long)]
pub token_url: Option<String>,
/// Scopes to request (comma-separated)
#[arg(long)]
pub scopes: Option<String>,
/// Server description
#[arg(long)]
pub description: Option<String>,
}
#[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<String>,
/// OAuth authorization URL (optional, can be discovered)
#[arg(long)]
auth_url: Option<String>,
/// OAuth token URL (optional, can be discovered)
#[arg(long)]
token_url: Option<String>,
/// Scopes to request (comma-separated)
#[arg(long)]
scopes: Option<String>,
/// Server description
#[arg(long)]
description: Option<String>,
},
Add(Box<McpAddArgs>),
/// 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<String>,
auth_url: Option<String>,
token_url: Option<String>,
scopes: Option<String>,
description: Option<String>,
) -> 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<String, String> = 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<String, String> = 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::<Vec<_>>()
.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"));
}
}
src/cli/mod.rs
+1 -1
View File
@@ -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<McpCommand>),
/// Query and manage workspace memory
#[command(
src/extensions/manager.rs
+2 -1
View File
@@ -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
src/main.rs
+4 -1
View File
@@ -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
src/tools/mcp/auth.rs
+570 -23
View File
@@ -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<String, AuthError> {
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::<IpAddr>()
&& 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::<IpAddr>().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<String> {
// 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<ProtectedResourceMetadata, AuthError> {
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<AuthorizationServerMetadata, AuthError> {
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<AuthorizationServerMetadata, AuthError> {
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<ProtectedResourceMetadata, AuthError> {
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<AuthorizationServerMetadata, AuthError> {
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<AuthorizationServerMetadata, AuthError> {
// 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<ClientRegistrationResponse, AuthError> {
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(&registration_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<String, String>,
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<AccessToken, AuthError> {
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(&params)
@@ -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="));
}
}
src/tools/mcp/client.rs
+242 -297
View File
@@ -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<dyn McpTransport>,
/// 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<McpServerConfig>,
/// Custom headers to include in every request.
custom_headers: HashMap<String, String>,
}
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<String>) -> 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<String>, server_url: impl Into<String>) -> 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<dyn SecretsStore + Send + Sync>,
user_id: impl Into<String>,
) -> 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<String>,
transport: Arc<dyn McpTransport>,
session_manager: Option<Arc<McpSessionManager>>,
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
user_id: impl Into<String>,
server_config: Option<McpServerConfig>,
) -> 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<Option<String>, 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<HashMap<String, String>, 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<McpResponse, ToolError> {
// 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<McpResponse, ToolError> {
// 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::<McpResponse>(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<InitializeResult, ToolError> {
// 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<Vec<McpTool>, 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<CallToolResult, ToolError> {
// 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<Vec<Arc<dyn Tool>>, 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<McpClient>,
}
@@ -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<ToolOutput, ToolError> {
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::<Vec<_>>()
.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 `<html` or `<!DOCTYPE`) and
/// strips all tags, collapsing whitespace. Non-HTML bodies are left
/// intact. In both cases the result is truncated to 200 *characters*
/// (char-boundary safe) so that large payloads don't bloat error messages.
///
/// See #263 — raw HTML error pages were propagating through the error
/// chain into the web UI, causing a white screen.
fn sanitize_error_body(body: &str) -> 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("<html") || lower.contains("<!doctype");
let text = if is_html_document {
let stripped = body
.chars()
.fold((String::new(), false), |(mut out, in_tag), c| {
if c == '<' {
(out, true)
} else if c == '>' {
(out, false)
} else if !in_tag {
out.push(c);
(out, false)
} else {
(out, true)
}
})
.0;
stripped.split_whitespace().collect::<Vec<_>>().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#"<!DOCTYPE html><html><body><h1>422 Error</h1><p>Invalid token</p></body></html>"#;
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<Vec<McpResponse>>,
recorded_headers: std::sync::Mutex<Vec<HashMap<String, String>>>,
}
#[test]
fn test_sanitize_error_body_truncates_large_html_page() {
let html = format!(
"<html><body><p>{}</p></body></html>",
"error detail ".repeat(50)
impl MockTransport {
fn new(supports_http: bool, responses: Vec<McpResponse>) -> Self {
Self {
supports_http,
responses: std::sync::Mutex::new(responses),
recorded_headers: std::sync::Mutex::new(Vec::new()),
}
}
fn recorded_headers(&self) -> Vec<HashMap<String, String>> {
self.recorded_headers.lock().unwrap().clone()
}
}
#[async_trait]
impl McpTransport for MockTransport {
async fn send(
&self,
_request: &McpRequest,
headers: &HashMap<String, String>,
) -> Result<McpResponse, ToolError> {
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 = "<HTML><BODY><H1>500 Internal Server Error</H1></BODY></HTML>";
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());
}
}
src/tools/mcp/config.rs
+382 -12
View File
@@ -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<String>,
#[serde(default)]
env: HashMap<String, String>,
},
/// 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<McpTransportConfig>,
/// Custom headers to include in every HTTP request.
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub headers: HashMap<String, String>,
/// OAuth configuration (if server requires authentication).
#[serde(skip_serializing_if = "Option::is_none")]
pub oauth: Option<OAuthConfig>,
@@ -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<String>,
command: impl Into<String>,
args: Vec<String>,
env: HashMap<String, String>,
) -> 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<String>, socket_path: impl Into<String>) -> 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<String, String>) -> 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<String, String>,
},
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");
}
}
src/tools/mcp/http_transport.rs
+386
View File
@@ -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<Arc<McpSessionManager>>,
custom_headers: HashMap<String, String>,
}
impl HttpMcpTransport {
/// Create a new HTTP transport for the given server URL.
pub fn new(server_url: impl Into<String>, server_name: impl Into<String>) -> 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<McpSessionManager>) -> 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<String, String>) -> 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<McpSessionManager>> {
self.session_manager.as_ref()
}
}
#[async_trait]
impl McpTransport for HttpMcpTransport {
async fn send(
&self,
request: &McpRequest,
headers: &HashMap<String, String>,
) -> Result<McpResponse, ToolError> {
// 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<McpResponse, ToolError> {
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::<McpResponse>(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::<McpResponse>(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 (`<html` or `<!doctype`),
/// strips all tags, collapsing whitespace. Non-HTML bodies are left
/// intact. In both cases the result is truncated to 200 *characters*
/// (char-boundary safe) so that large payloads don't bloat error messages.
///
/// See #263 — raw HTML error pages were propagating through the error
/// chain into the web UI, causing a white screen.
pub(crate) fn sanitize_error_body(body: &str) -> 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("<html") || lower.contains("<!doctype");
let text = if is_html_document {
let stripped = body
.chars()
.fold((String::new(), false), |(mut out, in_tag), c| {
if c == '<' {
(out, true)
} else if c == '>' {
(out, false)
} else if !in_tag {
out.push(c);
(out, false)
} else {
(out, true)
}
})
.0;
stripped.split_whitespace().collect::<Vec<_>>().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#"<!DOCTYPE html><html><body><h1>422 Error</h1><p>Invalid token</p></body></html>"#;
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!(
"<html><body><p>{}</p></body></html>",
"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 = "<HTML><BODY><H1>500 Internal Server Error</H1></BODY></HTML>";
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");
}
}
src/tools/mcp/mod.rs
+10
View File
@@ -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;
src/tools/mcp/process.rs
+206
View File
@@ -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<String>,
pub env: HashMap<String, String>,
}
/// Manages stdio MCP server processes.
///
/// Handles spawning, tracking, and shutdown of child processes.
pub struct McpProcessManager {
transports: RwLock<HashMap<String, Arc<StdioMcpTransport>>>,
configs: RwLock<HashMap<String, StdioSpawnConfig>>,
}
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<String>,
command: impl Into<String>,
args: Vec<String>,
env: HashMap<String, String>,
) -> Result<Arc<StdioMcpTransport>, 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<Arc<StdioMcpTransport>> {
self.transports.read().await.get(name).cloned()
}
/// Shut down all managed transports.
pub async fn shutdown_all(&self) {
let transports: Vec<(String, Arc<StdioMcpTransport>)> = {
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<Arc<StdioMcpTransport>, 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<String> {
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;
}
}
src/tools/mcp/protocol.rs
+6 -1
View File
@@ -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<u64>`).
/// 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,
}
src/tools/mcp/stdio_transport.rs
+263
View File
@@ -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<Mutex<tokio::process::ChildStdin>>,
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<McpResponse>>>>,
reader_handle: Mutex<Option<JoinHandle<()>>>,
stderr_handle: Mutex<Option<JoinHandle<()>>>,
child: Arc<Mutex<Child>>,
}
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<String>,
command: &str,
args: impl IntoIterator<Item = impl AsRef<std::ffi::OsStr>>,
env: impl IntoIterator<Item = (impl AsRef<std::ffi::OsStr>, impl AsRef<std::ffi::OsStr>)>,
) -> Result<Self, ToolError> {
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<Mutex<HashMap<u64, oneshot::Sender<McpResponse>>>> =
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<String, String>,
) -> Result<McpResponse, ToolError> {
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<String, String> = 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<String, String> = 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<String, String> = 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());
}
}
src/tools/mcp/transport.rs
+196
View File
@@ -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<String, String>,
) -> Result<McpResponse, ToolError>;
/// 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<R: AsyncBufRead + Unpin + Send + 'static>(
reader: R,
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<McpResponse>>>>,
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::<McpResponse>(&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<Mutex<HashMap<u64, oneshot::Sender<McpResponse>>>> =
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<Mutex<HashMap<u64, oneshot::Sender<McpResponse>>>> =
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");
}
}
src/tools/mcp/unix_transport.rs
+269
View File
@@ -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<Mutex<tokio::io::WriteHalf<UnixStream>>>,
pending: Arc<Mutex<HashMap<u64, oneshot::Sender<McpResponse>>>>,
reader_handle: Mutex<Option<JoinHandle<()>>>,
}
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<String>,
socket_path: impl AsRef<Path>,
) -> Result<Self, ToolError> {
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<Mutex<HashMap<u64, oneshot::Sender<McpResponse>>>> =
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<String, String>,
) -> Result<McpResponse, ToolError> {
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");
}
}