feat: add setup_proxy(app) hook for CustomLogger callbacks

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

litellm/integrations/custom_logger.py

@@ -82,6 +82,18 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
         self.turn_off_message_logging = turn_off_message_logging
         pass
 
+    def setup_proxy(self, app: Any) -> None:
+        """
+        Called once during proxy startup after all callbacks are initialized.
+
+        Override this method to interact with the FastAPI app during startup,
+        e.g., to add middleware, mount additional routes, or register shutdown hooks.
+
+        Args:
+            app: The FastAPI application instance.
+        """
+        pass
+
     @staticmethod
     def get_callback_env_vars(callback_name: Optional[str] = None) -> List[str]:
         """

litellm/proxy/proxy_server.py

@@ -749,6 +749,27 @@ async def _initialize_shared_aiohttp_session():
         return None
 
 
+def _invoke_callback_setup_proxy(app: FastAPI) -> None:
+    """
+    Call setup_proxy(app) on all registered CustomLogger callbacks.
+
+    This gives callbacks an opportunity to interact with the FastAPI app
+    during startup, e.g., to add middleware or mount additional routes.
+    """
+    from litellm.integrations.custom_logger import CustomLogger
+
+    for callback in litellm.callbacks:
+        if isinstance(callback, CustomLogger) and hasattr(callback, "setup_proxy"):
+            try:
+                callback.setup_proxy(app)
+            except Exception as e:
+                verbose_proxy_logger.warning(
+                    "Error calling setup_proxy on callback %s: %s",
+                    type(callback).__name__,
+                    str(e),
+                )
+
+
 @asynccontextmanager
 async def proxy_startup_event(app: FastAPI):  # noqa: PLR0915
     global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session
@@ -813,6 +834,9 @@ async def proxy_startup_event(app: FastAPI):  # noqa: PLR0915
             if isinstance(worker_config, dict):
                 await initialize(**worker_config)
 
+    # Call setup_proxy(app) on all registered CustomLogger callbacks
+    _invoke_callback_setup_proxy(app)
+
     # check if DATABASE_URL in environment - load from there
     if prisma_client is None:
         _db_url: Optional[str] = get_secret("DATABASE_URL", None)  # type: ignore

tests/test_litellm/proxy/test_callback_setup_proxy.py

@@ -0,0 +1,117 @@
+"""
+Tests for the CustomLogger.setup_proxy(app) hook.
+
+Verifies that:
+1. CustomLogger has a setup_proxy method (no-op default)
+2. _invoke_callback_setup_proxy calls setup_proxy on all CustomLogger callbacks
+3. Non-CustomLogger callbacks are skipped
+4. Exceptions in setup_proxy are caught and logged (don't crash startup)
+"""
+
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy.proxy_server import _invoke_callback_setup_proxy
+
+
+class TestCustomLoggerSetupProxy:
+    """Tests for the setup_proxy method on CustomLogger."""
+
+    def test_custom_logger_has_setup_proxy_method(self):
+        """CustomLogger base class should have a setup_proxy method."""
+        logger = CustomLogger()
+        assert hasattr(logger, "setup_proxy")
+        assert callable(logger.setup_proxy)
+
+    def test_setup_proxy_default_is_noop(self):
+        """Default setup_proxy should do nothing and not raise."""
+        logger = CustomLogger()
+        mock_app = MagicMock()
+        logger.setup_proxy(mock_app)  # Should not raise
+
+
+class TestInvokeCallbackSetupProxy:
+    """Tests for _invoke_callback_setup_proxy function."""
+
+    def test_calls_setup_proxy_on_custom_logger_callbacks(self):
+        """setup_proxy should be called on each CustomLogger in litellm.callbacks."""
+
+        class MyCallback(CustomLogger):
+            def __init__(self):
+                super().__init__()
+                self.setup_called_with = None
+
+            def setup_proxy(self, app):
+                self.setup_called_with = app
+
+        callback = MyCallback()
+        mock_app = MagicMock()
+
+        original_callbacks = litellm.callbacks
+        litellm.callbacks = [callback]
+        try:
+            _invoke_callback_setup_proxy(mock_app)
+            assert callback.setup_called_with is mock_app
+        finally:
+            litellm.callbacks = original_callbacks
+
+    def test_skips_non_custom_logger_callbacks(self):
+        """String callbacks and non-CustomLogger objects should be skipped."""
+        mock_app = MagicMock()
+
+        original_callbacks = litellm.callbacks
+        litellm.callbacks = ["langfuse", "sentry", 42]
+        try:
+            # Should not raise
+            _invoke_callback_setup_proxy(mock_app)
+        finally:
+            litellm.callbacks = original_callbacks
+
+    def test_exception_in_setup_proxy_is_caught(self):
+        """If setup_proxy raises, it should be caught and logged, not crash."""
+
+        class BadCallback(CustomLogger):
+            def setup_proxy(self, app):
+                raise RuntimeError("setup failed")
+
+        bad_callback = BadCallback()
+        good_callback = CustomLogger()
+        mock_app = MagicMock()
+
+        original_callbacks = litellm.callbacks
+        litellm.callbacks = [bad_callback, good_callback]
+        try:
+            # Should not raise even though bad_callback.setup_proxy raises
+            _invoke_callback_setup_proxy(mock_app)
+        finally:
+            litellm.callbacks = original_callbacks
+
+    def test_multiple_callbacks_all_called(self):
+        """All CustomLogger callbacks should have setup_proxy called."""
+        call_order = []
+
+        class Callback1(CustomLogger):
+            def setup_proxy(self, app):
+                call_order.append("cb1")
+
+        class Callback2(CustomLogger):
+            def setup_proxy(self, app):
+                call_order.append("cb2")
+
+        mock_app = MagicMock()
+        original_callbacks = litellm.callbacks
+        litellm.callbacks = [Callback1(), "langfuse", Callback2()]
+        try:
+            _invoke_callback_setup_proxy(mock_app)
+            assert call_order == ["cb1", "cb2"]
+        finally:
+            litellm.callbacks = original_callbacks