From 52feff6878418a3a84f8a8860298fb52d7d89478 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sun, 22 Mar 2026 10:43:44 -0400
Subject: [PATCH] Transparently refresh upstream token in
 OAuthProxy.load_access_token() (#3584)

* Transparently refresh upstream token in OAuthProxy.load_access_token()

When upstream token validation fails during load_access_token, attempt
to refresh using the stored refresh token before returning None. This
prevents premature 401s that force clients into expensive full re-auth
flows when the upstream token expires.

Co-authored-by: Claude <noreply@anthropic.com>

* Gate transparent refresh on token expiry, add advisory lock

Only attempt upstream refresh when the token is actually expired, not
on any validation failure (scope mismatch, revocation, etc.). Add
per-token advisory lock to prevent concurrent async tasks from racing
to refresh the same upstream token.

* Re-check expiry inside lock, reload from storage after refresh failure

---------

Co-authored-by: Claude <noreply@anthropic.com>
---
 src/fastmcp/server/auth/oauth_proxy/proxy.py | 168 ++++++++++-
 tests/server/auth/oauth_proxy/test_tokens.py | 294 ++++++++++++++++++-
 2 files changed, 460 insertions(+), 2 deletions(-)

diff --git a/src/fastmcp/server/auth/oauth_proxy/proxy.py b/src/fastmcp/server/auth/oauth_proxy/proxy.py
index a78e6dc1b..6868a1897 100644
--- a/src/fastmcp/server/auth/oauth_proxy/proxy.py
+++ b/src/fastmcp/server/auth/oauth_proxy/proxy.py
@@ -25,6 +25,7 @@ from base64 import urlsafe_b64encode
 from typing import Any, Literal
 from urllib.parse import urlencode, urlparse, urlunparse
 
+import anyio
 import httpx
 from authlib.common.security import generate_token
 from authlib.integrations.httpx_client import AsyncOAuth2Client
@@ -546,6 +547,13 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
                 allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
             )
 
