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 {
95 text: String,
97 },
98 #[serde(rename = "image")]
100 Image {
101 data: String,
103 #[serde(rename = "mimeType")]
105 mime_type: String,
106 },
107 #[serde(rename = "audio")]
109 Audio {
110 data: String,
112 #[serde(rename = "mimeType")]
114 mime_type: String,
115 },
116 #[serde(rename = "resource_link")]
118 ResourceLink {
119 uri: String,
121 #[serde(default)]
123 name: String,
124 #[serde(default, skip_serializing_if = "Option::is_none")]
126 description: Option<String>,
127 #[serde(rename = "mimeType", default, skip_serializing_if = "Option::is_none")]
129 mime_type: Option<String>,
130 },
131 #[serde(rename = "resource")]
133 Resource {
134 resource: EmbeddedResource,
136 },
137 #[serde(other)]
145 Unknown,
146}
147pub const SUPPORTED_PROTOCOL_VERSIONS: &[&str] =
155 &["2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05"];
156
157pub const PREFERRED_PROTOCOL_VERSION: &str = SUPPORTED_PROTOCOL_VERSIONS[0];
159
160pub struct MCPClient {
162 transport: Box<dyn Transport>,
164 next_id: AtomicU64,
166 request_timeout: Duration,
168 capabilities: Option<ServerCapabilities>,
170 protocol_version: Option<String>,
172 cached_tools: Vec<ToolMetadata>,
174}
175
176impl MCPClient {
177 pub(crate) fn new(transport: Box<dyn Transport>) -> Self {
179 Self {
180 transport,
181 next_id: AtomicU64::new(1),
182 request_timeout: DEFAULT_REQUEST_TIMEOUT,
183 capabilities: None,
184 protocol_version: None,
185 cached_tools: Vec::new(),
186 }
187 }
188
189 pub async fn spawn(
191 command: &str,
192 args: &[&str],
193 env: &HashMap<String, String>,
194 ) -> anyhow::Result<Self> {
195 let transport = StdioTransport::spawn(command, args, env).await?;
196 Ok(Self::new(Box::new(transport)))
197 }
198
199 pub fn connect_http(
201 url: &str,
202 headers: &HashMap<String, String>,
203 allow_env: &[String],
204 ) -> anyhow::Result<Self> {
205 let transport = HttpTransport::new(url, headers, allow_env)?;
206 Ok(Self::new(Box::new(transport)))
207 }
208
209 pub fn set_refresher(
215 &mut self,
216 refresher: std::sync::Arc<dyn crate::transport::BearerRefresher>,
217 ) {
218 self.transport.set_bearer_refresher(refresher);
219 }
220
221 pub async fn from_config(config: &MCPServerConfig) -> anyhow::Result<Self> {
227 Self::from_config_with_auth(config, None, &[]).await
228 }
229
230 pub async fn from_config_with_auth(
239 config: &MCPServerConfig,
240 auth_header: Option<(String, String)>,
241 allow_env: &[String],
242 ) -> anyhow::Result<Self> {
243 match config.resolve()? {
244 ResolvedTransport::Stdio { command, args, env } => {
245 let args: Vec<&str> = args.iter().map(String::as_str).collect();
246 Self::spawn(command, &args, env).await
247 }
248 ResolvedTransport::Http { url, headers } => {
249 let mut headers = headers.clone();
250 if let Some((name, value)) = auth_header {
251 headers.insert(name, value);
252 }
253 Self::connect_http(url, &headers, allow_env)
254 }
255 }
256 }
257
258 pub async fn connect(&mut self) -> anyhow::Result<()> {
260 tracing::info!("Initializing MCP connection");
261
262 let init_params = serde_json::json!({
263 "protocolVersion": PREFERRED_PROTOCOL_VERSION,
264 "capabilities": {},
265 "clientInfo": {
266 "name": "leviath",
267 "version": env!("CARGO_PKG_VERSION")
268 }
269 });
270
271 let result = self
275 .request_with_timeout("initialize", init_params, DEFAULT_CONNECT_TIMEOUT)
276 .await?;
277
278 let capabilities: ServerCapabilities = if let Some(caps) = result.get("capabilities") {
280 serde_json::from_value(caps.clone()).unwrap_or_default()
281 } else {
282 ServerCapabilities::default()
283 };
284 self.capabilities = Some(capabilities);
285 let version = negotiated_version(result.get("protocolVersion"));
286 self.protocol_version = Some(version.clone());
287
288 self.send_notification("notifications/initialized", serde_json::json!({}))
290 .await?;
291
292 tracing::info!(version = %version, "MCP connection established");
293 Ok(())
294 }
295
296 pub async fn list_tools(&mut self) -> anyhow::Result<Vec<ToolMetadata>> {
304 tracing::debug!("Listing MCP tools");
305
306 let mut tools: Vec<ToolMetadata> = Vec::new();
307 let mut cursor: Option<String> = None;
308
309 for page in 0..MAX_TOOL_PAGES {
310 let params = match &cursor {
311 Some(c) => serde_json::json!({ "cursor": c }),
312 None => serde_json::json!({}),
313 };
314 let result = self.send_request("tools/list", params).await?;
315
316 let tools_value = result.get("tools").cloned().unwrap_or(Value::Array(vec![]));
317 let page_tools: Vec<ToolMetadata> = serde_json::from_value(tools_value)
318 .map_err(|e| anyhow::anyhow!("Failed to parse tools list: {}", e))?;
319 tools.extend(page_tools);
320
321 cursor = result
322 .get("nextCursor")
323 .and_then(Value::as_str)
324 .map(str::to_string);
325 if cursor.is_none() {
326 break;
327 }
328 if page + 1 == MAX_TOOL_PAGES {
329 tracing::warn!(
330 pages = MAX_TOOL_PAGES,
331 "MCP server still returned a tools/list cursor at the page \
332 limit - stopping; some tools may be missing"
333 );
334 }
335 }
336
337 self.cached_tools = tools.clone();
338 let count = tools.len();
342 tracing::debug!(count, "Discovered MCP tools");
343 Ok(tools)
344 }
345
346 pub async fn call_tool(&mut self, name: &str, arguments: Value) -> anyhow::Result<ToolResult> {
348 tracing::debug!(tool = %name, "Calling MCP tool");
349
350 let params = serde_json::json!({
351 "name": name,
352 "arguments": arguments,
353 });
354
355 let result = self.send_request("tools/call", params).await?;
356
357 let tool_result: ToolResult = serde_json::from_value(result)
358 .map_err(|e| anyhow::anyhow!("Failed to parse tool result: {}", e))?;
359
360 Ok(tool_result)
361 }
362
363 pub async fn shutdown(&mut self) -> anyhow::Result<()> {
367 tracing::info!("Shutting down MCP server");
368 let _ = self.transport.close().await;
369 Ok(())
370 }
371
372 pub fn capabilities(&self) -> Option<&ServerCapabilities> {
374 self.capabilities.as_ref()
375 }
376
377 pub fn protocol_version(&self) -> Option<&str> {
379 self.protocol_version.as_deref()
380 }
381
382 pub fn cached_tools(&self) -> &[ToolMetadata] {
384 &self.cached_tools
385 }
386
387 async fn send_request(&mut self, method: &str, params: Value) -> anyhow::Result<Value> {
389 self.request_with_timeout(method, params, self.request_timeout)
390 .await
391 }
392
393 async fn request_with_timeout(
395 &mut self,
396 method: &str,
397 params: Value,
398 timeout: Duration,
399 ) -> anyhow::Result<Value> {
400 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
401 let request = JsonRpcRequest::request(id, method, params);
402 self.transport
403 .send_request(&request, timeout)
404 .await?
405 .into_result()
406 }
407
408 async fn send_notification(&mut self, method: &str, params: Value) -> anyhow::Result<()> {
410 let request = JsonRpcRequest::notification(method, params);
411 self.transport.send_notification(&request).await
412 }
413}
414
415fn negotiated_version(echoed: Option<&Value>) -> String {
423 match echoed.and_then(Value::as_str) {
424 Some(version) => {
425 if !SUPPORTED_PROTOCOL_VERSIONS.contains(&version) {
426 tracing::warn!(
427 version = %version,
428 "MCP server negotiated an unrecognized protocol revision - continuing"
429 );
430 }
431 version.to_string()
432 }
433 None => {
434 tracing::debug!("MCP server echoed no protocolVersion - assuming the offered one");
436 PREFERRED_PROTOCOL_VERSION.to_string()
437 }
438 }
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444 use crate::test_support::always_on_tracing_guard;
445
446 #[test]
449 fn test_server_capabilities_default() {
450 let caps = ServerCapabilities::default();
451 assert!(caps.tools.is_none());
452 }
453
454 #[test]
455 fn test_tools_capability_default() {
456 let cap = ToolsCapability::default();
457 assert!(cap.list_changed.is_none());
458 }
459
460 #[test]
461 fn test_server_capabilities_serialization() {
462 let caps = ServerCapabilities {
463 tools: Some(ToolsCapability {
464 list_changed: Some(true),
465 }),
466 };
467 let json = serde_json::to_string(&caps).unwrap();
468 assert!(json.contains("listChanged"));
469 assert!(json.contains("true"));
470
471 let deserialized: ServerCapabilities = serde_json::from_str(&json).unwrap();
472 assert!(deserialized.tools.unwrap().list_changed.unwrap());
473 }
474
475 #[test]
476 fn test_tool_result_serialization() {
477 let result = ToolResult {
478 content: vec![ToolResultContent::Text {
479 text: "Hello".to_string(),
480 }],
481 structured_content: None,
482 is_error: false,
483 };
484 let json = serde_json::to_string(&result).unwrap();
485 assert!(json.contains("Hello"));
486 assert!(json.contains("\"isError\":false"), "got: {json}");
489 assert!(!json.contains("is_error"), "got: {json}");
490 }
491
492 #[test]
493 fn test_tool_result_with_error() {
494 let result = ToolResult {
495 content: vec![ToolResultContent::Text {
496 text: "Something went wrong".to_string(),
497 }],
498 structured_content: None,
499 is_error: true,
500 };
501 assert!(result.is_error);
502 assert_eq!(result.content.len(), 1);
503 }
504
505 #[test]
506 fn test_tool_result_content_text() {
507 let content = ToolResultContent::Text {
508 text: "result text".to_string(),
509 };
510 let json = serde_json::to_string(&content).unwrap();
511 assert!(json.contains("result text"));
512 assert!(json.contains("\"type\":\"text\""));
513 }
514
515 #[test]
516 fn test_tool_result_content_image() {
517 let content = ToolResultContent::Image {
518 data: "base64data".to_string(),
519 mime_type: "image/png".to_string(),
520 };
521 let json = serde_json::to_string(&content).unwrap();
522 assert!(json.contains("base64data"));
523 assert!(json.contains("image/png"));
524 }
525
526 #[test]
527 fn test_tool_result_content_resource() {
528 let content = ToolResultContent::Resource {
529 resource: EmbeddedResource {
530 uri: "file:///tmp/test.txt".to_string(),
531 text: Some("file contents".to_string()),
532 blob: None,
533 mime_type: None,
534 },
535 };
536 let json = serde_json::to_string(&content).unwrap();
537 assert!(json.contains("file:///tmp/test.txt"));
538 assert!(json.contains("file contents"));
539 }
540
541 #[test]
542 fn test_tool_result_content_resource_no_text() {
543 let content = ToolResultContent::Resource {
544 resource: EmbeddedResource {
545 uri: "file:///tmp/test.txt".to_string(),
546 text: None,
547 blob: None,
548 mime_type: None,
549 },
550 };
551 let json = serde_json::to_string(&content).unwrap();
552 assert!(json.contains("file:///tmp/test.txt"));
553 }
554
555 #[test]
556 fn test_tool_result_deserialization() {
557 let json = r#"{"content":[{"type":"text","text":"Hello"}],"is_error":false}"#;
558 let result: ToolResult = serde_json::from_str(json).unwrap();
559 assert!(!result.is_error);
560 assert_eq!(result.content.len(), 1);
561 }
562
563 #[test]
564 fn test_tool_result_deserialization_missing_is_error() {
565 let json = r#"{"content":[{"type":"text","text":"Hello"}]}"#;
566 let result: ToolResult = serde_json::from_str(json).unwrap();
567 assert!(!result.is_error); }
569
570 #[test]
571 fn test_tool_result_multiple_content() {
572 let result = ToolResult {
573 content: vec![
574 ToolResultContent::Text {
575 text: "line 1".to_string(),
576 },
577 ToolResultContent::Text {
578 text: "line 2".to_string(),
579 },
580 ],
581 structured_content: None,
582 is_error: false,
583 };
584 assert_eq!(result.content.len(), 2);
585 }
586
587 #[test]
588 fn test_tool_result_clone() {
589 let result = ToolResult {
590 content: vec![ToolResultContent::Text {
591 text: "test".to_string(),
592 }],
593 structured_content: None,
594 is_error: true,
595 };
596 let cloned = result.clone();
597 assert!(cloned.is_error);
598 assert_eq!(cloned.content.len(), 1);
599 }
600
601 #[test]
602 fn test_server_capabilities_clone() {
603 let caps = ServerCapabilities {
604 tools: Some(ToolsCapability {
605 list_changed: Some(true),
606 }),
607 };
608 let cloned = caps.clone();
609 assert!(cloned.tools.unwrap().list_changed.unwrap());
610 }
611
612 #[test]
615 fn test_tool_result_empty_content() {
616 let result = ToolResult {
617 content: vec![],
618 structured_content: None,
619 is_error: false,
620 };
621 assert!(result.content.is_empty());
622 assert!(!result.is_error);
623 }
624
625 #[test]
626 fn test_tool_result_mixed_content_types() {
627 let result = ToolResult {
628 content: vec![
629 ToolResultContent::Text {
630 text: "hello".to_string(),
631 },
632 ToolResultContent::Image {
633 data: "base64".to_string(),
634 mime_type: "image/jpeg".to_string(),
635 },
636 ToolResultContent::Resource {
637 resource: EmbeddedResource {
638 uri: "file:///test".to_string(),
639 text: Some("content".to_string()),
640 blob: None,
641 mime_type: None,
642 },
643 },
644 ],
645 structured_content: None,
646 is_error: false,
647 };
648 assert_eq!(result.content.len(), 3);
649 let json = serde_json::to_string(&result).unwrap();
650 let back: ToolResult = serde_json::from_str(&json).unwrap();
651 assert_eq!(back.content.len(), 3);
652 }
653
654 #[test]
657 fn test_server_capabilities_from_empty_json() {
658 let caps: ServerCapabilities = serde_json::from_str("{}").unwrap();
659 assert!(caps.tools.is_none());
660 }
661
662 #[test]
663 fn test_server_capabilities_with_tools() {
664 let json = r#"{"tools":{"listChanged":false}}"#;
665 let caps: ServerCapabilities = serde_json::from_str(json).unwrap();
666 let tools = caps.tools.unwrap();
667 assert_eq!(tools.list_changed, Some(false));
668 }
669
670 #[test]
671 fn test_server_capabilities_with_null_tools() {
672 let json = r#"{"tools":null}"#;
673 let caps: ServerCapabilities = serde_json::from_str(json).unwrap();
674 assert!(caps.tools.is_none());
675 }
676
677 #[test]
680 fn test_tools_capability_with_list_changed_true() {
681 let cap = ToolsCapability {
682 list_changed: Some(true),
683 };
684 let json = serde_json::to_string(&cap).unwrap();
685 assert!(json.contains("true"));
686 let back: ToolsCapability = serde_json::from_str(&json).unwrap();
687 assert_eq!(back.list_changed, Some(true));
688 }
689
690 #[test]
691 fn test_tools_capability_no_list_changed() {
692 let cap = ToolsCapability { list_changed: None };
693 let json = serde_json::to_string(&cap).unwrap();
694 let back: ToolsCapability = serde_json::from_str(&json).unwrap();
695 assert!(back.list_changed.is_none());
696 }
697
698 #[test]
701 fn test_tool_result_content_text_deserialization() {
702 let json = r#"{"type":"text","text":"hello world"}"#;
703 let content: ToolResultContent = serde_json::from_str(json).unwrap();
704 assert_eq!(
705 content,
706 ToolResultContent::Text {
707 text: "hello world".to_string()
708 }
709 );
710 }
711
712 #[test]
713 fn test_tool_result_content_image_deserialization() {
714 let json = r#"{"type":"image","data":"abc123","mimeType":"image/png"}"#;
715 let content: ToolResultContent = serde_json::from_str(json).unwrap();
716 assert_eq!(
717 content,
718 ToolResultContent::Image {
719 data: "abc123".to_string(),
720 mime_type: "image/png".to_string(),
721 }
722 );
723 }
724
725 #[test]
726 fn test_tool_result_content_resource_deserialization() {
727 let json = r#"{"type":"resource","resource":{"uri":"file:///tmp/x","text":"data"}}"#;
730 let content: ToolResultContent = serde_json::from_str(json).unwrap();
731 assert_eq!(
732 content,
733 ToolResultContent::Resource {
734 resource: EmbeddedResource {
735 uri: "file:///tmp/x".to_string(),
736 text: Some("data".to_string()),
737 blob: None,
738 mime_type: None,
739 },
740 }
741 );
742 }
743
744 #[test]
753 fn test_tool_result_json_roundtrip() {
754 let result = ToolResult {
755 content: vec![
756 ToolResultContent::Text {
757 text: "line1".to_string(),
758 },
759 ToolResultContent::Image {
760 data: "abc".to_string(),
761 mime_type: "image/png".to_string(),
762 },
763 ToolResultContent::Resource {
764 resource: EmbeddedResource {
765 uri: "file:///x".to_string(),
766 text: Some("data".to_string()),
767 blob: None,
768 mime_type: None,
769 },
770 },
771 ],
772 structured_content: None,
773 is_error: false,
774 };
775 let json = serde_json::to_string(&result).unwrap();
776 let back: ToolResult = serde_json::from_str(&json).unwrap();
777 assert_eq!(back.content.len(), 3);
778 assert!(!back.is_error);
779 }
780
781 #[test]
786 fn test_server_capabilities_empty_tools_object() {
787 let json = r#"{"tools":{}}"#;
788 let caps: ServerCapabilities = serde_json::from_str(json).unwrap();
789 let tools = caps.tools.unwrap();
790 assert!(tools.list_changed.is_none());
791 }
792
793 #[test]
796 fn test_tool_result_content_text_empty() {
797 let content = ToolResultContent::Text {
798 text: "".to_string(),
799 };
800 let json = serde_json::to_string(&content).unwrap();
801 let back: ToolResultContent = serde_json::from_str(&json).unwrap();
802 assert_eq!(
803 back,
804 ToolResultContent::Text {
805 text: "".to_string()
806 }
807 );
808 }
809
810 #[test]
811 fn test_tool_result_content_image_empty_data() {
812 let content = ToolResultContent::Image {
813 data: "".to_string(),
814 mime_type: "".to_string(),
815 };
816 let json = serde_json::to_string(&content).unwrap();
817 assert!(json.contains("\"type\":\"image\""));
818 }
819
820 #[test]
821 fn test_tool_result_content_resource_empty_uri() {
822 let content = ToolResultContent::Resource {
823 resource: EmbeddedResource {
824 uri: "".to_string(),
825 text: None,
826 blob: None,
827 mime_type: None,
828 },
829 };
830 let json = serde_json::to_string(&content).unwrap();
831 assert!(json.contains("\"type\":\"resource\""));
832 }
833
834 async fn spawn_stub_client(script: &str) -> MCPClient {
841 MCPClient::spawn("python3", &["-c", script], &HashMap::new())
842 .await
843 .expect("Failed to spawn stub MCP server")
844 }
845
846 const STUB_INIT_LIST_CALL: &str = r#"
848import sys, json
849
850def respond(id, result):
851 msg = json.dumps({"jsonrpc": "2.0", "id": id, "result": result})
852 sys.stdout.write(msg + "\n")
853 sys.stdout.flush()
854
855for line in sys.stdin:
856 line = line.strip()
857 if not line:
858 continue
859 req = json.loads(line)
860 method = req.get("method", "")
861 id_ = req.get("id")
862 if method == "initialize":
863 respond(id_, {"capabilities": {"tools": {"listChanged": True}}, "protocolVersion": "2024-11-05"})
864 elif method == "notifications/initialized":
865 pass # notification -- no response
866 elif method == "tools/list":
867 respond(id_, {"tools": [{"name": "echo", "description": "echo tool", "inputSchema": {}}]})
868 elif method == "tools/call":
869 respond(id_, {"content": [{"type": "text", "text": "hello from tool"}], "isError": False})
870 elif method == "notifications/cancelled":
871 pass
872 else:
873 respond(id_, {"error": {"code": -32601, "message": "method not found"}})
874"#;
875
876 const STUB_ERROR_SERVER: &str = r#"
878import sys, json
879
880for line in sys.stdin:
881 line = line.strip()
882 if not line:
883 continue
884 req = json.loads(line)
885 id_ = req.get("id")
886 if id_ is not None:
887 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "error": {"code": -32600, "message": "server error"}})
888 sys.stdout.write(msg + "\n")
889 sys.stdout.flush()
890"#;
891
892 const STUB_CLOSE_IMMEDIATELY: &str = r#"
894import sys
895sys.stdout.close()
896"#;
897
898 const STUB_INIT_THEN_CLOSE_STDIN: &str = r#"
915import sys, json, os
916for line in sys.stdin:
917 line = line.strip()
918 if not line:
919 continue
920 req = json.loads(line)
921 if req.get("method") == "initialize":
922 id_ = req.get("id")
923 os.close(0)
924 result = {"capabilities": {}, "protocolVersion": "2024-11-05"}
925 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": result})
926 sys.stdout.write(msg + "\n")
927 sys.stdout.flush()
928 import time; time.sleep(10)
929 break
930"#;
931
932 #[tokio::test]
933 async fn test_mcp_client_spawn_succeeds() {
934 let _guard = always_on_tracing_guard();
935 let _client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
936 }
938
939 #[tokio::test]
940 async fn test_mcp_client_connect_parses_capabilities() {
941 let _guard = always_on_tracing_guard();
942 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
943 client.connect().await.expect("connect should succeed");
944
945 let caps = client.capabilities().expect("should have capabilities");
946 assert!(caps.tools.is_some());
947 assert_eq!(caps.tools.as_ref().unwrap().list_changed, Some(true));
948 }
949
950 #[tokio::test]
951 async fn test_mcp_client_capabilities_before_connect_is_none() {
952 let client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
953 assert!(client.capabilities().is_none());
955 }
956
957 #[tokio::test]
958 async fn test_connect_fails_when_notification_write_errors() {
959 let _guard = always_on_tracing_guard();
960 let mut client = spawn_stub_client(STUB_INIT_THEN_CLOSE_STDIN).await;
964 let result = client.connect().await;
965 assert!(result.is_err());
966 }
967
968 #[tokio::test]
969 async fn test_mcp_client_cached_tools_before_list_is_empty() {
970 let client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
971 assert!(client.cached_tools().is_empty());
972 }
973
974 #[tokio::test]
975 async fn test_mcp_client_list_tools_returns_tools() {
976 let _guard = always_on_tracing_guard();
977 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
978 client.connect().await.unwrap();
979
980 let tools = client
981 .list_tools()
982 .await
983 .expect("list_tools should succeed");
984 assert_eq!(tools.len(), 1);
985 assert_eq!(tools[0].name, "echo");
986
987 assert_eq!(client.cached_tools().len(), 1);
989 assert_eq!(client.cached_tools()[0].name, "echo");
990 }
991
992 #[tokio::test]
993 async fn test_mcp_client_call_tool_returns_result() {
994 let _guard = always_on_tracing_guard();
995 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
996 client.connect().await.unwrap();
997 client.list_tools().await.unwrap();
999
1000 let result = client
1001 .call_tool("echo", serde_json::json!({"msg": "hi"}))
1002 .await
1003 .expect("call_tool should succeed");
1004
1005 assert_eq!(result.content.len(), 1);
1006 assert_eq!(
1007 result.content[0],
1008 ToolResultContent::Text {
1009 text: "hello from tool".to_string()
1010 }
1011 );
1012 }
1013
1014 #[tokio::test]
1015 async fn test_mcp_client_shutdown_succeeds() {
1016 let _guard = always_on_tracing_guard();
1017 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
1018 client.connect().await.unwrap();
1019 client.shutdown().await.expect("shutdown should succeed");
1021 }
1022
1023 #[tokio::test]
1024 async fn test_mcp_client_shutdown_with_dead_process() {
1025 let mut client = spawn_stub_client(STUB_CLOSE_IMMEDIATELY).await;
1026 client
1028 .shutdown()
1029 .await
1030 .expect("shutdown should be graceful");
1031 }
1032
1033 #[tokio::test]
1044 async fn test_mcp_client_server_error_propagates() {
1045 let mut client = spawn_stub_client(STUB_ERROR_SERVER).await;
1046 let err = client.connect().await;
1048 assert!(err.is_err());
1049 assert!(err.unwrap_err().to_string().contains("server error"));
1050 }
1051
1052 #[tokio::test]
1053 async fn test_mcp_client_connect_with_no_capabilities_field() {
1054 let script = r#"
1056import sys, json
1057
1058for line in sys.stdin:
1059 line = line.strip()
1060 if not line:
1061 continue
1062 req = json.loads(line)
1063 method = req.get("method", "")
1064 id_ = req.get("id")
1065 if method == "initialize":
1066 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"protocolVersion": "2024-11-05"}})
1067 sys.stdout.write(msg + "\n")
1068 sys.stdout.flush()
1069 elif method == "notifications/initialized":
1070 pass
1071 elif method == "notifications/cancelled":
1072 pass
1073"#;
1074 let mut client = spawn_stub_client(script).await;
1075 client.connect().await.expect("connect should succeed");
1077 let caps = client.capabilities().unwrap();
1078 assert!(caps.tools.is_none());
1079 }
1080
1081 #[tokio::test]
1082 async fn test_mcp_client_list_tools_missing_tools_key_returns_empty() {
1083 let script = r#"
1085import sys, json
1086
1087for line in sys.stdin:
1088 line = line.strip()
1089 if not line:
1090 continue
1091 req = json.loads(line)
1092 method = req.get("method", "")
1093 id_ = req.get("id")
1094 if method == "initialize":
1095 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"capabilities": {}}})
1096 sys.stdout.write(msg + "\n")
1097 sys.stdout.flush()
1098 elif method == "notifications/initialized":
1099 pass
1100 elif method == "tools/list":
1101 # Return result without "tools" key
1102 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {}})
1103 sys.stdout.write(msg + "\n")
1104 sys.stdout.flush()
1105 elif method == "notifications/cancelled":
1106 pass
1107"#;
1108 let mut client = spawn_stub_client(script).await;
1109 client.connect().await.unwrap();
1110 let tools = client
1111 .list_tools()
1112 .await
1113 .expect("list_tools should succeed");
1114 assert!(tools.is_empty());
1115 }
1116
1117 #[tokio::test]
1118 async fn test_mcp_client_list_tools_malformed_response_is_error() {
1119 let script = r#"
1121import sys, json
1122
1123for line in sys.stdin:
1124 line = line.strip()
1125 if not line:
1126 continue
1127 req = json.loads(line)
1128 method = req.get("method", "")
1129 id_ = req.get("id")
1130 if method == "initialize":
1131 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"capabilities": {}}})
1132 sys.stdout.write(msg + "\n")
1133 sys.stdout.flush()
1134 elif method == "notifications/initialized":
1135 pass
1136 elif method == "tools/list":
1137 # Return tools as a string instead of an array -- triggers parse error
1138 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"tools": "not_an_array"}})
1139 sys.stdout.write(msg + "\n")
1140 sys.stdout.flush()
1141 elif method == "notifications/cancelled":
1142 pass
1143"#;
1144 let mut client = spawn_stub_client(script).await;
1145 client.connect().await.unwrap();
1146 let err = client.list_tools().await;
1147 assert!(err.is_err());
1148 assert!(err.unwrap_err().to_string().contains("parse"));
1149 }
1150
1151 #[tokio::test]
1152 async fn test_mcp_client_call_tool_malformed_result_is_error() {
1153 let script = r#"
1154import sys, json
1155
1156for line in sys.stdin:
1157 line = line.strip()
1158 if not line:
1159 continue
1160 req = json.loads(line)
1161 method = req.get("method", "")
1162 id_ = req.get("id")
1163 if method == "initialize":
1164 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": {"capabilities": {}}})
1165 sys.stdout.write(msg + "\n")
1166 sys.stdout.flush()
1167 elif method == "notifications/initialized":
1168 pass
1169 elif method == "tools/call":
1170 # Return a result that can't be parsed as ToolResult
1171 msg = json.dumps({"jsonrpc": "2.0", "id": id_, "result": "bad_tool_result"})
1172 sys.stdout.write(msg + "\n")
1173 sys.stdout.flush()
1174 elif method == "notifications/cancelled":
1175 pass
1176"#;
1177 let mut client = spawn_stub_client(script).await;
1178 client.connect().await.unwrap();
1179 let err = client.call_tool("broken", serde_json::json!({})).await;
1180 assert!(err.is_err());
1181 assert!(err.unwrap_err().to_string().contains("parse"));
1182 }
1183
1184 #[tokio::test]
1185 async fn test_mcp_client_response_with_no_result_is_error() {
1186 let script = r#"
1188import sys, json
1189
1190for line in sys.stdin:
1191 line = line.strip()
1192 if not line:
1193 continue
1194 req = json.loads(line)
1195 id_ = req.get("id")
1196 if id_ is not None:
1197 # No "result", no "error"
1198 msg = json.dumps({"jsonrpc": "2.0", "id": id_})
1199 sys.stdout.write(msg + "\n")
1200 sys.stdout.flush()
1201"#;
1202 let mut client = spawn_stub_client(script).await;
1203 let err = client.connect().await;
1204 assert!(err.is_err());
1205 assert!(err.unwrap_err().to_string().contains("no result"));
1206 }
1207
1208 #[test]
1216 fn tool_result_reads_is_error_from_camel_case_wire_name() {
1217 let json = r#"{"content":[{"type":"text","text":"boom"}],"isError":true}"#;
1222 let result: ToolResult = serde_json::from_str(json).unwrap();
1223 assert!(result.is_error, "isError:true must deserialize as an error");
1224 }
1225
1226 #[test]
1227 fn tool_result_snake_case_is_error_is_not_honored() {
1228 let json = r#"{"content":[],"is_error":true}"#;
1232 let result: ToolResult = serde_json::from_str(json).unwrap();
1233 assert!(!result.is_error);
1234 }
1235
1236 #[test]
1237 fn tool_result_content_defaults_to_empty() {
1238 let json = r#"{"structuredContent":{"ok":true}}"#;
1239 let result: ToolResult = serde_json::from_str(json).unwrap();
1240 assert!(result.content.is_empty());
1241 assert_eq!(
1242 result.structured_content,
1243 Some(serde_json::json!({"ok": true}))
1244 );
1245 }
1246
1247 #[test]
1248 fn tool_result_structured_content_absent_is_none_and_omitted() {
1249 let result = ToolResult {
1250 content: vec![],
1251 structured_content: None,
1252 is_error: false,
1253 };
1254 let json = serde_json::to_string(&result).unwrap();
1255 assert!(!json.contains("structuredContent"), "got: {json}");
1256 }
1257
1258 #[test]
1259 fn tool_result_content_image_uses_mime_type_wire_name() {
1260 let content = ToolResultContent::Image {
1261 data: "abc".to_string(),
1262 mime_type: "image/png".to_string(),
1263 };
1264 let json = serde_json::to_string(&content).unwrap();
1265 assert!(json.contains("\"mimeType\":\"image/png\""), "got: {json}");
1266 assert!(!json.contains("mime_type"), "got: {json}");
1267 }
1268
1269 #[test]
1270 fn tool_result_content_audio_roundtrip() {
1271 let json = r#"{"type":"audio","data":"YWJj","mimeType":"audio/wav"}"#;
1272 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1273 assert_eq!(
1274 content,
1275 ToolResultContent::Audio {
1276 data: "YWJj".to_string(),
1277 mime_type: "audio/wav".to_string(),
1278 }
1279 );
1280 let back: ToolResultContent =
1281 serde_json::from_str(&serde_json::to_string(&content).unwrap()).unwrap();
1282 assert_eq!(back, content);
1283 }
1284
1285 #[test]
1286 fn tool_result_content_resource_link_full() {
1287 let json = r#"{"type":"resource_link","uri":"file:///m.rs","name":"m.rs",
1288 "description":"entry point","mimeType":"text/x-rust"}"#;
1289 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1290 assert_eq!(
1291 content,
1292 ToolResultContent::ResourceLink {
1293 uri: "file:///m.rs".to_string(),
1294 name: "m.rs".to_string(),
1295 description: Some("entry point".to_string()),
1296 mime_type: Some("text/x-rust".to_string()),
1297 }
1298 );
1299 }
1300
1301 #[test]
1302 fn tool_result_content_resource_link_minimal_omits_optionals() {
1303 let content: ToolResultContent =
1304 serde_json::from_str(r#"{"type":"resource_link","uri":"file:///x"}"#).unwrap();
1305 assert_eq!(
1306 content,
1307 ToolResultContent::ResourceLink {
1308 uri: "file:///x".to_string(),
1309 name: String::new(),
1310 description: None,
1311 mime_type: None,
1312 }
1313 );
1314 let json = serde_json::to_string(&content).unwrap();
1315 assert!(!json.contains("description"), "got: {json}");
1316 assert!(!json.contains("mimeType"), "got: {json}");
1317 }
1318
1319 #[test]
1320 fn tool_result_content_embedded_resource_is_nested() {
1321 let json = r#"{"type":"resource","resource":{"uri":"file:///m.rs",
1322 "mimeType":"text/x-rust","text":"fn main() {}"}}"#;
1323 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1324 assert_eq!(
1325 content,
1326 ToolResultContent::Resource {
1327 resource: EmbeddedResource {
1328 uri: "file:///m.rs".to_string(),
1329 text: Some("fn main() {}".to_string()),
1330 blob: None,
1331 mime_type: Some("text/x-rust".to_string()),
1332 }
1333 }
1334 );
1335 }
1336
1337 #[test]
1338 fn tool_result_content_embedded_resource_binary_blob() {
1339 let json = r#"{"type":"resource","resource":{"uri":"file:///a.png","blob":"YWJj"}}"#;
1340 let content: ToolResultContent = serde_json::from_str(json).unwrap();
1341 assert_eq!(
1342 content,
1343 ToolResultContent::Resource {
1344 resource: EmbeddedResource {
1345 uri: "file:///a.png".to_string(),
1346 text: None,
1347 blob: Some("YWJj".to_string()),
1348 mime_type: None,
1349 }
1350 }
1351 );
1352 let json = serde_json::to_string(&content).unwrap();
1354 assert!(!json.contains("text"), "got: {json}");
1355 assert!(!json.contains("mimeType"), "got: {json}");
1356 }
1357
1358 #[test]
1359 fn unknown_content_type_degrades_instead_of_failing_the_result() {
1360 let json = r#"{"content":[
1364 {"type":"text","text":"keep me"},
1365 {"type":"hologram","payload":{"deeply":["nested"]}}
1366 ],"isError":false}"#;
1367 let result: ToolResult = serde_json::from_str(json).unwrap();
1368 assert_eq!(result.content.len(), 2);
1369 assert_eq!(
1370 result.content[0],
1371 ToolResultContent::Text {
1372 text: "keep me".to_string()
1373 }
1374 );
1375 assert_eq!(result.content[1], ToolResultContent::Unknown);
1376 }
1377
1378 #[test]
1379 fn unknown_content_serializes_without_panicking() {
1380 let json = serde_json::to_string(&ToolResultContent::Unknown).unwrap();
1384 assert!(json.contains("Unknown"), "got: {json}");
1385 }
1386
1387 fn paginated_stub(pages: usize) -> String {
1392 format!(
1393 r#"
1394import sys, json
1395PAGES = {pages}
1396for line in sys.stdin:
1397 line = line.strip()
1398 if not line:
1399 continue
1400 req = json.loads(line)
1401 method, id_ = req.get("method", ""), req.get("id")
1402 if method == "initialize":
1403 res = {{"capabilities": {{}}, "protocolVersion": "2024-11-05"}}
1404 elif method == "tools/list":
1405 cursor = int((req.get("params") or {{}}).get("cursor", "0"))
1406 res = {{"tools": [{{"name": "tool%d" % cursor, "inputSchema": {{}}}}]}}
1407 if cursor + 1 < PAGES:
1408 res["nextCursor"] = str(cursor + 1)
1409 else:
1410 continue
1411 sys.stdout.write(json.dumps({{"jsonrpc": "2.0", "id": id_, "result": res}}) + "\n")
1412 sys.stdout.flush()
1413"#
1414 )
1415 }
1416
1417 #[tokio::test]
1418 async fn list_tools_follows_next_cursor_across_pages() {
1419 let _guard = always_on_tracing_guard();
1420 let script = paginated_stub(3);
1421 let mut client = spawn_stub_client(&script).await;
1422 client.connect().await.unwrap();
1423
1424 let tools = client
1425 .list_tools()
1426 .await
1427 .expect("list_tools should succeed");
1428 let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
1430 assert_eq!(names, vec!["tool0", "tool1", "tool2"]);
1431 assert_eq!(client.cached_tools().len(), 3);
1432 }
1433
1434 #[tokio::test]
1435 async fn list_tools_stops_at_the_page_limit() {
1436 let _guard = always_on_tracing_guard();
1437 let script = paginated_stub(MAX_TOOL_PAGES + 10);
1439 let mut client = spawn_stub_client(&script).await;
1440 client.connect().await.unwrap();
1441
1442 let tools = client
1443 .list_tools()
1444 .await
1445 .expect("list_tools should succeed");
1446 assert_eq!(tools.len(), MAX_TOOL_PAGES);
1447 }
1448
1449 #[tokio::test]
1450 async fn transport_failure_propagates_out_of_a_request() {
1451 let _guard = always_on_tracing_guard();
1452 let mut client = spawn_stub_client("import sys\nsys.exit(0)\n").await;
1456 assert!(client.connect().await.is_err());
1457 }
1458
1459 #[test]
1462 fn preferred_version_is_the_newest_supported() {
1463 assert_eq!(PREFERRED_PROTOCOL_VERSION, SUPPORTED_PROTOCOL_VERSIONS[0]);
1464 let mut sorted = SUPPORTED_PROTOCOL_VERSIONS.to_vec();
1466 sorted.sort_unstable_by(|a, b| b.cmp(a));
1467 assert_eq!(sorted, SUPPORTED_PROTOCOL_VERSIONS);
1468 }
1469
1470 #[test]
1471 fn negotiation_adopts_a_recognized_echo() {
1472 let _guard = always_on_tracing_guard();
1473 let echoed = serde_json::json!("2025-06-18");
1474 assert_eq!(negotiated_version(Some(&echoed)), "2025-06-18");
1475 }
1476
1477 #[test]
1478 fn negotiation_honors_an_unrecognized_echo() {
1479 let _guard = always_on_tracing_guard();
1480 let echoed = serde_json::json!("2099-01-01");
1483 assert_eq!(negotiated_version(Some(&echoed)), "2099-01-01");
1484 }
1485
1486 #[test]
1487 fn negotiation_falls_back_when_the_server_omits_the_field() {
1488 let _guard = always_on_tracing_guard();
1489 assert_eq!(negotiated_version(None), PREFERRED_PROTOCOL_VERSION);
1490 }
1491
1492 #[test]
1493 fn negotiation_falls_back_when_the_echo_is_not_a_string() {
1494 let _guard = always_on_tracing_guard();
1495 let echoed = serde_json::json!(20251125);
1496 assert_eq!(
1497 negotiated_version(Some(&echoed)),
1498 PREFERRED_PROTOCOL_VERSION
1499 );
1500 }
1501
1502 #[tokio::test]
1503 async fn connect_offers_the_preferred_version_and_records_the_echo() {
1504 let _guard = always_on_tracing_guard();
1505 let script = r#"
1508import sys, json
1509for line in sys.stdin:
1510 line = line.strip()
1511 if not line:
1512 continue
1513 req = json.loads(line)
1514 if req.get("method") == "initialize":
1515 offered = req["params"]["protocolVersion"]
1516 assert offered == "2025-11-25", offered
1517 sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": req.get("id"),
1518 "result": {"capabilities": {}, "protocolVersion": "2025-03-26"}}) + "\n")
1519 sys.stdout.flush()
1520"#;
1521 let mut client = spawn_stub_client(script).await;
1522 client.connect().await.expect("connect should succeed");
1523 assert_eq!(client.protocol_version(), Some("2025-03-26"));
1524 }
1525
1526 #[tokio::test]
1527 async fn protocol_version_is_none_before_connect() {
1528 let client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
1529 assert!(client.protocol_version().is_none());
1530 }
1531
1532 #[test]
1535 fn connect_http_builds_a_client_without_touching_the_network() {
1536 assert!(MCPClient::connect_http("http://127.0.0.1:1/mcp", &HashMap::new(), &[]).is_ok());
1539 }
1540
1541 struct NoopRefresher;
1542 #[async_trait::async_trait]
1543 impl crate::transport::BearerRefresher for NoopRefresher {
1544 async fn refresh(&self) -> anyhow::Result<String> {
1545 Ok("Bearer x".to_string())
1546 }
1547 }
1548
1549 #[tokio::test]
1550 async fn set_refresher_on_stdio_is_a_noop() {
1551 let mut client = spawn_stub_client(STUB_INIT_LIST_CALL).await;
1554 client.set_refresher(std::sync::Arc::new(NoopRefresher));
1555 }
1556
1557 #[tokio::test]
1558 async fn set_refresher_on_http_is_accepted() {
1559 let mut client =
1560 MCPClient::connect_http("http://127.0.0.1:1/mcp", &HashMap::new(), &[]).unwrap();
1561 use crate::transport::BearerRefresher as _;
1563 assert_eq!(NoopRefresher.refresh().await.unwrap(), "Bearer x");
1564 client.set_refresher(std::sync::Arc::new(NoopRefresher));
1565 }
1566
1567 #[test]
1568 fn connect_http_rejects_an_unparseable_url() {
1569 assert!(MCPClient::connect_http("not a url", &HashMap::new(), &[]).is_err());
1570 }
1571
1572 #[tokio::test]
1573 async fn from_config_builds_the_stdio_transport() {
1574 let _guard = always_on_tracing_guard();
1575 let config = MCPServerConfig::stdio(
1576 "s",
1577 "python3",
1578 vec!["-c".into(), STUB_INIT_LIST_CALL.into()],
1579 );
1580 let mut client = MCPClient::from_config(&config)
1581 .await
1582 .expect("stdio config should connect");
1583 client.connect().await.expect("handshake should succeed");
1584 assert_eq!(client.list_tools().await.unwrap().len(), 1);
1585 }
1586
1587 #[tokio::test]
1588 async fn from_config_builds_the_http_transport() {
1589 let config = MCPServerConfig::http("s", "http://127.0.0.1:1/mcp");
1590 assert!(MCPClient::from_config(&config).await.is_ok());
1591 }
1592
1593 #[tokio::test]
1594 async fn from_config_with_auth_injects_a_bearer_for_http() {
1595 let config = MCPServerConfig::http("s", "http://127.0.0.1:1/mcp");
1598 let header = Some(("Authorization".to_string(), "Bearer tok".to_string()));
1599 assert!(
1600 MCPClient::from_config_with_auth(&config, header, &[])
1601 .await
1602 .is_ok()
1603 );
1604 }
1605
1606 #[tokio::test]
1607 async fn from_config_rejects_an_unresolvable_entry() {
1608 let config = MCPServerConfig {
1609 name: "broken".to_string(),
1610 ..Default::default()
1611 };
1612 let err = MCPClient::from_config(&config)
1613 .await
1614 .err()
1615 .expect("an entry with neither command nor url cannot connect");
1616 assert!(err.to_string().contains("either a `command`"), "got: {err}");
1617 }
1618}