diff --git a/miles/rollout/session/session_server.py b/miles/rollout/session/session_server.py index 07b90417ca..6e48c10b35 100644 --- a/miles/rollout/session/session_server.py +++ b/miles/rollout/session/session_server.py @@ -15,7 +15,6 @@ import setproctitle import uvicorn from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse from starlette.responses import Response from miles.rollout.session.sessions import get_worker_stats, setup_session_routes @@ -95,25 +94,21 @@ async def do_proxy( def build_proxy_response(self, result: dict) -> Response: content = result["response_body"] status_code = result["status_code"] - # Strip framing headers so JSONResponse / Response recompute them - # from the actual rendered body. Forwarding upstream's content-length - # verbatim breaks uvicorn h11 with "Too much data for declared - # Content-Length" whenever our re-serialization differs in even one - # byte. Mirrors the strip already done on the request path in do_proxy. - # Also strip "server": uvicorn adds its own Server header; passing - # the upstream one through produces two Server headers, which strict - # HTTP parsers (aiohttp/llhttp via litellm) reject as malformed. + headers = { k: v for k, v in result["headers"].items() if k.lower() not in ("content-length", "transfer-encoding", "server") } - content_type = headers.get("content-type", "") - try: - data = orjson.loads(content) - return JSONResponse(content=data, status_code=status_code, headers=headers) - except (orjson.JSONDecodeError, UnicodeDecodeError): - return Response(content=content, status_code=status_code, headers=headers, media_type=content_type) + + content_type = headers.get("content-type") + + return Response( + content=content, + status_code=status_code, + headers=headers, + media_type=content_type, + ) async def _stats_logger_loop(worker_port, interval_seconds: float = 30.0): diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index ae6c866361..aee310ad56 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -296,7 +296,7 @@ async def chat_completions(request: Request, session_id: str): else: return backend.build_proxy_response(result) - response = orjson.loads(result["response_body"]) + response = await asyncio.to_thread(orjson.loads, result["response_body"]) choice = response.get("choices", [{}])[0]