Add persistent storage for OAuth client registrations (#1879)

This commit is contained in:
Jeremiah Lowin 2025-09-21 10:56:08 -04:00 committed by GitHub
commit 7176d4f293
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 704 additions and 52 deletions

View file

@ -161,14 +161,25 @@ mcp = FastMCP(name="My Server", auth=auth)
<ParamField body="extra_token_params" type="dict[str, str] | None">
Additional parameters to forward to the upstream token endpoint during code exchange and token refresh. Useful for provider-specific requirements during token operations.
For example, some providers require additional context during token exchange:
```python
extra_token_params={"audience": "https://api.example.com"}
```
These parameters are included in all token requests to the upstream provider.
</ParamField>
<ParamField body="client_storage" type="KVStorage | None">
Storage backend for persisting OAuth client registrations. By default, clients are automatically persisted to disk in `~/.config/fastmcp/oauth-proxy-clients/`, allowing them to survive server restarts as long as the filesystem remains accessible. This means MCP clients only need to register once and can reconnect seamlessly after your server restarts.
```python
from fastmcp.utilities.storage import InMemoryStorage
# Use in-memory storage for testing (clients lost on restart)
auth = OAuthProxy(..., client_storage=InMemoryStorage())
```
</ParamField>
</Card>
### Provider-Specific Parameters

View file

@ -1,7 +1,6 @@
from __future__ import annotations
import asyncio
import json
import webbrowser
from asyncio import Future
from collections.abc import AsyncGenerator
@ -29,6 +28,7 @@ from fastmcp.client.oauth_callback import (
)
from fastmcp.utilities.http import find_available_port
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import JSONFileStorage
__all__ = ["OAuth"]
@ -62,13 +62,14 @@ class FileTokenStorage(TokenStorage):
Implements the mcp.client.auth.TokenStorage protocol.
Each instance is tied to a specific server URL for proper token isolation.
Uses JSONFileStorage internally for consistent file handling.
"""
def __init__(self, server_url: str, cache_dir: Path | None = None):
"""Initialize storage for a specific server URL."""
self.server_url = server_url
self.cache_dir = cache_dir or default_cache_dir()
self.cache_dir.mkdir(exist_ok=True, parents=True)
# Use JSONFileStorage for actual file operations
self._storage = JSONFileStorage(cache_dir or default_cache_dir())
@staticmethod
def get_base_url(url: str) -> str:
@ -76,28 +77,33 @@ class FileTokenStorage(TokenStorage):
parsed = urlparse(url)
return f"{parsed.scheme}://{parsed.netloc}"
def get_cache_key(self) -> str:
"""Generate a safe filesystem key from the server's base URL."""
def _get_storage_key(self, file_type: Literal["client_info", "tokens"]) -> str:
"""Get the storage key for the specified data type.
JSONFileStorage will handle making the key filesystem-safe.
"""
base_url = self.get_base_url(self.server_url)
return (
base_url.replace("://", "_")
.replace(".", "_")
.replace("/", "_")
.replace(":", "_")
)
return f"{base_url}_{file_type}"
def _get_file_path(self, file_type: Literal["client_info", "tokens"]) -> Path:
"""Get the file path for the specified cache file type."""
key = self.get_cache_key()
return self.cache_dir / f"{key}_{file_type}.json"
"""Get the file path for the specified cache file type.
This method is kept for backward compatibility with tests that access _get_file_path.
"""
key = self._get_storage_key(file_type)
return self._storage._get_file_path(key)
async def get_tokens(self) -> OAuthToken | None:
"""Load tokens from file storage."""
path = self._get_file_path("tokens")
key = self._get_storage_key("tokens")
data = await self._storage.get(key)
if data is None:
return None
try:
# Parse JSON and validate as StoredToken
stored = stored_token_adapter.validate_json(path.read_text())
# Parse and validate as StoredToken
stored = stored_token_adapter.validate_python(data)
# Check if token is expired
if stored.expires_at is not None:
@ -117,15 +123,15 @@ class FileTokenStorage(TokenStorage):
return stored.token_payload
except (FileNotFoundError, ValidationError) as e:
except ValidationError as e:
logger.debug(
f"Could not load tokens for {self.get_base_url(self.server_url)}: {e}"
f"Could not validate tokens for {self.get_base_url(self.server_url)}: {e}"
)
return None
async def set_tokens(self, tokens: OAuthToken) -> None:
"""Save tokens to file storage."""
path = self._get_file_path("tokens")
key = self._get_storage_key("tokens")
# Calculate absolute expiry time if expires_in is present
expires_at = None
@ -134,19 +140,22 @@ class FileTokenStorage(TokenStorage):
seconds=tokens.expires_in
)
# Create StoredToken and save using Pydantic serialization
# Create StoredToken and save using storage
# Note: JSONFileStorage will wrap this in {"data": ..., "timestamp": ...}
stored = StoredToken(token_payload=tokens, expires_at=expires_at)
path.write_text(stored.model_dump_json(indent=2))
await self._storage.set(key, stored.model_dump(mode="json"))
logger.debug(f"Saved tokens for {self.get_base_url(self.server_url)}")
async def get_client_info(self) -> OAuthClientInformationFull | None:
"""Load client information from file storage."""
path = self._get_file_path("client_info")
key = self._get_storage_key("client_info")
data = await self._storage.get(key)
if data is None:
return None
try:
client_info = OAuthClientInformationFull.model_validate_json(
path.read_text()
)
client_info = OAuthClientInformationFull.model_validate(data)
# Check if we have corresponding valid tokens
# If no tokens exist, the OAuth flow was incomplete and we should
# force a fresh client registration
@ -157,27 +166,31 @@ class FileTokenStorage(TokenStorage):
"OAuth flow may have been incomplete. Clearing client info to force fresh registration."
)
# Clear the incomplete client info
client_info_path = self._get_file_path("client_info")
client_info_path.unlink(missing_ok=True)
await self._storage.delete(key)
return None
return client_info
except (FileNotFoundError, json.JSONDecodeError, ValidationError) as e:
except ValidationError as e:
logger.debug(
f"Could not load client info for {self.get_base_url(self.server_url)}: {e}"
f"Could not validate client info for {self.get_base_url(self.server_url)}: {e}"
)
return None
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
"""Save client information to file storage."""
path = self._get_file_path("client_info")
path.write_text(client_info.model_dump_json(indent=2))
key = self._get_storage_key("client_info")
await self._storage.set(key, client_info.model_dump(mode="json"))
logger.debug(f"Saved client info for {self.get_base_url(self.server_url)}")
def clear(self) -> None:
"""Clear all cached data for this server."""
"""Clear all cached data for this server.
Note: This is a synchronous method for backward compatibility.
Uses direct file operations instead of async storage methods.
"""
file_types: list[Literal["client_info", "tokens"]] = ["client_info", "tokens"]
for file_type in file_types:
# Use the file path directly for synchronous deletion
path = self._get_file_path(file_type)
path.unlink(missing_ok=True)
logger.debug(f"Cleared OAuth cache for {self.get_base_url(self.server_url)}")

View file

@ -45,9 +45,11 @@ from starlette.requests import Request
from starlette.responses import RedirectResponse
from starlette.routing import Route
import fastmcp
from fastmcp.server.auth.auth import OAuthProvider, TokenVerifier
from fastmcp.server.auth.redirect_validation import validate_redirect_uri
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import JSONFileStorage, KVStorage
if TYPE_CHECKING:
pass
@ -254,6 +256,8 @@ class OAuthProxy(OAuthProvider):
extra_authorize_params: dict[str, str] | None = None,
# Extra parameters to forward to token endpoint
extra_token_params: dict[str, str] | None = None,
# Client storage
client_storage: KVStorage | None = None,
):
"""Initialize the OAuth proxy provider.
@ -287,6 +291,9 @@ class OAuthProxy(OAuthProvider):
Example: {"audience": "https://api.example.com"}
extra_token_params: Additional parameters to forward to the upstream token endpoint.
Useful for provider-specific parameters during token exchange.
client_storage: Storage implementation for OAuth client registrations.
Defaults to file-based storage in ~/.fastmcp/oauth-proxy-clients/ if not specified.
Pass any KVStorage implementation for custom storage backends.
"""
# Always enable DCR since we implement it locally for MCP clients
client_registration_options = ClientRegistrationOptions(
@ -335,8 +342,13 @@ class OAuthProxy(OAuthProvider):
self._extra_authorize_params = extra_authorize_params or {}
self._extra_token_params = extra_token_params or {}
# Local state for DCR and token bookkeeping
self._clients: dict[str, OAuthClientInformationFull] = {}
# Initialize client storage (default to file-based if not provided)
if client_storage is None:
cache_dir = fastmcp.settings.home / "oauth-proxy-clients"
client_storage = JSONFileStorage(cache_dir)
self._client_storage = client_storage
# Local state for token bookkeeping only (no client caching)
self._access_tokens: dict[str, AccessToken] = {}
self._refresh_tokens: dict[str, RefreshToken] = {}
@ -387,9 +399,20 @@ class OAuthProxy(OAuthProvider):
For unregistered clients, returns None (which will raise an error in the SDK).
"""
client = self._clients.get(client_id)
# Load from storage
data = await self._client_storage.get(client_id)
if not data:
return None
return client
if client_data := data.get("client", None):
return ProxyDCRClient(
allowed_redirect_uri_patterns=data.get(
"allowed_redirect_uri_patterns", self._allowed_client_redirect_uris
),
**client_data,
)
return None
async def register_client(self, client_info: OAuthClientInformationFull) -> None:
"""Register a client locally
@ -412,8 +435,12 @@ class OAuthProxy(OAuthProvider):
allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
)
# Store the ProxyDCRClient
self._clients[client_info.client_id] = proxy_client
# Store as structured dict with all needed metadata
storage_data = {
"client": proxy_client.model_dump(mode="json"),
"allowed_redirect_uri_patterns": self._allowed_client_redirect_uris,
}
await self._client_storage.set(client_info.client_id, storage_data)
# Log redirect URIs to help users discover what patterns they might need
if client_info.redirect_uris:

