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    /// Static access token — if set, skips the OAuth exchange flow.
555    /// Useful for long-lived tokens or service accounts.
556    #[serde(skip_serializing_if = "Option::is_none")]
557    pub access_token: Option<String>,
558}
559
560// ============================================================================
561// Tests
562// ============================================================================
563
564#[cfg(test)]
565mod tests {
566    use super::*;
567
568    #[test]
569    fn test_json_rpc_request_serialize() {
570        let req = JsonRpcRequest::new(1, "initialize", Some(serde_json::json!({"test": true})));
571        let json = serde_json::to_string(&req).unwrap();
572        assert!(json.contains("\"jsonrpc\":\"2.0\""));
573        assert!(json.contains("\"id\":1"));
574        assert!(json.contains("\"method\":\"initialize\""));
575    }
576
577    #[test]
578    fn test_json_rpc_response_deserialize() {
579        let json = r#"{"jsonrpc":"2.0","id":1,"result":{"success":true}}"#;
580        let resp: JsonRpcResponse = serde_json::from_str(json).unwrap();
581        assert_eq!(resp.id, Some(1));
582        assert!(resp.result.is_some());
583        assert!(resp.error.is_none());
584    }
585
586    #[test]
587    fn test_json_rpc_error_deserialize() {
588        let json =
589            r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"Invalid Request"}}"#;
590        let resp: JsonRpcResponse = serde_json::from_str(json).unwrap();
591        assert!(resp.error.is_some());
592        let err = resp.error.unwrap();
593        assert_eq!(err.code, -32600);
594    }
595
596    #[test]
597    fn test_mcp_tool_deserialize() {
598        let json = r#"{
599            "name": "create_issue",
600            "description": "Create a GitHub issue",
601            "inputSchema": {
602                "type": "object",
603                "properties": {
604                    "title": {"type": "string"},
605                    "body": {"type": "string"}
606                },
607                "required": ["title"]
608            }
609        }"#;
610        let tool: McpTool = serde_json::from_str(json).unwrap();
611        assert_eq!(tool.name, "create_issue");
612        assert!(tool.description.is_some());
613    }
614
615    #[test]
616    fn test_tool_content_text() {
617        let content = ToolContent::Text {
618            text: "Hello".to_string(),
619        };
620        let json = serde_json::to_string(&content).unwrap();
621        assert!(json.contains("\"type\":\"text\""));
622        assert!(json.contains("\"text\":\"Hello\""));
623    }
624
625    #[test]
626    fn test_mcp_transport_config_stdio() {
627        let json = r#"{
628            "type": "stdio",
629            "command": "npx",
630            "args": ["-y", "@modelcontextprotocol/server-github"]
631        }"#;
632        let config: McpTransportConfig = serde_json::from_str(json).unwrap();
633        match config {
634            McpTransportConfig::Stdio { command, args } => {
635                assert_eq!(command, "npx");
636                assert_eq!(args.len(), 2);
637            }
638            _ => panic!("Expected Stdio transport"),
639        }
640    }
641
642    #[test]
643    fn test_mcp_transport_config_http() {
644        let json = r#"{
645            "type": "http",
646            "url": "https://mcp.example.com/api",
647            "headers": {"Authorization": "Bearer token"}
648        }"#;
649        let config: McpTransportConfig = serde_json::from_str(json).unwrap();
650        match config {
651            McpTransportConfig::Http { url, headers } => {
652                assert_eq!(url, "https://mcp.example.com/api");
653                assert!(headers.contains_key("Authorization"));
654            }
655            _ => panic!("Expected Http transport"),
656        }
657    }
658
659    #[test]
660    fn test_mcp_notification_parse() {
661        let notification = JsonRpcNotification::new("notifications/tools/list_changed", None);
662        let mcp_notif = McpNotification::from_json_rpc(&notification);
663        match mcp_notif {
664            McpNotification::ToolsListChanged => {}
665            _ => panic!("Expected ToolsListChanged"),
666        }
667    }
668    #[test]
669    fn test_json_rpc_request_new_with_params() {
670        let req = JsonRpcRequest::new(1, "initialize", Some(serde_json::json!({"test": true})));
671        assert_eq!(req.jsonrpc, "2.0");
672        assert_eq!(req.id, 1);
673        assert_eq!(req.method, "initialize");
674        assert!(req.params.is_some());
675    }
676
677    #[test]
678    fn test_json_rpc_request_new_without_params() {
679        let req = JsonRpcRequest::new(2, "ping", None);
680        assert_eq!(req.jsonrpc, "2.0");
681        assert_eq!(req.id, 2);
682        assert_eq!(req.method, "ping");
683        assert!(req.params.is_none());
684    }
685
686    #[test]
687    fn test_json_rpc_request_serialization() {
688        let req = JsonRpcRequest::new(1, "test_method", Some(serde_json::json!({"key": "value"})));
689        let json = serde_json::to_string(&req).unwrap();
690        assert!(json.contains("\"jsonrpc\":\"2.0\""));
691        assert!(json.contains("\"id\":1"));
692        assert!(json.contains("\"method\":\"test_method\""));
693        assert!(json.contains("\"params\""));
694    }
695
696    #[test]
697    fn test_json_rpc_response_with_result() {
698        let resp = JsonRpcResponse {
699            jsonrpc: "2.0".to_string(),
700            id: Some(1),
701            result: Some(serde_json::json!({"success": true})),
702            error: None,
703        };
704        assert!(resp.result.is_some());
705        assert!(resp.error.is_none());
706    }
707
708    #[test]
709    fn test_json_rpc_response_with_error() {
710        let resp = JsonRpcResponse {
711            jsonrpc: "2.0".to_string(),
712            id: Some(1),
713            result: None,
714            error: Some(JsonRpcError {
715                code: -32600,
716                message: "Invalid Request".to_string(),
717                data: None,
718            }),
719        };
720        assert!(resp.result.is_none());
721        assert!(resp.error.is_some());
722    }
723
724    #[test]
725    fn test_json_rpc_response_both_none() {
726        let resp = JsonRpcResponse {
727            jsonrpc: "2.0".to_string(),
728            id: Some(1),
729            result: None,
730            error: None,
731        };
732        assert!(resp.result.is_none());
733        assert!(resp.error.is_none());
734    }
735
736    #[test]
737    fn test_json_rpc_response_serialization() {
738        let resp = JsonRpcResponse {
739            jsonrpc: "2.0".to_string(),
740            id: Some(1),
741            result: Some(serde_json::json!({"data": "test"})),
742            error: None,
743        };
744        let json = serde_json::to_string(&resp).unwrap();
745        assert!(json.contains("\"jsonrpc\":\"2.0\""));
746        assert!(json.contains("\"id\":1"));
747        assert!(json.contains("\"result\""));
748    }
749
750    #[test]
751    fn test_json_rpc_notification_new_with_params() {
752        let notif =
753            JsonRpcNotification::new("notification", Some(serde_json::json!({"msg": "hello"})));
754        assert_eq!(notif.jsonrpc, "2.0");
755        assert_eq!(notif.method, "notification");
756        assert!(notif.params.is_some());
757    }
758
759    #[test]
760    fn test_json_rpc_notification_new_without_params() {
761        let notif = JsonRpcNotification::new("ping", None);
762        assert_eq!(notif.jsonrpc, "2.0");
763        assert_eq!(notif.method, "ping");
764        assert!(notif.params.is_none());
765    }
766
767    #[test]
768    fn test_json_rpc_notification_serialization() {
769        let notif = JsonRpcNotification::new(
770            "test_notification",
771            Some(serde_json::json!({"key": "value"})),
772        );
773        let json = serde_json::to_string(&notif).unwrap();
774        assert!(json.contains("\"jsonrpc\":\"2.0\""));
775        assert!(json.contains("\"method\":\"test_notification\""));
776        assert!(!json.contains("\"id\""));
777    }
778
779    #[test]
780    fn test_mcp_tool_serialize() {
781        let tool = McpTool {
782            name: "test_tool".to_string(),
783            description: Some("A test tool".to_string()),
784            input_schema: serde_json::json!({"type": "object"}),
785        };
786        let json = serde_json::to_string(&tool).unwrap();
787        assert!(json.contains("\"name\":\"test_tool\""));
788        assert!(json.contains("\"description\":\"A test tool\""));
789    }
790
791    #[test]
792    fn test_mcp_tool_without_description() {
793        let json = r#"{"name":"tool","inputSchema":{"type":"object"}}"#;
794        let tool: McpTool = serde_json::from_str(json).unwrap();
795        assert_eq!(tool.name, "tool");
796        assert!(tool.description.is_none());
797    }
798
799    #[test]
800    fn test_mcp_resource_serialize() {
801        let resource = McpResource {
802            uri: "file:///test.txt".to_string(),
803            name: "test.txt".to_string(),
804            description: Some("Test file".to_string()),
805            mime_type: Some("text/plain".to_string()),
806        };
807        let json = serde_json::to_string(&resource).unwrap();
808        assert!(json.contains("\"uri\":\"file:///test.txt\""));
809        assert!(json.contains("\"name\":\"test.txt\""));
810    }
811
812    #[test]
813    fn test_mcp_resource_deserialize() {
814        let json = r#"{"uri":"file:///doc.md","name":"doc.md","mimeType":"text/markdown"}"#;
815        let resource: McpResource = serde_json::from_str(json).unwrap();
816        assert_eq!(resource.uri, "file:///doc.md");
817        assert_eq!(resource.name, "doc.md");
818        assert_eq!(resource.mime_type, Some("text/markdown".to_string()));
819    }
820
821    #[test]
822    fn test_initialize_params_serialization() {
823        let params = InitializeParams {
824            protocol_version: PROTOCOL_VERSION.to_string(),
825            capabilities: ClientCapabilities::default(),
826            client_info: ClientInfo {
827                name: "test-client".to_string(),
828                version: "1.0.0".to_string(),
829            },
830        };
831        let json = serde_json::to_string(&params).unwrap();
832        assert!(json.contains("\"protocolVersion\""));
833        assert!(json.contains("\"clientInfo\""));
834    }
835
836    #[test]
837    fn test_initialize_result_serialization() {
838        let result = InitializeResult {
839            protocol_version: PROTOCOL_VERSION.to_string(),
840            capabilities: ServerCapabilities::default(),
841            server_info: ServerInfo {
842                name: "test-server".to_string(),
843                version: "1.0.0".to_string(),
844            },
845        };
846        let json = serde_json::to_string(&result).unwrap();
847        assert!(json.contains("\"protocolVersion\""));
848        assert!(json.contains("\"serverInfo\""));
849    }
850
851    #[test]
852    fn test_call_tool_params_serialization() {
853        let params = CallToolParams {
854            name: "test_tool".to_string(),
855            arguments: Some(serde_json::json!({"arg1": "value1"})),
856        };
857        let json = serde_json::to_string(&params).unwrap();
858        assert!(json.contains("\"name\":\"test_tool\""));
859        assert!(json.contains("\"arguments\""));
860    }
861
862    #[test]
863    fn test_call_tool_params_without_arguments() {
864        let params = CallToolParams {
865            name: "simple_tool".to_string(),
866            arguments: None,
867        };
868        let json = serde_json::to_string(&params).unwrap();
869        assert!(json.contains("\"name\":\"simple_tool\""));
870        assert!(!json.contains("\"arguments\""));
871    }
872
873    #[test]
874    fn test_call_tool_result_serialization() {
875        let result = CallToolResult {
876            content: vec![ToolContent::Text {
877                text: "Result".to_string(),
878            }],
879            is_error: false,
880        };
881        let json = serde_json::to_string(&result).unwrap();
882        assert!(json.contains("\"content\""));
883        assert!(json.contains("\"isError\":false"));
884    }
885
886    #[test]
887    fn test_call_tool_result_error_flag() {
888        let result = CallToolResult {
889            content: vec![ToolContent::Text {
890                text: "Error occurred".to_string(),
891            }],
892            is_error: true,
893        };
894        assert!(result.is_error);
895    }
896
897    #[test]
898    fn test_call_tool_result_default() {
899        let json = r#"{"content":[]}"#;
900        let result: CallToolResult = serde_json::from_str(json).unwrap();
901        assert!(!result.is_error);
902    }
903
904    #[test]
905    fn test_read_resource_params_serialization() {
906        let params = ReadResourceParams {
907            uri: "file:///test.txt".to_string(),
908        };
909        let json = serde_json::to_string(&params).unwrap();
910        assert!(json.contains("\"uri\":\"file:///test.txt\""));
911    }
912
913    #[test]
914    fn test_read_resource_result_serialization() {
915        let result = ReadResourceResult {
916            contents: vec![ResourceContent {
917                uri: "file:///test.txt".to_string(),
918                mime_type: Some("text/plain".to_string()),
919                text: Some("Hello".to_string()),
920                blob: None,
921            }],
922        };
923        let json = serde_json::to_string(&result).unwrap();
924        assert!(json.contains("\"contents\""));
925        assert!(json.contains("\"uri\""));
926    }
927
928    #[test]
929    fn test_list_tools_result_serialization() {
930        let result = ListToolsResult {
931            tools: vec![McpTool {
932                name: "tool1".to_string(),
933                description: None,
934                input_schema: serde_json::json!({"type": "object"}),
935            }],
936        };
937        let json = serde_json::to_string(&result).unwrap();
938        assert!(json.contains("\"tools\""));
939    }
940
941    #[test]
942    fn test_list_resources_result_serialization() {
943        let result = ListResourcesResult {
944            resources: vec![McpResource {
945                uri: "file:///test.txt".to_string(),
946                name: "test.txt".to_string(),
947                description: None,
948                mime_type: None,
949            }],
950        };
951        let json = serde_json::to_string(&result).unwrap();
952        assert!(json.contains("\"resources\""));
953    }
954
955    #[test]
956    fn test_server_capabilities_default() {
957        let caps = ServerCapabilities::default();
958        assert!(caps.tools.is_none());
959        assert!(caps.resources.is_none());
960        assert!(caps.prompts.is_none());
961        assert!(caps.logging.is_none());
962    }
963
964    #[test]
965    fn test_server_capabilities_all_fields() {
966        let caps = ServerCapabilities {
967            tools: Some(ToolsCapability { list_changed: true }),
968            resources: Some(ResourcesCapability {
969                subscribe: true,
970                list_changed: true,
971            }),
972            prompts: Some(PromptsCapability { list_changed: true }),
973            logging: Some(LoggingCapability {}),
974        };
975        assert!(caps.tools.is_some());
976        assert!(caps.resources.is_some());
977        assert!(caps.prompts.is_some());
978        assert!(caps.logging.is_some());
979    }
980
981    #[test]
982    fn test_client_capabilities_default() {
983        let caps = ClientCapabilities::default();
984        assert!(caps.roots.is_none());
985        assert!(caps.sampling.is_none());
986    }
987
988    #[test]
989    fn test_client_capabilities_all_fields() {
990        let caps = ClientCapabilities {
991            roots: Some(RootsCapability { list_changed: true }),
992            sampling: Some(SamplingCapability {}),
993        };
994        assert!(caps.roots.is_some());
995        assert!(caps.sampling.is_some());
996    }
997
998    #[test]
999    fn test_mcp_notification_tools_list_changed() {
1000        let notif = JsonRpcNotification::new("notifications/tools/list_changed", None);
1001        let mcp_notif = McpNotification::from_json_rpc(&notif);
1002        match mcp_notif {
1003            McpNotification::ToolsListChanged => {}
1004            _ => panic!("Expected ToolsListChanged"),
1005        }
1006    }
1007
1008    #[test]
1009    fn test_mcp_notification_resources_list_changed() {
1010        let notif = JsonRpcNotification::new("notifications/resources/list_changed", None);
1011        let mcp_notif = McpNotification::from_json_rpc(&notif);
1012        match mcp_notif {
1013            McpNotification::ResourcesListChanged => {}
1014            _ => panic!("Expected ResourcesListChanged"),
1015        }
1016    }
1017
1018    #[test]
1019    fn test_mcp_notification_prompts_list_changed() {
1020        let notif = JsonRpcNotification::new("notifications/prompts/list_changed", None);
1021        let mcp_notif = McpNotification::from_json_rpc(&notif);
1022        match mcp_notif {
1023            McpNotification::PromptsListChanged => {}
1024            _ => panic!("Expected PromptsListChanged"),
1025        }
1026    }
1027
1028    #[test]
1029    fn test_mcp_notification_progress() {
1030        let notif = JsonRpcNotification::new(
1031            "notifications/progress",
1032            Some(serde_json::json!({
1033                "progressToken": "token-123",
1034                "progress": 50.0,
1035                "total": 100.0
1036            })),
1037        );
1038        let mcp_notif = McpNotification::from_json_rpc(&notif);
1039        match mcp_notif {
1040            McpNotification::Progress {
1041                progress_token,
1042                progress,
1043                total,
1044            } => {
1045                assert_eq!(progress_token, "token-123");
1046                assert_eq!(progress, 50.0);
1047                assert_eq!(total, Some(100.0));
1048            }
1049            _ => panic!("Expected Progress"),
1050        }
1051    }
1052
1053    #[test]
1054    fn test_mcp_notification_log() {
1055        let notif = JsonRpcNotification::new(
1056            "notifications/message",
1057            Some(serde_json::json!({
1058                "level": "error",
1059                "logger": "test-logger",
1060                "data": {"message": "test"}
1061            })),
1062        );
1063        let mcp_notif = McpNotification::from_json_rpc(&notif);
1064        match mcp_notif {
1065            McpNotification::Log {
1066                level,
1067                logger,
1068                data,
1069            } => {
1070                assert_eq!(level, "error");
1071                assert_eq!(logger, Some("test-logger".to_string()));
1072                assert!(data.is_object());
1073            }
1074            _ => panic!("Expected Log"),
1075        }
1076    }
1077
1078    #[test]
1079    fn test_mcp_notification_log_edge_case_no_logger() {
1080        let notif = JsonRpcNotification::new(
1081            "notifications/message",
1082            Some(serde_json::json!({
1083                "level": "info",
1084                "data": "simple message"
1085            })),
1086        );
1087        let mcp_notif = McpNotification::from_json_rpc(&notif);
1088        match mcp_notif {
1089            McpNotification::Log { level, logger, .. } => {
1090                assert_eq!(level, "info");
1091                assert!(logger.is_none());
1092            }
1093            _ => panic!("Expected Log"),
1094        }
1095    }
1096
1097    #[test]
1098    fn test_mcp_notification_log_edge_case_default_level() {
1099        let notif = JsonRpcNotification::new(
1100            "notifications/message",
1101            Some(serde_json::json!({
1102                "data": "message"
1103            })),
1104        );
1105        let mcp_notif = McpNotification::from_json_rpc(&notif);
1106        match mcp_notif {
1107            McpNotification::Log { level, .. } => {
1108                assert_eq!(level, "info");
1109            }
1110            _ => panic!("Expected Log"),
1111        }
1112    }
1113
1114    #[test]
1115    fn test_mcp_notification_unknown() {
1116        let notif = JsonRpcNotification::new(
1117            "unknown/notification",
1118            Some(serde_json::json!({"key": "value"})),
1119        );
1120        let mcp_notif = McpNotification::from_json_rpc(&notif);
1121        match mcp_notif {
1122            McpNotification::Unknown { method, params } => {
1123                assert_eq!(method, "unknown/notification");
1124                assert!(params.is_some());
1125            }
1126            _ => panic!("Expected Unknown"),
1127        }
1128    }
1129
1130    #[test]
1131    fn test_tool_content_image() {
1132        let content = ToolContent::Image {
1133            data: "base64data".to_string(),
1134            mime_type: "image/png".to_string(),
1135        };
1136        let json = serde_json::to_string(&content).unwrap();
1137        assert!(json.contains("\"type\":\"image\""));
1138        assert!(json.contains("\"data\":\"base64data\""));
1139        assert!(json.contains("\"mimeType\":\"image/png\""));
1140    }
1141
1142    #[test]
1143    fn test_tool_content_resource() {
1144        let content = ToolContent::Resource {
1145            resource: ResourceContent {
1146                uri: "file:///test.txt".to_string(),
1147                mime_type: Some("text/plain".to_string()),
1148                text: Some("content".to_string()),
1149                blob: None,
1150            },
1151        };
1152        let json = serde_json::to_string(&content).unwrap();
1153        assert!(json.contains("\"type\":\"resource\""));
1154        assert!(json.contains("\"uri\":\"file:///test.txt\""));
1155    }
1156
1157    #[test]
1158    fn test_mcp_server_config_default() {
1159        let config = McpServerConfig {
1160            name: "test-server".to_string(),
1161            transport: McpTransportConfig::Stdio {
1162                command: "node".to_string(),
1163                args: vec!["server.js".to_string()],
1164            },
1165            enabled: true,
1166            env: HashMap::new(),
1167            oauth: None,
1168            tool_timeout_secs: 60,
1169        };
1170        assert!(config.enabled);
1171        assert!(config.oauth.is_none());
1172    }
1173
1174    #[test]
1175    fn test_mcp_server_config_with_env() {
1176        let mut env = HashMap::new();
1177        env.insert("API_KEY".to_string(), "secret".to_string());
1178        let config = McpServerConfig {
1179            name: "test-server".to_string(),
1180            transport: McpTransportConfig::Stdio {
1181                command: "node".to_string(),
1182                args: vec![],
1183            },
1184            enabled: true,
1185            env,
1186            oauth: None,
1187            tool_timeout_secs: 60,
1188        };
1189        assert!(config.env.contains_key("API_KEY"));
1190    }
1191
1192    #[test]
1193    fn test_mcp_server_config_with_oauth() {
1194        let config = McpServerConfig {
1195            name: "test-server".to_string(),
1196            transport: McpTransportConfig::Http {
1197                url: "https://api.example.com".to_string(),
1198                headers: HashMap::new(),
1199            },
1200            enabled: true,
1201            env: HashMap::new(),
1202            oauth: Some(OAuthConfig {
1203                auth_url: "https://auth.example.com".to_string(),
1204                token_url: "https://token.example.com".to_string(),
1205                client_id: "client-123".to_string(),
1206                client_secret: Some("secret".to_string()),
1207                scopes: vec!["read".to_string(), "write".to_string()],
1208                redirect_uri: "http://localhost:8080/callback".to_string(),
1209                access_token: None,
1210            }),
1211            tool_timeout_secs: 60,
1212        };
1213        assert!(config.oauth.is_some());
1214    }
1215
1216    #[test]
1217    fn test_mcp_transport_config_stdio_variant() {
1218        let transport = McpTransportConfig::Stdio {
1219            command: "python".to_string(),
1220            args: vec!["-m".to_string(), "server".to_string()],
1221        };
1222        match transport {
1223            McpTransportConfig::Stdio { command, args } => {
1224                assert_eq!(command, "python");
1225                assert_eq!(args.len(), 2);
1226            }
1227            _ => panic!("Expected Stdio"),
1228        }
1229    }
1230
1231    #[test]
1232    fn test_mcp_transport_config_http_variant() {
1233        let mut headers = HashMap::new();
1234        headers.insert("Authorization".to_string(), "Bearer token".to_string());
1235        let transport = McpTransportConfig::Http {
1236            url: "https://mcp.example.com".to_string(),
1237            headers,
1238        };
1239        match transport {
1240            McpTransportConfig::Http { url, headers } => {
1241                assert_eq!(url, "https://mcp.example.com");
1242                assert!(headers.contains_key("Authorization"));
1243            }
1244            _ => panic!("Expected Http"),
1245        }
1246    }
1247
1248    #[test]
1249    fn test_mcp_prompt_serialize() {
1250        let prompt = McpPrompt {
1251            name: "test_prompt".to_string(),
1252            description: Some("A test prompt".to_string()),
1253            arguments: Some(vec![PromptArgument {
1254                name: "arg1".to_string(),
1255                description: Some("First argument".to_string()),
1256                required: true,
1257            }]),
1258        };
1259        let json = serde_json::to_string(&prompt).unwrap();
1260        assert!(json.contains("\"name\":\"test_prompt\""));
1261        assert!(json.contains("\"arguments\""));
1262    }
1263
1264    #[test]
1265    fn test_prompt_argument_default() {
1266        let json = r#"{"name":"arg"}"#;
1267        let arg: PromptArgument = serde_json::from_str(json).unwrap();
1268        assert_eq!(arg.name, "arg");
1269        assert!(!arg.required);
1270    }
1271
1272    #[test]
1273    fn test_oauth_config_with_static_token() {
1274        let json = r#"{
1275            "auth_url": "https://auth.example.com/authorize",
1276            "token_url": "https://auth.example.com/token",
1277            "client_id": "my-client",
1278            "scopes": ["read", "write"],
1279            "redirect_uri": "http://localhost/callback",
1280            "access_token": "static-token-abc123"
1281        }"#;
1282        let config: OAuthConfig = serde_json::from_str(json).unwrap();
1283        assert_eq!(config.client_id, "my-client");
1284        assert_eq!(config.access_token, Some("static-token-abc123".to_string()));
1285    }
1286
1287    #[test]
1288    fn test_oauth_config_without_static_token() {
1289        let json = r#"{
1290            "auth_url": "https://auth.example.com/authorize",
1291            "token_url": "https://auth.example.com/token",
1292            "client_id": "my-client",
1293            "scopes": [],
1294            "redirect_uri": "http://localhost/callback"
1295        }"#;
1296        let config: OAuthConfig = serde_json::from_str(json).unwrap();
1297        assert!(config.access_token.is_none());
1298    }
1299
1300    #[test]
1301    fn test_oauth_config_static_token_not_serialized_when_absent() {
1302        let config = OAuthConfig {
1303            auth_url: "https://example.com/auth".to_string(),
1304            token_url: "https://example.com/token".to_string(),
1305            client_id: "client".to_string(),
1306            client_secret: None,
1307            scopes: vec![],
1308            redirect_uri: "http://localhost/cb".to_string(),
1309            access_token: None,
1310        };
1311        let json = serde_json::to_string(&config).unwrap();
1312        assert!(!json.contains("access_token"));
1313    }
1314}
a3s_code_core/mcp/protocol.rs

a3s_code_core/mcp/
protocol.rs
