fix: propagate upstream_claims in load_access_token

[kvdhanush06]

Description

Closes #3723

FastMCP's OAuthProxy discarded custom claims embedded in the JWT under the upstream_claims key during the load_access_token process. While these claims were correctly signed into the JWT at issuance (via _extract_upstream_claims), they were lost during the token swap where the FastMCP JWT is exchanged for the upstream provider's token.

Root Cause

The load_access_token method verified the FastMCP JWT to retrieve the jti for mapping lookup but failed to preserve the verified payload. It then validated the upstream token and returned a fresh AccessToken object that only contained the provider's claims, losing the contextually relevant upstream_claims that the server had previously verified and signed.

Fix

Updated load_access_token in src/fastmcp/server/auth/oauth_proxy/proxy.py to:

1. Extract upstream_claims from the verified FastMCP JWT payload.
2. Merge these claims into the returned AccessToken.claims dictionary.

This ensures that any custom data intended to survive the token swap (e.g. cross-provider user IDs or internal metadata) is accessible to the downstream MCP server.

Verification

Added a regression test suite TestUpstreamClaimsPropagation in tests/server/auth/test_oidc_proxy_token.py covering:

• Correct extraction of upstream_claims from the FastMCP JWT.
• Successful merging of these claims into the final AccessToken after the upstream token swap.

# The regression test ensures the final AccessToken includes the propagated claims
result = await proxy.load_access_token(fastmcp_jwt)
assert "upstream_claims" in result.claims
assert result.claims["upstream_claims"] == {"sub": "idp-user-123"}

Contribution type

• Bug fix (simple, well-scoped fix for a clearly broken behavior)
• Documentation improvement
• Enhancement (maintainers typically implement enhancements — see CONTRIBUTING.md)

Checklist

• This PR addresses an existing issue (or fixes a self-evident bug)
• I have read CONTRIBUTING.md
• I have added tests that cover my changes
• I have run uv run prek run --all-files and all checks pass
• I have self-reviewed my changes
• If I used an LLM, it followed the repo's contributing conventions (not generic output)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f8f89e7b72

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Copy claims before injecting upstream_claims

load_access_token mutates validated.claims in place, but some verifiers (notably IntrospectionTokenVerifier) return cached AccessToken instances by reference (src/fastmcp/server/auth/providers/introspection.py:197-200,292). In that setup, this write can leak/overwrite claim state across requests sharing the same upstream token object; e.g., a request with upstream_claims can persist data that is then returned for a later token where upstream_claims is absent (the branch is skipped, so stale data remains). Copying the token/claims before mutation avoids cross-request contamination.

Useful? React with 👍 / 👎.

[Copilot]

Pull request overview

This PR fixes FastMCP OAuth proxy token swapping so that custom upstream_claims embedded (and signed) into the FastMCP JWT are preserved and surfaced on the AccessToken returned by load_access_token, addressing #3723.

Changes:

• Preserve upstream_claims from the verified FastMCP JWT payload during load_access_token.
• Merge propagated upstream_claims into the final returned AccessToken.claims.
• Add a regression test covering upstream-claims propagation through the token swap.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.

File	Description
src/fastmcp/server/auth/oauth_proxy/proxy.py	Extracts upstream_claims from verified FastMCP JWT and injects them into the returned AccessToken after upstream validation.
tests/server/auth/test_oidc_proxy_token.py	Adds an async regression test ensuring upstream_claims survive the swap and appear on AccessToken.claims.

[Copilot]

load_access_token mutates validated.claims in-place when injecting upstream_claims. If the configured _token_validator caches and reuses AccessToken instances (e.g., IntrospectionTokenVerifier caches results), this can leak/retain upstream_claims across requests/tokens. Prefer returning a copied AccessToken with merged claims (e.g., via model_copy/new dict) to avoid mutating shared/cached objects and to ensure the claims dict isn’t shared (shallow copies can still share nested dicts).

[Copilot]

The module-level docstring at the top of this test file was removed. Most test modules in this repo start with a brief module docstring; restoring it helps keep documentation/style consistent across the test suite.

[Copilot]

Using JTIMapping.model_construct(...) bypasses Pydantic validation, which can mask schema/type issues and make the test less representative. Since all required fields are provided here, prefer constructing JTIMapping(...) normally so the test will fail if the model contract changes.

[Copilot]

Similarly, UpstreamTokenSet.model_construct(...) bypasses validation and can hide issues in the test setup. Since this test provides all required fields with valid types, prefer UpstreamTokenSet(...) to ensure validation stays exercised and future model changes are caught.

[Copilot]

AccessToken.model_construct(...) creates an AccessToken without fields that real verifiers typically set (e.g., client_id) and bypasses validation entirely, so the test can pass with an invalid object shape. Prefer constructing a valid AccessToken(...) (including client_id) so the regression test exercises realistic behavior and will fail if the AccessToken contract changes.

[marvin-context-protocol]

Test Failure Analysis

Summary: The static analysis (ty type checker) failed because two # ty: ignore[invalid-assignment] suppression comments added in the new test code are flagged as unused — ty no longer considers those assignments to be type errors, so the suppressions should be removed.

