app/backend/tests/mcp/test_mcp_http.py

"""The MCP endpoint: who it lets in, and what happens once they are in."""

import asyncio
import json
import uuid
from collections.abc import Awaitable, Callable, Generator
from datetime import timedelta
from typing import Any

import httpx
import pytest
from fastapi.testclient import TestClient
from sqlmodel import Session

from fluksio.api.deps import CODE_IS_SUPERUSERS
from fluksio.core import security
from fluksio.core.config import settings
from fluksio.main import app
from fluksio.models import OAuthClient

MCP_HEADERS = {"Accept": "application/json, text/event-stream"}


def mcp_token(user_id: uuid.UUID, client_id: uuid.UUID) -> str:
    return security.create_oauth_access_token(user_id, client_id, timedelta(minutes=5))


@pytest.fixture
def agent(db: Session) -> Generator[uuid.UUID, None, None]:
    """A registered agent, because a token naming one that is gone is refused."""
    row = OAuthClient(client_name="Test agent", redirect_uris=[])
    db.add(row)
    db.commit()
    yield row.id
    left = db.get(OAuthClient, row.id)
    if left is not None:
        db.delete(left)
        db.commit()


@pytest.fixture
def over_mcp(monkeypatch: pytest.MonkeyPatch):
    """Run a block of calls against a live MCP endpoint.

    The session manager and the requests have to share one event loop, so the
    whole exchange happens inside a single ``asyncio.run``.
    """
    monkeypatch.setattr(settings, "MCP_ENABLED", True)
    from fluksio.mcp import http as mcp_http
    from fluksio.mcp import server as mcp_server

    # A FastMCP instance enters its session manager once, and another test in
    # the run may have built one already.
    monkeypatch.setattr(mcp_server.mcp, "_session_manager", None, raising=False)
    mcp_app = mcp_http.build_http_app(app)

    def run(block: Callable[[httpx.AsyncClient], Awaitable[Any]]) -> Any:
        async def main() -> Any:
            async with mcp_server.mcp.session_manager.run():
                async with httpx.AsyncClient(
                    transport=httpx.ASGITransport(app=mcp_app),
                    base_url=settings.oauth_issuer,
                ) as client:
                    return await block(client)

        return asyncio.run(main())

    return run


def rpc(token: str | None, method: str, **params: Any) -> dict[str, Any]:
    headers = dict(MCP_HEADERS)
    if token:
        headers["Authorization"] = f"Bearer {token}"
    body: dict[str, Any] = {"jsonrpc": "2.0", "id": 1, "method": method}
    if params:
        body["params"] = params
    return {"json": body, "headers": headers}


def test_an_unauthenticated_call_says_where_to_authenticate(over_mcp) -> None:
    async def block(client: httpx.AsyncClient) -> httpx.Response:
        return await client.post("/mcp", **rpc(None, "tools/list"))

    response = over_mcp(block)
    assert response.status_code == 401
    assert "resource_metadata=" in response.headers.get("www-authenticate", "")


def test_the_resource_metadata_names_the_authorization_server(over_mcp) -> None:
    async def block(client: httpx.AsyncClient) -> httpx.Response:
        return await client.get("/.well-known/oauth-protected-resource/mcp")

    document = over_mcp(block).json()
    # The SDK normalises the URL, so compare without the trailing slash.
    listed = [url.rstrip("/") for url in document["authorization_servers"]]
    assert settings.oauth_issuer in listed


@pytest.mark.usefixtures("enrolled")
def test_through_the_hub_an_agent_is_shown_the_hubs_addresses(
    over_mcp,
    client: TestClient,
    superuser_token_headers: dict[str, str],
    agent: uuid.UUID,
) -> None:
    """What the connector sends: its marker, and the hub as the Host."""
    from fluksio.cloud import config as cloud_config
    from tests.utils.portal import INSTANCE_ID, ISSUER

    via = {cloud_config.VIA_HEADER: cloud_config.VIA_PORTAL, "Host": "hub.example.test"}
    base = f"{ISSUER}/i/{INSTANCE_ID}"
    me = client.get(
        f"{settings.API_V1_STR}/users/me", headers=superuser_token_headers
    ).json()
    token = mcp_token(uuid.UUID(me["id"]), agent)

    async def block(http: httpx.AsyncClient) -> list[httpx.Response]:
        anonymous = rpc(None, "tools/list")
        signed = rpc(token, "tools/list")
        return [
            await http.post(
                "/mcp", json=anonymous["json"], headers={**anonymous["headers"], **via}
            ),
            await http.get("/.well-known/oauth-protected-resource/mcp", headers=via),
            await http.post(
                "/mcp", json=signed["json"], headers={**signed["headers"], **via}
            ),
            # Any other name is still refused.
            await http.post(
                "/mcp",
                json=signed["json"],
                headers={**signed["headers"], "Host": "elsewhere.test"},
            ),
        ]

    unauthenticated, document, listed, elsewhere = over_mcp(block)
    assert unauthenticated.status_code == 401
    assert (
        f'resource_metadata="{ISSUER}/.well-known/oauth-protected-resource/i/{INSTANCE_ID}/mcp"'
        in unauthenticated.headers["www-authenticate"]
    )
    assert document.json()["resource"] == f"{base}/mcp"
    assert document.json()["authorization_servers"] == [base]
    assert listed.status_code == 200
    assert elsewhere.status_code == 421


