Merge pull request #643 from jlowin/concurrency-2

Ensure close() cleans up client context appropriately

src/fastmcp/client/client.py

@@ -150,13 +150,6 @@ class Client(Generic[ClientTransportT]):
         self.transport = cast(ClientTransportT, infer_transport(transport))
         if auth is not None:
             self.transport._set_auth(auth)
-        self._session: ClientSession | None = None
-        self._exit_stack: AsyncExitStack | None = None
-        self._nesting_counter: int = 0
-        self._context_lock = anyio.Lock()
-        self._session_task: asyncio.Task | None = None
-        self._ready_event = asyncio.Event()
-        self._stop_event = asyncio.Event()
         self._initialize_result: mcp.types.InitializeResult | None = None
 
         if log_handler is None:
@@ -196,7 +189,15 @@ class Client(Generic[ClientTransportT]):
             self._session_kwargs["sampling_callback"] = create_sampling_callback(
                 sampling_handler
             )
-        # self._session_manager = self._context_manager()
+
+        # session context management
+        self._session: ClientSession | None = None
+        self._exit_stack: AsyncExitStack | None = None
+        self._nesting_counter: int = 0
+        self._context_lock = anyio.Lock()
+        self._session_task: asyncio.Task | None = None
+        self._ready_event = anyio.Event()
+        self._stop_event = anyio.Event()
 
     @property
     def session(self) -> ClientSession:
@@ -252,6 +253,14 @@ class Client(Generic[ClientTransportT]):
                     self._initialize_result = None
 
     async def __aenter__(self):
+        await self._connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._disconnect()
+
+    async def _connect(self):
+        # ensure only one session is running at a time to avoid race conditions
         async with self._context_lock:
             need_to_start = self._session_task is None or self._session_task.done()
             if need_to_start:
@@ -262,19 +271,37 @@ class Client(Generic[ClientTransportT]):
             self._nesting_counter += 1
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def _disconnect(self, force: bool = False):
+        # ensure only one session is running at a time to avoid race conditions
         async with self._context_lock:
-            self._nesting_counter -= 1
-            if self._nesting_counter != 0:
+            # if we are forcing a disconnect, reset the nesting counter
+            if force:
+                self._nesting_counter = 0
+
+            # otherwise decrement to check if we are done nesting
+            else:
+                self._nesting_counter = max(0, self._nesting_counter - 1)
+
+            # if we are still nested, return
+            if self._nesting_counter > 0:
+                return
+
+            # stop the active seesion
+            if self._session_task is None:
                 return
             self._stop_event.set()
             runner_task = self._session_task
             self._session_task = None
+
+        # wait for the session to finish
         if runner_task:
             await runner_task
+
         # Reset for future reconnects
         self._stop_event = anyio.Event()
         self._ready_event = anyio.Event()
+        self._session = None
+        self._initialize_result = None
 
     async def _session_runner(self):
         async with AsyncExitStack() as stack:
@@ -289,9 +316,8 @@ class Client(Generic[ClientTransportT]):
                 self._ready_event.set()
 
     async def close(self):
+        await self._disconnect(force=True)
         await self.transport.close()
-        self._session = None
-        self._initialize_result = None
 
     # --- MCP Client Methods ---
 

src/fastmcp/client/transports.py

@@ -18,6 +18,7 @@ from typing import (
     overload,
 )
 
+import anyio
 import httpx
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.session import (
@@ -327,8 +328,8 @@ class StdioTransport(ClientTransport):
 
         self._session: ClientSession | None = None
         self._connect_task: asyncio.Task | None = None
-        self._ready_event = asyncio.Event()
-        self._stop_event = asyncio.Event()
+        self._ready_event = anyio.Event()
+        self._stop_event = anyio.Event()
 
     @contextlib.asynccontextmanager
     async def connect_session(
@@ -391,8 +392,8 @@ class StdioTransport(ClientTransport):
 
         # reset variables and events for potential future reconnects
         self._connect_task = None
-        self._stop_event = asyncio.Event()
-        self._ready_event = asyncio.Event()
+        self._stop_event = anyio.Event()
+        self._ready_event = anyio.Event()
 
     async def close(self):
         await self.disconnect()

tests/client/test_client.py

@@ -342,6 +342,52 @@ async def test_client_nested_context_manager(fastmcp_server):
     assert client._session is None
 
 
+async def test_concurrent_client_context_managers():
+    """
+    Test that concurrent client usage doesn't cause cross-task cancel scope issues.
+    https://github.com/jlowin/fastmcp/pull/643
+    """
+    # Create a simple server
+    server = FastMCP("Test Server")
+
+    @server.tool()
+    def echo(text: str) -> str:
+        """Echo tool"""
+        return text
+
+    # Create client
+    client = Client(server)
+
+    # Track results
+    results = {}
+    errors = []
+
+    async def use_client(task_id: str, delay: float = 0):
+        """Use the client with a small delay to ensure overlap"""
+        try:
+            async with client:
+                # Add a small delay to ensure contexts overlap
+                await asyncio.sleep(delay)
+                # Make an actual call to exercise the session
+                tools = await client.list_tools()
+                results[task_id] = len(tools)
+        except Exception as e:
+            errors.append((task_id, str(e)))
+
+    # Run multiple tasks concurrently
+    # The key is having them enter and exit the context at different times
+    await asyncio.gather(
+        use_client("task1", 0.0),
+        use_client("task2", 0.01),  # Slight delay to ensure overlap
+        use_client("task3", 0.02),
+        return_exceptions=False,
+    )
+
+    assert len(errors) == 0, f"Errors occurred: {errors}"
+    assert len(results) == 3
+    assert all(count == 1 for count in results.values())  # All should see 1 tool
+
+
 async def test_resource_template(fastmcp_server):
     """Test using a resource template with InMemoryClient."""
     client = Client(transport=FastMCPTransport(fastmcp_server))

tests/server/test_logging.py

@@ -2,6 +2,7 @@ import asyncio
 import logging
 from unittest.mock import AsyncMock, Mock, patch
 
+import anyio
 import pytest
 
 from fastmcp.server.server import FastMCP
@@ -27,7 +28,7 @@ async def test_uvicorn_logging_default_level(
     """Tests that FastMCP passes log_level to uvicorn.Config if no log_config is given."""
     mock_server_instance = AsyncMock()
     mock_uvicorn_server_constructor.return_value = mock_server_instance
-    serve_finished_event = asyncio.Event()
+    serve_finished_event = anyio.Event()
     mock_server_instance.serve.side_effect = serve_finished_event.wait
 
     test_log_level = "warning"
@@ -63,7 +64,7 @@ async def test_uvicorn_logging_with_custom_log_config(
     """Tests that FastMCP passes log_config to uvicorn.Config and not log_level."""
     mock_server_instance = AsyncMock()
     mock_uvicorn_server_constructor.return_value = mock_server_instance
-    serve_finished_event = asyncio.Event()
+    serve_finished_event = anyio.Event()
     mock_server_instance.serve.side_effect = serve_finished_event.wait
 
     sample_log_config = {
@@ -123,7 +124,7 @@ async def test_uvicorn_logging_custom_log_config_overrides_log_level_param(
     """Tests log_config precedence if log_level is also passed to run_http_async."""
     mock_server_instance = AsyncMock()
     mock_uvicorn_server_constructor.return_value = mock_server_instance
-    serve_finished_event = asyncio.Event()
+    serve_finished_event = anyio.Event()
     mock_server_instance.serve.side_effect = serve_finished_event.wait
 
     sample_log_config = {