diff --git a/fastmcp_slim/fastmcp/server/providers/proxy.py b/fastmcp_slim/fastmcp/server/providers/proxy.py index 931c49594..2d18f635d 100644 --- a/fastmcp_slim/fastmcp/server/providers/proxy.py +++ b/fastmcp_slim/fastmcp/server/providers/proxy.py @@ -1804,10 +1804,12 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]): return cast(StatefulProxyClient[ClientTransportT], super().new()) async def __aexit__(self, exc_type, exc_value, traceback) -> None: # type: ignore[override] # ty:ignore[invalid-method-override] - """The stateful proxy client will be forced disconnected when the session is exited. - - So we do nothing here. - """ + """Release this context without disconnecting the persistent session.""" + with anyio.CancelScope(shield=True): + async with self._session_state.lock: + self._session_state.nesting_counter = max( + 0, self._session_state.nesting_counter - 1 + ) async def clear(self): """Clear all cached clients and force disconnect them.""" diff --git a/tests/server/providers/proxy/test_stateful_proxy_client.py b/tests/server/providers/proxy/test_stateful_proxy_client.py index 0bf832440..19614ce91 100644 --- a/tests/server/providers/proxy/test_stateful_proxy_client.py +++ b/tests/server/providers/proxy/test_stateful_proxy_client.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import weakref from dataclasses import dataclass from unittest.mock import MagicMock @@ -89,6 +90,32 @@ async def stateless_server(stateful_proxy_server: FastMCP): class TestStatefulProxyClient: + async def test_reconnects_after_persistent_session_ends(self): + """A completed request must not prevent a dead session from reconnecting.""" + backend = FastMCP("backend") + + @backend.tool + def echo(value: str) -> str: + return value + + client = StatefulProxyClient(backend) + try: + async with client: + result = await client.call_tool("echo", {"value": "first"}) + assert result.data == "first" + + session_task = client._session_state.session_task + assert session_task is not None + session_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await session_task + + async with client: + result = await client.call_tool("echo", {"value": "second"}) + assert result.data == "second" + finally: + await client.close() + async def test_concurrent_log_requests_no_mixing( self, stateful_proxy_server: FastMCP ):