diff --git a/pyproject.toml b/pyproject.toml index c105a51e1..00e4b7d09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "pydantic[email]>=2.11.7", "pyperclip>=1.9.0", "openapi-core>=0.19.5", + "py-key-value-aio[disk,memory]>=0.2.1", "websockets>=15.0.1", ] diff --git a/src/fastmcp/client/auth/oauth.py b/src/fastmcp/client/auth/oauth.py index 86c83b7ab..0e088953d 100644 --- a/src/fastmcp/client/auth/oauth.py +++ b/src/fastmcp/client/auth/oauth.py @@ -1,34 +1,33 @@ from __future__ import annotations import asyncio +import time import webbrowser from asyncio import Future from collections.abc import AsyncGenerator -from datetime import datetime, timedelta, timezone -from pathlib import Path -from typing import Any, Literal +from typing import Any from urllib.parse import urlparse import anyio import httpx +from key_value.aio.adapters.pydantic import PydanticAdapter +from key_value.aio.protocols import AsyncKeyValue +from key_value.aio.stores.memory import MemoryStore from mcp.client.auth import OAuthClientProvider, TokenStorage from mcp.shared.auth import ( OAuthClientInformationFull, OAuthClientMetadata, + OAuthToken, ) -from mcp.shared.auth import ( - OAuthToken as OAuthToken, -) -from pydantic import AnyHttpUrl, BaseModel, TypeAdapter, ValidationError +from pydantic import AnyHttpUrl +from typing_extensions import override from uvicorn.server import Server -from fastmcp import settings as fastmcp_global_settings from fastmcp.client.oauth_callback import ( create_oauth_callback_server, ) from fastmcp.utilities.http import find_available_port from fastmcp.utilities.logging import get_logger -from fastmcp.utilities.storage import JSONFileStorage __all__ = ["OAuth"] @@ -41,174 +40,6 @@ class ClientNotFoundError(Exception): pass -class StoredToken(BaseModel): - """Token storage format with absolute expiry time.""" - - token_payload: OAuthToken - expires_at: datetime | None - - -# Create TypeAdapter at module level for efficient parsing -stored_token_adapter = TypeAdapter(StoredToken) - - -def default_cache_dir() -> Path: - return fastmcp_global_settings.home / "oauth-mcp-client-cache" - - -class FileTokenStorage(TokenStorage): - """ - File-based token storage implementation for OAuth credentials and tokens. - 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 - # Use JSONFileStorage for actual file operations - self._storage = JSONFileStorage(cache_dir or default_cache_dir()) - - @staticmethod - def get_base_url(url: str) -> str: - """Extract the base URL (scheme + host) from a URL.""" - parsed = urlparse(url) - return f"{parsed.scheme}://{parsed.netloc}" - - 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 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. - - 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.""" - key = self._get_storage_key("tokens") - data = await self._storage.get(key) - - if data is None: - return None - - try: - # Parse and validate as StoredToken - stored = stored_token_adapter.validate_python(data) - - # Check if token is expired - if stored.expires_at is not None: - now = datetime.now(timezone.utc) - if now >= stored.expires_at: - logger.debug( - f"Token expired for {self.get_base_url(self.server_url)}" - ) - return None - - # Recalculate expires_in to be correct relative to now - if stored.token_payload.expires_in is not None: - remaining = stored.expires_at - now - stored.token_payload.expires_in = max( - 0, int(remaining.total_seconds()) - ) - - return stored.token_payload - - except ValidationError as e: - logger.debug( - 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.""" - key = self._get_storage_key("tokens") - - # Calculate absolute expiry time if expires_in is present - expires_at = None - if tokens.expires_in is not None: - expires_at = datetime.now(timezone.utc) + timedelta( - seconds=tokens.expires_in - ) - - # Create StoredToken and save using storage - # Note: JSONFileStorage will wrap this in {"data": ..., "timestamp": ...} - stored = StoredToken(token_payload=tokens, expires_at=expires_at) - 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.""" - 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(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 - tokens = await self.get_tokens() - if tokens is None: - logger.debug( - f"No tokens found for client info at {self.get_base_url(self.server_url)}. " - "OAuth flow may have been incomplete. Clearing client info to force fresh registration." - ) - # Clear the incomplete client info - await self._storage.delete(key) - return None - - return client_info - except ValidationError as e: - logger.debug( - 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.""" - 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. - - 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)}") - - @classmethod - def clear_all(cls, cache_dir: Path | None = None) -> None: - """Clear all cached data for all servers.""" - cache_dir = cache_dir or default_cache_dir() - if not cache_dir.exists(): - return - - file_types: list[Literal["client_info", "tokens"]] = ["client_info", "tokens"] - for file_type in file_types: - for file in cache_dir.glob(f"*_{file_type}.json"): - file.unlink(missing_ok=True) - logger.info("Cleared all OAuth client cache data.") - - async def check_if_auth_required( mcp_url: str, httpx_kwargs: dict[str, Any] | None = None ) -> bool: @@ -239,6 +70,70 @@ async def check_if_auth_required( return True +class TokenStorageAdapter(TokenStorage): + _server_url: str + _key_value_store: AsyncKeyValue + _storage_oauth_token: PydanticAdapter[OAuthToken] + _storage_client_info: PydanticAdapter[OAuthClientInformationFull] + + def __init__(self, async_key_value: AsyncKeyValue, server_url: str): + self._server_url = server_url + self._key_value_store = async_key_value + self._storage_oauth_token = PydanticAdapter[OAuthToken]( + default_collection="mcp-oauth-token", + key_value=async_key_value, + pydantic_model=OAuthToken, + raise_on_validation_error=True, + ) + self._storage_client_info = PydanticAdapter[OAuthClientInformationFull]( + default_collection="mcp-oauth-client-info", + key_value=async_key_value, + pydantic_model=OAuthClientInformationFull, + raise_on_validation_error=True, + ) + + def _get_token_cache_key(self) -> str: + return f"{self._server_url}/tokens" + + def _get_client_info_cache_key(self) -> str: + return f"{self._server_url}/client_info" + + async def clear(self) -> None: + await self._storage_oauth_token.delete(key=self._get_token_cache_key()) + await self._storage_client_info.delete(key=self._get_client_info_cache_key()) + + @override + async def get_tokens(self) -> OAuthToken | None: + return await self._storage_oauth_token.get(key=self._get_token_cache_key()) + + @override + async def set_tokens(self, tokens: OAuthToken) -> None: + await self._storage_oauth_token.put( + key=self._get_token_cache_key(), + value=tokens, + ttl=tokens.expires_in, + ) + + @override + async def get_client_info(self) -> OAuthClientInformationFull | None: + return await self._storage_client_info.get( + key=self._get_client_info_cache_key() + ) + + @override + async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + ttl: int | None = None + + if client_info.client_secret_expires_at: + ttl = client_info.client_secret_expires_at - int(time.time()) + + await self._storage_client_info.put( + key=self._get_client_info_cache_key(), + value=client_info, + ttl=ttl, + ) + + class OAuth(OAuthClientProvider): """ OAuth client provider for MCP servers with browser-based authentication. @@ -252,7 +147,7 @@ def __init__( mcp_url: str, scopes: str | list[str] | None = None, client_name: str = "FastMCP Client", - token_storage_cache_dir: Path | None = None, + token_storage: AsyncKeyValue | None = None, additional_client_metadata: dict[str, Any] | None = None, callback_port: int | None = None, ): @@ -264,7 +159,7 @@ def __init__( scopes: OAuth scopes to request. Can be a space-separated string or a list of strings. client_name: Name for this client during registration - token_storage_cache_dir: Directory for FileTokenStorage + token_storage: An AsyncKeyValue-compatible token store, tokens are stored in memory if not provided additional_client_metadata: Extra fields for OAuthClientMetadata callback_port: Fixed port for OAuth callback (default: random available port) """ @@ -294,8 +189,10 @@ def __init__( ) # Create server-specific token storage - storage = FileTokenStorage( - server_url=server_base_url, cache_dir=token_storage_cache_dir + token_storage = token_storage or MemoryStore() + + self.token_storage_adapter: TokenStorageAdapter = TokenStorageAdapter( + async_key_value=token_storage, server_url=server_base_url ) # Store server_base_url for use in callback_handler @@ -305,7 +202,7 @@ def __init__( super().__init__( server_url=server_base_url, client_metadata=client_metadata, - storage=storage, + storage=self.token_storage_adapter, redirect_handler=self.redirect_handler, callback_handler=self.callback_handler, ) @@ -399,23 +296,7 @@ async def async_auth_flow( # Clear cached state and retry once self._initialized = False - - # Try to clear storage if it supports it - if hasattr(self.context.storage, "clear"): - try: - self.context.storage.clear() - except Exception as e: - logger.warning(f"Failed to clear OAuth storage cache: {e}") - # Can't retry without clearing cache, re-raise original error - raise ClientNotFoundError( - "OAuth client not found and cache could not be cleared" - ) from e - else: - logger.warning( - "Storage does not support clear() - cannot retry with fresh credentials" - ) - # Can't retry without clearing cache, re-raise original error - raise + await self.token_storage_adapter.clear() gen = super().async_auth_flow(request) response = None diff --git a/src/fastmcp/server/auth/oauth_proxy.py b/src/fastmcp/server/auth/oauth_proxy.py index 8400aa68a..89b9c63a5 100644 --- a/src/fastmcp/server/auth/oauth_proxy.py +++ b/src/fastmcp/server/auth/oauth_proxy.py @@ -28,6 +28,9 @@ import httpx from authlib.common.security import generate_token from authlib.integrations.httpx_client import AsyncOAuth2Client +from key_value.aio.adapters.pydantic import PydanticAdapter +from key_value.aio.protocols import AsyncKeyValue +from key_value.aio.stores.memory import MemoryStore from mcp.server.auth.handlers.token import TokenErrorResponse, TokenSuccessResponse from mcp.server.auth.handlers.token import TokenHandler as _SDKTokenHandler from mcp.server.auth.json_response import PydanticJSONResponse @@ -45,16 +48,14 @@ RevocationOptions, ) from mcp.shared.auth import OAuthClientInformationFull, OAuthToken -from pydantic import AnyHttpUrl, AnyUrl, SecretStr +from pydantic import AnyHttpUrl, AnyUrl, Field, SecretStr 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 @@ -88,21 +89,7 @@ class ProxyDCRClient(OAuthClientInformationFull): arise from accepting arbitrary redirect URIs. """ - def __init__( - self, - *args: Any, - allowed_redirect_uri_patterns: list[str] | None = None, - **kwargs: Any, - ): - """Initialize with allowed redirect URI patterns. - - Args: - allowed_redirect_uri_patterns: List of allowed redirect URI patterns with wildcard support. - If None, defaults to localhost-only patterns. - If empty list, allows all redirect URIs. - """ - super().__init__(*args, **kwargs) - self._allowed_redirect_uri_patterns = allowed_redirect_uri_patterns + allowed_redirect_uri_patterns: list[str] | None = Field(default=None) def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl: """Validate redirect URI against allowed patterns. @@ -114,7 +101,10 @@ def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl: """ if redirect_uri is not None: # Validate against allowed patterns - if validate_redirect_uri(redirect_uri, self._allowed_redirect_uri_patterns): + if validate_redirect_uri( + redirect_uri=redirect_uri, + allowed_patterns=self.allowed_redirect_uri_patterns, + ): return redirect_uri # Fall back to normal validation if not in allowed patterns return super().validate_redirect_uri(redirect_uri) @@ -258,7 +248,6 @@ class OAuthProxy(OAuthProvider): State Management --------------- The proxy maintains minimal but crucial state: - - _clients: DCR registrations (all use ProxyDCRClient for flexibility) - _oauth_transactions: Active authorization flows with client context - _client_codes: Authorization codes with PKCE challenges and upstream tokens - _access_tokens, _refresh_tokens: Token storage for revocation @@ -314,7 +303,7 @@ def __init__( # Extra parameters to forward to token endpoint extra_token_params: dict[str, str] | None = None, # Client storage - client_storage: KVStorage | None = None, + client_storage: AsyncKeyValue | None = None, ): """Initialize the OAuth proxy provider. @@ -348,9 +337,7 @@ def __init__( 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. + client_storage: An AsyncKeyValue-compatible store for client registrations, registrations are stored in memory if not provided """ # Always enable DCR since we implement it locally for MCP clients client_registration_options = ClientRegistrationOptions( @@ -399,11 +386,14 @@ def __init__( self._extra_authorize_params = extra_authorize_params or {} self._extra_token_params = extra_token_params or {} - # 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 + self._client_storage: AsyncKeyValue = client_storage or MemoryStore() + + self._client_store = PydanticAdapter[ProxyDCRClient]( + key_value=self._client_storage, + pydantic_model=ProxyDCRClient, + default_collection="mcp-oauth-proxy-clients", + raise_on_validation_error=True, + ) # Local state for token bookkeeping only (no client caching) self._access_tokens: dict[str, AccessToken] = {} @@ -457,19 +447,13 @@ async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: For unregistered clients, returns None (which will raise an error in the SDK). """ # Load from storage - data = await self._client_storage.get(client_id) - if not data: + if not (client := await self._client_store.get(key=client_id)): return None - 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, - ) + if client.allowed_redirect_uri_patterns is None: + client.allowed_redirect_uri_patterns = self._allowed_client_redirect_uris - return None + return client async def register_client(self, client_info: OAuthClientInformationFull) -> None: """Register a client locally @@ -481,7 +465,7 @@ async def register_client(self, client_info: OAuthClientInformationFull) -> None """ # Create a ProxyDCRClient with configured redirect URI validation - proxy_client = ProxyDCRClient( + proxy_client: ProxyDCRClient = ProxyDCRClient( client_id=client_info.client_id, client_secret=client_info.client_secret, redirect_uris=client_info.redirect_uris or [AnyUrl("http://localhost")], @@ -492,12 +476,10 @@ async def register_client(self, client_info: OAuthClientInformationFull) -> None allowed_redirect_uri_patterns=self._allowed_client_redirect_uris, ) - # 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) + await self._client_store.put( + key=client_info.client_id, + value=proxy_client, + ) # 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 e7f24dcc9..589e0e2d3 100644 --- a/src/fastmcp/server/auth/oidc_proxy.py +++ b/src/fastmcp/server/auth/oidc_proxy.py @@ -12,6 +12,7 @@ from collections.abc import Sequence import httpx +from key_value.aio.protocols import AsyncKeyValue from pydantic import AnyHttpUrl, BaseModel, model_validator from typing_extensions import Self @@ -19,7 +20,6 @@ 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__) @@ -213,7 +213,7 @@ def __init__( redirect_path: str | None = None, # Client configuration allowed_client_redirect_uris: list[str] | None = None, - client_storage: KVStorage | None = None, + client_storage: AsyncKeyValue | None = None, # Token validation configuration token_endpoint_auth_method: str | None = None, ) -> None: @@ -236,8 +236,7 @@ def __init__( 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. + client_storage: An AsyncKeyValue-compatible store for client registrations, registrations are stored in memory if not provided 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"). diff --git a/src/fastmcp/server/auth/providers/auth0.py b/src/fastmcp/server/auth/providers/auth0.py index 512a70068..67ddf8b49 100644 --- a/src/fastmcp/server/auth/providers/auth0.py +++ b/src/fastmcp/server/auth/providers/auth0.py @@ -21,13 +21,13 @@ ``` """ +from key_value.aio.protocols import AsyncKeyValue from pydantic import AnyHttpUrl, SecretStr, field_validator 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__) @@ -92,7 +92,7 @@ def __init__( required_scopes: list[str] | NotSetT = NotSet, redirect_path: str | NotSetT = NotSet, allowed_client_redirect_uris: list[str] | NotSetT = NotSet, - client_storage: KVStorage | None = None, + client_storage: AsyncKeyValue | None = None, ) -> None: """Initialize Auth0 OAuth provider. @@ -106,8 +106,7 @@ def __init__( 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. + client_storage: An AsyncKeyValue-compatible store for client registrations, registrations are stored in memory if not provided """ settings = Auth0ProviderSettings.model_validate( { diff --git a/src/fastmcp/server/auth/providers/azure.py b/src/fastmcp/server/auth/providers/azure.py index 2f1cee8fa..a621a9ea3 100644 --- a/src/fastmcp/server/auth/providers/azure.py +++ b/src/fastmcp/server/auth/providers/azure.py @@ -8,9 +8,7 @@ from typing import TYPE_CHECKING -if TYPE_CHECKING: # pragma: no cover - from mcp.server.auth.provider import AuthorizationParams - from mcp.shared.auth import OAuthClientInformationFull +from key_value.aio.protocols import AsyncKeyValue from pydantic import SecretStr, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict @@ -18,9 +16,12 @@ 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 +if TYPE_CHECKING: + from mcp.server.auth.provider import AuthorizationParams + from mcp.shared.auth import OAuthClientInformationFull + logger = get_logger(__name__) @@ -104,7 +105,7 @@ def __init__( required_scopes: list[str] | None | NotSetT = NotSet, additional_authorize_scopes: list[str] | None | NotSetT = NotSet, allowed_client_redirect_uris: list[str] | NotSetT = NotSet, - client_storage: KVStorage | None = None, + client_storage: AsyncKeyValue | None = None, ) -> None: """Initialize Azure OAuth provider. @@ -124,8 +125,7 @@ def __init__( permissions. These are not used for token validation. 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. + client_storage: An AsyncKeyValue-compatible store for client registrations, registrations are stored in memory if not provided """ settings = AzureProviderSettings.model_validate( { diff --git a/src/fastmcp/server/auth/providers/github.py b/src/fastmcp/server/auth/providers/github.py index e08e2f97c..89692fbf1 100644 --- a/src/fastmcp/server/auth/providers/github.py +++ b/src/fastmcp/server/auth/providers/github.py @@ -22,6 +22,7 @@ from __future__ import annotations import httpx +from key_value.aio.protocols import AsyncKeyValue from pydantic import AnyHttpUrl, SecretStr, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict @@ -30,7 +31,6 @@ 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__) @@ -202,7 +202,7 @@ def __init__( required_scopes: list[str] | NotSetT = NotSet, timeout_seconds: int | NotSetT = NotSet, allowed_client_redirect_uris: list[str] | NotSetT = NotSet, - client_storage: KVStorage | None = None, + client_storage: AsyncKeyValue | None = None, ): """Initialize GitHub OAuth provider. @@ -215,8 +215,7 @@ def __init__( 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. + client_storage: An AsyncKeyValue-compatible store for client registrations, registrations are stored in memory if not provided """ settings = GitHubProviderSettings.model_validate( diff --git a/src/fastmcp/server/auth/providers/google.py b/src/fastmcp/server/auth/providers/google.py index c6fbeb6cd..721613d92 100644 --- a/src/fastmcp/server/auth/providers/google.py +++ b/src/fastmcp/server/auth/providers/google.py @@ -24,6 +24,7 @@ import time import httpx +from key_value.aio.protocols import AsyncKeyValue from pydantic import AnyHttpUrl, SecretStr, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict @@ -32,7 +33,6 @@ 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__) @@ -218,7 +218,7 @@ def __init__( required_scopes: list[str] | NotSetT = NotSet, timeout_seconds: int | NotSetT = NotSet, allowed_client_redirect_uris: list[str] | NotSetT = NotSet, - client_storage: KVStorage | None = None, + client_storage: AsyncKeyValue | None = None, ): """Initialize Google OAuth provider. @@ -234,8 +234,7 @@ def __init__( 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. + client_storage: An AsyncKeyValue-compatible store for client registrations, registrations are stored in memory if not provided """ settings = GoogleProviderSettings.model_validate( diff --git a/src/fastmcp/server/auth/providers/workos.py b/src/fastmcp/server/auth/providers/workos.py index b0c5f1817..e1b0631b7 100644 --- a/src/fastmcp/server/auth/providers/workos.py +++ b/src/fastmcp/server/auth/providers/workos.py @@ -11,6 +11,7 @@ from __future__ import annotations import httpx +from key_value.aio.protocols import AsyncKeyValue from pydantic import AnyHttpUrl, SecretStr, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict from starlette.responses import JSONResponse @@ -21,7 +22,6 @@ 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__) @@ -168,7 +168,7 @@ def __init__( 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, + client_storage: AsyncKeyValue | None = None, ): """Initialize WorkOS OAuth provider. @@ -182,8 +182,7 @@ def __init__( 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. + client_storage: An AsyncKeyValue-compatible store for client registrations, registrations are stored in memory if not provided """ settings = WorkOSProviderSettings.model_validate( diff --git a/src/fastmcp/settings.py b/src/fastmcp/settings.py index 95e9f6290..8844ea6ed 100644 --- a/src/fastmcp/settings.py +++ b/src/fastmcp/settings.py @@ -23,6 +23,8 @@ DuplicateBehavior = Literal["warn", "error", "replace", "ignore"] +TEN_MB_IN_BYTES = 1024 * 1024 * 10 + if TYPE_CHECKING: from fastmcp.server.auth.auth import AuthProvider diff --git a/src/fastmcp/utilities/storage.py b/src/fastmcp/utilities/storage.py deleted file mode 100644 index 2eb21cead..000000000 --- a/src/fastmcp/utilities/storage.py +++ /dev/null @@ -1,204 +0,0 @@ -"""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 deleted file mode 100644 index 1ec17d38c..000000000 --- a/tests/client/auth/test_oauth_token_expiry.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Test OAuth token expiry handling with absolute timestamps.""" - -import json -from datetime import datetime, timedelta, timezone -from pathlib import Path - -import pytest -from mcp.shared.auth import OAuthToken - -from fastmcp.client.auth.oauth import FileTokenStorage - - -@pytest.mark.asyncio -async def test_token_storage_with_expiry(tmp_path: Path): - """Test that tokens are stored with absolute expiry time and loaded correctly.""" - storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) - - # Create a token with 3600 seconds expiry - token = OAuthToken( - access_token="test_token", - token_type="Bearer", - expires_in=3600, - refresh_token="refresh_token", - ) - - # Save the token - await storage.set_tokens(token) - - # Check that the file contains the dataclass format - # JSONFileStorage wraps data in {"data": ..., "timestamp": ...} - token_file = storage._get_file_path("tokens") - 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 - assert data["expires_at"] is not None - # expires_at should be approximately now + 3600 seconds - expires_at = datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00")) - expected = datetime.now(timezone.utc) + timedelta(seconds=3600) - assert abs((expires_at - expected).total_seconds()) < 2 - - # Load the token back - loaded_token = await storage.get_tokens() - assert loaded_token is not None - assert loaded_token.access_token == "test_token" - # expires_in should be recalculated to be approximately 3600 (minus loading time) - assert loaded_token.expires_in is not None - assert 3595 <= loaded_token.expires_in <= 3600 - - -@pytest.mark.asyncio -async def test_expired_token_returns_none(tmp_path: Path): - """Test that expired tokens return None when loaded.""" - storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) - - # Manually create an already-expired token file - token_file = storage._get_file_path("tokens") - past_expiry = datetime.now(timezone.utc) - timedelta( - seconds=10 - ) # Expired 10 seconds ago - - expired_token = { - "token_payload": { - "access_token": "test_token", - "token_type": "Bearer", - "expires_in": 3600, - "refresh_token": "refresh_token", - }, - "expires_at": past_expiry.isoformat(), - } - token_file.write_text(json.dumps(expired_token, indent=2, default=str)) - - # Load the token - should return None since it's expired - loaded_token = await storage.get_tokens() - assert loaded_token is None - - -@pytest.mark.asyncio -async def test_token_without_expiry(tmp_path: Path): - """Test that tokens without expires_in are handled correctly.""" - storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) - - # Create a token without expires_in (perpetual token) - token = OAuthToken( - access_token="test_token", - token_type="Bearer", - expires_in=None, - refresh_token="refresh_token", - ) - - # Save the token - await storage.set_tokens(token) - - # Check that expires_at is None in the file - # JSONFileStorage wraps data in {"data": ..., "timestamp": ...} - token_file = storage._get_file_path("tokens") - 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 - loaded_token = await storage.get_tokens() - assert loaded_token is not None - assert loaded_token.access_token == "test_token" - assert loaded_token.expires_in is None - - -@pytest.mark.asyncio -async def test_invalid_format_returns_none(tmp_path: Path): - """Test that invalid token format returns None.""" - storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) - - # Manually write an invalid format token file (missing required fields) - token_file = storage._get_file_path("tokens") - invalid_token = { - "access_token": "invalid_token", - "token_type": "Bearer", - "expires_in": 3600, - "refresh_token": "refresh_token", - } - token_file.write_text(json.dumps(invalid_token, indent=2)) - - # Try to load - should return None - loaded_token = await storage.get_tokens() - assert loaded_token is None - - -@pytest.mark.asyncio -async def test_token_expiry_recalculated_on_load(tmp_path: Path): - """Test that expires_in is correctly recalculated when loading tokens.""" - storage = FileTokenStorage("http://test.example.com", cache_dir=tmp_path) - - # Manually create a token file with a specific expires_at - token_file = storage._get_file_path("tokens") - future_expiry = datetime.now(timezone.utc) + timedelta( - seconds=1800 - ) # 30 minutes from now - - # JSONFileStorage expects wrapped format - stored_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(), - }, - "timestamp": datetime.now(timezone.utc).timestamp(), - } - token_file.write_text(json.dumps(stored_token, indent=2, default=str)) - - # Load the token - loaded_token = await storage.get_tokens() - assert loaded_token is not None - # expires_in should be recalculated to approximately 1800 seconds - assert loaded_token.expires_in is not None - assert 1795 <= loaded_token.expires_in <= 1800 diff --git a/tests/server/auth/test_oauth_proxy_redirect_validation.py b/tests/server/auth/test_oauth_proxy_redirect_validation.py index 8a2ab8564..a4185538f 100644 --- a/tests/server/auth/test_oauth_proxy_redirect_validation.py +++ b/tests/server/auth/test_oauth_proxy_redirect_validation.py @@ -176,7 +176,7 @@ async def test_proxy_register_client_uses_patterns(self): "new-client" ) # Use the client ID we registered assert isinstance(registered, ProxyDCRClient) - assert registered._allowed_redirect_uri_patterns == custom_patterns + assert registered.allowed_redirect_uri_patterns == custom_patterns @pytest.mark.asyncio async def test_proxy_unregistered_client_returns_none(self): diff --git a/tests/server/auth/test_oauth_proxy_storage.py b/tests/server/auth/test_oauth_proxy_storage.py index 6ceef86aa..629708427 100644 --- a/tests/server/auth/test_oauth_proxy_storage.py +++ b/tests/server/auth/test_oauth_proxy_storage.py @@ -1,14 +1,18 @@ """Tests for OAuth proxy with persistent storage.""" +from collections.abc import AsyncGenerator from pathlib import Path from unittest.mock import AsyncMock, Mock import pytest +from diskcache.core import tempfile +from inline_snapshot import snapshot +from key_value.aio.stores.disk import MultiDiskStore +from key_value.aio.stores.memory import MemoryStore 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: @@ -23,14 +27,17 @@ def jwt_verifier(self): return verifier @pytest.fixture - def temp_storage(self, tmp_path: Path) -> JSONFileStorage: + async def temp_storage(self) -> AsyncGenerator[MultiDiskStore, None]: """Create file-based storage for testing.""" - return JSONFileStorage(tmp_path / "oauth-clients") + with tempfile.TemporaryDirectory() as temp_dir: + disk_store = MultiDiskStore(base_directory=Path(temp_dir)) + yield disk_store + await disk_store.close() @pytest.fixture - def memory_storage(self) -> InMemoryStorage: + def memory_storage(self) -> MemoryStore: """Create in-memory storage for testing.""" - return InMemoryStorage() + return MemoryStore() def create_proxy(self, jwt_verifier, storage=None) -> OAuthProxy: """Create an OAuth proxy with specified storage.""" @@ -48,7 +55,7 @@ def create_proxy(self, jwt_verifier, storage=None) -> OAuthProxy: 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) + assert isinstance(proxy._client_storage, MemoryStore) async def test_register_and_get_client(self, jwt_verifier, temp_storage): """Test registering and retrieving a client.""" @@ -132,7 +139,7 @@ async def test_proxy_dcr_client_redirect_validation( async def test_in_memory_storage_option(self, jwt_verifier): """Test using in-memory storage explicitly.""" - storage = InMemoryStorage() + storage = MemoryStore() proxy = self.create_proxy(jwt_verifier, storage=storage) client_info = OAuthClientInformationFull( @@ -151,7 +158,7 @@ async def test_in_memory_storage_option(self, jwt_verifier): assert client2 is not None # But new storage instance won't have it - proxy3 = self.create_proxy(jwt_verifier, storage=InMemoryStorage()) + proxy3 = self.create_proxy(jwt_verifier, storage=MemoryStore()) client3 = await proxy3.get_client("memory-client") assert client3 is None @@ -167,47 +174,31 @@ async def test_storage_data_structure(self, jwt_verifier, temp_storage): 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")], + raw_data = await temp_storage.get( + collection="mcp-oauth-proxy-clients", key="structured-client" ) - 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 raw_data is not None + assert raw_data == snapshot( + { + "redirect_uris": ["http://localhost:8080/callback"], + "token_endpoint_auth_method": "none", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "scope": "read write", + "client_name": None, + "client_uri": None, + "logo_uri": None, + "contacts": None, + "tos_uri": None, + "policy_uri": None, + "jwks_uri": None, + "jwks": None, + "software_id": None, + "software_version": None, + "client_id": "structured-client", + "client_secret": "secret", + "client_id_issued_at": None, + "client_secret_expires_at": None, + "allowed_redirect_uri_patterns": None, + } ) - 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 deleted file mode 100644 index 1c53b6637..000000000 --- a/tests/utilities/test_storage.py +++ /dev/null @@ -1,143 +0,0 @@ -"""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 diff --git a/uv.lock b/uv.lock index 82b662287..72c90e253 100644 --- a/uv.lock +++ b/uv.lock @@ -69,6 +69,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/59/76ab57e3fe74484f48a53f8e337171b4a2349e506eabe136d7e01d059086/backports_asyncio_runner-1.2.0-py3-none-any.whl", hash = "sha256:0da0a936a8aeb554eccb426dc55af3ba63bcdc69fa1a600b5bb305413a4477b5", size = 12313, upload-time = "2025-07-02T02:27:14.263Z" }, ] +[[package]] +name = "cachetools" +version = "6.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/61/e4fad8155db4a04bfb4734c7c8ff0882f078f24294d42798b3568eb63bff/cachetools-6.2.0.tar.gz", hash = "sha256:38b328c0889450f05f5e120f56ab68c8abaf424e1275522b138ffc93253f7e32", size = 30988, upload-time = "2025-08-25T18:57:30.924Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/56/3124f61d37a7a4e7cc96afc5492c78ba0cb551151e530b54669ddd1436ef/cachetools-6.2.0-py3-none-any.whl", hash = "sha256:1c76a8960c0041fcc21097e357f882197c79da0dbff766e7317890a65d7d8ba6", size = 11276, upload-time = "2025-08-25T18:57:29.684Z" }, +] + [[package]] name = "certifi" version = "2025.8.3" @@ -400,6 +409,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/0c/03cc99bf3b6328604b10829de3460f2b2ad3373200c45665c38508e550c6/dirty_equals-0.9.0-py3-none-any.whl", hash = "sha256:ff4d027f5cfa1b69573af00f7ba9043ea652dbdce3fe5cbe828e478c7346db9c", size = 28226, upload-time = "2025-01-11T23:23:37.489Z" }, ] +[[package]] +name = "diskcache" +version = "5.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3f/21/1c1ffc1a039ddcc459db43cc108658f32c57d271d7289a2794e401d0fdb6/diskcache-5.6.3.tar.gz", hash = "sha256:2c3a3fa2743d8535d832ec61c2054a1641f41775aa7c556758a109941e33e4fc", size = 67916, upload-time = "2023-08-31T06:12:00.316Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" }, +] + [[package]] name = "distlib" version = "0.4.0" @@ -526,6 +544,7 @@ dependencies = [ { name = "mcp" }, { name = "openapi-core" }, { name = "openapi-pydantic" }, + { name = "py-key-value-aio", extra = ["disk", "memory"] }, { name = "pydantic", extra = ["email"] }, { name = "pyperclip" }, { name = "python-dotenv" }, @@ -575,6 +594,7 @@ requires-dist = [ { name = "openai", marker = "extra == 'openai'", specifier = ">=1.102.0" }, { name = "openapi-core", specifier = ">=0.19.5" }, { name = "openapi-pydantic", specifier = ">=0.5.1" }, + { name = "py-key-value-aio", extras = ["disk", "memory"], specifier = ">=0.2.1" }, { name = "pydantic", extras = ["email"], specifier = ">=2.11.7" }, { name = "pyperclip", specifier = ">=1.9.0" }, { name = "python-dotenv", specifier = ">=1.1.0" }, @@ -1175,6 +1195,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7d/eb/b6260b31b1a96386c0a880edebe26f89669098acea8e0318bff6adb378fd/pathable-0.4.4-py3-none-any.whl", hash = "sha256:5ae9e94793b6ef5a4cbe0a7ce9dbbefc1eec38df253763fd0aeeacf2762dbbc2", size = 9592, upload-time = "2025-01-10T18:43:11.88Z" }, ] +[[package]] +name = "pathvalidate" +version = "3.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fa/2a/52a8da6fe965dea6192eb716b357558e103aea0a1e9a8352ad575a8406ca/pathvalidate-3.3.1.tar.gz", hash = "sha256:b18c07212bfead624345bb8e1d6141cdcf15a39736994ea0b94035ad2b1ba177", size = 63262, upload-time = "2025-06-15T09:07:20.736Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/70/875f4a23bfc4731703a5835487d0d2fb999031bd415e7d17c0ae615c18b7/pathvalidate-3.3.1-py3-none-any.whl", hash = "sha256:5263baab691f8e1af96092fa5137ee17df5bdfbd6cff1fcac4d6ef4bc2e1735f", size = 24305, upload-time = "2025-06-15T09:07:19.117Z" }, +] + [[package]] name = "pdbpp" version = "0.11.7" @@ -1279,6 +1308,39 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, ] +[[package]] +name = "py-key-value-aio" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "py-key-value-shared" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f9/bf/7237a1d41b4afc33a8c0f71c991d95a6bb6719cd5ccab8d1628b72fbe03c/py_key_value_aio-0.2.1.tar.gz", hash = "sha256:79c8c835451b61d4abd863c65d33870612f3a80dc312120b2d1445269764d625", size = 19440, upload-time = "2025-10-09T03:26:28.357Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/82/41b5574270fbed7171d34a9b7c9b1b18fd86c31421e45dc935ba354eb42f/py_key_value_aio-0.2.1-py3-none-any.whl", hash = "sha256:5f0bc1bb3f886578a88ed2b61858658142db35c59dd3ccd9ec727184c540288a", size = 41564, upload-time = "2025-10-09T03:26:26.174Z" }, +] + +[package.optional-dependencies] +disk = [ + { name = "diskcache" }, + { name = "pathvalidate" }, +] +memory = [ + { name = "cachetools" }, +] + +[[package]] +name = "py-key-value-shared" +version = "0.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/70/09/7c76aa82e5e41c6ad5e0e43bcc2072b48d84e03439dbb25b3e184773b553/py_key_value_shared-0.2.0.tar.gz", hash = "sha256:ee6d9a9101b54f228876c61b2f2f83a951c9c52233d8271599532c069fa26052", size = 6285, upload-time = "2025-09-29T02:27:46.252Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/f8/6c6cf5abcb78d103006ea1bec6137c9859611ffc50093684b5130c5642c1/py_key_value_shared-0.2.0-py3-none-any.whl", hash = "sha256:84cb4f6b6bed97a32feebc512ce1e333097ce5768c7198abcd7d4bd3c5f1de06", size = 10437, upload-time = "2025-09-29T02:27:45.281Z" }, +] + [[package]] name = "pycparser" version = "2.22" @@ -2082,11 +2144,11 @@ wheels = [ [[package]] name = "typing-extensions" -version = "4.14.1" +version = "4.15.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/98/5a/da40306b885cc8c09109dc2e1abd358d5684b1425678151cdaed4731c822/typing_extensions-4.14.1.tar.gz", hash = "sha256:38b39f4aeeab64884ce9f74c94263ef78f3c22467c8724005483154c26648d36", size = 107673, upload-time = "2025-07-04T13:28:34.16Z" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/00/d631e67a838026495268c2f6884f3711a15a9a2a96cd244fdaea53b823fb/typing_extensions-4.14.1-py3-none-any.whl", hash = "sha256:d1e1e3b58374dc93031d6eda2420a48ea44a36c2b4766a4fdeb3710755731d76", size = 43906, upload-time = "2025-07-04T13:28:32.743Z" }, + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] [[package]]