329 lines
12 KiB
Python
329 lines
12 KiB
Python
"""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 mcp import ClientSession
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
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())
|
|
|
|
# For a test that opens connections of its own, as a real client does.
|
|
run.app = mcp_app # type: ignore[attr-defined]
|
|
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}
|
|
|
|
|
|
def test_a_real_client_can_connect_and_call_a_tool(
|
|
over_mcp,
|
|
client: TestClient,
|
|
superuser_token_headers: dict[str, str],
|
|
agent: uuid.UUID,
|
|
) -> None:
|
|
"""The SDK's own client, over the transport an agent uses.
|
|
|
|
It does what the raw calls above leave out: initialize, the initialized
|
|
notification and the standalone-stream probe, with the bearer on every hop.
|
|
"""
|
|
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, Any]:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=over_mcp.app),
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
) as http_client:
|
|
async with streamable_http_client(
|
|
f"{settings.oauth_issuer}/mcp", http_client=http_client
|
|
) as (read, write, _):
|
|
async with ClientSession(read, write) as session:
|
|
initialized = await session.initialize()
|
|
tools = await session.list_tools()
|
|
called = await session.call_tool("list_flows", {})
|
|
return initialized, tools, called
|
|
|
|
initialized, tools, called = over_mcp(block)
|
|
assert initialized.serverInfo.name == "fluksio"
|
|
assert {"list_flows", "run_flow"} <= {tool.name for tool in tools.tools}
|
|
assert called.isError is False
|
|
|
|
|
|
@pytest.mark.parametrize("revision", ["2025-03-26", "2025-06-18", "2025-11-25"])
|
|
def test_clients_of_older_protocol_revisions_still_connect(
|
|
over_mcp,
|
|
client: TestClient,
|
|
superuser_token_headers: dict[str, str],
|
|
agent: uuid.UUID,
|
|
revision: str,
|
|
) -> None:
|
|
"""Agents in the field speak whatever revision they shipped with."""
|
|
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]:
|
|
opening = rpc(
|
|
token,
|
|
"initialize",
|
|
protocolVersion=revision,
|
|
capabilities={},
|
|
clientInfo={"name": "test", "version": "1"},
|
|
)
|
|
initialized = await http.post("/mcp", **opening)
|
|
later = {**opening["headers"], "MCP-Protocol-Version": revision}
|
|
listed = await http.post(
|
|
"/mcp",
|
|
json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
|
|
headers=later,
|
|
)
|
|
return initialized, listed
|
|
|
|
initialized, listed = over_mcp(block)
|
|
assert initialized.json()["result"]["protocolVersion"] == revision
|
|
assert listed.status_code == 200
|
|
assert "list_flows" in {tool["name"] for tool in listed.json()["result"]["tools"]}
|