diff --git a/mcp/app/auth.py b/mcp/app/auth.py index 4633ef2..786751c 100644 --- a/mcp/app/auth.py +++ b/mcp/app/auth.py @@ -1,21 +1,44 @@ import hmac -from starlette.middleware.base import BaseHTTPMiddleware -from starlette.requests import Request -from starlette.responses import JSONResponse +import json +from starlette.types import ASGIApp, Receive, Scope, Send from .config import settings +_BYPASS_PATHS = {"/health", "/register"} -class ApiKeyMiddleware(BaseHTTPMiddleware): - async def dispatch(self, request: Request, call_next): - # Health check bypass - if request.url.path == "/health": - return await call_next(request) - key = request.headers.get("X-API-Key", "") +class ApiKeyMiddleware: + """Pure ASGI middleware — compatible with SSE/streaming responses. + + BaseHTTPMiddleware buffers the full response body and breaks SSE streams. + This implementation operates at the ASGI scope level and never touches + the response stream. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + path: str = scope.get("path", "") + + if path in _BYPASS_PATHS or path.startswith("/.well-known/"): + await self.app(scope, receive, send) + return + + headers = dict(scope.get("headers", [])) + key = headers.get(b"x-api-key", b"").decode() expected = settings.mcp_api_key if not key or not hmac.compare_digest(key.encode(), expected.encode()): - return JSONResponse({"detail": "Invalid or missing X-API-Key"}, status_code=401) + body = json.dumps({"detail": "Invalid or missing X-API-Key"}).encode() + await send({"type": "http.response.start", "status": 401, + "headers": [(b"content-type", b"application/json"), + (b"content-length", str(len(body)).encode())]}) + await send({"type": "http.response.body", "body": body, "more_body": False}) + return - return await call_next(request) + await self.app(scope, receive, send) diff --git a/mcp/app/main.py b/mcp/app/main.py index c55019d..ca40953 100644 --- a/mcp/app/main.py +++ b/mcp/app/main.py @@ -1,7 +1,7 @@ from contextlib import asynccontextmanager -from fastapi import FastAPI +from fastapi import FastAPI, Request from mcp.server import Server -from mcp.server.sse import SseServerTransport +from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from .auth import ApiKeyMiddleware from .backend_client import backend @@ -13,29 +13,28 @@ mcp_server = Server("homelable") register_resources(mcp_server) register_tools(mcp_server) +session_manager = StreamableHTTPSessionManager( + app=mcp_server, + json_response=False, + stateless=True, +) + @asynccontextmanager async def lifespan(app: FastAPI): await backend.start() - yield + async with session_manager.run(): + yield await backend.stop() app = FastAPI(title="Homelable MCP", lifespan=lifespan) app.add_middleware(ApiKeyMiddleware) -sse = SseServerTransport("/mcp/messages") - -@app.get("/mcp") -async def mcp_sse(request): - async with sse.connect_sse(request.scope, request.receive, request._send) as streams: - await mcp_server.run(streams[0], streams[1], mcp_server.create_initialization_options()) - - -@app.post("/mcp/messages") -async def mcp_messages(request): - await sse.handle_post_message(request.scope, request.receive, request._send) +@app.api_route("/mcp", methods=["GET", "POST", "DELETE"]) +async def mcp_endpoint(request: Request): + await session_manager.handle_request(request.scope, request.receive, request._send) @app.get("/health") diff --git a/mcp/app/tools.py b/mcp/app/tools.py index 3bf3225..56fb933 100644 --- a/mcp/app/tools.py +++ b/mcp/app/tools.py @@ -71,6 +71,18 @@ def register_tools(server: Server): "required": ["id"], "properties": {"id": {"type": "string"}}, }), + Tool(name="get_canvas", description="Get the full canvas: all nodes and edges in the homelab topology", inputSchema={ + "type": "object", + "properties": {}, + }), + Tool(name="list_nodes", description="List all nodes (devices) in the homelab", inputSchema={ + "type": "object", + "properties": {}, + }), + Tool(name="list_pending_devices", description="List devices discovered by scan but not yet approved or hidden", inputSchema={ + "type": "object", + "properties": {}, + }), ] @server.call_tool() @@ -79,6 +91,27 @@ def register_tools(server: Server): return [TextContent(type="text", text=json.dumps(result, indent=2))] +def _slim_canvas(raw: dict) -> dict: + """Strip React Flow layout/style fields — keep only semantic data for AI use.""" + NODE_KEEP = {"id", "type", "label", "ip", "hostname", "status", "services", "description", "parentId"} + EDGE_KEEP = {"id", "source", "target", "type", "label"} + + def slim_node(n: dict) -> dict: + data = n.get("data", {}) + out = {k: v for k, v in data.items() if k in NODE_KEEP and v not in (None, "", [])} + out["id"] = n.get("id") + out["node_type"] = n.get("type") + return out + + def slim_edge(e: dict) -> dict: + return {k: v for k, v in e.items() if k in EDGE_KEEP and v not in (None, "")} + + return { + "nodes": [slim_node(n) for n in raw.get("nodes", [])], + "edges": [slim_edge(e) for e in raw.get("edges", [])], + } + + async def _dispatch(name: str, args: dict) -> dict: if name == "create_node": return await backend.post("/api/v1/nodes", args) @@ -107,4 +140,14 @@ async def _dispatch(name: str, args: dict) -> dict: if name == "hide_device": return await backend.post(f"/api/v1/scan/pending/{args['id']}/hide", {}) + if name == "get_canvas": + raw = await backend.get("/api/v1/canvas") + return _slim_canvas(raw) + + if name == "list_nodes": + return await backend.get("/api/v1/nodes") + + if name == "list_pending_devices": + return await backend.get("/api/v1/scan/pending") + raise ValueError(f"Unknown tool: {name}") diff --git a/mcp/tests/test_auth.py b/mcp/tests/test_auth.py index c34e8ad..29cdfc2 100644 --- a/mcp/tests/test_auth.py +++ b/mcp/tests/test_auth.py @@ -1,4 +1,5 @@ import pytest +from unittest.mock import AsyncMock, patch @pytest.mark.anyio @@ -21,6 +22,7 @@ async def test_wrong_api_key(client): @pytest.mark.anyio async def test_valid_api_key_passes(client, api_key): - # SSE endpoint returns 200 with valid key (even if stream is incomplete in tests) - resp = await client.get("/mcp", headers={"X-API-Key": api_key}) + # Auth passes — mock handle_request so we don't need a live MCP session + with patch("app.main.session_manager.handle_request", new_callable=AsyncMock): + resp = await client.get("/mcp", headers={"X-API-Key": api_key}) assert resp.status_code != 401 diff --git a/mcp/tests/test_tools.py b/mcp/tests/test_tools.py index a97ea1a..80a6cee 100644 --- a/mcp/tests/test_tools.py +++ b/mcp/tests/test_tools.py @@ -9,6 +9,7 @@ def mock_backend(): m.post = AsyncMock(return_value={"id": "1"}) m.patch = AsyncMock(return_value={"id": "1"}) m.delete = AsyncMock(return_value={}) + m.get = AsyncMock(return_value=[]) yield m @@ -67,6 +68,48 @@ async def test_hide_device(mock_backend): mock_backend.post.assert_called_once_with("/api/v1/scan/pending/5/hide", {}) +@pytest.mark.anyio +async def test_get_canvas(mock_backend): + mock_backend.get = AsyncMock(return_value={ + "nodes": [ + { + "id": "n1", + "type": "router", + "position": {"x": 100, "y": 200}, + "width": 160, + "height": 80, + "data": {"label": "Freebox", "ip": "192.168.1.1", "status": "online"}, + } + ], + "edges": [ + {"id": "e1", "source": "n1", "target": "n2", "type": "ethernet", "animated": True, "style": {"stroke": "#fff"}}, + ], + "viewport": {"x": 0, "y": 0, "zoom": 1}, + }) + result = await _dispatch("get_canvas", {}) + mock_backend.get.assert_called_once_with("/api/v1/canvas") + # Layout/style fields stripped, only semantic data kept + assert result["nodes"] == [{"id": "n1", "node_type": "router", "label": "Freebox", "ip": "192.168.1.1", "status": "online"}] + assert result["edges"] == [{"id": "e1", "source": "n1", "target": "n2", "type": "ethernet"}] + assert "viewport" not in result + + +@pytest.mark.anyio +async def test_list_nodes(mock_backend): + mock_backend.get = AsyncMock(return_value=[{"id": "1", "label": "Freebox"}]) + result = await _dispatch("list_nodes", {}) + mock_backend.get.assert_called_once_with("/api/v1/nodes") + assert result == [{"id": "1", "label": "Freebox"}] + + +@pytest.mark.anyio +async def test_list_pending_devices(mock_backend): + mock_backend.get = AsyncMock(return_value=[{"id": "p1", "ip": "192.168.1.50"}]) + result = await _dispatch("list_pending_devices", {}) + mock_backend.get.assert_called_once_with("/api/v1/scan/pending") + assert result == [{"id": "p1", "ip": "192.168.1.50"}] + + @pytest.mark.anyio async def test_unknown_tool(): with pytest.raises(ValueError, match="Unknown tool"):