mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-17 11:09:11 +02:00
Compare commits
4 commits
main
...
quickstart
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a42de8d429 | ||
|
|
3977ac447f | ||
|
|
4115de6311 | ||
|
|
2fa5e49e41 |
2 changed files with 91 additions and 55 deletions
|
|
@ -4,6 +4,7 @@ import asyncio
|
||||||
import json
|
import json
|
||||||
import webbrowser
|
import webbrowser
|
||||||
from asyncio import Future
|
from asyncio import Future
|
||||||
|
from collections.abc import AsyncGenerator
|
||||||
from datetime import datetime, timedelta, timezone
|
from datetime import datetime, timedelta, timezone
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Literal
|
from typing import Any, Literal
|
||||||
|
|
@ -34,6 +35,12 @@ __all__ = ["OAuth"]
|
||||||
logger = get_logger(__name__)
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ClientNotFoundError(Exception):
|
||||||
|
"""Raised when OAuth client credentials are not found on the server."""
|
||||||
|
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
class StoredToken(BaseModel):
|
class StoredToken(BaseModel):
|
||||||
"""Token storage format with absolute expiry time."""
|
"""Token storage format with absolute expiry time."""
|
||||||
|
|
||||||
|
|
@ -173,7 +180,7 @@ class FileTokenStorage(TokenStorage):
|
||||||
for file_type in file_types:
|
for file_type in file_types:
|
||||||
path = self._get_file_path(file_type)
|
path = self._get_file_path(file_type)
|
||||||
path.unlink(missing_ok=True)
|
path.unlink(missing_ok=True)
|
||||||
logger.info(f"Cleared OAuth cache for {self.get_base_url(self.server_url)}")
|
logger.debug(f"Cleared OAuth cache for {self.get_base_url(self.server_url)}")
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def clear_all(cls, cache_dir: Path | None = None) -> None:
|
def clear_all(cls, cache_dir: Path | None = None) -> None:
|
||||||
|
|
@ -300,7 +307,23 @@ class OAuth(OAuthClientProvider):
|
||||||
self.context.update_token_expiry(self.context.current_tokens)
|
self.context.update_token_expiry(self.context.current_tokens)
|
||||||
|
|
||||||
async def redirect_handler(self, authorization_url: str) -> None:
|
async def redirect_handler(self, authorization_url: str) -> None:
|
||||||
"""Open browser for authorization."""
|
"""Open browser for authorization, with pre-flight check for invalid client."""
|
||||||
|
# Pre-flight check to detect invalid client_id before opening browser
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
response = await client.get(authorization_url, follow_redirects=False)
|
||||||
|
|
||||||
|
# Check for client not found error (400 typically means bad client_id)
|
||||||
|
if response.status_code == 400:
|
||||||
|
raise ClientNotFoundError(
|
||||||
|
"OAuth client not found - cached credentials may be stale"
|
||||||
|
)
|
||||||
|
|
||||||
|
# For any non-redirect response, something is wrong
|
||||||
|
if response.status_code not in (302, 303, 307, 308):
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Unexpected authorization response: {response.status_code}"
|
||||||
|
)
|
||||||
|
|
||||||
logger.info(f"OAuth authorization URL: {authorization_url}")
|
logger.info(f"OAuth authorization URL: {authorization_url}")
|
||||||
webbrowser.open(authorization_url)
|
webbrowser.open(authorization_url)
|
||||||
|
|
||||||
|
|
@ -336,3 +359,56 @@ class OAuth(OAuthClientProvider):
|
||||||
tg.cancel_scope.cancel()
|
tg.cancel_scope.cancel()
|
||||||
|
|
||||||
raise RuntimeError("OAuth callback handler could not be started")
|
raise RuntimeError("OAuth callback handler could not be started")
|
||||||
|
|
||||||
|
async def async_auth_flow(
|
||||||
|
self, request: httpx.Request
|
||||||
|
) -> AsyncGenerator[httpx.Request, httpx.Response]:
|
||||||
|
"""HTTPX auth flow with automatic retry on stale cached credentials.
|
||||||
|
|
||||||
|
If the OAuth flow fails due to invalid/stale client credentials,
|
||||||
|
clears the cache and retries once with fresh registration.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# First attempt with potentially cached credentials
|
||||||
|
gen = super().async_auth_flow(request)
|
||||||
|
response = None
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
yielded_request = await gen.asend(response)
|
||||||
|
response = yield yielded_request
|
||||||
|
except StopAsyncIteration:
|
||||||
|
break
|
||||||
|
|
||||||
|
except ClientNotFoundError:
|
||||||
|
logger.debug(
|
||||||
|
"OAuth client not found on server, clearing cache and retrying..."
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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
|
||||||
|
|
||||||
|
gen = super().async_auth_flow(request)
|
||||||
|
response = None
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
yielded_request = await gen.asend(response)
|
||||||
|
response = yield yielded_request
|
||||||
|
except StopAsyncIteration:
|
||||||
|
break
|
||||||
|
|
|
||||||
|
|
@ -172,7 +172,6 @@ class OAuthProxy(OAuthProvider):
|
||||||
1. Client Registration (DCR):
|
1. Client Registration (DCR):
|
||||||
- Accept any client registration request
|
- Accept any client registration request
|
||||||
- Store ProxyDCRClient that accepts dynamic redirect URIs
|
- Store ProxyDCRClient that accepts dynamic redirect URIs
|
||||||
- Return shared upstream credentials to all clients
|
|
||||||
|
|
||||||
2. Authorization:
|
2. Authorization:
|
||||||
- Store transaction mapping client details to proxy flow
|
- Store transaction mapping client details to proxy flow
|
||||||
|
|
@ -323,67 +322,28 @@ class OAuthProxy(OAuthProvider):
|
||||||
# -------------------------------------------------------------------------
|
# -------------------------------------------------------------------------
|
||||||
|
|
||||||
async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
|
async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
|
||||||
"""Get client information by ID.
|
"""Get client information by ID. This is generally the random ID
|
||||||
|
provided to the DCR client during registration, not the upstream client ID.
|
||||||
|
|
||||||
For unregistered clients, returns a ProxyDCRClient that accepts
|
For unregistered clients, returns None (which will raise an error in the SDK).
|
||||||
any localhost redirect URI for DCR clients.
|
|
||||||
|
|
||||||
Even registered clients use ProxyDCRClient to ensure they can
|
|
||||||
authenticate with different dynamic ports on reconnection. This
|
|
||||||
handles the case where a client with cached tokens reconnects
|
|
||||||
on a different port.
|
|
||||||
"""
|
"""
|
||||||
client = self._clients.get(client_id)
|
client = self._clients.get(client_id)
|
||||||
|
|
||||||
if client is None:
|
|
||||||
# For unregistered DCR clients, create a permissive client
|
|
||||||
# that will accept any localhost redirect URI
|
|
||||||
# We need at least one URI for Pydantic validation, but our custom
|
|
||||||
# validate_redirect_uri will accept any localhost URI
|
|
||||||
client = ProxyDCRClient(
|
|
||||||
client_id=client_id,
|
|
||||||
client_secret=None,
|
|
||||||
redirect_uris=[
|
|
||||||
AnyUrl("http://localhost")
|
|
||||||
], # Placeholder, validation uses allowed_patterns
|
|
||||||
grant_types=["authorization_code", "refresh_token"],
|
|
||||||
scope=self._default_scope_str,
|
|
||||||
token_endpoint_auth_method="none",
|
|
||||||
allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
|
|
||||||
)
|
|
||||||
logger.debug("Created ProxyDCRClient for unregistered client %s", client_id)
|
|
||||||
|
|
||||||
return client
|
return client
|
||||||
|
|
||||||
async def register_client(self, client_info: OAuthClientInformationFull) -> None:
|
async def register_client(self, client_info: OAuthClientInformationFull) -> None:
|
||||||
"""Register a client locally using fixed upstream credentials.
|
"""Register a client locally
|
||||||
|
|
||||||
This implementation always uses the upstream client_id and client_secret
|
When a client registers, we create a ProxyDCRClient that is more
|
||||||
regardless of what the client requests. It modifies the client_info object
|
forgiving about validating redirect URIs, since the DCR client's
|
||||||
in place since the MCP framework ignores return values.
|
redirect URI will likely be localhost or unknown to the proxied IDP. The
|
||||||
|
proxied IDP only knows about this server's fixed redirect URI.
|
||||||
This ensures all clients use the same credentials that are registered
|
|
||||||
with the upstream server.
|
|
||||||
|
|
||||||
Implementation Detail:
|
|
||||||
We store a ProxyDCRClient (not the original client_info) to ensure
|
|
||||||
the client can reconnect with different dynamic redirect URIs. This is
|
|
||||||
essential for cached token scenarios where the client port changes.
|
|
||||||
|
|
||||||
The flow:
|
|
||||||
1. Client provides its desired redirect URIs (dynamic localhost ports)
|
|
||||||
2. We create a ProxyDCRClient that will accept ANY localhost URI
|
|
||||||
3. We store this flexible client for future authentications
|
|
||||||
4. When client reconnects with a different port, ProxyDCRClient accepts it
|
|
||||||
"""
|
"""
|
||||||
# Always use the upstream credentials
|
|
||||||
upstream_id = self._upstream_client_id
|
|
||||||
upstream_secret = self._upstream_client_secret.get_secret_value()
|
|
||||||
|
|
||||||
# Create a ProxyDCRClient with configured redirect URI validation
|
# Create a ProxyDCRClient with configured redirect URI validation
|
||||||
proxy_client = ProxyDCRClient(
|
proxy_client = ProxyDCRClient(
|
||||||
client_id=upstream_id,
|
client_id=client_info.client_id,
|
||||||
client_secret=upstream_secret,
|
client_secret=client_info.client_secret,
|
||||||
redirect_uris=client_info.redirect_uris or [AnyUrl("http://localhost")],
|
redirect_uris=client_info.redirect_uris or [AnyUrl("http://localhost")],
|
||||||
grant_types=client_info.grant_types
|
grant_types=client_info.grant_types
|
||||||
or ["authorization_code", "refresh_token"],
|
or ["authorization_code", "refresh_token"],
|
||||||
|
|
@ -392,8 +352,8 @@ class OAuthProxy(OAuthProvider):
|
||||||
allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
|
allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Store the ProxyDCRClient using the upstream ID
|
# Store the ProxyDCRClient
|
||||||
self._clients[upstream_id] = proxy_client
|
self._clients[client_info.client_id] = proxy_client
|
||||||
|
|
||||||
# Log redirect URIs to help users discover what patterns they might need
|
# Log redirect URIs to help users discover what patterns they might need
|
||||||
if client_info.redirect_uris:
|
if client_info.redirect_uris:
|
||||||
|
|
@ -406,7 +366,7 @@ class OAuthProxy(OAuthProvider):
|
||||||
|
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"Registered client %s with %d redirect URIs",
|
"Registered client %s with %d redirect URIs",
|
||||||
upstream_id,
|
client_info.client_id,
|
||||||
len(proxy_client.redirect_uris),
|
len(proxy_client.redirect_uris),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue