ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-21 13:04:18 +02:00
Code Issues Projects Releases Packages Wiki Activity

Transparently refresh upstream token in OAuthProxy.load_access_token() (#3584)

Browse source 
* 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>
This commit is contained in:
Jeremiah Lowin 2026-03-22 10:43:44 -04:00 • committed by GitHub
commit 52feff6878
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 460 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

168
src/fastmcp/server/auth/oauth_proxy/proxy.py
View file

@ -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

294
tests/server/auth/oauth_proxy/test_tokens.py
View file

@ -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"