Add callback server

src/fastmcp/client/auth.py

@@ -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(

src/fastmcp/client/oauth_callback.py

@@ -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,
+    )