fix: improve JSON-RPC notification handling to match specification (#…

…1124)
block · Feb 13, 2025 · 5c44967 · 5c44967
1 parent 89d3748
commit 5c44967
Showing 1 changed file with 52 additions and 1 deletion.
diff --git a/crates/mcp-core/src/protocol.rs b/crates/mcp-core/src/protocol.rs
@@ -96,7 +96,7 @@ impl TryFrom<JsonRpcRaw> for JsonRpcMessage {
 
         // If we have a method, it's either a notification or request
         if let Some(method) = raw.method {
-            if method.starts_with("notifications/") {
+            if raw.id.is_none() {
                 return Ok(JsonRpcMessage::Notification(JsonRpcNotification {
                     jsonrpc: raw.jsonrpc,
                     method,
@@ -236,3 +236,54 @@ pub struct GetPromptResult {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct EmptyResult {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn test_notification_conversion() {
+        let raw = JsonRpcRaw {
+            jsonrpc: "2.0".to_string(),
+            id: None,
+            method: Some("notify".to_string()),
+            params: Some(json!({"key": "value"})),
+            result: None,
+            error: None,
+        };
+
+        let message = JsonRpcMessage::try_from(raw).unwrap();
+        match message {
+            JsonRpcMessage::Notification(n) => {
+                assert_eq!(n.jsonrpc, "2.0");
+                assert_eq!(n.method, "notify");
+                assert_eq!(n.params.unwrap(), json!({"key": "value"}));
+            }
+            _ => panic!("Expected Notification"),
+        }
+    }
+
+    #[test]
+    fn test_request_conversion() {
+        let raw = JsonRpcRaw {
+            jsonrpc: "2.0".to_string(),
+            id: Some(1),
+            method: Some("request".to_string()),
+            params: Some(json!({"key": "value"})),
+            result: None,
+            error: None,
+        };
+
+        let message = JsonRpcMessage::try_from(raw).unwrap();
+        match message {
+            JsonRpcMessage::Request(r) => {
+                assert_eq!(r.jsonrpc, "2.0");
+                assert_eq!(r.id, Some(1));
+                assert_eq!(r.method, "request");
+                assert_eq!(r.params.unwrap(), json!({"key": "value"}));
+            }
+            _ => panic!("Expected Request"),
+        }
+    }
+}