From 5fe8e2fe6b457d13a91e22ffdbba846712851468 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 29 Aug 2025 21:14:56 -0400 Subject: [PATCH] Fix OAuth token expiry handling (#1649) (#1671) --- src/fastmcp/client/auth/oauth.py | 65 ++++++-- tests/client/auth/test_oauth_token_expiry.py | 152 +++++++++++++++++++ 2 files changed, 208 insertions(+), 9 deletions(-) create mode 100644 tests/client/auth/test_oauth_token_expiry.py diff --git a/src/fastmcp/client/auth/oauth.py b/src/fastmcp/client/auth/oauth.py index 88a8a9d33..b589c2873 100644 --- a/src/fastmcp/client/auth/oauth.py +++ b/src/fastmcp/client/auth/oauth.py @@ -4,6 +4,7 @@ import asyncio import json import webbrowser from asyncio import Future +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Literal from urllib.parse import urlparse @@ -18,7 +19,7 @@ from mcp.shared.auth import ( from mcp.shared.auth import ( OAuthToken as OAuthToken, ) -from pydantic import AnyHttpUrl, ValidationError +from pydantic import AnyHttpUrl, BaseModel, TypeAdapter, ValidationError from uvicorn.server import Server from fastmcp import settings as fastmcp_global_settings @@ -33,6 +34,17 @@ __all__ = ["OAuth"] logger = get_logger(__name__) +class StoredToken(BaseModel): + """Token storage format with absolute expiry time.""" + + token_payload: OAuthToken + expires_at: datetime | None + + +# Create TypeAdapter at module level for efficient parsing +stored_token_adapter = TypeAdapter(StoredToken) + + def default_cache_dir() -> Path: return fastmcp_global_settings.home / "oauth-mcp-client-cache" @@ -77,13 +89,28 @@ class FileTokenStorage(TokenStorage): path = self._get_file_path("tokens") try: - tokens = OAuthToken.model_validate_json(path.read_text()) - # now = datetime.datetime.now(datetime.timezone.utc) - # if tokens.expires_at is not None and tokens.expires_at <= now: - # logger.debug(f"Token expired for {self.get_base_url(self.server_url)}") - # return None - return tokens - except (FileNotFoundError, json.JSONDecodeError, ValidationError) as e: + # Parse JSON and validate as StoredToken + stored = stored_token_adapter.validate_json(path.read_text()) + + # Check if token is expired + if stored.expires_at is not None: + now = datetime.now(timezone.utc) + if now >= stored.expires_at: + logger.debug( + f"Token expired for {self.get_base_url(self.server_url)}" + ) + return None + + # Recalculate expires_in to be correct relative to now + if stored.token_payload.expires_in is not None: + remaining = stored.expires_at - now + stored.token_payload.expires_in = max( + 0, int(remaining.total_seconds()) + ) + + return stored.token_payload + + except (FileNotFoundError, ValidationError) as e: logger.debug( f"Could not load tokens for {self.get_base_url(self.server_url)}: {e}" ) @@ -92,7 +119,18 @@ class FileTokenStorage(TokenStorage): async def set_tokens(self, tokens: OAuthToken) -> None: """Save tokens to file storage.""" path = self._get_file_path("tokens") - path.write_text(tokens.model_dump_json(indent=2)) + + # Calculate absolute expiry time if expires_in is present + expires_at = None + if tokens.expires_in is not None: + expires_at = datetime.now(timezone.utc) + timedelta( + seconds=tokens.expires_in + ) + + # Create StoredToken and save using Pydantic serialization + stored = StoredToken(token_payload=tokens, expires_at=expires_at) + + path.write_text(stored.model_dump_json(indent=2)) logger.debug(f"Saved tokens for {self.get_base_url(self.server_url)}") async def get_client_info(self) -> OAuthClientInformationFull | None: @@ -252,6 +290,15 @@ class OAuth(OAuthClientProvider): callback_handler=self.callback_handler, ) + async def _initialize(self) -> None: + """Load stored tokens and client info, properly setting token expiry.""" + # Call parent's _initialize to load tokens and client info + await super()._initialize() + + # If tokens were loaded and have expires_in, update the context's token_expiry_time + if self.context.current_tokens and self.context.current_tokens.expires_in: + self.context.update_token_expiry(self.context.current_tokens) + async def redirect_handler(self, authorization_url: str) -> None: """Open browser for authorization.""" logger.info(f"OAuth authorization URL: {authorization_url}") diff --git a/tests/client/auth/test_oauth_token_expiry.py b/tests/client/auth/test_oauth_token_expiry.py new file mode 100644 index 000000000..77e5d552e --- /dev/null +++ b/tests/client/auth/test_oauth_token_expiry.py @@ -0,0 +1,152 @@ +"""Test OAuth token expiry handling with absolute timestamps.""" + +import json +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest +from mcp.shared.auth import OAuthToken + +from fastmcp.client.auth.oauth import FileTokenStorage + + +@pytest.mark.asyncio +async def test_token_storage_with_expiry(tmp_path: Path): + """Test that tokens are stored with absolute expiry time and loaded correctly.""" + storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) + + # Create a token with 3600 seconds expiry + token = OAuthToken( + access_token="test_token", + token_type="Bearer", + expires_in=3600, + refresh_token="refresh_token", + ) + + # Save the token + await storage.set_tokens(token) + + # Check that the file contains the dataclass format + token_file = storage._get_file_path("tokens") + data = json.loads(token_file.read_text()) + + assert "token_payload" in data + assert "expires_at" in data + assert data["expires_at"] is not None + # expires_at should be approximately now + 3600 seconds + expires_at = datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00")) + expected = datetime.now(timezone.utc) + timedelta(seconds=3600) + assert abs((expires_at - expected).total_seconds()) < 2 + + # Load the token back + loaded_token = await storage.get_tokens() + assert loaded_token is not None + assert loaded_token.access_token == "test_token" + # expires_in should be recalculated to be approximately 3600 (minus loading time) + assert loaded_token.expires_in is not None + assert 3595 <= loaded_token.expires_in <= 3600 + + +@pytest.mark.asyncio +async def test_expired_token_returns_none(tmp_path: Path): + """Test that expired tokens return None when loaded.""" + storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) + + # Manually create an already-expired token file + token_file = storage._get_file_path("tokens") + past_expiry = datetime.now(timezone.utc) - timedelta( + seconds=10 + ) # Expired 10 seconds ago + + expired_token = { + "token_payload": { + "access_token": "test_token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "refresh_token", + }, + "expires_at": past_expiry.isoformat(), + } + token_file.write_text(json.dumps(expired_token, indent=2, default=str)) + + # Load the token - should return None since it's expired + loaded_token = await storage.get_tokens() + assert loaded_token is None + + +@pytest.mark.asyncio +async def test_token_without_expiry(tmp_path: Path): + """Test that tokens without expires_in are handled correctly.""" + storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) + + # Create a token without expires_in (perpetual token) + token = OAuthToken( + access_token="test_token", + token_type="Bearer", + expires_in=None, + refresh_token="refresh_token", + ) + + # Save the token + await storage.set_tokens(token) + + # Check that expires_at is None in the file + token_file = storage._get_file_path("tokens") + data = json.loads(token_file.read_text()) + assert data["expires_at"] is None + + # Load the token back - should work since no expiry + loaded_token = await storage.get_tokens() + assert loaded_token is not None + assert loaded_token.access_token == "test_token" + assert loaded_token.expires_in is None + + +@pytest.mark.asyncio +async def test_invalid_format_returns_none(tmp_path: Path): + """Test that invalid token format returns None.""" + storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) + + # Manually write an invalid format token file (missing required fields) + token_file = storage._get_file_path("tokens") + invalid_token = { + "access_token": "invalid_token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "refresh_token", + } + token_file.write_text(json.dumps(invalid_token, indent=2)) + + # Try to load - should return None + loaded_token = await storage.get_tokens() + assert loaded_token is None + + +@pytest.mark.asyncio +async def test_token_expiry_recalculated_on_load(tmp_path: Path): + """Test that expires_in is correctly recalculated when loading tokens.""" + storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) + + # Manually create a token file with a specific expires_at + token_file = storage._get_file_path("tokens") + future_expiry = datetime.now(timezone.utc) + timedelta( + seconds=1800 + ) # 30 minutes from now + + stored_token = { + "token_payload": { + "access_token": "test_token", + "token_type": "Bearer", + "expires_in": 3600, # Original value (will be recalculated) + "refresh_token": "refresh_token", + }, + "expires_at": future_expiry.isoformat(), + } + token_file.write_text(json.dumps(stored_token, indent=2, default=str)) + + # Load the token + loaded_token = await storage.get_tokens() + assert loaded_token is not None + # expires_in should be recalculated to approximately 1800 seconds + assert loaded_token.expires_in is not None + assert 1795 <= loaded_token.expires_in <= 1800