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

a3s_code_core/mcp/
protocol.rs
