modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Wrap JSONRPC messages with SessionMessage for metadata support (#590)

Browse source
This commit is contained in:
ihrpr 2025-05-02 14:29:00 +01:00 • committed by GitHub
commit da0cf22355
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 286 additions and 173 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/mcp/server/lowlevel/server.py
View file

@ -84,6 +84,7 @@ from mcp.server.session import ServerSession
from mcp.server.stdio import stdio_server as stdio_server
from mcp.shared.context import RequestContext
from mcp.shared.exceptions import McpError
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
logger = logging.getLogger(__name__)
@ -471,8 +472,8 @@ class Server(Generic[LifespanResultT]):
async def run(
self,
read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
initialization_options: InitializationOptions,
# When False, exceptions are returned as messages to the client.
# When True, exceptions are raised, which will cause the server to shut down

5
src/mcp/server/session.py
View file

@ -47,6 +47,7 @@ from pydantic import AnyUrl
import mcp.types as types
from mcp.server.models import InitializationOptions
from mcp.shared.message import SessionMessage
from mcp.shared.session import (
BaseSession,
RequestResponder,
@ -82,8 +83,8 @@ class ServerSession(
def __init__(
self,
read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
init_options: InitializationOptions,
stateless: bool = False,
) -> None:

24
src/mcp/server/sse.py
View file

@ -46,6 +46,7 @@ from starlette.responses import Response
from starlette.types import Receive, Scope, Send
import mcp.types as types
from mcp.shared.message import SessionMessage
logger = logging.getLogger(__name__)
@ -63,9 +64,7 @@ class SseServerTransport:
"""
_endpoint: str
_read_stream_writers: dict[
UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]
]
_read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]]
def __init__(self, endpoint: str) -> None:
"""
@ -85,11 +84,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: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@ -109,12 +108,12 @@ class SseServerTransport:
await sse_stream_writer.send({"event": "endpoint", "data": session_uri})
logger.debug(f"Sent endpoint event: {session_uri}")
async for message in write_stream_reader:
logger.debug(f"Sending message via SSE: {message}")
async for session_message in write_stream_reader:
logger.debug(f"Sending message via SSE: {session_message}")
await sse_stream_writer.send(
{
"event": "message",
"data": message.model_dump_json(
"data": session_message.message.model_dump_json(
by_alias=True, exclude_none=True
),
}
@ -169,7 +168,8 @@ class SseServerTransport:
await writer.send(err)
return
logger.debug(f"Sending message to writer: {message}")
session_message = SessionMessage(message)
logger.debug(f"Sending session message to writer: {session_message}")
response = Response("Accepted", status_code=202)
await response(scope, receive, send)
await writer.send(message)
await writer.send(session_message)

18
src/mcp/server/stdio.py
View file

@ -27,6 +27,7 @@ import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
import mcp.types as types
from mcp.shared.message import SessionMessage
@asynccontextmanager
@ -47,11 +48,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: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@ -66,15 +67,18 @@ async def stdio_server(
await read_stream_writer.send(exc)
continue
await read_stream_writer.send(message)
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
except anyio.ClosedResourceError:
await anyio.lowlevel.checkpoint()
async def stdout_writer():
try:
async with write_stream_reader:
async for message in write_stream_reader:
json = message.model_dump_json(by_alias=True, exclude_none=True)
async for session_message in write_stream_reader:
json = session_message.message.model_dump_json(
by_alias=True, exclude_none=True
)
await stdout.write(json + "\n")
await stdout.flush()
except anyio.ClosedResourceError:

25
src/mcp/server/streamable_http.py
View file

@ -24,6 +24,7 @@ from starlette.requests import Request
from starlette.responses import Response
from starlette.types import Receive, Scope, Send
from mcp.shared.message import SessionMessage
from mcp.types import (
INTERNAL_ERROR,
INVALID_PARAMS,
@ -125,10 +126,10 @@ class StreamableHTTPServerTransport:
"""
# Server notification streams for POST requests as well as standalone SSE stream
_read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception] | None = (
_read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception] | None = (
None
)
_write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage] | None = None
_write_stream_reader: MemoryObjectReceiveStream[SessionMessage] | None = None
def __init__(
self,
@ -378,7 +379,8 @@ class StreamableHTTPServerTransport:
await response(scope, receive, send)
# Process the message after sending the response
await writer.send(message)
session_message = SessionMessage(message)
await writer.send(session_message)
return
@ -394,7 +396,8 @@ class StreamableHTTPServerTransport:
if self.is_json_response_enabled:
# Process the message
await writer.send(message)
session_message = SessionMessage(message)
await writer.send(session_message)
try:
# Process messages from the request-specific stream
# We need to collect all messages until we get a response
@ -500,7 +503,8 @@ class StreamableHTTPServerTransport:
async with anyio.create_task_group() as tg:
tg.start_soon(response, scope, receive, send)
# Then send the message to be processed by the server
await writer.send(message)
session_message = SessionMessage(message)
await writer.send(session_message)
except Exception:
logger.exception("SSE response error")
# Clean up the request stream if something goes wrong
@ -792,8 +796,8 @@ class StreamableHTTPServerTransport:
self,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[JSONRPCMessage | Exception],
MemoryObjectSendStream[JSONRPCMessage],
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
],
None,
]:
@ -806,10 +810,10 @@ class StreamableHTTPServerTransport:
# Create the memory streams for this connection
read_stream_writer, read_stream = anyio.create_memory_object_stream[
JSONRPCMessage | Exception
SessionMessage | Exception
](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[
JSONRPCMessage
SessionMessage
](0)
# Store the streams
@ -821,8 +825,9 @@ class StreamableHTTPServerTransport:
# Create a message router that distributes messages to request streams
async def message_router():
try:
async for message in write_stream_reader:
async for session_message in write_stream_reader:
# Determine which request stream(s) should receive this message
message = session_message.message
target_request_id = None
if isinstance(
message.root, JSONRPCNotification | JSONRPCRequest

18
src/mcp/server/websocket.py
View file

@ -8,6 +8,7 @@ from starlette.types import Receive, Scope, Send
from starlette.websockets import WebSocket
import mcp.types as types
from mcp.shared.message import SessionMessage
logger = logging.getLogger(__name__)
@ -22,11 +23,11 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
websocket = WebSocket(scope, receive, send)
await websocket.accept(subprotocol="mcp")
read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@ -41,15 +42,18 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
await read_stream_writer.send(exc)
continue
await read_stream_writer.send(client_message)
session_message = SessionMessage(client_message)
await read_stream_writer.send(session_message)
except anyio.ClosedResourceError:
await websocket.close()
async def ws_writer():
try:
async with write_stream_reader:
async for message in write_stream_reader:
obj = message.model_dump_json(by_alias=True, exclude_none=True)
async for session_message in write_stream_reader:
obj = session_message.message.model_dump_json(
by_alias=True, exclude_none=True
)
await websocket.send_text(obj)
except anyio.ClosedResourceError:
await websocket.close()