fix(client): forward __aenter__ to wrapped client in StatefulProxyClient (refs #4054)

src/fastmcp/server/providers/proxy.py

@@ -932,9 +932,9 @@ async def default_proxy_progress_handler(
     await ctx.report_progress(progress, total, message)
 
 
-def _restore_request_context(
+async def _restore_request_context(
     rc_ref: list[Any],
-) -> None:
+) -> Context | None:
     """Set the ``request_ctx`` and ``_current_context`` ContextVars from stashed values.
 
     Called at the start of proxy handler invocations in
@@ -948,12 +948,15 @@ def _restore_request_context(
     ContextVar-dependent and would resolve stale values in the receive
     loop.  Instead we construct a fresh ``Context`` here after restoring
     ``request_ctx``, so its property accesses read the correct values.
-    """
-    from fastmcp.server.context import Context, _current_context
 
+    The returned ``Context`` (when one is entered) must have ``__aexit__``
+    called on it by the caller to release the ``_current_context`` /
+    ``_current_server`` tokens it pushed; ``_make_restoring_handler``
+    handles that via ``try/finally``.
+    """
     stashed = rc_ref[0]
     if stashed is None:
-        return
+        return None
 
     rc, fastmcp_ref = stashed
     try:
@@ -962,13 +965,18 @@ def _restore_request_context(
         request_ctx.set(rc)
         fastmcp = fastmcp_ref()
         if fastmcp is not None:
-            _current_context.set(Context(fastmcp))
-        return
+            # Use ``__aenter__`` (not a raw ``_current_context.set``) so
+            # ``Context._tokens``, ``_current_server`` and ``SharedContext``
+            # are all populated; otherwise the matching ``__aexit__`` would
+            # be a no-op and tokens would leak.
+            return await Context(fastmcp).__aenter__()
+        return None
     if current_rc.session is rc.session and current_rc.request_id != rc.request_id:
         request_ctx.set(rc)
         fastmcp = fastmcp_ref()
         if fastmcp is not None:
-            _current_context.set(Context(fastmcp))
+            return await Context(fastmcp).__aenter__()
+    return None
 
 
 def _make_restoring_handler(handler: Callable, rc_ref: list[Any]) -> Callable:
@@ -980,8 +988,12 @@ def _make_restoring_handler(handler: Callable, rc_ref: list[Any]) -> Callable:
     """
 
     async def wrapper(*args: Any, **kwargs: Any) -> Any:
-        _restore_request_context(rc_ref)
-        return await handler(*args, **kwargs)
+        ctx = await _restore_request_context(rc_ref)
+        try:
+            return await handler(*args, **kwargs)
+        finally:
+            if ctx is not None:
+                await ctx.__aexit__(None, None, None)
 
     return wrapper
 

tests/server/providers/proxy/test_stateful_proxy_client.py

@@ -199,3 +199,114 @@ async def elicitation_handler(message, response_type, params, ctx):
                 # one that would hang without the fix.
                 result2 = await client.call_tool("ask_name", {})
                 assert result2.data == "Hello, Alice!"
+
+
+class TestRestoreRequestContext:
+    """Regression tests for the Context lifecycle in `_restore_request_context`.
+
+    Previously `_restore_request_context` called `_current_context.set(Context(fastmcp))`
+    directly, bypassing `Context.__aenter__`.  That left `Context._tokens` empty,
+    `_current_server` unset, and the matching `__aexit__` a no-op (refs #4054).
+    """
+
+    async def test_restore_enters_context_and_sets_current_server(self):
+        """The fresh Context must be entered so `_current_context` and
+        `_current_server` are both populated, and `_tokens` is non-empty so
+        the eventual `__aexit__` actually releases the ContextVar."""
+        import weakref
+        from unittest.mock import MagicMock
+
+        from mcp.server.lowlevel.server import request_ctx
+
+        from fastmcp import FastMCP
+        from fastmcp.server.context import _current_context
+        from fastmcp.server.dependencies import _current_server
+        from fastmcp.server.providers.proxy import _restore_request_context
+
+        fastmcp = FastMCP("test")
+
+        # Stand-in RequestContext — only identity / attribute access is exercised.
+        rc = MagicMock()
+        rc.session = MagicMock()
+        rc.request_id = "req-1"
+
+        rc_ref: list = [(rc, weakref.ref(fastmcp))]
+
+        # No prior request_ctx in this task; _restore_request_context should
+        # take the LookupError branch, set request_ctx and enter the Context.
+        ctx = await _restore_request_context(rc_ref)
+        try:
+            assert ctx is not None, (
+                "Context should be entered when fastmcp weakref is alive"
+            )
+            # __aenter__ must have populated _tokens (so __aexit__ does work)
+            assert ctx._tokens, (
+                "Context.__aenter__ was not called: _tokens is empty"
+            )
+            # _current_context must point at the freshly entered Context
+            assert _current_context.get() is ctx
+            # _current_server must be set (only happens via __aenter__)
+            srv_ref = _current_server.get()
+            assert srv_ref is not None and srv_ref() is fastmcp
+            # And request_ctx must have been restored from the stash.
+            assert request_ctx.get() is rc
+        finally:
+            await ctx.__aexit__(None, None, None) if ctx else None
+
+        # __aexit__ must release the ContextVars cleanly.
+        assert _current_context.get(None) is not ctx
+
+    async def test_make_restoring_handler_awaits_aenter_and_aexit(self):
+        """The wrapper produced by `_make_restoring_handler` must await
+        `Context.__aenter__` before the handler runs and `__aexit__` after,
+        so handlers see a fully-initialised Context and tokens don't leak.
+        """
+        import weakref
+        from unittest.mock import MagicMock
+
+        from mcp.server.lowlevel.server import request_ctx
+
+        from fastmcp import FastMCP
+        from fastmcp.server.context import Context, _current_context
+        from fastmcp.server.providers.proxy import _make_restoring_handler
+
+        fastmcp = FastMCP("test")
+
+        rc = MagicMock()
+        rc.session = MagicMock()
+        rc.request_id = "req-2"
+
+        rc_ref: list = [(rc, weakref.ref(fastmcp))]
+
+        seen: dict = {}
+
+        async def handler(*args, **kwargs):
+            # Inside the handler the fresh Context must already be active.
+            cur = _current_context.get(None)
+            seen["context_active"] = isinstance(cur, Context)
+            seen["tokens_populated"] = bool(cur and cur._tokens)
+            return "ok"
+
+        wrapped = _make_restoring_handler(handler, rc_ref)
+
+        result = await wrapped()
+
+        assert result == "ok"
+        assert seen.get("context_active"), (
+            "handler did not see an active Context (aenter was not awaited)"
+        )
+        assert seen.get("tokens_populated"), (
+            "Context._tokens was empty inside handler (aenter bypassed)"
+        )
+        # After the wrapper returns, __aexit__ should have released the token,
+        # so _current_context no longer points at our restored Context.
+        assert _current_context.get(None) is None
+
+    async def test_restore_returns_none_when_rc_ref_empty(self):
+        """If nothing is stashed in `rc_ref`, no Context should be entered
+        (and therefore nothing for the caller to clean up)."""
+        from fastmcp.server.providers.proxy import _restore_request_context
+
+        rc_ref: list = [None]
+        ctx = await _restore_request_context(rc_ref)
+        assert ctx is None