Update app/middleware/profiler.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: grillazz

app/middleware/profiler.py

@@ -2,27 +2,38 @@
 
 from fastapi import Request
 from pyinstrument import Profiler
-from starlette.middleware.base import (
-    BaseHTTPMiddleware,
-    RequestResponseEndpoint,
-)
 from starlette.responses import HTMLResponse, Response
 
 
-class ProfilingMiddleware(BaseHTTPMiddleware):
-    async def dispatch(
-        self, request: Request, call_next: RequestResponseEndpoint
-    ) -> Response:
+class ProfilingMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send) -> Response:
+        # Only handle HTTP requests; pass through other scope types (e.g. WebSocket).
+        if scope.get("type") != "http":
+            return await self.app(scope, receive, send)
+
+        request = Request(scope, receive=receive)
+
         if request.query_params.get("pyprofile") == "true":
             profiler = Profiler(interval=0.001, async_mode="enabled")
             profiler.start()
 
-            await call_next(request)
+            async def _inner_send(message):
+                # Swallow downstream response messages; we replace the response
+                # entirely with the profiler HTML output, matching the previous behavior.
+                return
+
+            await self.app(scope, receive, _inner_send)
 
             profiler.stop()
-            return HTMLResponse(
+            response = HTMLResponse(
                 profiler.output_html(),
-                headers={"Content-Disposition": "attachment; filename=profile.html"},
+                headers={
+                    "Content-Disposition": "attachment; filename=profile.html"
+                },
             )
+            return await response(scope, receive, send)
 
-        return await call_next(request)
+        return await self.app(scope, receive, send)