mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-23 08:44:22 +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:
committed by
GitHub
parent
ad7f7a5473
commit
9d0f2daddb
@@ -38,7 +38,6 @@ from urllib.parse import quote
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import anyio
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||
from pydantic import ValidationError
|
||||
from sse_starlette import EventSourceResponse
|
||||
from starlette.requests import Request
|
||||
@@ -46,6 +45,13 @@ from starlette.responses import Response
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
import mcp.types as types
|
||||
from mcp.shared.session import (
|
||||
ReadStream,
|
||||
ReadStreamWriter,
|
||||
WriteStream,
|
||||
WriteStreamReader,
|
||||
)
|
||||
from mcp.types import MessageFrame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -63,9 +69,7 @@ class SseServerTransport:
|
||||
"""
|
||||
|
||||
_endpoint: str
|
||||
_read_stream_writers: dict[
|
||||
UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]
|
||||
]
|
||||
_read_stream_writers: dict[UUID, ReadStreamWriter]
|
||||
|
||||
def __init__(self, endpoint: str) -> None:
|
||||
"""
|
||||
@@ -85,11 +89,11 @@ class SseServerTransport:
|
||||
raise ValueError("connect_sse can only handle HTTP requests")
|
||||
|
||||
logger.debug("Setting up SSE connection")
|
||||
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)
|
||||
@@ -172,4 +176,4 @@ class SseServerTransport:
|
||||
logger.debug(f"Sending message to writer: {message}")
|
||||
response = Response("Accepted", status_code=202)
|
||||
await response(scope, receive, send)
|
||||
await writer.send(message)
|
||||
await writer.send(MessageFrame(message=message, raw=request))
|
||||
|
||||
Reference in New Issue
Block a user