Merge 3687b4aff1 into 9ffb82e139

litellm/integrations/custom_logger.py

@@ -28,7 +28,7 @@ from litellm.types.utils import (
 )
 
 if TYPE_CHECKING:
-    from fastapi import HTTPException
+    from fastapi import FastAPI, HTTPException
     from opentelemetry.trace import Span as _Span
 
     from litellm.caching.caching import DualCache
@@ -100,6 +100,21 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
         self.message_logging = message_logging
         self.turn_off_message_logging = turn_off_message_logging
 
+    def setup_proxy(self, app: "FastAPI") -> None:
+        """Initialize routes or app state once per proxy worker lifespan.
+
+        Runs after startup callback initialization. Exceptions abort startup.
+        The ASGI middleware stack is already built: do not add middleware here.
+        Hooks must tolerate another lifespan entry on the same application.
+        """
+
+    async def async_close(self) -> None:
+        """Release callback resources after the proxy has drained request logging.
+
+        Implementations should be idempotent and close loop-bound resources on
+        their owning loop. SDK-only users remain responsible for calling this.
+        """
+
     @staticmethod
     def get_callback_env_vars(callback_name: str | None = None) -> list[str]:
         """

litellm/proxy/common_utils/callback_lifecycle.py

@@ -0,0 +1,30 @@
+"""Startup hooks for configured callback instances."""
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+from litellm._logging import verbose_proxy_logger
+from litellm.integrations.custom_logger import CustomLogger
+
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
+
+def setup_proxy_callbacks(app: "FastAPI", callbacks: Iterable[object]) -> None:
+    """Run each configured logger once; a failed setup must fail startup."""
+    for callback in _unique_callbacks(callbacks):
+        callback.setup_proxy(app)
+    app.openapi_schema = None  # rebind-ok: routes added by startup hooks invalidate FastAPI's cached schema
+
+
+async def close_proxy_callbacks(callbacks: Iterable[object]) -> None:
+    """One failing callback must not prevent the remaining cleanup hooks."""
+    for callback in _unique_callbacks(callbacks):
+        try:
+            await callback.async_close()
+        except Exception:  # noqa: BLE001  # Third-party cleanup failures must not skip other callbacks
+            verbose_proxy_logger.exception("Error closing callback %s", type(callback).__name__)
+
+
+def _unique_callbacks(callbacks: Iterable[object]) -> tuple[CustomLogger, ...]:
+    return tuple({id(callback): callback for callback in callbacks if isinstance(callback, CustomLogger)}.values())

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -313,6 +313,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
         self.presidio_analyze_chunk_size_bytes: int = self._coerce_analyze_chunk_size(presidio_analyze_chunk_size_bytes)
         # Shared HTTP session to prevent memory leaks (issue #14540)
         self._http_session: aiohttp.ClientSession | None = None
+        self._http_session_loop: asyncio.AbstractEventLoop | None = None
         # Lock to prevent race conditions when creating session under concurrent load
         # Note: asyncio.Lock() can be created without an event loop; it only needs one when awaited
         self._session_lock: asyncio.Lock = asyncio.Lock()
@@ -401,6 +402,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             async with self._session_lock:
                 if self._http_session is None or self._http_session.closed:
                     self._http_session = aiohttp.ClientSession()
+                    self._http_session_loop = current_loop
                 yield self._http_session
         else:
             # Background thread/loop -> use loop-bound session cache
@@ -421,6 +423,28 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
                 await session.close()
         self._loop_sessions.clear()
 
+    @staticmethod
+    async def _close_owned_session(session: aiohttp.ClientSession, loop: asyncio.AbstractEventLoop | None) -> None:
+        if session.closed:
+            return
+        if loop is None or loop is asyncio.get_running_loop():
+            await session.close()
+        elif loop.is_running():
+            await asyncio.wrap_future(asyncio.run_coroutine_threadsafe(session.close(), loop))
+        else:
+            raise RuntimeError("Presidio session owner loop stopped before callback cleanup")
+
+    async def async_close(self) -> None:
+        """Release sessions after requests/logging drain, on their owning loops."""
+        async with self._session_lock:
+            if self._http_session is not None:
+                await self._close_owned_session(self._http_session, self._http_session_loop)
+                self._http_session = None
+                self._http_session_loop = None
+        for loop, session in tuple(self._loop_sessions.items()):
+            await self._close_owned_session(session, loop)
+            self._loop_sessions.pop(loop, None)
+
     def __del__(self):
         """Cleanup: we try to close, but doing async cleanup in __del__ is risky."""
 

litellm/proxy/proxy_server.py

@@ -517,6 +517,7 @@ from litellm.proxy.common_utils.anthropic_error_payload import (
 from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
     AuthCacheInvalidationSubscriber,
 )
+from litellm.proxy.common_utils.callback_lifecycle import close_proxy_callbacks, setup_proxy_callbacks
 from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
 from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body
 from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes
@@ -1973,6 +1974,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
             async with AsyncExitStack() as admin_mcp_stack:
                 try:
                     await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app))
+                    setup_proxy_callbacks(app, litellm.callbacks)
                     gc.freeze()
                     verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count())
                     yield state
@@ -2018,6 +2020,8 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
 
             await proxy_config.stop_auth_cache_invalidation_subscriber()
 
+            await close_proxy_callbacks(litellm.callbacks)
+            await litellm.close_litellm_async_clients()
             await proxy_shutdown_event(worker_heartbeat=worker_heartbeat)
 
             if prometheus_multiproc_dir:

tests/unit/proxy/common_utils/test_callback_lifecycle.py

@@ -0,0 +1,92 @@
+from contextlib import asynccontextmanager
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy.common_utils.callback_lifecycle import close_proxy_callbacks, setup_proxy_callbacks
+
+
+class RouteCallback(CustomLogger):
+    def setup_proxy(self, app):
+        @app.get("/callback-route")
+        def callback_route():
+            return {"source": "callback"}
+
+
+def test_lifespan_hook_adds_route_and_refreshes_openapi():
+    callback = RouteCallback()
+
+    @asynccontextmanager
+    async def lifespan(app):
+        setup_proxy_callbacks(app, [callback, "unresolved-name", callback])
+        yield
+
+    app = FastAPI(lifespan=lifespan)
+    assert "/callback-route" not in app.openapi()["paths"]
+    with TestClient(app) as client:
+        assert client.get("/callback-route").json() == {"source": "callback"}
+        assert "/callback-route" in client.get("/openapi.json").json()["paths"]
+    assert sum(route.path == "/callback-route" for route in app.routes) == 1
+
+
+def test_failed_hook_aborts_lifespan_startup():
+    class FailedCallback(CustomLogger):
+        def setup_proxy(self, app):
+            raise ValueError("invalid callback configuration")
+
+    @asynccontextmanager
+    async def lifespan(app):
+        setup_proxy_callbacks(app, [FailedCallback()])
+        yield
+
+    with pytest.raises(ValueError, match="invalid callback configuration"):
+        with TestClient(FastAPI(lifespan=lifespan)):
+            pytest.fail("startup should not succeed")
+
+
+@pytest.mark.asyncio
+async def test_cleanup_closes_real_sessions_after_another_callback_fails():
+    from litellm.proxy.guardrails.guardrail_hooks.presidio import OPTIONAL_PresidioPIIMasking
+
+    class FailedCallback(CustomLogger):
+        async def async_close(self):
+            raise ValueError("cleanup failed")
+
+    guardrail = OPTIONAL_PresidioPIIMasking(mock_testing=True)
+    async with guardrail._get_session_iterator() as session:
+        assert not session.closed
+    await close_proxy_callbacks([FailedCallback(), guardrail, guardrail, "name"])
+    assert session.closed
+    assert guardrail._http_session is None
+    await close_proxy_callbacks([guardrail])
+    assert guardrail._http_session is None
+
+
+@pytest.mark.asyncio
+async def test_presidio_cleanup_uses_background_session_owner_loop():
+    import asyncio
+    import threading
+
+    from litellm.proxy.guardrails.guardrail_hooks.presidio import OPTIONAL_PresidioPIIMasking
+
+    guardrail = OPTIONAL_PresidioPIIMasking(mock_testing=True)
+    loop = asyncio.new_event_loop()
+    thread = threading.Thread(target=loop.run_forever)
+    thread.start()
+
+    async def get_session():
+        async with guardrail._get_session_iterator() as session:
+            return session
+
+    try:
+        session = await asyncio.wrap_future(asyncio.run_coroutine_threadsafe(get_session(), loop))
+        assert not session.closed
+        await guardrail.async_close()
+        assert session.closed
+        assert guardrail._loop_sessions == {}
+    finally:
+        loop.call_soon_threadsafe(loop.stop)
+        thread.join()
+        loop.close()