fastmcp 2.5.1__py3-none-any.whl → 2.6.0__py3-none-any.whl

fastmcp/client/client.py

@@ -1,14 +1,17 @@
+import asyncio
 import datetime
 from contextlib import AsyncExitStack, asynccontextmanager
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, Generic, Literal, cast, overload
 
 import anyio
+import httpx
 import mcp.types
 from exceptiongroup import catch
 from mcp import ClientSession
 from pydantic import AnyUrl
 
+import fastmcp
 from fastmcp.client.logging import (
     LogHandler,
     MessageHandler,
@@ -27,10 +30,22 @@ from fastmcp.server import FastMCP
 from fastmcp.utilities.exceptions import get_catch_handlers
 from fastmcp.utilities.mcp_config import MCPConfig
 
-from .transports import ClientTransport, SessionKwargs, infer_transport
+from .transports import (
+    ClientTransportT,
+    FastMCP1Server,
+    FastMCPTransport,
+    MCPConfigTransport,
+    NodeStdioTransport,
+    PythonStdioTransport,
+    SessionKwargs,
+    SSETransport,
+    StreamableHttpTransport,
+    infer_transport,
+)
 
 __all__ = [
     "Client",
+    "SessionKwargs",
     "RootsHandler",
     "RootsList",
     "LogHandler",
@@ -40,13 +55,13 @@ __all__ = [
 ]
 
 
-class Client:
+class Client(Generic[ClientTransportT]):
     """
     MCP client that delegates connection management to a Transport instance.
 
     The Client class is responsible for MCP protocol logic, while the Transport
-    handles connection establishment and management. Client provides methods
-    for working with resources, prompts, tools and other MCP capabilities.
+    handles connection establishment and management. Client provides methods for
+    working with resources, prompts, tools and other MCP capabilities.
 
     Args:
         transport: Connection source specification, which can be:
@@ -62,24 +77,60 @@ class Client:
         message_handler: Optional handler for protocol messages
         progress_handler: Optional handler for progress notifications
         timeout: Optional timeout for requests (seconds or timedelta)
+        init_timeout: Optional timeout for initial connection (seconds or timedelta).
+            Set to 0 to disable. If None, uses the value in the FastMCP global settings.
 
     Examples:
-        ```python
-        # Connect to FastMCP server
-        client = Client("http://localhost:8080")
+        ```python # Connect to FastMCP server client =
+        Client("http://localhost:8080")
 
         async with client:
-            # List available resources
-            resources = await client.list_resources()
+            # List available resources resources = await client.list_resources()
 
-            # Call a tool
-            result = await client.call_tool("my_tool", {"param": "value"})
+            # Call a tool result = await client.call_tool("my_tool", {"param":
+            "value"})
         ```
     """
 
+    @overload
+    def __new__(
+        cls,
+        transport: ClientTransportT,
+        **kwargs: Any,
+    ) -> "Client[ClientTransportT]": ...
+
+    @overload
+    def __new__(
+        cls, transport: AnyUrl, **kwargs
+    ) -> "Client[SSETransport|StreamableHttpTransport]": ...
+
+    @overload
+    def __new__(
+        cls, transport: FastMCP | FastMCP1Server, **kwargs
+    ) -> "Client[FastMCPTransport]": ...
+
+    @overload
+    def __new__(
+        cls, transport: Path, **kwargs
+    ) -> "Client[PythonStdioTransport|NodeStdioTransport]": ...
+
+    @overload
+    def __new__(
+        cls, transport: MCPConfig | dict[str, Any], **kwargs
+    ) -> "Client[MCPConfigTransport]": ...
+
+    @overload
+    def __new__(
+        cls, transport: str, **kwargs
+    ) -> "Client[PythonStdioTransport|NodeStdioTransport|SSETransport|StreamableHttpTransport]": ...
+
+    def __new__(cls, transport, **kwargs) -> "Client":
+        instance = super().__new__(cls)
+        return instance
+
     def __init__(
         self,
-        transport: ClientTransport
+        transport: ClientTransportT
         | FastMCP
         | AnyUrl
         | Path
@@ -93,11 +144,12 @@ class Client:
         message_handler: MessageHandler | None = None,
         progress_handler: ProgressHandler | None = None,
         timeout: datetime.timedelta | float | int | None = None,
+        init_timeout: datetime.timedelta | float | int | None = None,
+        auth: httpx.Auth | Literal["oauth"] | str | None = None,
     ):
-        self.transport = infer_transport(transport)
-        self._session: ClientSession | None = None
-        self._exit_stack: AsyncExitStack | None = None
-        self._nesting_counter: int = 0
+        self.transport = cast(ClientTransportT, infer_transport(transport))
+        if auth is not None:
+            self.transport._set_auth(auth)
         self._initialize_result: mcp.types.InitializeResult | None = None
 
         if log_handler is None:
@@ -111,6 +163,17 @@ class Client:
         if isinstance(timeout, int | float):
             timeout = datetime.timedelta(seconds=timeout)
 
+        # handle init handshake timeout
+        if init_timeout is None:
+            init_timeout = fastmcp.settings.settings.client_init_timeout
+        if isinstance(init_timeout, datetime.timedelta):
+            init_timeout = init_timeout.total_seconds()
+        elif not init_timeout:
+            init_timeout = None
+        else:
+            init_timeout = float(init_timeout)
+        self._init_timeout = init_timeout
+
         self._session_kwargs: SessionKwargs = {
             "sampling_callback": None,
             "list_roots_callback": None,
@@ -127,6 +190,15 @@ class Client:
                 sampling_handler
             )
 
+        # session context management
+        self._session: ClientSession | None = None
+        self._exit_stack: AsyncExitStack | None = None
+        self._nesting_counter: int = 0
+        self._context_lock = anyio.Lock()
+        self._session_task: asyncio.Task | None = None
+        self._ready_event = anyio.Event()
+        self._stop_event = anyio.Event()
+
     @property
     def session(self) -> ClientSession:
         """Get the current active session. Raises RuntimeError if not connected."""
@@ -134,6 +206,7 @@ class Client:
             raise RuntimeError(
                 "Client is not connected. Use the 'async with client:' context manager first."
             )
+
         return self._session
 
     @property
@@ -168,40 +241,83 @@ class Client:
                 self._session = session
                 # Initialize the session
                 try:
-                    with anyio.fail_after(1):
+                    with anyio.fail_after(self._init_timeout):
                         self._initialize_result = await self._session.initialize()
                     yield
+                except anyio.ClosedResourceError:
+                    raise RuntimeError("Server session was closed unexpectedly")
                 except TimeoutError:
                     raise RuntimeError("Failed to initialize server session")
                 finally:
-                    self._exit_stack = None
                     self._session = None
                     self._initialize_result = None
 
     async def __aenter__(self):
-        if self._nesting_counter == 0:
-            # Create exit stack to manage both context managers
-            stack = AsyncExitStack()
-            await stack.__aenter__()
-
-            await stack.enter_async_context(self._context_manager())
-
-            self._exit_stack = stack
-
-        self._nesting_counter += 1
-
+        await self._connect()
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        self._nesting_counter -= 1
+        await self._disconnect()
+
+    async def _connect(self):
+        # ensure only one session is running at a time to avoid race conditions
+        async with self._context_lock:
+            need_to_start = self._session_task is None or self._session_task.done()
+            if need_to_start:
+                self._stop_event = anyio.Event()
+                self._ready_event = anyio.Event()
+                self._session_task = asyncio.create_task(self._session_runner())
+            await self._ready_event.wait()
+            self._nesting_counter += 1
+        return self
 
-        if self._nesting_counter == 0:
-            # Exit the stack which will handle cleaning up the session
-            if self._exit_stack is not None:
-                try:
-                    await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
-                finally:
-                    self._exit_stack = None
+    async def _disconnect(self, force: bool = False):
+        # ensure only one session is running at a time to avoid race conditions
+        async with self._context_lock:
+            # if we are forcing a disconnect, reset the nesting counter
+            if force:
+                self._nesting_counter = 0
+
+            # otherwise decrement to check if we are done nesting
+            else:
+                self._nesting_counter = max(0, self._nesting_counter - 1)
+
+            # if we are still nested, return
+            if self._nesting_counter > 0:
+                return
+
+            # stop the active seesion
+            if self._session_task is None:
+                return
+            self._stop_event.set()
+            runner_task = self._session_task
+            self._session_task = None
+
+        # wait for the session to finish
+        if runner_task:
+            await runner_task
+
+        # Reset for future reconnects
+        self._stop_event = anyio.Event()
+        self._ready_event = anyio.Event()
+        self._session = None
+        self._initialize_result = None
+
+    async def _session_runner(self):
+        async with AsyncExitStack() as stack:
+            try:
+                await stack.enter_async_context(self._context_manager())
+                # Session/context is now ready
+                self._ready_event.set()
+                # Wait until disconnect/stop is requested
+                await self._stop_event.wait()
+            finally:
+                # On exit, ensure ready event is set (idempotent)
+                self._ready_event.set()
+
+    async def close(self):
+        await self._disconnect(force=True)
+        await self.transport.close()
 
     # --- MCP Client Methods ---
 

fastmcp/client/oauth_callback.py

@@ -0,0 +1,310 @@
+"""
+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
+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
+
+from fastmcp.utilities.http import find_available_port
+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>
+    """
+
+
+@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",
+    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: Request):
+        """Handle OAuth callback requests with proper HTML responses."""
+        query_params = dict(request.query_params)
+        callback_response = CallbackResponse.from_dict(query_params)
+
+        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: {callback_response.error} - {error_desc}"
+                    )
+                )
+
+            return HTMLResponse(
+                create_callback_html(
+                    f"FastMCP OAuth Error: {callback_response.error}<br>{error_desc}",
+                    is_success=False,
+                ),
+                status_code=400,
+            )
+
+        if not callback_response.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(
+                    "FastMCP 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(
+                (callback_response.code, callback_response.state)
+            )
+
+        return HTMLResponse(
+            create_callback_html("FastMCP 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,
+    )