a3s_code_core/mcp/
protocol.rs

1//! MCP Protocol Type Definitions
2//!
3//! Defines the core types for the Model Context Protocol (MCP).
4//! Based on the MCP specification: https://spec.modelcontextprotocol.io/
5
6use serde::{Deserialize, Deserializer, Serialize};
7use std::collections::HashMap;
8
9/// MCP protocol version
10pub const PROTOCOL_VERSION: &str = "2024-11-05";
11
12/// JSON-RPC request
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct JsonRpcRequest {
15    pub jsonrpc: String,
16    pub id: u64,
17    pub method: String,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub params: Option<serde_json::Value>,
20}
21
22impl JsonRpcRequest {
23    pub fn new(id: u64, method: &str, params: Option<serde_json::Value>) -> Self {
24        Self {
25            jsonrpc: "2.0".to_string(),
26            id,
27            method: method.to_string(),
28            params,
29        }
30    }
31}
32
33/// JSON-RPC response
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct JsonRpcResponse {
36    pub jsonrpc: String,
37    pub id: Option<u64>,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub result: Option<serde_json::Value>,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub error: Option<JsonRpcError>,
42}
43
44/// JSON-RPC error
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct JsonRpcError {
47    pub code: i32,
48    pub message: String,
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub data: Option<serde_json::Value>,
51}
52
53/// JSON-RPC notification (no id)
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct JsonRpcNotification {
56    pub jsonrpc: String,
57    pub method: String,
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub params: Option<serde_json::Value>,
60}
61
62impl JsonRpcNotification {
63    pub fn new(method: &str, params: Option<serde_json::Value>) -> Self {
64        Self {
65            jsonrpc: "2.0".to_string(),
66            method: method.to_string(),
67            params,
68        }
69    }
70}
71
72// ============================================================================
73// MCP Initialize
74// ============================================================================
75
76/// Client capabilities
77#[derive(Debug, Clone, Default, Serialize, Deserialize)]
78pub struct ClientCapabilities {
79    #[serde(skip_serializing_if = "Option::is_none")]
80    pub roots: Option<RootsCapability>,
81    #[serde(skip_serializing_if = "Option::is_none")]
82    pub sampling: Option<SamplingCapability>,
83}
84
85#[derive(Debug, Clone, Default, Serialize, Deserialize)]
86pub struct RootsCapability {
87    #[serde(default)]
88    pub list_changed: bool,
89}
90
91#[derive(Debug, Clone, Default, Serialize, Deserialize)]
92pub struct SamplingCapability {}
93
94/// Client info
95#[derive(Debug, Clone, Serialize, Deserialize)]
96pub struct ClientInfo {
97    pub name: String,
98    pub version: String,
99}
100
101/// Initialize request params
102#[derive(Debug, Clone, Serialize, Deserialize)]
103#[serde(rename_all = "camelCase")]
104pub struct InitializeParams {
105    pub protocol_version: String,
106    pub capabilities: ClientCapabilities,
107    pub client_info: ClientInfo,
108}
109
110/// Server capabilities
111#[derive(Debug, Clone, Default, Serialize, Deserialize)]
112pub struct ServerCapabilities {
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub tools: Option<ToolsCapability>,
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub resources: Option<ResourcesCapability>,
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub prompts: Option<PromptsCapability>,
119    #[serde(skip_serializing_if = "Option::is_none")]
120    pub logging: Option<LoggingCapability>,
121}
122
123#[derive(Debug, Clone, Default, Serialize, Deserialize)]
124#[serde(rename_all = "camelCase")]
125pub struct ToolsCapability {
126    #[serde(default)]
127    pub list_changed: bool,
128}
129
130#[derive(Debug, Clone, Default, Serialize, Deserialize)]
131#[serde(rename_all = "camelCase")]
132pub struct ResourcesCapability {
133    #[serde(default)]
134    pub subscribe: bool,
135    #[serde(default)]
136    pub list_changed: bool,
137}
138
139#[derive(Debug, Clone, Default, Serialize, Deserialize)]
140#[serde(rename_all = "camelCase")]
141pub struct PromptsCapability {
142    #[serde(default)]
143    pub list_changed: bool,
144}
145
146#[derive(Debug, Clone, Default, Serialize, Deserialize)]
147pub struct LoggingCapability {}
148
149/// Server info
150#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct ServerInfo {
152    pub name: String,
153    pub version: String,
154}
155
156/// Initialize result
157#[derive(Debug, Clone, Serialize, Deserialize)]
158#[serde(rename_all = "camelCase")]
159pub struct InitializeResult {
160    pub protocol_version: String,
161    pub capabilities: ServerCapabilities,
162    pub server_info: ServerInfo,
163}
164
165// ============================================================================
166// MCP Tools
167// ============================================================================
168
169/// MCP tool definition
170#[derive(Debug, Clone, Serialize, Deserialize)]
171#[serde(rename_all = "camelCase")]
172pub struct McpTool {
173    pub name: String,
174    #[serde(skip_serializing_if = "Option::is_none")]
175    pub description: Option<String>,
176    pub input_schema: serde_json::Value,
177}
178
179/// List tools result
180#[derive(Debug, Clone, Serialize, Deserialize)]
181pub struct ListToolsResult {
182    pub tools: Vec<McpTool>,
183}
184
185/// Call tool params
186#[derive(Debug, Clone, Serialize, Deserialize)]
187pub struct CallToolParams {
188    pub name: String,
189    #[serde(skip_serializing_if = "Option::is_none")]
190    pub arguments: Option<serde_json::Value>,
191}
192
193/// Tool content types
194#[derive(Debug, Clone, Serialize, Deserialize)]
195#[serde(tag = "type", rename_all = "lowercase")]
196pub enum ToolContent {
197    Text {
198        text: String,
199    },
200    Image {
201        data: String,
202        #[serde(rename = "mimeType")]
203        mime_type: String,
204    },
205    Resource {
206        resource: ResourceContent,
207    },
208}
209
210/// Resource content
211#[derive(Debug, Clone, Serialize, Deserialize)]
212#[serde(rename_all = "camelCase")]
213pub struct ResourceContent {
214    pub uri: String,
215    #[serde(skip_serializing_if = "Option::is_none")]
216    pub mime_type: Option<String>,
217    #[serde(skip_serializing_if = "Option::is_none")]
218    pub text: Option<String>,
219    #[serde(skip_serializing_if = "Option::is_none")]
220    pub blob: Option<String>,
221}
222
223/// Call tool result
224#[derive(Debug, Clone, Serialize, Deserialize)]
225#[serde(rename_all = "camelCase")]
226pub struct CallToolResult {
227    pub content: Vec<ToolContent>,
228    #[serde(default)]
229    pub is_error: bool,
230}
231
232// ============================================================================
233// MCP Resources
234// ============================================================================
235
236/// MCP resource definition
237#[derive(Debug, Clone, Serialize, Deserialize)]
238#[serde(rename_all = "camelCase")]
239pub struct McpResource {
240    pub uri: String,
241    pub name: String,
242    #[serde(skip_serializing_if = "Option::is_none")]
243    pub description: Option<String>,
244    #[serde(skip_serializing_if = "Option::is_none")]
245    pub mime_type: Option<String>,
246}
247
248/// List resources result
249#[derive(Debug, Clone, Serialize, Deserialize)]
250pub struct ListResourcesResult {
251    pub resources: Vec<McpResource>,
252}
253
254/// Read resource params
255#[derive(Debug, Clone, Serialize, Deserialize)]
256pub struct ReadResourceParams {
257    pub uri: String,
258}
259
260/// Read resource result
261#[derive(Debug, Clone, Serialize, Deserialize)]
262pub struct ReadResourceResult {
263    pub contents: Vec<ResourceContent>,
264}
265
266// ============================================================================
267// MCP Prompts
268// ============================================================================
269
270/// MCP prompt definition
271#[derive(Debug, Clone, Serialize, Deserialize)]
272pub struct McpPrompt {
273    pub name: String,
274    #[serde(skip_serializing_if = "Option::is_none")]
275    pub description: Option<String>,
276    #[serde(skip_serializing_if = "Option::is_none")]
277    pub arguments: Option<Vec<PromptArgument>>,
278}
279
280/// Prompt argument
281#[derive(Debug, Clone, Serialize, Deserialize)]
282pub struct PromptArgument {
283    pub name: String,
284    #[serde(skip_serializing_if = "Option::is_none")]
285    pub description: Option<String>,
286    #[serde(default)]
287    pub required: bool,
288}
289
290/// List prompts result
291#[derive(Debug, Clone, Serialize, Deserialize)]
292pub struct ListPromptsResult {
293    pub prompts: Vec<McpPrompt>,
294}
295
296// ============================================================================
297// MCP Notifications
298// ============================================================================
299
300/// MCP notification types
301#[derive(Debug, Clone)]
302pub enum McpNotification {
303    ToolsListChanged,
304    ResourcesListChanged,
305    PromptsListChanged,
306    Progress {
307        progress_token: String,
308        progress: f64,
309        total: Option<f64>,
310    },
311    Log {
312        level: String,
313        logger: Option<String>,
314        data: serde_json::Value,
315    },
316    Unknown {
317        method: String,
318        params: Option<serde_json::Value>,
319    },
320}
321
322impl McpNotification {
323    pub fn from_json_rpc(notification: &JsonRpcNotification) -> Self {
324        match notification.method.as_str() {
325            "notifications/tools/list_changed" => McpNotification::ToolsListChanged,
326            "notifications/resources/list_changed" => McpNotification::ResourcesListChanged,
327            "notifications/prompts/list_changed" => McpNotification::PromptsListChanged,
328            "notifications/progress" => {
329                if let Some(params) = &notification.params {
330                    let progress_token = params
331                        .get("progressToken")
332                        .and_then(|v| v.as_str())
333                        .unwrap_or("")
334                        .to_string();
335                    let progress = params
336                        .get("progress")
337                        .and_then(|v| v.as_f64())
338                        .unwrap_or(0.0);
339                    let total = params.get("total").and_then(|v| v.as_f64());
340                    McpNotification::Progress {
341                        progress_token,
342                        progress,
343                        total,
344                    }
345                } else {
346                    McpNotification::Unknown {
347                        method: notification.method.clone(),
348                        params: notification.params.clone(),
349                    }
350                }
351            }
352            "notifications/message" => {
353                if let Some(params) = &notification.params {
354                    let level = params
355                        .get("level")
356                        .and_then(|v| v.as_str())
357                        .unwrap_or("info")
358                        .to_string();
359                    let logger = params
360                        .get("logger")
361                        .and_then(|v| v.as_str())
362                        .map(|s| s.to_string());
363                    let data = params
364                        .get("data")
365                        .cloned()
366                        .unwrap_or(serde_json::Value::Null);
367                    McpNotification::Log {
368                        level,
369                        logger,
370                        data,
371                    }
372                } else {
373                    McpNotification::Unknown {
374                        method: notification.method.clone(),
375                        params: notification.params.clone(),
376                    }
377                }
378            }
379            _ => McpNotification::Unknown {
380                method: notification.method.clone(),
381                params: notification.params.clone(),
382            },
383        }
384    }
385}
386
387// ============================================================================
388// Configuration Types
389// ============================================================================
390
391/// MCP server configuration
392#[derive(Debug, Clone, Serialize)]
393pub struct McpServerConfig {
394    /// Server name (used for tool prefix)
395    pub name: String,
396    /// Transport configuration
397    pub transport: McpTransportConfig,
398    /// Whether enabled
399    #[serde(default = "default_true")]
400    pub enabled: bool,
401    /// Environment variables
402    #[serde(default)]
403    pub env: HashMap<String, String>,
404    /// OAuth configuration (optional)
405    #[serde(skip_serializing_if = "Option::is_none")]
406    pub oauth: Option<OAuthConfig>,
407    /// Per-tool execution timeout in seconds (default: 60)
408    #[serde(default = "default_tool_timeout")]
409    pub tool_timeout_secs: u64,
410}
411
412impl<'de> Deserialize<'de> for McpServerConfig {
413    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
414        use serde::de::Error;
415        use serde_json::Value;
416
417        let mut map = serde_json::Map::deserialize(deserializer)?;
418
419        // Build transport from flat HCL fields (transport = "stdio", command = "...", args = [...])
420        // or from a nested transport object ({ type = "stdio", command = "..." })
421        let transport = if let Some(t) = map.remove("transport") {
422            match &t {
423                Value::String(kind) => {
424                    // Flat HCL format: transport = "stdio", command = "...", args = [...]
425                    match kind.as_str() {
426                        "stdio" => {
427                            let command = map
428                                .remove("command")
429                                .and_then(|v| v.as_str().map(String::from))
430                                .ok_or_else(|| D::Error::missing_field("command"))?;
431                            let args = map
432                                .remove("args")
433                                .and_then(|v| serde_json::from_value(v).ok())
434                                .unwrap_or_default();
435                            McpTransportConfig::Stdio { command, args }
436                        }
437                        "http" => {
438                            let url = map
439                                .remove("url")
440                                .and_then(|v| v.as_str().map(String::from))
441                                .ok_or_else(|| D::Error::missing_field("url"))?;
442                            let headers = map
443                                .remove("headers")
444                                .and_then(|v| serde_json::from_value(v).ok())
445                                .unwrap_or_default();
446                            McpTransportConfig::Http { url, headers }
447                        }
448                        "streamable-http" | "streamable_http" => {
449                            let url = map
450                                .remove("url")
451                                .and_then(|v| v.as_str().map(String::from))
452                                .ok_or_else(|| D::Error::missing_field("url"))?;
453                            let headers = map
454                                .remove("headers")
455                                .and_then(|v| serde_json::from_value(v).ok())
456                                .unwrap_or_default();
457                            McpTransportConfig::StreamableHttp { url, headers }
458                        }
459                        other => {
460                            return Err(D::Error::unknown_variant(
461                                other,
462                                &["stdio", "http", "streamable-http"],
463                            ));
464                        }
465                    }
466                }
467                // Nested object format: transport { type = "stdio", command = "..." }
468                Value::Object(_) => serde_json::from_value(t).map_err(D::Error::custom)?,
469                _ => return Err(D::Error::custom("transport must be a string or object")),
470            }
471        } else {
472            return Err(D::Error::missing_field("transport"));
473        };
474
475        let name = map
476            .remove("name")
477            .and_then(|v| v.as_str().map(String::from))
478            .ok_or_else(|| D::Error::missing_field("name"))?;
479        let enabled = map
480            .remove("enabled")
481            .and_then(|v| v.as_bool())
482            .unwrap_or(true);
483        let env = map
484            .remove("env")
485            .and_then(|v| serde_json::from_value(v).ok())
486            .unwrap_or_default();
487        let oauth = map
488            .remove("oauth")
489            .and_then(|v| serde_json::from_value(v).ok());
490        let tool_timeout_secs = map
491            .remove("tool_timeout_secs")
492            .and_then(|v| v.as_u64())
493            .unwrap_or(60);
494
495        Ok(McpServerConfig {
496            name,
497            transport,
498            enabled,
499            env,
500            oauth,
501            tool_timeout_secs,
502        })
503    }
504}
505
506#[allow(dead_code)] // used by serde default = "default_tool_timeout"
507fn default_tool_timeout() -> u64 {
508    60
509}
510
511#[allow(dead_code)]
512fn default_true() -> bool {
513    true
514}
515
516/// Transport configuration
517#[derive(Debug, Clone, Serialize, Deserialize)]
518#[serde(tag = "type", rename_all = "kebab-case")]
519pub enum McpTransportConfig {
520    /// Local process (stdio)
521    Stdio {
522        command: String,
523        #[serde(default)]
524        args: Vec<String>,
525    },
526    /// Remote HTTP + SSE (legacy, pre-2025-03-26)
527    Http {
528        url: String,
529        #[serde(default)]
530        headers: HashMap<String, String>,
531    },
532    /// Streamable HTTP (MCP 2025-03-26 spec)
533    ///
534    /// Single endpoint handles all communication.
535    /// POST with `Accept: application/json, text/event-stream`.
536    StreamableHttp {
537        url: String,
538        #[serde(default)]
539        headers: HashMap<String, String>,
540    },
541}
542
543/// OAuth configuration
544#[derive(Debug, Clone, Serialize, Deserialize)]
545pub struct OAuthConfig {
546    pub auth_url: String,
547    pub token_url: String,
548    pub client_id: String,
549    #[serde(skip_serializing_if = "Option::is_none")]
550    pub client_secret: Option<String>,
551    #[serde(default)]
552    pub scopes: Vec<String>,
553    pub redirect_uri: String,
554}
555
556// ============================================================================
557// Tests
558// ============================================================================
559
560#[cfg(test)]
561mod tests {
562    use super::*;
563
564    #[test]
565    fn test_json_rpc_request_serialize() {
566        let req = JsonRpcRequest::new(1, "initialize", Some(serde_json::json!({"test": true})));
567        let json = serde_json::to_string(&req).unwrap();
568        assert!(json.contains("\"jsonrpc\":\"2.0\""));
569        assert!(json.contains("\"id\":1"));
570        assert!(json.contains("\"method\":\"initialize\""));
571    }
572
573    #[test]
574    fn test_json_rpc_response_deserialize() {
575        let json = r#"{"jsonrpc":"2.0","id":1,"result":{"success":true}}"#;
576        let resp: JsonRpcResponse = serde_json::from_str(json).unwrap();
577        assert_eq!(resp.id, Some(1));
578        assert!(resp.result.is_some());
579        assert!(resp.error.is_none());
580    }
581
582    #[test]
583    fn test_json_rpc_error_deserialize() {
584        let json =
585            r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"Invalid Request"}}"#;
586        let resp: JsonRpcResponse = serde_json::from_str(json).unwrap();
587        assert!(resp.error.is_some());
588        let err = resp.error.unwrap();
589        assert_eq!(err.code, -32600);
590    }
591
592    #[test]
593    fn test_mcp_tool_deserialize() {
594        let json = r#"{
595            "name": "create_issue",
596            "description": "Create a GitHub issue",
597            "inputSchema": {
598                "type": "object",
599                "properties": {
600                    "title": {"type": "string"},
601                    "body": {"type": "string"}
602                },
603                "required": ["title"]
604            }
605        }"#;
606        let tool: McpTool = serde_json::from_str(json).unwrap();
607        assert_eq!(tool.name, "create_issue");
608        assert!(tool.description.is_some());
609    }
610
611    #[test]
612    fn test_tool_content_text() {
613        let content = ToolContent::Text {
614            text: "Hello".to_string(),
615        };
616        let json = serde_json::to_string(&content).unwrap();
617        assert!(json.contains("\"type\":\"text\""));
618        assert!(json.contains("\"text\":\"Hello\""));
619    }
620
621    #[test]
622    fn test_mcp_transport_config_stdio() {
623        let json = r#"{
624            "type": "stdio",
625            "command": "npx",
626            "args": ["-y", "@modelcontextprotocol/server-github"]
627        }"#;
628        let config: McpTransportConfig = serde_json::from_str(json).unwrap();
629        match config {
630            McpTransportConfig::Stdio { command, args } => {
631                assert_eq!(command, "npx");
632                assert_eq!(args.len(), 2);
633            }
634            _ => panic!("Expected Stdio transport"),
635        }
636    }
637
638    #[test]
639    fn test_mcp_transport_config_http() {
640        let json = r#"{
641            "type": "http",
642            "url": "https://mcp.example.com/api",
643            "headers": {"Authorization": "Bearer token"}
644        }"#;
645        let config: McpTransportConfig = serde_json::from_str(json).unwrap();
646        match config {
647            McpTransportConfig::Http { url, headers } => {
648                assert_eq!(url, "https://mcp.example.com/api");
649                assert!(headers.contains_key("Authorization"));
650            }
651            _ => panic!("Expected Http transport"),
652        }
653    }
654
655    #[test]
656    fn test_mcp_notification_parse() {
657        let notification = JsonRpcNotification::new("notifications/tools/list_changed", None);
658        let mcp_notif = McpNotification::from_json_rpc(&notification);
659        match mcp_notif {
660            McpNotification::ToolsListChanged => {}
661            _ => panic!("Expected ToolsListChanged"),
662        }
663    }
664    #[test]
665    fn test_json_rpc_request_new_with_params() {
666        let req = JsonRpcRequest::new(1, "initialize", Some(serde_json::json!({"test": true})));
667        assert_eq!(req.jsonrpc, "2.0");
668        assert_eq!(req.id, 1);
669        assert_eq!(req.method, "initialize");
670        assert!(req.params.is_some());
671    }
672
673    #[test]
674    fn test_json_rpc_request_new_without_params() {
675        let req = JsonRpcRequest::new(2, "ping", None);
676        assert_eq!(req.jsonrpc, "2.0");
677        assert_eq!(req.id, 2);
678        assert_eq!(req.method, "ping");
679        assert!(req.params.is_none());
680    }
681
682    #[test]
683    fn test_json_rpc_request_serialization() {
684        let req = JsonRpcRequest::new(1, "test_method", Some(serde_json::json!({"key": "value"})));
685        let json = serde_json::to_string(&req).unwrap();
686        assert!(json.contains("\"jsonrpc\":\"2.0\""));
687        assert!(json.contains("\"id\":1"));
688        assert!(json.contains("\"method\":\"test_method\""));
689        assert!(json.contains("\"params\""));
690    }
691
692    #[test]
693    fn test_json_rpc_response_with_result() {
694        let resp = JsonRpcResponse {
695            jsonrpc: "2.0".to_string(),
696            id: Some(1),
697            result: Some(serde_json::json!({"success": true})),
698            error: None,
699        };
700        assert!(resp.result.is_some());
701        assert!(resp.error.is_none());
702    }
703
704    #[test]
705    fn test_json_rpc_response_with_error() {
706        let resp = JsonRpcResponse {
707            jsonrpc: "2.0".to_string(),
708            id: Some(1),
709            result: None,
710            error: Some(JsonRpcError {
711                code: -32600,
712                message: "Invalid Request".to_string(),
713                data: None,
714            }),
715        };
716        assert!(resp.result.is_none());
717        assert!(resp.error.is_some());
718    }
719
720    #[test]
721    fn test_json_rpc_response_both_none() {
722        let resp = JsonRpcResponse {
723            jsonrpc: "2.0".to_string(),
724            id: Some(1),
725            result: None,
726            error: None,
727        };
728        assert!(resp.result.is_none());
729        assert!(resp.error.is_none());
730    }
731
732    #[test]
733    fn test_json_rpc_response_serialization() {
734        let resp = JsonRpcResponse {
735            jsonrpc: "2.0".to_string(),
736            id: Some(1),
737            result: Some(serde_json::json!({"data": "test"})),
738            error: None,
739        };
740        let json = serde_json::to_string(&resp).unwrap();
741        assert!(json.contains("\"jsonrpc\":\"2.0\""));
742        assert!(json.contains("\"id\":1"));
743        assert!(json.contains("\"result\""));
744    }
745
746    #[test]
747    fn test_json_rpc_notification_new_with_params() {
748        let notif =
749            JsonRpcNotification::new("notification", Some(serde_json::json!({"msg": "hello"})));
750        assert_eq!(notif.jsonrpc, "2.0");
751        assert_eq!(notif.method, "notification");
752        assert!(notif.params.is_some());
753    }
754
755    #[test]
756    fn test_json_rpc_notification_new_without_params() {
757        let notif = JsonRpcNotification::new("ping", None);
758        assert_eq!(notif.jsonrpc, "2.0");
759        assert_eq!(notif.method, "ping");
760        assert!(notif.params.is_none());
761    }
762
763    #[test]
764    fn test_json_rpc_notification_serialization() {
765        let notif = JsonRpcNotification::new(
766            "test_notification",
767            Some(serde_json::json!({"key": "value"})),
768        );
769        let json = serde_json::to_string(&notif).unwrap();
770        assert!(json.contains("\"jsonrpc\":\"2.0\""));
771        assert!(json.contains("\"method\":\"test_notification\""));
772        assert!(!json.contains("\"id\""));
773    }
774
775    #[test]
776    fn test_mcp_tool_serialize() {
777        let tool = McpTool {
778            name: "test_tool".to_string(),
779            description: Some("A test tool".to_string()),
780            input_schema: serde_json::json!({"type": "object"}),
781        };
782        let json = serde_json::to_string(&tool).unwrap();
783        assert!(json.contains("\"name\":\"test_tool\""));
784        assert!(json.contains("\"description\":\"A test tool\""));
785    }
786
787    #[test]
788    fn test_mcp_tool_without_description() {
789        let json = r#"{"name":"tool","inputSchema":{"type":"object"}}"#;
790        let tool: McpTool = serde_json::from_str(json).unwrap();
791        assert_eq!(tool.name, "tool");
792        assert!(tool.description.is_none());
793    }
794
795    #[test]
796    fn test_mcp_resource_serialize() {
797        let resource = McpResource {
798            uri: "file:///test.txt".to_string(),
799            name: "test.txt".to_string(),
800            description: Some("Test file".to_string()),
801            mime_type: Some("text/plain".to_string()),
802        };
803        let json = serde_json::to_string(&resource).unwrap();
804        assert!(json.contains("\"uri\":\"file:///test.txt\""));
805        assert!(json.contains("\"name\":\"test.txt\""));
806    }
807
808    #[test]
809    fn test_mcp_resource_deserialize() {
810        let json = r#"{"uri":"file:///doc.md","name":"doc.md","mimeType":"text/markdown"}"#;
811        let resource: McpResource = serde_json::from_str(json).unwrap();
812        assert_eq!(resource.uri, "file:///doc.md");
813        assert_eq!(resource.name, "doc.md");
814        assert_eq!(resource.mime_type, Some("text/markdown".to_string()));
815    }
816
817    #[test]
818    fn test_initialize_params_serialization() {
819        let params = InitializeParams {
820            protocol_version: PROTOCOL_VERSION.to_string(),
821            capabilities: ClientCapabilities::default(),
822            client_info: ClientInfo {
823                name: "test-client".to_string(),
824                version: "1.0.0".to_string(),
825            },
826        };
827        let json = serde_json::to_string(&params).unwrap();
828        assert!(json.contains("\"protocolVersion\""));
829        assert!(json.contains("\"clientInfo\""));
830    }
831
832    #[test]
833    fn test_initialize_result_serialization() {
834        let result = InitializeResult {
835            protocol_version: PROTOCOL_VERSION.to_string(),
836            capabilities: ServerCapabilities::default(),
837            server_info: ServerInfo {
838                name: "test-server".to_string(),
839                version: "1.0.0".to_string(),
840            },
841        };
842        let json = serde_json::to_string(&result).unwrap();
843        assert!(json.contains("\"protocolVersion\""));
844        assert!(json.contains("\"serverInfo\""));
845    }
846
847    #[test]
848    fn test_call_tool_params_serialization() {
849        let params = CallToolParams {
850            name: "test_tool".to_string(),
851            arguments: Some(serde_json::json!({"arg1": "value1"})),
852        };
853        let json = serde_json::to_string(&params).unwrap();
854        assert!(json.contains("\"name\":\"test_tool\""));
855        assert!(json.contains("\"arguments\""));
856    }
857
858    #[test]
859    fn test_call_tool_params_without_arguments() {
860        let params = CallToolParams {
861            name: "simple_tool".to_string(),
862            arguments: None,
863        };
864        let json = serde_json::to_string(&params).unwrap();
865        assert!(json.contains("\"name\":\"simple_tool\""));
866        assert!(!json.contains("\"arguments\""));
867    }
868
869    #[test]
870    fn test_call_tool_result_serialization() {
871        let result = CallToolResult {
872            content: vec![ToolContent::Text {
873                text: "Result".to_string(),
874            }],
875            is_error: false,
876        };
877        let json = serde_json::to_string(&result).unwrap();
878        assert!(json.contains("\"content\""));
879        assert!(json.contains("\"isError\":false"));
880    }
881
882    #[test]
883    fn test_call_tool_result_error_flag() {
884        let result = CallToolResult {
885            content: vec![ToolContent::Text {
886                text: "Error occurred".to_string(),
887            }],
888            is_error: true,
889        };
890        assert!(result.is_error);
891    }
892
893    #[test]
894    fn test_call_tool_result_default() {
895        let json = r#"{"content":[]}"#;
896        let result: CallToolResult = serde_json::from_str(json).unwrap();
897        assert!(!result.is_error);
898    }
899
900    #[test]
901    fn test_read_resource_params_serialization() {
902        let params = ReadResourceParams {
903            uri: "file:///test.txt".to_string(),
904        };
905        let json = serde_json::to_string(&params).unwrap();
906        assert!(json.contains("\"uri\":\"file:///test.txt\""));
907    }
908
909    #[test]
910    fn test_read_resource_result_serialization() {
911        let result = ReadResourceResult {
912            contents: vec![ResourceContent {
913                uri: "file:///test.txt".to_string(),
914                mime_type: Some("text/plain".to_string()),
915                text: Some("Hello".to_string()),
916                blob: None,
917            }],
918        };
919        let json = serde_json::to_string(&result).unwrap();
920        assert!(json.contains("\"contents\""));
921        assert!(json.contains("\"uri\""));
922    }
923
924    #[test]
925    fn test_list_tools_result_serialization() {
926        let result = ListToolsResult {
927            tools: vec![McpTool {
928                name: "tool1".to_string(),
929                description: None,
930                input_schema: serde_json::json!({"type": "object"}),
931            }],
932        };
933        let json = serde_json::to_string(&result).unwrap();
934        assert!(json.contains("\"tools\""));
935    }
936
937    #[test]
938    fn test_list_resources_result_serialization() {
939        let result = ListResourcesResult {
940            resources: vec![McpResource {
941                uri: "file:///test.txt".to_string(),
942                name: "test.txt".to_string(),
943                description: None,
944                mime_type: None,
945            }],
946        };
947        let json = serde_json::to_string(&result).unwrap();
948        assert!(json.contains("\"resources\""));
949    }
950
951    #[test]
952    fn test_server_capabilities_default() {
953        let caps = ServerCapabilities::default();
954        assert!(caps.tools.is_none());
955        assert!(caps.resources.is_none());
956        assert!(caps.prompts.is_none());
957        assert!(caps.logging.is_none());
958    }
959
960    #[test]
961    fn test_server_capabilities_all_fields() {
962        let caps = ServerCapabilities {
963            tools: Some(ToolsCapability { list_changed: true }),
964            resources: Some(ResourcesCapability {
965                subscribe: true,
966                list_changed: true,
967            }),
968            prompts: Some(PromptsCapability { list_changed: true }),
969            logging: Some(LoggingCapability {}),
970        };
971        assert!(caps.tools.is_some());
972        assert!(caps.resources.is_some());
973        assert!(caps.prompts.is_some());
974        assert!(caps.logging.is_some());
975    }
976
977    #[test]
978    fn test_client_capabilities_default() {
979        let caps = ClientCapabilities::default();
980        assert!(caps.roots.is_none());
981        assert!(caps.sampling.is_none());
982    }
983
984    #[test]
985    fn test_client_capabilities_all_fields() {
986        let caps = ClientCapabilities {
987            roots: Some(RootsCapability { list_changed: true }),
988            sampling: Some(SamplingCapability {}),
989        };
990        assert!(caps.roots.is_some());
991        assert!(caps.sampling.is_some());
992    }
993
994    #[test]
995    fn test_mcp_notification_tools_list_changed() {
996        let notif = JsonRpcNotification::new("notifications/tools/list_changed", None);
997        let mcp_notif = McpNotification::from_json_rpc(&notif);
998        match mcp_notif {
999            McpNotification::ToolsListChanged => {}
1000            _ => panic!("Expected ToolsListChanged"),
1001        }
1002    }
1003
1004    #[test]
1005    fn test_mcp_notification_resources_list_changed() {
1006        let notif = JsonRpcNotification::new("notifications/resources/list_changed", None);
1007        let mcp_notif = McpNotification::from_json_rpc(&notif);
1008        match mcp_notif {
1009            McpNotification::ResourcesListChanged => {}
1010            _ => panic!("Expected ResourcesListChanged"),
1011        }
1012    }
1013
1014    #[test]
1015    fn test_mcp_notification_prompts_list_changed() {
1016        let notif = JsonRpcNotification::new("notifications/prompts/list_changed", None);
1017        let mcp_notif = McpNotification::from_json_rpc(&notif);
1018        match mcp_notif {
1019            McpNotification::PromptsListChanged => {}
1020            _ => panic!("Expected PromptsListChanged"),
1021        }
1022    }
1023
1024    #[test]
1025    fn test_mcp_notification_progress() {
1026        let notif = JsonRpcNotification::new(
1027            "notifications/progress",
1028            Some(serde_json::json!({
1029                "progressToken": "token-123",
1030                "progress": 50.0,
1031                "total": 100.0
1032            })),
1033        );
1034        let mcp_notif = McpNotification::from_json_rpc(&notif);
1035        match mcp_notif {
1036            McpNotification::Progress {
1037                progress_token,
1038                progress,
1039                total,
1040            } => {
1041                assert_eq!(progress_token, "token-123");
1042                assert_eq!(progress, 50.0);
1043                assert_eq!(total, Some(100.0));
1044            }
1045            _ => panic!("Expected Progress"),
1046        }
1047    }
1048
1049    #[test]
1050    fn test_mcp_notification_log() {
1051        let notif = JsonRpcNotification::new(
1052            "notifications/message",
1053            Some(serde_json::json!({
1054                "level": "error",
1055                "logger": "test-logger",
1056                "data": {"message": "test"}
1057            })),
1058        );
1059        let mcp_notif = McpNotification::from_json_rpc(&notif);
1060        match mcp_notif {
1061            McpNotification::Log {
1062                level,
1063                logger,
1064                data,
1065            } => {
1066                assert_eq!(level, "error");
1067                assert_eq!(logger, Some("test-logger".to_string()));
1068                assert!(data.is_object());
1069            }
1070            _ => panic!("Expected Log"),
1071        }
1072    }
1073
1074    #[test]
1075    fn test_mcp_notification_log_edge_case_no_logger() {
1076        let notif = JsonRpcNotification::new(
1077            "notifications/message",
1078            Some(serde_json::json!({
1079                "level": "info",
1080                "data": "simple message"
1081            })),
1082        );
1083        let mcp_notif = McpNotification::from_json_rpc(&notif);
1084        match mcp_notif {
1085            McpNotification::Log { level, logger, .. } => {
1086                assert_eq!(level, "info");
1087                assert!(logger.is_none());
1088            }
1089            _ => panic!("Expected Log"),
1090        }
1091    }
1092
1093    #[test]
1094    fn test_mcp_notification_log_edge_case_default_level() {
1095        let notif = JsonRpcNotification::new(
1096            "notifications/message",
1097            Some(serde_json::json!({
1098                "data": "message"
1099            })),
1100        );
1101        let mcp_notif = McpNotification::from_json_rpc(&notif);
1102        match mcp_notif {
1103            McpNotification::Log { level, .. } => {
1104                assert_eq!(level, "info");
1105            }
1106            _ => panic!("Expected Log"),
1107        }
1108    }
1109
1110    #[test]
1111    fn test_mcp_notification_unknown() {
1112        let notif = JsonRpcNotification::new(
1113            "unknown/notification",
1114            Some(serde_json::json!({"key": "value"})),
1115        );
1116        let mcp_notif = McpNotification::from_json_rpc(&notif);
1117        match mcp_notif {
1118            McpNotification::Unknown { method, params } => {
1119                assert_eq!(method, "unknown/notification");
1120                assert!(params.is_some());
1121            }
1122            _ => panic!("Expected Unknown"),
1123        }
1124    }
1125
1126    #[test]
1127    fn test_tool_content_image() {
1128        let content = ToolContent::Image {
1129            data: "base64data".to_string(),
1130            mime_type: "image/png".to_string(),
1131        };
1132        let json = serde_json::to_string(&content).unwrap();
1133        assert!(json.contains("\"type\":\"image\""));
1134        assert!(json.contains("\"data\":\"base64data\""));
1135        assert!(json.contains("\"mimeType\":\"image/png\""));
1136    }
1137
1138    #[test]
1139    fn test_tool_content_resource() {
1140        let content = ToolContent::Resource {
1141            resource: ResourceContent {
1142                uri: "file:///test.txt".to_string(),
1143                mime_type: Some("text/plain".to_string()),
1144                text: Some("content".to_string()),
1145                blob: None,
1146            },
1147        };
1148        let json = serde_json::to_string(&content).unwrap();
1149        assert!(json.contains("\"type\":\"resource\""));
1150        assert!(json.contains("\"uri\":\"file:///test.txt\""));
1151    }
1152
1153    #[test]
1154    fn test_mcp_server_config_default() {
1155        let config = McpServerConfig {
1156            name: "test-server".to_string(),
1157            transport: McpTransportConfig::Stdio {
1158                command: "node".to_string(),
1159                args: vec!["server.js".to_string()],
1160            },
1161            enabled: true,
1162            env: HashMap::new(),
1163            oauth: None,
1164            tool_timeout_secs: 60,
1165        };
1166        assert!(config.enabled);
1167        assert!(config.oauth.is_none());
1168    }
1169
1170    #[test]
1171    fn test_mcp_server_config_with_env() {
1172        let mut env = HashMap::new();
1173        env.insert("API_KEY".to_string(), "secret".to_string());
1174        let config = McpServerConfig {
1175            name: "test-server".to_string(),
1176            transport: McpTransportConfig::Stdio {
1177                command: "node".to_string(),
1178                args: vec![],
1179            },
1180            enabled: true,
1181            env,
1182            oauth: None,
1183            tool_timeout_secs: 60,
1184        };
1185        assert!(config.env.contains_key("API_KEY"));
1186    }
1187
1188    #[test]
1189    fn test_mcp_server_config_with_oauth() {
1190        let config = McpServerConfig {
1191            name: "test-server".to_string(),
1192            transport: McpTransportConfig::Http {
1193                url: "https://api.example.com".to_string(),
1194                headers: HashMap::new(),
1195            },
1196            enabled: true,
1197            env: HashMap::new(),
1198            oauth: Some(OAuthConfig {
1199                auth_url: "https://auth.example.com".to_string(),
1200                token_url: "https://token.example.com".to_string(),
1201                client_id: "client-123".to_string(),
1202                client_secret: Some("secret".to_string()),
1203                scopes: vec!["read".to_string(), "write".to_string()],
1204                redirect_uri: "http://localhost:8080/callback".to_string(),
1205            }),
1206            tool_timeout_secs: 60,
1207        };
1208        assert!(config.oauth.is_some());
1209    }
1210
1211    #[test]
1212    fn test_mcp_transport_config_stdio_variant() {
1213        let transport = McpTransportConfig::Stdio {
1214            command: "python".to_string(),
1215            args: vec!["-m".to_string(), "server".to_string()],
1216        };
1217        match transport {
1218            McpTransportConfig::Stdio { command, args } => {
1219                assert_eq!(command, "python");
1220                assert_eq!(args.len(), 2);
1221            }
1222            _ => panic!("Expected Stdio"),
1223        }
1224    }
1225
1226    #[test]
1227    fn test_mcp_transport_config_http_variant() {
1228        let mut headers = HashMap::new();
1229        headers.insert("Authorization".to_string(), "Bearer token".to_string());
1230        let transport = McpTransportConfig::Http {
1231            url: "https://mcp.example.com".to_string(),
1232            headers,
1233        };
1234        match transport {
1235            McpTransportConfig::Http { url, headers } => {
1236                assert_eq!(url, "https://mcp.example.com");
1237                assert!(headers.contains_key("Authorization"));
1238            }
1239            _ => panic!("Expected Http"),
1240        }
1241    }
1242
1243    #[test]
1244    fn test_mcp_prompt_serialize() {
1245        let prompt = McpPrompt {
1246            name: "test_prompt".to_string(),
1247            description: Some("A test prompt".to_string()),
1248            arguments: Some(vec![PromptArgument {
1249                name: "arg1".to_string(),
1250                description: Some("First argument".to_string()),
1251                required: true,
1252            }]),
1253        };
1254        let json = serde_json::to_string(&prompt).unwrap();
1255        assert!(json.contains("\"name\":\"test_prompt\""));
1256        assert!(json.contains("\"arguments\""));
1257    }
1258
1259    #[test]
1260    fn test_prompt_argument_default() {
1261        let json = r#"{"name":"arg"}"#;
1262        let arg: PromptArgument = serde_json::from_str(json).unwrap();
1263        assert_eq!(arg.name, "arg");
1264        assert!(!arg.required);
1265    }
1266}
a3s_code_core/mcp/protocol.rs

a3s_code_core/mcp/
protocol.rs
