Skip to content

Commit 7b8cb42

Browse files
CopilotScooletzeiriktsarpalis
committed
Optimize JsonRpcMessage deserialization
Co-authored-by: scooletz <scooletz@gmail.com> Co-authored-by: eiriktsarpalis <2813363+eiriktsarpalis@users.noreply.github.com>
1 parent b2afc2f commit 7b8cb42

4 files changed

Lines changed: 848 additions & 30 deletions

File tree

‎src/ModelContextProtocol.Analyzers/XmlToDescriptionGenerator.cs‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,14 @@ public sealed class XmlToDescriptionGenerator : IIncrementalGenerator
1919
{
2020
private const string GeneratedFileName = "ModelContextProtocol.Descriptions.g.cs";
2121

22+
/// <summary>
23+
/// A display format that produces fully-qualified type names with "global::" prefix
24+
/// and includes nullability annotations.
25+
/// </summary>
26+
private static readonly SymbolDisplayFormat s_fullyQualifiedFormatWithNullability =
27+
SymbolDisplayFormat.FullyQualifiedFormat.AddMiscellaneousOptions(
28+
SymbolDisplayMiscellaneousOptions.IncludeNullableReferenceTypeModifier);
29+
2230
public void Initialize(IncrementalGeneratorInitializationContext context)
2331
{
2432
// Extract method information for all MCP tools, prompts, and resources.
@@ -125,7 +133,7 @@ private static MethodToGenerate ExtractMethodInfo(
125133
.Where(m => !m.IsKind(SyntaxKind.AsyncKeyword))
126134
.Select(m => m.Text);
127135
string modifiersStr = string.Join(" ", modifiers);
128-
string returnType = methodSymbol.ReturnType.ToDisplayString(SymbolDisplayFormat.MinimallyQualifiedFormat);
136+
string returnType = methodSymbol.ReturnType.ToDisplayString(s_fullyQualifiedFormatWithNullability);
129137
string methodName = methodSymbol.Name;
130138

131139
// Extract parameters
@@ -137,7 +145,7 @@ private static MethodToGenerate ExtractMethodInfo(
137145
var paramSyntax = i < parameterSyntaxList.Count ? parameterSyntaxList[i] : null;
138146

139147
parameters[i] = new ParameterInfo(
140-
ParameterType: param.Type.ToDisplayString(SymbolDisplayFormat.MinimallyQualifiedFormat),
148+
ParameterType: param.Type.ToDisplayString(s_fullyQualifiedFormatWithNullability),
141149
Name: param.Name,
142150
HasDescriptionAttribute: descriptionAttribute is not null && HasAttribute(param, descriptionAttribute),
143151
XmlDescription: xmlDocs?.Parameters.TryGetValue(param.Name, out var pd) == true && !string.IsNullOrWhiteSpace(pd) ? pd : null,

‎src/ModelContextProtocol.Core/Protocol/JsonRpcMessage.cs‎

Lines changed: 149 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
using ModelContextProtocol.Server;
22
using System.ComponentModel;
3+
using System.Diagnostics;
34
using System.Text.Json;
5+
using System.Text.Json.Nodes;
46
using System.Text.Json.Serialization;
57

68
namespace ModelContextProtocol.Protocol;
@@ -70,59 +72,83 @@ private protected JsonRpcMessage()
7072
[EditorBrowsable(EditorBrowsableState.Never)]
7173
public sealed class Converter : JsonConverter<JsonRpcMessage>
7274
{
75+
private const string JsonRpcVersion = "2.0";
76+
7377
/// <inheritdoc/>
7478
public override JsonRpcMessage? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
7579
{
76-
if (reader.TokenType != JsonTokenType.StartObject)
77-
{
78-
throw new JsonException("Expected StartObject token");
79-
}
80-
81-
using var doc = JsonDocument.ParseValue(ref reader);
82-
var root = doc.RootElement;
80+
var union = ParseUnion(ref reader, options);
8381

8482
// All JSON-RPC messages must have a jsonrpc property with value "2.0"
85-
if (!root.TryGetProperty("jsonrpc", out var versionProperty) ||
86-
versionProperty.GetString() != "2.0")
83+
if (union.JsonRpc != JsonRpcVersion)
8784
{
8885
throw new JsonException("Invalid or missing jsonrpc version");
8986
}
9087

91-
// Determine the message type based on the presence of id, method, and error properties
92-
bool hasId = root.TryGetProperty("id", out _);
93-
bool hasMethod = root.TryGetProperty("method", out _);
94-
bool hasError = root.TryGetProperty("error", out _);
95-
96-
var rawText = root.GetRawText();
97-
9888
// Messages with an id but no method are responses
99-
if (hasId && !hasMethod)
89+
if (union.HasId && !union.HasMethod)
10090
{
10191
// Messages with an error property are error responses
102-
if (hasError)
92+
if (union.HasError)
10393
{
104-
return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcError>());
94+
if (union.Error is null)
95+
{
96+
throw new JsonException("Error property cannot be null");
97+
}
98+
99+
return new JsonRpcError
100+
{
101+
JsonRpc = union.JsonRpc,
102+
Id = union.Id,
103+
Error = union.Error
104+
};
105105
}
106106

107107
// Messages with a result property are success responses
108-
if (root.TryGetProperty("result", out _))
108+
if (union.HasResult)
109109
{
110-
return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcResponse>());
110+
return new JsonRpcResponse
111+
{
112+
JsonRpc = union.JsonRpc,
113+
Id = union.Id,
114+
Result = union.Result
115+
};
111116
}
112117

113118
throw new JsonException("Response must have either result or error");
114119
}
115120

116121
// Messages with a method but no id are notifications
117-
if (hasMethod && !hasId)
122+
if (union.HasMethod && !union.HasId)
118123
{
119-
return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcNotification>());
124+
if (union.Method is null)
125+
{
126+
throw new JsonException("Method property cannot be null");
127+
}
128+
129+
return new JsonRpcNotification
130+
{
131+
JsonRpc = union.JsonRpc,
132+
Method = union.Method,
133+
Params = union.Params
134+
};
120135
}
121136

122137
// Messages with both method and id are requests
123-
if (hasMethod && hasId)
138+
if (union.HasMethod && union.HasId)
124139
{
125-
return JsonSerializer.Deserialize(rawText, options.GetTypeInfo<JsonRpcRequest>());
140+
if (union.Method is null)
141+
{
142+
throw new JsonException("Method property cannot be null");
143+
}
144+
145+
return new JsonRpcRequest
146+
{
147+
JsonRpc = union.JsonRpc,
148+
Id = union.Id,
149+
Method = union.Method,
150+
Params = union.Params
151+
};
126152
}
127153

128154
throw new JsonException("Invalid JSON-RPC message format");
@@ -149,5 +175,103 @@ public override void Write(Utf8JsonWriter writer, JsonRpcMessage value, JsonSeri
149175
throw new JsonException($"Unknown JSON-RPC message type: {value.GetType()}");
150176
}
151177
}
178+
179+
/// <summary>
180+
/// Manually parses a JSON-RPC message from the reader into the Union struct.
181+
/// </summary>
182+
private static Union ParseUnion(ref Utf8JsonReader reader, JsonSerializerOptions options)
183+
{
184+
var union = new Union
185+
{
186+
JsonRpc = string.Empty // Initialize to avoid null reference warnings
187+
};
188+
189+
if (reader.TokenType != JsonTokenType.StartObject)
190+
{
191+
throw new JsonException("Expected StartObject token");
192+
}
193+
194+
while (true)
195+
{
196+
bool success = reader.Read();
197+
Debug.Assert(success, "custom converters are guaranteed to be passed fully buffered objects");
198+
199+
if (reader.TokenType is JsonTokenType.EndObject)
200+
{
201+
break;
202+
}
203+
204+
Debug.Assert(reader.TokenType is JsonTokenType.PropertyName);
205+
string propertyName = reader.GetString()!;
206+
207+
success = reader.Read();
208+
Debug.Assert(success, "custom converters are guaranteed to be passed fully buffered objects");
209+
210+
switch (propertyName)
211+
{
212+
case "jsonrpc":
213+
union.JsonRpc = reader.GetString() ?? string.Empty;
214+
break;
215+
216+
case "id":
217+
union.Id = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<RequestId>());
218+
union.HasId = true;
219+
break;
220+
221+
case "method":
222+
union.Method = reader.GetString();
223+
union.HasMethod = union.Method is not null;
224+
break;
225+
226+
case "params":
227+
union.Params = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
228+
break;
229+
230+
case "error":
231+
union.Error = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonRpcErrorDetail>());
232+
union.HasError = union.Error is not null;
233+
break;
234+
235+
case "result":
236+
union.Result = JsonSerializer.Deserialize(ref reader, options.GetTypeInfo<JsonNode>());
237+
union.HasResult = true;
238+
break;
239+
240+
default:
241+
// Skip unknown properties
242+
reader.Skip();
243+
break;
244+
}
245+
}
246+
247+
return union;
248+
}
249+
250+
/// <summary>
251+
/// Private struct to hold parsed JSON-RPC message data during deserialization.
252+
/// </summary>
253+
private struct Union
254+
{
255+
/// <summary>The JSON-RPC protocol version (must be "2.0").</summary>
256+
public string JsonRpc;
257+
/// <summary>The message identifier for requests and responses.</summary>
258+
public RequestId Id;
259+
/// <summary>The method name for requests and notifications.</summary>
260+
public string? Method;
261+
/// <summary>The parameters for requests and notifications.</summary>
262+
public JsonNode? Params;
263+
/// <summary>The error details for error responses.</summary>
264+
public JsonRpcErrorDetail? Error;
265+
/// <summary>The result for successful responses.</summary>
266+
public JsonNode? Result;
267+
/// <summary>Indicates whether an 'id' property was present.</summary>
268+
public bool HasId;
269+
/// <summary>Indicates whether a 'method' property was present.</summary>
270+
public bool HasMethod;
271+
/// <summary>Indicates whether an 'error' property was present.</summary>
272+
public bool HasError;
273+
/// <summary>Indicates whether a 'result' property was present.</summary>
274+
public bool HasResult;
275+
}
152276
}
153277
}

0 commit comments

Comments
 (0)