Refactor for key-value 0.2.0

This commit is contained in:
William Easton 2025-09-28 21:35:22 -05:00
commit ccd12c4b1b
No known key found for this signature in database
12 changed files with 138 additions and 75 deletions

View file

@ -15,7 +15,7 @@ dependencies = [
"pydantic[email]>=2.11.7",
"pyperclip>=1.9.0",
"openapi-core>=0.19.5",
"kv-store-adapter[disk,memory]>=0.1.2",
"py-key-value-aio[disk,memory]>=0.2.0",
]
requires-python = ">=3.10"

View file

@ -10,8 +10,8 @@ from urllib.parse import urlparse
import anyio
import httpx
from kv_store_adapter.adapters.pydantic import PydanticAdapter
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.adapters.pydantic import PydanticAdapter
from key_value.aio.protocols import AsyncKeyValue
from mcp.client.auth import OAuthClientProvider, TokenStorage
from mcp.shared.auth import (
OAuthClientInformationFull,
@ -71,18 +71,18 @@ async def check_if_auth_required(
class TokenStorageAdapter(TokenStorage):
_server_url: str
_kv_store_protocol: KVStoreProtocol
_key_value_store: AsyncKeyValue
_storage_oauth_token: PydanticAdapter[OAuthToken]
_storage_client_info: PydanticAdapter[OAuthClientInformationFull]
def __init__(self, kv_store_protocol: KVStoreProtocol, server_url: str):
def __init__(self, async_key_value: AsyncKeyValue, server_url: str):
self._server_url = server_url
self._kv_store_protocol = kv_store_protocol
self._key_value_store = async_key_value
self._storage_oauth_token = PydanticAdapter[OAuthToken](
store_protocol=kv_store_protocol, pydantic_model=OAuthToken
key_value=async_key_value, pydantic_model=OAuthToken
)
self._storage_client_info = PydanticAdapter[OAuthClientInformationFull](
store_protocol=kv_store_protocol, pydantic_model=OAuthClientInformationFull
key_value=async_key_value, pydantic_model=OAuthClientInformationFull
)
def _get_token_cache_key(self) -> str:
@ -144,7 +144,7 @@ class OAuth(OAuthClientProvider):
mcp_url: str,
scopes: str | list[str] | None = None,
client_name: str = "FastMCP Client",
token_storage: KVStoreProtocol | None = None,
token_storage: AsyncKeyValue | None = None,
additional_client_metadata: dict[str, Any] | None = None,
callback_port: int | None = None,
):
@ -156,7 +156,7 @@ class OAuth(OAuthClientProvider):
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: KVStoreProtocol for token storage, the default disk store is used if not provided
token_storage: AsyncKeyValue for token storage, the default disk store is used if not provided
additional_client_metadata: Extra fields for OAuthClientMetadata
callback_port: Fixed port for OAuth callback (default: random available port)
"""
@ -186,10 +186,10 @@ class OAuth(OAuthClientProvider):
)
# Create server-specific token storage
token_storage = token_storage or settings.data_store
token_storage = token_storage or settings.key_value_store
self.token_storage_adapter: TokenStorageAdapter = TokenStorageAdapter(
kv_store_protocol=token_storage, server_url=server_base_url
async_key_value=token_storage, server_url=server_base_url
)
# Store server_base_url for use in callback_handler

View file

@ -22,14 +22,15 @@ import hashlib
import secrets
import time
from base64 import urlsafe_b64encode
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final
from urllib.parse import urlencode
import httpx
from authlib.common.security import generate_token
from authlib.integrations.httpx_client import AsyncOAuth2Client
from kv_store_adapter.adapters.pydantic import PydanticAdapter
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.adapters.pydantic import PydanticAdapter
from key_value.aio.protocols import AsyncKeyValue
from mcp.server.auth.provider import (
AccessToken,
AuthorizationCode,
@ -42,7 +43,7 @@ from mcp.server.auth.settings import (
RevocationOptions,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
from pydantic import AnyHttpUrl, AnyUrl, Field, SecretStr
from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, SecretStr
from starlette.requests import Request
from starlette.responses import RedirectResponse
from starlette.routing import Route
@ -115,6 +116,12 @@ DEFAULT_AUTH_CODE_EXPIRY_SECONDS: Final[int] = 5 * 60 # 5 minutes
HTTP_TIMEOUT_SECONDS: Final[int] = 30
@dataclass
class RelatedTokens(BaseModel):
access_token: str
refresh_token: str
class OAuthProxy(OAuthProvider):
"""OAuth provider that presents a DCR-compliant interface while proxying to non-DCR IDPs.
@ -194,7 +201,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
@ -250,7 +256,7 @@ class OAuthProxy(OAuthProvider):
# Extra parameters to forward to token endpoint
extra_token_params: dict[str, str] | None = None,
# Client storage
client_storage: KVStoreProtocol | None = None,
client_storage: AsyncKeyValue | None = None,
):
"""Initialize the OAuth proxy provider.
@ -333,12 +339,13 @@ class OAuthProxy(OAuthProvider):
self._extra_authorize_params = extra_authorize_params or {}
self._extra_token_params = extra_token_params or {}
self._client_storage: KVStoreProtocol = client_storage or settings.data_store
self._client_storage_collection = "oauth-proxy-clients"
self._client_storage: AsyncKeyValue = client_storage or settings.key_value_store
self._client_storage_adapter: PydanticAdapter[ProxyDCRClient] = PydanticAdapter[
ProxyDCRClient
](store_protocol=self._client_storage, pydantic_model=ProxyDCRClient)
self._client_store = PydanticAdapter[ProxyDCRClient](
key_value=self._client_storage,
pydantic_model=ProxyDCRClient,
default_collection="oauth-proxy-clients",
)
# Local state for token bookkeeping only (no client caching)
self._access_tokens: dict[str, AccessToken] = {}
@ -392,11 +399,7 @@ class OAuthProxy(OAuthProvider):
For unregistered clients, returns None (which will raise an error in the SDK).
"""
# Load from storage
if not (
client := await self._client_storage_adapter.get(
collection=self._client_storage_collection, key=client_id
)
):
if not (client := await self._client_store.get(key=client_id)):
return None
if client.allowed_redirect_uri_patterns is None:
@ -425,8 +428,7 @@ class OAuthProxy(OAuthProvider):
allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
)
await self._client_storage_adapter.put(
collection=self._client_storage_collection,
await self._client_store.put(
key=client_info.client_id,
value=proxy_client,
)

View file

@ -12,7 +12,7 @@ This implementation is based on:
from collections.abc import Sequence
import httpx
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.protocols import AsyncKeyValue
from pydantic import AnyHttpUrl, BaseModel, model_validator
from typing_extensions import Self
@ -213,7 +213,7 @@ class OIDCProxy(OAuthProxy):
redirect_path: str | None = None,
# Client configuration
allowed_client_redirect_uris: list[str] | None = None,
client_storage: KVStoreProtocol | None = None,
client_storage: AsyncKeyValue | None = None,
# Token validation configuration
token_endpoint_auth_method: str | None = None,
) -> None:

View file

@ -21,7 +21,7 @@ Example:
```
"""
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.protocols import AsyncKeyValue
from pydantic import AnyHttpUrl, SecretStr, field_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
@ -92,7 +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: KVStoreProtocol | None = None,
client_storage: AsyncKeyValue | None = None,
) -> None:
"""Initialize Auth0 OAuth provider.

View file

@ -7,7 +7,7 @@ using the OAuth Proxy pattern for non-DCR OAuth flows.
from __future__ import annotations
import httpx
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.protocols import AsyncKeyValue
from pydantic import SecretStr, field_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
@ -161,7 +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: KVStoreProtocol | None = None,
client_storage: AsyncKeyValue | None = None,
):
"""Initialize Azure OAuth provider.

View file

@ -22,7 +22,7 @@ Example:
from __future__ import annotations
import httpx
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.protocols import AsyncKeyValue
from pydantic import AnyHttpUrl, SecretStr, field_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
@ -202,7 +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: KVStoreProtocol | None = None,
client_storage: AsyncKeyValue | None = None,
):
"""Initialize GitHub OAuth provider.

View file

@ -24,7 +24,7 @@ from __future__ import annotations
import time
import httpx
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.protocols import AsyncKeyValue
from pydantic import AnyHttpUrl, SecretStr, field_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
@ -218,7 +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: KVStoreProtocol | None = None,
client_storage: AsyncKeyValue | None = None,
):
"""Initialize Google OAuth provider.

View file

@ -13,7 +13,7 @@ from __future__ import annotations
from typing import Any
import httpx
from kv_store_adapter.types import KVStoreProtocol
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
@ -170,7 +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: KVStoreProtocol | None = None,
client_storage: AsyncKeyValue | None = None,
):
"""Initialize WorkOS OAuth provider.

View file

@ -6,7 +6,7 @@ from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any, Literal
from kv_store_adapter.types import KVStoreProtocol
from key_value.aio.protocols import AsyncKeyValue
from pydantic import Field, ImportString, field_validator
from pydantic.fields import FieldInfo
from pydantic_settings import (
@ -62,6 +62,27 @@ class ExtendedSettingsConfigDict(SettingsConfigDict, total=False):
env_prefixes: list[str] | None
class DiskStorageSettings(BaseSettings):
model_config = SettingsConfigDict(
env_prefix="FASTMCP_STORAGE_",
extra="ignore",
)
directory: Annotated[
str | None,
Field(
description="A custom path to store data in. If set to `None` (default), a folder called `data` will be created in the `home` directory."
),
] = None
max_collection_size: Annotated[
int,
Field(
description="The maximum size for each collection in the storage directory, in bytes."
),
] = 1024 * 1024 * 10 # 10MB
class ExperimentalSettings(BaseSettings):
model_config = SettingsConfigDict(
env_prefix="FASTMCP_EXPERIMENTAL_",
@ -151,7 +172,12 @@ class Settings(BaseSettings):
home: Path = Path.home() / ".fastmcp"
data_path: Path | None = home / "data.db"
storage: Annotated[
DiskStorageSettings | None,
Field(
description="The default storage settings for the server. Defaults to disk storage, if set to None, data will be stored in memory by default."
),
] = DiskStorageSettings()
test_mode: bool = False
@ -386,15 +412,20 @@ class Settings(BaseSettings):
return auth_class
@cached_property
def data_store(self) -> KVStoreProtocol:
if not self.data_path:
from kv_store_adapter.stores.memory import MemoryStore
def key_value_store(self) -> AsyncKeyValue:
"""A default data store that can be leveraged as a fallback for components that require storage."""
if not self.storage:
from key_value.aio.stores.memory import MemoryStore
return MemoryStore()
from kv_store_adapter.stores.disk import DiskStore
from key_value.aio.stores.disk.multi_store import MultiDiskStore
return DiskStore(path=str(self.data_path), size_limit=TEN_MB_IN_BYTES)
base_directory: Path = self.storage.directory or (self.home / "data")
return MultiDiskStore(
base_directory=base_directory, max_size=self.storage.max_collection_size
)
def __getattr__(name: str):

View file

@ -1,12 +1,14 @@
"""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 kv_store_adapter.stores.disk import DiskStore
from kv_store_adapter.stores.memory import MemoryStore
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
@ -25,9 +27,12 @@ class TestOAuthProxyStorage:
return verifier
@pytest.fixture
def temp_storage(self, tmp_path: Path) -> DiskStore:
async def temp_storage(
self, tmp_path: Path
) -> AsyncGenerator[MultiDiskStore, None]:
"""Create file-based storage for testing."""
return DiskStore(path=str(tmp_path / "oauth-clients"))
with tempfile.TemporaryDirectory() as temp_dir:
yield MultiDiskStore(base_directory=Path(temp_dir))
@pytest.fixture
def memory_storage(self) -> MemoryStore:
@ -50,7 +55,7 @@ class TestOAuthProxyStorage:
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, DiskStore)
assert isinstance(proxy._client_storage, MultiDiskStore)
async def test_register_and_get_client(self, jwt_verifier, temp_storage):
"""Test registering and retrieving a client."""

69
uv.lock generated
View file

@ -541,10 +541,10 @@ dependencies = [
{ name = "cyclopts" },
{ name = "exceptiongroup" },
{ name = "httpx" },
{ name = "kv-store-adapter", extra = ["disk", "memory"] },
{ 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" },
@ -591,11 +591,11 @@ requires-dist = [
{ name = "cyclopts", specifier = ">=3.0.0" },
{ name = "exceptiongroup", specifier = ">=1.2.2" },
{ name = "httpx", specifier = ">=0.28.1" },
{ name = "kv-store-adapter", extras = ["disk", "memory"], specifier = ">=0.1.2" },
{ name = "mcp", specifier = ">=1.12.4,<2.0.0" },
{ 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.0" },
{ name = "pydantic", extras = ["email"], specifier = ">=2.11.7" },
{ name = "pyperclip", specifier = ">=1.9.0" },
{ name = "python-dotenv", specifier = ">=1.1.0" },
@ -929,23 +929,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/01/0e/b27cdbaccf30b890c40ed1da9fd4a3593a5cf94dae54fb34f8a4b74fcd3f/jsonschema_specifications-2025.4.1-py3-none-any.whl", hash = "sha256:4653bffbd6584f7de83a67e0d620ef16900b390ddc7939d56684d6c81e33f1af", size = 18437, upload-time = "2025-04-23T12:34:05.422Z" },
]
[[package]]
name = "kv-store-adapter"
version = "0.1.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/b9/ef/86cd8a32d13e24707422005738f6c28795b405f956358e0000d8dd471629/kv_store_adapter-0.1.2.tar.gz", hash = "sha256:790c9354d44962458716d7c157d12b8d7ab39f5ccb51c2be28abc28f6f73320f", size = 95486, upload-time = "2025-09-25T01:41:55.87Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/4d/97/516043daff470074deddb58de206b99aa7ee4ea1f881e219836ec6604630/kv_store_adapter-0.1.2-py3-none-any.whl", hash = "sha256:c24bd2e0b02b5779187b66e5d28609c7ffcb727773a515c304db64edcb872d4f", size = 29792, upload-time = "2025-09-25T01:41:52.854Z" },
]
[package.optional-dependencies]
disk = [
{ name = "diskcache" },
]
memory = [
{ name = "cachetools" },
]
[[package]]
name = "lazy-object-proxy"
version = "1.11.0"
@ -1212,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"
@ -1316,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.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "py-key-value-shared" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b1/96/b1c6b8ca95f411725449ccae5a0d4554ccc98c785026e722a5faba33625d/py_key_value_aio-0.2.0.tar.gz", hash = "sha256:d8276ff0cac0eec313c6961854087e476dbb76b97eff7310df30a7228a7939d8", size = 19328, upload-time = "2025-09-29T02:00:33.664Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/d9/17/c962abd354d6a302d415b7fdd68fa2149597c5f8e2b477f25bff2a9ec6d7/py_key_value_aio-0.2.0-py3-none-any.whl", hash = "sha256:775cf30d26fe958499757410b190eca4d338ce87613e27f2f537526b508fffb1", size = 41397, upload-time = "2025-09-29T02:00:32.701Z" },
]
[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"
@ -2107,11 +2132,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]]