From 91e3c86a12a7540ad76dea62dcd44aaad6770a0a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 9 Mar 2026 11:16:35 -0700 Subject: [PATCH] fix: address 6 critical sidecar bugs from code review - Fix connection leak: use release() instead of close() in SidecarResponseStream.aclose() - Fix concurrent session leak: add asyncio.Lock with double-check in _get_session() - Fix TOCTOU race: use DashMap entry() API in Rust sidecar get_or_create_client() - Fix broken streaming: detect stream from request body JSON instead of Accept header - Fix event loop stall: replace blocking subprocess.wait() with async run_in_executor() - Fix type contract: decode orjson.dumps() bytes to str for consistency with json.dumps() Co-Authored-By: Claude Opus 4.6 --- litellm-sidecar/src/main.rs | 24 ++++----- litellm/llms/custom_httpx/llm_http_handler.py | 10 ++-- .../llms/custom_httpx/sidecar_transport.py | 54 ++++++++++++------- litellm/proxy/sidecar_client.py | 10 +++- 4 files changed, 59 insertions(+), 39 deletions(-) diff --git a/litellm-sidecar/src/main.rs b/litellm-sidecar/src/main.rs index febf6ba668e..bd1287c6091 100644 --- a/litellm-sidecar/src/main.rs +++ b/litellm-sidecar/src/main.rs @@ -33,18 +33,18 @@ impl Sidecar { } fn get_or_create_client(&self, host: &str) -> Client { - if let Some(client) = self.pools.get(host) { - return client.clone(); - } - let client = Client::builder() - .pool_max_idle_per_host(200) - .pool_idle_timeout(std::time::Duration::from_secs(90)) - .tcp_keepalive(std::time::Duration::from_secs(60)) - .tcp_nodelay(true) - .build() - .expect("Failed to build reqwest client"); - self.pools.insert(host.to_string(), client.clone()); - client + self.pools + .entry(host.to_string()) + .or_insert_with(|| { + Client::builder() + .pool_max_idle_per_host(200) + .pool_idle_timeout(std::time::Duration::from_secs(90)) + .tcp_keepalive(std::time::Duration::from_secs(60)) + .tcp_nodelay(true) + .build() + .expect("Failed to build reqwest client") + }) + .clone() } } diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7eccc8fe71c..4a06169a3f7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,5 +1,4 @@ import json -import orjson import ssl from typing import ( TYPE_CHECKING, @@ -16,6 +15,7 @@ from typing import ( ) import httpx # type: ignore +import orjson from openai.types.file_deleted import FileDeleted import litellm @@ -23,9 +23,7 @@ import litellm.litellm_core_utils import litellm.types import litellm.types.utils from litellm._logging import verbose_logger -from litellm.anthropic_beta_headers_manager import ( - update_headers_with_filtered_beta, -) +from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -181,7 +179,7 @@ class BaseLLMHTTPHandler: data=( signed_json_body if signed_json_body is not None - else orjson.dumps(data) + else orjson.dumps(data).decode() ), timeout=timeout, stream=stream, @@ -241,7 +239,7 @@ class BaseLLMHTTPHandler: data=( signed_json_body if signed_json_body is not None - else orjson.dumps(data) + else orjson.dumps(data).decode() ), timeout=timeout, stream=stream, diff --git a/litellm/llms/custom_httpx/sidecar_transport.py b/litellm/llms/custom_httpx/sidecar_transport.py index c4dfe5fa51b..501ff825da7 100644 --- a/litellm/llms/custom_httpx/sidecar_transport.py +++ b/litellm/llms/custom_httpx/sidecar_transport.py @@ -7,6 +7,8 @@ logging, callbacks, and retry logic still run in Python — only the TCP connection pooling and HTTP round-trip move to Rust. """ +import asyncio +import json import typing from typing import Optional @@ -27,7 +29,7 @@ class SidecarResponseStream(httpx.AsyncByteStream): yield chunk async def aclose(self) -> None: - self._response.close() + self._response.release() class LiteLLMSidecarTransport(httpx.AsyncBaseTransport): @@ -42,18 +44,23 @@ class LiteLLMSidecarTransport(httpx.AsyncBaseTransport): def __init__(self, sidecar_url: str = "http://127.0.0.1:8787"): self._sidecar_url = sidecar_url self._session: Optional[aiohttp.ClientSession] = None + self._session_lock = asyncio.Lock() - def _get_session(self) -> aiohttp.ClientSession: - if self._session is None or self._session.closed: - self._session = aiohttp.ClientSession( - connector=aiohttp.TCPConnector( - limit=0, - keepalive_timeout=90, - enable_cleanup_closed=True, - ), - timeout=aiohttp.ClientTimeout(total=None, connect=5), - ) - return self._session + async def _get_session(self) -> aiohttp.ClientSession: + if self._session is not None and not self._session.closed: + return self._session + async with self._session_lock: + # Double-check after acquiring lock + if self._session is None or self._session.closed: + self._session = aiohttp.ClientSession( + connector=aiohttp.TCPConnector( + limit=0, + keepalive_timeout=90, + enable_cleanup_closed=True, + ), + timeout=aiohttp.ClientTimeout(total=None, connect=5), + ) + return self._session async def handle_async_request(self, request: httpx.Request) -> httpx.Response: # Extract the provider host from the full URL for the sidecar headers @@ -69,8 +76,20 @@ class LiteLLMSidecarTransport(httpx.AsyncBaseTransport): if auth.lower().startswith("bearer "): api_key = auth[7:] - # Determine if streaming from content-type or accept headers - is_stream = "text/event-stream" in request.headers.get("accept", "") + try: + body = request.content + except httpx.RequestNotRead: + body = b"" + + # Determine if streaming from the request body's "stream" field + is_stream = False + try: + if body: + parsed_body = json.loads(body) + if isinstance(parsed_body, dict): + is_stream = parsed_body.get("stream", False) is True + except (json.JSONDecodeError, UnicodeDecodeError): + pass # Get timeout from request extensions timeout_config = request.extensions.get("timeout", {}) @@ -80,11 +99,6 @@ class LiteLLMSidecarTransport(httpx.AsyncBaseTransport): else: timeout_secs = 300 - try: - body = request.content - except httpx.RequestNotRead: - body = b"" - headers = { "X-LiteLLM-Provider-URL": provider_base, "X-LiteLLM-API-Key": api_key, @@ -100,7 +114,7 @@ class LiteLLMSidecarTransport(httpx.AsyncBaseTransport): if val: headers[f"X-LiteLLM-Fwd-{key}"] = val - session = self._get_session() + session = await self._get_session() resp = await session.post( f"{self._sidecar_url}/forward", data=body, diff --git a/litellm/proxy/sidecar_client.py b/litellm/proxy/sidecar_client.py index eeeff90cf22..15503cf9dce 100644 --- a/litellm/proxy/sidecar_client.py +++ b/litellm/proxy/sidecar_client.py @@ -142,7 +142,15 @@ class SidecarClient: self._session = None if self._process: self._process.terminate() - self._process.wait(timeout=5) + try: + await asyncio.wait_for( + asyncio.get_event_loop().run_in_executor( + None, self._process.wait + ), + timeout=5, + ) + except asyncio.TimeoutError: + self._process.kill() self._process = None self._healthy = False