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
5 changed files with 310 additions and 94 deletions
Download patch file Download diff file Expand all files Collapse all files

6
src/mcp/client/__main__.py
View file

@ -6,9 +6,9 @@ from functools import partial
from urllib.parse import urlparse
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import types
from mcp.client._transport import ReadStream, WriteStream
from mcp.client.session import ClientSession
from mcp.client.sse import sse_client
from mcp.client.stdio import StdioServerParameters, stdio_client
@ -33,8 +33,8 @@ async def message_handler(
async def run_session(
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
client_info: types.Implementation | None = None,
):
async with ClientSession(

7
src/mcp/client/_transport.py
View file

@ -5,11 +5,12 @@ from __future__ import annotations
from contextlib import AbstractAsyncContextManager
from typing import Protocol
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import SessionMessage
TransportStreams = tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
__all__ = ["ReadStream", "WriteStream", "Transport", "TransportStreams"]
TransportStreams = tuple[ReadStream[SessionMessage | Exception], WriteStream[SessionMessage]]
class Transport(AbstractAsyncContextManager[TransportStreams], Protocol):

6
src/mcp/client/session.py
View file

@ -4,10 +4,10 @@ import logging
from typing import Any, Protocol
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import TypeAdapter
from mcp import types
from mcp.client._transport import ReadStream, WriteStream
from mcp.client.experimental import ExperimentalClientFeatures
from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
from mcp.shared._context import RequestContext
@ -109,8 +109,8 @@ class ClientSession(
):
def __init__(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
read_timeout_seconds: float | None = None,
sampling_callback: SamplingFnT | None = None,
elicitation_callback: ElicitationFnT | None = None,

23
src/mcp/client/sse.py
View file

@ -7,11 +7,11 @@ from urllib.parse import parse_qs, urljoin, urlparse
import anyio
import httpx
from anyio.abc import TaskStatus
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx_sse import aconnect_sse
from httpx_sse._exceptions import SSEError
from mcp import types
from mcp.shared._context_streams import create_context_streams
from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
from mcp.shared.message import SessionMessage
@ -51,12 +51,6 @@ async def sse_client(
auth: Optional HTTPX authentication handler.
on_session_created: Optional callback invoked with the session ID when received.
"""
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}")
async with httpx_client_factory(
headers=headers, auth=auth, timeout=httpx.Timeout(timeout, read=sse_read_timeout)
@ -65,8 +59,8 @@ async def sse_client(
event_source.response.raise_for_status()
logger.debug("SSE connection established")
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 sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
try:
@ -124,7 +118,8 @@ async def sse_client(
async def post_writer(endpoint_url: str):
try:
async with write_stream_reader, write_stream:
async for session_message in write_stream_reader:
async def _send_message(session_message: SessionMessage) -> None:
logger.debug(f"Sending client message: {session_message}")
response = await client.post(
endpoint_url,
@ -136,6 +131,14 @@ async def sse_client(
)
response.raise_for_status()
logger.debug(f"Client message sent successfully: {response.status_code}")
async for session_message in write_stream_reader:
sender_ctx = write_stream_reader.last_context
if sender_ctx is not None:
async with anyio.create_task_group() as tg:
sender_ctx.run(tg.start_soon, _send_message, session_message)
else:
await _send_message(session_message) # pragma: no cover
except Exception: # pragma: lax no cover
logger.exception("Error in post_writer")

23
src/mcp/client/streamable_http.py
View file

@ -11,11 +11,11 @@ from dataclasses import dataclass
import anyio
import httpx
from anyio.abc import TaskGroup
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
from pydantic import ValidationError
from mcp.client._transport import TransportStreams
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
from mcp.shared._httpx_utils import create_mcp_http_client
from mcp.shared.message import ClientMessageMetadata, SessionMessage
from mcp.types import (
@ -38,8 +38,8 @@ logger = logging.getLogger(__name__)
# TODO(Marcelo): Put the TransportStreams in a module under shared, so we can import here.
SessionMessageOrError = SessionMessage | Exception
StreamWriter = MemoryObjectSendStream[SessionMessageOrError]
StreamReader = MemoryObjectReceiveStream[SessionMessage]
StreamWriter = ContextSendStream[SessionMessageOrError]
StreamReader = ContextReceiveStream[SessionMessage]
MCP_SESSION_ID = "mcp-session-id"
MCP_PROTOCOL_VERSION = "mcp-protocol-version"
@ -434,14 +434,15 @@ class StreamableHTTPTransport:
client: httpx.AsyncClient,
write_stream_reader: StreamReader,
read_stream_writer: StreamWriter,
write_stream: MemoryObjectSendStream[SessionMessage],
write_stream: ContextSendStream[SessionMessage],
start_get_stream: Callable[[], None],
tg: TaskGroup,
) -> None:
"""Handle writing requests to the server."""
try:
async with write_stream_reader, read_stream_writer, write_stream:
async for session_message in write_stream_reader:
async def _handle_message(session_message: SessionMessage) -> None:
message = session_message.message
metadata = (
session_message.metadata
@ -478,6 +479,14 @@ class StreamableHTTPTransport:
else:
await handle_request_async()
async for session_message in write_stream_reader:
sender_ctx = write_stream_reader.last_context
if sender_ctx is not None:
async with anyio.create_task_group() as tg_local:
sender_ctx.run(tg_local.start_soon, _handle_message, session_message)
else:
await _handle_message(session_message) # pragma: no cover
except Exception: # pragma: lax no cover
logger.exception("Error in post_writer")
@ -547,8 +556,8 @@ async def streamable_http_client(
if not client_provided:
await stack.enter_async_context(client)
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)
async with (
read_stream_writer,