From 0cd3690aefa9dde6b5071e92595bb8e3a61ff4e7 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sun, 14 Dec 2025 09:11:30 -0500
Subject: [PATCH] Consolidate session state reset logic and improve
 cancellation cleanup (#2615)

---
 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 d67b7de8e..6dd13d892 100644
--- a/src/fastmcp/client/client.py
+++ b/src/fastmcp/client/client.py
@@ -6,7 +6,7 @@ import datetime
 import secrets
 import uuid
 import weakref
-from contextlib import AsyncExitStack, asynccontextmanager, suppress
+from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Generic, Literal, TypeVar, cast, overload
@@ -324,6 +324,20 @@ class Client(Generic[ClientTransportT]):
             str, weakref.ref[ToolTask | PromptTask | ResourceTask]
         ] = {}
 
+    def _reset_session_state(self, full: bool = False) -> None:
+        """Reset session state after disconnect or cancellation.
+
+        Args:
+            full: If True, also resets session_task and nesting_counter.
+                  Use full=True for cancellation cleanup where the session
+                  task was started but never completed normally.
+        """
+        self._session_state.session = None
+        self._session_state.initialize_result = None
+        if full:
+            self._session_state.session_task = None
+            self._session_state.nesting_counter = 0
+
     @property
     def session(self) -> ClientSession:
         """Get the current active session. Raises RuntimeError if not connected."""
@@ -404,8 +418,7 @@ class Client(Generic[ClientTransportT]):
                 except anyio.ClosedResourceError as e:
                     raise RuntimeError("Server session was closed unexpectedly") from e
                 finally:
-                    self._session_state.session = None
-                    self._session_state.initialize_result = None
+                    self._reset_session_state()
 
     async def initialize(
         self,
@@ -510,25 +523,38 @@ class Client(Generic[ClientTransportT]):
                     # when __aenter__ is cancelled. Since we hold the session lock here
                     # and we know we started the session task, it's safe to tear it down
                     # without impacting other active contexts.
+                    #
+                    # Note: session_task is an asyncio.Task (not anyio) because it needs
+                    # to outlive individual context manager scopes - anyio's structured
+                    # concurrency doesn't allow tasks to escape their task group.
                     session_task = self._session_state.session_task
                     if session_task is not None:
                         # Request a graceful stop if the runner has already reached
                         # its stop_event wait.
                         self._session_state.stop_event.set()
                         session_task.cancel()
-                        # Preserve the original cancellation.
-                        with suppress(asyncio.CancelledError, Exception):
-                            await asyncio.shield(session_task)
+                        with anyio.CancelScope(shield=True):
+                            with anyio.move_on_after(3):
+                                try:
+                                    await session_task
+                                except asyncio.CancelledError:
+                                    pass
+                                except Exception as e:
+                                    logger.debug(
+                                        f"Error during cancelled session cleanup: {e}"
+                                    )
 
                     # Reset session state so future callers can reconnect cleanly.
-                    self._session_state.session_task = None
-                    self._session_state.session = None
-                    self._session_state.initialize_result = None
-                    self._session_state.nesting_counter = 0
+                    self._reset_session_state(full=True)
 
-                    # Preserve the original cancellation.
-                    with suppress(asyncio.CancelledError, Exception):
-                        await asyncio.shield(self.transport.close())
+                    with anyio.CancelScope(shield=True):
+                        with anyio.move_on_after(3):
+                            try:
+                                await self.transport.close()
+                            except Exception as e:
+                                logger.debug(
+                                    f"Error closing transport after cancellation: {e}"
+                                )
 
                     raise