fastmcp/fastmcp_slim/fastmcp/server/mixins/transport.py

"""Transport-related methods for FastMCP Server."""

from __future__ import annotations

import socket
from collections.abc import Awaitable, Callable
from functools import partial
from typing import TYPE_CHECKING, Any, Literal

import anyio
import uvicorn
from mcp.server.lowlevel.server import NotificationOptions
from mcp.server.stdio import stdio_server
from mcp.server.streamable_http import EventStore
from starlette.middleware import Middleware as ASGIMiddleware
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import BaseRoute, Route

import fastmcp
from fastmcp.server.http import (
    HostOriginProtection,
    StarletteWithLifespan,
    _is_loopback_host,
    create_sse_app,
    create_streamable_http_app,
)
from fastmcp.server.providers.base import Provider
from fastmcp.server.providers.fastmcp_provider import FastMCPProvider
from fastmcp.server.providers.wrapped_provider import _WrappedProvider
from fastmcp.utilities.logging import get_logger, temporary_log_level

if TYPE_CHECKING:
    from fastmcp.server.server import FastMCP, Transport

logger = get_logger(__name__)


def _format_host_for_url(host: str) -> str:
    """Format a host for inclusion in a URL, bracketing IPv6 addresses.

    A bare IPv6 address like ``::1`` must be wrapped in brackets when placed
    before a ``:port`` suffix, otherwise the result (``http://::1:8000``) is an
    invalid URL. Hostnames and IPv4 addresses are returned unchanged, as are
    addresses that are already bracketed.
    """
    if ":" in host and not host.startswith("["):
        return f"[{host}]"
    return host


def _resolve_allowed_hosts_for_run(
    *,
    host: str,
    host_origin_protection: HostOriginProtection,
    allowed_hosts: list[str] | None,
    configured_allowed_hosts: list[str] | None,
) -> list[str] | None:
    if allowed_hosts is not None:
        return allowed_hosts

    if host_origin_protection == "auto" and _is_loopback_host(host):
        return [*(configured_allowed_hosts or []), host]

    return configured_allowed_hosts


