mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 14:04:18 +02:00
Add persistent storage for OAuth client registrations (#1879)
This commit is contained in:
parent
92b7e92e05
commit
7176d4f293
14 changed files with 704 additions and 52 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)}")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
204
src/fastmcp/utilities/storage.py
Normal file
204
src/fastmcp/utilities/storage.py
Normal 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)
|
||||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
213
tests/server/auth/test_oauth_proxy_storage.py
Normal file
213
tests/server/auth/test_oauth_proxy_storage.py
Normal 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
|
||||
143
tests/utilities/test_storage.py
Normal file
143
tests/utilities/test_storage.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue