Fix OAuth token expiry handling (#1649) (#1671)

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}")

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