Support different transports in Client

Author: Kludex

README.md

@@ -2213,11 +2213,7 @@ from mcp.client.streamable_http import streamable_http_client
 
 async def main():
     # Connect to a streamable HTTP server
-    async with streamable_http_client("http://localhost:8000/mcp") as (
-        read_stream,
-        write_stream,
-        _,
-    ):
+    async with streamable_http_client("http://localhost:8000/mcp") as (read_stream, write_stream):
         # Create a session using the client streams
         async with ClientSession(read_stream, write_stream) as session:
             # Initialize the connection
@@ -2395,7 +2391,7 @@ async def main():
     )
 
     async with httpx.AsyncClient(auth=oauth_auth, follow_redirects=True) as custom_client:
-        async with streamable_http_client("http://localhost:8001/mcp", http_client=custom_client) as (read, write, _):
+        async with streamable_http_client("http://localhost:8001/mcp", http_client=custom_client) as (read, write):
             async with ClientSession(read, write) as session:
                 await session.initialize()
 

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

@@ -14,7 +14,7 @@
 import time
 import webbrowser
 from http.server import BaseHTTPRequestHandler, HTTPServer
-from typing import Any, Callable
+from typing import Any
 from urllib.parse import parse_qs, urlparse
 
 import httpx
@@ -223,15 +223,15 @@ async def _default_redirect_handler(authorization_url: str) -> None:
                     auth=oauth_auth,
                     timeout=60.0,
                 ) as (read_stream, write_stream):
-                    await self._run_session(read_stream, write_stream, None)
+                    await self._run_session(read_stream, write_stream)
             else:
                 print("📡 Opening StreamableHTTP transport connection with auth...")
                 async with httpx.AsyncClient(auth=oauth_auth, follow_redirects=True) as custom_client:
-                    async with streamable_http_client(
-                        url=self.server_url,
-                        http_client=custom_client,
-                    ) as (read_stream, write_stream, get_session_id):
-                        await self._run_session(read_stream, write_stream, get_session_id)
+                    async with streamable_http_client(url=self.server_url, http_client=custom_client) as (
+                        read_stream,
+                        write_stream,
+                    ):
+                        await self._run_session(read_stream, write_stream)
 
         except Exception as e:
             print(f"❌ Failed to connect: {e}")
@@ -243,7 +243,6 @@ async def _run_session(
         self,
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
         write_stream: MemoryObjectSendStream[SessionMessage],
-        get_session_id: Callable[[], str | None] | None = None,
     ):
         """Run the MCP session with the given streams."""
         print("🤝 Initializing MCP session...")
