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

tests/client/conftest.py

@@ -4,15 +4,15 @@ from typing import Any
 from unittest.mock import patch
 
 import pytest
-from anyio.streams.memory import MemoryObjectSendStream
 
 import mcp.shared.memory
+from mcp.client._transport import WriteStream
 from mcp.shared.message import SessionMessage
 from mcp.types import JSONRPCNotification, JSONRPCRequest
 
 
 class SpyMemoryObjectSendStream:
-    def __init__(self, original_stream: MemoryObjectSendStream[SessionMessage]):
+    def __init__(self, original_stream: WriteStream[SessionMessage]):
         self.original_stream = original_stream
         self.sent_messages: list[SessionMessage] = []
 

tests/client/test_client.py

@@ -2,6 +2,9 @@
 
 from __future__ import annotations
 
+import contextvars
+from collections.abc import Iterator
+from contextlib import contextmanager
 from unittest.mock import patch
 
 import anyio
@@ -320,3 +323,33 @@ async def test_client_uses_transport_directly(app: MCPServer):
                 structured_content={"result": "Hello, Transport!"},
             )
         )
+
+
+_TEST_CONTEXTVAR = contextvars.ContextVar("test_var", default="initial")
+
+
+@contextmanager
+def _set_test_contextvar(value: str) -> Iterator[None]:
+    token = _TEST_CONTEXTVAR.set(value)
+    try:
+        yield
+    finally:
+        _TEST_CONTEXTVAR.reset(token)
+
+
+async def test_context_propagation():
+    """Sender's contextvars.Context is propagated to the server handler."""
+    server = MCPServer("test")
+
+    @server.tool()
+    async def check_context() -> str:
+        """Return the contextvar value visible to the handler."""
+        return _TEST_CONTEXTVAR.get()
+
+    async with Client(server) as client:
+        with _set_test_contextvar("client_value"):
+            result = await client.call_tool("check_context", {})
+
+    assert result.content[0].text == "client_value", (  # type: ignore[union-attr]
+        "Server handler did not see the sender's contextvars.Context"
+    )

tests/shared/test_streamable_http.py

@@ -45,6 +45,7 @@ from mcp.server.streamable_http import (
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared._context import RequestContext
+from mcp.shared._context_streams import create_context_streams
 from mcp.shared._httpx_utils import (
     MCP_DEFAULT_SSE_READ_TIMEOUT,
     MCP_DEFAULT_TIMEOUT,
@@ -1783,8 +1784,8 @@ async def test_handle_sse_event_skips_empty_data():
     # Create a mock SSE event with empty data (keep-alive ping)
     mock_sse = ServerSentEvent(event="message", data="", id=None, retry=None)
 
-    # Create a mock stream writer
-    write_stream, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+    # Create a context-aware stream writer (matches StreamWriter type alias)
+    write_stream, read_stream = create_context_streams[SessionMessage | Exception](1)
 
     try:
         # Call _handle_sse_event with empty data - should return False and not raise
@@ -1794,8 +1795,9 @@ async def test_handle_sse_event_skips_empty_data():
         assert result is False
 
         # Nothing should have been written to the stream
-        # Check buffer is empty (statistics().current_buffer_used returns buffer size)
-        assert write_stream.statistics().current_buffer_used == 0
+        with pytest.raises(TimeoutError):
+            with anyio.fail_after(0):
+                await read_stream.receive()
     finally:
         await write_stream.aclose()
         await read_stream.aclose()