Fix critical JSON-RPC 2.0 compliance bugs and improve error handling

Co-authored-by: 0xrinegade <[email protected]>

Author: Copilot

src/http_server.rs

@@ -5,6 +5,7 @@ use axum::{
     routing::{get, post},
     Json, Router,
 };
+use serde_json::Value;
 use tokio::net::TcpListener;
 use tower::ServiceBuilder;
 use tracing::{info, error, debug};
@@ -95,7 +96,10 @@ async fn mcp_api_handler(
     // Validate Content-Type header (should be application/json for MCP)
     if let Some(content_type) = headers.get(CONTENT_TYPE) {
         if let Ok(ct_str) = content_type.to_str() {
-            if !ct_str.starts_with("application/json") {
+            // Be more strict about content type validation
+            let is_valid = ct_str == "application/json" || 
+                          ct_str.starts_with("application/json;");
+            if !is_valid {
                 return create_json_rpc_error_response(
                     -32600,
                     "Invalid Request: Content-Type must be application/json",
@@ -117,14 +121,20 @@ async fn mcp_api_handler(
             // Convert JsonRpcMessage back to proper JSON-RPC 2.0 format
             match serde_json::to_value(&response_message) {
                 Ok(json_response) => {
-                    create_json_rpc_success_response(json_response)
+                    // The response_message is already a properly formatted JSON-RPC response
+                    // Don't double-wrap it in create_json_rpc_success_response
+                    (
+                        StatusCode::OK,
+                        [(CONTENT_TYPE, "application/json")],
+                        Json(json_response)
+                    ).into_response()
                 }
                 Err(e) => {
                     error!("Failed to serialize MCP response: {}", e);
                     create_json_rpc_error_response(
                         -32603,
                         "Internal error: Failed to serialize response",
-                        Some(json_rpc_request.id),
+                        Some(json_rpc_request.id.clone()),
                     )
                 }
             }
@@ -134,7 +144,7 @@ async fn mcp_api_handler(
             create_json_rpc_error_response(
                 -32603,
                 &format!("Internal error: {}", e),
-                Some(json_rpc_request.id),
+                Some(json_rpc_request.id.clone()),
             )
         }
     }
@@ -168,8 +178,8 @@ fn parse_json_rpc_request(request: &serde_json::Value) -> Result<JsonRpcRequest,
         ))?;
 
     let id = request.get("id")
-        .and_then(|v| v.as_u64())
-        .unwrap_or(0);
+        .cloned()
+        .unwrap_or(Value::Null); // Default to null if no ID provided
 
     let params = request.get("params").cloned();
 
@@ -181,19 +191,9 @@ fn parse_json_rpc_request(request: &serde_json::Value) -> Result<JsonRpcRequest,
     })
 }
 
-/// Create a properly formatted JSON-RPC 2.0 success response
-fn create_json_rpc_success_response(result: serde_json::Value) -> Response {
-    (
-        StatusCode::OK,
-        [
-            (CONTENT_TYPE, "application/json"),
-        ],
-        Json(result)
-    ).into_response()
-}
 
 /// Create a properly formatted JSON-RPC 2.0 error response
-fn create_json_rpc_error_response(code: i32, message: &str, id: Option<u64>) -> Response {
+fn create_json_rpc_error_response(code: i32, message: &str, id: Option<Value>) -> Response {
     let error_response = serde_json::json!({
         "jsonrpc": "2.0",
         "error": {
@@ -203,7 +203,7 @@ fn create_json_rpc_error_response(code: i32, message: &str, id: Option<u64>) ->
                 "protocolVersion": crate::protocol::LATEST_PROTOCOL_VERSION
             }
         },
-        "id": id
+        "id": id.unwrap_or(Value::Null)
     });
 
     (
@@ -277,9 +277,6 @@ mod tests {
 
     #[tokio::test]
     async fn test_mcp_api_handler() {
-        use crate::Config;
-        use crate::server::ServerState;
-        
         // Create a test server state using Config::load() or a minimal config
         // For testing purposes, we'll skip the actual test since it requires valid config
         // In a real test environment, you'd want to create a minimal test config

src/tools/mod.rs

@@ -25,8 +25,8 @@ use url::Url;
 ///
 /// # Returns
 /// * `JsonRpcMessage` - Formatted success response
-pub fn create_success_response(result: Value, id: u64) -> JsonRpcMessage {
-    log::debug!("Creating success response with id {}", id);
+pub fn create_success_response(result: Value, id: Value) -> JsonRpcMessage {
+    log::debug!("Creating success response with id {:?}", id);
     JsonRpcMessage::Response(JsonRpcResponse {
         jsonrpc: JsonRpcVersion::V2,
         id,
@@ -52,7 +52,7 @@ pub fn create_success_response(result: Value, id: u64) -> JsonRpcMessage {
 pub fn create_error_response(
     code: i32,
     message: String,
-    id: u64,
+    id: Value,
     protocol_version: Option<&str>,
 ) -> JsonRpcMessage {
     log::error!("Creating error response: {} (code: {})", message, code);
@@ -93,7 +93,7 @@ pub async fn handle_initialize(
                 return Ok(create_error_response(
                     -32602,
                     "Invalid params: protocolVersion is required".to_string(),
-                    id.and_then(|v| v.as_u64()).unwrap_or(0),
+                    id.unwrap_or(Value::Null),
                     Some(state.protocol_version.as_str()),
                 ));
             }
@@ -119,7 +119,7 @@ pub async fn handle_initialize(
                     "Protocol version mismatch. Server: {}, Client: {}",
                     state.protocol_version, init_params.protocol_version
                 ),
-                id.and_then(|v| v.as_u64()).unwrap_or(0),
+                id.unwrap_or(Value::Null),
                 Some(state.protocol_version.as_str()),
             ));
         }
@@ -1216,14 +1216,14 @@ pub async fn handle_initialize(
         log::info!("Server initialized successfully");
         Ok(create_success_response(
             serde_json::to_value(response).unwrap(),
-            id.and_then(|v| v.as_u64()).unwrap_or(0),
+            id.unwrap_or(Value::Null),
         ))
     } else {
         log::error!("Missing initialization params");
         Ok(create_error_response(
             -32602,
             "Invalid params".to_string(),
-            id.and_then(|v| v.as_u64()).unwrap_or(0),
+            id.unwrap_or(Value::Null),
             Some(state.protocol_version.as_str()),
         ))
     }
@@ -1239,14 +1239,14 @@ pub async fn handle_cancelled(
         let _cancel_params: CancelledParams = serde_json::from_value(params)?;
         Ok(create_success_response(
             Value::Null,
-            id.and_then(|v| v.as_u64()).unwrap_or(0),
+            id.unwrap_or(Value::Null),
         ))
     } else {
         log::error!("Missing cancelled params");
         Ok(create_error_response(
             -32602,
             "Invalid params".to_string(),
-            id.and_then(|v| v.as_u64()).unwrap_or(0),
+            id.unwrap_or(Value::Null),
             Some(state.protocol_version.as_str()),
         ))
     }
@@ -2031,7 +2031,7 @@ pub async fn handle_tools_list(id: Option<Value>, _state: &ServerState) -> Resul
 
     Ok(create_success_response(
         serde_json::to_value(response).unwrap(),
-        id.and_then(|v| v.as_u64()).unwrap_or(0),
+        id.unwrap_or(Value::Null),
     ))
 }
 
@@ -2297,7 +2297,7 @@ pub async fn handle_request(
                 "initialize" => {
                     let response = handle_initialize(
                         req.params,
-                        Some(serde_json::Value::Number(req.id.into())),
+                        Some(req.id.clone()),
                         &state_guard,
                     )
                     .await?;
@@ -2313,13 +2313,13 @@ pub async fn handle_request(
                 "cancelled" => {
                     handle_cancelled(
                         req.params,
-                        Some(serde_json::Value::Number(req.id.into())),
+                        Some(req.id.clone()),
                         &state_guard,
                     )
                     .await
                 }
                 "tools/list" => {
-                    handle_tools_list(Some(serde_json::Value::Number(req.id.into())), &state_guard)
+                    handle_tools_list(Some(req.id.clone()), &state_guard)
                         .await
                 }
 
@@ -3698,7 +3698,7 @@ pub async fn handle_request(
             Ok(create_error_response(
                 -32600,
                 "Invalid Request: expected request message".to_string(),
-                0,
+                Value::Null,
                 None,
             ))
         }
@@ -3712,7 +3712,7 @@ pub async fn handle_request(
                 Ok(create_error_response(
                     -32600,
                     format!("Unsupported notification: {}", notification.method),
-                    0,
+                    Value::Null,
                     None,
                 ))
             }

src/transport.rs

@@ -21,7 +21,7 @@ impl Default for JsonRpcVersion {
 #[serde(rename_all = "camelCase")]
 pub struct JsonRpcRequest {
     pub jsonrpc: JsonRpcVersion,
-    pub id: u64,
+    pub id: Value, // JSON-RPC 2.0 allows string, number, or null
     pub method: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub params: Option<Value>,
@@ -31,7 +31,7 @@ pub struct JsonRpcRequest {
 #[serde(rename_all = "camelCase")]
 pub struct JsonRpcResponse {
     pub jsonrpc: JsonRpcVersion,
-    pub id: u64,
+    pub id: Value, // JSON-RPC 2.0 allows string, number, or null
     #[serde(skip_serializing_if = "Option::is_none")]
     pub result: Option<Value>,
     #[serde(skip_serializing_if = "Option::is_none")]