refactor: reorganize message handling for better type safety and clarity (#239)
Some checks failed
Main branch checks / checks (push) Has been cancelled

* refactor: improve typing with memory stream type aliases

Move memory stream type definitions to models.py and use them throughout
the codebase for better type safety and maintainability.

GitHub-Issue:#201

* refactor: move streams to ParsedMessage

* refactor: update test files to use ParsedMessage

Updates test files to work with the ParsedMessage stream type aliases
and fixes a line length issue in test_201_client_hangs_on_logging.py.

Github-Issue:#201

* refactor: rename ParsedMessage to MessageFrame for clarity

🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>

* refactor: move MessageFrame class to types.py for better code organization

🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>

* fix pyright

* refactor: update websocket client to use MessageFrame

Modified the websocket client to work with the new MessageFrame type,
preserving raw message text and properly extracting the root JSON-RPC
message when sending.

Github-Issue:#204

* fix: use NoneType instead of None for type parameters in MessageFrame

🤖 Generated with [Claude Code](https://claude.ai/code)

* refactor: rename root to message
This commit is contained in:
David Soria Parra 2025-03-13 13:44:55 +00:00 • committed by GitHub
commit 9d0f2daddb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -11,6 +11,7 @@ from mcp.types import (
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
MessageFrame,
NotificationParams,
)
@ -64,7 +65,9 @@ async def test_request_id_match() -> None:
jsonrpc="2.0",
)
await client_writer.send(JSONRPCMessage(root=init_req))
await client_writer.send(
MessageFrame(message=JSONRPCMessage(root=init_req), raw=None)
)
await server_reader.receive() # Get init response but don't need to check it
# Send initialized notification
@ -73,21 +76,27 @@ async def test_request_id_match() -> None:
params=NotificationParams().model_dump(by_alias=True, exclude_none=True),
jsonrpc="2.0",
)
await client_writer.send(JSONRPCMessage(root=initialized_notification))
await client_writer.send(
MessageFrame(
message=JSONRPCMessage(root=initialized_notification), raw=None
)
)
# Send ping request with custom ID
ping_request = JSONRPCRequest(
id=custom_request_id, method="ping", params={}, jsonrpc="2.0"
)
await client_writer.send(JSONRPCMessage(root=ping_request))
await client_writer.send(
MessageFrame(message=JSONRPCMessage(root=ping_request), raw=None)
)
# Read response
response = await server_reader.receive()
# Verify response ID matches request ID
assert (
response.root.id == custom_request_id
response.message.root.id == custom_request_id
), "Response ID should match request ID"
# Cancel server task