stroblme
f569f325a7
The connector now serves /mcp, its OAuth documents and /hooks/ besides /api/v1/, refuses any dot segment (httpx resolves them, which carried /api/v1/../../mcp onto the MCP mount), addresses calls to the hub's host, and bridges served sockets by calling the app with a websocket scope. A hook without a secret answers the tunnel with the 404 (or 1008) of a wrong secret. Requests through the hub get the hub's issuer, resource and consent page in the OAuth metadata and the 401, and the MCP host check admits the hub while enrolled.
252 lines
8.7 KiB
Python
252 lines
8.7 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 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="http://api.localhost",
|
|
) 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}
|