Merge e2a632f8da into 3746ba58d7

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
 

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
 

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