ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-20 12:34:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

Consolidate session state reset logic and improve cancellation cleanup (#2615)

Browse source
This commit is contained in:
Jeremiah Lowin 2025-12-14 09:11:30 -05:00 • committed by GitHub
commit 0cd3690aef
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 40 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

52
src/fastmcp/client/client.py
View file

@ -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