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

@ -24,9 +24,15 @@ from io import TextIOWrapper
import anyio
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
import mcp.types as types
from mcp.shared.session import (
ReadStream,
ReadStreamWriter,
WriteStream,
WriteStreamReader,
)
from mcp.types import MessageFrame
@asynccontextmanager
@ -47,11 +53,11 @@ async def stdio_server(
if not stdout:
stdout = anyio.wrap_file(TextIOWrapper(sys.stdout.buffer, encoding="utf-8"))
read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
read_stream: ReadStream
read_stream_writer: ReadStreamWriter
write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
write_stream: WriteStream
write_stream_reader: WriteStreamReader
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@ -66,7 +72,9 @@ async def stdio_server(
await read_stream_writer.send(exc)
continue
await read_stream_writer.send(message)
await read_stream_writer.send(
MessageFrame(message=message, raw=line)
)
except anyio.ClosedResourceError:
await anyio.lowlevel.checkpoint()
@ -74,6 +82,7 @@ async def stdio_server(
try:
async with write_stream_reader:
async for message in write_stream_reader:
# Extract the inner JSONRPCRequest/JSONRPCResponse from MessageFrame
json = message.model_dump_json(by_alias=True, exclude_none=True)
await stdout.write(json + "\n")
await stdout.flush()