fix: fail fast when session is not started

Author: Epochex

src/mcp/shared/session.py

@@ -7,6 +7,7 @@
 from typing import Any, Generic, Protocol, TypeVar
 
 import anyio
+from anyio.abc import TaskGroup
 from anyio.streams.memory import MemoryObjectSendStream
 from opentelemetry.trace import SpanKind
 from pydantic import BaseModel, TypeAdapter
@@ -156,10 +157,22 @@ def __init__(
         self._session_read_timeout_seconds = read_timeout_seconds
         self._progress_callbacks = {}
         self._exit_stack = AsyncExitStack()
+        self._task_group: TaskGroup = anyio.create_task_group()
+        self._started = False
+
+    def _require_started(self) -> None:
+        if not self._started:
+            raise RuntimeError(
+                "Session is not running. Use it as an async context manager "
+                "(e.g. `async with ClientSession(...) as session:`)."
+            )
 
     async def __aenter__(self) -> Self:
+        if self._started:
+            raise RuntimeError("Session is already running")
         self._task_group = anyio.create_task_group()
         await self._task_group.__aenter__()
+        self._started = True
         self._task_group.start_soon(self._receive_loop)
         return self
 
@@ -174,9 +187,12 @@ async def __aexit__(
         # would be very surprising behavior), so make sure to cancel the tasks
         # in the task group.
         self._task_group.cancel_scope.cancel()
-        result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
-        await resync_tracer()
-        return result
+        try:
+            result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
+            await resync_tracer()
+            return result
+        finally:
+            self._started = False
 
     async def send_request(
         self,
@@ -193,6 +209,7 @@ async def send_request(
 
         Do not use this method to emit notifications! Use send_notification() instead.
         """
+        self._require_started()
         request_id = self._request_id
         self._request_id = request_id + 1
 
@@ -255,6 +272,7 @@ async def send_notification(
         related_request_id: RequestId | None = None,
     ) -> None:
         """Emits a notification, which is a one-way message that does not expect a response."""
+        self._require_started()
         # Some transport implementations may need to set the related_request_id
         # to attribute to the notifications to the request that triggered them.
         jsonrpc_notification = JSONRPCNotification(

tests/client/test_session.py

@@ -140,6 +140,23 @@ async def message_handler(  # pragma: no cover
     assert isinstance(initialized_notification, InitializedNotification)
 
 
+@pytest.mark.anyio
+async def test_client_session_requires_context_manager():
+    client_to_server_send, _client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    _server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+
+    async with (
+        client_to_server_send,
+        _client_to_server_receive,
+        _server_to_client_send,
+        server_to_client_receive,
+    ):
+        session = ClientSession(server_to_client_receive, client_to_server_send)
+
+        with pytest.raises(RuntimeError, match="async context manager"):
+            await session.initialize()
+
+
 @pytest.mark.anyio
 async def test_client_session_custom_client_info():
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)