mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 21:44:18 +02:00
Add callback server
This commit is contained in:
parent
67c36b0c11
commit
c438987d4e
2 changed files with 343 additions and 102 deletions
|
|
@ -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(
|
||||
|
|
|
|||
296
src/fastmcp/client/oauth_callback.py
Normal file
296
src/fastmcp/client/oauth_callback.py
Normal 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,
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue