From 10c868274beab8cd24a9283e272aad1bc57f4446 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sat, 31 May 2025 18:03:14 -0400
Subject: [PATCH] Ensure close() cleans up clients

---
 src/fastmcp/client/client.py | 52 +++++++++++++++++++++++++++---------
 1 file changed, 39 insertions(+), 13 deletions(-)

diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py
index b908fbbd7..15e72c7f1 100644
--- a/src/fastmcp/client/client.py
+++ b/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 = asyncio.Event()
+        self._stop_event = asyncio.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 ---