refactor: address PR review comments for Client class

- Make InMemoryTransport private by moving to _memory.py module
- Improve exception handling in __aenter__: use AsyncExitStack.pop_all()
  pattern instead of manual try/except for cleaner resource cleanup
- Add reentry check: raise RuntimeError if Client is entered twice
- Add test for reentry behavior
- Remove low-level transport documentation (InMemoryTransport is now private)

Author: felixweinberger

docs/testing.md

@@ -69,26 +69,3 @@ async def test_call_add_tool():
 1. If you are using `trio`, you should set `"trio"` as the `anyio_backend`. Check more information in the [anyio documentation](https://anyio.readthedocs.io/en/stable/testing.html#specifying-the-backends-to-run-on).
 
 There you go! You can now extend your tests to cover more scenarios.
-
-## Advanced: Low-Level Transport Access
-
-For tests that need direct access to the underlying `ClientSession`, use `InMemoryTransport`:
-
-```python
-from mcp.client.session import ClientSession
-from mcp.client.transports import InMemoryTransport
-
-from server import app
-
-
-async def test_with_low_level_access():
-    transport = InMemoryTransport(app, raise_exceptions=True)
-    async with transport.connect() as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-            # Direct access to the full ClientSession API
-            result = await session.call_tool("add", {"a": 1, "b": 2})
-```
-
-The `Client` class is built on top of `InMemoryTransport` and handles initialization automatically.
-Use `InMemoryTransport` directly only when you need fine-grained control over the session lifecycle.

src/mcp/client/__init__.py

@@ -2,10 +2,8 @@
 
 from mcp.client.client import Client
 from mcp.client.session import ClientSession
-from mcp.client.transports.memory import InMemoryTransport
 
 __all__ = [
     "Client",
     "ClientSession",
-    "InMemoryTransport",
 ]

src/mcp/client/transports/memory.py src/mcp/client/_memory.pysrc/mcp/client/transports/memory.py renamed to src/mcp/client/_memory.py

@@ -9,6 +9,7 @@
 from pydantic import AnyUrl
 
 import mcp.types as types
+from mcp.client._memory import InMemoryTransport
 from mcp.client.session import (
     ClientSession,
     ElicitationFnT,
@@ -17,7 +18,6 @@
     MessageHandlerFnT,
     SamplingFnT,
 )
-from mcp.client.transports.memory import InMemoryTransport
 from mcp.server import Server
 from mcp.server.fastmcp import FastMCP
 from mcp.shared.session import ProgressFnT
@@ -92,16 +92,16 @@ def __init__(
 
     async def __aenter__(self) -> Client:
         """Enter the async context manager."""
-        self._exit_stack = AsyncExitStack()
-        await self._exit_stack.__aenter__()
+        if self._session is not None:
+            raise RuntimeError("Client is already entered; cannot reenter")
 
-        try:
+        async with AsyncExitStack() as exit_stack:
             # Create transport and connect
             transport = InMemoryTransport(self._server, raise_exceptions=self._raise_exceptions)
-            read_stream, write_stream = await self._exit_stack.enter_async_context(transport.connect())
+            read_stream, write_stream = await exit_stack.enter_async_context(transport.connect())
 
             # Create session
-            self._session = await self._exit_stack.enter_async_context(
+            self._session = await exit_stack.enter_async_context(
                 ClientSession(
                     read_stream=read_stream,
                     write_stream=write_stream,
@@ -118,10 +118,9 @@ async def __aenter__(self) -> Client:
             # Initialize the session
             await self._session.initialize()
 
+            # Transfer ownership to self for __aexit__ to handle
+            self._exit_stack = exit_stack.pop_all()
             return self
-        except Exception:
-            await self._exit_stack.__aexit__(None, None, None)
-            raise
 
     async def __aexit__(
         self,

src/mcp/client/client.py

@@ -1,5 +1,3 @@
 """Transport implementations for MCP clients."""
 
-from mcp.client.transports.memory import InMemoryTransport
-
-__all__ = ["InMemoryTransport"]
+__all__: list[str] = []

src/mcp/client/transports/__init__.py

@@ -125,7 +125,7 @@ async def patched_create_streams():
     # Apply the patch for the duration of the test
     # Patch both locations since InMemoryTransport imports it directly
     with patch("mcp.shared.memory.create_client_server_memory_streams", patched_create_streams):
-        with patch("mcp.client.transports.memory.create_client_server_memory_streams", patched_create_streams):
+        with patch("mcp.client._memory.create_client_server_memory_streams", patched_create_streams):
             # Return a collection with helper methods
             def get_spy_collection() -> StreamSpyCollection:
                 assert client_spy is not None, "client_spy was not initialized"

tests/client/conftest.py

@@ -220,6 +220,13 @@ def test_session_property_before_enter(app: FastMCP):
         _ = client.session
 
 
+async def test_reentry_raises_runtime_error(app: FastMCP):
+    """Test that reentering a client raises RuntimeError."""
+    async with Client(app) as client:
+        with pytest.raises(RuntimeError, match="Client is already entered"):
+            await client.__aenter__()
+
+
 async def test_cleanup_on_init_failure(app: FastMCP):
     """Test that resources are cleaned up if initialization fails."""
     with patch("mcp.client.client.ClientSession") as mock_session_class:
@@ -230,9 +237,16 @@ async def test_cleanup_on_init_failure(app: FastMCP):
         mock_session_class.return_value = mock_session
 
         client = Client(app)
-        with pytest.raises(RuntimeError, match="Session init failed"):
+        with pytest.raises(BaseException) as exc_info:
             await client.__aenter__()
 
+        # The error should contain our message (may be wrapped in ExceptionGroup)
+        exc = exc_info.value
+        if hasattr(exc, "exceptions"):  # ExceptionGroup
+            assert any("Session init failed" in str(e) for e in exc.exceptions)
+        else:
+            assert "Session init failed" in str(exc)
+
         # Verify the client is in a clean state (session should be None)
         assert client._session is None
 

tests/client/test_client.py

@@ -3,8 +3,8 @@
 import pytest
 
 import mcp.types as types
+from mcp.client._memory import InMemoryTransport
 from mcp.client.session import ClientSession
-from mcp.client.transports import InMemoryTransport
 from mcp.server import Server
 from mcp.server.fastmcp import FastMCP
 from mcp.types import ListToolsRequest, ListToolsResult

tests/client/test_list_methods_cursor.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from mcp.client.transports.memory import InMemoryTransport
+from mcp.client._memory import InMemoryTransport
 from mcp.server import Server
 from mcp.server.fastmcp import FastMCP
 from mcp.types import Resource

tests/client/transports/test_memory.py

@@ -6,8 +6,8 @@
 import pytest
 
 import mcp.types as types
+from mcp.client._memory import InMemoryTransport
 from mcp.client.session import ClientSession
-from mcp.client.transports import InMemoryTransport
 from mcp.server.lowlevel.server import Server
 from mcp.shared.exceptions import McpError
 from mcp.types import (