View file

@ -19,6 +19,7 @@ from fastmcp.server.auth import TokenVerifier
from fastmcp.server.auth.oauth_proxy import OAuthProxy
from fastmcp.server.auth.providers.jwt import JWTVerifier
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import KVStorage
logger = get_logger(__name__)
@ -212,6 +213,7 @@ class OIDCProxy(OAuthProxy):
redirect_path: str | None = None,
# Client configuration
allowed_client_redirect_uris: list[str] | None = None,
client_storage: KVStorage | None = None,
# Token validation configuration
token_endpoint_auth_method: str | None = None,
) -> None:
@ -234,6 +236,8 @@ class OIDCProxy(OAuthProxy):
If None (default), only localhost redirect URIs are allowed.
If empty list, all redirect URIs are allowed (not recommended for production).
These are for MCP clients performing loopback redirects, NOT for the upstream OAuth app.
client_storage: Storage implementation for OAuth client registrations.
Defaults to file-based storage if not specified.
token_endpoint_auth_method: Token endpoint authentication method for upstream server.
Common values: "client_secret_basic", "client_secret_post", "none".
If None, authlib will use its default (typically "client_secret_basic").
@ -288,6 +292,7 @@ class OIDCProxy(OAuthProxy):
"base_url": base_url,
"service_documentation_url": self.oidc_config.service_documentation,
"allowed_client_redirect_uris": allowed_client_redirect_uris,
"client_storage": client_storage,
"token_endpoint_auth_method": token_endpoint_auth_method,
}

View file

