Merge fa5709fde5 into da80552897

litellm/llms/custom_httpx/http_handler.py

@@ -665,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,
@@ -1460,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:
@@ -1762,7 +1788,16 @@ def get_async_httpx_client(
             except Exception:
                 pass
 
-    _cache_key_name: Final = "async_httpx_client" + _params_key_name + llm_provider
+    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
 
     # Lazily initialize the global in-memory client cache to avoid relying on
     # litellm globals being fully populated during import time.
@@ -1777,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(  # rebind-ok: branch-specific handler initialization
+            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
@@ -1792,7 +1832,7 @@ def get_async_httpx_client(
         key=_cache_key_name,
         value=_new_client,
         ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
-        litellm_owned_client=True,
+        litellm_owned_client=_new_client.owns_client,
     )
     return _new_client
 
@@ -1812,7 +1852,13 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
             except Exception:
                 pass
 
-    _cache_key_name: Final = "httpx_client" + _params_key_name
+    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
 
     # Lazily initialize the global in-memory client cache to avoid relying on
     # litellm globals being fully populated during import time.
@@ -1827,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(  # rebind-ok: branch-specific handler initialization
+            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)
@@ -1838,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=True,
+        litellm_owned_client=_new_client.owns_client,
     )
     return _new_client
 

tests/unit/llms/custom_httpx/test_http_handler.py

@@ -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,
@@ -2172,3 +2173,101 @@ async def test_async_ollama_ssl_verify_false_builds_a_client_that_skips_certific
     assert route.call_count == 1
     assert text == "Sunny"
     assert _cached_clients_verify_modes(client_cache, AsyncHTTPHandler) == (ssl.CERT_NONE,)
+
+
+@pytest.mark.asyncio
+async def test_aclient_session_is_used_by_http_handler_providers_and_not_closed(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    hits: Final[list[str]] = []  # mutable-ok: transport capture buffer
+
+    async def handle(request: httpx.Request) -> httpx.Response:
+        hits.append(str(request.url))
+        return httpx.Response(
+            400,
+            request=request,
+            json={"error": {"message": "mock", "type": "invalid_request_error"}},
+        )
+
+    client_cache: Final = LLMClientCache()
+    monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache)
+    monkeypatch.setattr(litellm, "num_retries", 0)
+
+    default_handler: Final = get_async_httpx_client(llm_provider="anthropic")
+    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
+
+        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(
+                model="anthropic/claude-sonnet-4-20250514",
+                messages=[{"role": "user", "content": "hi"}],
+                api_key="sk-test",
+                api_base="http://127.0.0.1:9",
+                timeout=3,
+            )
+
+        assert hits == ["http://127.0.0.1:9/v1/messages"]
+        await session_handler.close()
+        assert not session.is_closed
+    finally:
+        await session.aclose()
+        await default_handler.close()
+
+
+def test_client_session_is_used_by_sync_http_handler_providers_and_not_closed(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    hits: Final[list[str]] = []  # mutable-ok: transport capture buffer
+
+    def handle(request: httpx.Request) -> httpx.Response:
+        hits.append(str(request.url))
+        return httpx.Response(
+            400,
+            request=request,
+            json={"error": {"message": "mock", "type": "invalid_request_error"}},
+        )
+
+    client_cache: Final = LLMClientCache()
+    monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache)
+    monkeypatch.setattr(litellm, "num_retries", 0)
+
+    default_handler: Final = get_httpx_client()
+    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
+
+        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(
+                model="anthropic/claude-sonnet-4-20250514",
+                messages=[{"role": "user", "content": "hi"}],
+                api_key="sk-test",
+                api_base="http://127.0.0.1:9",
+                timeout=3,
+            )
+
+        assert hits == ["http://127.0.0.1:9/v1/messages"]
+        session_handler.close()
+        assert not session.is_closed
+    finally:
+        session.close()
+        default_handler.close()