From 7176d4f293c537a2cecf961ad4b2811a98d052c6 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sun, 21 Sep 2025 10:56:08 -0400
Subject: [PATCH] Add persistent storage for OAuth client registrations (#1879)
---
docs/servers/auth/oauth-proxy.mdx | 15 +-
src/fastmcp/client/auth/oauth.py | 81 ++++---
src/fastmcp/server/auth/oauth_proxy.py | 39 +++-
src/fastmcp/server/auth/oidc_proxy.py | 5 +
src/fastmcp/server/auth/providers/auth0.py | 5 +
src/fastmcp/server/auth/providers/azure.py | 5 +
src/fastmcp/server/auth/providers/github.py | 5 +
src/fastmcp/server/auth/providers/google.py | 5 +
src/fastmcp/server/auth/providers/workos.py | 5 +
src/fastmcp/utilities/storage.py | 204 +++++++++++++++++
tests/client/auth/test_oauth_token_expiry.py | 27 ++-
tests/server/auth/test_oauth_proxy.py | 4 +-
tests/server/auth/test_oauth_proxy_storage.py | 213 ++++++++++++++++++
tests/utilities/test_storage.py | 143 ++++++++++++
14 files changed, 704 insertions(+), 52 deletions(-)
create mode 100644 src/fastmcp/utilities/storage.py
create mode 100644 tests/server/auth/test_oauth_proxy_storage.py
create mode 100644 tests/utilities/test_storage.py
diff --git a/docs/servers/auth/oauth-proxy.mdx b/docs/servers/auth/oauth-proxy.mdx
index 32b894a58..d532a8108 100644
--- a/docs/servers/auth/oauth-proxy.mdx
+++ b/docs/servers/auth/oauth-proxy.mdx
@@ -161,14 +161,25 @@ mcp = FastMCP(name="My Server", auth=auth)
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.
+
+
+ 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())
+ ```
+
### Provider-Specific Parameters
diff --git a/src/fastmcp/client/auth/oauth.py b/src/fastmcp/client/auth/oauth.py
index e47ab7778..86c83b7ab 100644
--- a/src/fastmcp/client/auth/oauth.py
+++ b/src/fastmcp/client/auth/oauth.py
@@ -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)}")
diff --git a/src/fastmcp/server/auth/oauth_proxy.py b/src/fastmcp/server/auth/oauth_proxy.py
index f7e7e92da..a4e1d9d34 100644
--- a/src/fastmcp/server/auth/oauth_proxy.py
+++ b/src/fastmcp/server/auth/oauth_proxy.py
@@ -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:
diff --git a/src/fastmcp/server/auth/oidc_proxy.py b/src/fastmcp/server/auth/oidc_proxy.py
index 1043f32b2..e7f24dcc9 100644
--- a/src/fastmcp/server/auth/oidc_proxy.py
+++ b/src/fastmcp/server/auth/oidc_proxy.py
@@ -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,
}
diff --git a/src/fastmcp/server/auth/providers/auth0.py b/src/fastmcp/server/auth/providers/auth0.py
index 61b2ed6a7..512a70068 100644
--- a/src/fastmcp/server/auth/providers/auth0.py
+++ b/src/fastmcp/server/auth/providers/auth0.py
@@ -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)
diff --git a/src/fastmcp/server/auth/providers/azure.py b/src/fastmcp/server/auth/providers/azure.py
index 84a3d4791..cd6a2ded9 100644
--- a/src/fastmcp/server/auth/providers/azure.py
+++ b/src/fastmcp/server/auth/providers/azure.py
@@ -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(
diff --git a/src/fastmcp/server/auth/providers/github.py b/src/fastmcp/server/auth/providers/github.py
index 2a9490ff9..e08e2f97c 100644
--- a/src/fastmcp/server/auth/providers/github.py
+++ b/src/fastmcp/server/auth/providers/github.py
@@ -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(
diff --git a/src/fastmcp/server/auth/providers/google.py b/src/fastmcp/server/auth/providers/google.py
index a8b0312c3..c6fbeb6cd 100644
--- a/src/fastmcp/server/auth/providers/google.py
+++ b/src/fastmcp/server/auth/providers/google.py
@@ -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(
diff --git a/src/fastmcp/server/auth/providers/workos.py b/src/fastmcp/server/auth/providers/workos.py
index 363f96f7d..d68e16c71 100644
--- a/src/fastmcp/server/auth/providers/workos.py
+++ b/src/fastmcp/server/auth/providers/workos.py
@@ -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(
diff --git a/src/fastmcp/utilities/storage.py b/src/fastmcp/utilities/storage.py
new file mode 100644
index 000000000..2eb21cead
--- /dev/null
+++ b/src/fastmcp/utilities/storage.py
@@ -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)
diff --git a/tests/client/auth/test_oauth_token_expiry.py b/tests/client/auth/test_oauth_token_expiry.py
index 77e5d552e..1ec17d38c 100644
--- a/tests/client/auth/test_oauth_token_expiry.py
+++ b/tests/client/auth/test_oauth_token_expiry.py
@@ -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))
diff --git a/tests/server/auth/test_oauth_proxy.py b/tests/server/auth/test_oauth_proxy.py
index 4fffb1cf8..f7ed5d900 100644
--- a/tests/server/auth/test_oauth_proxy.py
+++ b/tests/server/auth/test_oauth_proxy.py
@@ -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"
diff --git a/tests/server/auth/test_oauth_proxy_storage.py b/tests/server/auth/test_oauth_proxy_storage.py
new file mode 100644
index 000000000..6ceef86aa
--- /dev/null
+++ b/tests/server/auth/test_oauth_proxy_storage.py
@@ -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
diff --git a/tests/utilities/test_storage.py b/tests/utilities/test_storage.py
new file mode 100644
index 000000000..1c53b6637
--- /dev/null
+++ b/tests/utilities/test_storage.py
@@ -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