This commit is contained in:
Jeremiah Lowin 2025-05-30 20:08:42 -04:00
commit 62eea7a20b
3 changed files with 47 additions and 27 deletions

View file

@ -186,7 +186,7 @@ class FileTokenStorage(TokenStorage):
async def set_tokens(self, tokens: _MCPOAuthToken) -> None:
"""Save tokens to file storage."""
# Convert to custom model with expiration datetime
tokens = OAuthToken.model_validate(tokens)
tokens = OAuthToken.model_validate(tokens.model_dump())
path = self._get_file_path("tokens")
path.write_text(tokens.model_dump_json(indent=2))
logger.debug(f"Saved tokens for {self.get_base_url(self.server_url)}")
@ -266,7 +266,7 @@ async def discover_oauth_metadata(
async def check_if_auth_required(
mcp_endpoint_url: str, httpx_kwargs: dict[str, Any] | None = None
mcp_url: str, httpx_kwargs: dict[str, Any] | None = None
) -> bool:
"""
Check if the MCP endpoint requires authentication by making a test request.
@ -277,7 +277,7 @@ async def check_if_auth_required(
async with httpx.AsyncClient(**(httpx_kwargs or {})) as client:
try:
# Try a simple request to the endpoint
response = await client.get(mcp_endpoint_url, timeout=5.0)
response = await client.get(mcp_url, timeout=5.0)
# If we get 401/403, auth is likely required
if response.status_code in (401, 403):
@ -296,7 +296,7 @@ async def check_if_auth_required(
def OAuth(
mcp_endpoint_url: str,
mcp_url: str,
scopes: str | list[str] | None = None,
client_name: str = "FastMCP Client",
token_storage_cache_dir: Path | None = None,
@ -309,17 +309,18 @@ def OAuth(
httpx.AsyncClient (or appropriate FastMCP client/transport instance)
Args:
mcp_endpoint_url: Full URL to the MCP endpoint (e.g.,
"http://host/mcp/sse") 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 additional_client_metadata: Extra fields for
OAuthClientMetadata
mcp_url: Full URL to the MCP endpoint (e.g.,
"http://host/mcp/sse")
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
additional_client_metadata: Extra fields for OAuthClientMetadata
Returns:
OAuthClientProvider
"""
parsed_url = urlparse(mcp_endpoint_url)
parsed_url = urlparse(mcp_url)
server_base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
# Setup OAuth client
@ -347,7 +348,7 @@ def OAuth(
# Define OAuth handlers
async def redirect_handler(authorization_url: str) -> None:
"""Open browser for authorization."""
logger.info(f"Opening browser for OAuth authorization: {authorization_url}")
logger.info(f"OAuth authorization URL: {authorization_url}")
webbrowser.open(authorization_url)
async def callback_handler() -> tuple[str, str | None]:
@ -363,19 +364,19 @@ def OAuth(
)
# Run server until response is received with timeout logic
async with anyio.create_task_group() as tg:
tg.start_soon(server.serve)
logger.info(
f"🎧 OAuth callback server started on http://127.0.0.1:{redirect_port}"
)
TIMEOUT = 300.0 # 5 minute timeout
try:
with anyio.fail_after(300.0): # 5 minute timeout
with anyio.fail_after(TIMEOUT):
auth_code, state = await response_future
return auth_code, state
except TimeoutError:
raise TimeoutError("OAuth callback timed out after 300 seconds")
raise TimeoutError(f"OAuth callback timed out after {TIMEOUT} seconds")
finally:
server.should_exit = True
await asyncio.sleep(0.1) # Allow server to shutdown gracefully

View file

@ -9,8 +9,10 @@ from __future__ import annotations
import asyncio
import socket
from dataclasses import dataclass
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import HTMLResponse
from starlette.routing import Route
from uvicorn import Config, Server
@ -184,6 +186,21 @@ def find_available_port() -> int:
return s.getsockname()[1]
@dataclass
class CallbackResponse:
code: str | None = None
state: str | None = None
error: str | None = None
error_description: str | None = None
@classmethod
def from_dict(cls, data: dict[str, str]) -> CallbackResponse:
return cls(**{k: v for k, v in data.items() if k in cls.__annotations__})
def to_dict(self) -> dict[str, str]:
return {k: v for k, v in self.__dict__.items() if v is not None}
def create_oauth_callback_server(
port: int,
callback_path: str = "/callback",
@ -203,30 +220,31 @@ def create_oauth_callback_server(
Configured uvicorn Server instance (not yet running)
"""
async def callback_handler(request):
async def callback_handler(request: Request):
"""Handle OAuth callback requests with proper HTML responses."""
query_params = dict(request.query_params)
auth_code = query_params.get("code")
state = query_params.get("state")
error = query_params.get("error")
callback_response = CallbackResponse.from_dict(query_params)
if error:
error_desc = query_params.get("error_description", "Unknown error")
if callback_response.error:
error_desc = callback_response.error_description or "Unknown error"
# Resolve future with exception if provided
if response_future and not response_future.done():
response_future.set_exception(
RuntimeError(f"OAuth error: {error} - {error_desc}")
RuntimeError(
f"OAuth error: {callback_response.error} - {error_desc}"
)
)
return HTMLResponse(
create_callback_html(
f"FastMCP OAuth Error: {error}<br>{error_desc}", is_success=False
f"FastMCP OAuth Error: {callback_response.error}<br>{error_desc}",
is_success=False,
),
status_code=400,
)
if not auth_code:
if not callback_response.code:
# Resolve future with exception if provided
if response_future and not response_future.done():
response_future.set_exception(
@ -243,7 +261,9 @@ def create_oauth_callback_server(
# Success case
if response_future and not response_future.done():
response_future.set_result((auth_code, state))
response_future.set_result(
(callback_response.code, callback_response.state)
)
return HTMLResponse(
create_callback_html("FastMCP OAuth login complete!", server_url=server_url)

View file

@ -25,8 +25,7 @@ from fastmcp.server.auth.auth import (
# Default expiration times (in seconds)
DEFAULT_AUTH_CODE_EXPIRY_SECONDS = 5 * 60 # 5 minutes
DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS = 60 * 60 # 1 hour
# Refresh tokens often have longer or no expiry; let's make them non-expiring for simplicity
DEFAULT_REFRESH_TOKEN_EXPIRY_SECONDS = None
DEFAULT_REFRESH_TOKEN_EXPIRY_SECONDS = None # No expiry
class InMemoryOAuthProvider(OAuthProvider):