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

[MukundaKatta]

Why

Issue #4054 catalogues four unrelated auth-handling bugs. The StatefulProxyClient Context lifecycle bug (Bug 4) is the last one outstanding:

• Bug 1 (auth-denial error message): not addressed yet.
• Bug 2 (OIDC discovery caching): not addressed yet.
• Bug 3 (refresh_expires_in=0 handling): tracked in fix: treat Keycloak refresh_expires_in=0 as never-expires sentinel #4060.
• Bug 4 (this PR): _restore_request_context skipped Context.__aenter__.

This PR closes out Bug 4 only.

For visibility, two adjacent issues from the same audit batch are being addressed in #4076 (parse_qs blank values, refs #4056) and #4077 (ToolResult.isError + dedupe meta deep-copy, refs #4055).

What

src/fastmcp/server/providers/proxy.py — _restore_request_context previously did:

_current_context.set(Context(fastmcp))

That bypasses Context.__aenter__, so:

• Context._tokens stays empty, which means the matching Context.__aexit__ is a no-op.
• _current_server is never set, so dependency-injection helpers that rely on it see stale values.
• SharedContext is never created when running without a Docket lifespan, so Shared() dependencies inside the handler get a stale shared store.

The fresh Context therefore looks half-initialised inside the proxy handlers triggered from StatefulProxyClient's receive loop, and its tokens leak as the wrapper returns.

The fix:

1. Make _restore_request_context async and await Context(fastmcp).__aenter__() instead of touching _current_context directly. The function returns the entered Context | None so the caller can pair the enter with an exit.
2. Update _make_restoring_handler to wrap the handler in try/finally and call await ctx.__aexit__(None, None, None) after the handler runs, so the _current_context token (and any _current_server / _shared_context / docket tokens) are released cleanly.

The request_ctx restoration logic and the staleness-detection branch (only override when same session, different request_id) are unchanged — only the _current_context side now goes through the proper enter/exit.

Tests

tests/server/providers/proxy/test_stateful_proxy_client.py — new TestRestoreRequestContext class covering:

• test_restore_enters_context_and_sets_current_server — the returned Context has populated _tokens, _current_context.get() is the fresh Context, _current_server resolves back to the same FastMCP, and request_ctx was restored to the stashed value. After __aexit__, _current_context no longer points at the entered Context.
• test_make_restoring_handler_awaits_aenter_and_aexit — the handler invoked through the wrapper sees an active Context with non-empty _tokens, and after the wrapper returns _current_context is back to None (no token leak).
• test_restore_returns_none_when_rc_ref_empty — pins down the rc_ref[0] is None shortcut.

The existing TestStatefulProxyClient suite (concurrent log routing, stateful state, multi-proxy isolation, elicitation-over-HTTP) should continue to pass since the request_ctx and routing behaviour is unchanged.

Notes

• Diff is intentionally minimal: one production change to proxy.py, one new test class.
• _restore_request_context becoming async is internal — it's only called from _make_restoring_handler (also internal).
• Refs Multiple auth handling bugs #4054. The other three bugs in that issue are out of scope for this PR.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1ddb8c74f3

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Keep restored context alive for same-request callbacks

_make_restoring_handler now always calls ctx.__aexit__() after each handler invocation, but _restore_request_context only re-enters a Context when request_id changes (or request_ctx is missing). If an upstream server emits multiple callbacks for the same request (for example, several progress or log notifications), the second callback runs without a freshly entered context, so get_context() in the default handlers can resolve a stale context or raise RuntimeError("No active context found."). This regresses multi-notification flows within a single request.

Useful? React with 👍 / 👎.

[marvin-context-protocol]

tl;dr: ruff-check and ruff-format failed on tests/server/providers/proxy/test_stateful_proxy_client.py. Run uv run prek run --all-files locally and push.

Root Cause: The new test file has an unused import (from mcp.server.lowlevel.server import request_ctx) and a multi-line assert that ruff wants formatted as a single line. ruff-check auto-fixed both issues (2 errors fixed), and ruff-format reformatted the assert. Both hooks fail in CI when they modify files.

Fix: Run the following locally and push the result:
```bash
uv run prek run --all-files
```

Log excerpts

ruff check...............................................................Failed
- hook id: ruff-check
- exit code: 1
- files were modified by this hook

  Found 2 errors (2 fixed, 0 remaining).

ruff format..............................................................Failed
- hook id: ruff-format
- files were modified by this hook

  1 file reformatted, 753 files left unchanged

Diff applied by the hooks to tests/server/providers/proxy/test_stateful_proxy_client.py:

