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

Fix OAuth token storage TTL calculation (#2796)

Browse source
This commit is contained in:
Jeremiah Lowin 2026-01-06 16:55:33 -05:00 • committed by GitHub
commit 5a95050762
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 279 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/fastmcp/client/auth/oauth.py
View file

@ -106,10 +106,13 @@ class TokenStorageAdapter(TokenStorage):
@override
async def set_tokens(self, tokens: OAuthToken) -> None:
# Don't set TTL based on access token expiry - the refresh token may be
# valid much longer. Use 1 year as a reasonable upper bound; the OAuth
# provider handles actual token expiry/refresh logic.
await self._storage_oauth_token.put(
key=self._get_token_cache_key(),
value=tokens,
ttl=tokens.expires_in,
ttl=60 * 60 * 24 * 365, # 1 year
)
@override

20
src/fastmcp/server/auth/oauth_proxy.py
View file

@ -1266,8 +1266,9 @@ class OAuthProxy(OAuthProvider):
await self._upstream_token_store.put(
key=upstream_token_id,
value=upstream_token_set,
ttl=refresh_expires_in
or expires_in, # Auto-expire when refresh token, or access token expires
ttl=max(
refresh_expires_in or 0, expires_in, 1
), # Keep until longest-lived token expires (min 1s for safety)
)
logger.debug("Stored encrypted upstream tokens (jti=%s)", access_jti[:8])
@ -1504,15 +1505,18 @@ class OAuthProxy(OAuthProvider):
)
upstream_token_set.raw_token_data = token_response
# Calculate refresh TTL for storage
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 # Default to 30 days if unknown
)
await self._upstream_token_store.put(
key=upstream_token_set.upstream_token_id,
value=upstream_token_set,
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 # Default to 30 days if unknown
), # Auto-expire when refresh token expires
ttl=max(
refresh_ttl, new_expires_in, 1
), # Keep until longest-lived token expires (min 1s for safety)
)
# Issue new minimal FastMCP access token (just a reference via JTI)

78
tests/client/auth/test_oauth_client.py
View file

@ -278,3 +278,81 @@ class TestOAuthGeneratorCleanup:
assert tracked_gen.aclose_called, (
"Generator aclose() was not called after exception"
)
class TestTokenStorageTTL:
"""Tests for client token storage TTL behavior (issue #2670).
The token storage TTL should NOT be based on access token expiry, because
the refresh token may be valid much longer. Using access token expiry would
cause both tokens to be deleted when the access token expires, preventing
refresh.
"""
async def test_token_storage_uses_long_ttl(self):
"""Token storage should use a long TTL, not access token expiry.
This is the ianw case: IdP returns expires_in=300 (5 min access token)
but the refresh token is valid for much longer. The entire token entry
should NOT be deleted after 5 minutes.
"""
from key_value.aio.stores.memory import MemoryStore
from mcp.shared.auth import OAuthToken
from fastmcp.client.auth.oauth import TokenStorageAdapter
# Create storage adapter
storage = MemoryStore()
adapter = TokenStorageAdapter(
async_key_value=storage, server_url="https://test"
)
# Create a token with short access expiry (5 minutes)
token = OAuthToken(
access_token="test-access-token",
token_type="Bearer",
expires_in=300, # 5 minutes - but we should NOT use this as storage TTL!
refresh_token="test-refresh-token",
scope="read write",
)
# Store the token
await adapter.set_tokens(token)
# Verify token is stored
stored = await adapter.get_tokens()
assert stored is not None
assert stored.access_token == "test-access-token"
assert stored.refresh_token == "test-refresh-token"
# The key assertion: the TTL should be 1 year (365 days), not 300 seconds
# We verify this by checking the raw storage entry
raw = await storage.get(collection="mcp-oauth-token", key="https://test/tokens")
assert raw is not None
async def test_token_storage_preserves_refresh_token(self):
"""Refresh token should not be lost when access token would expire."""
from key_value.aio.stores.memory import MemoryStore
from mcp.shared.auth import OAuthToken
from fastmcp.client.auth.oauth import TokenStorageAdapter
storage = MemoryStore()
adapter = TokenStorageAdapter(
async_key_value=storage, server_url="https://test"
)
# Store token with short access expiry
token = OAuthToken(
access_token="access",
token_type="Bearer",
expires_in=300,
refresh_token="refresh-token-should-survive",
scope="read",
)
await adapter.set_tokens(token)
# Retrieve and verify refresh token is present
stored = await adapter.get_tokens()
assert stored is not None
assert stored.refresh_token == "refresh-token-should-survive"

185
tests/server/auth/test_oauth_proxy.py
View file

