ultrafast_mcp_core/protocol/
jsonrpc.rs

1use crate::protocol::constants::{
2    JSONRPC_VERSION, MAX_REQUEST_ID_LENGTH, MAX_REQUEST_ID_NUMBER, MIN_REQUEST_ID_NUMBER,
3};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::borrow::Cow;
7use std::collections::HashMap;
8
9/// Standard JSON-RPC 2.0 error codes
10pub mod error_codes {
11    /// Parse error (invalid JSON)
12    pub const PARSE_ERROR: i32 = -32700;
13    /// Invalid request (malformed request)
14    pub const INVALID_REQUEST: i32 = -32600;
15    /// Method not found
16    pub const METHOD_NOT_FOUND: i32 = -32601;
17    /// Invalid parameters
18    pub const INVALID_PARAMS: i32 = -32602;
19    /// Internal error
20    pub const INTERNAL_ERROR: i32 = -32603;
21    /// Server error (implementation-defined)
22    pub const SERVER_ERROR_START: i32 = -32000;
23    pub const SERVER_ERROR_END: i32 = -32099;
24}
25
26/// MCP-specific error codes
27pub mod mcp_error_codes {
28    /// Initialization failed
29    pub const INITIALIZATION_FAILED: i32 = -32000;
30    /// Capability not supported
31    pub const CAPABILITY_NOT_SUPPORTED: i32 = -32001;
32    /// Resource not found
33    pub const RESOURCE_NOT_FOUND: i32 = -32002;
34    /// Tool execution error
35    pub const TOOL_EXECUTION_ERROR: i32 = -32003;
36    /// Invalid URI
37    pub const INVALID_URI: i32 = -32004;
38    /// Access denied
39    pub const ACCESS_DENIED: i32 = -32005;
40    /// Request timeout
41    pub const REQUEST_TIMEOUT: i32 = -32006;
42    /// Protocol version not supported
43    pub const PROTOCOL_VERSION_NOT_SUPPORTED: i32 = -32007;
44}
45
46/// JSON-RPC 2.0 request ID can be string or number
47#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
48#[serde(untagged)]
49pub enum RequestId {
50    /// String-based request ID
51    String(String),
52    /// Number-based request ID
53    Number(i64),
54}
55
56impl RequestId {
57    /// Create a new string-based request ID
58    pub fn string(s: impl Into<String>) -> Self {
59        Self::String(s.into())
60    }
61
62    /// Create a new number-based request ID
63    pub fn number(n: i64) -> Self {
64        Self::Number(n)
65    }
66
67    /// Validate the request ID
68    pub fn validate(&self) -> Result<(), crate::error::ProtocolError> {
69        match self {
70            RequestId::String(s) => {
71                if s.is_empty() {
72                    return Err(crate::error::ProtocolError::InvalidRequestId(
73                        "Request ID string cannot be empty".to_string(),
74                    ));
75                }
76                if s.len() > MAX_REQUEST_ID_LENGTH {
77                    return Err(crate::error::ProtocolError::InvalidRequestId(format!(
78                        "Request ID string too long (max {MAX_REQUEST_ID_LENGTH} characters)"
79                    )));
80                }
81            }
82            RequestId::Number(n) => {
83                if *n < MIN_REQUEST_ID_NUMBER || *n > MAX_REQUEST_ID_NUMBER {
84                    return Err(crate::error::ProtocolError::InvalidRequestId(format!(
85                        "Request ID number out of range ({MIN_REQUEST_ID_NUMBER} to {MAX_REQUEST_ID_NUMBER})"
86                    )));
87                }
88            }
89        }
90        Ok(())
91    }
92}
93
94impl std::fmt::Display for RequestId {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        match self {
97            RequestId::String(s) => write!(f, "{s}"),
98            RequestId::Number(n) => write!(f, "{n}"),
99        }
100    }
101}
102
103impl From<String> for RequestId {
104    fn from(s: String) -> Self {
105        RequestId::String(s)
106    }
107}
108
109impl From<i64> for RequestId {
110    fn from(n: i64) -> Self {
111        RequestId::Number(n)
112    }
113}
114
115impl From<u64> for RequestId {
116    fn from(n: u64) -> Self {
117        RequestId::Number(n as i64)
118    }
119}
120
121/// JSON-RPC 2.0 Request with optimized memory usage
122#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
123pub struct JsonRpcRequest {
124    /// JSON-RPC version string (always "2.0")
125    #[serde(rename = "jsonrpc")]
126    pub jsonrpc: Cow<'static, str>,
127    /// The name of the method to be invoked
128    pub method: String,
129    /// Optional parameters for the method
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub params: Option<Value>,
132    /// Optional request ID (null for notifications)
133    #[serde(skip_serializing_if = "Option::is_none")]
134    pub id: Option<RequestId>,
135    /// Additional metadata fields (flattened)
136    #[serde(flatten)]
137    pub meta: HashMap<String, Value>,
138}
139
140impl JsonRpcRequest {
141    pub fn new(method: String, params: Option<Value>, id: Option<RequestId>) -> Self {
142        Self {
143            jsonrpc: Cow::Borrowed(JSONRPC_VERSION),
144            method,
145            params,
146            id,
147            meta: HashMap::new(),
148        }
149    }
150
151    pub fn notification(method: String, params: Option<Value>) -> Self {
152        Self::new(method, params, None)
153    }
154
155    pub fn is_notification(&self) -> bool {
156        self.id.is_none()
157    }
158
159    pub fn with_meta(mut self, key: String, value: Value) -> Self {
160        self.meta.insert(key, value);
161        self
162    }
163}
164
165/// JSON-RPC 2.0 Response with optimized memory usage
166#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
167pub struct JsonRpcResponse {
168    #[serde(rename = "jsonrpc")]
169    pub jsonrpc: Cow<'static, str>,
170    #[serde(skip_serializing_if = "Option::is_none")]
171    pub result: Option<Value>,
172    #[serde(skip_serializing_if = "Option::is_none")]
173    pub error: Option<JsonRpcError>,
174    pub id: Option<RequestId>,
175    #[serde(flatten)]
176    pub meta: HashMap<String, Value>,
177}
178
179impl JsonRpcResponse {
180    pub fn success(result: Value, id: Option<RequestId>) -> Self {
181        Self {
182            jsonrpc: Cow::Borrowed(JSONRPC_VERSION),
183            result: Some(result),
184            error: None,
185            id,
186            meta: HashMap::new(),
187        }
188    }
189
190    pub fn error(error: JsonRpcError, id: Option<RequestId>) -> Self {
191        Self {
192            jsonrpc: Cow::Borrowed(JSONRPC_VERSION),
193            result: None,
194            error: Some(error),
195            id,
196            meta: HashMap::new(),
197        }
198    }
199
200    pub fn with_meta(mut self, key: String, value: Value) -> Self {
201        self.meta.insert(key, value);
202        self
203    }
204}
205
206/// JSON-RPC 2.0 Error with improved error handling and optimized memory usage
207#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
208pub struct JsonRpcError {
209    pub code: i32,
210    pub message: String,
211    #[serde(skip_serializing_if = "Option::is_none")]
212    pub data: Option<Value>,
213}
214
215impl JsonRpcError {
216    pub fn new(code: i32, message: String) -> Self {
217        Self {
218            code,
219            message,
220            data: None,
221        }
222    }
223
224    pub fn with_data(mut self, data: Value) -> Self {
225        self.data = Some(data);
226        self
227    }
228
229    pub fn parse_error(message: Option<String>) -> Self {
230        let msg = message.unwrap_or_else(|| "Parse error".to_string());
231        Self::new(error_codes::PARSE_ERROR, msg)
232    }
233
234    pub fn invalid_request(message: Option<String>) -> Self {
235        let msg = message.unwrap_or_else(|| "Invalid request".to_string());
236        Self::new(error_codes::INVALID_REQUEST, msg)
237    }
238
239    pub fn method_not_found(method: String) -> Self {
240        Self::new(
241            error_codes::METHOD_NOT_FOUND,
242            format!("Method not found: {method}"),
243        )
244    }
245
246    pub fn invalid_params(message: Option<String>) -> Self {
247        let msg = message.unwrap_or_else(|| "Invalid parameters".to_string());
248        Self::new(error_codes::INVALID_PARAMS, msg)
249    }
250
251    pub fn internal_error(message: Option<String>) -> Self {
252        let msg = message.unwrap_or_else(|| "Internal error".to_string());
253        Self::new(error_codes::INTERNAL_ERROR, msg)
254    }
255
256    pub fn initialization_failed(message: String) -> Self {
257        Self::new(mcp_error_codes::INITIALIZATION_FAILED, message)
258    }
259
260    pub fn capability_not_supported(capability: String) -> Self {
261        Self::new(
262            mcp_error_codes::CAPABILITY_NOT_SUPPORTED,
263            format!("Capability not supported: {capability}"),
264        )
265    }
266
267    pub fn resource_not_found(uri: String) -> Self {
268        Self::new(
269            mcp_error_codes::RESOURCE_NOT_FOUND,
270            format!("Resource not found: {uri}"),
271        )
272    }
273
274    pub fn tool_execution_error(tool_name: String, error: String) -> Self {
275        Self::new(
276            mcp_error_codes::TOOL_EXECUTION_ERROR,
277            format!("Tool execution error for '{tool_name}': {error}"),
278        )
279    }
280
281    pub fn invalid_uri(uri: String) -> Self {
282        Self::new(mcp_error_codes::INVALID_URI, format!("Invalid URI: {uri}"))
283    }
284
285    pub fn access_denied(resource: String) -> Self {
286        Self::new(
287            mcp_error_codes::ACCESS_DENIED,
288            format!("Access denied to resource: {resource}"),
289        )
290    }
291
292    pub fn request_timeout() -> Self {
293        Self::new(
294            mcp_error_codes::REQUEST_TIMEOUT,
295            "Request timeout".to_string(),
296        )
297    }
298
299    pub fn protocol_version_not_supported(version: String) -> Self {
300        Self::new(
301            mcp_error_codes::PROTOCOL_VERSION_NOT_SUPPORTED,
302            format!("Protocol version not supported: {version}"),
303        )
304    }
305}
306
307/// JSON-RPC 2.0 Message
308#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
309#[serde(untagged)]
310pub enum JsonRpcMessage {
311    Request(JsonRpcRequest),
312    Response(JsonRpcResponse),
313    Notification(JsonRpcRequest), // Notification is a request without an ID
314}
315
316impl JsonRpcMessage {
317    pub fn get_id(&self) -> Option<&RequestId> {
318        match self {
319            JsonRpcMessage::Request(req) => req.id.as_ref(),
320            JsonRpcMessage::Response(resp) => resp.id.as_ref(),
321            JsonRpcMessage::Notification(_) => None,
322        }
323    }
324
325    pub fn is_notification(&self) -> bool {
326        matches!(self, JsonRpcMessage::Notification(_))
327    }
328}
329
330/// Validate JSON-RPC message format with improved error handling
331pub fn validate_jsonrpc_message(
332    message: &JsonRpcMessage,
333) -> Result<(), crate::error::ProtocolError> {
334    match message {
335        JsonRpcMessage::Request(request) => {
336            if request.jsonrpc != JSONRPC_VERSION {
337                return Err(crate::error::ProtocolError::InvalidVersion(format!(
338                    "Expected JSON-RPC version {}, got {}",
339                    JSONRPC_VERSION, request.jsonrpc
340                )));
341            }
342
343            if request.method.is_empty() {
344                return Err(crate::error::ProtocolError::InvalidRequest(
345                    "Method name cannot be empty".to_string(),
346                ));
347            }
348
349            if let Some(ref id) = request.id {
350                id.validate()?;
351            }
352        }
353        JsonRpcMessage::Response(response) => {
354            if response.jsonrpc != JSONRPC_VERSION {
355                return Err(crate::error::ProtocolError::InvalidVersion(format!(
356                    "Expected JSON-RPC version {}, got {}",
357                    JSONRPC_VERSION, response.jsonrpc
358                )));
359            }
360
361            if response.result.is_some() && response.error.is_some() {
362                return Err(crate::error::ProtocolError::InvalidResponse(
363                    "Response cannot have both result and error".to_string(),
364                ));
365            }
366
367            if response.result.is_none() && response.error.is_none() {
368                return Err(crate::error::ProtocolError::InvalidResponse(
369                    "Response must have either result or error".to_string(),
370                ));
371            }
372
373            if let Some(ref id) = response.id {
374                id.validate()?;
375            }
376        }
377        JsonRpcMessage::Notification(notification) => {
378            if notification.jsonrpc != JSONRPC_VERSION {
379                return Err(crate::error::ProtocolError::InvalidVersion(format!(
380                    "Expected JSON-RPC version {}, got {}",
381                    JSONRPC_VERSION, notification.jsonrpc
382                )));
383            }
384
385            if notification.method.is_empty() {
386                return Err(crate::error::ProtocolError::InvalidRequest(
387                    "Method name cannot be empty".to_string(),
388                ));
389            }
390
391            if notification.id.is_some() {
392                return Err(crate::error::ProtocolError::InvalidRequest(
393                    "Notification cannot have an ID".to_string(),
394                ));
395            }
396        }
397    }
398    Ok(())
399}
400
401#[cfg(test)]
402mod tests {
403    use super::*;
404
405    #[test]
406    fn test_request_serialization() {
407        let request = JsonRpcRequest::new(
408            "test_method".to_string(),
409            Some(serde_json::json!({"param": "value"})),
410            Some(RequestId::number(1)),
411        );
412
413        let serialized = serde_json::to_string(&request).expect("Failed to serialize request");
414        let deserialized: JsonRpcRequest =
415            serde_json::from_str(&serialized).expect("Failed to deserialize request");
416
417        assert_eq!(request, deserialized);
418    }
419
420    #[test]
421    fn test_notification() {
422        let notification = JsonRpcRequest::notification(
423            "test_notification".to_string(),
424            Some(serde_json::json!({"data": "value"})),
425        );
426
427        assert!(notification.is_notification());
428        assert_eq!(notification.id, None);
429    }
430
431    #[test]
432    fn test_response_success() {
433        let response = JsonRpcResponse::success(
434            serde_json::json!({"result": "success"}),
435            Some(RequestId::number(1)),
436        );
437
438        assert!(response.result.is_some());
439        assert!(response.error.is_none());
440        assert_eq!(response.id, Some(RequestId::number(1)));
441    }
442
443    #[test]
444    fn test_response_error() {
445        let error = JsonRpcError::method_not_found("unknown_method".to_string());
446        let response = JsonRpcResponse::error(error, Some(RequestId::number(1)));
447
448        assert!(response.result.is_none());
449        assert!(response.error.is_some());
450        assert_eq!(response.id, Some(RequestId::number(1)));
451    }
452
453    #[test]
454    fn test_request_id_validation() {
455        // Valid string ID
456        let valid_string = RequestId::string("valid_id");
457        assert!(valid_string.validate().is_ok());
458
459        // Valid number ID
460        let valid_number = RequestId::number(123);
461        assert!(valid_number.validate().is_ok());
462
463        // Invalid empty string
464        let invalid_string = RequestId::string("");
465        assert!(invalid_string.validate().is_err());
466
467        // Invalid number out of range
468        let invalid_number = RequestId::number(9999999999);
469        assert!(invalid_number.validate().is_err());
470    }
471
472    #[test]
473    fn test_request_id_to_string() {
474        let string_id = RequestId::string("test_id");
475        assert_eq!(string_id.to_string(), "test_id");
476
477        let number_id = RequestId::number(123);
478        assert_eq!(number_id.to_string(), "123");
479    }
480
481    #[test]
482    fn test_jsonrpc_message_validation() {
483        // Valid request
484        let valid_request = JsonRpcMessage::Request(JsonRpcRequest::new(
485            "test_method".to_string(),
486            None,
487            Some(RequestId::number(1)),
488        ));
489        assert!(validate_jsonrpc_message(&valid_request).is_ok());
490
491        // Invalid request with empty method
492        let invalid_request = JsonRpcMessage::Request(JsonRpcRequest::new(
493            "".to_string(),
494            None,
495            Some(RequestId::number(1)),
496        ));
497        assert!(validate_jsonrpc_message(&invalid_request).is_err());
498    }
499}
ultrafast_mcp_core/protocol/jsonrpc.rs

ultrafast_mcp_core/protocol/
jsonrpc.rs
