From 27873fcf36119e7620f7f34d144bdedb5d556cf9 Mon Sep 17 00:00:00 2001 From: Jaeyeon Kim Date: Sun, 1 Nov 2026 22:07:24 +0000 Subject: [PATCH 1/2] feat(proxy): run callback setup and cleanup hooks in lifespan Signed-off-by: Jaeyeon Kim --- litellm/integrations/custom_logger.py | 17 +++- .../proxy/common_utils/callback_lifecycle.py | 33 +++++++ .../guardrails/guardrail_hooks/presidio.py | 24 +++++ litellm/proxy/proxy_server.py | 4 + .../common_utils/test_callback_lifecycle.py | 92 +++++++++++++++++++ 5 files changed, 169 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/common_utils/callback_lifecycle.py create mode 100644 tests/unit/proxy/common_utils/test_callback_lifecycle.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 0579eb15baf..371b8dd4b03 100644 --- a/litellm/integrations/custom_logger.py +++ b/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]: """ diff --git a/litellm/proxy/common_utils/callback_lifecycle.py b/litellm/proxy/common_utils/callback_lifecycle.py new file mode 100644 index 00000000000..a05e64d2981 --- /dev/null +++ b/litellm/proxy/common_utils/callback_lifecycle.py @@ -0,0 +1,33 @@ +"""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.""" + seen: set[int] = set() + for callback in tuple(callbacks): + if isinstance(callback, CustomLogger) and id(callback) not in seen: + seen.add(id(callback)) + callback.setup_proxy(app) + # A hook may have added routes after another startup component built OpenAPI. + app.openapi_schema = None + + +async def close_proxy_callbacks(callbacks: Iterable[object]) -> None: + """One failing callback must not prevent the remaining cleanup hooks.""" + seen: set[int] = set() + for callback in tuple(callbacks): + if isinstance(callback, CustomLogger) and id(callback) not in seen: + seen.add(id(callback)) + try: + await callback.async_close() + except Exception: + verbose_proxy_logger.exception("Error closing callback %s", type(callback).__name__) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 458528e7e13..a80dcf3786f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/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.""" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5fcbfb2564..f258cfb5f5a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -482,6 +482,7 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l 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 @@ -1910,6 +1911,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 @@ -1955,6 +1957,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: diff --git a/tests/unit/proxy/common_utils/test_callback_lifecycle.py b/tests/unit/proxy/common_utils/test_callback_lifecycle.py new file mode 100644 index 00000000000..af4930c3554 --- /dev/null +++ b/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() From 3687b4aff1243b965568f345b6dbc282b7a56a44 Mon Sep 17 00:00:00 2001 From: Jaeyeon Kim Date: Sun, 1 Nov 2026 22:28:15 +0000 Subject: [PATCH 2/2] fix(proxy): satisfy callback lifecycle lint rules Signed-off-by: Jaeyeon Kim --- .../proxy/common_utils/callback_lifecycle.py | 27 +++++++++---------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/common_utils/callback_lifecycle.py b/litellm/proxy/common_utils/callback_lifecycle.py index a05e64d2981..84c6319c3a1 100644 --- a/litellm/proxy/common_utils/callback_lifecycle.py +++ b/litellm/proxy/common_utils/callback_lifecycle.py @@ -12,22 +12,19 @@ if TYPE_CHECKING: def setup_proxy_callbacks(app: "FastAPI", callbacks: Iterable[object]) -> None: """Run each configured logger once; a failed setup must fail startup.""" - seen: set[int] = set() - for callback in tuple(callbacks): - if isinstance(callback, CustomLogger) and id(callback) not in seen: - seen.add(id(callback)) - callback.setup_proxy(app) - # A hook may have added routes after another startup component built OpenAPI. - app.openapi_schema = None + 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.""" - seen: set[int] = set() - for callback in tuple(callbacks): - if isinstance(callback, CustomLogger) and id(callback) not in seen: - seen.add(id(callback)) - try: - await callback.async_close() - except Exception: - verbose_proxy_logger.exception("Error closing callback %s", type(callback).__name__) + 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())