turul_mcp_json_rpc_server/
dispatch.rs

1use serde_json::Value;
2
3use crate::{
4    error::JsonRpcError,
5    notification::JsonRpcNotification,
6    request::JsonRpcRequest,
7    response::JsonRpcResponse,
8    types::RequestId,
9};
10
11/// Enum representing different types of JSON-RPC messages
12#[derive(Debug, Clone)]
13pub enum JsonRpcMessage {
14    Request(JsonRpcRequest),
15    Notification(JsonRpcNotification),
16}
17
18/// Result of parsing and processing a JSON-RPC message
19#[derive(Debug, Clone)]
20pub enum JsonRpcMessageResult {
21    /// A response to a request
22    Response(JsonRpcResponse),
23    /// An error response
24    Error(JsonRpcError),
25    /// No response needed (for notifications)
26    NoResponse,
27}
28
29impl JsonRpcMessageResult {
30    /// Convert to JSON string if there's a response to send
31    pub fn to_json_string(&self) -> Option<String> {
32        match self {
33            JsonRpcMessageResult::Response(response) => {
34                serde_json::to_string(response).ok()
35            }
36            JsonRpcMessageResult::Error(error) => {
37                serde_json::to_string(error).ok()
38            }
39            JsonRpcMessageResult::NoResponse => None,
40        }
41    }
42
43    /// Check if this result represents an error
44    pub fn is_error(&self) -> bool {
45        matches!(self, JsonRpcMessageResult::Error(_))
46    }
47
48    /// Check if this result needs a response
49    pub fn needs_response(&self) -> bool {
50        !matches!(self, JsonRpcMessageResult::NoResponse)
51    }
52}
53
54/// Parse a JSON string into a JSON-RPC message
55pub fn parse_json_rpc_message(json_str: &str) -> Result<JsonRpcMessage, JsonRpcError> {
56    let value: Value = serde_json::from_str(json_str)
57        .map_err(|_| JsonRpcError::parse_error())?;
58
59    // Check if it's a valid JSON-RPC message
60    if !value.is_object() {
61        return Err(JsonRpcError::invalid_request(None));
62    }
63
64    let obj = value.as_object().unwrap();
65
66    // Check JSON-RPC version
67    match obj.get("jsonrpc") {
68        Some(version) if version == "2.0" => {}
69        _ => return Err(JsonRpcError::invalid_request(None)),
70    }
71
72    // Check if it has an ID (request) or not (notification)
73    if obj.contains_key("id") {
74        // It's a request
75        serde_json::from_value::<JsonRpcRequest>(value.clone())
76            .map(JsonRpcMessage::Request)
77            .map_err(|_| {
78                // Try to extract ID for error response
79                let id = obj.get("id")
80                    .and_then(|v| {
81                        match v {
82                            Value::String(s) => Some(RequestId::String(s.clone())),
83                            Value::Number(n) => n.as_i64().map(RequestId::Number),
84                            _ => None,
85                        }
86                    });
87                JsonRpcError::invalid_request(id)
88            })
89    } else {
90        // It's a notification
91        serde_json::from_value::<JsonRpcNotification>(value)
92            .map(JsonRpcMessage::Notification)
93            .map_err(|_| JsonRpcError::invalid_request(None))
94    }
95}
96
97/// Utility functions for working with JSON-RPC messages
98impl JsonRpcMessage {
99    /// Get the method name
100    pub fn method(&self) -> &str {
101        match self {
102            JsonRpcMessage::Request(req) => &req.method,
103            JsonRpcMessage::Notification(notif) => &notif.method,
104        }
105    }
106
107    /// Check if this is a request (has ID)
108    pub fn is_request(&self) -> bool {
109        matches!(self, JsonRpcMessage::Request(_))
110    }
111
112    /// Check if this is a notification (no ID)
113    pub fn is_notification(&self) -> bool {
114        matches!(self, JsonRpcMessage::Notification(_))
115    }
116
117    /// Get the request ID if this is a request
118    pub fn request_id(&self) -> Option<&RequestId> {
119        match self {
120            JsonRpcMessage::Request(req) => Some(&req.id),
121            JsonRpcMessage::Notification(_) => None,
122        }
123    }
124}
125
126/// Parse multiple JSON-RPC messages from a single JSON string
127/// This handles both single messages and potential future batch support
128pub fn parse_json_rpc_messages(json_str: &str) -> Vec<Result<JsonRpcMessage, JsonRpcError>> {
129    // For now, we only support single messages (JSON-RPC 2.0 removed batch support)
130    vec![parse_json_rpc_message(json_str)]
131}
132
133/// Create a simple success response
134pub fn create_success_response(id: RequestId, result: Value) -> JsonRpcMessageResult {
135    JsonRpcMessageResult::Response(JsonRpcResponse::success(id, result))
136}
137
138/// Create a simple error response
139pub fn create_error_response(id: Option<RequestId>, code: i64, message: &str) -> JsonRpcMessageResult {
140    let error_obj = crate::error::JsonRpcErrorObject {
141        code,
142        message: message.to_string(),
143        data: None,
144    };
145    JsonRpcMessageResult::Error(JsonRpcError::new(id, error_obj))
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use serde_json::json;
152
153    #[test]
154    fn test_parse_valid_request() {
155        let json = r#"{"jsonrpc": "2.0", "method": "test", "id": 1}"#;
156        let message = parse_json_rpc_message(json).unwrap();
157
158        assert!(message.is_request());
159        assert_eq!(message.method(), "test");
160        assert_eq!(message.request_id(), Some(&RequestId::Number(1)));
161    }
162
163    #[test]
164    fn test_parse_valid_notification() {
165        let json = r#"{"jsonrpc": "2.0", "method": "notify"}"#;
166        let message = parse_json_rpc_message(json).unwrap();
167
168        assert!(message.is_notification());
169        assert_eq!(message.method(), "notify");
170        assert_eq!(message.request_id(), None);
171    }
172
173    #[test]
174    fn test_parse_invalid_json() {
175        let json = r#"{"jsonrpc": "2.0", "method": "test""#; // Invalid JSON
176        let result = parse_json_rpc_message(json);
177
178        assert!(result.is_err());
179        let error = result.unwrap_err();
180        assert_eq!(error.error.code, -32700); // Parse error
181    }
182
183    #[test]
184    fn test_parse_invalid_version() {
185        let json = r#"{"jsonrpc": "1.0", "method": "test", "id": 1}"#;
186        let result = parse_json_rpc_message(json);
187
188        assert!(result.is_err());
189        let error = result.unwrap_err();
190        assert_eq!(error.error.code, -32600); // Invalid request
191    }
192
193    #[test]
194    fn test_message_result_to_json() {
195        let response = create_success_response(
196            RequestId::Number(1),
197            json!({"result": "success"}),
198        );
199
200        let json_str = response.to_json_string().unwrap();
201        assert!(json_str.contains("\"result\""));
202        assert!(json_str.contains("\"jsonrpc\":\"2.0\""));
203    }
204
205    #[test]
206    fn test_message_result_properties() {
207        let success = create_success_response(RequestId::Number(1), json!({}));
208        let error = create_error_response(Some(RequestId::Number(1)), -32601, "Not found");
209        let no_response = JsonRpcMessageResult::NoResponse;
210
211        assert!(!success.is_error());
212        assert!(success.needs_response());
213
214        assert!(error.is_error());
215        assert!(error.needs_response());
216
217        assert!(!no_response.is_error());
218        assert!(!no_response.needs_response());
219    }
220}