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:
Pouzor
2026-03-13 17:28:38 +01:00
parent e1d16b86e3
commit e41dbe579c
5 changed files with 137 additions and 27 deletions
+34 -11
View File
@@ -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)