refactor: reorganize message handling for better type safety and clarity (#239)
Some checks failed
Main branch checks / checks (push) Has been cancelled
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:
parent
ad7f7a5473
commit
9d0f2daddb
1 changed files with 283 additions and 151 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue