Skip to content

Commit e9f325d

Browse files
Copilotstephentoub
andcommitted
Inline ParseUnion and remove Union struct, optimize jsonrpc validation
Major refactoring to simplify the code based on reviewer feedback: 1. Inlined ParseUnion method directly into Read() - eliminates method call overhead and makes the parsing logic more straightforward 2. Replaced Union struct with local variables in Read() - removes the need for an intermediate data structure, reducing memory overhead 3. Optimized "jsonrpc" validation using reader.ValueTextEquals("2.0"u8) instead of allocating a string - validates raw UTF-8 bytes directly without string allocation 4. Track jsonrpc presence with hasJsonRpc boolean flag, validate after parsing completes 5. Replaced switch expression pattern matching with simple if statements for clarity since we no longer have a struct Benefits: - No Union struct to copy around - No string allocation for "jsonrpc" property - Cleaner, more straightforward code - All the parsing logic in one place - Still maintains single-pass deserialization All 1035 tests pass. Co-authored-by: Scooletz <scooletz@users.noreply.github.com> Co-authored-by: eiriktsarpalis <2813363+eiriktsarpalis@users.noreply.github.com> Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com> Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com>
1 parent bbb6535 commit e9f325d

1 file changed

Lines changed: 103 additions & 107 deletions

File tree

src/ModelContextProtocol.Core/Protocol/JsonRpcMessage.cs

Lines changed: 103 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -77,93 +77,22 @@ public sealed class Converter : JsonConverter<JsonRpcMessage>
7777
/// <inheritdoc/>
7878
public override JsonRpcMessage? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
7979
{
80-
ParseUnion(ref reader, options, out Union union);
81-
82-
// All JSON-RPC messages must have a jsonrpc property with value "2.0"
83-
if (union.JsonRpc != JsonRpcVersion)
84-
{
85-
throw new JsonException("Invalid or missing jsonrpc version");
86-
}
87-
88-
// Determine message type based on presence of id and method properties
89-
return union switch
90-
{
91-
// Messages with both method and id are requests
92-
{ Method: not null, Id.Id: not null } => new JsonRpcRequest
93-
{
94-
JsonRpc = union.JsonRpc,
95-
Id = union.Id,
96-
Method = union.Method,
97-
Params = union.Params
98-
},
99-
100-
// Messages with a method but no id are notifications
101-
{ Method: not null } => new JsonRpcNotification
102-
{
103-
JsonRpc = union.JsonRpc,
104-
Method = union.Method,
105-
Params = union.Params
106-
},
107-
108-
// Messages with a result and id are success responses
109-
{ HasResult: true, Id.Id: not null } => new JsonRpcResponse
110-
{
111-
JsonRpc = union.JsonRpc,
112-
Id = union.Id,
113-
Result = union.Result
114-
},
115-
116-
// Messages with an error and id are error responses
117-
{ Error: not null, Id.Id: not null } => new JsonRpcError
118-
{
119-
JsonRpc = union.JsonRpc,
120-
Id = union.Id,
121-
Error = union.Error
122-
},
123-
124-
// Error: Messages with an id but no method, error, or result are invalid
125-
{ Id.Id: not null } => throw new JsonException("Response must have either result or error"),
126-
127-
// Error: Messages with neither id nor method are invalid
128-
_ => throw new JsonException("Invalid JSON-RPC message format")
129-
};
130-
}
131-
132-
/// <inheritdoc/>
133-
public override void Write(Utf8JsonWriter writer, JsonRpcMessage value, JsonSerializerOptions options)
134-
{
135-
switch (value)
136-
{
137-
case JsonRpcRequest request:
138-
JsonSerializer.Serialize(writer, request, options.GetTypeInfo<JsonRpcRequest>());
139-
break;
140-
case JsonRpcNotification notification:
141-
JsonSerializer.Serialize(writer, notification, options.GetTypeInfo<JsonRpcNotification>());
142-
break;
143-
case JsonRpcResponse response:
144-
JsonSerializer.Serialize(writer, response, options.GetTypeInfo<JsonRpcResponse>());
145-
break;
146-
case JsonRpcError error:
147-
JsonSerializer.Serialize(writer, error, options.GetTypeInfo<JsonRpcError>());
148-
break;
149-
default:
150-
throw new JsonException($"Unknown JSON-RPC message type: {value.GetType()}");
151-
}
152-
}
153-
154-
/// <summary>
155-
/// Manually parses a JSON-RPC message from the reader into the Union struct.
156-
/// </summary>
157-
private static void ParseUnion(ref Utf8JsonReader reader, JsonSerializerOptions options, out Union union)
158-
{
159-
union = default;
160-
union.JsonRpc = string.Empty; // Initialize to avoid null reference warnings
161-
16280
if (reader.TokenType != JsonTokenType.StartObject)
16381
{
16482
throw new JsonException("Expected StartObject token");
16583
}
16684

85+
// Track whether we've seen the required jsonrpc property
86+
bool hasJsonRpc = false;
87+
88+
// Local variables for parsed message data
89+
RequestId id = default;
90+
string? method = null;
91+
JsonNode? @params = null;
92+
JsonRpcErrorDetail? error = null;
93+
JsonNode? result = null;
94+
bool hasResult = false;
95+
16796
while (true)
16897
{
16998
bool success = reader.Read();
@@ -183,28 +112,33 @@ private static void ParseUnion(ref Utf8JsonReader reader, JsonSerializerOptions
183112
switch (propertyName)
184113
{
185114
case "jsonrpc":
186-
union.JsonRpc = reader.GetString() ?? string.Empty;
115+
// Validate that the value is "2.0" without allocating a string
116+
if (!reader.ValueTextEquals("2.0"u8))
117+
{
118+
throw new JsonException("Invalid jsonrpc version");
119+
}
120+
hasJsonRpc = true;
187121
break;
188122

189123
case "id":
190-
union.Id = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<RequestId>());
124+
id = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<RequestId>());
191125
break;
192126

193127
case "method":
194-
union.Method = reader.GetString();
128+
method = reader.GetString();
195129
break;
196130

197131
case "params":
198-
union.Params = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
132+
@params = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
199133
break;
200134

201135
case "error":
202-
union.Error = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonRpcErrorDetail>());
136+
error = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonRpcErrorDetail>());
203137
break;
204138