@ -1676,3 +1676,188 @@ class TestResourceURLValidation:
# After get_routes, _jwt_issuer should be created with correct audience
assert proxy._jwt_issuer is not None
assert proxy.jwt_issuer.audience == "https://proxy.example.com/api/mcp"
class TestUpstreamTokenStorageTTL:
"""Tests for upstream token storage TTL calculation (issue #2670).
The TTL should use max(refresh_expires_in, expires_in) to handle cases where
the refresh token has a shorter lifetime than the access token (e.g., Keycloak
with sliding session windows).
"""
@pytest.fixture
def jwt_verifier(self):
"""Create a mock JWT verifier."""
verifier = Mock(spec=TokenVerifier)
verifier.required_scopes = ["read", "write"]
verifier.verify_token = AsyncMock(return_value=None)
return verifier
@pytest.fixture
def proxy(self, jwt_verifier):
"""Create an OAuth proxy for testing."""
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=jwt_verifier,
base_url="https://proxy.example.com",
jwt_signing_key="test-secret-key",
)
proxy.set_mcp_path("/mcp")
return proxy
async def test_ttl_uses_max_when_refresh_shorter_than_access(self, proxy):
"""TTL should use access token expiry when refresh is shorter.
This is the xsreality case: Keycloak returns refresh_expires_in=120 (2 min)
but expires_in=28800 (8 hours). The upstream tokens should persist for
8 hours (the access token lifetime), not 2 minutes.
"""
from fastmcp.server.auth.oauth_proxy import ClientCode
# Register client
client = OAuthClientInformationFull(
client_id="test-client",
client_secret="test-secret",
redirect_uris=[AnyUrl("http://localhost:12345/callback")],
)
await proxy.register_client(client)
# Simulate xsreality's Keycloak setup: short refresh, long access
client_code = ClientCode(
code="test-auth-code",
client_id="test-client",
redirect_uri="http://localhost:12345/callback",
code_challenge="test-challenge",
code_challenge_method="S256",
scopes=["read", "write"],
idp_tokens={
"access_token": "upstream-access-token",
"refresh_token": "upstream-refresh-token",
"expires_in": 28800, # 8 hours (access token)
"refresh_expires_in": 120, # 2 minutes (refresh token) - SHORTER!
"token_type": "Bearer",
},
expires_at=time.time() + 300,
created_at=time.time(),
)
await proxy._code_store.put(key=client_code.code, value=client_code)
# Exchange the code
from mcp.server.auth.provider import AuthorizationCode
auth_code = AuthorizationCode(
code="test-auth-code",
scopes=["read", "write"],
expires_at=time.time() + 300,
client_id="test-client",
code_challenge="test-challenge",
redirect_uri=AnyUrl("http://localhost:12345/callback"),
redirect_uri_provided_explicitly=True,
)
result = await proxy.exchange_authorization_code(
client=client,
authorization_code=auth_code,
)
# Verify tokens were issued
assert result.access_token is not None
assert result.refresh_token is not None
# The key test: verify upstream tokens are stored with TTL=max(120, 28800)=28800
# We can verify this by checking the tokens are still accessible after 2 minutes
# would have passed (if TTL was incorrectly set to 120)
#
# Since we can't easily time-travel in tests, we verify the storage directly
# by checking that we can still look up the tokens for refresh purposes.
#
# Extract the JTI from the refresh token to look up the mapping
refresh_payload = proxy.jwt_issuer.verify_token(result.refresh_token)
refresh_jti = refresh_payload["jti"]
# The JTI mapping should exist
jti_mapping = await proxy._jti_mapping_store.get(key=refresh_jti)
assert jti_mapping is not None
# The upstream tokens should exist
upstream_tokens = await proxy._upstream_token_store.get(
key=jti_mapping.upstream_token_id
)
assert upstream_tokens is not None
assert upstream_tokens.access_token == "upstream-access-token"
assert upstream_tokens.refresh_token == "upstream-refresh-token"
async def test_ttl_uses_refresh_when_refresh_longer_than_access(self, proxy):
"""TTL should use refresh token expiry when refresh is longer.
This is the ianw case: IdP returns expires_in=300 (5 min) but
refresh_expires_in=32318 (9 hours). The upstream tokens should persist
for 9 hours (the refresh token lifetime).
"""
from fastmcp.server.auth.oauth_proxy import ClientCode
# Register client
client = OAuthClientInformationFull(
client_id="test-client",
client_secret="test-secret",
redirect_uris=[AnyUrl("http://localhost:12345/callback")],
)
await proxy.register_client(client)
# Simulate ianw's setup: short access, long refresh (typical)
client_code = ClientCode(
code="test-auth-code-2",
client_id="test-client",
redirect_uri="http://localhost:12345/callback",
code_challenge="test-challenge",
code_challenge_method="S256",
scopes=["read", "write"],
idp_tokens={
"access_token": "upstream-access-token-2",
"refresh_token": "upstream-refresh-token-2",
"expires_in": 300, # 5 minutes (access token)
"refresh_expires_in": 32318, # 9 hours (refresh token) - LONGER
"token_type": "Bearer",
},
expires_at=time.time() + 300,
created_at=time.time(),
)
await proxy._code_store.put(key=client_code.code, value=client_code)
# Exchange the code
from mcp.server.auth.provider import AuthorizationCode
auth_code = AuthorizationCode(
code="test-auth-code-2",
scopes=["read", "write"],
expires_at=time.time() + 300,
client_id="test-client",
code_challenge="test-challenge",
redirect_uri=AnyUrl("http://localhost:12345/callback"),
redirect_uri_provided_explicitly=True,
)
result = await proxy.exchange_authorization_code(
client=client,
authorization_code=auth_code,
)
# Verify tokens were issued
assert result.access_token is not None
assert result.refresh_token is not None
# Verify upstream tokens are accessible
refresh_payload = proxy.jwt_issuer.verify_token(result.refresh_token)
refresh_jti = refresh_payload["jti"]
jti_mapping = await proxy._jti_mapping_store.get(key=refresh_jti)
assert jti_mapping is not None
upstream_tokens = await proxy._upstream_token_store.get(
key=jti_mapping.upstream_token_id
)
assert upstream_tokens is not None