Merge branch 'main' into fix/stateless-task-group-leak

Author: wiggzz

.github/workflows/claude-code-review.yml

@@ -14,7 +14,7 @@ on:
 jobs:
   claude:
     if: |
-      (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude') && !startsWith(github.event.comment.body, '@claude review')) ||
       (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
       (github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
       (github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))

.github/workflows/claude.yml

@@ -18,7 +18,7 @@
 from urllib.parse import parse_qs, urlparse
 
 import httpx
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.client._transport import ReadStream, WriteStream
 from mcp.client.auth import OAuthClientProvider, TokenStorage
 from mcp.client.session import ClientSession
 from mcp.client.sse import sse_client
@@ -241,8 +241,8 @@ async def _default_redirect_handler(authorization_url: str) -> None:
 
     async def _run_session(
         self,
-        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
-        write_stream: MemoryObjectSendStream[SessionMessage],
+        read_stream: ReadStream[SessionMessage | Exception],
+        write_stream: WriteStream[SessionMessage],
     ):
         """Run the MCP session with the given streams."""
         print("🤝 Initializing MCP session...")

examples/clients/simple-auth-client/mcp_simple_auth_client/main.py

@@ -26,7 +26,7 @@ classifiers = [
 ]
 dependencies = [
     "anyio>=4.9",
-    "httpx>=0.27.1",
+    "httpx>=0.27.1,<1.0.0",
     "httpx-sse>=0.4",
     "pydantic>=2.12.0",
     "starlette>=0.48.0; python_version >= '3.14'",
@@ -40,6 +40,7 @@ dependencies = [
     "pyjwt[crypto]>=2.10.1",
     "typing-extensions>=4.13.0",
     "typing-inspection>=0.4.1",
+    "opentelemetry-api>=1.28.0",
 ]
 
 [project.optional-dependencies]
@@ -71,6 +72,7 @@ dev = [
     "coverage[toml]>=7.10.7,<=7.13",
     "pillow>=12.0",
     "strict-no-cover",
+    "logfire>=3.0.0",
 ]
 docs = [
     "mkdocs>=1.6.1",
@@ -219,13 +221,10 @@ skip_covered = true
 show_missing = true
 ignore_errors = true
 precision = 2
-exclude_lines = [
-    "pragma: no cover",
+exclude_also = [
     "pragma: lax no cover",
-    "if TYPE_CHECKING:",
     "@overload",
     "raise NotImplementedError",
-    "^\\s*\\.\\.\\.\\s*$",
 ]
 
 # https://coverage.readthedocs.io/en/latest/config.html#paths

pyproject.toml

@@ -6,9 +6,9 @@
 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(

src/mcp/client/__main__.py

@@ -5,11 +5,12 @@
 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):

src/mcp/client/_transport.py

@@ -4,10 +4,10 @@
 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,

src/mcp/client/session.py

@@ -7,11 +7,11 @@
 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_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
             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 post_writer(endpoint_url: str):
                             )
                             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")
 

src/mcp/client/sse.py

@@ -11,11 +11,11 @@
 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 @@
 
 # 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 @@ async def post_writer(
         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 @@ async def handle_request_async():
                     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,