Skip to content

Commit 1c53fc2

Browse files
dsp-antdsp-ant
committed
refactor: rename root to message
1 parent 4a341a6 commit 1c53fc2

File tree

11 files changed

+72
-40
lines changed

11 files changed

+72
-40
lines changed

‎src/mcp/client/sse.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ async def sse_reader(
9191
case "message":
9292
try:
9393
message = MessageFrame(
94-
root=types.JSONRPCMessage.model_validate_json( # noqa: E501
94+
message=types.JSONRPCMessage.model_validate_json( # noqa: E501
9595
sse.data
9696
),
9797
raw=sse,

‎src/mcp/client/websocket.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ async def ws_reader():
5858
raw_text
5959
)
6060
# Create MessageFrame with JSON message as root
61-
message = MessageFrame(root=json_message, raw=raw_text)
61+
message = MessageFrame(message=json_message, raw=raw_text)
6262
await read_stream_writer.send(message)
6363
except ValidationError as exc:
6464
# If JSON parse or model validation fails, send the exception
@@ -72,7 +72,7 @@ async def ws_writer():
7272
async with write_stream_reader:
7373
async for message in write_stream_reader:
7474
# Extract the JSON-RPC message from MessageFrame and convert to JSON
75-
msg_dict = message.root.model_dump(
75+
msg_dict = message.message.model_dump(
7676
by_alias=True, mode="json", exclude_none=True
7777
)
7878
await ws.send(json.dumps(msg_dict))

‎src/mcp/server/sse.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,4 +176,4 @@ async def handle_post_message(
176176
logger.debug(f"Sending message to writer: {message}")
177177
response = Response("Accepted", status_code=202)
178178
await response(scope, receive, send)
179-
await writer.send(MessageFrame(root=message, raw=request))
179+
await writer.send(MessageFrame(message=message, raw=request))

‎src/mcp/server/stdio.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,9 @@ async def stdin_reader():
7272
await read_stream_writer.send(exc)
7373
continue
7474

75-
await read_stream_writer.send(MessageFrame(root=message, raw=line))
75+
await read_stream_writer.send(
76+
MessageFrame(message=message, raw=line)
77+
)
7678
except anyio.ClosedResourceError:
7779
await anyio.lowlevel.checkpoint()
7880

‎src/mcp/server/websocket.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ async def ws_reader():
4747
continue
4848

4949
await read_stream_writer.send(
50-
MessageFrame(root=client_message, raw=message)
50+
MessageFrame(message=client_message, raw=message)
5151
)
5252
except anyio.ClosedResourceError:
5353
await websocket.close()

‎src/mcp/shared/session.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -249,7 +249,7 @@ async def send_request(
249249
# TODO: Support progress callbacks
250250

251251
await self._write_stream.send(
252-
MessageFrame(root=JSONRPCMessage(jsonrpc_request), raw=None)
252+
MessageFrame(message=JSONRPCMessage(jsonrpc_request), raw=None)
253253
)
254254

255255
try:
@@ -287,7 +287,7 @@ async def send_notification(self, notification: SendNotificationT) -> None:
287287
)
288288

289289
await self._write_stream.send(
290-
MessageFrame(root=JSONRPCMessage(jsonrpc_notification), raw=None)
290+
MessageFrame(message=JSONRPCMessage(jsonrpc_notification), raw=None)
291291
)
292292

293293
async def _send_response(
@@ -296,7 +296,7 @@ async def _send_response(
296296
if isinstance(response, ErrorData):
297297
jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
298298
await self._write_stream.send(
299-
MessageFrame(root=JSONRPCMessage(jsonrpc_error), raw=None)
299+
MessageFrame(message=JSONRPCMessage(jsonrpc_error), raw=None)
300300
)
301301
else:
302302
jsonrpc_response = JSONRPCResponse(
@@ -307,7 +307,7 @@ async def _send_response(
307307
),
308308
)
309309
await self._write_stream.send(
310-
MessageFrame(root=JSONRPCMessage(jsonrpc_response), raw=None)
310+
MessageFrame(message=JSONRPCMessage(jsonrpc_response), raw=None)
311311
)
312312

313313
async def _receive_loop(self) -> None:
@@ -321,7 +321,7 @@ async def _receive_loop(self) -> None:
321321
await self._incoming_message_stream_writer.send(raw_message)
322322
continue
323323

324-
message = raw_message.root
324+
message = raw_message.message
325325
if isinstance(message.root, JSONRPCRequest):
326326
validated_request = self._receive_request_type.model_validate(
327327
message.root.model_dump(

‎src/mcp/types.py‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -184,15 +184,43 @@ class JSONRPCMessage(
184184

185185

186186
class MessageFrame(BaseModel, Generic[RawT]):
187-
root: JSONRPCMessage
187+
"""
188+
A wrapper around the general message received that contains both the parsed message
189+
and the raw message.
190+
191+
This class serves as an encapsulation for JSON-RPC messages, providing access to
192+
both the parsed structure (root) and the original raw data. This design is
193+
particularly useful for Server-Sent Events (SSE) consumers who may need to access
194+
additional metadata or headers associated with the message.
195+
196+
The 'root' attribute contains the parsed JSONRPCMessage, which could be a request,
197+
notification, response, or error. The 'raw' attribute preserves the original
198+
message as received, allowing access to any additional context or metadata that
199+
might be lost in parsing.
200+
201+
This dual representation allows for flexible handling of messages, where consumers
202+
can work with the structured data for standard operations, but still have the
203+
option to examine or utilize the raw data when needed, such as for debugging,
204+
logging, or accessing transport-specific information.
205+
"""
206+
207+
message: JSONRPCMessage
188208
raw: RawT | None = None
189209
model_config = ConfigDict(extra="allow")
190210

191211
def model_dump(self, *args, **kwargs):
192-
return self.root.model_dump(*args, **kwargs)
212+
"""
213+
Dumps the model to a dictionary, delegating to the root JSONRPCMessage.
214+
This method allows for consistent serialization of the parsed message.
215+
"""
216+
return self.message.model_dump(*args, **kwargs)
193217

194218
def model_dump_json(self, *args, **kwargs):
195-
return self.root.model_dump_json(*args, **kwargs)
219+
"""
220+
Dumps the model to a JSON string, delegating to the root JSONRPCMessage.
221+
This method provides a convenient way to serialize the parsed message to JSON.
222+
"""
223+
return self.message.model_dump_json(*args, **kwargs)
196224

197225

198226
class EmptyResult(Result):

‎tests/client/test_session.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,13 @@ async def mock_server():
5858
)
5959

6060
async with server_to_client_send:
61-
assert isinstance(jsonrpc_request.root.root, JSONRPCRequest)
61+
assert isinstance(jsonrpc_request.message.root, JSONRPCRequest)
6262
await server_to_client_send.send(
6363
MessageFrame(
64-
root=JSONRPCMessage(
64+
message=JSONRPCMessage(
6565
JSONRPCResponse(
6666
jsonrpc="2.0",
67-
id=jsonrpc_request.root.root.id,
67+
id=jsonrpc_request.message.root.id,
6868
result=result.model_dump(
6969
by_alias=True, mode="json", exclude_none=True
7070
),
@@ -74,9 +74,9 @@ async def mock_server():
7474
)
7575
)
7676
jsonrpc_notification = await client_to_server_receive.receive()
77-
assert isinstance(jsonrpc_notification.root, JSONRPCMessage)
77+
assert isinstance(jsonrpc_notification.message, JSONRPCMessage)
7878
initialized_notification = ClientNotification.model_validate(
79-
jsonrpc_notification.root.model_dump(
79+
jsonrpc_notification.message.model_dump(
8080
by_alias=True, mode="json", exclude_none=True
8181
)
8282
)

‎tests/issues/test_192_request_id.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ async def run_server():
6666
)
6767

6868
await client_writer.send(
69-
MessageFrame(root=JSONRPCMessage(root=init_req), raw=None)
69+
MessageFrame(message=JSONRPCMessage(root=init_req), raw=None)
7070
)
7171
await server_reader.receive() # Get init response but don't need to check it
7272

@@ -77,7 +77,9 @@ async def run_server():
7777
jsonrpc="2.0",
7878
)
7979
await client_writer.send(
80-
MessageFrame(root=JSONRPCMessage(root=initialized_notification), raw=None)
80+
MessageFrame(
81+
message=JSONRPCMessage(root=initialized_notification), raw=None
82+
)
8183
)
8284

8385
# Send ping request with custom ID
@@ -86,15 +88,15 @@ async def run_server():
8688
)
8789

8890
await client_writer.send(
89-
MessageFrame(root=JSONRPCMessage(root=ping_request), raw=None)
91+
MessageFrame(message=JSONRPCMessage(root=ping_request), raw=None)
9092
)
9193

9294
# Read response
9395
response = await server_reader.receive()
9496

9597
# Verify response ID matches request ID
9698
assert (
97-
response.root.root.id == custom_request_id
99+
response.message.root.id == custom_request_id
98100
), "Response ID should match request ID"
99101

100102
# Cancel server task

‎tests/server/test_lifespan.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ async def run_server():
8484
)
8585
await send_stream1.send(
8686
MessageFrame(
87-
root=JSONRPCMessage(
87+
message=JSONRPCMessage(
8888
root=JSONRPCRequest(
8989
jsonrpc="2.0",
9090
id=1,
@@ -100,7 +100,7 @@ async def run_server():
100100
# Send initialized notification
101101
await send_stream1.send(
102102
MessageFrame(
103-
root=JSONRPCMessage(
103+
message=JSONRPCMessage(
104104
root=JSONRPCNotification(
105105
jsonrpc="2.0",
106106
method="notifications/initialized",
@@ -113,7 +113,7 @@ async def run_server():
113113
# Call the tool to verify lifespan context
114114
await send_stream1.send(
115115
MessageFrame(
116-
root=JSONRPCMessage(
116+
message=JSONRPCMessage(
117117
root=JSONRPCRequest(
118118
jsonrpc="2.0",
119119
id=2,
@@ -127,7 +127,7 @@ async def run_server():
127127

128128
# Get response and verify
129129
response = await receive_stream2.receive()
130-
assert response.root.root.result["content"][0]["text"] == "true"
130+
assert response.message.root.result["content"][0]["text"] == "true"
131131

132132
# Cancel server task
133133
tg.cancel_scope.cancel()
@@ -189,7 +189,7 @@ async def run_server():
189189
)
190190
await send_stream1.send(
191191
MessageFrame(
192-
root=JSONRPCMessage(
192+
message=JSONRPCMessage(
193193
root=JSONRPCRequest(
194194
jsonrpc="2.0",
195195
id=1,
@@ -205,7 +205,7 @@ async def run_server():
205205
# Send initialized notification
206206
await send_stream1.send(
207207
MessageFrame(
208-
root=JSONRPCMessage(
208+
message=JSONRPCMessage(
209209
root=JSONRPCNotification(
210210
jsonrpc="2.0",
211211
method="notifications/initialized",
@@ -218,7 +218,7 @@ async def run_server():
218218
# Call the tool to verify lifespan context
219219
await send_stream1.send(
220220
MessageFrame(
221-
root=JSONRPCMessage(
221+
message=JSONRPCMessage(
222222
root=JSONRPCRequest(
223223
jsonrpc="2.0",
224224
id=2,
@@ -232,7 +232,7 @@ async def run_server():
232232

233233
# Get response and verify
234234
response = await receive_stream2.receive()
235-
assert response.root.root.result["content"][0]["text"] == "true"
235+
assert response.message.root.result["content"][0]["text"] == "true"
236236

237237
# Cancel server task
238238
tg.cancel_scope.cancel()

0 commit comments

Comments
 (0)