From 6d9b85a4eccad970a80d4cfa898fcb2f1aeab59a Mon Sep 17 00:00:00 2001 From: hwangshincheol Date: Mon, 12 Jan 2026 16:48:13 +0900 Subject: [PATCH 1/7] Fix: Prevent hang when ClientSession is initialized without context manager --- src/mcp/client/session.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index b61bf0b03..a671476f4 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -1,5 +1,6 @@ import logging from typing import Any, Protocol, overload +from types import TracebackType import anyio.lowlevel from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream @@ -140,11 +141,33 @@ 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." + ) + 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 From 6ba5675124fa973ecf6d5e113716f85a0142e26b Mon Sep 17 00:00:00 2001 From: hwangshincheol Date: Mon, 12 Jan 2026 17:06:48 +0900 Subject: [PATCH 2/7] Fix: Apply formatting and add type hint --- src/mcp/client/session.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index a671476f4..b503c2782 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -109,6 +109,8 @@ class ClientSession( types.ServerNotification, ] ): + _entered: bool + def __init__( self, read_stream: MemoryObjectReceiveStream[SessionMessage | Exception], @@ -162,9 +164,7 @@ async def __aexit__( def _check_is_active(self) -> None: if not self._entered: - raise RuntimeError( - "ClientSession must be used within an 'async with' block." - ) + raise RuntimeError("ClientSession must be used within an 'async with' block.") async def initialize(self) -> types.InitializeResult: self._check_is_active() From 05f0b1753808837b2b2255f3bc3124c1241dd493 Mon Sep 17 00:00:00 2001 From: hwangshincheol Date: Mon, 12 Jan 2026 17:09:04 +0900 Subject: [PATCH 3/7] Test: Add coverage for invalid session usage --- tests/client/test_session.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/client/test_session.py b/tests/client/test_session.py index eb2683fbd..e9137336b 100644 --- a/tests/client/test_session.py +++ b/tests/client/test_session.py @@ -768,3 +768,20 @@ 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. + """ + # Create dummy streams for session initialization + read_stream, write_stream = anyio.create_memory_object_stream(0) + + async with read_stream, write_stream: + # Instantiate the session object directly + session = ClientSession(read_stream, write_stream) + + # Verify that calling initialize() raises a RuntimeError + with pytest.raises(RuntimeError, match="must be used within"): + await session.initialize() From 13c694e915b183efb1b5a7ea6e2b1a9b38d220fc Mon Sep 17 00:00:00 2001 From: hwangshincheol Date: Mon, 12 Jan 2026 17:11:27 +0900 Subject: [PATCH 4/7] Fix: Apply formatting and fix test resource cleanup --- src/mcp/client/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index b503c2782..04562a0b7 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -110,7 +110,7 @@ class ClientSession( ] ): _entered: bool - + def __init__( self, read_stream: MemoryObjectReceiveStream[SessionMessage | Exception], From fb11722840680e768284bbd92e0afaae9823f669 Mon Sep 17 00:00:00 2001 From: hwangshincheol Date: Mon, 12 Jan 2026 17:13:33 +0900 Subject: [PATCH 5/7] Fix: Sort imports to satisfy ruff --- src/mcp/client/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 04562a0b7..9f5bb1774 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -1,6 +1,6 @@ import logging -from typing import Any, Protocol, overload from types import TracebackType +from typing import Any, Protocol, overload import anyio.lowlevel from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream From 82e0f2f65b1cb2c2d7980b15ab7de86b8d1b83eb Mon Sep 17 00:00:00 2001 From: hwangshincheol Date: Mon, 12 Jan 2026 17:27:17 +0900 Subject: [PATCH 6/7] Test: Modify test_initialize_without_context_manager_raises_error --- tests/client/test_session.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/client/test_session.py b/tests/client/test_session.py index e9137336b..316462ca1 100644 --- a/tests/client/test_session.py +++ b/tests/client/test_session.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, cast import anyio import pytest @@ -775,13 +775,13 @@ async def test_initialize_without_context_manager_raises_error(): """ Test that calling initialize() without entering the context manager raises RuntimeError. """ - # Create dummy streams for session initialization - read_stream, write_stream = anyio.create_memory_object_stream(0) + send_stream, receive_stream = anyio.create_memory_object_stream[Any](0) - async with read_stream, write_stream: - # Instantiate the session object directly + read_stream = cast(Any, receive_stream) + write_stream = cast(Any, send_stream) + + async with send_stream, receive_stream: session = ClientSession(read_stream, write_stream) - # Verify that calling initialize() raises a RuntimeError with pytest.raises(RuntimeError, match="must be used within"): await session.initialize() From ba178a1c0264ed54c4605092706baa2cffda4684 Mon Sep 17 00:00:00 2001 From: hwangshincheol Date: Tue, 13 Jan 2026 18:21:14 +0900 Subject: [PATCH 7/7] Check Restart