From 9be727de68ae22bcc1a3c63474884416d7c04523 Mon Sep 17 00:00:00 2001 From: Peter Lodri Date: Fri, 10 Jul 2026 00:09:33 +0200 Subject: [PATCH] fix(litellm): inherit CustomLogger so future hooks don't crash proxy (#1114) (#1391) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Fixes #1114 — LiteLLM 1.89.x added `async_post_call_success_hook` and started calling it after every successful completion. `HeadroomCallback` was a plain `object` subclass with no such method, causing: ``` AttributeError: type object 'HeadroomCallback' has no attribute 'async_post_call_success_hook' ``` This crashed the LiteLLM proxy on every successful API call. ### Root cause ```python class HeadroomCallback: # plain object — no-op hooks not inherited ... ``` ### Fix Inherit from `litellm.integrations.custom_logger.CustomLogger` which provides no-op defaults for every hook it defines. Future additions to `CustomLogger` will be covered automatically. ```python try: from litellm.integrations.custom_logger import CustomLogger as _CustomLogger except ImportError: _CustomLogger = object # litellm not installed — graceful fallback class HeadroomCallback(_CustomLogger): ... def __init__(self, ...): super().__init__() ... ``` ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/integrations/litellm_callback.py` — inherit `CustomLogger`; add `super().__init__()` - `tests/test_litellm_callback.py` — 7 tests: instantiation, `async_post_call_success_hook` present + callable + no-op, all current hooks present, pre-call hook still works ## Real behavior proof **Setup:** Python 3.13, litellm 1.89.1, headroom-ai 0.27.0-dev **Steps after patch:** ```bash python3 -c " from headroom.integrations.litellm_callback import HeadroomCallback import asyncio cb = HeadroomCallback() # Simulate what litellm proxy calls on success asyncio.run(cb.async_post_call_success_hook(data={}, user_api_key_dict={}, response=None)) print('OK — no AttributeError') " ``` **After-fix evidence:** Runs without exception. Before fix: `AttributeError: type object 'HeadroomCallback' has no attribute 'async_post_call_success_hook'`. **What I did not test:** Live LiteLLM proxy with YAML config (no LiteLLM proxy running in test env); tested via unit tests and direct Python instantiation. ## Test Results ``` tests/test_litellm_callback.py 7/7 passed ``` 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- headroom/integrations/litellm_callback.py | 13 +++- tests/test_litellm_callback.py | 72 +++++++++++++++++++++++ 2 files changed, 83 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm_callback.py diff --git a/headroom/integrations/litellm_callback.py b/headroom/integrations/litellm_callback.py index 3d5e137a1..4c2e4e271 100644 --- a/headroom/integrations/litellm_callback.py +++ b/headroom/integrations/litellm_callback.py @@ -24,10 +24,18 @@ logger = logging.getLogger(__name__) _DEFAULT_CLOUD_URL = "https://api.headroomlabs.ai" -class HeadroomCallback: +try: + from litellm.integrations.custom_logger import CustomLogger as _CustomLogger +except ImportError: # litellm not installed — fall back to plain object + _CustomLogger = object # type: ignore[assignment,misc] + + +class HeadroomCallback(_CustomLogger): """LiteLLM callback that compresses messages before each API call. - Implements LiteLLM's CustomLogger interface (async_pre_call_hook). + Inherits from litellm.integrations.custom_logger.CustomLogger so that + any hook LiteLLM adds in future versions (e.g. async_post_call_success_hook + added in 1.89.x) has a no-op default and won't raise AttributeError (#1114). Two modes: - Local (default): Compresses in-process using headroom.compress(). @@ -56,6 +64,7 @@ class HeadroomCallback: api_key: str | None = None, api_url: str | None = None, ) -> None: + super().__init__() self._min_tokens = min_tokens self._model_limit = model_limit self._hooks = hooks diff --git a/tests/test_litellm_callback.py b/tests/test_litellm_callback.py new file mode 100644 index 000000000..5ad40d265 --- /dev/null +++ b/tests/test_litellm_callback.py @@ -0,0 +1,72 @@ +"""Tests for HeadroomCallback LiteLLM integration. + +Regression for #1114: HeadroomCallback did not inherit CustomLogger, so any +hook LiteLLM added post-1.89.x (e.g. async_post_call_success_hook) raised +AttributeError and crashed the LiteLLM proxy. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from tests._dotenv import importorskip_no_env_leak + +importorskip_no_env_leak("litellm") + +from headroom.integrations.litellm_callback import HeadroomCallback # noqa: E402 + + +class TestHeadroomCallbackCustomLoggerInheritance: + def test_instantiates_without_error(self) -> None: + cb = HeadroomCallback() + assert cb is not None + + def test_has_async_post_call_success_hook(self) -> None: + """Regression: AttributeError: 'HeadroomCallback' has no attr 'async_post_call_success_hook'.""" + cb = HeadroomCallback() + assert hasattr(cb, "async_post_call_success_hook"), ( + "async_post_call_success_hook must exist (added in litellm 1.89.x)" + ) + + def test_async_post_call_success_hook_is_callable(self) -> None: + """LiteLLM must be able to await the hook without exception.""" + cb = HeadroomCallback() + hook = cb.async_post_call_success_hook + assert callable(hook) + + def test_async_post_call_success_hook_does_not_raise(self) -> None: + """Calling the hook (no-op from CustomLogger) must not raise.""" + cb = HeadroomCallback() + + async def _run() -> None: + await cb.async_post_call_success_hook( + data={"model": "gpt-4o", "messages": []}, + user_api_key_dict={}, + response=None, + ) + + asyncio.run(_run()) + + def test_all_current_litellm_async_hooks_present(self) -> None: + """HeadroomCallback must expose every async hook CustomLogger defines.""" + from litellm.integrations.custom_logger import CustomLogger + + cb = HeadroomCallback() + missing = [ + name + for name in dir(CustomLogger) + if name.startswith("async_") and not hasattr(cb, name) + ] + assert not missing, f"Missing CustomLogger hooks: {missing}" + + def test_async_pre_call_hook_still_works(self) -> None: + """Inheritance must not break the existing compression hook.""" + cb = HeadroomCallback() + assert hasattr(cb, "async_pre_call_hook") + assert callable(cb.async_pre_call_hook) + + def test_total_tokens_saved_property(self) -> None: + cb = HeadroomCallback() + assert cb.total_tokens_saved == 0