Root Cause: In tests/server/auth/test_oidc_proxy_token.py, the PR added # ty: ignore[invalid-assignment] to two lines:

• Line 549: proxy._jti_mapping_store = MagicMock() # ty: ignore[invalid-assignment]
• Line 557: proxy._upstream_token_store = MagicMock() # ty: ignore[invalid-assignment]

ty treats unused suppression comments as errors by default (warning[unused-ignore-comment]), and since these assignments are no longer type errors (perhaps due to how ty infers the types), the suppression comments are unnecessary and must be removed.

Suggested Solution: Remove the two # ty: ignore[invalid-assignment] comments from tests/server/auth/test_oidc_proxy_token.py:

# Line 549 — change to:
proxy._jti_mapping_store = MagicMock()

# Line 557 — change to:
proxy._upstream_token_store = MagicMock()

Note: other # ty: ignore[invalid-assignment] comments in the same test file (e.g., for proxy._token_validator.verify_token) did not trigger warnings, so only these two need to be removed.

Detailed Analysis

prek ran three hooks: ruff check ✅, ruff format ✅, ty check ❌

The ty output:

warning[unused-ignore-comment]: Unused `ty: ignore` directive
 --> tests/server/auth/test_oidc_proxy_token.py:549:53
548 |             # 2. Mock storage for first request
549 |             proxy._jti_mapping_store = MagicMock()  # ty: ignore[invalid-assignment]
    |                                                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
help: Remove the unused suppression comment

warning[unused-ignore-comment]: Unused `ty: ignore` directive
 --> tests/server/auth/test_oidc_proxy_token.py:557:56
557 |             proxy._upstream_token_store = MagicMock()  # ty: ignore[invalid-assignment]
    |                                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
help: Remove the unused suppression comment

Found 2 diagnostics

ty's own suggested fix is to remove the comments — the assignments are not actually type errors in its view.

Related Files

• tests/server/auth/test_oidc_proxy_token.py — contains the unused suppression comments at lines 549 and 557 that need to be removed.

---

Posted by marvin, the triage bot

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[jlowin]

Thanks for the fix — the behavior change itself is correct and well-reasoned. The upstream_claims should survive load_access_token, and the model_copy(deep=True) to avoid mutating a cached verifier result is the right call.

However, the tests need to be simplified. The existing TestTransparentUpstreamRefresh class in tests/server/auth/oauth_proxy/test_tokens.py demonstrates the pattern we use for testing load_access_token: a real OAuthProxy with real in-memory stores, populated directly, no mocked stores. Your tests mock every internal store with MagicMock, which is more brittle and harder to read.

Since the propagation logic lives in OAuthProxy.load_access_token (not OIDCProxy), the tests should live in test_tokens.py alongside the other load_access_token tests, not in test_oidc_proxy_token.py. This also eliminates the OIDC configuration boilerplate.

Following the existing pattern, you'd add a setup helper and two concise tests:

async def _setup_session_with_claims(
    self, proxy, *, upstream_claims=None
):
    """Set up a proxy JWT pointing at a valid upstream token, with optional upstream_claims."""
    upstream_token_id = "upstream-tok-id"
    access_jti = "test-claims-jti"

    upstream_token_set = UpstreamTokenSet(
        upstream_token_id=upstream_token_id,
        access_token="valid-upstream-access",
        refresh_token=None,
        refresh_token_expires_at=None,
        expires_at=time.time() + 3600,
        token_type="Bearer",
        scope="read",
        client_id="test-client",
        created_at=time.time(),
    )
    await proxy._upstream_token_store.put(
        key=upstream_token_id, value=upstream_token_set, ttl=3600,
    )
    await proxy._jti_mapping_store.put(
        key=access_jti,
        value=JTIMapping(
            jti=access_jti, upstream_token_id=upstream_token_id, created_at=time.time(),
        ),
        ttl=3600,
    )
    return proxy.jwt_issuer.issue_access_token(
        client_id="test-client", scopes=["read"], jti=access_jti,
        expires_in=3600, upstream_claims=upstream_claims,
    )

Then the tests themselves become straightforward:

async def test_upstream_claims_propagated(self, proxy):
    jwt = await self._setup_session_with_claims(
        proxy, upstream_claims={"sub": "user-123"}
    )
    result = await proxy.load_access_token(jwt)
    assert result is not None
    assert result.claims["upstream_claims"] == {"sub": "user-123"}

async def test_upstream_claims_not_mutated_on_cached_token(self, proxy, mock_verifier):
    jwt = await self._setup_session_with_claims(
        proxy, upstream_claims={"sub": "user-123"}
    )
    result = await proxy.load_access_token(jwt)
    assert result is not None
    assert result.claims["upstream_claims"] == {"sub": "user-123"}
    # Original verifier result must not be mutated
    for call in mock_verifier.verify_token.call_args_list:
        returned = await mock_verifier.verify_token(call.args[0])
        if returned:
            assert "upstream_claims" not in returned.claims

A few other small things:

• Drop @pytest.mark.asyncio — asyncio_mode = "auto" is configured globally
• The cast(AccessToken, result) after assert result is not None is redundant — the assert already narrows the type