@@ -240,9 +240,7 @@ class TestRestoreRequestContext:
-            assert ctx._tokens, (
-                "Context.__aenter__ was not called: _tokens is empty"
-            )
+            assert ctx._tokens, "Context.__aenter__ was not called: _tokens is empty"
             # _current_context must point at the freshly entered Context

@@ -264,8 +262,6 @@ class TestRestoreRequestContext:
         import weakref
         from unittest.mock import MagicMock

-        from mcp.server.lowlevel.server import request_ctx
-
         from fastmcp import FastMCP

[jlowin]

Three things before merge:

1. Use a context manager for Context lifecycle. The current shape (ctx = await _restore_request_context(...) + try/finally: await ctx.__aexit__(None, None, None)) is hand-rolling lifecycle management that the language already has syntax for. It also loses exception info: when the handler raises, __aexit__ gets (None, None, None) instead of the actual exception tuple. Two ways to fix. Make _restore_request_context an @asynccontextmanager and have the wrapper do async with _restore_request_context(rc_ref): return await handler(...), or take an AsyncExitStack as a parameter and push the Context onto it (matches how _lifespan_manager handles conditional entries). Either is fine.

2. Retitle. The title says "forward __aenter__ to wrapped client in StatefulProxyClient" but the change is in _restore_request_context / _make_restoring_handler. Nothing to do with StatefulProxyClient.__aenter__. Something like fix(proxy): enter Context properly in _restore_request_context.

3. Hoist test imports to module level. Every method in TestRestoreRequestContext has its imports inside the function body. Move them to the top of tests/server/providers/proxy/test_stateful_proxy_client.py.

[strawgate]

This PR targets Bug 4 from #4054, and that gap is real: _restore_request_context never sets _current_server, so any user-supplied sampling/elicitation/roots callback that calls get_server() or uses a DI dependency raises "No FastMCP server instance in context". That's worth fixing and stays tracked on #4054.

The reason we're closing this rather than iterating on it: the bug is that the path doesn't call __aexit__ — but the fix isn't to call __aexit__. It's that this path shouldn't be opening a context-manager scope in the first place.

The StatefulProxyClient receive loop is a single long-lived task. _restore_request_context isn't opening a scope around a unit of work — it's repairing stale ContextVars in place on a task that keeps running and handles the next message, where the next repair just overwrites. A context manager models enter → bounded work → exit; there is no natural "exit" here. Bolting one on (here, or via the @asynccontextmanager refactor suggested in review) invents a per-handler scope around something deliberately unscoped, and that invented scope is what generates the follow-on problems:

• Exception info loss: finally: await ctx.__aexit__(None, None, None) drops the real exception. Context.__aexit__ forwards the exc tuple to SharedContext.__aexit__ → its AsyncExitStack, so error-sensitive Shared() CMs see success on a failed handler.
• Cross-session _request_state bleed: Context.__aenter__ does self._request_state = parent._request_state from the current contextvar. In the shared-ref / copy.copy concurrent case this function's own docstring warns about, that "parent" is another session's Context, so the entered Context grafts a foreign session's request state.
• Per-notification churn: enter/exit now run on every restored callback, including each forwarded log/progress, allocating+entering+exiting a SharedContext (or churning docket/worker tokens) per message.

All three exist because a scope was opened that then has to be unwound — they aren't independent bugs to patch.

Decomposing what __aenter__ does on this path: _current_context.set, _current_server.set, and the docket/worker sets are pure ContextVar writes that are correct to do set-only on a long-lived task — ContextVars don't leak, they're overwritten, and not reset()-ing is the point (we never want the stale prior value back). _request_state inheritance must be skipped, not performed. The only side-effect that genuinely wants a scoped lifetime is SharedContext, and a proxy receive-loop callback using a Shared() dependency is exotic, never worked before this PR, and isn't a regression to leave unsupported.

So the minimal, principled fix keeps the original set-only model and adds the one missing line:

request_ctx.set(rc)
fastmcp = fastmcp_ref()
if fastmcp is not None:
    _current_context.set(Context(fastmcp))
    _current_server.set(weakref.ref(fastmcp))   # the actual Bug 4 gap

No __aenter__/__aexit__, so the exc-info issue, the cross-session bleed, and the per-call churn don't exist — not because they're handled, but because there's no scope to get wrong. If Shared() in receive-loop callbacks ever becomes a real requirement, the right move is to scope only a SharedContext per handler call, not to CM-wrap the whole Context.

This is also a good illustration of why CONTRIBUTING.md asks for the planned approach to be discussed on the issue before a PR — the lifecycle direction here would have been worth aligning on up front, before the implementation and review cycles. Closing this in favor of the set-only fix tracked on #4054.