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,11 +11,11 @@ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStre
|
|||
|
||||
from mcp.client.session import ClientSession, ListRootsFnT, SamplingFnT
|
||||
from mcp.server import Server
|
||||
from mcp.types import JSONRPCMessage
|
||||
from mcp.types import MessageFrame
|
||||
|
||||
MessageStream = tuple[
|
||||
MemoryObjectReceiveStream[JSONRPCMessage | Exception],
|
||||
MemoryObjectSendStream[JSONRPCMessage],
|
||||
MemoryObjectReceiveStream[MessageFrame | Exception],
|
||||
MemoryObjectSendStream[MessageFrame],
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -32,10 +32,10 @@ async def create_client_server_memory_streams() -> (
|
|||
"""
|
||||
# Create streams for both directions
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
|
||||
JSONRPCMessage | Exception
|
||||
MessageFrame | Exception
|
||||
](1)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
|
||||
JSONRPCMessage | Exception
|
||||
MessageFrame | Exception
|
||||
](1)
|
||||
|
||||
client_streams = (server_to_client_receive, client_to_server_send)
|
||||
|
|
@ -60,12 +60,9 @@ async def create_connected_server_and_client_session(
|
|||
) -> AsyncGenerator[ClientSession, None]:
|
||||
"""Creates a ClientSession that is connected to a running MCP server."""
|
||||
async with create_client_server_memory_streams() as (
|
||||
client_streams,
|
||||
server_streams,
|
||||
(client_read, client_write),
|
||||
(server_read, server_write),
|
||||
):
|
||||
client_read, client_write = client_streams
|
||||
server_read, server_write = server_streams
|
||||
|
||||
# Create a cancel scope for the server task
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(
|
||||
|
|
|
|||
Loading…
Reference in a new issue