1use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::time::Duration;
12
13use crate::discovery::{MCPServerConfig, ResolvedTransport, ToolMetadata};
14use crate::transport::http::HttpTransport;
15use crate::transport::stdio::StdioTransport;
16use crate::transport::{
17 DEFAULT_CONNECT_TIMEOUT, DEFAULT_REQUEST_TIMEOUT, JsonRpcRequest, Transport,
18};
19
20const MAX_TOOL_PAGES: usize = 100;
23
24#[derive(Debug, Clone, Serialize, Deserialize, Default)]
26pub struct ServerCapabilities {
27 pub tools: Option<ToolsCapability>,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, Default)]
33pub struct ToolsCapability {
34 #[serde(rename = "listChanged")]
36 pub list_changed: Option<bool>,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct ToolResult {
42 #[serde(default)]
47 pub content: Vec<ToolResultContent>,
48 #[serde(
51 rename = "structuredContent",
52 default,
53 skip_serializing_if = "Option::is_none"
54 )]
55 pub structured_content: Option<Value>,
56 #[serde(rename = "isError", default)]
62 pub is_error: bool,
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
70pub struct EmbeddedResource {
71 pub uri: String,
73 #[serde(default, skip_serializing_if = "Option::is_none")]
75 pub text: Option<String>,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
78 pub blob: Option<String>,
79 #[serde(rename = "mimeType", default, skip_serializing_if = "Option::is_none")]
81 pub mime_type: Option<String>,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
90#[serde(tag = "type")]
91pub enum ToolResultContent {
92 #[serde(rename = "text")]
94 Text { text: String },
95 #[serde(rename = "image")]
97 Image {
98 data: String,
99 #[serde(rename = "mimeType")]
100 mime_type: String,
101 },
102 #[serde(rename = "audio")]
104 Audio {
105 data: String,
106 #[serde(rename = "mimeType")]
107 mime_type: String,
108 },
109 #[serde(rename = "resource_link")]
111 ResourceLink {
112 uri: String,
113 #[serde(default)]
114 name: String,
115 #[serde(default, skip_serializing_if = "Option::is_none")]
116 description: Option<String>,
117 #[serde(rename = "mimeType", default, skip_serializing_if = "Option::is_none")]
118 mime_type: Option<String>,
119 },
120 #[serde(rename = "resource")]
122 Resource { resource: EmbeddedResource },
123 #[serde(other)]
131 Unknown,
132}
133pub const SUPPORTED_PROTOCOL_VERSIONS: &[&str] =
141 &["2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05"];
142
143pub const PREFERRED_PROTOCOL_VERSION: &str = SUPPORTED_PROTOCOL_VERSIONS[0];
145
146pub struct MCPClient {
148 transport: Box<dyn Transport>,
150 next_id: AtomicU64,
152 request_timeout: Duration,
154 capabilities: Option<ServerCapabilities>,
156 protocol_version: Option<String>,
158 cached_tools: Vec<ToolMetadata>,
160}
161
162impl MCPClient {
163 pub(crate) fn new(transport: Box<dyn Transport>) -> Self {
165 Self {
166 transport,
167 next_id: AtomicU64::new(1),
168 request_timeout: DEFAULT_REQUEST_TIMEOUT,
169 capabilities: None,
170 protocol_version: None,
171 cached_tools: Vec::new(),
172 }
173 }
174
175 pub async fn spawn(
177 command: &str,
178 args: &[&str],
179 env: &HashMap<String, String>,
180 ) -> anyhow::Result<Self> {
181 let transport = StdioTransport::spawn(command, args, env).await?;
182 Ok(Self::new(Box::new(transport)))
183 }
184
185 pub fn connect_http(
187 url: &str,
188 headers: &HashMap<String, String>,
189 allow_env: &[String],
190 ) -> anyhow::Result<Self> {
191 let transport = HttpTransport::new(url, headers, allow_env)?;
192 Ok(Self::new(Box::new(transport)))
193 }
194
195 pub fn set_refresher(
201 &mut self,
202 refresher: std::sync::Arc<dyn crate::transport::BearerRefresher>,
203 ) {
204 self.transport.set_bearer_refresher(refresher);
205 }
206
207 pub async fn from_config(config: &MCPServerConfig) -> anyhow::Result<Self> {
213 Self::from_config_with_auth(config, None, &[]).await
214 }
215
216 pub async fn from_config_with_auth(
225 config: &MCPServerConfig,
226 auth_header: Option<(String, String)>,
227 allow_env: &[String],
228 ) -> anyhow::Result<Self> {
229 match config.resolve()? {
230 ResolvedTransport::Stdio { command, args, env } => {
231 let args: Vec<&str> = args.iter().map(String::as_str).collect();
232 Self::spawn(command, &args, env).await
233 }
234 ResolvedTransport::Http { url, headers } => {
235 let mut headers = headers.clone();
236 if let Some((name, value)) = auth_header {
237 headers.insert(name, value);
238 }
239 Self::connect_http(url, &headers, allow_env)
240 }
241 }
242 }
243
244 pub async fn connect(&mut self) -> anyhow::Result<()> {
246 tracing::info!("Initializing MCP connection");
247
248 let init_params = serde_json::json!({
249 "protocolVersion": PREFERRED_PROTOCOL_VERSION,
250 "capabilities": {},
251 "clientInfo": {
252 "name": "leviath",
253 "version": env!("CARGO_PKG_VERSION")
254 }
255 });
256
257 let result = self
261 .request_with_timeout("initialize", init_params, DEFAULT_CONNECT_TIMEOUT)
262 .await?;
263
264 let capabilities: ServerCapabilities = if let Some(caps) = result.get("capabilities") {
266 serde_json::from_value(caps.clone()).unwrap_or_default()
267 } else {
268 ServerCapabilities::default()
269 };
270 self.capabilities = Some(capabilities);
271 let version = negotiated_version(result.get("protocolVersion"));
272 self.protocol_version = Some(version.clone());
273
274 self.send_notification("notifications/initialized", serde_json::json!({}))
276 .await?;
277
278 tracing::info!(version = %version, "MCP connection established");
279 Ok(())
280 }
281
282 pub async fn list_tools(&mut self) -> anyhow::Result<Vec<ToolMetadata>> {
290 tracing::debug!("Listing MCP tools");
291
292 let mut tools: Vec<ToolMetadata> = Vec::new();
293 let mut cursor: Option<String> = None;
294
295 for page in 0..MAX_TOOL_PAGES {
296 let params = match &cursor {
297 Some(c) => serde_json::json!({ "cursor": c }),
298 None => serde_json::json!({}),
299 };
300 let result = self.send_request("tools/list", params).await?;
301
302 let tools_value = result.get("tools").cloned().unwrap_or(Value::Array(vec![]));
303 let page_tools: Vec<ToolMetadata> = serde_json::from_value(tools_value)
304 .map_err(|e| anyhow::anyhow!("Failed to parse tools list: {}", e))?;
305 tools.extend(page_tools);
306
307 cursor = result
308 .get("nextCursor")
309 .and_then(Value::as_str)
310 .map(str::to_string);
311 if cursor.is_none() {
312 break;
313 }
314 if page + 1 == MAX_TOOL_PAGES {
315 tracing::warn!(
316 pages = MAX_TOOL_PAGES,
317 "MCP server still returned a tools/list cursor at the page \
318 limit - stopping; some tools may be missing"
319 );
320 }
321 }
322
323 self.cached_tools = tools.clone();
324 let count = tools.len();
328 tracing::debug!(count, "Discovered MCP tools");
329 Ok(tools)
330 }
331
332 pub async fn call_tool(&mut self, name: &str, arguments: Value) -> anyhow::Result<ToolResult> {
334 tracing::debug!(tool = %name, "Calling MCP tool");
335
336 let params = serde_json::json!({
337 "name": name,
338 "arguments": arguments,
339 });
340
341 let result = self.send_request("tools/call", params).await?;
342
343 let tool_result: ToolResult = serde_json::from_value(result)
344 .map_err(|e| anyhow::anyhow!("Failed to parse tool result: {}", e))?;
345
346 Ok(tool_result)
347 }
348
349 pub async fn shutdown(&mut self) -> anyhow::Result<()> {
353 tracing::info!("Shutting down MCP server");
354 let _ = self.transport.close().await;
355 Ok(())
356 }
357
358 pub fn capabilities(&self) -> Option<&ServerCapabilities> {
360 self.capabilities.as_ref()
361 }
362
363 pub fn protocol_version(&self) -> Option<&str> {
365 self.protocol_version.as_deref()
366 }
367
368 pub fn cached_tools(&self) -> &[ToolMetadata] {
370 &self.cached_tools
371 }
372
373 async fn send_request(&mut self, method: &str, params: Value) -> anyhow::Result<Value> {
375 self.request_with_timeout(method, params, self.request_timeout)
376 .await
377 }
378
379 async fn request_with_timeout(
381 &mut self,
382 method: &str,
383 params: Value,
384 timeout: Duration,
385 ) -> anyhow::Result<Value> {
386 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
387 let request = JsonRpcRequest::request(id, method, params);
388 self.transport
389 .send_request(&request, timeout)
390 .await?
391 .into_result()
392 }
393
394 async fn send_notification(&mut self, method: &str, params: Value) -> anyhow::Result<()> {
396 let request = JsonRpcRequest::notification(method, params);
397 self.transport.send_notification(&request).await
398 }
399}
400
401fn negotiated_version(echoed: Option<&Value>) -> String {
409 match echoed.and_then(Value::as_str) {
410 Some(version) => {
411 if !SUPPORTED_PROTOCOL_VERSIONS.contains(&version) {
412 tracing::warn!(
413 version = %version,
414 "MCP server negotiated an unrecognized protocol revision - continuing"
415 );
416 }
417 version.to_string()
418 }
419 None => {
420 tracing::debug!("MCP server echoed no protocolVersion - assuming the offered one");
422 PREFERRED_PROTOCOL_VERSION.to_string()
423 }
424 }
425}
426
427#[cfg(test)]
428mod tests {
429 use super::*;
430 use crate::test_support::always_on_tracing_guard;
431
432 #[test]
435 fn test_server_capabilities_default() {
436 let caps = ServerCapabilities::default();
437 assert!(caps.tools.is_none());
438 }
439
440 #[test]
441 fn test_tools_capability_default() {
442 let cap = ToolsCapability::default();
443 assert!(cap.list_changed.is_none());
444 }
445
446 #[test]
447 fn test_server_capabilities_serialization() {
448 let caps = ServerCapabilities {
449 tools: Some(ToolsCapability {
450 list_changed: Some(true),
451 }),
452 };
453 let json = serde_json::to_string(&caps).unwrap();
454 assert!(json.contains("listChanged"));
455 assert!(json.contains("true"));
456
457 let deserialized: ServerCapabilities = serde_json::from_str(&json).unwrap();
458 assert!(deserialized.tools.unwrap().list_changed.unwrap());
459 }
460
461 #[test]
462 fn test_tool_result_serialization() {
463 let result = ToolResult {
464 content: vec![ToolResultContent::Text {
465 text: "Hello".to_string(),
466 }],
467 structured_content: None,
468 is_error: false,
469 };
470 let json = serde_json::to_string(&result).unwrap();
471 assert!(json.contains("Hello"));
472 assert!(json.contains("\"isError\":false"), "got: {json}");
475 assert!(!json.contains("is_error"), "got: {json}");
476 }
477
478 #[test]
479 fn test_tool_result_with_error() {
480 let result = ToolResult {
481 content: vec![ToolResultContent::Text {
482 text: "Something went wrong".to_string(),
483 }],
484 structured_content: None,
485 is_error: true,
486 };
487 assert!(result.is_error);
488 assert_eq!(result.content.len(), 1);
489 }
490
491 #[test]
492 fn test_tool_result_content_text() {
493 let content = ToolResultContent::Text {
494 text: "result text".to_string(),
495 };
496 let json = serde_json::to_string(&content).unwrap();
497 assert!(json.contains("result text"));
498 assert!(json.contains("\"type\":\"text\""));
499 }
500
501 #[test]
502 fn test_tool_result_content_image() {
503 let content = ToolResultContent::Image {
504 data: "base64data".to_string(),
505 mime_type: "image/png".to_string(),
506 };
507 let json = serde_json::to_string(&content).unwrap();
508 assert!(json.contains("base64data"));
509 assert!(json.contains("image/png"));
510 }
511
512 #[test]
513 fn test_tool_result_content_resource() {
514 let content = ToolResultContent::Resource {
515 resource: EmbeddedResource {
516 uri: "file:///tmp/test.txt".to_string(),
517 text: Some("file contents".to_string()),
518 blob: None,
519 mime_type: None,
520 },
521 };
522 let json = serde_json::to_string(&content).unwrap();
523 assert!(json.contains("file:///tmp/test.txt"));
524 assert!(json.contains("file contents"));
525 }
526
527 #[test]
528 fn test_tool_result_content_resource_no_text() {
529 let content = ToolResultContent::Resource {
530 resource: EmbeddedResource {
531 uri: "file:///tmp/test.txt".to_string(),
532 text: None,
533 blob: None,
534 mime_type: None,
535 },
536 };
537 let json = serde_json::to_string(&content).unwrap();
538 assert!(json.contains("file:///tmp/test.txt"));
539 }
540
541 #[test]
542 fn test_tool_result_deserialization() {
543 let json = r#"{"content":[{"type":"text","text":"Hello"}],"is_error":false}"#;
544 let result: ToolResult = serde_json::from_str(json).unwrap();
545 assert!(!result.is_error);
546 assert_eq!(result.content.len(), 1);
547 }
548
549 #[test]
550 fn test_tool_result_deserialization_missing_is_error() {
551 let json = r#"{"content":[{"type":"text","text":"Hello"}]}"#;
552 let result: ToolResult = serde_json::from_str(json).unwrap();
553 assert!(!result.is_error); }
555
556 #[test]
557 fn test_tool_result_multiple_content() {
558 let result = ToolResult {
559 content: vec![
560 ToolResultContent::Text {
561 text: "line 1".to_string(),
562 },
563 ToolResultContent::Text {
564 text: "line 2".to_string(),
565 },
566 ],
567 structured_content: None,
568 is_error: false,
569 };
570 assert_eq!(result.content.len(), 2);
571 }
572
573 #[test]
574 fn test_tool_result_clone() {
575 let result = ToolResult {
576 content: vec![ToolResultContent::Text {
577 text: "test".to_string(),
578 }],
579 structured_content: None,
580 is_error: true,
581 };
582 let cloned = result.clone();
583 assert!(cloned.is_error);
584 assert_eq!(cloned.content.len(), 1);
585 }
586
587 #[test]
588 fn test_server_capabilities_clone() {
589 let caps = ServerCapabilities {
590 tools: Some(ToolsCapability {
591 list_changed: Some(true),
592 }),
593 };
594 let cloned = caps.clone();
595 assert!(cloned.tools.unwrap().list_changed.unwrap());
596 }
597
598 #[test]
601 fn test_tool_result_empty_content() {
602 let result = ToolResult {
603 content: vec![],
604 structured_content: None,
605 is_error: false,
606 };
607 assert!(result.content.is_empty());
608 assert!(!result.is_error);
609 }
610
611 #[test]
612 fn test_tool_result_mixed_content_types() {
613 let result = ToolResult {
614 content: vec![
615 ToolResultContent::Text {
616 text: "hello".to_string(),
617 },
618 ToolResultContent::Image {
619 data: "base64".to_string(),
620 mime_type: "image/jpeg".to_string(),
621 },
622 ToolResultContent::Resource {
623 resource: EmbeddedResource {
624 uri: "file:///test".to_string(),
625 text: Some("content".to_string()),
626 blob: None,
627 mime_type: None,
628 },
629 },
630 ],
631 structured_content: None,
632 is_error: false,
633 };
634 assert_eq!(result.content.len(), 3);
635 let json = serde_json::to_string(&result).unwrap();
636 let back: ToolResult = serde_json::from_str(&json).unwrap();
637 assert_eq!(back.content.len(), 3);
638 }
639
640 #[test]
643 fn test_server_capabilities_from_empty_json() {
644 let caps: ServerCapabilities = serde_json::from_str("{}").unwrap();
645 assert!(caps.tools.is_none());
646 }
647
648 #[test]
649 fn test_server_capabilities_with_tools() {
650 let json = r#"{"tools":{"listChanged":false}}"#;
651 let caps: ServerCapabilities = serde_json::from_str(json).unwrap();
652 let tools = caps.tools.unwrap();
653 assert_eq!(tools.list_changed, Some(false));
654 }
655
656 #[test]
657 fn test_server_capabilities_with_null_tools() {
658 let json = r#"{"tools":null}"#;
659 let caps: ServerCapabilities = serde_json::from_str(json).unwrap();
660 assert!(caps.tools.is_none());
661 }
662
663 #[test]
666 fn test_tools_capability_with_list_changed_true() {
667 let cap = ToolsCapability {
668 list_changed: Some(true),
669 };
670 let json = serde_json::to_string(&cap).unwrap();
671 assert!(json.contains("true"));
672 let back: ToolsCapability = serde_json::from_str(&json).unwrap();
673 assert_eq!(back.list_changed, Some(true));
674 }
675
676 #[test]
677 fn test_tools_capability_no_list_changed() {
678 let cap = ToolsCapability { list_changed: None };
679 let json = serde_json::to_string(&cap).unwrap();
680 let back: ToolsCapability = serde_json::from_str(&json).unwrap();
681 assert!(back.list_changed.is_none());
682 }
683
684 #[test]
687 fn test_tool_result_content_text_deserialization() {
688 let json = r#"{"type":"text","text":"hello world"}"#;
689 let content: ToolResultContent = serde_json::from_str(json).unwrap();
690 assert_eq!(
691 content,
692 ToolResultContent::Text {
693 text: "hello world".to_string()
694 }
695 );
696 }
697
698 #[test]
699 fn test_tool_result_content_image_deserialization() {
700 let json = r#"{"type":"image","data":"abc123","mimeType":"image/png"}"#;
701 let content: ToolResultContent = serde_json::from_str(json).unwrap();
702 assert_eq!(
703 content,
704 ToolResultContent::Image {
705 data: "abc123".to_string(),
706 mime_type: "image/png".to_string(),
707 }
708 );
709 }
710
711 #[test]
712 fn test_tool_result_content_resource_deserialization() {
713 let json = r#"{"type":"resource","resource":{"uri":"file:///tmp/x","text":"data"}}"#;
716 let content: ToolResultContent = serde_json::from_str(json).unwrap();
717 assert_eq!(
718 content,
719 ToolResultContent::Resource {
720 resource: EmbeddedResource {
721 uri: "file:///tmp/x".to_string(),
722 text: Some("data".to_string()),
723 blob: None,
724 mime_type: None,
725 },
726 }
727 );
728 }
729
730 #[test]
739 fn test_tool_result_json_roundtrip() {
740 let result = ToolResult {
741 content: vec![
742 ToolResultContent::Text {
743 text: "line1".to_string(),
744 },
745 ToolResultContent::Image {
746 data: "abc".to_string(),
747 mime_type: "image/png".to_string(),
748 },
749 ToolResultContent::Resource {
750 resource: EmbeddedResource {
751 uri: "file:///x".to_string(),
752 text: Some("data".to_string()),
753 blob: None,
754 mime_type: None,
755 },
756 },
757 ],
758 structured_content: None,
759 is_error: false,
760 };
761 let json = serde_json::to_string(&result).unwrap();
762 let back: ToolResult = serde_json::from_str(&json).unwrap();
763 assert_eq!(back.content.len(), 3);
764 assert!(!back.is_error);
765 }
766
767 #[test]
772 fn test_server_capabilities_empty_tools_object() {
773 let json = r#"{"tools":{}}"#;
774 let caps: ServerCapabilities = serde_json::from_str(json).unwrap();
775 let tools = caps.tools.unwrap();
776 assert!(tools.list_changed.is_none());
777 }
778
779 #[test]
782 fn test_tool_result_content_text_empty() {
783 let content = ToolResultContent::Text {
784 text: "".to_string(),
785 };
786 let json = serde_json::to_string(&content).unwrap();
787 let back: ToolResultContent = serde_json::from_str(&json).unwrap();
788 assert_eq!(
789 back,
790 ToolResultContent::Text {
791 text: "".to_string()
792 }
793 );
794 }
795
796 #[test]
797 fn test_tool_result_content_image_empty_data() {
798 let content = ToolResultContent::Image {
799 data: "".to_string(),
800 mime_type: "".to_string(),
801 };
802 let json = serde_json::to_string(&content).unwrap();
803 assert!(json.contains("\"type\":\"image\""));
804 }
805
806 #[test]
807 fn test_tool_result_content_resource_empty_uri() {
808 let content = ToolResultContent::Resource {
809 resource: EmbeddedResource {
810 uri: "".to_string(),
811 text: None,
812 blob: None,
813 mime_type: None,
814 },
815 };
816 let json = serde_json::to_string(&content).unwrap();
817 assert!(json.contains("\"type\":\"resource\""));
818 }
819
820 async fn spawn_stub_client(script: &str) -> MCPClient {
827 MCPClient::spawn("python3", &["-c", script], &HashMap::new())
828 .await
829 .expect("Failed to spawn stub MCP server")
830 }
831
832 const STUB_INIT_LIST_CALL: &str = r#"
834import sys, json
835
836def respond(id, result):
837 msg = json.dumps({"jsonrpc": "2.0", "id": id, "result": result})
838 sys.stdout.write(msg + "\n")
839 sys.stdout.flush()
840
841for line in sys.stdin:
842 line = line.strip()
843 if not line:
844 continue
845 req = json.loads(line)
846 method = req.get("method", "")
847 id_ = req.get("id")
848 if method == "initialize":
849 respond(id_, {"capabilities": {"tools": {"listChanged": True}}, "protocolVersion": "2024-11-05"})
850 elif method == "notifications/initialized":
851 pass # notification -- no response
852 elif method == "tools/list":
853 respond(id_, {"tools": [{"name": "echo", "description": "echo tool", "inputSchema": {}}]})
854 elif method == "tools/call":
855 respond(id_, {"content": [{"type": "text", "text": "hello from tool"}], "isError": False})
856 elif method == "notifications/cancelled":
857 pass
858 else:
859 respond(id_, {"error": {"code": -32601, "message": "method not found"}})
860"#;
861
862 const STUB_ERROR_SERVER: &str = r#"
864import sys, json
865
866for line in sys.stdin:
867 line = line.strip()
868 if not line:
869 continue
870 req = json.loads(line)
871 id_ = req.get("id")
872 if id_ is not None:
873 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "error": {"code": -32600, "message": "server error"}})
874 sys.stdout.write(msg + "\n")
875 sys.stdout.flush()
876"#;
877
878 const STUB_CLOSE_IMMEDIATELY: &str = r#"
880import sys
881sys.stdout.close()
882"#;
883
884 const STUB_INIT_THEN_CLOSE_STDIN: &str = r#"
901import sys, json, os
902for line in sys.stdin:
903 line = line.strip()
904 if not line:
905 continue
906 req = json.loads(line)
907 if req.get("method") == "initialize":
908 id_ = req.get("id")
909 os.close(0)
910 result = {"capabilities": {}, "protocolVersion": "2024-11-05"}
911 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": result})
912 sys.stdout.write(msg + "\n")
913 sys.stdout.flush()
914 import time; time.sleep(10)
915 break
916"#;
917
918 #[tokio::test]
919 async fn test_mcp_client_spawn_succeeds() {
920 let _guard = always_on_tracing_guard();
921 let _client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
922 }
924
925 #[tokio::test]
926 async fn test_mcp_client_connect_parses_capabilities() {
927 let _guard = always_on_tracing_guard();
928 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
929 client.connect().await.expect("connect should succeed");
930
931 let caps = client.capabilities().expect("should have capabilities");
932 assert!(caps.tools.is_some());
933 assert_eq!(caps.tools.as_ref().unwrap().list_changed, Some(true));
934 }
935
936 #[tokio::test]
937 async fn test_mcp_client_capabilities_before_connect_is_none() {
938 let client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
939 assert!(client.capabilities().is_none());
941 }
942
943 #[tokio::test]
944 async fn test_connect_fails_when_notification_write_errors() {
945 let _guard = always_on_tracing_guard();
946 let mut client = spawn_stub_client(STUB_INIT_THEN_CLOSE_STDIN).await;
950 let result = client.connect().await;
951 assert!(result.is_err());
952 }
953
954 #[tokio::test]
955 async fn test_mcp_client_cached_tools_before_list_is_empty() {
956 let client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
957 assert!(client.cached_tools().is_empty());
958 }
959
960 #[tokio::test]
961 async fn test_mcp_client_list_tools_returns_tools() {
962 let _guard = always_on_tracing_guard();
963 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
964 client.connect().await.unwrap();
965
966 let tools = client
967 .list_tools()
968 .await
969 .expect("list_tools should succeed");
970 assert_eq!(tools.len(), 1);
971 assert_eq!(tools[0].name, "echo");
972
973 assert_eq!(client.cached_tools().len(), 1);
975 assert_eq!(client.cached_tools()[0].name, "echo");
976 }
977
978 #[tokio::test]
979 async fn test_mcp_client_call_tool_returns_result() {
980 let _guard = always_on_tracing_guard();
981 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
982 client.connect().await.unwrap();
983 client.list_tools().await.unwrap();
985
986 let result = client
987 .call_tool("echo", serde_json::json!({"msg": "hi"}))
988 .await
989 .expect("call_tool should succeed");
990
991 assert_eq!(result.content.len(), 1);
992 assert_eq!(
993 result.content[0],
994 ToolResultContent::Text {
995 text: "hello from tool".to_string()
996 }
997 );
998 }
999
1000 #[tokio::test]
1001 async fn test_mcp_client_shutdown_succeeds() {
1002 let _guard = always_on_tracing_guard();
1003 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
1004 client.connect().await.unwrap();
1005 client.shutdown().await.expect("shutdown should succeed");
1007 }
1008
1009 #[tokio::test]
1010 async fn test_mcp_client_shutdown_with_dead_process() {
1011 let mut client = spawn_stub_client(STUB_CLOSE_IMMEDIATELY).await;
1012 client
1014 .shutdown()
1015 .await
1016 .expect("shutdown should be graceful");
1017 }
1018
1019 #[tokio::test]
1030 async fn test_mcp_client_server_error_propagates() {
1031 let mut client = spawn_stub_client(STUB_ERROR_SERVER).await;
1032 let err = client.connect().await;
1034 assert!(err.is_err());
1035 assert!(err.unwrap_err().to_string().contains("server error"));
1036 }
1037
1038 #[tokio::test]
1039 async fn test_mcp_client_connect_with_no_capabilities_field() {
1040 let script = r#"
1042import sys, json
1043
1044for line in sys.stdin:
1045 line = line.strip()
1046 if not line:
1047 continue
1048 req = json.loads(line)
1049 method = req.get("method", "")
1050 id_ = req.get("id")
1051 if method == "initialize":
1052 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"protocolVersion": "2024-11-05"}})
1053 sys.stdout.write(msg + "\n")
1054 sys.stdout.flush()
1055 elif method == "notifications/initialized":
1056 pass
1057 elif method == "notifications/cancelled":
1058 pass
1059"#;
1060 let mut client = spawn_stub_client(script).await;
1061 client.connect().await.expect("connect should succeed");
1063 let caps = client.capabilities().unwrap();
1064 assert!(caps.tools.is_none());
1065 }
1066
1067 #[tokio::test]
1068 async fn test_mcp_client_list_tools_missing_tools_key_returns_empty() {
1069 let script = r#"
1071import sys, json
1072
1073for line in sys.stdin:
1074 line = line.strip()
1075 if not line:
1076 continue
1077 req = json.loads(line)
1078 method = req.get("method", "")
1079 id_ = req.get("id")
1080 if method == "initialize":
1081 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"capabilities": {}}})
1082 sys.stdout.write(msg + "\n")
1083 sys.stdout.flush()
1084 elif method == "notifications/initialized":
1085 pass
1086 elif method == "tools/list":
1087 # Return result without "tools" key
1088 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {}})
1089 sys.stdout.write(msg + "\n")
1090 sys.stdout.flush()
1091 elif method == "notifications/cancelled":
1092 pass
1093"#;
1094 let mut client = spawn_stub_client(script).await;
1095 client.connect().await.unwrap();
1096 let tools = client
1097 .list_tools()
1098 .await
1099 .expect("list_tools should succeed");
1100 assert!(tools.is_empty());
1101 }
1102
1103 #[tokio::test]
1104 async fn test_mcp_client_list_tools_malformed_response_is_error() {
1105 let script = r#"
1107import sys, json
1108
1109for line in sys.stdin:
1110 line = line.strip()
1111 if not line:
1112 continue
1113 req = json.loads(line)
1114 method = req.get("method", "")
1115 id_ = req.get("id")
1116 if method == "initialize":
1117 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"capabilities": {}}})
1118 sys.stdout.write(msg + "\n")
1119 sys.stdout.flush()
1120 elif method == "notifications/initialized":
1121 pass
1122 elif method == "tools/list":
1123 # Return tools as a string instead of an array -- triggers parse error
1124 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"tools": "not_an_array"}})
1125 sys.stdout.write(msg + "\n")
1126 sys.stdout.flush()
1127 elif method == "notifications/cancelled":
1128 pass
1129"#;
1130 let mut client = spawn_stub_client(script).await;
1131 client.connect().await.unwrap();
1132 let err = client.list_tools().await;
1133 assert!(err.is_err());
1134 assert!(err.unwrap_err().to_string().contains("parse"));
1135 }
1136
1137 #[tokio::test]
1138 async fn test_mcp_client_call_tool_malformed_result_is_error() {
1139 let script = r#"
1140import sys, json
1141
1142for line in sys.stdin:
1143 line = line.strip()
1144 if not line:
1145 continue
1146 req = json.loads(line)
1147 method = req.get("method", "")
1148 id_ = req.get("id")
1149 if method == "initialize":
1150 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"capabilities": {}}})
1151 sys.stdout.write(msg + "\n")
1152 sys.stdout.flush()
1153 elif method == "notifications/initialized":
1154 pass
1155 elif method == "tools/call":
1156 # Return a result that can't be parsed as ToolResult
1157 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": "bad_tool_result"})
1158 sys.stdout.write(msg + "\n")
1159 sys.stdout.flush()
1160 elif method == "notifications/cancelled":
1161 pass
1162"#;
1163 let mut client = spawn_stub_client(script).await;
1164 client.connect().await.unwrap();
1165 let err = client.call_tool("broken", serde_json::json!({})).await;
1166 assert!(err.is_err());
1167 assert!(err.unwrap_err().to_string().contains("parse"));
1168 }
1169
1170 #[tokio::test]
1171 async fn test_mcp_client_response_with_no_result_is_error() {
1172 let script = r#"
1174import sys, json
1175
1176for line in sys.stdin:
1177 line = line.strip()
1178 if not line:
1179 continue
1180 req = json.loads(line)
1181 id_ = req.get("id")
1182 if id_ is not None:
1183 # No "result", no "error"
1184 msg = json.dumps({"jsonrpc": "2.0", "id": id_})
1185 sys.stdout.write(msg + "\n")
1186 sys.stdout.flush()
1187"#;
1188 let mut client = spawn_stub_client(script).await;
1189 let err = client.connect().await;
1190 assert!(err.is_err());
1191 assert!(err.unwrap_err().to_string().contains("no result"));
1192 }
1193
1194 #[test]
1202 fn tool_result_reads_is_error_from_camel_case_wire_name() {
1203 let json = r#"{"content":[{"type":"text","text":"boom"}],"isError":true}"#;
1208 let result: ToolResult = serde_json::from_str(json).unwrap();
1209 assert!(result.is_error, "isError:true must deserialize as an error");
1210 }
1211
1212 #[test]
1213 fn tool_result_snake_case_is_error_is_not_honored() {
1214 let json = r#"{"content":[],"is_error":true}"#;
1218 let result: ToolResult = serde_json::from_str(json).unwrap();
1219 assert!(!result.is_error);
1220 }
1221
1222 #[test]
1223 fn tool_result_content_defaults_to_empty() {
1224 let json = r#"{"structuredContent":{"ok":true}}"#;
1225 let result: ToolResult = serde_json::from_str(json).unwrap();
1226 assert!(result.content.is_empty());
1227 assert_eq!(
1228 result.structured_content,
1229 Some(serde_json::json!({"ok": true}))
1230 );
1231 }
1232
1233 #[test]
1234 fn tool_result_structured_content_absent_is_none_and_omitted() {
1235 let result = ToolResult {
1236 content: vec![],
1237 structured_content: None,
1238 is_error: false,
1239 };
1240 let json = serde_json::to_string(&result).unwrap();
1241 assert!(!json.contains("structuredContent"), "got: {json}");
1242 }
1243
1244 #[test]
1245 fn tool_result_content_image_uses_mime_type_wire_name() {
1246 let content = ToolResultContent::Image {
1247 data: "abc".to_string(),
1248 mime_type: "image/png".to_string(),
1249 };
1250 let json = serde_json::to_string(&content).unwrap();
1251 assert!(json.contains("\"mimeType\":\"image/png\""), "got: {json}");
1252 assert!(!json.contains("mime_type"), "got: {json}");
1253 }
1254
1255 #[test]
1256 fn tool_result_content_audio_roundtrip() {
1257 let json = r#"{"type":"audio","data":"YWJj","mimeType":"audio/wav"}"#;
1258 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1259 assert_eq!(
1260 content,
1261 ToolResultContent::Audio {
1262 data: "YWJj".to_string(),
1263 mime_type: "audio/wav".to_string(),
1264 }
1265 );
1266 let back: ToolResultContent =
1267 serde_json::from_str(&serde_json::to_string(&content).unwrap()).unwrap();
1268 assert_eq!(back, content);
1269 }
1270
1271 #[test]
1272 fn tool_result_content_resource_link_full() {
1273 let json = r#"{"type":"resource_link","uri":"file:///m.rs","name":"m.rs",
1274 "description":"entry point","mimeType":"text/x-rust"}"#;
1275 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1276 assert_eq!(
1277 content,
1278 ToolResultContent::ResourceLink {
1279 uri: "file:///m.rs".to_string(),
1280 name: "m.rs".to_string(),
1281 description: Some("entry point".to_string()),
1282 mime_type: Some("text/x-rust".to_string()),
1283 }
1284 );
1285 }
1286
1287 #[test]
1288 fn tool_result_content_resource_link_minimal_omits_optionals() {
1289 let content: ToolResultContent =
1290 serde_json::from_str(r#"{"type":"resource_link","uri":"file:///x"}"#).unwrap();
1291 assert_eq!(
1292 content,
1293 ToolResultContent::ResourceLink {
1294 uri: "file:///x".to_string(),
1295 name: String::new(),
1296 description: None,
1297 mime_type: None,
1298 }
1299 );
1300 let json = serde_json::to_string(&content).unwrap();
1301 assert!(!json.contains("description"), "got: {json}");
1302 assert!(!json.contains("mimeType"), "got: {json}");
1303 }
1304
1305 #[test]
1306 fn tool_result_content_embedded_resource_is_nested() {
1307 let json = r#"{"type":"resource","resource":{"uri":"file:///m.rs",
1308 "mimeType":"text/x-rust","text":"fn main() {}"}}"#;
1309 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1310 assert_eq!(
1311 content,
1312 ToolResultContent::Resource {
1313 resource: EmbeddedResource {
1314 uri: "file:///m.rs".to_string(),
1315 text: Some("fn main() {}".to_string()),
1316 blob: None,
1317 mime_type: Some("text/x-rust".to_string()),
1318 }
1319 }
1320 );
1321 }
1322
1323 #[test]
1324 fn tool_result_content_embedded_resource_binary_blob() {
1325 let json = r#"{"type":"resource","resource":{"uri":"file:///a.png","blob":"YWJj"}}"#;
1326 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1327 assert_eq!(
1328 content,
1329 ToolResultContent::Resource {
1330 resource: EmbeddedResource {
1331 uri: "file:///a.png".to_string(),
1332 text: None,
1333 blob: Some("YWJj".to_string()),
1334 mime_type: None,
1335 }
1336 }
1337 );
1338 let json = serde_json::to_string(&content).unwrap();
1340 assert!(!json.contains("text"), "got: {json}");
1341 assert!(!json.contains("mimeType"), "got: {json}");
1342 }
1343
1344 #[test]
1345 fn unknown_content_type_degrades_instead_of_failing_the_result() {
1346 let json = r#"{"content":[
1350 {"type":"text","text":"keep me"},
1351 {"type":"hologram","payload":{"deeply":["nested"]}}
1352 ],"isError":false}"#;
1353 let result: ToolResult = serde_json::from_str(json).unwrap();
1354 assert_eq!(result.content.len(), 2);
1355 assert_eq!(
1356 result.content[0],
1357 ToolResultContent::Text {
1358 text: "keep me".to_string()
1359 }
1360 );
1361 assert_eq!(result.content[1], ToolResultContent::Unknown);
1362 }
1363
1364 #[test]
1365 fn unknown_content_serializes_without_panicking() {
1366 let json = serde_json::to_string(&ToolResultContent::Unknown).unwrap();
1370 assert!(json.contains("Unknown"), "got: {json}");
1371 }
1372
1373 fn paginated_stub(pages: usize) -> String {
1378 format!(
1379 r#"
1380import sys, json
1381PAGES = {pages}
1382for line in sys.stdin:
1383 line = line.strip()
1384 if not line:
1385 continue
1386 req = json.loads(line)
1387 method, id_ = req.get("method", ""), req.get("id")
1388 if method == "initialize":
1389 res = {{"capabilities": {{}}, "protocolVersion": "2024-11-05"}}
1390 elif method == "tools/list":
1391 cursor = int((req.get("params") or {{}}).get("cursor", "0"))
1392 res = {{"tools": [{{"name": "tool%d" % cursor, "inputSchema": {{}}}}]}}
1393 if cursor + 1 < PAGES:
1394 res["nextCursor"] = str(cursor + 1)
1395 else:
1396 continue
1397 sys.stdout.write(json.dumps({{"jsonrpc": "2.0", "id": id_, "result": res}}) + "\n")
1398 sys.stdout.flush()
1399"#
1400 )
1401 }
1402
1403 #[tokio::test]
1404 async fn list_tools_follows_next_cursor_across_pages() {
1405 let _guard = always_on_tracing_guard();
1406 let script = paginated_stub(3);
1407 let mut client = spawn_stub_client(&script).await;
1408 client.connect().await.unwrap();
1409
1410 let tools = client
1411 .list_tools()
1412 .await
1413 .expect("list_tools should succeed");
1414 let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
1416 assert_eq!(names, vec!["tool0", "tool1", "tool2"]);
1417 assert_eq!(client.cached_tools().len(), 3);
1418 }
1419
1420 #[tokio::test]
1421 async fn list_tools_stops_at_the_page_limit() {
1422 let _guard = always_on_tracing_guard();
1423 let script = paginated_stub(MAX_TOOL_PAGES + 10);
1425 let mut client = spawn_stub_client(&script).await;
1426 client.connect().await.unwrap();
1427
1428 let tools = client
1429 .list_tools()
1430 .await
1431 .expect("list_tools should succeed");
1432 assert_eq!(tools.len(), MAX_TOOL_PAGES);
1433 }
1434
1435 #[tokio::test]
1436 async fn transport_failure_propagates_out_of_a_request() {
1437 let _guard = always_on_tracing_guard();
1438 let mut client = spawn_stub_client("import sys\nsys.exit(0)\n").await;
1442 assert!(client.connect().await.is_err());
1443 }
1444
1445 #[test]
1448 fn preferred_version_is_the_newest_supported() {
1449 assert_eq!(PREFERRED_PROTOCOL_VERSION, SUPPORTED_PROTOCOL_VERSIONS[0]);
1450 let mut sorted = SUPPORTED_PROTOCOL_VERSIONS.to_vec();
1452 sorted.sort_unstable_by(|a, b| b.cmp(a));
1453 assert_eq!(sorted, SUPPORTED_PROTOCOL_VERSIONS);
1454 }
1455
1456 #[test]
1457 fn negotiation_adopts_a_recognized_echo() {
1458 let _guard = always_on_tracing_guard();
1459 let echoed = serde_json::json!("2025-06-18");
1460 assert_eq!(negotiated_version(Some(&echoed)), "2025-06-18");
1461 }
1462
1463 #[test]
1464 fn negotiation_honors_an_unrecognized_echo() {
1465 let _guard = always_on_tracing_guard();
1466 let echoed = serde_json::json!("2099-01-01");
1469 assert_eq!(negotiated_version(Some(&echoed)), "2099-01-01");
1470 }
1471
1472 #[test]
1473 fn negotiation_falls_back_when_the_server_omits_the_field() {
1474 let _guard = always_on_tracing_guard();
1475 assert_eq!(negotiated_version(None), PREFERRED_PROTOCOL_VERSION);
1476 }
1477
1478 #[test]
1479 fn negotiation_falls_back_when_the_echo_is_not_a_string() {
1480 let _guard = always_on_tracing_guard();
1481 let echoed = serde_json::json!(20251125);
1482 assert_eq!(
1483 negotiated_version(Some(&echoed)),
1484 PREFERRED_PROTOCOL_VERSION
1485 );
1486 }
1487
1488 #[tokio::test]
1489 async fn connect_offers_the_preferred_version_and_records_the_echo() {
1490 let _guard = always_on_tracing_guard();
1491 let script = r#"
1494import sys, json
1495for line in sys.stdin:
1496 line = line.strip()
1497 if not line:
1498 continue
1499 req = json.loads(line)
1500 if req.get("method") == "initialize":
1501 offered = req["params"]["protocolVersion"]
1502 assert offered == "2025-11-25", offered
1503 sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": req.get("id"),
1504 "result": {"capabilities": {}, "protocolVersion": "2025-03-26"}}) + "\n")
1505 sys.stdout.flush()
1506"#;
1507 let mut client = spawn_stub_client(script).await;
1508 client.connect().await.expect("connect should succeed");
1509 assert_eq!(client.protocol_version(), Some("2025-03-26"));
1510 }
1511
1512 #[tokio::test]
1513 async fn protocol_version_is_none_before_connect() {
1514 let client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
1515 assert!(client.protocol_version().is_none());
1516 }
1517
1518 #[test]
1521 fn connect_http_builds_a_client_without_touching_the_network() {
1522 assert!(MCPClient::connect_http("http://127.0.0.1:1/mcp", &HashMap::new(), &[]).is_ok());
1525 }
1526
1527 struct NoopRefresher;
1528 #[async_trait::async_trait]
1529 impl crate::transport::BearerRefresher for NoopRefresher {
1530 async fn refresh(&self) -> anyhow::Result<String> {
1531 Ok("Bearer x".to_string())
1532 }
1533 }
1534
1535 #[tokio::test]
1536 async fn set_refresher_on_stdio_is_a_noop() {
1537 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
1540 client.set_refresher(std::sync::Arc::new(NoopRefresher));
1541 }
1542
1543 #[tokio::test]
1544 async fn set_refresher_on_http_is_accepted() {
1545 let mut client =
1546 MCPClient::connect_http("http://127.0.0.1:1/mcp", &HashMap::new(), &[]).unwrap();
1547 use crate::transport::BearerRefresher as _;
1549 assert_eq!(NoopRefresher.refresh().await.unwrap(), "Bearer x");
1550 client.set_refresher(std::sync::Arc::new(NoopRefresher));
1551 }
1552
1553 #[test]
1554 fn connect_http_rejects_an_unparseable_url() {
1555 assert!(MCPClient::connect_http("not a url", &HashMap::new(), &[]).is_err());
1556 }
1557
1558 #[tokio::test]
1559 async fn from_config_builds_the_stdio_transport() {
1560 let _guard = always_on_tracing_guard();
1561 let config = MCPServerConfig::stdio(
1562 "s",
1563 "python3",
1564 vec!["-c".into(), STUB_INIT_LIST_CALL.into()],
1565 );
1566 let mut client = MCPClient::from_config(&config)
1567 .await
1568 .expect("stdio config should connect");
1569 client.connect().await.expect("handshake should succeed");
1570 assert_eq!(client.list_tools().await.unwrap().len(), 1);
1571 }
1572
1573 #[tokio::test]
1574 async fn from_config_builds_the_http_transport() {
1575 let config = MCPServerConfig::http("s", "http://127.0.0.1:1/mcp");
1576 assert!(MCPClient::from_config(&config).await.is_ok());
1577 }
1578
1579 #[tokio::test]
1580 async fn from_config_with_auth_injects_a_bearer_for_http() {
1581 let config = MCPServerConfig::http("s", "http://127.0.0.1:1/mcp");
1584 let header = Some(("Authorization".to_string(), "Bearer tok".to_string()));
1585 assert!(
1586 MCPClient::from_config_with_auth(&config, header, &[])
1587 .await
1588 .is_ok()
1589 );
1590 }
1591
1592 #[tokio::test]
1593 async fn from_config_rejects_an_unresolvable_entry() {
1594 let config = MCPServerConfig {
1595 name: "broken".to_string(),
1596 ..Default::default()
1597 };
1598 let err = MCPClient::from_config(&config)
1599 .await
1600 .err()
1601 .expect("an entry with neither command nor url cannot connect");
1602 assert!(err.to_string().contains("either a `command`"), "got: {err}");
1603 }
1604}