fix(litellm): close shared cloud client

## Description

Adds an explicit, idempotent async cleanup lifecycle for the LiteLLM
callback's shared cloud HTTP client.

Fixes #2894

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Added `HeadroomCallback.aclose()` to close the lazily-created
`httpx.AsyncClient` and clear its reference.
- Made cleanup safe when cloud mode was never used and when shutdown
cleanup is invoked more than once.
- Added regression coverage for initialized-client cleanup, reference
clearing, and repeated/no-op cleanup.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
python -m pytest -q tests/test_integrations/test_litellm_callback.py
5 passed

ruff check .
All checks passed!

ruff format --check .
1382 files already formatted

python -m mypy headroom
Success: no issues found in 515 source files

python -m pytest -q
Collection blocked in this Windows environment by 174 errors, primarily missing compiled headroom._core; one unrelated test also lacks respx.
```

## Real Behavior Proof

- Environment: Windows, Python 3.12, loopback HTTP server, real
`httpx.AsyncClient`.
- Exact command / steps: Started a local HTTP server, configured
`HeadroomCallback(api_key="hdr_test",
api_url="http://127.0.0.1:<port>")`, ran `_cloud_compress()` against it,
saved the created client, awaited `callback.aclose()`, then awaited
`callback.aclose()` again.
- Observed result: The real cloud request succeeded; the client was open
during the request, reported closed after `aclose()`, the callback
reference became `None`, and repeated cleanup was harmless.
- Who maintains it: Headroom Labs maintains this active upstream
repository and its LiteLLM integration.
- Install surface: No dependencies or install behavior changed. Cloud
mode continues to use the existing optional `httpx` dependency; no
native code or runtime network access is introduced by this fix.
- Not tested: The complete test suite could not run past collection
because the local Windows environment lacks the compiled
`headroom._core` extension.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] Documentation changes are not required; `aclose()` is documented
in its public docstring and the host owns shutdown sequencing
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes (full
suite blocked by missing native extension; targeted tests pass)
- [x] I did not edit `CHANGELOG.md` - it is generated by release-please
from my Conventional Commit PR title.

## Screenshots (if applicable)

Not applicable.

## Additional Notes

The callback exposes `aclose()` for the host application's async
shutdown lifecycle, matching the existing ASGI integration pattern.
This commit is contained in:
Abhinav Kumar Singh 2026-08-11 22:43:10 +05:30 committed by GitHub
parent 99f07e7bbd
commit 07d89a751d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 40 additions and 0 deletions

View file

@ -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,

View file

@ -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