def test_a_browser_token_is_not_an_agent_token(
    over_mcp, superuser_token_headers: dict[str, str]
) -> None:
    session_token = superuser_token_headers["Authorization"].removeprefix("Bearer ")

    async def block(client: httpx.AsyncClient) -> httpx.Response:
        return await client.post("/mcp", **rpc(session_token, "tools/list"))

    # It validates perfectly well against the API; it is refused here because
    # it was issued for a person's session, not for an agent.
    assert over_mcp(block).status_code == 401


def test_an_agent_can_list_and_call_tools(
    over_mcp,
    client: TestClient,
    superuser_token_headers: dict[str, str],
    agent: uuid.UUID,
) -> None:
    me = client.get(
        f"{settings.API_V1_STR}/users/me", headers=superuser_token_headers
    ).json()
    token = mcp_token(uuid.UUID(me["id"]), agent)

    async def block(http: httpx.AsyncClient) -> tuple[Any, Any]:
        listed = await http.post("/mcp", **rpc(token, "tools/list"))
        called = await http.post(
            "/mcp",
            **rpc(token, "tools/call", name="list_flows", arguments={}),
        )
        return listed, called

    listed, called = over_mcp(block)

    assert listed.status_code == 200
    names = {tool["name"] for tool in listed.json()["result"]["tools"]}
    assert {"list_flows", "save_flow", "publish_flow", "run_flow"} <= names

    # The call reached the real API, carrying the agent's own token.
    assert called.status_code == 200
    assert "error" not in called.json()
    assert called.json()["result"]["isError"] is False


def test_revoking_an_agent_stops_the_token_it_already_holds(
    over_mcp,
    client: TestClient,
    db: Session,
    superuser_token_headers: dict[str, str],
    agent: uuid.UUID,
) -> None:
    """Withdrawing an agent has to bite now, not whenever its token expires.

    An access token is a stateless JWT valid for its whole life, so the
    registration it names is the only thing deleting a client takes away.

    Both calls share one block because the session manager runs once.
    """
    me = client.get(
        f"{settings.API_V1_STR}/users/me", headers=superuser_token_headers
    ).json()
    token = mcp_token(uuid.UUID(me["id"]), agent)

    async def block(http: httpx.AsyncClient) -> tuple[Any, Any]:
        before = await http.post("/mcp", **rpc(token, "tools/list"))
        row = db.get(OAuthClient, agent)
        assert row is not None
        db.delete(row)
        db.commit()
        after = await http.post("/mcp", **rpc(token, "tools/list"))
        return before, after

    before, after = over_mcp(block)
    assert before.status_code == 200
    assert after.status_code == 401


def test_an_agent_cannot_write_code_its_person_may_not(
    over_mcp,
    client: TestClient,
    normal_user_token_headers: dict[str, str],
    agent: uuid.UUID,
) -> None:
    """The tools call the same routes, so an agent inherits the code gate."""
    me = client.get(
        f"{settings.API_V1_STR}/users/me", headers=normal_user_token_headers
    ).json()
    assert me["is_superuser"] is False
    token = mcp_token(uuid.UUID(me["id"]), agent)

    async def block(http: httpx.AsyncClient) -> httpx.Response:
        return await http.post(
            "/mcp",
            **rpc(
                token,
                "tools/call",
                name="save_node_source",
                arguments={"name": "any", "node_id": "fn", "code": "import os"},
            ),
        )

    answer = json.loads(over_mcp(block).json()["result"]["content"][0]["text"])
    assert answer == {"error": CODE_IS_SUPERUSERS, "status": 403}