app/backend/fluksio/mcp/http.py

"""Mounting the MCP endpoint, and deciding whose tokens it accepts.

The transport is streamable HTTP, stateless, one JSON response per call: no
sticky sessions to route and nothing for a proxy to buffer. The SDK serves both
``/mcp`` and the protected-resource metadata beside it, which is why the sub-app
is mounted at the root rather than under ``/mcp``.

An enrolled instance also answers at ``/i/{id}/mcp`` on its hub, so an agent
that came through the tunnel is shown the hub's addresses instead.
"""

from __future__ import annotations

import logging
from urllib.parse import urlparse

import jwt
from fastapi.concurrency import run_in_threadpool
from mcp.server.auth.provider import AccessToken
from mcp.server.auth.routes import build_resource_metadata_url
from mcp.server.auth.settings import AuthSettings
from mcp.server.transport_security import TransportSecuritySettings
from pydantic import AnyHttpUrl
from starlette.datastructures import Headers
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from fluksio.api.deps import oauth_client_lives
from fluksio.cloud import config as cloud_config
from fluksio.core import security
from fluksio.core.config import settings
from fluksio.mcp import server

logger = logging.getLogger(__name__)


class _JWTVerifier:
    """Accept only tokens minted for the MCP channel, by agents still registered.

    A perfectly valid browser token is refused: it was issued for a person's
    session, and honouring it here would make agent traffic indistinguishable
    from theirs.

    The client row is looked up because the token itself cannot be withdrawn:
    it is stateless and good until ``MCP_TOKEN_EXPIRE_MINUTES`` runs out, so
    the registration it names is the thing revoking an agent actually removes.
    Refusing at the door rather than leaving it to the API the tools call means
    a revoked agent gets the 401 that sends it back to authorize, instead of a
    tool listing that works and a tool call that does not.
    """

    async def verify_token(self, token: str) -> AccessToken | None:
        try:
            payload = security.decode_oauth_token(token)
        except jwt.InvalidTokenError as exc:
            logger.debug("MCP token rejected: %s", exc)
            return None
        if payload.get("mcp") is not True:
            logger.debug("MCP token rejected: not an MCP-channel token")
            return None
        client_id = str(payload.get("client_id", ""))
        # Off the event loop: the lookup is SQLite, like every other read this
        # process makes, and the session it opens is blocking.
        if not await run_in_threadpool(oauth_client_lives, client_id):
            logger.debug(
                "MCP token rejected: agent %s is no longer registered", client_id
            )
            return None
        return AccessToken(
            token=token,
            client_id=client_id,
            scopes=[security.MCP_SCOPE],
            subject=str(payload.get("sub", "")),
            # `decode_oauth_token` has checked the audience; the SDK compares
            # this against its own resource URL as well.
            resource=settings.mcp_resource,
        )


class _AndTheHub(list[str]):
    """The hosts named here, plus the hub's while this instance is enrolled.

    A call through the tunnel is addressed to the hub, as its Host says. The
    hub is looked up per request, since enrolment comes and goes while the
    endpoint is up. Its origin is not needed: the hub never forwards one.
    """

    def __contains__(self, host: object) -> bool:
        config = cloud_config.load()
        hub = urlparse(config.issuer).netloc if config is not None else None
        return super().__contains__(host) or (hub is not None and host == hub)


def _transport_security() -> TransportSecuritySettings:
    """Which Host headers to trust.

    The SDK defaults to localhost only and answers 421 to anything arriving
    through a reverse proxy under the real hostname, so the deployment's own
    host is named here rather than the protection being switched off.
    """
    host = urlparse(settings.oauth_issuer).netloc
    allowed_hosts = [host, f"{host}:*"]
    allowed_origins = [f"https://{host}", f"http://{host}"]
    if settings.ENVIRONMENT == "local":
        # The test client and a direct uvicorn run do not go through Traefik.
        allowed_hosts += [
            "localhost",
            "localhost:*",
            "127.0.0.1",
            "127.0.0.1:*",
            "testserver",
        ]
        allowed_origins += ["http://localhost", "http://127.0.0.1"]
    # Unvalidated, since validating would copy the hosts into a plain list.
    return TransportSecuritySettings.model_construct(
        enable_dns_rebinding_protection=True,
        allowed_hosts=_AndTheHub(allowed_hosts),
        allowed_origins=allowed_origins,
    )


class _ThroughThePortal:
    """Show an agent that came through the hub the hub's addresses.

    The SDK builds its metadata once, for the API host, and names it in every
    401. An agent at ``/i/{id}/mcp`` on the hub cannot reach that host, so its
    discovery has to lead to the hub's documents instead.
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        base = None
        if scope["type"] == "http":
            base = cloud_config.portal_base(Headers(scope=scope))
        if base is None:
            await self.app(scope, receive, send)
            return
        if scope["path"] == "/.well-known/oauth-protected-resource/mcp":
            document = {
                "resource": f"{base}/mcp",
                "authorization_servers": [base],
                "bearer_methods_supported": ["header"],
            }
            headers = {"Cache-Control": "public, max-age=3600"}
            await JSONResponse(document, headers=headers)(scope, receive, send)
            return
        here = str(build_resource_metadata_url(AnyHttpUrl(settings.mcp_resource)))
        there = str(build_resource_metadata_url(AnyHttpUrl(f"{base}/mcp")))

        async def send_theirs(message: Message) -> None:
            if message["type"] == "http.response.start":
                message["headers"] = [
                    (
                        key,
                        value.replace(here.encode(), there.encode())
                        if key == b"www-authenticate"
                        else value,
                    )
                    for key, value in message["headers"]
                ]
            await send(message)

        await self.app(scope, receive, send_theirs)


def build_http_app(app: object) -> ASGIApp:
    """The MCP sub-app, wired to this app's API and its OAuth server."""
    server.mcp.settings.auth = AuthSettings(
        issuer_url=AnyHttpUrl(settings.oauth_issuer),
        resource_server_url=AnyHttpUrl(settings.mcp_resource),
        required_scopes=None,
        validate_token_resource=True,
    )
    # The SDK only takes a verifier through its constructor, and the instance is
    # module-level so the tools can be declared at import time.
    server.mcp._token_verifier = _JWTVerifier()  # noqa: SLF001

    server.connect(app)
    return _ThroughThePortal(
        server.mcp.streamable_http_app(
            stateless_http=True,
            json_response=True,
            transport_security=_transport_security(),
        )
    )