Skip to content

Commit f848fbe

Browse files
author
Jianke LIN
committed
fix: fail fast when session is not started
1 parent 9bdc03d commit f848fbe

2 files changed

Lines changed: 29 additions & 2 deletions

File tree

‎src/mcp/client/session.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,7 @@ def __init__(
365365
self._task_group: anyio.abc.TaskGroup | None = None
366366
# subscriptions/listen demux routes; membership decides ack consumption (raw listens are never registered)
367367
self._listen_routes: dict[RequestId, ListenRoute] = {}
368+
self._entered = False
368369
if dispatcher is not None:
369370
if read_stream is not None or write_stream is not None:
370371
raise ValueError("pass read_stream/write_stream or dispatcher, not both")
@@ -386,6 +387,8 @@ def __init__(
386387
)
387388

388389
async def __aenter__(self) -> Self:
390+
if self._entered:
391+
raise RuntimeError("Session is already running")
389392
self._task_group = anyio.create_task_group()
390393
await self._task_group.__aenter__()
391394
try:
@@ -414,6 +417,7 @@ async def __aenter__(self) -> Self:
414417
finally:
415418
self._close_binding_queues()
416419
raise
420+
self._entered = True
417421
return self
418422

419423
async def __aexit__(
@@ -427,11 +431,12 @@ async def __aexit__(
427431
self._task_group.cancel_scope.cancel()
428432
try:
429433
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
434+
await resync_tracer()
435+
return result
430436
finally:
437+
self._entered = False
431438
self._close_binding_queues()
432439
self._settle_listen_routes_closed()
433-
await resync_tracer()
434-
return result
435440

436441
def _close_binding_queues(self) -> None:
437442
# Unclosed memory object streams warn at garbage collection; close is idempotent.
@@ -472,6 +477,11 @@ async def send_request(
472477
pydantic.ValidationError: The server returned a result that does not
473478
conform to the negotiated protocol version.
474479
"""
480+
if self._task_group is None:
481+
raise RuntimeError(
482+
"Session is not running. Use it as an async context manager "
483+
"(e.g. `async with ClientSession(...) as session:`)."
484+
)
475485
data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
476486
method: str = data["method"]
477487
opts: CallOptions = {}

‎tests/client/test_session.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,23 @@ async def message_handler( # pragma: no cover
155155
assert isinstance(initialized_notification, InitializedNotification)
156156

157157

158+
@pytest.mark.anyio
159+
async def test_client_session_requires_context_manager():
160+
client_to_server_send, _client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
161+
_server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
162+
163+
async with (
164+
client_to_server_send,
165+
_client_to_server_receive,
166+
_server_to_client_send,
167+
server_to_client_receive,
168+
):
169+
session = ClientSession(server_to_client_receive, client_to_server_send)
170+
171+
with pytest.raises(RuntimeError, match="async context manager"):
172+
await session.initialize()
173+
174+
158175
@pytest.mark.anyio
159176
async def test_client_session_custom_client_info():
160177
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)

0 commit comments

Comments
 (0)