@ -27,6 +27,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
from fastmcp.server.auth.oidc_proxy import OIDCProxy
from fastmcp.utilities.auth import parse_scopes
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import KVStorage
from fastmcp.utilities.types import NotSet, NotSetT
logger = get_logger(__name__)
@ -91,6 +92,7 @@ class Auth0Provider(OIDCProxy):
required_scopes: list[str] | NotSetT = NotSet,
redirect_path: str | NotSetT = NotSet,
allowed_client_redirect_uris: list[str] | NotSetT = NotSet,
client_storage: KVStorage | None = None,
) -> None:
"""Initialize Auth0 OAuth provider.
@ -104,6 +106,8 @@ class Auth0Provider(OIDCProxy):
redirect_path: Redirect path configured in Auth0 application
allowed_client_redirect_uris: List of allowed redirect URI patterns for MCP clients.
If None (default), all URIs are allowed. If empty list, no URIs are allowed.
client_storage: Storage implementation for OAuth client registrations.
Defaults to file-based storage if not specified.
"""
settings = Auth0ProviderSettings.model_validate(
{
@ -158,6 +162,7 @@ class Auth0Provider(OIDCProxy):
"redirect_path": settings.redirect_path,
"required_scopes": auth0_required_scopes,
"allowed_client_redirect_uris": settings.allowed_client_redirect_uris,
"client_storage": client_storage,
}
super().__init__(**init_kwargs)

View file

@ -14,6 +14,7 @@ from fastmcp.server.auth import AccessToken, TokenVerifier
from fastmcp.server.auth.oauth_proxy import OAuthProxy
from fastmcp.utilities.auth import parse_scopes
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import KVStorage
from fastmcp.utilities.types import NotSet, NotSetT
logger = get_logger(__name__)
@ -160,6 +161,7 @@ class AzureProvider(OAuthProxy):
required_scopes: list[str] | None | NotSetT = NotSet,
timeout_seconds: int | NotSetT = NotSet,
allowed_client_redirect_uris: list[str] | NotSetT = NotSet,
client_storage: KVStorage | None = None,
):
"""Initialize Azure OAuth provider.
@ -173,6 +175,8 @@ class AzureProvider(OAuthProxy):
timeout_seconds: HTTP request timeout for Azure API calls
allowed_client_redirect_uris: List of allowed redirect URI patterns for MCP clients.
If None (default), all URIs are allowed. If empty list, no URIs are allowed.
client_storage: Storage implementation for OAuth client registrations.
Defaults to file-based storage if not specified.
"""
settings = AzureProviderSettings.model_validate(
{
@ -251,6 +255,7 @@ class AzureProvider(OAuthProxy):
redirect_path=settings.redirect_path,
issuer_url=settings.base_url,
allowed_client_redirect_uris=allowed_client_redirect_uris_final,
client_storage=client_storage,
)
logger.info(

View file

@ -30,6 +30,7 @@ from fastmcp.server.auth.auth import AccessToken
from fastmcp.server.auth.oauth_proxy import OAuthProxy
from fastmcp.utilities.auth import parse_scopes
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import KVStorage
from fastmcp.utilities.types import NotSet, NotSetT
logger = get_logger(__name__)
@ -201,6 +202,7 @@ class GitHubProvider(OAuthProxy):
required_scopes: list[str] | NotSetT = NotSet,
timeout_seconds: int | NotSetT = NotSet,
allowed_client_redirect_uris: list[str] | NotSetT = NotSet,
client_storage: KVStorage | None = None,
):
"""Initialize GitHub OAuth provider.
@ -213,6 +215,8 @@ class GitHubProvider(OAuthProxy):
timeout_seconds: HTTP request timeout for GitHub API calls
allowed_client_redirect_uris: List of allowed redirect URI patterns for MCP clients.
If None (default), all URIs are allowed. If empty list, no URIs are allowed.
client_storage: Storage implementation for OAuth client registrations.
Defaults to file-based storage if not specified.
"""
settings = GitHubProviderSettings.model_validate(
@ -269,6 +273,7 @@ class GitHubProvider(OAuthProxy):
redirect_path=settings.redirect_path,
issuer_url=settings.base_url, # We act as the issuer for client registration
allowed_client_redirect_uris=allowed_client_redirect_uris_final,
client_storage=client_storage,
)
logger.info(

View file

@ -32,6 +32,7 @@ from fastmcp.server.auth.auth import AccessToken
from fastmcp.server.auth.oauth_proxy import OAuthProxy
from fastmcp.utilities.auth import parse_scopes
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import KVStorage
from fastmcp.utilities.types import NotSet, NotSetT
logger = get_logger(__name__)
@ -217,6 +218,7 @@ class GoogleProvider(OAuthProxy):
required_scopes: list[str] | NotSetT = NotSet,
timeout_seconds: int | NotSetT = NotSet,
allowed_client_redirect_uris: list[str] | NotSetT = NotSet,
client_storage: KVStorage | None = None,
):
"""Initialize Google OAuth provider.
@ -232,6 +234,8 @@ class GoogleProvider(OAuthProxy):
timeout_seconds: HTTP request timeout for Google API calls
allowed_client_redirect_uris: List of allowed redirect URI patterns for MCP clients.
If None (default), all URIs are allowed. If empty list, no URIs are allowed.
client_storage: Storage implementation for OAuth client registrations.
Defaults to file-based storage if not specified.
"""
settings = GoogleProviderSettings.model_validate(
@ -288,6 +292,7 @@ class GoogleProvider(OAuthProxy):
redirect_path=settings.redirect_path,
issuer_url=settings.base_url, # We act as the issuer for client registration
allowed_client_redirect_uris=allowed_client_redirect_uris_final,
client_storage=client_storage,
)
logger.info(

View file

@ -23,6 +23,7 @@ from fastmcp.server.auth.oauth_proxy import OAuthProxy
from fastmcp.server.auth.providers.jwt import JWTVerifier
from fastmcp.utilities.auth import parse_scopes
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.storage import KVStorage
from fastmcp.utilities.types import NotSet, NotSetT
logger = get_logger(__name__)
@ -169,6 +170,7 @@ class WorkOSProvider(OAuthProxy):
required_scopes: list[str] | None | NotSetT = NotSet,
timeout_seconds: int | NotSetT = NotSet,
allowed_client_redirect_uris: list[str] | NotSetT = NotSet,
client_storage: KVStorage | None = None,
):
"""Initialize WorkOS OAuth provider.
@ -182,6 +184,8 @@ class WorkOSProvider(OAuthProxy):
timeout_seconds: HTTP request timeout for WorkOS API calls
allowed_client_redirect_uris: List of allowed redirect URI patterns for MCP clients.
If None (default), all URIs are allowed. If empty list, no URIs are allowed.
client_storage: Storage implementation for OAuth client registrations.
Defaults to file-based storage if not specified.
"""
settings = WorkOSProviderSettings.model_validate(
@ -247,6 +251,7 @@ class WorkOSProvider(OAuthProxy):
redirect_path=settings.redirect_path,
issuer_url=settings.base_url,
allowed_client_redirect_uris=allowed_client_redirect_uris_final,
client_storage=client_storage,
)
logger.info(

View file

@ -0,0 +1,204 @@
"""Key-value storage utilities for persistent data management."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Protocol
import pydantic_core
from fastmcp.utilities.logging import get_logger
logger = get_logger(__name__)
class KVStorage(Protocol):
"""Protocol for key-value storage of JSON data."""
async def get(self, key: str) -> dict[str, Any] | None:
"""Get a JSON dict by key."""
...
async def set(self, key: str, value: dict[str, Any]) -> None:
"""Store a JSON dict by key."""
...
async def delete(self, key: str) -> None:
"""Delete a value by key."""
...
class JSONFileStorage:
"""File-based key-value storage for JSON data with automatic metadata tracking.
Each key-value pair is stored as a separate JSON file on disk.
Keys are sanitized to be filesystem-safe.
The storage automatically wraps all data with metadata:
- timestamp: Timestamp when the entry was last written
Args:
cache_dir: Directory for storing JSON files
"""
def __init__(self, cache_dir: Path):
"""Initialize JSON file storage."""
self.cache_dir = cache_dir
self.cache_dir.mkdir(exist_ok=True, parents=True)
def _get_safe_key(self, key: str) -> str:
"""Convert key to filesystem-safe string."""
safe_key = key
# Replace problematic characters with underscores
for char in [".", "/", "\\", ":", "*", "?", '"', "<", ">", "|", " "]:
safe_key = safe_key.replace(char, "_")
# Compress multiple underscores into one
while "__" in safe_key:
safe_key = safe_key.replace("__", "_")
# Strip leading and trailing underscores
safe_key = safe_key.strip("_")
return safe_key
def _get_file_path(self, key: str) -> Path:
"""Get the file path for a given key."""
safe_key = self._get_safe_key(key)
return self.cache_dir / f"{safe_key}.json"
async def get(self, key: str) -> dict[str, Any] | None:
"""Get a JSON dict from storage by key.
Args:
key: The key to retrieve
Returns:
The stored dict or None if not found
"""
path = self._get_file_path(key)
try:
wrapper = json.loads(path.read_text())
# Expect wrapped format with metadata
if not isinstance(wrapper, dict) or "data" not in wrapper:
logger.warning(f"Invalid storage format for key '{key}'")
return None
logger.debug(f"Loaded data for key '{key}'")
return wrapper["data"]
except FileNotFoundError:
logger.debug(f"No data found for key '{key}'")
return None
except json.JSONDecodeError as e:
logger.warning(f"Failed to load data for key '{key}': {e}")
return None
async def set(self, key: str, value: dict[str, Any]) -> None:
"""Store a JSON dict with metadata.
Args:
key: The key to store under
value: The dict to store
"""
import time
path = self._get_file_path(key)
current_time = time.time()
# Create wrapper with metadata
wrapper = {
"data": value,
"timestamp": current_time,
}
# Use pydantic_core for consistent JSON serialization
json_data = pydantic_core.to_json(wrapper, fallback=str)
path.write_bytes(json_data)
logger.debug(f"Saved data for key '{key}'")
async def delete(self, key: str) -> None:
"""Delete a value from storage.
Args:
key: The key to delete
"""
path = self._get_file_path(key)
if path.exists():
path.unlink()
logger.debug(f"Deleted data for key '{key}'")
async def cleanup_old_entries(
self,
max_age_seconds: int = 30 * 24 * 60 * 60, # 30 days default
) -> int:
"""Remove entries older than the specified age.
Uses the timestamp field to determine age.
Args:
max_age_seconds: Maximum age in seconds (default 30 days)
Returns:
Number of entries removed
"""
import time
current_time = time.time()
removed_count = 0
for json_file in self.cache_dir.glob("*.json"):
try:
# Read the file and check timestamp
wrapper = json.loads(json_file.read_text())
# Check wrapped format
if not isinstance(wrapper, dict) or "data" not in wrapper:
continue # Invalid format, skip
if "timestamp" not in wrapper:
continue # No timestamp field, skip
entry_age = current_time - wrapper["timestamp"]
if entry_age > max_age_seconds:
json_file.unlink()
removed_count += 1
logger.debug(
f"Removed old entry '{json_file.stem}' (age: {entry_age:.0f}s)"
)
except (json.JSONDecodeError, KeyError) as e:
logger.debug(f"Error reading {json_file.name}: {e}")
continue
if removed_count > 0:
logger.info(f"Cleaned up {removed_count} old entries from storage")
return removed_count
class InMemoryStorage:
"""In-memory key-value storage for JSON data.
Simple dict-based storage that doesn't persist across restarts.
Useful for testing or environments where file storage isn't available.
"""
def __init__(self):
"""Initialize in-memory storage."""
self._data: dict[str, dict[str, Any]] = {}
async def get(self, key: str) -> dict[str, Any] | None:
"""Get a JSON dict from memory by key."""
return self._data.get(key)
async def set(self, key: str, value: dict[str, Any]) -> None:
"""Store a JSON dict in memory."""
self._data[key] = value
async def delete(self, key: str) -> None:
"""Delete a value from memory."""
self._data.pop(key, None)

View file

@ -27,8 +27,13 @@ async def test_token_storage_with_expiry(tmp_path: Path):
await storage.set_tokens(token)
# Check that the file contains the dataclass format
# JSONFileStorage wraps data in {"data": ..., "timestamp": ...}
token_file = storage._get_file_path("tokens")
data = json.loads(token_file.read_text())
wrapper = json.loads(token_file.read_text())
assert "data" in wrapper
assert "timestamp" in wrapper
data = wrapper["data"]
assert "token_payload" in data
assert "expires_at" in data
@ -91,8 +96,10 @@ async def test_token_without_expiry(tmp_path: Path):
await storage.set_tokens(token)
# Check that expires_at is None in the file
# JSONFileStorage wraps data in {"data": ..., "timestamp": ...}
token_file = storage._get_file_path("tokens")
data = json.loads(token_file.read_text())
wrapper = json.loads(token_file.read_text())
data = wrapper["data"]
assert data["expires_at"] is None
# Load the token back - should work since no expiry
@ -133,14 +140,18 @@ async def test_token_expiry_recalculated_on_load(tmp_path: Path):
seconds=1800
) # 30 minutes from now
# JSONFileStorage expects wrapped format
stored_token = {
"token_payload": {
"access_token": "test_token",
"token_type": "Bearer",
"expires_in": 3600, # Original value (will be recalculated)
"refresh_token": "refresh_token",
"data": {
"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(),
},
"expires_at": future_expiry.isoformat(),
"timestamp": datetime.now(timezone.utc).timestamp(),
}
token_file.write_text(json.dumps(stored_token, indent=2, default=str))

View file

@ -395,8 +395,8 @@ class TestOAuthProxyClientRegistration:
await oauth_proxy.register_client(client_info)
# Client should be stored with original credentials
stored = oauth_proxy._clients.get("original-client")
# Client should be retrievable with original credentials
stored = await oauth_proxy.get_client("original-client")
assert stored is not None
assert stored.client_id == "original-client"
assert stored.client_secret == "original-secret"

View file

@ -0,0 +1,213 @@
"""Tests for OAuth proxy with persistent storage."""
from pathlib import Path
from unittest.mock import AsyncMock, Mock
import pytest
from mcp.shared.auth import OAuthClientInformationFull
from pydantic import AnyUrl
from fastmcp.server.auth.oauth_proxy import OAuthProxy
from fastmcp.utilities.storage import InMemoryStorage, JSONFileStorage
class TestOAuthProxyStorage:
"""Tests for OAuth proxy client storage functionality."""
@pytest.fixture
def jwt_verifier(self):
"""Create a mock JWT verifier."""
verifier = Mock()
verifier.required_scopes = ["read", "write"]
verifier.verify_token = AsyncMock(return_value=None)
return verifier
@pytest.fixture
def temp_storage(self, tmp_path: Path) -> JSONFileStorage:
"""Create file-based storage for testing."""
return JSONFileStorage(tmp_path / "oauth-clients")
@pytest.fixture
def memory_storage(self) -> InMemoryStorage:
"""Create in-memory storage for testing."""
return InMemoryStorage()
def create_proxy(self, jwt_verifier, storage=None) -> OAuthProxy:
"""Create an OAuth proxy with specified storage."""
return OAuthProxy(
upstream_authorization_endpoint="https://github.com/login/oauth/authorize",
upstream_token_endpoint="https://github.com/login/oauth/access_token",
upstream_client_id="test-client-id",
upstream_client_secret="test-client-secret",
token_verifier=jwt_verifier,
base_url="https://myserver.com",
redirect_path="/auth/callback",
client_storage=storage,
)
async def test_default_storage_is_file_based(self, jwt_verifier):
"""Test that proxy defaults to file-based storage."""
proxy = self.create_proxy(jwt_verifier, storage=None)
assert isinstance(proxy._client_storage, JSONFileStorage)
async def test_register_and_get_client(self, jwt_verifier, temp_storage):
"""Test registering and retrieving a client."""
proxy = self.create_proxy(jwt_verifier, storage=temp_storage)
# Register client
client_info = OAuthClientInformationFull(
client_id="test-client-123",
client_secret="secret-456",
redirect_uris=[AnyUrl("http://localhost:8080/callback")],
grant_types=["authorization_code", "refresh_token"],
scope="read write",
)
await proxy.register_client(client_info)
# Get client back
client = await proxy.get_client("test-client-123")
assert client is not None
assert client.client_id == "test-client-123"
assert client.client_secret == "secret-456"
assert client.scope == "read write"
async def test_client_persists_across_proxy_instances(
self, jwt_verifier, temp_storage
):
"""Test that clients persist when proxy is recreated."""
# First proxy registers client
proxy1 = self.create_proxy(jwt_verifier, storage=temp_storage)
client_info = OAuthClientInformationFull(
client_id="persistent-client",
client_secret="persistent-secret",
redirect_uris=[AnyUrl("http://localhost:9999/callback")],
scope="openid profile",
)
await proxy1.register_client(client_info)
# Second proxy can retrieve it
proxy2 = self.create_proxy(jwt_verifier, storage=temp_storage)
client = await proxy2.get_client("persistent-client")
assert client is not None
assert client.client_secret == "persistent-secret"
assert client.scope == "openid profile"
async def test_nonexistent_client_returns_none(self, jwt_verifier, temp_storage):
"""Test that requesting non-existent client returns None."""
proxy = self.create_proxy(jwt_verifier, storage=temp_storage)
client = await proxy.get_client("does-not-exist")
assert client is None
async def test_proxy_dcr_client_redirect_validation(
self, jwt_verifier, temp_storage
):
"""Test that ProxyDCRClient is created with redirect URI patterns."""
proxy = OAuthProxy(
upstream_authorization_endpoint="https://github.com/login/oauth/authorize",
upstream_token_endpoint="https://github.com/login/oauth/access_token",
upstream_client_id="test-client-id",
upstream_client_secret="test-client-secret",
token_verifier=jwt_verifier,
base_url="https://myserver.com",
allowed_client_redirect_uris=["http://localhost:*"],
client_storage=temp_storage,
)
client_info = OAuthClientInformationFull(
client_id="test-proxy-client",
client_secret="secret",
redirect_uris=[AnyUrl("http://localhost:8080/callback")],
)
await proxy.register_client(client_info)
# Get client back - should be ProxyDCRClient
client = await proxy.get_client("test-proxy-client")
assert client is not None
# ProxyDCRClient should validate dynamic localhost ports
validated = client.validate_redirect_uri(
AnyUrl("http://localhost:12345/callback")
)
assert validated is not None
async def test_in_memory_storage_option(self, jwt_verifier):
"""Test using in-memory storage explicitly."""
storage = InMemoryStorage()
proxy = self.create_proxy(jwt_verifier, storage=storage)
client_info = OAuthClientInformationFull(
client_id="memory-client",
client_secret="memory-secret",
redirect_uris=[AnyUrl("http://localhost:8080/callback")],
)
await proxy.register_client(client_info)
client = await proxy.get_client("memory-client")
assert client is not None
# Create new proxy with same storage instance
proxy2 = self.create_proxy(jwt_verifier, storage=storage)
client2 = await proxy2.get_client("memory-client")
assert client2 is not None
# But new storage instance won't have it
proxy3 = self.create_proxy(jwt_verifier, storage=InMemoryStorage())
client3 = await proxy3.get_client("memory-client")
assert client3 is None
async def test_storage_data_structure(self, jwt_verifier, temp_storage):
"""Test that storage uses proper structured format."""
proxy = self.create_proxy(jwt_verifier, storage=temp_storage)
client_info = OAuthClientInformationFull(
client_id="structured-client",
client_secret="secret",
redirect_uris=[AnyUrl("http://localhost:8080/callback")],
)
await proxy.register_client(client_info)
# Check raw storage data
raw_data = await temp_storage.get("structured-client")
assert raw_data is not None
assert "client" in raw_data
assert "allowed_redirect_uri_patterns" in raw_data
async def test_cleanup_old_clients(self, jwt_verifier, temp_storage):
"""Test cleanup of old clients using storage's cleanup method."""
import json
import time
proxy = self.create_proxy(jwt_verifier, storage=temp_storage)
# Register some clients
client1 = OAuthClientInformationFull(
client_id="old-client",
client_secret="secret1",
redirect_uris=[AnyUrl("http://localhost:8080/callback")],
)
await proxy.register_client(client1)
client2 = OAuthClientInformationFull(
client_id="recent-client",
client_secret="secret2",
redirect_uris=[AnyUrl("http://localhost:9090/callback")],
)
await proxy.register_client(client2)
# Manually make the first client old by modifying the file directly
old_client_path = temp_storage._get_file_path("old-client")
wrapper = json.loads(old_client_path.read_text())
wrapper["timestamp"] = time.time() - (35 * 24 * 60 * 60) # 35 days old
old_client_path.write_text(json.dumps(wrapper))
# Run cleanup directly on storage
removed_count = await temp_storage.cleanup_old_entries(
max_age_seconds=30 * 24 * 60 * 60
)
assert removed_count == 1
# Old client should be gone
assert await proxy.get_client("old-client") is None
# Recent client should still exist
assert await proxy.get_client("recent-client") is not None

View file

@ -0,0 +1,143 @@
"""Tests for KVStorage implementations."""
from pathlib import Path
import pytest
from fastmcp.utilities.storage import InMemoryStorage, JSONFileStorage
class TestJSONFileStorage:
"""Tests for file-based JSON storage."""
@pytest.fixture
def temp_storage(self, tmp_path: Path) -> JSONFileStorage:
"""Create a JSONFileStorage with temp directory."""
return JSONFileStorage(tmp_path / "storage")
async def test_basic_get_set_delete(self, temp_storage: JSONFileStorage):
"""Test basic storage operations."""
# Initially empty
assert await temp_storage.get("key1") is None
# Set a value
data = {"name": "test", "value": 123}
await temp_storage.set("key1", data)
# Get it back
loaded = await temp_storage.get("key1")
assert loaded == data
# Delete it
await temp_storage.delete("key1")
assert await temp_storage.get("key1") is None
async def test_special_characters_in_keys(self, temp_storage: JSONFileStorage):
"""Test that special characters in keys are handled safely."""
key = "user/123:test.json?query=value"
data = {"test": "data"}
await temp_storage.set(key, data)
loaded = await temp_storage.get(key)
assert loaded == data
# Verify the file was created with safe name
files = list(temp_storage.cache_dir.glob("*.json"))
assert len(files) == 1
assert "/" not in files[0].name
assert ":" not in files[0].name
assert "?" not in files[0].name
async def test_multiple_keys(self, temp_storage: JSONFileStorage):
"""Test storing multiple keys."""
data1 = {"id": 1}
data2 = {"id": 2}
data3 = {"id": 3}
await temp_storage.set("key1", data1)
await temp_storage.set("key2", data2)
await temp_storage.set("key3", data3)
assert await temp_storage.get("key1") == data1
assert await temp_storage.get("key2") == data2
assert await temp_storage.get("key3") == data3
# Delete one
await temp_storage.delete("key2")
assert await temp_storage.get("key1") == data1
assert await temp_storage.get("key2") is None
assert await temp_storage.get("key3") == data3
async def test_overwrite_existing(self, temp_storage: JSONFileStorage):
"""Test overwriting existing values."""
await temp_storage.set("key", {"version": 1})
await temp_storage.set("key", {"version": 2})
loaded = await temp_storage.get("key")
assert loaded == {"version": 2}
async def test_persistence_across_instances(self, tmp_path: Path):
"""Test that data persists across storage instances."""
storage_dir = tmp_path / "persistent"
# First instance
storage1 = JSONFileStorage(storage_dir)
data = {"persistent": True, "value": 42}
await storage1.set("mykey", data)
# New instance, same directory
storage2 = JSONFileStorage(storage_dir)
loaded = await storage2.get("mykey")
assert loaded == data
async def test_delete_nonexistent(self, temp_storage: JSONFileStorage):
"""Test deleting non-existent key doesn't error."""
# Should not raise
await temp_storage.delete("nonexistent")
class TestInMemoryStorage:
"""Tests for in-memory storage."""
@pytest.fixture
def memory_storage(self) -> InMemoryStorage:
"""Create an InMemoryStorage instance."""
return InMemoryStorage()
async def test_basic_operations(self, memory_storage: InMemoryStorage):
"""Test basic storage operations."""
# Initially empty
assert await memory_storage.get("key1") is None
# Set and get
data = {"name": "test", "value": 123}
await memory_storage.set("key1", data)
assert await memory_storage.get("key1") == data
# Delete
await memory_storage.delete("key1")
assert await memory_storage.get("key1") is None
async def test_no_persistence(self):
"""Test that data doesn't persist across instances."""
storage1 = InMemoryStorage()
await storage1.set("key", {"value": 1})
storage2 = InMemoryStorage()
assert await storage2.get("key") is None
async def test_isolation_between_keys(self, memory_storage: InMemoryStorage):
"""Test that keys are isolated from each other."""
data1 = {"id": 1, "nested": {"value": "a"}}
data2 = {"id": 2, "nested": {"value": "b"}}
await memory_storage.set("key1", data1)
await memory_storage.set("key2", data2)
# Modify retrieved data shouldn't affect stored
retrieved = await memory_storage.get("key1")
if retrieved:
retrieved["modified"] = True
# Original should be unchanged
assert await memory_storage.get("key1") == data1