litellm_rs/core/mcp/
protocol.rs

1//! MCP Protocol Types
2//!
3//! JSON-RPC 2.0 message types for MCP communication.
4
5use super::error::McpError;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8
9/// JSON-RPC 2.0 version constant
10pub const JSONRPC_VERSION: &str = "2.0";
11
12/// MCP protocol version supported by this client
13pub const SUPPORTED_PROTOCOL_VERSION: &str = "2024-11-05";
14
15/// JSON-RPC 2.0 Request
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct JsonRpcRequest {
18    /// JSON-RPC version (must be "2.0")
19    pub jsonrpc: String,
20
21    /// Request method
22    pub method: String,
23
24    /// Request parameters
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub params: Option<Value>,
27
28    /// Request ID (for matching responses)
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub id: Option<Value>,
31}
32
33impl JsonRpcRequest {
34    /// Create a new JSON-RPC request
35    pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
36        Self {
37            jsonrpc: JSONRPC_VERSION.to_string(),
38            method: method.into(),
39            params,
40            id: Some(Value::Number(1.into())),
41        }
42    }
43
44    /// Create a request with a specific ID
45    pub fn with_id(mut self, id: impl Into<Value>) -> Self {
46        self.id = Some(id.into());
47        self
48    }
49
50    /// Create a notification (no ID, no response expected)
51    pub fn notification(method: impl Into<String>, params: Option<Value>) -> Self {
52        Self {
53            jsonrpc: JSONRPC_VERSION.to_string(),
54            method: method.into(),
55            params,
56            id: None,
57        }
58    }
59
60    /// Check if this is a notification
61    pub fn is_notification(&self) -> bool {
62        self.id.is_none()
63    }
64}
65
66/// JSON-RPC 2.0 Response
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct JsonRpcResponse {
69    /// JSON-RPC version (must be "2.0")
70    pub jsonrpc: String,
71
72    /// Response result (mutually exclusive with error)
73    #[serde(skip_serializing_if = "Option::is_none")]
74    pub result: Option<Value>,
75
76    /// Response error (mutually exclusive with result)
77    #[serde(skip_serializing_if = "Option::is_none")]
78    pub error: Option<JsonRpcError>,
79
80    /// Request ID this is responding to
81    #[serde(skip_serializing_if = "Option::is_none")]
82    pub id: Option<Value>,
83}
84
85impl JsonRpcResponse {
86    /// Create a success response
87    pub fn success(result: Value, id: Value) -> Self {
88        Self {
89            jsonrpc: JSONRPC_VERSION.to_string(),
90            result: Some(result),
91            error: None,
92            id: Some(id),
93        }
94    }
95
96    /// Create an error response
97    pub fn error(error: JsonRpcError, id: Option<Value>) -> Self {
98        Self {
99            jsonrpc: JSONRPC_VERSION.to_string(),
100            result: None,
101            error: Some(error),
102            id,
103        }
104    }
105
106    /// Check if this response is an error
107    pub fn is_error(&self) -> bool {
108        self.error.is_some()
109    }
110
111    /// Check if this response is successful
112    pub fn is_success(&self) -> bool {
113        self.result.is_some()
114    }
115}
116
117/// JSON-RPC 2.0 Error
118#[derive(Debug, Clone, Serialize, Deserialize)]
119pub struct JsonRpcError {
120    /// Error code
121    pub code: i32,
122
123    /// Error message
124    pub message: String,
125
126    /// Additional error data
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub data: Option<Value>,
129}
130
131impl JsonRpcError {
132    /// Create a new error
133    pub fn new(code: i32, message: impl Into<String>) -> Self {
134        Self {
135            code,
136            message: message.into(),
137            data: None,
138        }
139    }
140
141    /// Create an error with additional data
142    pub fn with_data(mut self, data: Value) -> Self {
143        self.data = Some(data);
144        self
145    }
146
147    /// Parse error (-32700)
148    pub fn parse_error() -> Self {
149        Self::new(-32700, "Parse error")
150    }
151
152    /// Invalid request (-32600)
153    pub fn invalid_request() -> Self {
154        Self::new(-32600, "Invalid Request")
155    }
156
157    /// Method not found (-32601)
158    pub fn method_not_found() -> Self {
159        Self::new(-32601, "Method not found")
160    }
161
162    /// Invalid params (-32602)
163    pub fn invalid_params() -> Self {
164        Self::new(-32602, "Invalid params")
165    }
166
167    /// Internal error (-32603)
168    pub fn internal_error() -> Self {
169        Self::new(-32603, "Internal error")
170    }
171
172    /// Server error (-32000 to -32099)
173    pub fn server_error(code: i32, message: impl Into<String>) -> Self {
174        let code = code.clamp(-32099, -32000);
175        Self::new(code, message)
176    }
177
178    /// Build protocol error payload from typed MCP errors.
179    pub fn from_mcp_error(error: &McpError) -> Self {
180        use crate::utils::error::canonical::CanonicalError;
181
182        let code = match error {
183            McpError::ServerNotFound { .. } | McpError::ToolNotFound { .. } => -32001,
184            McpError::AuthenticationError { .. } | McpError::AuthorizationError { .. } => -32004,
185            McpError::RateLimitExceeded { .. } => -32029,
186            McpError::Timeout { .. } => -32008,
187            McpError::ConnectionError { .. } | McpError::TransportError { .. } => -32010,
188            McpError::InvalidUrl { .. } | McpError::ProtocolError { .. } => -32600,
189            McpError::ConfigurationError { .. } => -32602,
190            McpError::ToolExecutionError { .. } | McpError::SerializationError { .. } => -32603,
191            McpError::ServerAlreadyExists { .. } => -32009,
192            McpError::ToolDefinitionChanged { .. } => -32009,
193            McpError::ValidationError { .. } => -32602,
194        };
195
196        let mut rpc_error = Self::new(code, error.to_string());
197        rpc_error.data = Some(serde_json::json!({
198            "canonical_code": error.canonical_code().as_str(),
199            "retryable": error.canonical_retryable(),
200        }));
201        rpc_error
202    }
203}
204
205impl std::fmt::Display for JsonRpcError {
206    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
207        write!(f, "[{}] {}", self.code, self.message)
208    }
209}
210
211impl std::error::Error for JsonRpcError {}
212
213/// MCP-specific message types
214#[derive(Debug, Clone, Serialize, Deserialize)]
215#[serde(untagged)]
216pub enum McpMessage {
217    /// JSON-RPC request
218    Request(JsonRpcRequest),
219
220    /// JSON-RPC response
221    Response(JsonRpcResponse),
222
223    /// Batch of requests/responses
224    Batch(Vec<McpMessage>),
225}
226
227impl McpMessage {
228    /// Try to parse from JSON string
229    pub fn parse(s: &str) -> Result<Self, serde_json::Error> {
230        serde_json::from_str(s)
231    }
232
233    /// Convert to JSON string
234    pub fn to_json(&self) -> Result<String, serde_json::Error> {
235        serde_json::to_string(self)
236    }
237}
238
239/// MCP Method names as constants
240pub mod methods {
241    /// Initialize the connection
242    pub const INITIALIZE: &str = "initialize";
243
244    /// Notify the server that initialization completed
245    pub const INITIALIZED: &str = "notifications/initialized";
246
247    /// List available tools
248    pub const LIST_TOOLS: &str = "tools/list";
249
250    /// Call a tool
251    pub const CALL_TOOL: &str = "tools/call";
252
253    /// List available resources
254    pub const LIST_RESOURCES: &str = "resources/list";
255
256    /// Read a resource
257    pub const READ_RESOURCE: &str = "resources/read";
258
259    /// List available prompts
260    pub const LIST_PROMPTS: &str = "prompts/list";
261
262    /// Get a prompt
263    pub const GET_PROMPT: &str = "prompts/get";
264
265    /// Complete (text completion)
266    pub const COMPLETE: &str = "completion/complete";
267
268    /// Set logging level
269    pub const SET_LOGGING_LEVEL: &str = "logging/setLevel";
270
271    /// Ping (health check)
272    pub const PING: &str = "ping";
273}
274
275/// MCP Capability types
276#[derive(Debug, Clone, Serialize, Deserialize, Default)]
277pub struct McpCapabilities {
278    /// Supported tools capability
279    #[serde(skip_serializing_if = "Option::is_none")]
280    pub tools: Option<ToolsCapability>,
281
282    /// Supported resources capability
283    #[serde(skip_serializing_if = "Option::is_none")]
284    pub resources: Option<ResourcesCapability>,
285
286    /// Supported prompts capability
287    #[serde(skip_serializing_if = "Option::is_none")]
288    pub prompts: Option<PromptsCapability>,
289
290    /// Supported logging capability
291    #[serde(skip_serializing_if = "Option::is_none")]
292    pub logging: Option<LoggingCapability>,
293}
294
295/// Tools capability
296#[derive(Debug, Clone, Serialize, Deserialize, Default)]
297#[serde(rename_all = "camelCase")]
298pub struct ToolsCapability {
299    /// Whether list_changed notifications are supported
300    #[serde(default)]
301    pub list_changed: bool,
302}
303
304/// Resources capability
305#[derive(Debug, Clone, Serialize, Deserialize, Default)]
306#[serde(rename_all = "camelCase")]
307pub struct ResourcesCapability {
308    /// Whether subscribe is supported
309    #[serde(default)]
310    pub subscribe: bool,
311
312    /// Whether list_changed notifications are supported
313    #[serde(default)]
314    pub list_changed: bool,
315}
316
317/// Prompts capability
318#[derive(Debug, Clone, Serialize, Deserialize, Default)]
319#[serde(rename_all = "camelCase")]
320pub struct PromptsCapability {
321    /// Whether list_changed notifications are supported
322    #[serde(default)]
323    pub list_changed: bool,
324}
325
326/// Logging capability
327#[derive(Debug, Clone, Serialize, Deserialize, Default)]
328pub struct LoggingCapability {}
329
330/// Initialize request params
331#[derive(Debug, Clone, Serialize, Deserialize)]
332#[serde(rename_all = "camelCase")]
333pub struct InitializeParams {
334    /// Protocol version
335    pub protocol_version: String,
336
337    /// Client capabilities
338    pub capabilities: McpCapabilities,
339
340    /// Client info
341    pub client_info: ClientInfo,
342}
343
344/// Initialize response result
345#[derive(Debug, Clone, Serialize, Deserialize)]
346#[serde(rename_all = "camelCase")]
347pub struct InitializeResult {
348    /// Protocol version selected by the server
349    pub protocol_version: String,
350
351    /// Server capabilities
352    pub capabilities: McpCapabilities,
353
354    /// Server implementation information
355    pub server_info: ClientInfo,
356}
357
358/// Client info for initialization
359#[derive(Debug, Clone, Serialize, Deserialize)]
360pub struct ClientInfo {
361    /// Client name
362    pub name: String,
363
364    /// Client version
365    pub version: String,
366}
367
368impl Default for ClientInfo {
369    fn default() -> Self {
370        Self {
371            name: "litellm-rs".to_string(),
372            version: env!("CARGO_PKG_VERSION").to_string(),
373        }
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380
381    #[test]
382    fn test_jsonrpc_request_new() {
383        let req = JsonRpcRequest::new("test_method", Some(serde_json::json!({"key": "value"})));
384        assert_eq!(req.jsonrpc, "2.0");
385        assert_eq!(req.method, "test_method");
386        assert!(req.params.is_some());
387        assert!(req.id.is_some());
388    }
389
390    #[test]
391    fn test_jsonrpc_request_notification() {
392        let req = JsonRpcRequest::notification("test_method", None);
393        assert!(req.is_notification());
394        assert!(req.id.is_none());
395    }
396
397    #[test]
398    fn test_jsonrpc_response_success() {
399        let resp =
400            JsonRpcResponse::success(serde_json::json!({"result": "ok"}), Value::Number(1.into()));
401        assert!(resp.is_success());
402        assert!(!resp.is_error());
403    }
404
405    #[test]
406    fn test_jsonrpc_response_error() {
407        let resp = JsonRpcResponse::error(
408            JsonRpcError::method_not_found(),
409            Some(Value::Number(1.into())),
410        );
411        assert!(resp.is_error());
412        assert!(!resp.is_success());
413    }
414
415    #[test]
416    fn test_jsonrpc_error_codes() {
417        assert_eq!(JsonRpcError::parse_error().code, -32700);
418        assert_eq!(JsonRpcError::invalid_request().code, -32600);
419        assert_eq!(JsonRpcError::method_not_found().code, -32601);
420        assert_eq!(JsonRpcError::invalid_params().code, -32602);
421        assert_eq!(JsonRpcError::internal_error().code, -32603);
422    }
423
424    #[test]
425    fn test_jsonrpc_error_server_error_clamping() {
426        let err = JsonRpcError::server_error(-99999, "test");
427        assert!(err.code >= -32099 && err.code <= -32000);
428    }
429
430    #[test]
431    fn test_jsonrpc_error_from_mcp_error_includes_canonical_data() {
432        let error = McpError::RateLimitExceeded {
433            server_name: "github".to_string(),
434            retry_after_ms: Some(1000),
435        };
436
437        let rpc_error = JsonRpcError::from_mcp_error(&error);
438        assert_eq!(rpc_error.code, -32029);
439        let data = rpc_error.data.expect("canonical data should exist");
440        assert_eq!(data["canonical_code"], "RATE_LIMITED");
441        assert_eq!(data["retryable"], true);
442    }
443
444    #[test]
445    fn test_jsonrpc_error_display() {
446        let err = JsonRpcError::method_not_found();
447        assert!(err.to_string().contains("-32601"));
448        assert!(err.to_string().contains("Method not found"));
449    }
450
451    #[test]
452    fn test_mcp_message_parse() {
453        let json = r#"{"jsonrpc":"2.0","method":"test","id":1}"#;
454        let msg = McpMessage::parse(json).unwrap();
455        match msg {
456            McpMessage::Request(req) => {
457                assert_eq!(req.method, "test");
458            }
459            _ => panic!("Expected request"),
460        }
461    }
462
463    #[test]
464    fn test_request_serialization() {
465        let req = JsonRpcRequest::new("tools/list", None);
466        let json = serde_json::to_string(&req).unwrap();
467        assert!(json.contains("tools/list"));
468        assert!(json.contains("2.0"));
469    }
470
471    #[test]
472    fn test_client_info_default() {
473        let info = ClientInfo::default();
474        assert_eq!(info.name, "litellm-rs");
475        assert!(!info.version.is_empty());
476    }
477
478    #[test]
479    fn test_capabilities_default() {
480        let caps = McpCapabilities::default();
481        assert!(caps.tools.is_none());
482        assert!(caps.resources.is_none());
483    }
484
485    #[test]
486    fn test_initialize_params_use_mcp_camel_case_fields() {
487        let params = InitializeParams {
488            protocol_version: SUPPORTED_PROTOCOL_VERSION.to_string(),
489            capabilities: McpCapabilities {
490                tools: Some(ToolsCapability { list_changed: true }),
491                resources: Some(ResourcesCapability {
492                    subscribe: true,
493                    list_changed: true,
494                }),
495                prompts: Some(PromptsCapability { list_changed: true }),
496                logging: None,
497            },
498            client_info: ClientInfo::default(),
499        };
500
501        let value = serde_json::to_value(params).expect("initialize params should serialize");
502        assert_eq!(value["protocolVersion"], SUPPORTED_PROTOCOL_VERSION);
503        assert!(value.get("clientInfo").is_some());
504        assert!(value.get("protocol_version").is_none());
505        assert!(value.get("client_info").is_none());
506        assert_eq!(value["capabilities"]["tools"]["listChanged"], true);
507        assert_eq!(value["capabilities"]["resources"]["listChanged"], true);
508        assert_eq!(value["capabilities"]["prompts"]["listChanged"], true);
509        assert!(value["capabilities"]["tools"].get("list_changed").is_none());
510    }
511
512    #[test]
513    fn test_initialize_result_requires_all_mcp_fields() {
514        let valid = serde_json::json!({
515            "protocolVersion": SUPPORTED_PROTOCOL_VERSION,
516            "capabilities": {"tools": {"listChanged": true}},
517            "serverInfo": {"name": "test-server", "version": "1.0.0"}
518        });
519
520        let result: InitializeResult =
521            serde_json::from_value(valid.clone()).expect("valid result should parse");
522        assert_eq!(result.protocol_version, SUPPORTED_PROTOCOL_VERSION);
523        assert!(
524            result
525                .capabilities
526                .tools
527                .expect("tools capability")
528                .list_changed
529        );
530
531        for required_field in ["protocolVersion", "capabilities", "serverInfo"] {
532            let mut invalid = valid.clone();
533            invalid
534                .as_object_mut()
535                .expect("fixture should be an object")
536                .remove(required_field);
537            assert!(
538                serde_json::from_value::<InitializeResult>(invalid).is_err(),
539                "missing {required_field} should fail"
540            );
541        }
542    }
543
544    #[test]
545    fn test_method_constants() {
546        assert_eq!(methods::INITIALIZE, "initialize");
547        assert_eq!(methods::INITIALIZED, "notifications/initialized");
548        assert_eq!(methods::LIST_TOOLS, "tools/list");
549        assert_eq!(methods::CALL_TOOL, "tools/call");
550    }
551}
litellm_rs/core/mcp/protocol.rs

litellm_rs/core/mcp/
protocol.rs
