mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 06:54:18 +01:00
refactor: reorganize message handling for better type safety and clarity (#239)
* 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
committed by
GitHub
GitHub
parent
ad7f7a5473
commit
9d0f2daddb
@@ -17,6 +17,7 @@ from mcp.types import (
|
||||
JSONRPCMessage,
|
||||
JSONRPCNotification,
|
||||
JSONRPCRequest,
|
||||
MessageFrame,
|
||||
)
|
||||
|
||||
|
||||
@@ -64,7 +65,7 @@ async def test_lowlevel_server_lifespan():
|
||||
send_stream2,
|
||||
InitializationOptions(
|
||||
server_name="test",
|
||||
server_version="0.1.0",
|
||||
server_version="1.0.0",
|
||||
capabilities=server.get_capabilities(
|
||||
notification_options=NotificationOptions(),
|
||||
experimental_capabilities={},
|
||||
@@ -82,42 +83,51 @@ async def test_lowlevel_server_lifespan():
|
||||
clientInfo=Implementation(name="test-client", version="0.1.0"),
|
||||
)
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
MessageFrame(
|
||||
message=JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
),
|
||||
raw=None,
|
||||
)
|
||||
)
|
||||
response = await receive_stream2.receive()
|
||||
|
||||
# Send initialized notification
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
MessageFrame(
|
||||
message=JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
),
|
||||
raw=None,
|
||||
)
|
||||
)
|
||||
|
||||
# Call the tool to verify lifespan context
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
MessageFrame(
|
||||
message=JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
),
|
||||
raw=None,
|
||||
)
|
||||
)
|
||||
|
||||
# Get response and verify
|
||||
response = await receive_stream2.receive()
|
||||
assert response.root.result["content"][0]["text"] == "true"
|
||||
assert response.message.root.result["content"][0]["text"] == "true"
|
||||
|
||||
# Cancel server task
|
||||
tg.cancel_scope.cancel()
|
||||
@@ -178,42 +188,51 @@ async def test_fastmcp_server_lifespan():
|
||||
clientInfo=Implementation(name="test-client", version="0.1.0"),
|
||||
)
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
MessageFrame(
|
||||
message=JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=1,
|
||||
method="initialize",
|
||||
params=TypeAdapter(InitializeRequestParams).dump_python(params),
|
||||
)
|
||||
),
|
||||
raw=None,
|
||||
)
|
||||
)
|
||||
response = await receive_stream2.receive()
|
||||
|
||||
# Send initialized notification
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
MessageFrame(
|
||||
message=JSONRPCMessage(
|
||||
root=JSONRPCNotification(
|
||||
jsonrpc="2.0",
|
||||
method="notifications/initialized",
|
||||
)
|
||||
),
|
||||
raw=None,
|
||||
)
|
||||
)
|
||||
|
||||
# Call the tool to verify lifespan context
|
||||
await send_stream1.send(
|
||||
JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
MessageFrame(
|
||||
message=JSONRPCMessage(
|
||||
root=JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=2,
|
||||
method="tools/call",
|
||||
params={"name": "check_lifespan", "arguments": {}},
|
||||
)
|
||||
),
|
||||
raw=None,
|
||||
)
|
||||
)
|
||||
|
||||
# Get response and verify
|
||||
response = await receive_stream2.receive()
|
||||
assert response.root.result["content"][0]["text"] == "true"
|
||||
assert response.message.root.result["content"][0]["text"] == "true"
|
||||
|
||||
# Cancel server task
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
Reference in New Issue
Block a user