1#![warn(missing_docs)]
2
3pub mod circuit_breaker;
12pub mod reconnect;
13pub mod router;
14pub mod timeout;
15
16use std::borrow::Cow;
17use std::collections::HashMap;
18
19use anyhow::{Context, Result};
20use forge_sandbox::{ResourceDispatcher, ToolDispatcher};
21use rmcp::model::{CallToolRequestParams, CallToolResult, Content, RawContent};
22use rmcp::service::RunningService;
23use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
24use rmcp::transport::{ConfigureCommandExt, StreamableHttpClientTransport, TokioChildProcess};
25use rmcp::{RoleClient, ServiceExt};
26use serde_json::Value;
27use tokio::process::Command;
28
29pub use circuit_breaker::{
30 CircuitBreakerConfig, CircuitBreakerDispatcher, CircuitBreakerResourceDispatcher,
31};
32pub use reconnect::ReconnectingClient;
33pub use router::{RouterDispatcher, RouterResourceDispatcher};
34pub use timeout::{TimeoutDispatcher, TimeoutResourceDispatcher};
35
36#[derive(Debug, Clone)]
38#[non_exhaustive]
39pub enum TransportConfig {
40 Stdio {
42 command: String,
44 args: Vec<String>,
46 env: HashMap<String, String>,
48 },
49 Http {
51 url: String,
53 headers: HashMap<String, String>,
55 },
56}
57
58const STDIO_ENV_PASSTHROUGH: &[&str] = &[
59 "PATH",
60 "HOME",
61 "USERPROFILE",
62 "SYSTEMROOT",
63 "WINDIR",
64 "PATHEXT",
65 "TEMP",
66 "TMP",
67 "TMPDIR",
68 "SSL_CERT_FILE",
69 "SSL_CERT_DIR",
70];
71
72fn stdio_child_env(explicit: &HashMap<String, String>) -> HashMap<String, String> {
73 let mut env = HashMap::new();
74 for key in STDIO_ENV_PASSTHROUGH {
75 if let Ok(value) = std::env::var(key) {
76 env.insert((*key).to_string(), value);
77 }
78 }
79 for (key, value) in explicit {
80 env.insert(key.clone(), value.clone());
81 }
82 env
83}
84
85fn redact_stdio_args(args: &[String]) -> Vec<String> {
86 let mut redacted = Vec::with_capacity(args.len());
87 let mut redact_next = false;
88
89 for arg in args {
90 if redact_next {
91 redacted.push("[REDACTED]".to_string());
92 redact_next = false;
93 continue;
94 }
95
96 if let Some((name, _value)) = arg.split_once('=') {
97 if is_sensitive_arg_name(name) {
98 redacted.push(format!("{name}=[REDACTED]"));
99 continue;
100 }
101 }
102
103 if is_sensitive_arg_name(arg) {
104 redacted.push(arg.clone());
105 redact_next = true;
106 } else if looks_like_secret_value(arg) {
107 redacted.push("[REDACTED]".to_string());
108 } else {
109 redacted.push(arg.clone());
110 }
111 }
112
113 redacted
114}
115
116fn is_sensitive_arg_name(arg: &str) -> bool {
117 let name = arg
118 .trim_start_matches('-')
119 .to_ascii_lowercase()
120 .replace(['_', '-'], "");
121 name.contains("apikey")
122 || name.contains("accesstoken")
123 || name.contains("authtoken")
124 || name == "token"
125 || name.ends_with("token")
126 || name.contains("secret")
127 || name.contains("password")
128 || name.contains("credential")
129}
130
131fn looks_like_secret_value(arg: &str) -> bool {
132 let lower = arg.to_ascii_lowercase();
133 lower.starts_with("bearer ")
134 || lower.starts_with("sk-")
135 || lower.starts_with("ghp_")
136 || lower.starts_with("gho_")
137 || lower.starts_with("ghs_")
138 || lower.starts_with("ghr_")
139 || lower.starts_with("github_pat_")
140}
141
142fn redact_url_for_log(url: &str) -> Cow<'_, str> {
143 if let Some((base, _query)) = url.split_once('?') {
144 Cow::Owned(format!("{base}?[REDACTED]"))
145 } else {
146 Cow::Borrowed(url)
147 }
148}
149
150pub struct McpClient {
155 name: String,
156 inner: ClientInner,
157}
158
159enum ClientInner {
160 Stdio(RunningService<RoleClient, ()>),
161 Http(RunningService<RoleClient, ()>),
162}
163
164impl ClientInner {
165 fn peer(&self) -> &rmcp::Peer<RoleClient> {
166 match self {
167 ClientInner::Stdio(s) => s,
168 ClientInner::Http(s) => s,
169 }
170 }
171}
172
173#[derive(Debug, Clone)]
175pub struct ToolInfo {
176 pub name: String,
178 pub description: Option<String>,
180 pub input_schema: Value,
182}
183
184#[derive(Debug, Clone)]
186pub struct ResourceInfo {
187 pub uri: String,
189 pub name: String,
191 pub description: Option<String>,
193 pub mime_type: Option<String>,
195}
196
197impl McpClient {
198 pub async fn connect_stdio(
202 name: impl Into<String>,
203 command: &str,
204 args: &[&str],
205 ) -> Result<Self> {
206 Self::connect_stdio_with_env(name, command, args, HashMap::new()).await
207 }
208
209 pub async fn connect_stdio_with_env(
214 name: impl Into<String>,
215 command: &str,
216 args: &[&str],
217 env: HashMap<String, String>,
218 ) -> Result<Self> {
219 let name = name.into();
220 let args_owned: Vec<String> = args.iter().map(|s| s.to_string()).collect();
221 let child_env = stdio_child_env(&env);
222 let mut env_keys: Vec<&str> = child_env.keys().map(String::as_str).collect();
223 env_keys.sort_unstable();
224
225 tracing::info!(
226 server = %name,
227 command = %command,
228 arg_count = args_owned.len(),
229 env_keys = ?env_keys,
230 "connecting to downstream MCP server (stdio)"
231 );
232 tracing::debug!(
233 server = %name,
234 command = %command,
235 args = ?redact_stdio_args(&args_owned),
236 "stdio command arguments"
237 );
238
239 let transport = TokioChildProcess::new(Command::new(command).configure(|cmd| {
240 cmd.env_clear();
241 for (key, value) in &child_env {
242 cmd.env(key, value);
243 }
244 for arg in &args_owned {
245 cmd.arg(arg);
246 }
247 }))
248 .with_context(|| {
249 format!(
250 "failed to spawn stdio transport for server '{}' (command: {})",
251 name, command
252 )
253 })?;
254
255 let service: RunningService<RoleClient, ()> = ()
256 .serve(transport)
257 .await
258 .with_context(|| format!("MCP handshake failed for server '{}'", name))?;
259
260 tracing::info!(server = %name, "connected to downstream MCP server (stdio)");
261
262 Ok(Self {
263 name,
264 inner: ClientInner::Stdio(service),
265 })
266 }
267
268 pub async fn connect_http(
270 name: impl Into<String>,
271 url: &str,
272 headers: Option<HashMap<String, String>>,
273 ) -> Result<Self> {
274 let name = name.into();
275
276 if url.starts_with("http://") {
277 tracing::warn!(
278 server = %name,
279 url = %redact_url_for_log(url),
280 "connecting over plain HTTP — consider using HTTPS for production"
281 );
282 }
283
284 tracing::info!(
285 server = %name,
286 url = %redact_url_for_log(url),
287 "connecting to downstream MCP server (HTTP)"
288 );
289
290 let mut config = StreamableHttpClientTransportConfig::with_uri(url);
291
292 if let Some(ref hdrs) = headers {
294 check_http_credential_safety(url, hdrs)?;
295 }
296
297 let headers = headers.map(|mut h| {
299 sanitize_headers_for_transport(url, &mut h);
300 h
301 });
302
303 if let Some(hdrs) = &headers {
304 for (key, value) in hdrs {
305 if key.to_lowercase() == "authorization" {
306 tracing::debug!(server = %name, header = %key, "setting auth header (redacted)");
307 } else {
308 tracing::debug!(server = %name, header = %key, value = %value, "setting header");
309 }
310 }
311
312 let mut header_map = HashMap::new();
313 for (key, value) in hdrs {
314 let header_name = http::HeaderName::from_bytes(key.as_bytes())
315 .with_context(|| format!("invalid header name: {key}"))?;
316 let header_value = http::HeaderValue::from_str(value)
317 .with_context(|| format!("invalid header value for {key}"))?;
318 header_map.insert(header_name, header_value);
319 }
320 config = config.custom_headers(header_map);
321 }
322
323 let transport = StreamableHttpClientTransport::from_config(config);
324 let service: RunningService<RoleClient, ()> = ()
325 .serve(transport)
326 .await
327 .with_context(|| format!("MCP handshake failed for server '{}' (HTTP)", name))?;
328
329 tracing::info!(server = %name, "connected to downstream MCP server (HTTP)");
330
331 Ok(Self {
332 name,
333 inner: ClientInner::Http(service),
334 })
335 }
336
337 pub async fn connect(name: impl Into<String>, config: &TransportConfig) -> Result<Self> {
339 let name = name.into();
340 match config {
341 TransportConfig::Stdio { command, args, env } => {
342 let arg_refs: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
343 Self::connect_stdio_with_env(name, command, &arg_refs, env.clone()).await
344 }
345 TransportConfig::Http { url, headers } => {
346 let hdrs = if headers.is_empty() {
347 None
348 } else {
349 Some(headers.clone())
350 };
351 Self::connect_http(name, url, hdrs).await
352 }
353 }
354 }
355
356 pub async fn list_tools(&self) -> Result<Vec<ToolInfo>> {
358 let tools = self
359 .inner
360 .peer()
361 .list_all_tools()
362 .await
363 .with_context(|| format!("failed to list tools for server '{}'", self.name))?;
364
365 Ok(tools
366 .into_iter()
367 .map(|t| ToolInfo {
368 name: t.name.to_string(),
369 description: t.description.map(|d: Cow<'_, str>| d.to_string()),
370 input_schema: serde_json::to_value(&*t.input_schema)
371 .unwrap_or(Value::Object(Default::default())),
372 })
373 .collect())
374 }
375
376 pub async fn list_resources(&self) -> Result<Vec<ResourceInfo>> {
381 let result = self.inner.peer().list_resources(None).await;
382
383 match result {
384 Ok(list) => Ok(list
385 .resources
386 .into_iter()
387 .map(|r| ResourceInfo {
388 uri: r.uri.clone(),
389 name: r.name.clone(),
390 description: r.description.clone(),
391 mime_type: r.mime_type.clone(),
392 })
393 .collect()),
394 Err(e) => {
395 let err_str = e.to_string();
396 if err_str.contains("Method not found")
398 || err_str.contains("method not found")
399 || err_str.contains("-32601")
400 {
401 tracing::debug!(
402 server = %self.name,
403 "server does not support resources/list, returning empty"
404 );
405 Ok(vec![])
406 } else {
407 Err(anyhow::anyhow!(
408 "failed to list resources for server '{}': {}",
409 self.name,
410 e
411 ))
412 }
413 }
414 }
415 }
416
417 pub async fn read_resource(&self, uri: &str) -> Result<Value> {
419 let result = self
420 .inner
421 .peer()
422 .read_resource(rmcp::model::ReadResourceRequestParams::new(uri))
423 .await
424 .with_context(|| {
425 format!(
426 "resource read failed: server='{}', uri='{}'",
427 self.name, uri
428 )
429 })?;
430
431 let contents = result.contents;
433 if contents.is_empty() {
434 return Ok(Value::Null);
435 }
436
437 if contents.len() == 1 {
438 resource_content_to_value(&contents[0])
439 } else {
440 let values: Vec<Value> = contents
441 .iter()
442 .filter_map(|c| resource_content_to_value(c).ok())
443 .collect();
444 Ok(Value::Array(values))
445 }
446 }
447
448 pub fn name(&self) -> &str {
450 &self.name
451 }
452
453 pub async fn disconnect(self) -> Result<()> {
455 tracing::info!(server = %self.name, "disconnecting from downstream MCP server");
456 match self.inner {
457 ClientInner::Stdio(s) => {
458 let _ = s.cancel().await;
459 }
460 ClientInner::Http(s) => {
461 let _ = s.cancel().await;
462 }
463 }
464 Ok(())
465 }
466}
467
468#[async_trait::async_trait]
469impl ToolDispatcher for McpClient {
470 async fn call_tool(
471 &self,
472 _server: &str,
473 tool: &str,
474 args: Value,
475 ) -> Result<Value, forge_error::DispatchError> {
476 let arguments = args.as_object().cloned().or_else(|| {
477 if args.is_null() {
478 Some(serde_json::Map::new())
479 } else {
480 None
481 }
482 });
483
484 let result: CallToolResult = self
485 .inner
486 .peer()
487 .call_tool(
488 CallToolRequestParams::new(tool.to_string())
489 .with_arguments(arguments.unwrap_or_default()),
490 )
491 .await
492 .map_err(|e| {
493 let msg = format!("tool call failed: tool='{}': {}", tool, e);
494 let err_str = e.to_string();
495 if is_transport_dead(&err_str) {
496 forge_error::DispatchError::TransportDead {
497 server: self.name.clone(),
498 reason: msg,
499 }
500 } else {
501 forge_error::DispatchError::Upstream {
502 server: self.name.clone(),
503 message: msg,
504 }
505 }
506 })?;
507
508 if result.is_error == Some(true) && result.structured_content.is_none() {
512 let error_text = result
513 .content
514 .iter()
515 .filter_map(|c| match &c.raw {
516 RawContent::Text(t) => Some(t.text.as_str()),
517 _ => None,
518 })
519 .collect::<Vec<_>>()
520 .join("\n");
521 return Err(forge_error::DispatchError::ToolError {
522 server: self.name.clone(),
523 tool: tool.to_string(),
524 message: format!("tool returned error: {}", error_text),
525 });
526 }
527
528 call_tool_result_to_value(result).map_err(|e| forge_error::DispatchError::Upstream {
529 server: self.name.clone(),
530 message: e.to_string(),
531 })
532 }
533}
534
535#[async_trait::async_trait]
536impl ResourceDispatcher for McpClient {
537 async fn read_resource(
538 &self,
539 _server: &str,
540 uri: &str,
541 ) -> Result<Value, forge_error::DispatchError> {
542 self.read_resource(uri).await.map_err(|e| {
543 let msg = format!("resource read failed: uri='{}': {}", uri, e);
544 let err_str = e.to_string();
545 if is_transport_dead(&err_str) {
546 forge_error::DispatchError::TransportDead {
547 server: self.name.clone(),
548 reason: msg,
549 }
550 } else {
551 forge_error::DispatchError::Upstream {
552 server: self.name.clone(),
553 message: msg,
554 }
555 }
556 })
557 }
558}
559
560fn is_transport_dead(err_str: &str) -> bool {
565 err_str.contains("TransportClosed")
566 || err_str.contains("transport closed")
567 || err_str.contains("channel closed")
568 || err_str.contains("broken pipe")
569 || err_str.contains("Broken pipe")
570 || err_str.contains("BrokenPipe")
571}
572
573fn resource_content_to_value(content: &rmcp::model::ResourceContents) -> Result<Value> {
575 match content {
576 rmcp::model::ResourceContents::TextResourceContents { text, .. } => {
577 serde_json::from_str(text).or_else(|_| Ok(Value::String(text.clone())))
579 }
580 rmcp::model::ResourceContents::BlobResourceContents {
581 blob, mime_type, ..
582 } => Ok(serde_json::json!({
583 "_type": "blob",
584 "_encoding": "base64",
585 "data": blob,
586 "mime_type": mime_type.as_deref().unwrap_or("application/octet-stream"),
587 })),
588 }
589}
590
591fn call_tool_result_to_value(result: CallToolResult) -> Result<Value> {
593 if let Some(structured) = result.structured_content {
594 return Ok(structured);
595 }
596
597 if result.is_error == Some(true) {
598 let error_text = result
599 .content
600 .iter()
601 .filter_map(|c| match &c.raw {
602 RawContent::Text(t) => Some(t.text.as_str()),
603 _ => None,
604 })
605 .collect::<Vec<_>>()
606 .join("\n");
607 return Err(anyhow::anyhow!("tool returned error: {}", error_text));
608 }
609
610 if result.content.len() == 1 {
611 content_to_value(&result.content[0])
612 } else if result.content.is_empty() {
613 Ok(Value::Null)
614 } else {
615 let values: Vec<Value> = result
616 .content
617 .iter()
618 .filter_map(|c| content_to_value(c).ok())
619 .collect();
620 Ok(Value::Array(values))
621 }
622}
623
624const MAX_BINARY_CONTENT_SIZE: usize = 1_048_576; const MAX_TEXT_CONTENT_SIZE: usize = 10_485_760; fn floor_char_boundary(s: &str, max: usize) -> usize {
632 let mut end = max.min(s.len());
633 while end > 0 && !s.is_char_boundary(end) {
634 end -= 1;
635 }
636 end
637}
638
639fn content_to_value(content: &Content) -> Result<Value> {
644 match &content.raw {
645 RawContent::Text(t) => {
646 if t.text.len() > MAX_TEXT_CONTENT_SIZE {
647 let preview_end = floor_char_boundary(&t.text, 1024);
648 Ok(serde_json::json!({
649 "type": "text",
650 "truncated": true,
651 "original_size": t.text.len(),
652 "preview": &t.text[..preview_end],
653 }))
654 } else {
655 serde_json::from_str(&t.text).or_else(|_| Ok(Value::String(t.text.clone())))
656 }
657 }
658 RawContent::Image(img) => {
659 if img.data.len() > MAX_BINARY_CONTENT_SIZE {
660 Ok(serde_json::json!({
661 "type": "image",
662 "truncated": true,
663 "original_size": img.data.len(),
664 "mime_type": img.mime_type,
665 }))
666 } else {
667 Ok(serde_json::json!({
668 "type": "image",
669 "data": img.data,
670 "mime_type": img.mime_type,
671 }))
672 }
673 }
674 RawContent::Resource(r) => Ok(serde_json::json!({
675 "type": "resource",
676 "resource": serde_json::to_value(&r.resource).unwrap_or(Value::Null),
677 })),
678 RawContent::Audio(a) => {
679 if a.data.len() > MAX_BINARY_CONTENT_SIZE {
680 Ok(serde_json::json!({
681 "type": "audio",
682 "truncated": true,
683 "original_size": a.data.len(),
684 "mime_type": a.mime_type,
685 }))
686 } else {
687 Ok(serde_json::json!({
688 "type": "audio",
689 "data": a.data,
690 "mime_type": a.mime_type,
691 }))
692 }
693 }
694 _ => Ok(serde_json::json!({"type": "unknown"})),
695 }
696}
697
698const SENSITIVE_HEADER_PATTERNS: &[&str] = &[
701 "authorization",
702 "cookie",
703 "token",
704 "secret",
705 "key",
706 "credential",
707 "password",
708 "auth",
709];
710
711fn is_sensitive_header(name: &str) -> bool {
713 let lower = name.to_lowercase();
714 SENSITIVE_HEADER_PATTERNS
715 .iter()
716 .any(|pattern| lower.contains(pattern))
717}
718
719fn check_http_credential_safety(
725 url: &str,
726 headers: &HashMap<String, String>,
727) -> Result<(), anyhow::Error> {
728 if url.starts_with("http://") {
729 let sensitive: Vec<&String> = headers.keys().filter(|k| is_sensitive_header(k)).collect();
730 if !sensitive.is_empty() {
731 return Err(anyhow::anyhow!(
732 "refusing to send credentials over plain HTTP (headers: {}). \
733 Use HTTPS or remove sensitive headers.",
734 sensitive
735 .iter()
736 .map(|s| s.as_str())
737 .collect::<Vec<_>>()
738 .join(", ")
739 ));
740 }
741 }
742 Ok(())
743}
744
745fn sanitize_headers_for_transport(url: &str, headers: &mut HashMap<String, String>) {
752 if url.starts_with("http://") {
753 let removed: Vec<String> = headers
754 .keys()
755 .filter(|k| is_sensitive_header(k))
756 .cloned()
757 .collect();
758 for key in &removed {
759 headers.remove(key);
760 }
761 if !removed.is_empty() {
762 tracing::warn!(
763 url = %url,
764 removed_headers = ?removed,
765 "stripped sensitive headers from plain HTTP connection — use HTTPS to send credentials"
766 );
767 }
768 }
769}
770
771#[cfg(test)]
772mod tests {
773 use super::*;
774 use rmcp::model::{Content, RawContent};
775
776 #[test]
777 fn content_to_value_text_string() {
778 let content = Content::text("hello");
779 let val = content_to_value(&content).unwrap();
780 assert_eq!(val, Value::String("hello".into()));
781 }
782
783 #[test]
784 fn content_to_value_text_json() {
785 let content = Content::text(r#"{"k":"v"}"#);
786 let val = content_to_value(&content).unwrap();
787 assert_eq!(val, serde_json::json!({"k": "v"}));
788 }
789
790 #[test]
791 fn content_to_value_small_image_preserved() {
792 let small_data = "a".repeat(1024); let content = Content::image(small_data.clone(), "image/png");
794 let val = content_to_value(&content).unwrap();
795 assert_eq!(val["type"], "image");
796 assert_eq!(val["data"], small_data);
797 assert!(val.get("truncated").is_none());
798 }
799
800 #[test]
801 fn content_to_value_oversized_image_truncated() {
802 let large_data = "a".repeat(2 * 1024 * 1024); let content = Content::image(large_data, "image/png");
804 let val = content_to_value(&content).unwrap();
805 assert_eq!(val["type"], "image");
806 assert_eq!(val["truncated"], true);
807 assert!(val.get("data").is_none());
808 assert!(val["original_size"].as_u64().unwrap() > MAX_BINARY_CONTENT_SIZE as u64);
809 }
810
811 #[test]
812 fn content_to_value_oversized_audio_truncated() {
813 let large_data = "a".repeat(2 * 1024 * 1024); let content = Content {
815 raw: RawContent::Audio(rmcp::model::RawAudioContent {
816 data: large_data,
817 mime_type: "audio/wav".into(),
818 }),
819 annotations: None,
820 };
821 let val = content_to_value(&content).unwrap();
822 assert_eq!(val["type"], "audio");
823 assert_eq!(val["truncated"], true);
824 assert!(val.get("data").is_none());
825 }
826
827 #[test]
828 fn content_to_value_oversized_text_truncated() {
829 let large_text = "x".repeat(11 * 1024 * 1024); let content = Content::text(large_text);
831 let val = content_to_value(&content).unwrap();
832 assert_eq!(val["type"], "text");
833 assert_eq!(val["truncated"], true);
834 assert!(val["original_size"].as_u64().unwrap() > MAX_TEXT_CONTENT_SIZE as u64);
835 assert!(val["preview"].as_str().unwrap().len() <= 1024);
836 }
837
838 #[test]
839 fn content_to_value_oversized_text_truncates_on_char_boundary() {
840 let large_text = "€".repeat((MAX_TEXT_CONTENT_SIZE / 3) + 100);
841 let content = Content::text(large_text);
842 let val = content_to_value(&content).unwrap();
843 assert_eq!(val["type"], "text");
844 assert_eq!(val["truncated"], true);
845 assert!(val["preview"].as_str().unwrap().is_char_boundary(0));
846 }
847
848 #[test]
849 fn content_to_value_normal_text_not_truncated() {
850 let normal_text = "x".repeat(1024); let content = Content::text(normal_text.clone());
852 let val = content_to_value(&content).unwrap();
853 assert_eq!(val, Value::String(normal_text));
854 }
855
856 #[test]
857 fn sanitize_headers_strips_auth_on_http() {
858 let mut headers = HashMap::new();
859 headers.insert("Authorization".into(), "Bearer secret".into());
860 headers.insert("Content-Type".into(), "application/json".into());
861 sanitize_headers_for_transport("http://example.com/mcp", &mut headers);
862 assert!(!headers.contains_key("Authorization"));
863 assert!(headers.contains_key("Content-Type"));
864 }
865
866 #[test]
867 fn sanitize_headers_strips_api_key_on_http() {
868 let mut headers = HashMap::new();
869 headers.insert("X-Api-Key".into(), "sk-123".into());
870 headers.insert("Content-Type".into(), "application/json".into());
871 sanitize_headers_for_transport("http://example.com/mcp", &mut headers);
872 assert!(!headers.contains_key("X-Api-Key"));
873 assert!(headers.contains_key("Content-Type"));
874 }
875
876 #[test]
877 fn sanitize_headers_strips_cookie_on_http() {
878 let mut headers = HashMap::new();
879 headers.insert("Cookie".into(), "session=abc123".into());
880 sanitize_headers_for_transport("http://example.com/mcp", &mut headers);
881 assert!(!headers.contains_key("Cookie"));
882 }
883
884 #[test]
885 fn sanitize_headers_strips_custom_token_on_http() {
886 let mut headers = HashMap::new();
887 headers.insert("X-Auth-Token".into(), "tok_secret".into());
888 headers.insert("X-Secret-Key".into(), "s3cr3t".into());
889 headers.insert("X-Custom-Credential".into(), "cred".into());
890 headers.insert("X-Password".into(), "pass".into());
891 headers.insert("Accept".into(), "application/json".into());
892 sanitize_headers_for_transport("http://example.com/mcp", &mut headers);
893 assert!(!headers.contains_key("X-Auth-Token"));
894 assert!(!headers.contains_key("X-Secret-Key"));
895 assert!(!headers.contains_key("X-Custom-Credential"));
896 assert!(!headers.contains_key("X-Password"));
897 assert!(headers.contains_key("Accept"));
898 }
899
900 #[test]
901 fn sanitize_headers_preserves_all_on_https() {
902 let mut headers = HashMap::new();
903 headers.insert("Authorization".into(), "Bearer secret".into());
904 headers.insert("X-Api-Key".into(), "sk-123".into());
905 headers.insert("Cookie".into(), "session=abc".into());
906 sanitize_headers_for_transport("https://example.com/mcp", &mut headers);
907 assert!(headers.contains_key("Authorization"));
908 assert!(headers.contains_key("X-Api-Key"));
909 assert!(headers.contains_key("Cookie"));
910 }
911
912 #[test]
914 fn http_sec_01_rejects_creds_on_http() {
915 let mut headers = HashMap::new();
916 headers.insert("Authorization".into(), "Bearer secret".into());
917 let err = check_http_credential_safety("http://example.com/mcp", &headers);
918 assert!(err.is_err(), "should reject creds on HTTP");
919 let msg = err.unwrap_err().to_string();
920 assert!(
921 msg.contains("plain HTTP"),
922 "error should mention plain HTTP: {msg}"
923 );
924 }
925
926 #[test]
928 fn http_sec_02_allows_http_without_creds() {
929 let mut headers = HashMap::new();
930 headers.insert("Content-Type".into(), "application/json".into());
931 assert!(check_http_credential_safety("http://example.com/mcp", &headers).is_ok());
932 assert!(check_http_credential_safety("http://example.com/mcp", &HashMap::new()).is_ok());
934 }
935
936 #[test]
938 fn http_sec_03_allows_https_with_creds() {
939 let mut headers = HashMap::new();
940 headers.insert("Authorization".into(), "Bearer secret".into());
941 headers.insert("X-Api-Key".into(), "sk-123".into());
942 assert!(check_http_credential_safety("https://example.com/mcp", &headers).is_ok());
943 }
944
945 #[test]
946 fn is_sensitive_header_matches() {
947 assert!(is_sensitive_header("Authorization"));
948 assert!(is_sensitive_header("x-api-key"));
949 assert!(is_sensitive_header("Cookie"));
950 assert!(is_sensitive_header("X-Auth-Token"));
951 assert!(is_sensitive_header("X-Secret-Key"));
952 assert!(is_sensitive_header("X-Custom-Credential"));
953 assert!(is_sensitive_header("X-Password"));
954 assert!(!is_sensitive_header("Content-Type"));
955 assert!(!is_sensitive_header("Accept"));
956 assert!(!is_sensitive_header("User-Agent"));
957 }
958
959 #[test]
960 fn stdio_child_env_only_preserves_launch_allowlist_and_explicit_env() {
961 let mut explicit = HashMap::new();
962 explicit.insert("GITHUB_TOKEN".to_string(), "secret123".to_string());
963
964 let env = stdio_child_env(&explicit);
965
966 assert_eq!(
967 env.get("GITHUB_TOKEN").map(String::as_str),
968 Some("secret123")
969 );
970 for key in env.keys() {
971 assert!(
972 STDIO_ENV_PASSTHROUGH.contains(&key.as_str()) || key == "GITHUB_TOKEN",
973 "unexpected inherited env var: {key}"
974 );
975 }
976 }
977
978 #[test]
979 fn redact_stdio_args_removes_sensitive_values() {
980 let args = vec![
981 "--repos".to_string(),
982 ".".to_string(),
983 "--access-token".to_string(),
984 "secret-token".to_string(),
985 "--api-key=sk-live-123".to_string(),
986 "ghp_abcdefghijklmnopqrstuvwxyz".to_string(),
987 ];
988
989 let redacted = redact_stdio_args(&args);
990
991 assert_eq!(redacted[0], "--repos");
992 assert_eq!(redacted[1], ".");
993 assert_eq!(redacted[2], "--access-token");
994 assert_eq!(redacted[3], "[REDACTED]");
995 assert_eq!(redacted[4], "--api-key=[REDACTED]");
996 assert_eq!(redacted[5], "[REDACTED]");
997 }
998
999 #[test]
1000 fn redact_url_for_log_removes_query_string() {
1001 assert_eq!(
1002 redact_url_for_log("https://example.com/mcp?token=secret").as_ref(),
1003 "https://example.com/mcp?[REDACTED]"
1004 );
1005 assert_eq!(
1006 redact_url_for_log("https://example.com/mcp").as_ref(),
1007 "https://example.com/mcp"
1008 );
1009 }
1010
1011 #[test]
1014 fn call_tool_result_is_error_true_returns_err() {
1015 let result = CallToolResult::error(vec![Content::text(
1016 "Invalid params: missing field 'base_url'",
1017 )]);
1018 let err = call_tool_result_to_value(result);
1019 assert!(err.is_err());
1020 let msg = err.unwrap_err().to_string();
1021 assert!(
1022 msg.contains("Invalid params"),
1023 "expected error text, got: {msg}"
1024 );
1025 }
1026
1027 #[test]
1028 fn call_tool_result_success_returns_ok() {
1029 let result = CallToolResult::success(vec![Content::text(r#"{"status": "ok"}"#)]);
1030 let val = call_tool_result_to_value(result).unwrap();
1031 assert_eq!(val["status"], "ok");
1032 }
1033
1034 #[test]
1035 fn call_tool_result_structured_content_takes_priority_over_is_error() {
1036 let structured = serde_json::json!({"data": "important"});
1037 let mut result = CallToolResult::error(vec![Content::text("error text")]);
1038 result.structured_content = Some(structured.clone());
1039 let val = call_tool_result_to_value(result).unwrap();
1040 assert_eq!(val, structured);
1041 }
1042
1043 #[test]
1046 fn transport_dead_detects_transport_closed() {
1047 assert!(is_transport_dead("TransportClosed: channel full"));
1048 assert!(is_transport_dead("error: transport closed unexpectedly"));
1049 assert!(is_transport_dead("channel closed by peer"));
1050 assert!(is_transport_dead("broken pipe while writing"));
1051 assert!(is_transport_dead("Broken pipe (os error 32)"));
1052 assert!(is_transport_dead("BrokenPipe"));
1053 }
1054
1055 #[test]
1056 fn transport_dead_rejects_normal_errors() {
1057 assert!(!is_transport_dead("tool not found: echo"));
1058 assert!(!is_transport_dead("timeout after 5000ms"));
1059 assert!(!is_transport_dead("Invalid params: missing field"));
1060 assert!(!is_transport_dead("connection refused"));
1061 }
1062
1063 #[test]
1065 #[allow(unreachable_patterns)]
1066 fn ne_transport_config_is_non_exhaustive() {
1067 let config = TransportConfig::Stdio {
1068 command: "test".into(),
1069 args: vec![],
1070 env: HashMap::new(),
1071 };
1072 match config {
1073 TransportConfig::Stdio { .. } | TransportConfig::Http { .. } => {}
1074 _ => {}
1075 }
1076 }
1077}