From f7335351adb36cc690989dd5a4675c0db147bc92 Mon Sep 17 00:00:00 2001 From: Hoya <96610382+Jeongho0805@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:38:24 +0900 Subject: [PATCH 1/2] fix(logging): reuse the cached GenericAPILogger so config reloads stop leaking flush tasks --- .../generic_api/generic_api_callback.py | 27 ++++++++++++++++--- .../logging_callback_manager.py | 10 +++++-- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index dfedc3a3cc9..25648f82a03 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -7,6 +7,7 @@ Callback to log events to a Generic API Endpoint """ import asyncio +import contextlib import json import os import re @@ -158,7 +159,7 @@ class GenericAPILogger(CustomBatchLogger): "endpoint not set for GenericAPILogger, GENERIC_LOGGER_ENDPOINT not found in environment variables" ) - self.headers: dict[str, str] = self._get_headers(headers) + self.headers: dict[str, str] = self.resolve_headers(headers) self.endpoint: str = endpoint self.event_types: list[API_EVENT_TYPES] | None = event_types self.callback_name: str | None = callback_name @@ -190,10 +191,30 @@ class GenericAPILogger(CustomBatchLogger): ######################################################### self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) - asyncio.create_task(self.periodic_flush()) + self._flush_task = asyncio.create_task(self._run_periodic_flush()) self.log_queue: list[dict | StandardLoggingPayload] = [] - def _get_headers(self, headers: dict | None = None): + async def _run_periodic_flush(self) -> None: + """Run the periodic flush loop; on cancellation, make a best-effort final + flush of any buffered logs before exiting. Prevents dropping the last batch + when this logger is evicted from the cache and cancelled via shutdown().""" + try: + await self.periodic_flush() + except asyncio.CancelledError: + # suppress(Exception) leaves a nested CancelledError (a BaseException) free to + # propagate, so a second cancellation during the final flush is still honored. + with contextlib.suppress(Exception): + await self.flush_queue() + raise + + def shutdown(self) -> None: + """Cancel the background flush task. Called by LoggingCallbackManager when this + logger is evicted from the cache so its periodic_flush task does not leak.""" + if not self._flush_task.done(): + self._flush_task.cancel() + + @staticmethod + def resolve_headers(headers: dict | None = None) -> dict[str, str]: # mutable-ok: caller owns the returned headers """ Get headers for the Generic API Logger diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9b612993a69..423786e8850 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -204,15 +204,19 @@ class LoggingCallbackManager: if ( isinstance(cached_logger, GenericAPILogger) and cached_logger.endpoint == endpoint - and cached_logger.headers == headers + and cached_logger.headers == GenericAPILogger.resolve_headers(headers) and cached_logger.event_types == event_types - and cached_logger.log_format == log_format + and cached_logger.log_format == (log_format if log_format is not None else "json_array") and cached_logger.max_retries == max_retries and cached_logger.retry_delay == retry_delay and cached_logger.timeout == timeout ): return cached_logger + # Cache miss (config changed or first creation). Build the replacement + # first, so that an invalid config raising in the constructor cannot leave + # the still-cached logger cancelled. Only after a successful build do we + # retire the evicted logger's background flush task. new_logger = GenericAPILogger( endpoint=endpoint, headers=headers, @@ -222,6 +226,8 @@ class LoggingCallbackManager: retry_delay=retry_delay, timeout=timeout, ) + if isinstance(cached_logger, GenericAPILogger): + cached_logger.shutdown() _generic_api_logger_cache[callback] = new_logger return new_logger From e2a632f8dab86cabed7443b42f1729af6ebafcd0 Mon Sep 17 00:00:00 2001 From: Hoya <96610382+Jeongho0805@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:38:24 +0900 Subject: [PATCH 2/2] test(logging): cover GenericAPILogger cache reuse, shutdown, and header rotation --- .../test_logging_callback_manager.py | 171 ++++++++++++++++++ 1 file changed, 171 insertions(+) create mode 100644 tests/test_litellm/litellm_core_utils/test_logging_callback_manager.py diff --git a/tests/test_litellm/litellm_core_utils/test_logging_callback_manager.py b/tests/test_litellm/litellm_core_utils/test_logging_callback_manager.py new file mode 100644 index 00000000000..96b3f100828 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_logging_callback_manager.py @@ -0,0 +1,171 @@ +""" +Tests for the GenericAPILogger cache-resolution path in +LoggingCallbackManager._add_custom_callback_generic_api_str. + +Covers: + - Cache hit reuses the same logger instance across repeated resolutions + - Invalid empty log_format still raises ValueError + - Genuine config change recreates the logger and cancels the old flush task + - Header rotation is compared on the effective headers, not the raw config ones +""" + +import asyncio +import contextlib + +import pytest + +import litellm +from litellm.litellm_core_utils.logging_callback_manager import ( + GenericAPILogger, + LoggingCallbackManager, + _generic_api_logger_cache, +) + + +@pytest.fixture(autouse=True) +def callback_settings(monkeypatch): + settings = {} + monkeypatch.setattr(litellm, "callback_settings", settings) + _generic_api_logger_cache.clear() + yield settings + _generic_api_logger_cache.clear() + + +class TestGenericAPILoggerCaching: + @pytest.mark.asyncio + async def test_generic_api_logger_reused_on_repeated_resolution(self, callback_settings): + callback_settings["cb"] = { + "callback_type": "generic_api", + "endpoint": "http://127.0.0.1:9/x", + "headers": {"Authorization": "Bearer t"}, + } + + resolved = [LoggingCallbackManager._add_custom_callback_generic_api_str("cb") for _ in range(5)] + + try: + assert all(isinstance(logger, GenericAPILogger) for logger in resolved) + assert all(logger is resolved[0] for logger in resolved) + finally: + resolved[0].shutdown() + + @pytest.mark.asyncio + async def test_generic_api_logger_empty_log_format_still_raises(self, callback_settings): + callback_settings["cb"] = { + "callback_type": "generic_api", + "endpoint": "http://127.0.0.1:9/x", + "headers": {"Authorization": "Bearer t"}, + "log_format": "", + } + + with pytest.raises(ValueError, match="Invalid log_format"): + LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + @pytest.mark.asyncio + async def test_generic_api_logger_recreated_and_old_task_cancelled_on_config_change(self, callback_settings): + callback_settings["cb"] = { + "callback_type": "generic_api", + "endpoint": "http://127.0.0.1:9/a", + "headers": {"Authorization": "Bearer t"}, + } + logger_a = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + callback_settings["cb"]["endpoint"] = "http://127.0.0.1:9/b" + logger_b = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + try: + assert logger_a is not logger_b + + with contextlib.suppress(asyncio.CancelledError): + await logger_a._flush_task + assert logger_a._flush_task.cancelled() + finally: + logger_b.shutdown() + + @pytest.mark.asyncio + async def test_bad_replacement_config_does_not_evict_existing_logger(self, callback_settings): + callback_settings["cb"] = { + "callback_type": "generic_api", + "endpoint": "http://127.0.0.1:9/a", + "headers": {"Authorization": "Bearer t"}, + } + logger_a = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + flush_task_a = logger_a._flush_task + + callback_settings["cb"]["log_format"] = "bad_format" + + try: + with pytest.raises(ValueError, match="Invalid log_format"): + LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + assert flush_task_a.cancelled() is False + assert _generic_api_logger_cache["cb"] is logger_a + finally: + logger_a.shutdown() + + @pytest.mark.asyncio + async def test_env_header_rotation_recreates_logger(self, callback_settings, monkeypatch): + monkeypatch.setenv("GENERIC_LOGGER_HEADERS", "Authorization=Bearer old") + callback_settings["cb"] = { + "callback_type": "generic_api", + "endpoint": "http://127.0.0.1:9/a", + "headers": {"X-Static": "s"}, + } + logger_a = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + monkeypatch.setenv("GENERIC_LOGGER_HEADERS", "Authorization=Bearer new") + logger_b = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + try: + assert logger_a is not logger_b + assert logger_b.headers["Authorization"] == "Bearer new" + + with contextlib.suppress(asyncio.CancelledError): + await logger_a._flush_task + assert logger_a._flush_task.cancelled() + finally: + logger_b.shutdown() + + @pytest.mark.asyncio + async def test_env_header_shadowed_by_config_reuses_logger(self, callback_settings, monkeypatch): + monkeypatch.setenv("GENERIC_LOGGER_HEADERS", "Authorization=Bearer old") + callback_settings["cb"] = { + "callback_type": "generic_api", + "endpoint": "http://127.0.0.1:9/a", + "headers": {"Authorization": "Bearer from-config"}, + } + logger_a = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + monkeypatch.setenv("GENERIC_LOGGER_HEADERS", "Authorization=Bearer new") + logger_b = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + try: + assert logger_a is logger_b + finally: + logger_a.shutdown() + + @pytest.mark.asyncio + async def test_final_flush_error_on_cancellation_is_swallowed(self, callback_settings): + callback_settings["cb"] = { + "callback_type": "generic_api", + "endpoint": "http://127.0.0.1:9/a", + "headers": {"Authorization": "Bearer t"}, + } + logger = LoggingCallbackManager._add_custom_callback_generic_api_str("cb") + + flushed = asyncio.Event() + + async def _raise_on_flush(): + flushed.set() + raise Exception("boom") + + logger.flush_queue = _raise_on_flush + # Let the flush task reach its first await; cancelling a task that never started + # skips _run_periodic_flush entirely and the final-flush path goes untested. + await asyncio.sleep(0) + + logger.shutdown() + + with contextlib.suppress(asyncio.CancelledError): + await logger._flush_task + assert flushed.is_set() + assert logger._flush_task.cancelled() is True