fix(mcp): use pure ASGI middleware for client IP forwarding #1444

Author: XiaJunjie2020

backend/main.py

@@ -13,7 +13,7 @@
 from starlette.datastructures import MutableHeaders
 from starlette.exceptions import HTTPException as StarletteHTTPException
 from starlette.middleware.cors import CORSMiddleware
-from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp, Receive, Scope, Send
 
 from alembic import command
 from apps.api import api_router
@@ -98,18 +98,21 @@ def custom_generate_unique_id(route: APIRoute) -> str:
 )
 
 
-class McpClientIpForwardMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request: Request, call_next):
-        client_host = request.client.host if request.client else None
-        if client_host:
-            headers = MutableHeaders(scope=request.scope)
-            if not headers.get("x-real-ip"):
-                headers["x-real-ip"] = client_host
-            if not headers.get("x-forwarded-for"):
-                headers["x-forwarded-for"] = client_host
-            if not headers.get("x-client-ip"):
-                headers["x-client-ip"] = client_host
-        return await call_next(request)
+class McpClientIpForwardMiddleware:
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        # MCP SSE uses the ASGI send channel directly. Do not wrap its response
+        # in BaseHTTPMiddleware's response-body iterator.
+        if scope["type"] == "http":
+            client = scope.get("client")
+            if client:
+                headers = MutableHeaders(scope=scope)
+                for name in ("x-real-ip", "x-forwarded-for", "x-client-ip"):
+                    if not headers.get(name):
+                        headers[name] = client[0]
+        await self.app(scope, receive, send)
 
 # cache docs for different text
 _openapi_cache: Dict[str, Dict[str, Any]] = {}