Add raise on validation error and set uniform collection names

This commit is contained in:
William Easton 2025-10-08 23:40:58 -04:00
commit de0c39c1b3
No known key found for this signature in database
5 changed files with 25 additions and 21 deletions

View file

@ -15,7 +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.0",
"py-key-value-aio[disk,memory]>=0.2.1",
"websockets>=15.0.1",
]

View file

@ -20,6 +20,7 @@ from mcp.shared.auth import (
OAuthToken,
)
from pydantic import AnyHttpUrl
from typing_extensions import override
from uvicorn.server import Server
from fastmcp.client.oauth_callback import (
@ -79,10 +80,16 @@ class TokenStorageAdapter(TokenStorage):
self._server_url = server_url
self._key_value_store = async_key_value
self._storage_oauth_token = PydanticAdapter[OAuthToken](
key_value=async_key_value, pydantic_model=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](
key_value=async_key_value, pydantic_model=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:
@ -92,31 +99,28 @@ class TokenStorageAdapter(TokenStorage):
return f"{self._server_url}/client_info"
async def clear(self) -> None:
await self._storage_oauth_token.delete(
collection="oauth-mcp-client-cache", key=self._get_token_cache_key()
)
await self._storage_client_info.delete(
collection="oauth-mcp-client-cache", key=self._get_client_info_cache_key()
)
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(
collection="oauth-mcp-client-cache", key=self._get_token_cache_key()
)
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(
collection="oauth-mcp-client-cache",
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(
collection="oauth-mcp-client-cache", key=self._get_client_info_cache_key()
key=self._get_client_info_cache_key()
)
@override
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
ttl: int | None = None
@ -124,7 +128,6 @@ class TokenStorageAdapter(TokenStorage):
ttl = client_info.client_secret_expires_at - int(time.time())
await self._storage_client_info.put(
collection="oauth-mcp-client-cache",
key=self._get_client_info_cache_key(),
value=client_info,
ttl=ttl,

View file

@ -391,7 +391,8 @@ class OAuthProxy(OAuthProvider):
self._client_store = PydanticAdapter[ProxyDCRClient](
key_value=self._client_storage,
pydantic_model=ProxyDCRClient,
default_collection="oauth-proxy-clients",
default_collection="mcp-oauth-proxy-clients",
raise_on_validation_error=True,
)
# Local state for token bookkeeping only (no client caching)

View file

@ -175,7 +175,7 @@ class TestOAuthProxyStorage:
# Check raw storage data
raw_data = await temp_storage.get(
collection="oauth-proxy-clients", key="structured-client"
collection="mcp-oauth-proxy-clients", key="structured-client"
)
assert raw_data is not None
assert raw_data == snapshot(

8
uv.lock generated
View file

@ -594,7 +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.0" },
{ 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" },
@ -1310,14 +1310,14 @@ wheels = [
[[package]]
name = "py-key-value-aio"
version = "0.2.0"
version = "0.2.1"
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" }
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/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" },
{ 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]