205139
case "result":
206-
union.Result = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
207-
union.HasResult = true;
140+
result = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
141+
hasResult = true;
208142
break;
209143

210144
default:
@@ -213,27 +147,89 @@ private static void ParseUnion(ref Utf8JsonReader reader, JsonSerializerOptions
213147
break;
214148
}
215149
}
150+
151+
// All JSON-RPC messages must have a jsonrpc property with value "2.0"
152+
if (!hasJsonRpc)
153+
{
154+
throw new JsonException("Invalid or missing jsonrpc version");
155+
}
156+
157+
// Determine message type based on presence of id and method properties
158+
// Messages with both method and id are requests
159+
if (method is not null && id.Id is not null)
160+
{
161+
return new JsonRpcRequest
162+
{
163+
JsonRpc = JsonRpcVersion,
164+
Id = id,
165+
Method = method,
166+
Params = @params
167+
};
168+
}
169+
170+
// Messages with a method but no id are notifications
171+
if (method is not null)
172+
{
173+
return new JsonRpcNotification
174+
{
175+
JsonRpc = JsonRpcVersion,
176+
Method = method,
177+
Params = @params
178+
};
179+
}
180+
181+
// Messages with a result and id are success responses
182+
if (hasResult && id.Id is not null)
183+
{
184+
return new JsonRpcResponse
185+
{
186+
JsonRpc = JsonRpcVersion,
187+
Id = id,
188+
Result = result
189+
};
190+
}
191+
192+
// Messages with an error and id are error responses
193+
if (error is not null && id.Id is not null)
194+
{
195+
return new JsonRpcError
196+
{
197+
JsonRpc = JsonRpcVersion,
198+
Id = id,
199+
Error = error
200+
};
201+
}
202+
203+
// Error: Messages with an id but no method, error, or result are invalid
204+
if (id.Id is not null)
205+
{
206+
throw new JsonException("Response must have either result or error");
207+
}
208+
209+
// Error: Messages with neither id nor method are invalid
210+
throw new JsonException("Invalid JSON-RPC message format");
216211
}
217212

218-
/// <summary>
219-
/// Private struct to hold parsed JSON-RPC message data during deserialization.
220-
/// </summary>
221-
private struct Union
213+
/// <inheritdoc/>
214+
public override void Write(Utf8JsonWriter writer, JsonRpcMessage value, JsonSerializerOptions options)
222215
{
223-
/// <summary>The JSON-RPC protocol version (must be "2.0").</summary>
224-
public string JsonRpc;
225-
/// <summary>The message identifier for requests and responses.</summary>
226-
public RequestId Id;
227-
/// <summary>The method name for requests and notifications.</summary>
228-
public string? Method;
229-
/// <summary>The parameters for requests and notifications.</summary>
230-
public JsonNode? Params;
231-
/// <summary>The error details for error responses.</summary>
232-
public JsonRpcErrorDetail? Error;
233-
/// <summary>The result for successful responses.</summary>
234-
public JsonNode? Result;
235-
/// <summary>Indicates whether a 'result' property was present (result can be null).</summary>
236-
public bool HasResult;
216+
switch (value)
217+
{
218+
case JsonRpcRequest request:
219+
JsonSerializer.Serialize(writer, request, options.GetTypeInfo<JsonRpcRequest>());
220+
break;
221+
case JsonRpcNotification notification:
222+
JsonSerializer.Serialize(writer, notification, options.GetTypeInfo<JsonRpcNotification>());
223+
break;
224+
case JsonRpcResponse response:
225+
JsonSerializer.Serialize(writer, response, options.GetTypeInfo<JsonRpcResponse>());
226+
break;
227+
case JsonRpcError error:
228+
JsonSerializer.Serialize(writer, error, options.GetTypeInfo<JsonRpcError>());
229+
break;
230+
default:
231+
throw new JsonException($"Unknown JSON-RPC message type: {value.GetType()}");
232+
}
237233
}
238234
}
239235
}

0 commit comments

Comments
 (0)