Fix: Raise RuntimeError when ClientSession is used without context manager

This change ensures ClientSession is always used within an 'async with' block
by checking the context manager state before allowing operations that depend
on the background receive loop.

Without this check, calling initialize() outside of a context manager causes
the program to hang indefinitely because the _receive_loop task is only started
in __aenter__.

Changes:
- Add _entered flag to track context manager state
- Override __aenter__ and __aexit__ to manage the flag
- Add _check_is_active() method to enforce proper usage
- Call _check_is_active() in initialize() to prevent hangs
- Add test case for invalid usage

Fixes #1564

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: newbiehwang

src/mcp/client/session.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from types import TracebackType
 from typing import Any, Protocol
 
 import anyio.lowlevel
@@ -107,6 +108,8 @@ class ClientSession(
         types.ServerNotification,
     ]
 ):
+    _entered: bool
+
     def __init__(
         self,
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
@@ -133,10 +136,29 @@ def __init__(
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
         self._server_capabilities: types.ServerCapabilities | None = None
         self._experimental_features: ExperimentalClientFeatures | None = None
+        self._entered = False
 
         # Experimental: Task handlers (use defaults if not provided)
         self._task_handlers = experimental_task_handlers or ExperimentalTaskHandlers()
 
+    async def __aenter__(self) -> ClientSession:
+        self._entered = True
+        await super().__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self._entered = False
+        await super().__aexit__(exc_type, exc_value, traceback)
+
+    def _check_is_active(self) -> None:
+        if not self._entered:
+            raise RuntimeError("ClientSession must be used within an 'async with' block.")
+
     @property
     def _receive_request_adapter(self) -> TypeAdapter[types.ServerRequest]:
         return types.server_request_adapter
@@ -146,6 +168,7 @@ def _receive_notification_adapter(self) -> TypeAdapter[types.ServerNotification]
         return types.server_notification_adapter
 
     async def initialize(self) -> types.InitializeResult:
+        self._check_is_active()
         sampling = (
             (self._sampling_capabilities or types.SamplingCapability())
             if self._sampling_callback is not _default_sampling_callback

tests/client/test_session.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any, cast
+
 import anyio
 import pytest
 
@@ -706,3 +708,18 @@ async def mock_server():
         await session.initialize()
 
         await session.call_tool(name=mocked_tool.name, arguments={"foo": "bar"}, meta=meta)
+
+
+@pytest.mark.anyio
+async def test_initialize_without_context_manager_raises_error():
+    """Test that calling initialize() without entering the context manager raises RuntimeError."""
+    send_stream, receive_stream = anyio.create_memory_object_stream[Any](0)
+
+    read_stream = cast(Any, receive_stream)
+    write_stream = cast(Any, send_stream)
+
+    async with send_stream, receive_stream:
+        session = ClientSession(read_stream, write_stream)
+
+        with pytest.raises(RuntimeError, match="must be used within"):
+            await session.initialize()