fix(mcp): fix SSE streaming crash and reduce get_canvas token usage
- Replace BaseHTTPMiddleware with pure ASGI middleware in auth.py to fix the "Unexpected message: http.response.start" crash on SSE streams - Migrate from SseServerTransport to StreamableHTTPSessionManager in main.py - Add _slim_canvas() in tools.py to strip React Flow layout fields from get_canvas responses (60-80% payload reduction) - Update tests: assert canvas slimming, mock session_manager.handle_request in auth tests to avoid uninitialized task group errors
This commit is contained in:
+34
-11
@@ -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)
|
||||
|
||||
+13
-14
@@ -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")
|
||||
|
||||
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user