Add callback server

This commit is contained in:
Jeremiah Lowin 2025-05-30 13:03:01 -04:00
commit c438987d4e
2 changed files with 343 additions and 102 deletions

View file

@ -2,13 +2,11 @@ from __future__ import annotations
import asyncio
import json
import socket
import webbrowser
from pathlib import Path
from typing import Any, Literal, cast
from typing import Any, Literal
from urllib.parse import urljoin, urlparse
import anyio
import httpx
from mcp.client.auth import OAuthClientProvider as _MCPOAuthClientProvider
from mcp.client.auth import TokenStorage
@ -21,11 +19,11 @@ from mcp.shared.auth import (
OAuthMetadata as _MCPServerOAuthMetadata,
)
from pydantic import AnyHttpUrl, ValidationError
from starlette.applications import Starlette
from starlette.responses import PlainTextResponse
from starlette.routing import Route
from uvicorn import Config, Server
from fastmcp.client.oauth_callback import (
create_oauth_callback_server,
find_available_port,
)
from fastmcp.settings import settings as fastmcp_global_settings
from fastmcp.utilities.logging import get_logger
@ -186,11 +184,8 @@ class FileTokenStorage(TokenStorage):
def clear_cache(self) -> None:
"""Clear all cached data for this server."""
# Use explicit literals to satisfy type checker
for file_type in [
cast(Literal["client_info", "tokens"], "client_info"),
cast(Literal["client_info", "tokens"], "tokens"),
]:
file_types: list[Literal["client_info", "tokens"]] = ["client_info", "tokens"]
for file_type in file_types:
path = self._get_file_path(file_type)
path.unlink(missing_ok=True)
logger.info(f"Cleared OAuth cache for {self.get_base_url(self.server_url)}")
@ -253,95 +248,13 @@ class FileTokenStorage(TokenStorage):
if not cache_dir.exists():
return
# Use explicit literals to satisfy type checker
for file_type in [
cast(Literal["client_info", "tokens"], "client_info"),
cast(Literal["client_info", "tokens"], "tokens"),
]:
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.")
def find_available_port() -> int:
"""Find an available port by letting the OS assign one."""
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
async def _get_redirect_callback(
port: int, path: str = "/callback", timeout: float = 300.0
) -> tuple[str, str | None]:
"""
Start a temporary server to handle OAuth redirect and return auth code and state.
Returns:
Tuple of (authorization_code, state)
"""
response_future = asyncio.get_running_loop().create_future()
async def callback_handler(request):
if not response_future.done():
query_params = dict(request.query_params)
auth_code = query_params.get("code")
state = query_params.get("state")
error = query_params.get("error")
if error:
error_desc = query_params.get("error_description", "Unknown error")
response_future.set_exception(
RuntimeError(f"OAuth error: {error} - {error_desc}")
)
return PlainTextResponse(
f"❌ OAuth Error: {error}\n{error_desc}\nYou can close this tab.",
status_code=400,
)
if not auth_code:
response_future.set_exception(
RuntimeError("OAuth callback missing authorization code")
)
return PlainTextResponse(
"❌ OAuth Error: No authorization code received.\nYou can close this tab.",
status_code=400,
)
response_future.set_result((auth_code, state))
return PlainTextResponse(
"✅ FastMCP OAuth login complete!\nYou can close this tab now."
)
return PlainTextResponse("Callback already processed. You can close this tab.")
server = Server(
Config(
app=Starlette(routes=[Route(path, callback_handler)]),
host="127.0.0.1",
port=port,
lifespan="off",
log_level="warning",
)
)
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:{port}{path}"
)
try:
with anyio.fail_after(timeout):
auth_code, state = await response_future
return auth_code, state
except TimeoutError:
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
tg.cancel_scope.cancel()
async def discover_oauth_metadata(
server_base_url: str, httpx_kwargs: dict[str, Any] | None = None
) -> _MCPServerOAuthMetadata | None:
@ -417,12 +330,16 @@ def OAuth(
"""
Create an OAuthClientProvider for an MCP server.
This is intended to be provided to the `auth` parameter of an
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_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
Returns:
OAuthClientProvider
@ -460,7 +377,35 @@ def OAuth(
async def callback_handler() -> tuple[str, str | None]:
"""Handle OAuth callback and return (auth_code, state)."""
return await _get_redirect_callback(port=redirect_port)
# Create a future to capture the OAuth response
response_future = asyncio.get_running_loop().create_future()
# Create server with the future
server = create_oauth_callback_server(
port=redirect_port,
server_url=server_base_url,
response_future=response_future,
)
# Run server until response is received with timeout logic
import anyio
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}"
)
try:
with anyio.fail_after(300.0): # 5 minute timeout
auth_code, state = await response_future
return auth_code, state
except TimeoutError:
raise TimeoutError("OAuth callback timed out after 300 seconds")
finally:
server.should_exit = True
await asyncio.sleep(0.1) # Allow server to shutdown gracefully
tg.cancel_scope.cancel()
# Create OAuth provider
oauth_provider = OAuthClientProvider(

View file

@ -0,0 +1,296 @@
"""
OAuth callback server for handling authorization code flows.
This module provides a reusable callback server that can handle OAuth redirects
and display styled responses to users.
"""
from __future__ import annotations
import asyncio
import socket
from starlette.applications import Starlette
from starlette.responses import HTMLResponse
from starlette.routing import Route
from uvicorn import Config, Server
from fastmcp.utilities.logging import get_logger
logger = get_logger(__name__)
def create_callback_html(
message: str,
is_success: bool = True,
title: str = "FastMCP OAuth",
server_url: str | None = None,
) -> str:
"""Create a styled HTML response for OAuth callbacks."""
status_emoji = "" if is_success else ""
status_color = "#10b981" if is_success else "#ef4444" # emerald-500 / red-500
# Add server info for success cases
server_info = ""
if is_success and server_url:
server_info = f"""
<div class="server-info">
Connected to: <strong>{server_url}</strong>
</div>
"""
return f"""
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{title}</title>
<style>
body {{
font-family: 'SF Mono', 'Monaco', 'Consolas', 'Roboto Mono', monospace;
margin: 0;
padding: 0;
min-height: 100vh;
display: flex;
align-items: center;
justify-content: center;
background: linear-gradient(135deg, #0f0f23 0%, #1a1a2e 25%, #16213e 50%, #0f0f23 100%);
color: #e2e8f0;
overflow: hidden;
}}
body::before {{
content: '';
position: fixed;
top: 0;
left: 0;
width: 100%;
height: 100%;
background:
radial-gradient(circle at 20% 80%, rgba(120, 119, 198, 0.1) 0%, transparent 50%),
radial-gradient(circle at 80% 20%, rgba(16, 185, 129, 0.1) 0%, transparent 50%),
radial-gradient(circle at 40% 40%, rgba(14, 165, 233, 0.1) 0%, transparent 50%);
pointer-events: none;
z-index: -1;
}}
.container {{
background: rgba(30, 41, 59, 0.9);
backdrop-filter: blur(10px);
border: 1px solid rgba(71, 85, 105, 0.3);
padding: 3rem 2rem;
border-radius: 1rem;
box-shadow:
0 25px 50px -12px rgba(0, 0, 0, 0.7),
0 0 0 1px rgba(255, 255, 255, 0.05),
inset 0 1px 0 0 rgba(255, 255, 255, 0.1);
text-align: center;
max-width: 500px;
margin: 1rem;
position: relative;
}}
.container::before {{
content: '';
position: absolute;
top: 0;
left: 0;
right: 0;
height: 1px;
background: linear-gradient(90deg, transparent, rgba(16, 185, 129, 0.5), transparent);
}}
.status-icon {{
font-size: 4rem;
margin-bottom: 1rem;
display: block;
filter: drop-shadow(0 0 20px currentColor);
}}
.message {{
font-size: 1.25rem;
line-height: 1.6;
color: {status_color};
margin-bottom: 1.5rem;
font-weight: 600;
text-shadow: 0 0 10px rgba({
"16, 185, 129" if is_success else "239, 68, 68"
}, 0.3);
}}
.server-info {{
background: rgba(6, 182, 212, 0.1);
border: 1px solid rgba(6, 182, 212, 0.3);
border-radius: 0.75rem;
padding: 1rem;
margin: 1rem 0;
font-size: 0.9rem;
color: #67e8f9;
font-family: 'SF Mono', 'Monaco', 'Consolas', 'Roboto Mono', monospace;
text-shadow: 0 0 10px rgba(103, 232, 249, 0.3);
}}
.server-info strong {{
color: #22d3ee;
font-weight: 700;
}}
.subtitle {{
font-size: 1rem;
color: #94a3b8;
margin-top: 1rem;
}}
.close-instruction {{
background: rgba(51, 65, 85, 0.8);
border: 1px solid rgba(71, 85, 105, 0.4);
border-radius: 0.75rem;
padding: 1rem;
margin-top: 1.5rem;
font-size: 0.9rem;
color: #cbd5e1;
font-family: 'SF Mono', 'Monaco', 'Consolas', 'Roboto Mono', monospace;
}}
@keyframes glow {{
0%, 100% {{ opacity: 1; }}
50% {{ opacity: 0.7; }}
}}
.status-icon {{
animation: glow 2s ease-in-out infinite;
}}
</style>
</head>
<body>
<div class="container">
<span class="status-icon">{status_emoji}</span>
<div class="message">{message}</div>
{server_info}
<div class="close-instruction">
You can safely close this tab now.
</div>
</div>
</body>
</html>
"""
def find_available_port() -> int:
"""Find an available port by letting the OS assign one."""
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def create_oauth_callback_server(
port: int,
callback_path: str = "/callback",
server_url: str | None = None,
response_future: asyncio.Future | None = None,
) -> Server:
"""
Create an OAuth callback server.
Args:
port: The port to run the server on
callback_path: The path to listen for OAuth redirects on
server_url: Optional server URL to display in success messages
response_future: Optional future to resolve when OAuth callback is received
Returns:
Configured uvicorn Server instance (not yet running)
"""
async def callback_handler(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")
if error:
error_desc = query_params.get("error_description", "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}")
)
return HTMLResponse(
create_callback_html(
f"OAuth Error: {error}<br>{error_desc}", is_success=False
),
status_code=400,
)
if not auth_code:
# Resolve future with exception if provided
if response_future and not response_future.done():
response_future.set_exception(
RuntimeError("OAuth callback missing authorization code")
)
return HTMLResponse(
create_callback_html(
"OAuth Error: No authorization code received", is_success=False
),
status_code=400,
)
# Success case
if response_future and not response_future.done():
response_future.set_result((auth_code, state))
return HTMLResponse(
create_callback_html("OAuth login complete!", server_url=server_url)
)
app = Starlette(routes=[Route(callback_path, callback_handler)])
return Server(
Config(
app=app,
host="127.0.0.1",
port=port,
lifespan="off",
log_level="warning",
)
)
if __name__ == "__main__":
"""Run a test server when executed directly."""
import webbrowser
import uvicorn
port = find_available_port()
print("🎭 OAuth Callback Test Server")
print("📍 Test URLs:")
print(f" Success: http://localhost:{port}/callback?code=test123&state=xyz")
print(
f" Error: http://localhost:{port}/callback?error=access_denied&error_description=User%20denied"
)
print(f" Missing: http://localhost:{port}/callback")
print("🛑 Press Ctrl+C to stop")
print()
# Create test server without future (just for testing HTML responses)
server = create_oauth_callback_server(
port=port, server_url="https://fastmcp-test-server.example.com"
)
# Open browser to success example
webbrowser.open(f"http://localhost:{port}/callback?code=test123&state=xyz")
# Run with uvicorn directly
uvicorn.run(
server.config.app,
host="127.0.0.1",
port=port,
log_level="warning",
access_log=False,
)