+        # Advisory locks for transparent upstream token refresh, keyed by
+        # upstream_token_id. Prevents concurrent async tasks from racing to
+        # refresh the same token within a single process. Does not protect
+        # against cross-process races in distributed deployments — those are
+        # handled by re-reading from storage after refresh failure.
+        self._refresh_locks: dict[str, anyio.Lock] = {}
+
         logger.debug(
             "Initialized OAuth proxy provider with upstream server %s",
             self._upstream_authorization_endpoint,
@@ -1457,6 +1465,90 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
         """
         return False
 
+    async def _try_transparent_refresh(
+        self,
+        upstream_token_set: UpstreamTokenSet,
+    ) -> UpstreamTokenSet:
+        """Refresh the upstream token transparently and update storage.
+
+        Called during load_access_token when the upstream token has expired
+        but a refresh token is available. This avoids returning a 401 that
+        would force the client into a full re-authentication flow.
+
+        Mutates and returns the upstream_token_set with refreshed token data.
+        Raises on failure (caller should catch and fall through to None).
+        """
+        scopes = upstream_token_set.scope.split() if upstream_token_set.scope else []
+        upstream_scopes = self._prepare_scopes_for_upstream_refresh(scopes)
+        oauth_client = self._create_upstream_oauth_client()
+
+        token_response: dict[str, Any] = await oauth_client.refresh_token(
+            url=self._upstream_token_endpoint,
+            refresh_token=upstream_token_set.refresh_token,
+            scope=" ".join(upstream_scopes) if upstream_scopes else None,
+            **self._extra_token_params,
+        )
+        logger.debug(
+            "Transparent upstream refresh succeeded (token_id=%s)",
+            upstream_token_set.upstream_token_id[:8],
+        )
+
+        # Calculate new expiry
+        if "expires_in" in token_response:
+            new_expires_in = int(token_response["expires_in"])
+        elif self._fallback_access_token_expiry_seconds is not None:
+            new_expires_in = self._fallback_access_token_expiry_seconds
+        else:
+            new_expires_in = DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS
+
+        upstream_token_set.access_token = token_response["access_token"]
+        upstream_token_set.expires_at = time.time() + new_expires_in
+        upstream_token_set.scope = " ".join(
+            parse_scopes(token_response["scope"]) or []
+            if "scope" in token_response
+            else scopes
+        )
+
+        # Handle upstream refresh token rotation
+        new_refresh_expires_in = None
+        if new_upstream_refresh := token_response.get("refresh_token"):
+            if new_upstream_refresh != upstream_token_set.refresh_token:
+                upstream_token_set.refresh_token = new_upstream_refresh
+            if "refresh_expires_in" in token_response and int(
+                token_response["refresh_expires_in"]
+            ):
+                new_refresh_expires_in = int(token_response["refresh_expires_in"])
+                upstream_token_set.refresh_token_expires_at = (
+                    time.time() + new_refresh_expires_in
+                )
+            elif upstream_token_set.refresh_token_expires_at:
+                new_refresh_expires_in = int(
+                    upstream_token_set.refresh_token_expires_at - time.time()
+                )
+            else:
+                new_refresh_expires_in = 60 * 60 * 24 * 30
+                upstream_token_set.refresh_token_expires_at = (
+                    time.time() + new_refresh_expires_in
+                )
+
+        upstream_token_set.raw_token_data = {
+            **upstream_token_set.raw_token_data,
+            **token_response,
+        }
+
+        refresh_ttl = new_refresh_expires_in or (
+            int(upstream_token_set.refresh_token_expires_at - time.time())
+            if upstream_token_set.refresh_token_expires_at
+            else 60 * 60 * 24 * 30
+        )
+        await self._upstream_token_store.put(
+            key=upstream_token_set.upstream_token_id,
+            value=upstream_token_set,
+            ttl=max(refresh_ttl, new_expires_in, 1),
+        )
+
+        return upstream_token_set
+
     async def load_access_token(self, token: str) -> AccessToken | None:  # type: ignore[override]
         """Validate FastMCP JWT by swapping for upstream token.
 
@@ -1465,7 +1557,8 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
         2. Look up upstream token via JTI mapping
         3. Decrypt upstream token
         4. Validate upstream token with provider (GitHub API, JWT validation, etc.)
-        5. Return upstream validation result
+        5. If upstream validation fails, attempt transparent refresh
+        6. Return upstream validation result
 
         The FastMCP JWT is a reference token - all authorization data comes
         from validating the upstream token via the TokenVerifier.
@@ -1501,6 +1594,79 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
                 return None
             validated = await self._token_validator.verify_token(verification_token)
 
+            # 4. If upstream validation failed due to token expiry and we
+            # have a refresh token, attempt transparent refresh to avoid
+            # forcing the client into a full re-auth flow. Only refresh on
+            # expiry — other failures (scope mismatch, revocation) won't be
+            # helped by a refresh and would just burn tokens.
+            if (
+                not validated
+                and upstream_token_set.refresh_token
+                and upstream_token_set.expires_at <= time.time()
+            ):
+                try:
+                    token_id = upstream_token_set.upstream_token_id
+
+                    # Advisory lock prevents concurrent requests from racing
+                    # to refresh the same upstream token.
+                    if token_id not in self._refresh_locks:
+                        self._refresh_locks[token_id] = anyio.Lock()
+                    lock = self._refresh_locks[token_id]
+
+                    async with lock:
+                        # Re-read from storage — another task may have
+                        # already refreshed while we waited for the lock.
+                        upstream_token_set = (
+                            await self._upstream_token_store.get(key=token_id)
+                            or upstream_token_set
+                        )
+
+                        verification_token = self._get_verification_token(
+                            upstream_token_set
+                        )
+                        if verification_token is not None:
+                            validated = await self._token_validator.verify_token(
+                                verification_token
+                            )
+
+                        # Only refresh if the (possibly reloaded) token is
+                        # still expired — a non-expiry failure on a fresh
+                        # token (scope mismatch, revocation) won't be
+                        # helped by refreshing.
+                        if (
+                            not validated
+                            and upstream_token_set.expires_at <= time.time()
+                        ):
+                            upstream_token_set = await self._try_transparent_refresh(
+                                upstream_token_set
+                            )
+                            verification_token = self._get_verification_token(
+                                upstream_token_set
+                            )
+                            if verification_token is not None:
+                                validated = await self._token_validator.verify_token(
+                                    verification_token
+                                )
+                except Exception as e:
+                    logger.debug("Transparent upstream refresh failed: %s", e)
+                    # In a distributed deployment, another worker may have
+                    # already refreshed and rotated the token, causing our
+                    # stale refresh token to fail. Re-read and re-validate.
+                    try:
+                        reloaded = await self._upstream_token_store.get(
+                            key=upstream_token_set.upstream_token_id
+                        )
+                        if reloaded:
+                            verification_token = self._get_verification_token(reloaded)
+                            if verification_token is not None:
+                                validated = await self._token_validator.verify_token(
+                                    verification_token
+                                )
+                                if validated:
+                                    upstream_token_set = reloaded
+                    except Exception:
+                        pass
+
             if not validated:
                 logger.debug("Upstream token validation failed")
                 return None
diff --git a/tests/server/auth/oauth_proxy/test_tokens.py b/tests/server/auth/oauth_proxy/test_tokens.py
index f7579c714..739abf43e 100644
--- a/tests/server/auth/oauth_proxy/test_tokens.py
+++ b/tests/server/auth/oauth_proxy/test_tokens.py
@@ -7,7 +7,7 @@ import pytest
 from key_value.aio.stores.memory import MemoryStore
 from mcp.server.auth.handlers.token import TokenErrorResponse
 from mcp.server.auth.handlers.token import TokenHandler as SDKTokenHandler
-from mcp.server.auth.provider import AuthorizationCode
+from mcp.server.auth.provider import AccessToken, AuthorizationCode
 from mcp.shared.auth import OAuthClientInformationFull
 from pydantic import AnyUrl
 
@@ -17,6 +17,8 @@ from fastmcp.server.auth.oauth_proxy.models import (
     DEFAULT_ACCESS_TOKEN_EXPIRY_NO_REFRESH_SECONDS,
     DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS,
     ClientCode,
+    JTIMapping,
+    UpstreamTokenSet,
     _hash_token,
 )
 from fastmcp.server.auth.providers.jwt import JWTVerifier
@@ -563,3 +565,293 @@ class TestUpstreamTokenStorageTTL:
             key=_hash_token(result.refresh_token)
         )
         assert refresh_meta is not None
+
+
+class TestTransparentUpstreamRefresh:
+    """Tests for transparent upstream token refresh in load_access_token.
+
+    When the upstream token expires but a refresh token is available, the proxy
+    should transparently refresh the upstream token rather than returning a 401
+    that forces the client into a full re-authentication flow.
+    """
+
+    @pytest.fixture
+    def mock_verifier(self):
+        """Token verifier that rejects expired tokens, accepts refreshed ones."""
+        verifier = Mock(spec=TokenVerifier)
+        verifier.required_scopes = ["read"]
+
+        async def verify(token: str) -> AccessToken | None:
+            if token.startswith("refreshed-"):
+                return AccessToken(
+                    token=token,
+                    client_id="test-client",
+                    scopes=["read"],
+                    expires_at=int(time.time() + 3600),
+                )
+            return None
+
+        verifier.verify_token = AsyncMock(side_effect=verify)
+        return verifier
+
+    @pytest.fixture
+    def proxy(self, mock_verifier):
+        proxy = OAuthProxy(
+            upstream_authorization_endpoint="https://idp.example.com/authorize",
+            upstream_token_endpoint="https://idp.example.com/token",
+            upstream_client_id="test-client",
+            upstream_client_secret="test-secret",
+            token_verifier=mock_verifier,
+            base_url="https://proxy.example.com",
+            jwt_signing_key="test-secret-key",
+            client_storage=MemoryStore(),
+        )
+        proxy.set_mcp_path("/mcp")
+        return proxy
+
+    async def _setup_expired_session(
+        self,
+        proxy: OAuthProxy,
+        *,
+        upstream_refresh_token: str | None = "upstream-refresh-tok",
+    ) -> str:
+        """Set up a proxy JWT pointing at an expired upstream token.
+
+        Returns the proxy JWT (access token) that can be passed to
+        load_access_token.
+        """
+        upstream_token_id = "upstream-tok-id"
+        access_jti = "test-access-jti"
+
+        upstream_token_set = UpstreamTokenSet(
+            upstream_token_id=upstream_token_id,
+            access_token="expired-upstream-access",
+            refresh_token=upstream_refresh_token,
+            refresh_token_expires_at=time.time() + 86400
+            if upstream_refresh_token
+            else None,
+            expires_at=time.time() - 60,  # expired 1 minute ago
+            token_type="Bearer",
+            scope="read",
+            client_id="test-client",
+            created_at=time.time() - 3600,
+        )
+        await proxy._upstream_token_store.put(
+            key=upstream_token_id,
+            value=upstream_token_set,
+            ttl=86400,
+        )
+
+        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,
+        )
+
+        fastmcp_jwt = proxy.jwt_issuer.issue_access_token(
+            client_id="test-client",
+            scopes=["read"],
+            jti=access_jti,
+            expires_in=3600,
+        )
+        return fastmcp_jwt
+
+    async def test_transparent_refresh_on_expired_upstream(self, proxy):
+        """load_access_token refreshes upstream token when validation fails."""
+        fastmcp_jwt = await self._setup_expired_session(proxy)
+
+        mock_oauth_client = AsyncMock()
+        mock_oauth_client.refresh_token = AsyncMock(
+            return_value={
+                "access_token": "refreshed-upstream-access",
+                "token_type": "Bearer",
+                "expires_in": 3600,
+                "refresh_token": "upstream-refresh-tok",
+                "scope": "read",
+            }
+        )
+
+        with patch.object(
+            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
+        ):
+            result = await proxy.load_access_token(fastmcp_jwt)
+
+        assert result is not None
+        assert result.token == "refreshed-upstream-access"
+        mock_oauth_client.refresh_token.assert_called_once()
+
+    async def test_transparent_refresh_updates_stored_token(self, proxy):
+        """After transparent refresh, the stored upstream token is updated."""
+        fastmcp_jwt = await self._setup_expired_session(proxy)
+
+        mock_oauth_client = AsyncMock()
+        mock_oauth_client.refresh_token = AsyncMock(
+            return_value={
+                "access_token": "refreshed-upstream-access",
+                "token_type": "Bearer",
+                "expires_in": 7200,
+                "refresh_token": "upstream-refresh-tok",
+                "scope": "read",
+            }
+        )
+
+        with patch.object(
+            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
+        ):
+            await proxy.load_access_token(fastmcp_jwt)
+
+        stored = await proxy._upstream_token_store.get(key="upstream-tok-id")
+        assert stored is not None
+        assert stored.access_token == "refreshed-upstream-access"
+        assert stored.expires_at > time.time()
+
+    async def test_no_refresh_without_refresh_token(self, proxy):
+        """Without a refresh token, load_access_token returns None immediately."""
+        fastmcp_jwt = await self._setup_expired_session(
+            proxy, upstream_refresh_token=None
+        )
+
+        result = await proxy.load_access_token(fastmcp_jwt)
+
+        assert result is None
+
+    async def test_returns_none_when_refresh_fails(self, proxy):
+        """If the upstream refresh call fails, load_access_token returns None."""
+        fastmcp_jwt = await self._setup_expired_session(proxy)
+
+        mock_oauth_client = AsyncMock()
+        mock_oauth_client.refresh_token = AsyncMock(
+            side_effect=Exception("upstream refused refresh")
+        )
+
+        with patch.object(
+            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
+        ):
+            result = await proxy.load_access_token(fastmcp_jwt)
+
+        assert result is None
+
+    async def test_returns_none_when_refreshed_token_still_invalid(self, proxy):
+        """If the refreshed token also fails validation, returns None."""
+        fastmcp_jwt = await self._setup_expired_session(proxy)
+
+        mock_oauth_client = AsyncMock()
+        mock_oauth_client.refresh_token = AsyncMock(
+            return_value={
+                "access_token": "still-bad-token",
+                "token_type": "Bearer",
+                "expires_in": 3600,
+                "scope": "read",
+            }
+        )
+
+        with patch.object(
+            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
+        ):
+            result = await proxy.load_access_token(fastmcp_jwt)
+
+        # "still-bad-token" doesn't start with "refreshed-" so verifier rejects it
+        assert result is None
+
+    async def test_no_refresh_when_token_not_expired(self, proxy):
+        """Non-expiry validation failures (e.g. revocation) should not trigger refresh."""
+        upstream_token_id = "upstream-tok-id"
+        access_jti = "test-access-jti"
+
+        # Token is NOT expired — verification failure is for another reason
+        upstream_token_set = UpstreamTokenSet(
+            upstream_token_id=upstream_token_id,
+            access_token="revoked-upstream-access",
+            refresh_token="upstream-refresh-tok",
+            refresh_token_expires_at=time.time() + 86400,
+            expires_at=time.time() + 3600,  # still valid for 1 hour
+            token_type="Bearer",
+            scope="read",
+            client_id="test-client",
+            created_at=time.time() - 60,
+        )
+        await proxy._upstream_token_store.put(
+            key=upstream_token_id,
+            value=upstream_token_set,
+            ttl=86400,
+        )
+        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,
+        )
+        fastmcp_jwt = proxy.jwt_issuer.issue_access_token(
+            client_id="test-client",
+            scopes=["read"],
+            jti=access_jti,
+            expires_in=3600,
+        )
+
+        mock_oauth_client = AsyncMock()
+        mock_oauth_client.refresh_token = AsyncMock()
+
+        with patch.object(
+            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
+        ):
+            result = await proxy.load_access_token(fastmcp_jwt)
+
+        assert result is None
+        # Refresh should NOT have been attempted
+        mock_oauth_client.refresh_token.assert_not_called()
+
+    async def test_reload_from_storage_after_refresh_failure(self, proxy):
+        """If refresh fails, re-read from storage in case another worker refreshed."""
+        fastmcp_jwt = await self._setup_expired_session(proxy)
+
+        # Simulate: refresh fails (stale refresh token), but another worker
+        # already wrote a fresh upstream token to storage.
+        refreshed_token_set = UpstreamTokenSet(
+            upstream_token_id="upstream-tok-id",
+            access_token="refreshed-upstream-access",
+            refresh_token="new-refresh-tok",
+            refresh_token_expires_at=time.time() + 86400,
+            expires_at=time.time() + 3600,
+            token_type="Bearer",
+            scope="read",
+            client_id="test-client",
+            created_at=time.time(),
+        )
+
+        # The store is read three times:
+        # 1. Initial lookup in load_access_token
+        # 2. Re-read inside the advisory lock
+        # 3. Re-read after refresh failure (recovery path)
+        original_get = proxy._upstream_token_store.get
+        call_count = 0
+
+        async def mock_get(key: str) -> UpstreamTokenSet | None:
+            nonlocal call_count
+            call_count += 1
+            if call_count <= 2:
+                return await original_get(key)
+            # After refresh failure, return the "other worker's" refreshed token
+            return refreshed_token_set
+
+        mock_oauth_client = AsyncMock()
+        mock_oauth_client.refresh_token = AsyncMock(
+            side_effect=Exception("refresh token rotated by another worker")
+        )
+
+        with (
+            patch.object(
+                proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
+            ),
+            patch.object(proxy._upstream_token_store, "get", side_effect=mock_get),
+        ):
+            result = await proxy.load_access_token(fastmcp_jwt)
+
+        assert result is not None
+        assert result.token == "refreshed-upstream-access"