class TransportMixin:
    """Mixin providing transport-related methods for FastMCP.

    Includes HTTP/stdio/SSE transport handling and custom HTTP routes.
    """

    async def run_async(
        self: FastMCP,
        transport: Transport | None = None,
        show_banner: bool | None = None,
        **transport_kwargs: Any,
    ) -> None:
        """Run the FastMCP server asynchronously.

        Args:
            transport: Transport protocol to use ("stdio", "http", "sse", or "streamable-http")
            show_banner: Whether to display the server banner. If None, uses the
                FASTMCP_SHOW_SERVER_BANNER setting (default: True).
        """
        if show_banner is None:
            show_banner = fastmcp.settings.show_server_banner
        if transport is None:
            transport = fastmcp.settings.transport
        if transport not in {"stdio", "http", "sse", "streamable-http"}:
            raise ValueError(f"Unknown transport: {transport}")

        if transport == "stdio":
            await self.run_stdio_async(
                show_banner=show_banner,
                **transport_kwargs,
            )
        elif transport in {"http", "sse", "streamable-http"}:
            await self.run_http_async(
                transport=transport,
                show_banner=show_banner,
                **transport_kwargs,
            )
        else:
            raise ValueError(f"Unknown transport: {transport}")

    def run(
        self: FastMCP,
        transport: Transport | None = None,
        show_banner: bool | None = None,
        **transport_kwargs: Any,
    ) -> None:
        """Run the FastMCP server. Note this is a synchronous function.

        Args:
            transport: Transport protocol to use ("http", "stdio", "sse", or "streamable-http")
            show_banner: Whether to display the server banner. If None, uses the
                FASTMCP_SHOW_SERVER_BANNER setting (default: True).
        """

        anyio.run(
            partial(
                self.run_async,
                transport,
                show_banner=show_banner,
                **transport_kwargs,
            )
        )

    def custom_route(
        self: FastMCP,
        path: str,
        methods: list[str],
        name: str | None = None,
        include_in_schema: bool = True,
    ) -> Callable[
        [Callable[[Request], Awaitable[Response]]],
        Callable[[Request], Awaitable[Response]],
    ]:
        """
        Decorator to register a custom HTTP route on the FastMCP server.

        Allows adding arbitrary HTTP endpoints outside the standard MCP protocol,
        which can be useful for OAuth callbacks, health checks, or admin APIs.
        The handler function must be an async function that accepts a Starlette
        Request and returns a Response.

        Args:
            path: URL path for the route (e.g., "/auth/callback")
            methods: List of HTTP methods to support (e.g., ["GET", "POST"])
            name: Optional name for the route (to reference this route with
                Starlette's reverse URL lookup feature)
            include_in_schema: Whether to include in OpenAPI schema, defaults to True

        Example:
            Register a custom HTTP route for a health check endpoint:
            ```python
            @server.custom_route("/health", methods=["GET"])
            async def health_check(request: Request) -> Response:
                return JSONResponse({"status": "ok"})
            ```
        """

        def decorator(
            fn: Callable[[Request], Awaitable[Response]],
        ) -> Callable[[Request], Awaitable[Response]]:
            self._additional_http_routes.append(
                Route(
                    path,
                    endpoint=fn,
                    methods=methods,
                    name=name,
                    include_in_schema=include_in_schema,
                )
            )
            return fn

        return decorator

    def _get_additional_http_routes(self: FastMCP) -> list[BaseRoute]:
        """Get all additional HTTP routes including from mounted servers.

        Collects custom HTTP routes registered via ``@server.custom_route()``
        from this server **and** from any FastMCP servers reachable through
        mounted providers (recursively).  This ensures that routes defined on
        a child server are forwarded to the parent's HTTP app when using
        ``server.mount(child)``.

        Note:
            When path collisions occur between a parent and a mounted child,
            the parent's routes take precedence because they appear first in
            the returned list.

        Returns:
            List of Starlette Route objects
        """
        routes: list[BaseRoute] = list(self._additional_http_routes)

        def _unwrap_provider(provider: Provider) -> Provider:
            """Unwrap _WrappedProvider layers to find the inner provider."""
            while isinstance(provider, _WrappedProvider):
                provider = provider._inner
            return provider

        for provider in self.providers:
            inner = _unwrap_provider(provider)
            if isinstance(inner, FastMCPProvider):
                # Recurse into the mounted server to collect its routes
                # (and any routes from servers mounted on *it*).
                routes.extend(inner.server._get_additional_http_routes())

        return routes

    async def run_stdio_async(
        self: FastMCP,
        show_banner: bool = True,
        log_level: str | None = None,
        stateless: bool = False,
    ) -> None:
        """Run the server using stdio transport.

        Args:
            show_banner: Whether to display the server banner
            log_level: Log level for the server
            stateless: Whether to run in stateless mode (no session initialization)
        """
        from fastmcp.server.context import reset_transport, set_transport

        # Display server banner
        if show_banner:
            from fastmcp.utilities.cli import log_server_banner

            log_server_banner(server=self)

        token = set_transport("stdio")
        try:
            with temporary_log_level(log_level):
                async with self._lifespan_manager():
                    async with stdio_server() as (read_stream, write_stream):
                        mode = " (stateless)" if stateless else ""
                        logger.info(
                            f"Starting MCP server {self.name!r} with transport 'stdio'{mode}"
                        )

                        await self._mcp_server.run(
                            read_stream,
                            write_stream,
                            self._mcp_server.create_initialization_options(
                                notification_options=NotificationOptions(
                                    tools_changed=True
                                ),
                            ),
                        )
        finally:
            reset_transport(token)

    async def run_http_async(
        self: FastMCP,
        show_banner: bool = True,
        transport: Literal["http", "streamable-http", "sse"] = "http",
        host: str | None = None,
        port: int | None = None,
        log_level: str | None = None,
        path: str | None = None,
        uvicorn_config: dict[str, Any] | None = None,
        middleware: list[ASGIMiddleware] | None = None,
        json_response: bool | None = None,
        stateless_http: bool | None = None,
        stateless: bool | None = None,
        host_origin_protection: HostOriginProtection | None = None,
        allowed_hosts: list[str] | None = None,
        allowed_origins: list[str] | None = None,
        sockets: list[socket.socket] | None = None,
    ) -> None:
        """Run the server using HTTP transport.

        Args:
            transport: Transport protocol to use - "http" (default), "streamable-http", or "sse"
            host: Host address to bind to (defaults to settings.host)
            port: Port to bind to (defaults to settings.port)
            log_level: Log level for the server (defaults to settings.log_level)
            path: Path for the endpoint (defaults to settings.streamable_http_path or settings.sse_path)
            uvicorn_config: Additional configuration for the Uvicorn server
            middleware: A list of middleware to apply to the app
            json_response: Whether to use JSON response format (defaults to settings.json_response)
            stateless_http: Whether to use stateless HTTP (defaults to settings.stateless_http)
            stateless: Alias for stateless_http for CLI consistency
            host_origin_protection: Whether to validate Host and Origin headers
                before requests reach the MCP endpoint. Defaults to
                settings.http_host_origin_protection. "auto" protects
                localhost-bound servers and explicit host/origin allowlists.
            allowed_hosts: Additional hostnames that may appear in the Host header.
            allowed_origins: Additional browser origins trusted by the request guard.
                Configure CORS separately when browser JavaScript must read
                cross-origin responses.
            sockets: Pre-bound sockets to pass to Uvicorn
        """
        # Allow stateless as alias for stateless_http
        if stateless is not None and stateless_http is None:
            stateless_http = stateless

        # Resolve from settings/env var if not explicitly set
        if stateless_http is None:
            stateless_http = fastmcp.settings.stateless_http

        # SSE doesn't support stateless mode
        if stateless_http and transport == "sse":
            raise ValueError("SSE transport does not support stateless mode")

        host = host if host is not None else fastmcp.settings.host
        port = port if port is not None else fastmcp.settings.port
        resolved_host_origin_protection = (
            host_origin_protection
            if host_origin_protection is not None
            else fastmcp.settings.http_host_origin_protection
        )
        resolved_allowed_hosts = _resolve_allowed_hosts_for_run(
            host=host,
            host_origin_protection=resolved_host_origin_protection,
            allowed_hosts=allowed_hosts,
            configured_allowed_hosts=fastmcp.settings.http_allowed_hosts,
        )
        default_log_level_to_use = (
            log_level if log_level is not None else fastmcp.settings.log_level
        ).lower()

        app = self.http_app(
            path=path,
            transport=transport,
            middleware=middleware,
            json_response=json_response,
            stateless_http=stateless_http,
            host_origin_protection=resolved_host_origin_protection,
            allowed_hosts=resolved_allowed_hosts,
            allowed_origins=allowed_origins,
        )

        # Display server banner
        if show_banner:
            from fastmcp.utilities.cli import log_server_banner

            log_server_banner(server=self)
        uvicorn_config_from_user = uvicorn_config or {}

        config_kwargs: dict[str, Any] = {
            "timeout_graceful_shutdown": 2,
            "lifespan": "on",
            "ws": "websockets-sansio",
        }
        config_kwargs.update(uvicorn_config_from_user)

        if "log_config" not in config_kwargs and "log_level" not in config_kwargs:
            config_kwargs["log_level"] = default_log_level_to_use

        with temporary_log_level(log_level):
            async with self._lifespan_manager():
                config = uvicorn.Config(app, host=host, port=port, **config_kwargs)
                server = uvicorn.Server(config)
                path = getattr(app.state, "path", "").lstrip("/")
                mode = " (stateless)" if stateless_http else ""
                display_host = _format_host_for_url(host)
                logger.info(
                    f"Starting MCP server {self.name!r} with transport {transport!r}{mode} on http://{display_host}:{port}/{path}"
                )

                if sockets is not None:
                    await server.serve(sockets=sockets)
                else:
                    await server.serve()

    def http_app(
        self: FastMCP,
        path: str | None = None,
        middleware: list[ASGIMiddleware] | None = None,
        json_response: bool | None = None,
        stateless_http: bool | None = None,
        transport: Literal["http", "streamable-http", "sse"] = "http",
        event_store: EventStore | None = None,
        retry_interval: int | None = None,
        host_origin_protection: HostOriginProtection | None = None,
        allowed_hosts: list[str] | None = None,
        allowed_origins: list[str] | None = None,
        session_idle_timeout: float | None = None,
    ) -> StarletteWithLifespan:
        """Create a Starlette app using the specified HTTP transport.

        Args:
            path: The path for the HTTP endpoint
            middleware: A list of middleware to apply to the app
            json_response: Whether to use JSON response format
            stateless_http: Whether to use stateless mode (new transport per request)
            transport: Transport protocol to use - "http", "streamable-http", or "sse"
            event_store: Optional event store for SSE polling/resumability. When set,
                enables clients to reconnect and resume receiving events after
                server-initiated disconnections. Only used with streamable-http transport.
            retry_interval: Optional retry interval in milliseconds for SSE polling.
                Controls how quickly clients should reconnect after server-initiated
                disconnections. Requires event_store to be set. Only used with
                streamable-http transport.
            host_origin_protection: Whether to validate Host and Origin headers
                before requests reach the MCP endpoint. Defaults to
                settings.http_host_origin_protection. "auto" protects
                localhost-bound servers and explicit host/origin allowlists.
            allowed_hosts: Additional hostnames that may appear in the Host header.
            allowed_origins: Additional browser origins trusted by the request guard.
                Configure CORS separately when browser JavaScript must read
                cross-origin responses.
            session_idle_timeout: Maximum time in seconds a streamable-HTTP
                session may remain idle before it is terminated. When None,
                falls back to the ``http_session_idle_timeout`` setting.

        Returns:
            A Starlette application configured with the specified transport
        """

        if transport in ("streamable-http", "http"):
            return create_streamable_http_app(
                server=self,
                streamable_http_path=path
                if path is not None
                else fastmcp.settings.streamable_http_path,
                event_store=event_store,
                retry_interval=retry_interval,
                auth=self.auth,
                json_response=(
                    json_response
                    if json_response is not None
                    else fastmcp.settings.json_response
                ),
                stateless_http=(
                    stateless_http
                    if stateless_http is not None
                    else fastmcp.settings.stateless_http
                ),
                debug=fastmcp.settings.debug,
                middleware=middleware,
                host_origin_protection=(
                    host_origin_protection
                    if host_origin_protection is not None
                    else fastmcp.settings.http_host_origin_protection
                ),
                allowed_hosts=(
                    allowed_hosts
                    if allowed_hosts is not None
                    else fastmcp.settings.http_allowed_hosts
                ),
                allowed_origins=(
                    allowed_origins
                    if allowed_origins is not None
                    else fastmcp.settings.http_allowed_origins
                ),
                session_idle_timeout=(
                    session_idle_timeout
                    if session_idle_timeout is not None
                    else fastmcp.settings.http_session_idle_timeout
                ),
            )
        elif transport == "sse":
            return create_sse_app(
                server=self,
                message_path=fastmcp.settings.message_path,
                sse_path=path if path is not None else fastmcp.settings.sse_path,
                auth=self.auth,
                debug=fastmcp.settings.debug,
                middleware=middleware,
            )
        else:
            raise ValueError(f"Unknown transport: {transport}")