diff --git a/headroom/integrations/litellm_callback.py b/headroom/integrations/litellm_callback.py index d5d54cacb..c88078069 100644 --- a/headroom/integrations/litellm_callback.py +++ b/headroom/integrations/litellm_callback.py @@ -94,6 +94,18 @@ class HeadroomCallback(_CustomLogger): """Whether cloud compression is enabled.""" return self._api_key is not None + async def aclose(self) -> None: + """Close the shared cloud HTTP client, if it was initialized. + + Applications using LiteLLM should await this method during their async + shutdown lifecycle. It is safe to call when cloud mode was not used or + after the client has already been closed. + """ + client = self._client + self._client = None + if client is not None: + await client.aclose() + async def async_pre_call_hook( self, user_api_key_dict: Any = None, diff --git a/tests/test_integrations/test_litellm_callback.py b/tests/test_integrations/test_litellm_callback.py index 50f698316..ddf71d4da 100644 --- a/tests/test_integrations/test_litellm_callback.py +++ b/tests/test_integrations/test_litellm_callback.py @@ -5,6 +5,7 @@ from __future__ import annotations import importlib import inspect from pathlib import Path +from unittest.mock import AsyncMock, MagicMock import pytest @@ -52,3 +53,30 @@ class TestHeadroomCallbackPostCallSuccessHook: response=sentinel, ) assert result is sentinel + + +class TestHeadroomCallbackClientLifecycle: + """Cloud client cleanup must be explicit and safe to repeat.""" + + @pytest.mark.asyncio + async def test_aclose_closes_and_clears_initialized_client(self) -> None: + cb = HeadroomCallback(api_key="hdr_test") + client = MagicMock() + client.aclose = AsyncMock() + cb._client = client + + await cb.aclose() + + client.aclose.assert_awaited_once_with() + assert cb._client is None + + await cb.aclose() + client.aclose.assert_awaited_once_with() + + @pytest.mark.asyncio + async def test_aclose_without_initialized_client_is_a_noop(self) -> None: + cb = HeadroomCallback(api_key="hdr_test") + + await cb.aclose() + + assert cb._client is None