Remove unnecessary pragmas and fix generic type propagation

- Remove `# pragma: no branch` and `@runtime_checkable` from stream protocols
- Replace `_CreateContextStreams` with a proper generic class inheriting from
  `tuple[ContextSendStream[T], ContextReceiveStream[T]]`, matching anyio's
  `create_memory_object_stream` pattern so bracket syntax propagates types
- Update `MessageStream` type alias to match actual stream creation

Author: Kludex

src/mcp/shared/_context_streams.py

@@ -103,26 +103,17 @@ async def __aexit__(
         return None
 
 
-def _create_context_streams(
-    max_buffer_size: float = 0,
-) -> tuple[ContextSendStream[Any], ContextReceiveStream[Any]]:
-    raw_send: MemoryObjectSendStream[Any]
-    raw_receive: MemoryObjectReceiveStream[Any]
-    raw_send, raw_receive = anyio.create_memory_object_stream(max_buffer_size)
-    return ContextSendStream(raw_send), ContextReceiveStream(raw_receive)
+class create_context_streams(
+    tuple[ContextSendStream[T], ContextReceiveStream[T]],
+):
+    """Create context-aware memory object streams.
 
-
-class _CreateContextStreams:
-    """Callable that supports ``create_context_streams[T](n)`` bracket syntax.
-
-    Matches anyio's ``create_memory_object_stream`` API style.
+    Supports ``create_context_streams[T](n)`` bracket syntax,
+    matching anyio's ``create_memory_object_stream`` API style.
     """
 
-    def __getitem__(self, _item: Any) -> _CreateContextStreams:
-        return self
-
-    def __call__(self, max_buffer_size: float = 0) -> tuple[ContextSendStream[Any], ContextReceiveStream[Any]]:
-        return _create_context_streams(max_buffer_size)
-
-
-create_context_streams = _CreateContextStreams()
+    def __new__(cls, max_buffer_size: float = 0) -> tuple[ContextSendStream[T], ContextReceiveStream[T]]:  # type: ignore[type-var]
+        raw_send: MemoryObjectSendStream[Any]
+        raw_receive: MemoryObjectReceiveStream[Any]
+        raw_send, raw_receive = anyio.create_memory_object_stream(max_buffer_size)
+        return (ContextSendStream(raw_send), ContextReceiveStream(raw_receive))

src/mcp/shared/_stream_protocols.py

@@ -7,43 +7,41 @@
 from __future__ import annotations
 
 from types import TracebackType
-from typing import Protocol, TypeVar, runtime_checkable
+from typing import Protocol, TypeVar
 
 from typing_extensions import Self
 
 T_co = TypeVar("T_co", covariant=True)
 T_contra = TypeVar("T_contra", contravariant=True)
 
 
-@runtime_checkable
-class ReadStream(Protocol[T_co]):  # pragma: no branch
+class ReadStream(Protocol[T_co]):
     """Protocol for reading items from a stream.
 
     Consumers that need the sender's context should use
     ``getattr(stream, 'last_context', None)``.
     """
 
-    async def receive(self) -> T_co: ...  # pragma: no branch
-    async def aclose(self) -> None: ...  # pragma: no branch
-    def __aiter__(self) -> ReadStream[T_co]: ...  # pragma: no branch
-    async def __anext__(self) -> T_co: ...  # pragma: no branch
-    async def __aenter__(self) -> Self: ...  # pragma: no branch
-    async def __aexit__(  # pragma: no branch
+    async def receive(self) -> T_co: ...
+    async def aclose(self) -> None: ...
+    def __aiter__(self) -> ReadStream[T_co]: ...
+    async def __anext__(self) -> T_co: ...
+    async def __aenter__(self) -> Self: ...
+    async def __aexit__(
         self,
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> bool | None: ...
 
 
-@runtime_checkable
-class WriteStream(Protocol[T_contra]):  # pragma: no branch
+class WriteStream(Protocol[T_contra]):
     """Protocol for writing items to a stream."""
 
-    async def send(self, item: T_contra, /) -> None: ...  # pragma: no branch
-    async def aclose(self) -> None: ...  # pragma: no branch
-    async def __aenter__(self) -> Self: ...  # pragma: no branch
-    async def __aexit__(  # pragma: no branch
+    async def send(self, item: T_contra, /) -> None: ...
+    async def aclose(self) -> None: ...
+    async def __aenter__(self) -> Self: ...
+    async def __aexit__(
         self,
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,

src/mcp/shared/memory.py

@@ -8,7 +8,7 @@
 from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
 from mcp.shared.message import SessionMessage
 
-MessageStream = tuple[ContextReceiveStream[SessionMessage | Exception], ContextSendStream[SessionMessage]]
+MessageStream = tuple[ContextReceiveStream[SessionMessage | Exception], ContextSendStream[SessionMessage | Exception]]
 
 
 @asynccontextmanager