diff --git a/litellm/llms/custom_httpx/async_client_cleanup.py b/litellm/llms/custom_httpx/async_client_cleanup.py index 9c1f6af7e9c..f7f5c90d955 100644 --- a/litellm/llms/custom_httpx/async_client_cleanup.py +++ b/litellm/llms/custom_httpx/async_client_cleanup.py @@ -3,6 +3,7 @@ Utility functions for cleaning up async HTTP clients to prevent resource leaks. """ import asyncio +import inspect async def close_litellm_async_clients(): @@ -46,13 +47,22 @@ async def close_litellm_async_clients(): # Silently ignore errors during cleanup pass - # Handle any other objects with aclose method - elif hasattr(handler, "aclose"): - try: - await handler.aclose() - except Exception: - # Silently ignore errors during cleanup - pass + # Handle any other cached async client. Most expose aclose(), but some + # SDK clients -- notably the openai SDK's AsyncOpenAI / AsyncAzureOpenAI, + # which are cached here (set_cached_openai_client) and built on litellm's + # aiohttp transport -- name the async teardown close() instead. Accept + # either, guarding on iscoroutinefunction so sync close() methods + # (sync OpenAI / HTTPHandler) are left untouched. + else: + closer = getattr(handler, "aclose", None) or getattr( + handler, "close", None + ) + if closer is not None and inspect.iscoroutinefunction(closer): + try: + await closer() + except Exception: + # Silently ignore errors during cleanup + pass # Close the global base_llm_aiohttp_handler instance (issue #12443) # This is used by Gemini and other providers that use aiohttp diff --git a/tests/test_litellm/llms/custom_httpx/test_openai_session_leak.py b/tests/test_litellm/llms/custom_httpx/test_openai_session_leak.py new file mode 100644 index 00000000000..fde184626ac --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_openai_session_leak.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +""" +Test for the AsyncOpenAI cleanup gap: close_litellm_async_clients() must close +cached clients whose async teardown is named close() (the openai SDK's +AsyncOpenAI / AsyncAzureOpenAI), not only those exposing aclose(). + +Proposed path: tests/test_litellm/llms/custom_httpx/test_openai_session_leak.py +""" + +import pytest + +import litellm +from litellm.llms.custom_httpx.async_client_cleanup import close_litellm_async_clients + + +class _AsyncCloseOnlyClient: + """Mirrors openai.AsyncOpenAI: async close(), no aclose(), no .client.""" + + def __init__(self): + self.closed = False + + async def close(self): + self.closed = True + + +class _SyncClient: + """A sync client (close() is not a coroutine) -- must be left untouched.""" + + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + +@pytest.mark.asyncio +async def test_cleanup_closes_async_close_only_clients(): + async_client = _AsyncCloseOnlyClient() + sync_client = _SyncClient() + litellm.in_memory_llm_clients_cache.set_cache( + key="test_async_close_only", value=async_client, ttl=60 + ) + litellm.in_memory_llm_clients_cache.set_cache( + key="test_sync_close", value=sync_client, ttl=60 + ) + + await close_litellm_async_clients() + + assert async_client.closed, "async close()-only client should be closed" + assert not sync_client.closed, "sync client should be left untouched"