fix(http): isolate global sessions from module-level URL fetchers and satisfy strict type checks

This commit is contained in:
Muhammad Osama Amin 2026-10-18 13:44:28 +00:00
parent 740be08822
commit 963c86ec57
2 changed files with 75 additions and 53 deletions

View file

@ -633,7 +633,6 @@ class AsyncHTTPHandler:
shared_session: Optional["ClientSession"] = None,
transport: httpx.AsyncBaseTransport | None = None,
follow_redirects: bool = True,
client: httpx.AsyncClient | None = None,
):
self.timeout = timeout
self.event_hooks = event_hooks
@ -641,41 +640,17 @@ class AsyncHTTPHandler:
self.shared_session = shared_session
self.transport = transport
self.follow_redirects = follow_redirects
resolved_client: Final = (
client
if client is not None
else (
litellm.aclient_session
if isinstance(litellm.aclient_session, httpx.AsyncClient)
and transport is None
and shared_session is None
and event_hooks is None
else None
)
)
self._owns_client = resolved_client is None
self._client = (
self.create_client(
timeout=timeout,
event_hooks=event_hooks,
ssl_verify=ssl_verify,
shared_session=shared_session,
)
if resolved_client is None
else resolved_client
self._owns_client = True
self._client = self.create_client(
timeout=timeout,
event_hooks=event_hooks,
ssl_verify=ssl_verify,
shared_session=shared_session,
)
self.client_alias = client_alias
@property
def client(self) -> httpx.AsyncClient:
if (
self._owns_client
and self.transport is None
and self.shared_session is None
and self.event_hooks is None
and isinstance(litellm.aclient_session, httpx.AsyncClient)
):
return litellm.aclient_session
if self._owns_client and self._client.is_closed:
self._client = self.create_client(
timeout=self.timeout,
@ -690,6 +665,28 @@ class AsyncHTTPHandler:
self._client = client
self._owns_client = False
@property
def owns_client(self) -> bool:
return self._owns_client
@classmethod
def from_client(
cls,
client: httpx.AsyncClient,
timeout: float | httpx.Timeout | None = None,
) -> "AsyncHTTPHandler":
handler: Final = cls.__new__(cls)
handler.timeout = timeout
handler.event_hooks = None
handler.ssl_verify = None
handler.shared_session = None
handler.transport = None
handler.follow_redirects = True
handler._owns_client = False
handler._client = client
handler.client_alias = None
return handler
def create_client(
self,
timeout: float | httpx.Timeout | None,
@ -1444,14 +1441,9 @@ class HTTPHandler:
self.ssl_verify = ssl_verify
self.disable_default_headers = disable_default_headers
self.follow_redirects = follow_redirects
resolved_client: Final = (
client
if client is not None
else (litellm.client_session if isinstance(litellm.client_session, httpx.Client) else None)
)
self._owns_client = resolved_client is None
self._owns_client = client is None
self._heal_lock = threading.Lock()
self._client = self.create_client() if resolved_client is None else resolved_client
self._client = self.create_client() if client is None else client
def create_client(self) -> httpx.Client:
# Get unified SSL configuration
@ -1479,8 +1471,6 @@ class HTTPHandler:
@property
def client(self) -> httpx.Client:
if self._owns_client and isinstance(litellm.client_session, httpx.Client):
return litellm.client_session
if self._owns_client and self._client.is_closed:
with self._heal_lock:
if self._owns_client and self._client.is_closed:
@ -1492,6 +1482,10 @@ class HTTPHandler:
self._client = client
self._owns_client = False
@property
def owns_client(self) -> bool:
return self._owns_client
def close(self):
# Close the client when you're done with it
if self._owns_client:
@ -1794,7 +1788,14 @@ def get_async_httpx_client(
except Exception:
pass
aclient_session: Final = litellm.aclient_session if isinstance(litellm.aclient_session, httpx.AsyncClient) else None
aclient_session: Final = (
litellm.aclient_session
if isinstance(litellm.aclient_session, httpx.AsyncClient)
and shared_session is None
and llm_provider != "litellm_module_level_client"
and (params is None or ("transport" not in params and "event_hooks" not in params))
else None
)
_session_key_name: Final = f"_aclient_session_{id(aclient_session)}" if aclient_session is not None else ""
_cache_key_name: Final = "async_httpx_client" + _params_key_name + llm_provider + _session_key_name
@ -1811,7 +1812,12 @@ def get_async_httpx_client(
if _cached_client:
return _cached_client
if params is not None:
if aclient_session is not None:
_new_client = AsyncHTTPHandler.from_client(
client=aclient_session,
timeout=_default_cached_client_timeout(),
)
elif params is not None:
# Filter out params that are only used for cache key, not for AsyncHTTPHandler.__init__
handler_params: Final = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"}
handler_params["shared_session"] = shared_session
@ -1826,7 +1832,7 @@ def get_async_httpx_client(
key=_cache_key_name,
value=_new_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=_new_client._owns_client,
litellm_owned_client=_new_client.owns_client,
)
return _new_client
@ -1846,7 +1852,11 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
except Exception:
pass
client_session: Final = litellm.client_session if isinstance(litellm.client_session, httpx.Client) else None
client_session: Final = (
litellm.client_session
if isinstance(litellm.client_session, httpx.Client) and (params is None or "client" not in params)
else None
)
_session_key_name: Final = f"_client_session_{id(client_session)}" if client_session is not None else ""
_cache_key_name: Final = "httpx_client" + _params_key_name + _session_key_name
@ -1863,7 +1873,12 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
if _cached_client:
return _cached_client
if params is not None:
if client_session is not None:
_new_client = HTTPHandler(
timeout=_default_cached_client_timeout(),
client=client_session,
)
elif params is not None:
# Filter out params that are only used for cache key, not for HTTPHandler.__init__
handler_params: Final = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"}
_new_client = HTTPHandler(**handler_params)
@ -1874,7 +1889,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
key=_cache_key_name,
value=_new_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=_new_client._owns_client,
litellm_owned_client=_new_client.owns_client,
)
return _new_client

View file

@ -25,6 +25,7 @@ from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
MaskedHTTPStatusError,
get_async_httpx_client,
get_httpx_client,
get_shared_realtime_ssl_context,
get_ssl_configuration,
@ -2192,18 +2193,19 @@ async def test_aclient_session_is_used_by_http_handler_providers_and_not_closed(
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache)
monkeypatch.setattr(litellm, "num_retries", 0)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
# Pre-warm cache without aclient_session to verify setting it bypasses stale default handlers
default_handler: Final = get_async_httpx_client(llm_provider="anthropic")
assert default_handler._owns_client is True
assert default_handler.owns_client is True
session: Final = httpx.AsyncClient(transport=httpx.MockTransport(handle), trust_env=False)
try:
monkeypatch.setattr(litellm, "aclient_session", session)
session_handler: Final = get_async_httpx_client(llm_provider="anthropic")
assert session_handler.client is session
assert session_handler._owns_client is False
assert session_handler.owns_client is False
module_level_handler: Final = get_async_httpx_client(llm_provider="litellm_module_level_client")
assert module_level_handler.client is not session
assert module_level_handler.owns_client is True
with pytest.raises(litellm.BadRequestError):
await litellm.acompletion(
@ -2240,14 +2242,19 @@ def test_client_session_is_used_by_sync_http_handler_providers_and_not_closed(
monkeypatch.setattr(litellm, "num_retries", 0)
default_handler: Final = get_httpx_client()
assert default_handler._owns_client is True
assert default_handler.owns_client is True
session: Final = httpx.Client(transport=httpx.MockTransport(handle), trust_env=False)
try:
monkeypatch.setattr(litellm, "client_session", session)
session_handler: Final = get_httpx_client()
assert session_handler.client is session
assert session_handler._owns_client is False
assert session_handler.owns_client is False
module_level_sync: Final = HTTPHandler()
assert module_level_sync.client is not session
assert module_level_sync.owns_client is True
module_level_sync.close()
with pytest.raises(litellm.BadRequestError):
litellm.completion(