@@ -254,10 +253,6 @@ async def _run_session(
             print("✨ Session initialization complete!")
 
             print(f"\n✅ Connected to MCP server at {self.server_url}")
-            if get_session_id:
-                session_id = get_session_id()
-                if session_id:
-                    print(f"Session ID: {session_id}")
 
             # Run interactive loop
             await self.interactive_loop()

examples/clients/simple-task-client/mcp_simple_task_client/main.py

@@ -9,7 +9,7 @@
 
 
 async def run(url: str) -> None:
-    async with streamable_http_client(url) as (read, write, _):
+    async with streamable_http_client(url) as (read, write):
         async with ClientSession(read, write) as session:
             await session.initialize()
 

examples/clients/simple-task-interactive-client/mcp_simple_task_interactive_client/main.py

@@ -73,7 +73,7 @@ def get_text(result: CallToolResult) -> str:
 
 
 async def run(url: str) -> None:
-    async with streamable_http_client(url) as (read, write, _):
+    async with streamable_http_client(url) as (read, write):
         async with ClientSession(
             read,
             write,

examples/clients/sse-polling-client/mcp_sse_polling_client/main.py

@@ -31,7 +31,7 @@ async def run_demo(url: str, items: int, checkpoint_every: int) -> None:
     print(f"Processing {items} items with checkpoints every {checkpoint_every}")
     print(f"{'=' * 60}\n")
 
-    async with streamable_http_client(url) as (read_stream, write_stream, _):
+    async with streamable_http_client(url) as (read_stream, write_stream):
         async with ClientSession(read_stream, write_stream) as session:
             # Initialize the connection
             print("Initializing connection...")

examples/snippets/clients/oauth_client.py

@@ -69,7 +69,7 @@ async def main():
     )
 
     async with httpx.AsyncClient(auth=oauth_auth, follow_redirects=True) as custom_client:
-        async with streamable_http_client("http://localhost:8001/mcp", http_client=custom_client) as (read, write, _):
+        async with streamable_http_client("http://localhost:8001/mcp", http_client=custom_client) as (read, write):
             async with ClientSession(read, write) as session:
                 await session.initialize()
 

examples/snippets/clients/streamable_basic.py

@@ -10,11 +10,7 @@
 
 async def main():
     # Connect to a streamable HTTP server
-    async with streamable_http_client("http://localhost:8000/mcp") as (
-        read_stream,
-        write_stream,
-        _,
-    ):
+    async with streamable_http_client("http://localhost:8000/mcp") as (read_stream, write_stream):
         # Create a session using the client streams
         async with ClientSession(read_stream, write_stream) as session:
             # Initialize the connection

src/mcp/client/__init__.py

@@ -1,6 +1,7 @@
 """MCP Client module."""
 
+from mcp.client._transport import Transport
 from mcp.client.client import Client
 from mcp.client.session import ClientSession
 
-__all__ = ["Client", "ClientSession"]
+__all__ = ["Client", "ClientSession", "Transport"]

src/mcp/client/_memory.py

@@ -2,17 +2,17 @@
 
 from __future__ import annotations
 
-from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from types import TracebackType
 from typing import Any
 
 import anyio
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
+from mcp.client._transport import TransportStreams
 from mcp.server import Server
 from mcp.server.mcpserver import MCPServer
 from mcp.shared.memory import create_client_server_memory_streams
-from mcp.shared.message import SessionMessage
 
 
 class InMemoryTransport:
@@ -23,17 +23,17 @@ class InMemoryTransport:
     stopped when the context manager exits.
 
     Example:
-        server = MCPServer("test")
-        transport = InMemoryTransport(server)
+        ```python
+        from mcp.client import Client, ClientSession
+        from mcp.server.mcpserver import MCPServer
+        from mcp.client.memory import InMemoryTransport
 
-        async with transport.connect() as (read_stream, write_stream):
+        server = MCPServer("test")
+        async with InMemoryTransport(server) as (read_stream, write_stream):
             async with ClientSession(read_stream, write_stream) as session:
                 await session.initialize()
                 # Use the session...
-
-    Or more commonly, use with Client:
-        async with Client(server) as client:
-            result = await client.call_tool("my_tool", {...})
+        ```
     """
 
     def __init__(self, server: Server[Any] | MCPServer, *, raise_exceptions: bool = False) -> None:
@@ -45,26 +45,15 @@ def __init__(self, server: Server[Any] | MCPServer, *, raise_exceptions: bool =
         """
         self._server = server
         self._raise_exceptions = raise_exceptions
+        self._cm: AbstractAsyncContextManager[TransportStreams] | None = None
 
     @asynccontextmanager
-    async def connect(
-        self,
-    ) -> AsyncGenerator[
-        tuple[
-            MemoryObjectReceiveStream[SessionMessage | Exception],
-            MemoryObjectSendStream[SessionMessage],
-        ],
-        None,
-    ]:
-        """Connect to the server and return streams for communication.
-
-        Yields:
-            A tuple of (read_stream, write_stream) for bidirectional communication
-        """
+    async def _connect(self) -> AsyncIterator[TransportStreams]:
+        """Connect to the server and yield streams for communication."""
         # Unwrap MCPServer to get underlying Server
-        actual_server: Server[Any]
         if isinstance(self._server, MCPServer):
-            actual_server = self._server._lowlevel_server  # type: ignore[reportPrivateUsage]
+            # TODO(Marcelo): Make `lowlevel_server` public.
+            actual_server: Server[Any] = self._server._lowlevel_server  # type: ignore[reportPrivateUsage]
         else:
             actual_server = self._server
 
@@ -87,3 +76,16 @@ async def connect(
                     yield client_read, client_write
                 finally:
                     tg.cancel_scope.cancel()
+
+    async def __aenter__(self) -> TransportStreams:
+        """Connect to the server and return streams for communication."""
+        self._cm = self._connect()
+        return await self._cm.__aenter__()
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> None:
+        """Close the transport and stop the server."""
+        if self._cm is not None:
+            await self._cm.__aexit__(exc_type, exc_val, exc_tb)
+            self._cm = None

src/mcp/client/_transport.py

@@ -0,0 +1,20 @@
+"""Transport protocol for MCP clients."""
+
+from __future__ import annotations
+
+from contextlib import AbstractAsyncContextManager
+from typing import Protocol
+
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+from mcp.shared.message import SessionMessage
+
+TransportStreams = tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
+
+
+class Transport(AbstractAsyncContextManager[TransportStreams], Protocol):
+    """Protocol for MCP transports.
+
+    A transport is an async context manager that yields read and write streams
+    for bidirectional communication with an MCP server.
+    """