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

Propagate contextvars.Context through anyio streams without modifying SessionMessage (#2298)

Browse source
This commit is contained in:
Marcelo Trylesinski 2026-03-31 12:49:38 -04:00 • committed by GitHub
commit e6235d1667
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 310 additions and 94 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -36,6 +36,7 @@ handler callables by method string.
from __future__ import annotations
import contextvars
import logging
import warnings
from collections.abc import AsyncIterator, Awaitable, Callable
@ -44,7 +45,6 @@ from importlib.metadata import version as importlib_version
from typing import Any, Generic
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.authentication import AuthenticationMiddleware
@ -65,6 +65,7 @@ from mcp.server.session import ServerSession
from mcp.server.streamable_http import EventStore
from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
from mcp.server.transport_security import TransportSecuritySettings
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.exceptions import MCPError
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.shared.session import RequestResponder
@ -355,8 +356,8 @@ class Server(Generic[LifespanResultT]):
async def run(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[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
@ -391,7 +392,13 @@ class Server(Generic[LifespanResultT]):
async for message in session.incoming_messages:
logger.debug("Received message: %s", message)
tg.start_soon(
if isinstance(message, RequestResponder) and message.context is not None:
context = message.context
else:
context = contextvars.copy_context()
context.run(
tg.start_soon,
self._handle_message,
message,
session,

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

@ -33,13 +33,14 @@ from typing import Any, TypeVar, overload
import anyio
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from anyio.streams.memory import MemoryObjectReceiveStream
from pydantic import AnyUrl, TypeAdapter
from mcp import types
from mcp.server.experimental.session_features import ExperimentalServerSessionFeatures
from mcp.server.models import InitializationOptions
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.exceptions import StatelessModeNotSupported
from mcp.shared.experimental.tasks.capabilities import check_tasks_capability
from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY
@ -79,8 +80,8 @@ class ServerSession(
def __init__(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
init_options: InitializationOptions,
stateless: bool = False,
) -> None:

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

@ -43,7 +43,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
@ -55,6 +54,7 @@ from mcp.server.transport_security import (
TransportSecurityMiddleware,
TransportSecuritySettings,
)
from mcp.shared._context_streams import ContextSendStream, create_context_streams
from mcp.shared.message import ServerMessageMetadata, SessionMessage
logger = logging.getLogger(__name__)
@ -72,7 +72,7 @@ class SseServerTransport:
"""
_endpoint: str
_read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]]
_read_stream_writers: dict[UUID, ContextSendStream[SessionMessage | Exception]]
_security: TransportSecurityMiddleware
def __init__(self, endpoint: str, security_settings: TransportSecuritySettings | None = None) -> None:
@ -129,14 +129,9 @@ class SseServerTransport:
raise ValueError("Request validation failed")
logger.debug("Setting up SSE connection")
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
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)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
session_id = uuid4()
self._read_stream_writers[session_id] = read_stream_writer

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

@ -23,9 +23,9 @@ from io import TextIOWrapper
import anyio
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import types
from mcp.shared._context_streams import create_context_streams
from mcp.shared.message import SessionMessage
@ -43,14 +43,8 @@ async def stdio_server(stdin: anyio.AsyncFile[str] | None = None, stdout: anyio.
if not stdout:
stdout = anyio.wrap_file(TextIOWrapper(sys.stdout.buffer, encoding="utf-8"))
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
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)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
async def stdin_reader():
try:

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

@ -25,6 +25,8 @@ from starlette.responses import Response
from starlette.types import Receive, Scope, Send
from mcp.server.transport_security import TransportSecurityMiddleware, TransportSecuritySettings
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
from mcp.types import (
@ -119,10 +121,10 @@ class StreamableHTTPServerTransport:
"""
# Server notification streams for POST requests as well as standalone SSE stream
_read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception] | None = None
_read_stream: MemoryObjectReceiveStream[SessionMessage | Exception] | None = None
_write_stream: MemoryObjectSendStream[SessionMessage] | None = None
_write_stream_reader: MemoryObjectReceiveStream[SessionMessage] | None = None
_read_stream_writer: ContextSendStream[SessionMessage | Exception] | None = None
_read_stream: ContextReceiveStream[SessionMessage | Exception] | None = None
_write_stream: ContextSendStream[SessionMessage] | None = None
_write_stream_reader: ContextReceiveStream[SessionMessage] | None = None
_security: TransportSecurityMiddleware
def __init__(
@ -954,8 +956,8 @@ class StreamableHTTPServerTransport:
self,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
ReadStream[SessionMessage | Exception],
WriteStream[SessionMessage],
],
None,
]:
@ -967,8 +969,8 @@ class StreamableHTTPServerTransport:
# Create the memory streams for this connection
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
# Store the streams
self._read_stream_writer = read_stream_writer

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

@ -1,12 +1,12 @@
from contextlib import asynccontextmanager
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic_core import ValidationError
from starlette.types import Receive, Scope, Send
from starlette.websockets import WebSocket
from mcp import types
from mcp.shared._context_streams import create_context_streams
from mcp.shared.message import SessionMessage
@ -19,14 +19,8 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
websocket = WebSocket(scope, receive, send)
await websocket.accept(subprotocol="mcp")
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
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)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
async def ws_reader():
try: