Optimize JsonRpcMessage deserialization with single-pass parsing

src/ModelContextProtocol.Core/Protocol/JsonRpcMessage.cs

@@ -1,6 +1,8 @@
 using ModelContextProtocol.Server;
 using System.ComponentModel;
+using System.Diagnostics;
 using System.Text.Json;
+using System.Text.Json.Nodes;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol;
@@ -70,62 +72,61 @@ private protected JsonRpcMessage()
     [EditorBrowsable(EditorBrowsableState.Never)]
     public sealed class Converter : JsonConverter<JsonRpcMessage>
     {
+        private const string JsonRpcVersion = "2.0";
+
         /// <inheritdoc/>
         public override JsonRpcMessage? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
         {
-            if (reader.TokenType != JsonTokenType.StartObject)
-            {
-                throw new JsonException("Expected StartObject token");
-            }
-
-            using var doc = JsonDocument.ParseValue(ref reader);
-            var root = doc.RootElement;
+            ParseUnion(ref reader, options, out Union union);
 
             // All JSON-RPC messages must have a jsonrpc property with value "2.0"
-            if (!root.TryGetProperty("jsonrpc", out var versionProperty) ||
-                versionProperty.GetString() != "2.0")
+            if (union.JsonRpc != JsonRpcVersion)
             {
                 throw new JsonException("Invalid or missing jsonrpc version");
             }
 
-            // Determine the message type based on the presence of id, method, and error properties
-            bool hasId = root.TryGetProperty("id", out _);
-            bool hasMethod = root.TryGetProperty("method", out _);
-            bool hasError = root.TryGetProperty("error", out _);
-
-            var rawText = root.GetRawText();
-
-            // Messages with an id but no method are responses
-            if (hasId && !hasMethod)
+            // Determine message type based on presence of id and method properties
+            return union switch
             {
-                // Messages with an error property are error responses
-                if (hasError)
+                // Messages with both method and id are requests
+                { Method: not null, Id.Id: not null } => new JsonRpcRequest
                 {
-                    return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcError>());
-                }
-
-                // Messages with a result property are success responses
-                if (root.TryGetProperty("result", out _))
+                    JsonRpc = union.JsonRpc,
+                    Id = union.Id,
+                    Method = union.Method,
+                    Params = union.Params
+                },
+
+                // Messages with a method but no id are notifications
+                { Method: not null } => new JsonRpcNotification
                 {
-                    return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcResponse>());
-                }
+                    JsonRpc = union.JsonRpc,
+                    Method = union.Method,
+                    Params = union.Params
+                },
 
-                throw new JsonException("Response must have either result or error");
-            }
+                // Messages with a result and id are success responses
+                { HasResult: true, Id.Id: not null } => new JsonRpcResponse
+                {
+                    JsonRpc = union.JsonRpc,
+                    Id = union.Id,
+                    Result = union.Result
+                },
 
-            // Messages with a method but no id are notifications
-            if (hasMethod && !hasId)
-            {
-                return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcNotification>());
-            }
+                // Messages with an error and id are error responses
+                { Error: not null, Id.Id: not null } => new JsonRpcError
+                {
+                    JsonRpc = union.JsonRpc,
+                    Id = union.Id,
+                    Error = union.Error
+                },
 
-            // Messages with both method and id are requests
-            if (hasMethod && hasId)
-            {
-                return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcRequest>());
-            }
+                // Error: Messages with an id but no method, error, or result are invalid
+                { Id.Id: not null } => throw new JsonException("Response must have either result or error"),
 
-            throw new JsonException("Invalid JSON-RPC message format");
+                // Error: Messages with neither id nor method are invalid
+                _ => throw new JsonException("Invalid JSON-RPC message format")
+            };
         }
 
         /// <inheritdoc/>
@@ -149,5 +150,90 @@ public override void Write(Utf8JsonWriter writer, JsonRpcMessage value, JsonSeri
                     throw new JsonException($"Unknown JSON-RPC message type: {value.GetType()}");
             }
         }
+
+        /// <summary>
+        /// Manually parses a JSON-RPC message from the reader into the Union struct.
+        /// </summary>
+        private static void ParseUnion(ref Utf8JsonReader reader, JsonSerializerOptions options, out Union union)
+        {
+            union = default;
+            union.JsonRpc = string.Empty; // Initialize to avoid null reference warnings
+
+            if (reader.TokenType != JsonTokenType.StartObject)
+            {
+                throw new JsonException("Expected StartObject token");
+            }
+
+            while (true)
+            {
+                bool success = reader.Read();
+                Debug.Assert(success, "custom converters are guaranteed to be passed fully buffered objects");
+
+                if (reader.TokenType is JsonTokenType.EndObject)
+                {
+                    break;
+                }
+
+                Debug.Assert(reader.TokenType is JsonTokenType.PropertyName);
+                string propertyName = reader.GetString()!;
+
+                success = reader.Read();
+                Debug.Assert(success, "custom converters are guaranteed to be passed fully buffered objects");
+
+                switch (propertyName)
+                {
+                    case "jsonrpc":
+                        union.JsonRpc = reader.GetString() ?? string.Empty;
+                        break;
+
+                    case "id":
+                        union.Id = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<RequestId>());
+                        break;
+
+                    case "method":
+                        union.Method = reader.GetString();
+                        break;
+
+                    case "params":
+                        union.Params = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
+                        break;
+
+                    case "error":
+                        union.Error = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonRpcErrorDetail>());
+                        break;
+
+                    case "result":
+                        union.Result = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
+                        union.HasResult = true;
+                        break;
+
+                    default:
+                        // Skip unknown properties
+                        reader.Skip();
+                        break;
+                }
+            }
+        }
+
+        /// <summary>
+        /// Private struct to hold parsed JSON-RPC message data during deserialization.
+        /// </summary>
+        private struct Union
+        {
+            /// <summary>The JSON-RPC protocol version (must be "2.0").</summary>
+            public string JsonRpc;
+            /// <summary>The message identifier for requests and responses.</summary>
+            public RequestId Id;
+            /// <summary>The method name for requests and notifications.</summary>
+            public string? Method;
+            /// <summary>The parameters for requests and notifications.</summary>
+            public JsonNode? Params;
+            /// <summary>The error details for error responses.</summary>
+            public JsonRpcErrorDetail? Error;
+            /// <summary>The result for successful responses.</summary>
+            public JsonNode? Result;
+            /// <summary>Indicates whether a 'result' property was present (result can be null).</summary>
+            public bool HasResult;
+        }
     }
 }