diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index a42187b3a44..3cf0cd0be1b 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -19,6 +19,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_utils import _get_request_ip_address +from litellm.proxy.auth.resolvers.exceptions import NoDatabaseConnectionError from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes @@ -76,13 +77,18 @@ class UserAPIKeyAuthExceptionHandler: """ from litellm.proxy.proxy_server import ( general_settings, + is_prisma_initial_connect_recovery_pending, + prisma_client, proxy_logging_obj, ) - if ( - PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - and PrismaDBExceptionHandler.is_database_connection_error(e) - ): + configured_database_without_client: Final = isinstance(e, NoDatabaseConnectionError) and ( + is_prisma_initial_connect_recovery_pending() or prisma_client is not None + ) + database_unavailable: Final = ( + PrismaDBExceptionHandler.is_database_connection_error(e) or configured_database_without_client + ) + if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() and database_unavailable: # log this as a DB failure on prometheus proxy_logging_obj.service_logging_obj.service_failure_hook( service=ServiceTypes.DB, @@ -183,7 +189,7 @@ class UserAPIKeyAuthExceptionHandler: ) elif isinstance(e, ProxyException): raise e - if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e) or configured_database_without_client: raise ProxyException( message=( "Service Unavailable, the authentication database is " diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py index 832baf03432..c548ae2adf0 100644 --- a/litellm/proxy/auth/resolvers/store.py +++ b/litellm/proxy/auth/resolvers/store.py @@ -100,13 +100,13 @@ class IdentityStore: return principal.source_key async def _resolve_key(self, hashed_token: str) -> UserAPIKeyAuth: - if self._prisma is None: - raise NoDatabaseConnectionError() - cached: Final = await self._cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth) if cached is not None: return _copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + if self._prisma is None: + raise NoDatabaseConnectionError() + if self._check_cache_only: raise KeyNotInCacheError(hashed_token) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c23cde6052e..a7be4cdecf9 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -7,10 +7,11 @@ import time import traceback from collections.abc import Iterable, Mapping from datetime import datetime, timedelta -from typing import Any, Final, Literal, TypedDict, cast +from typing import Any, Final, Literal, NotRequired, TypedDict, cast import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger, verbose_proxy_logger @@ -1317,6 +1318,7 @@ async def health_license_endpoint( class DBHealthCache(TypedDict): status: str last_updated: datetime + fail_open_safe: NotRequired[ReadOnly[bool]] db_health_cache: DBHealthCache = {"status": "unknown", "last_updated": datetime.now()} @@ -1333,14 +1335,26 @@ async def _db_health_readiness_check(): return db_health_cache if prisma_client is None: - db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} + db_health_cache = { + "status": "disconnected", + "last_updated": datetime.now(), + "fail_open_safe": False, + } return db_health_cache await prisma_client.health_check() - db_health_cache = {"status": "connected", "last_updated": datetime.now()} + db_health_cache = { + "status": "connected", + "last_updated": datetime.now(), + "fail_open_safe": False, + } return db_health_cache except Exception as e: - db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} + db_health_cache = { + "status": "disconnected", + "last_updated": datetime.now(), + "fail_open_safe": PrismaDBExceptionHandler.is_database_connection_error(e), + } if PrismaDBExceptionHandler.is_database_transport_error(e): try: verbose_proxy_logger.warning("_db_health_readiness_check: health_check failed, attempting reconnect") @@ -1350,13 +1364,26 @@ async def _db_health_readiness_check(): db_health_cache = { "status": "connected", "last_updated": datetime.now(), + "fail_open_safe": False, } return db_health_cache - except Exception: + except Exception as reconnect_error: verbose_proxy_logger.error("_db_health_readiness_check: reconnect failed") + return { + "status": "disconnected", + "last_updated": db_health_cache["last_updated"], + "fail_open_safe": PrismaDBExceptionHandler.is_database_connection_error(reconnect_error), + } return db_health_cache +def _readiness_can_fail_open(db_health_status: DBHealthCache) -> bool: + return ( + db_health_status.get("fail_open_safe") is True + and PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + ) + + @router.get( "/settings", tags=["health"], @@ -1486,7 +1513,11 @@ async def _get_health_readiness_details( """ Detailed health payload for authenticated diagnostics. """ - from litellm.proxy.proxy_server import prisma_client, version + from litellm.proxy.proxy_server import ( + is_prisma_initial_connect_recovery_pending, + prisma_client, + version, + ) try: # get success callback @@ -1525,11 +1556,13 @@ async def _get_health_readiness_details( # check DB if prisma_client is not None: # if db passed in, check if it's connected db_health_status: Final = await _db_health_readiness_check() - # A configured DB that is not reachable means the worker cannot - # serve requests that depend on persisted state (keys, budgets, - # spend logs). Return 503 so orchestrators take this pod out of - # rotation; "Not connected" (no DB configured at all) stays 200. - if response is not None and db_health_status["status"] != "connected": + # A configured DB that is not reachable is unhealthy unless the + # operator explicitly opted into fail-open request handling. + if ( + response is not None + and db_health_status["status"] != "connected" + and _readiness_can_fail_open(db_health_status) is False + ): response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", @@ -1543,9 +1576,16 @@ async def _get_health_readiness_details( "show_no_redis_warning": show_no_redis_warning, } else: + initial_recovery_pending: Final = is_prisma_initial_connect_recovery_pending() + if ( + response is not None + and initial_recovery_pending + and PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() is False + ): + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", - "db": "Not connected", + "db": "disconnected" if initial_recovery_pending else "Not connected", "cache": cache_type, "litellm_version": version, "success_callbacks": success_callback_names, @@ -1614,13 +1654,16 @@ async def _resolve_public_readiness_db(response: Response) -> str: 503 when a configured DB is unreachable. Mirrors the legacy values: "Not connected" (no DB configured), "connected", "disconnected". """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import is_prisma_initial_connect_recovery_pending, prisma_client if prisma_client is None: - return "Not connected" + initial_recovery_pending: Final = is_prisma_initial_connect_recovery_pending() + if initial_recovery_pending and PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() is False: + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE + return "disconnected" if initial_recovery_pending else "Not connected" db_health_status: Final = await _db_health_readiness_check() - if db_health_status["status"] != "connected": + if db_health_status["status"] != "connected" and _readiness_can_fail_open(db_health_status) is False: response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return db_health_status["status"] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e55b254ab8e..dfa3473682a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16,7 +16,7 @@ import threading import time import traceback import warnings -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping, MutableMapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -872,6 +872,7 @@ def cleanup_router_config_variables(): health_check_interval = None health_check_concurrency = None prisma_client = None + _initial_prisma_connect_recovery_state.reset() async def _flush_spend_logs_queue_on_shutdown() -> None: @@ -1080,49 +1081,11 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: database_url=_db_url, proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, + recovered_client_initializer=ProxyStartupEvent.initialize_recovered_prisma_services, ) if prisma_client is not None: - - async def _run_pw_migration(): - try: - result: Final = await migrate_passwords_to_scrypt_async(prisma_client) - verbose_proxy_logger.info("Password migration: %s", result) - except Exception as e: - verbose_proxy_logger.warning("Password migration skipped: %s", e) - - asyncio.create_task(_run_pw_migration()) - - async def _run_agent_grant_id_migration() -> None: - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry, - object_permission_table, - ) - - for attempt in range(3): - try: - result = await global_agent_registry.migrate_legacy_grant_ids( - table=object_permission_table(prisma_client) - ) - if result.rewritten: - verbose_proxy_logger.info( - "Rewrote %s object_permission rows from legacy config agent ids", result.rewritten - ) - if result.missed == 0: - return - verbose_proxy_logger.warning( - "Legacy agent grant id migration attempt %s/3 left %s rows unmigrated", - attempt + 1, - result.missed, - ) - except Exception as e: # noqa: BLE001 # startup task must survive any DB error and retry - verbose_proxy_logger.warning( - "Legacy agent grant id migration attempt %s/3 failed: %s", attempt + 1, e - ) - if attempt < 2: - await asyncio.sleep(5) - - asyncio.create_task(_run_agent_grant_id_migration()) + ProxyStartupEvent.start_prisma_migrations(prisma_client) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before @@ -1255,11 +1218,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: # `llm_router.adaptive_routers` is empty. Per-router DB state is loaded # lazily by the flusher on first tick (see `_state_loaded` flag) so # hot-reloaded routers also get their persisted priors. - if llm_router is not None and getattr(llm_router, "adaptive_routers", None): - for _tagged_routers in llm_router.adaptive_routers.values(): - for _tagged in _tagged_routers: - await _tagged.strategy.load_state_from_db(prisma_client) - _tagged.strategy._state_loaded = True + await ProxyStartupEvent.load_adaptive_router_state(prisma_client) asyncio.create_task(_adaptive_router_flusher_loop()) ## [Optional] Initialize dd tracer @@ -1271,6 +1230,8 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: ## Initialize shared aiohttp session for connection reuse shared_aiohttp_session = await _initialize_shared_aiohttp_session() + ProxyStartupEvent.start_initial_prisma_connect_recovery() + # End of startup event yield @@ -1287,6 +1248,8 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error closing shared aiohttp session: %s", e) + await ProxyStartupEvent.stop_initial_prisma_connect_recovery() + # Shutdown event - stop RDS IAM token refresh background task if ( prisma_client is not None @@ -1311,7 +1274,9 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await proxy_config.stop_auth_cache_invalidation_subscriber() - await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) + await proxy_shutdown_event( + worker_heartbeat=worker_heartbeat or _initial_prisma_connect_recovery_state.recovered_worker_heartbeat + ) def _generate_stable_operation_id(route: "APIRoute") -> str: @@ -2244,7 +2209,30 @@ worker_config: Final = None master_key: str | None = None otel_logging = False prisma_client: PrismaClient | None = None + + +class _InitialPrismaConnectRecoveryState: + def __init__(self) -> None: + self.candidate: PrismaClient | None = None + self.task: asyncio.Task[None] | None = None + self.recovered_initializer: Callable[[PrismaClient], Awaitable[None]] | None = None + self.recovered_worker_heartbeat: ProxyWorkerHeartbeat | None = None + + def reset(self) -> None: + self.candidate = None + self.task = None + self.recovered_initializer = None + self.recovered_worker_heartbeat = None + + +_initial_prisma_connect_recovery_state: Final = _InitialPrismaConnectRecoveryState() shared_aiohttp_session: Optional["ClientSession"] = None # Global shared session for connection reuse + + +def is_prisma_initial_connect_recovery_pending() -> bool: + return _initial_prisma_connect_recovery_state.candidate is not None + + user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -9561,26 +9549,204 @@ class ProxyStartupEvent: ) await _scheduled_fallback_stats() + @classmethod + def start_prisma_migrations(cls, client: PrismaClient) -> None: + async def _run_pw_migration() -> None: + try: + result: Final = await migrate_passwords_to_scrypt_async(client) + verbose_proxy_logger.info("Password migration: %s", result) + except Exception as migration_error: + verbose_proxy_logger.warning("Password migration skipped: %s", migration_error) + + async def _run_agent_grant_id_migration() -> None: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + object_permission_table, + ) + + for attempt in range(3): + try: + result = await global_agent_registry.migrate_legacy_grant_ids(table=object_permission_table(client)) + if result.rewritten: + verbose_proxy_logger.info( + "Rewrote %s object_permission rows from legacy config agent ids", result.rewritten + ) + if result.missed == 0: + return + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 left %s rows unmigrated", + attempt + 1, + result.missed, + ) + except Exception as migration_error: # noqa: BLE001 # startup migration retries every DB failure + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 failed: %s", + attempt + 1, + migration_error, + ) + if attempt < 2: + await asyncio.sleep(5) + + asyncio.create_task(_run_pw_migration()) + asyncio.create_task(_run_agent_grant_id_migration()) + + @classmethod + async def initialize_recovered_prisma_services(cls, client: PrismaClient) -> None: + cls.start_prisma_migrations(client) + + db_coordination_redis_cache: Final = await cls._init_coordination_redis_from_db( + litellm_settings=proxy_config.get_config_state().get("litellm_settings") or {}, + llm_router=llm_router, + ) + if db_coordination_redis_cache is not None: + _set_redis_usage_cache(db_coordination_redis_cache) + proxy_logging_obj.update_values(redis_cache=db_coordination_redis_cache) + + proxy_logging_obj.add_missing_proxy_hooks(llm_router) + + if litellm.max_budget > 0: + cls._add_proxy_budget_to_db() + asyncio.create_task( + cls._warm_global_spend_cache( + user_api_key_cache=user_api_key_cache, + prisma_client=client, + ) + ) + + _initial_prisma_connect_recovery_state.recovered_worker_heartbeat = ( + await cls.initialize_scheduled_background_jobs( + general_settings=general_settings, + prisma_client=client, + proxy_budget_rescheduler_min_time=proxy_budget_rescheduler_min_time, + proxy_budget_rescheduler_max_time=proxy_budget_rescheduler_max_time, + proxy_batch_write_at=proxy_batch_write_at, + proxy_logging_obj=proxy_logging_obj, + ) + ) + await cls._update_default_team_member_budget() + await cls._sync_ui_settings_to_general_settings() + + await cls.load_adaptive_router_state(client) + + @classmethod + async def load_adaptive_router_state(cls, client: PrismaClient | None) -> None: + if llm_router is not None and getattr(llm_router, "adaptive_routers", None): + for tagged_routers in llm_router.adaptive_routers.values(): + for tagged in tagged_routers: + await tagged.strategy.load_state_from_db(client) + tagged.strategy._state_loaded = True + + @classmethod + async def _complete_prisma_client_setup(cls, prisma_client: PrismaClient) -> None: + if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): + await prisma_client.db.start_token_refresh_task() + + asyncio.create_task(prisma_client.check_view_exists()) + asyncio.create_task(prisma_client._set_spend_logs_row_count_in_proxy_state()) + + if hasattr(prisma_client, "start_db_health_watchdog_task"): + await prisma_client.start_db_health_watchdog_task() + + if get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) is not True: + await prisma_client.health_check() + + @classmethod + async def _recover_initial_prisma_connect(cls, candidate: PrismaClient) -> None: + global prisma_client # noqa: PLW0603 - publish the recovered module singleton + + retry_interval_value: Final = getattr(candidate, "_db_health_watchdog_interval_seconds", 30) + retry_interval_seconds: Final = ( + float(retry_interval_value) if isinstance(retry_interval_value, (int, float)) else 30.0 + ) + while _initial_prisma_connect_recovery_state.candidate is candidate: + try: + recovered = await candidate.attempt_db_reconnect( + reason="initial_prisma_connect_failure", + force=True, + ) + if recovered: + if _initial_prisma_connect_recovery_state.candidate is not candidate: + return + await cls._complete_prisma_client_setup(candidate) + cls._initialize_jwt_auth( + general_settings=general_settings, + prisma_client=candidate, + user_api_key_cache=user_api_key_cache, + ) + prisma_client = candidate + _initial_prisma_connect_recovery_state.candidate = None + recovered_initializer = _initial_prisma_connect_recovery_state.recovered_initializer + _initial_prisma_connect_recovery_state.recovered_initializer = None + if recovered_initializer is not None: + try: + await recovered_initializer(candidate) + except Exception as initialization_error: # noqa: BLE001 # keep the recovered client published + verbose_proxy_logger.exception( + "Prisma recovered, but deferred database startup failed: %s", + initialization_error, + ) + verbose_proxy_logger.info("Prisma connected after an initial startup failure") + return + except asyncio.CancelledError: + raise + except Exception as recovery_error: # noqa: BLE001 # retry transient DB and engine failures + verbose_proxy_logger.warning( + "Prisma initial connection recovery failed: %s", + recovery_error, + ) + await asyncio.sleep(retry_interval_seconds) + + @classmethod + def start_initial_prisma_connect_recovery(cls) -> None: + candidate: Final = _initial_prisma_connect_recovery_state.candidate + current_task: Final = _initial_prisma_connect_recovery_state.task + if candidate is None or (current_task is not None and current_task.done() is False): + return + _initial_prisma_connect_recovery_state.task = asyncio.create_task( + cls._recover_initial_prisma_connect(candidate) + ) + + @classmethod + async def stop_initial_prisma_connect_recovery(cls) -> None: + recovery_task: Final = _initial_prisma_connect_recovery_state.task + if recovery_task is not None and recovery_task.done() is False: + recovery_task.cancel() + try: + await recovery_task + except asyncio.CancelledError: + pass + _initial_prisma_connect_recovery_state.candidate = None + _initial_prisma_connect_recovery_state.task = None + _initial_prisma_connect_recovery_state.recovered_initializer = None + @classmethod async def _setup_prisma_client( cls, database_url: str | None, proxy_logging_obj: ProxyLogging, user_api_key_cache: UserApiKeyCache, + recovered_client_initializer: Callable[[PrismaClient], Awaitable[None]] | None = None, ) -> PrismaClient | None: """ - Sets up prisma client - Adds necessary views to proxy """ + candidate_client: PrismaClient | None = None connected_client: PrismaClient | None = None + current_recovery_task: Final = _initial_prisma_connect_recovery_state.task + if current_recovery_task is not None and current_recovery_task.done(): + _initial_prisma_connect_recovery_state.task = None + _initial_prisma_connect_recovery_state.candidate = None + _initial_prisma_connect_recovery_state.recovered_initializer = None + _initial_prisma_connect_recovery_state.recovered_worker_heartbeat = None try: if database_url is None: return None - prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) + candidate_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) try: - await prisma_client.connect() + await candidate_client.connect() except Exception as e: if "P3018" in str(e) or "P3009" in str(e): verbose_proxy_logger.debug("CRITICAL: DATABASE MIGRATION FAILED") @@ -9588,31 +9754,9 @@ class ProxyStartupEvent: verbose_proxy_logger.debug("FIX: Run 'prisma migrate resolve --applied '") raise e - connected_client = prisma_client - - ## Start RDS IAM token refresh background task if enabled ## - # This proactively refreshes IAM tokens before they expire, - # preventing the 15-minute connection failure bug (#16220) - if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): - await prisma_client.db.start_token_refresh_task() - - ## Add necessary views to proxy ## - asyncio.create_task( - prisma_client.check_view_exists() - ) # check if all necessary views exist. Don't block execution - - asyncio.create_task( - prisma_client._set_spend_logs_row_count_in_proxy_state() - ) # set the spend logs row count in proxy state. Don't block execution - - if hasattr(prisma_client, "start_db_health_watchdog_task"): - await prisma_client.start_db_health_watchdog_task() - - # run a health check to ensure the DB is ready - if get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) is not True: - await prisma_client.health_check() - - return prisma_client + connected_client = candidate_client + await cls._complete_prisma_client_setup(candidate_client) + return candidate_client except Exception as e: PrismaDBExceptionHandler.handle_db_exception(e) if connected_client is not None: @@ -9621,6 +9765,9 @@ class ProxyStartupEvent: "The DB health watchdog keeps probing and reconnects once the database recovers.", e, ) + elif candidate_client is not None: + _initial_prisma_connect_recovery_state.candidate = candidate_client + _initial_prisma_connect_recovery_state.recovered_initializer = recovered_client_initializer return connected_client @classmethod diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8cbf5b685fd..5f6512c627a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -783,13 +783,18 @@ class ProxyLogging: self.db_spend_update_writer.redis_update_buffer.redis_cache = redis_cache self.db_spend_update_writer.pod_lock_manager.redis_cache = redis_cache - def _add_proxy_hooks(self, llm_router: Router | None = None): + def add_missing_proxy_hooks(self, llm_router: Router | None = None): """ - Add proxy hooks to litellm.callbacks + Add proxy hooks that have not already been initialized. + + Database-dependent hooks can be skipped during fail-open startup and + installed later when the Prisma client recovers. """ from litellm.proxy.proxy_server import prisma_client for hook in PROXY_HOOKS: + if hook in self.proxy_hook_mapping: + continue proxy_hook = get_proxy_hook(hook) expected_args = inspect.getfullargspec(proxy_hook).args if "prisma_client" in expected_args and prisma_client is None: @@ -804,9 +809,17 @@ class ProxyLogging: passed_in_args["prisma_client"] = prisma_client proxy_hook_obj = cast(CustomLogger, proxy_hook(**passed_in_args)) litellm.logging_callback_manager.add_litellm_callback(proxy_hook_obj) + litellm.logging_callback_manager.add_litellm_success_callback(proxy_hook_obj) + litellm.logging_callback_manager.add_litellm_failure_callback(proxy_hook_obj) + litellm.logging_callback_manager.add_litellm_async_success_callback(proxy_hook_obj) + litellm.logging_callback_manager.add_litellm_async_failure_callback(proxy_hook_obj) self.proxy_hook_mapping[hook] = proxy_hook_obj + def _add_proxy_hooks(self, llm_router: Router | None = None): + """Backward-compatible entry point for existing proxy setup callers.""" + self.add_missing_proxy_hooks(llm_router) + def get_proxy_hook(self, hook: str) -> CustomLogger | None: """ Get a proxy hook from the proxy_hook_mapping @@ -814,7 +827,7 @@ class ProxyLogging: return self.proxy_hook_mapping.get(hook) def _init_litellm_callbacks(self, llm_router: Router | None = None): - self._add_proxy_hooks(llm_router) + self.add_missing_proxy_hooks(llm_router) litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # Track string callbacks and their initialized instances so we can diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 0cdf3500d50..80b39784ddd 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -465,7 +465,7 @@ def test_is_api_route_allowed(route, user_role, expected_result): @pytest.mark.asyncio -async def test_auth_not_connected_to_db(): +async def test_models_auth_uses_restricted_identity_when_not_connected_to_db(): """ ensure requests don't fail when `prisma_client` = None """ @@ -487,7 +487,7 @@ async def test_auth_not_connected_to_db(): ) request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") + request._url = URL(url="/v1/models") valid_token = await user_api_key_auth(request=request, api_key="Bearer " + user_key) print("got valid token", valid_token) diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 90b3b29d919..17fbd05ca42 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -27,8 +27,11 @@ from prisma.errors import ( from litellm._logging import verbose_proxy_logger from litellm.exceptions import BudgetExceededError -from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy import proxy_server as proxy_server_module +from litellm.proxy._types import LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth import auth_exception_handler as auth_exception_handler_module from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler +from litellm.proxy.auth.resolvers.exceptions import NoDatabaseConnectionError @pytest.mark.asyncio @@ -66,6 +69,92 @@ async def test_handle_authentication_error_db_unavailable_connectivity(db_error) assert result.token == "failed-to-connect-to-db" +@pytest.mark.asyncio +async def test_initial_prisma_recovery_nil_client_uses_fail_open_identity(monkeypatch): + handler = UserAPIKeyAuthExceptionHandler() + + monkeypatch.setattr(proxy_server_module, "general_settings", {"allow_requests_on_db_unavailable": True}) + monkeypatch.setattr(proxy_server_module, "is_prisma_initial_connect_recovery_pending", lambda: True) + result = await handler._handle_authentication_error( + NoDatabaseConnectionError(), + MagicMock(), + {}, + "/v1/models", + None, + "cached-jwt", + ) + + assert result.key_name == "failed-to-connect-to-db" + assert result.user_id == "__db_unavailable_fallback__" + assert result.user_role == LitellmUserRoles.INTERNAL_USER + + +@pytest.mark.asyncio +async def test_initial_prisma_recovery_nil_client_returns_503_when_fail_open_disabled(monkeypatch): + handler = UserAPIKeyAuthExceptionHandler() + + monkeypatch.setattr(auth_exception_handler_module, "seed_request_identity", MagicMock()) + monkeypatch.setattr(proxy_server_module.proxy_logging_obj, "post_call_failure_hook", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server_module, "general_settings", {"allow_requests_on_db_unavailable": False}) + monkeypatch.setattr(proxy_server_module, "is_prisma_initial_connect_recovery_pending", lambda: True) + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + NoDatabaseConnectionError(), + MagicMock(), + {}, + "/v1/models", + None, + "cached-jwt", + ) + + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + + +@pytest.mark.asyncio +async def test_initial_prisma_recovery_snapshot_stays_fail_open_after_recovery_completes(monkeypatch): + """A request can retain a nil client while recovery publishes the global client.""" + handler = UserAPIKeyAuthExceptionHandler() + + monkeypatch.setattr(proxy_server_module, "general_settings", {"allow_requests_on_db_unavailable": True}) + monkeypatch.setattr(proxy_server_module, "is_prisma_initial_connect_recovery_pending", lambda: False) + monkeypatch.setattr(proxy_server_module, "prisma_client", MagicMock()) + result = await handler._handle_authentication_error( + NoDatabaseConnectionError(), + MagicMock(), + {}, + "/v1/models", + None, + "cached-jwt", + ) + + assert result.user_id == "__db_unavailable_fallback__" + assert result.user_role == LitellmUserRoles.INTERNAL_USER + + +@pytest.mark.asyncio +async def test_nil_client_without_configured_database_does_not_get_fail_open_identity(monkeypatch): + handler = UserAPIKeyAuthExceptionHandler() + + monkeypatch.setattr(auth_exception_handler_module, "seed_request_identity", MagicMock()) + monkeypatch.setattr(proxy_server_module.proxy_logging_obj, "post_call_failure_hook", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server_module, "general_settings", {"allow_requests_on_db_unavailable": True}) + monkeypatch.setattr(proxy_server_module, "is_prisma_initial_connect_recovery_pending", lambda: False) + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + NoDatabaseConnectionError(), + MagicMock(), + {}, + "/v1/models", + None, + "cached-jwt", + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + @pytest.mark.asyncio @pytest.mark.parametrize( "prisma_error", diff --git a/tests/test_litellm/proxy/auth/test_resolvers_store.py b/tests/test_litellm/proxy/auth/test_resolvers_store.py index 5e644a084a5..b35ce380b14 100644 --- a/tests/test_litellm/proxy/auth/test_resolvers_store.py +++ b/tests/test_litellm/proxy/auth/test_resolvers_store.py @@ -16,8 +16,7 @@ from litellm.proxy.auth.resolvers.store import IdentityStore class _FakeCache: """Stands in for the DualCache get_key_object reads. It returns a cache hit - before the DB is touched, so seeding it exercises resolve without a database - (a non-None prisma client is still required; it is never reached on a hit).""" + before the DB is touched, so seeding it exercises resolve without a database.""" def __init__(self, entries: Optional[Dict[str, object]] = None) -> None: self._entries = entries or {} @@ -64,6 +63,18 @@ def test_key_from_principal_raises_when_no_source_key_is_carried(): IdentityStore.key_from_principal(bare) +async def test_resolve_returns_cached_key_without_a_db_connection(): + token_hash = hash_token("sk-cached-before-outage") + key = UserAPIKeyAuth(token=token_hash, user_id="u-cached", team_id="t-cached") + store = IdentityStore(None, _FakeCache({token_hash: key})) + + principal = await store.resolve(hashed_token=token_hash) + + assert principal.source_key is not None + assert principal.source_key.user_id == "u-cached" + assert principal.source_key.team_id == "t-cached" + + async def test_resolve_raises_without_a_db_connection(): store = IdentityStore(None, _FakeCache()) with pytest.raises(NoDatabaseConnectionError): diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6a117985820..424fa6f14e5 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -33,6 +33,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _ensure_parent_otel_span_on_request_state, _PendingAutoRegister, _matches_routing_override, + _resolve_jwt_to_virtual_key, _reserve_budget_after_common_checks, _route_requires_auth_despite_public, _routing_selector_matches_claim, @@ -51,6 +52,35 @@ class _RoutingRequest: self.state = SimpleNamespace() +@pytest.mark.asyncio +async def test_jwt_virtual_key_mapping_uses_cached_key_during_initial_db_recovery(): + token_hash = "cached-virtual-key-hash" + cache = DualCache() + await cache.async_set_cache("jwt_key_mapping:sub:jwt-user", token_hash) + await cache.async_set_cache( + token_hash, + UserAPIKeyAuth(token=token_hash, user_id="jwt-user", team_id="jwt-team"), + ) + jwt_handler = MagicMock() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + virtual_key_mapping_cache_ttl=300, + ) + + result = await _resolve_jwt_to_virtual_key( + jwt_claims={"sub": "jwt-user"}, + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.user_id == "jwt-user" + assert result.team_id == "jwt-team" + + def test_get_api_key(): bearer_token = "Bearer sk-12345678" api_key = "sk-12345678" diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 72c6f77a4f7..69d51bf2308 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -8,9 +8,11 @@ import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from prisma.engine.errors import BinaryNotFoundError from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module +import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -2349,6 +2351,170 @@ async def test_health_readiness_returns_503_when_db_disconnected(): assert result == {"status": "healthy", "db": "disconnected"} +@pytest.mark.asyncio +async def test_health_readiness_stays_ready_when_db_disconnected_and_fail_open_enabled(monkeypatch): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + mock_prisma.attempt_db_reconnect = AsyncMock(return_value=False) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + monkeypatch.setattr(proxy_server_module, "prisma_client", mock_prisma) + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result == {"status": "healthy", "db": "disconnected"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "permanent_error", + [ + PrismaError("query engine fault"), + BinaryNotFoundError("query engine binary not found"), + ], +) +async def test_health_readiness_fail_open_rejects_permanent_database_fault(monkeypatch, permanent_error): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=permanent_error) + mock_prisma.attempt_db_reconnect = AsyncMock() + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + monkeypatch.setattr(proxy_server_module, "prisma_client", mock_prisma) + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + response = Response() + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result == {"status": "healthy", "db": "disconnected"} + mock_prisma.attempt_db_reconnect.assert_not_called() + + +@pytest.mark.asyncio +async def test_health_readiness_uses_final_reconnect_failure_for_fail_open_policy(monkeypatch): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock( + side_effect=[ + httpx.ConnectError("connection refused"), + BinaryNotFoundError("query engine binary not found"), + ] + ) + mock_prisma.attempt_db_reconnect = AsyncMock(return_value=False) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + monkeypatch.setattr(proxy_server_module, "prisma_client", mock_prisma) + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + + response = Response() + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result == {"status": "healthy", "db": "disconnected"} + + +@pytest.mark.asyncio +async def test_health_readiness_distinguishes_initial_recovery_from_no_database(monkeypatch): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + monkeypatch.setattr(proxy_server_module, "is_prisma_initial_connect_recovery_pending", lambda: True) + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + + response = Response() + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result == {"status": "healthy", "db": "disconnected"} + + +@pytest.mark.asyncio +async def test_health_readiness_returns_503_if_fail_open_is_disabled_during_initial_recovery(monkeypatch): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + monkeypatch.setattr(proxy_server_module, "is_prisma_initial_connect_recovery_pending", lambda: True) + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": False}, + ) + + response = Response() + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result == {"status": "healthy", "db": "disconnected"} + + +@pytest.mark.asyncio +async def test_detailed_readiness_stays_ready_for_disconnected_db_when_fail_open_enabled(monkeypatch): + from fastapi import Response + + mock_prisma = MagicMock() + monkeypatch.setattr(proxy_server_module, "prisma_client", mock_prisma) + monkeypatch.setattr(proxy_server_module, "redis_usage_cache", MagicMock()) + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + monkeypatch.setattr( + _health_endpoints_module, + "_db_health_readiness_check", + AsyncMock(return_value={"status": "disconnected", "fail_open_safe": True}), + ) + + response = Response() + result = await _health_endpoints_module._get_health_readiness_details(response=response) + + assert response.status_code == 200 + assert result["db"] == "disconnected" + + @pytest.mark.asyncio async def test_health_readiness_returns_200_when_db_connected(): """Happy path: connected DB keeps the legacy 200.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index afc42e8db45..1802a25182a 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -18,6 +18,7 @@ import yaml from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient +from prisma.engine.errors import BinaryNotFoundError, MismatchedVersionsError import litellm @@ -1443,10 +1444,22 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): from fastapi import FastAPI # Import happens here - this is when the module probably reads the config path - from litellm.proxy.proxy_server import proxy_startup_event + from litellm.proxy.proxy_server import ProxyStartupEvent, proxy_startup_event # Mock the Prisma import monkeypatch.setattr("litellm.proxy.proxy_server.PrismaClient", MockPrisma) + start_initial_prisma_recovery = MagicMock() + stop_initial_prisma_recovery = AsyncMock() + monkeypatch.setattr( + ProxyStartupEvent, + "start_initial_prisma_connect_recovery", + start_initial_prisma_recovery, + ) + monkeypatch.setattr( + ProxyStartupEvent, + "stop_initial_prisma_connect_recovery", + stop_initial_prisma_recovery, + ) # Create test app app = FastAPI() @@ -1499,6 +1512,14 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): assert master_key == test_resolved_key + assert mock_prisma.await_count == 3 + assert all( + call.kwargs["recovered_client_initializer"] == ProxyStartupEvent.initialize_recovered_prisma_services + for call in mock_prisma.await_args_list + ) + assert start_initial_prisma_recovery.call_count == 3 + assert stop_initial_prisma_recovery.await_count == 3 + def test_team_info_masking(): """ @@ -11022,23 +11043,274 @@ async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(mon @pytest.mark.asyncio -async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkeypatch): - """Retaining only ever applies to a client that connected. If ``connect()`` - failed there is no usable client and no watchdog to recover it, so the caller - must still get ``None``.""" - monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") +@pytest.mark.parametrize( + "permanent_error", + [ + pytest.param( + BinaryNotFoundError("query engine binary not found"), + id="binary-not-found", + ), + pytest.param( + MismatchedVersionsError(expected="1", got="2"), + id="version-mismatch", + ), + ], +) +async def test_setup_prisma_client_does_not_recover_permanent_initial_fault( + monkeypatch, + permanent_error, +): + """Fail-open recovery is only for transient connection failures. + + A missing or incompatible query engine cannot recover when the database + returns, so retaining it would keep the proxy Ready in degraded auth mode + indefinitely. + """ + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + + mock_client = _mock_startup_prisma_client(connect_error=permanent_error) + + with pytest.raises(type(permanent_error)): + await _run_setup_prisma_client(mock_client) + + recovery_state = proxy_server_module._initial_prisma_connect_recovery_state + assert recovery_state.candidate is None + assert recovery_state.task is None + assert proxy_server_module.prisma_client is None + + +@pytest.mark.asyncio +async def test_setup_prisma_client_recovers_after_initial_connect_failure(monkeypatch): + """A failed initial connect must recover without publishing its disconnected client.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + + mock_client = _mock_startup_prisma_client(connect_error=httpx.ConnectError("connection refused")) + mock_client.attempt_db_reconnect = AsyncMock(return_value=True) + recovered_client_initializer = AsyncMock() + + from litellm.proxy.proxy_server import ProxyStartupEvent + + monkeypatch.setattr(proxy_server_module, "PrismaClient", MagicMock(return_value=mock_client)) + result = await ProxyStartupEvent._setup_prisma_client( + database_url="postgresql://litellm:litellm@localhost:5432/litellm", + proxy_logging_obj=MagicMock(), + user_api_key_cache=DualCache(), + recovered_client_initializer=recovered_client_initializer, + ) + + assert result is None + assert proxy_server_module.prisma_client is None + assert proxy_server_module._initial_prisma_connect_recovery_state.task is None + + ProxyStartupEvent.start_initial_prisma_connect_recovery() + recovery_task = proxy_server_module._initial_prisma_connect_recovery_state.task + if recovery_task is None: + pytest.fail("initial Prisma recovery task was not scheduled") + await recovery_task + + mock_client.attempt_db_reconnect.assert_awaited_once() + assert proxy_server_module.prisma_client is mock_client + assert mock_client.connect.await_count == 1 + mock_client.db.start_token_refresh_task.assert_awaited_once() + mock_client.start_db_health_watchdog_task.assert_awaited_once() + recovered_client_initializer.assert_awaited_once_with(mock_client) + + +@pytest.mark.asyncio +async def test_initial_prisma_recovery_keeps_client_when_deferred_startup_fails(monkeypatch): + """Deferred startup is not collectively idempotent, so publish the recovered + client and log one initializer failure instead of retrying duplicate jobs.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + + mock_client = _mock_startup_prisma_client(connect_error=httpx.ConnectError("connection refused")) + mock_client.attempt_db_reconnect = AsyncMock(return_value=True) + recovered_client_initializer = AsyncMock(side_effect=RuntimeError("deferred startup failed")) + + from litellm.proxy.proxy_server import ProxyStartupEvent + + monkeypatch.setattr(proxy_server_module, "PrismaClient", MagicMock(return_value=mock_client)) + await ProxyStartupEvent._setup_prisma_client( + database_url="postgresql://litellm:litellm@localhost:5432/litellm", + proxy_logging_obj=MagicMock(), + user_api_key_cache=DualCache(), + recovered_client_initializer=recovered_client_initializer, + ) + + ProxyStartupEvent.start_initial_prisma_connect_recovery() + recovery_task = proxy_server_module._initial_prisma_connect_recovery_state.task + if recovery_task is None: + pytest.fail("initial Prisma recovery task was not scheduled") + await recovery_task + + assert proxy_server_module.prisma_client is mock_client + mock_client.attempt_db_reconnect.assert_awaited_once() + recovered_client_initializer.assert_awaited_once_with(mock_client) + assert proxy_server_module.is_prisma_initial_connect_recovery_pending() is False + + +@pytest.mark.asyncio +async def test_initial_prisma_connect_recovery_retries_and_initializes_once(monkeypatch): + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + + mock_client = _mock_startup_prisma_client(connect_error=httpx.ConnectError("connection refused")) + mock_client._db_health_watchdog_interval_seconds = 0 + mock_client.attempt_db_reconnect = AsyncMock(side_effect=[False, True]) + recovered_client_initializer = AsyncMock() + + from litellm.proxy.proxy_server import ProxyStartupEvent + + monkeypatch.setattr(proxy_server_module, "PrismaClient", MagicMock(return_value=mock_client)) + await ProxyStartupEvent._setup_prisma_client( + database_url="postgresql://litellm:litellm@localhost:5432/litellm", + proxy_logging_obj=MagicMock(), + user_api_key_cache=DualCache(), + recovered_client_initializer=recovered_client_initializer, + ) + + ProxyStartupEvent.start_initial_prisma_connect_recovery() + recovery_task = proxy_server_module._initial_prisma_connect_recovery_state.task + if recovery_task is None: + pytest.fail("initial Prisma recovery task was not scheduled") + await recovery_task + + assert mock_client.attempt_db_reconnect.await_count == 2 + assert proxy_server_module.prisma_client is mock_client + recovered_client_initializer.assert_awaited_once_with(mock_client) + + +@pytest.mark.asyncio +async def test_initial_prisma_connect_recovery_is_singleton_and_stops_cleanly(monkeypatch): + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") monkeypatch.setattr( proxy_server_module, "general_settings", {"allow_requests_on_db_unavailable": True}, ) + reconnect_started = asyncio.Event() + + async def _wait_for_database(*args, **kwargs): + reconnect_started.set() + await asyncio.Event().wait() + mock_client = _mock_startup_prisma_client(connect_error=httpx.ConnectError("connection refused")) - result = await _run_setup_prisma_client(mock_client) + mock_client.attempt_db_reconnect = AsyncMock(side_effect=_wait_for_database) + mock_client.disconnect = AsyncMock() + + from litellm.proxy.proxy_server import ProxyStartupEvent + + monkeypatch.setattr(proxy_server_module, "PrismaClient", MagicMock(return_value=mock_client)) + await ProxyStartupEvent._setup_prisma_client( + database_url="postgresql://litellm:litellm@localhost:5432/litellm", + proxy_logging_obj=MagicMock(), + user_api_key_cache=DualCache(), + ) + + ProxyStartupEvent.start_initial_prisma_connect_recovery() + recovery_task = proxy_server_module._initial_prisma_connect_recovery_state.task + ProxyStartupEvent.start_initial_prisma_connect_recovery() + + assert recovery_task is not None + assert proxy_server_module._initial_prisma_connect_recovery_state.task is recovery_task + await reconnect_started.wait() + + await ProxyStartupEvent.stop_initial_prisma_connect_recovery() + + assert recovery_task.cancelled() + assert proxy_server_module.is_prisma_initial_connect_recovery_pending() is False + mock_client.disconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_setup_prisma_client_without_database_does_not_schedule_recovery(monkeypatch): + from litellm.proxy.proxy_server import ProxyStartupEvent + + prisma_constructor = MagicMock() + monkeypatch.setattr(proxy_server_module, "PrismaClient", prisma_constructor) + result = await ProxyStartupEvent._setup_prisma_client( + database_url=None, + proxy_logging_obj=MagicMock(), + user_api_key_cache=DualCache(), + ) + + ProxyStartupEvent.start_initial_prisma_connect_recovery() assert result is None - assert mock_client.start_db_health_watchdog_task.await_count == 0 - assert mock_client.health_check.await_count == 0 + prisma_constructor.assert_not_called() + assert proxy_server_module.is_prisma_initial_connect_recovery_pending() is False + assert proxy_server_module._initial_prisma_connect_recovery_state.task is None + + +@pytest.mark.asyncio +async def test_recovered_prisma_client_runs_database_dependent_startup(monkeypatch): + from litellm.proxy.proxy_server import ProxyStartupEvent + + mock_client = MagicMock() + coordination_cache = MagicMock() + worker_heartbeat = MagicMock() + start_migrations = MagicMock() + init_coordination_redis = AsyncMock(return_value=coordination_cache) + set_redis_usage_cache = MagicMock() + update_proxy_logging = MagicMock() + add_missing_proxy_hooks = MagicMock() + initialize_jobs = AsyncMock(return_value=worker_heartbeat) + update_default_budget = AsyncMock() + sync_ui_settings = AsyncMock() + + monkeypatch.setattr(litellm, "max_budget", 0) + monkeypatch.setattr(proxy_server_module, "llm_router", None) + monkeypatch.setattr( + proxy_server_module._initial_prisma_connect_recovery_state, + "recovered_worker_heartbeat", + None, + ) + monkeypatch.setattr(ProxyStartupEvent, "start_prisma_migrations", start_migrations) + monkeypatch.setattr(ProxyStartupEvent, "_init_coordination_redis_from_db", init_coordination_redis) + monkeypatch.setattr(proxy_server_module, "_set_redis_usage_cache", set_redis_usage_cache) + monkeypatch.setattr(proxy_server_module.proxy_logging_obj, "update_values", update_proxy_logging) + monkeypatch.setattr( + proxy_server_module.proxy_logging_obj, + "add_missing_proxy_hooks", + add_missing_proxy_hooks, + ) + monkeypatch.setattr(ProxyStartupEvent, "initialize_scheduled_background_jobs", initialize_jobs) + monkeypatch.setattr(ProxyStartupEvent, "_update_default_team_member_budget", update_default_budget) + monkeypatch.setattr(ProxyStartupEvent, "_sync_ui_settings_to_general_settings", sync_ui_settings) + + await ProxyStartupEvent.initialize_recovered_prisma_services(mock_client) + + start_migrations.assert_called_once_with(mock_client) + init_coordination_redis.assert_awaited_once() + set_redis_usage_cache.assert_called_once_with(coordination_cache) + update_proxy_logging.assert_called_once_with(redis_cache=coordination_cache) + add_missing_proxy_hooks.assert_called_once_with(None) + initialize_jobs.assert_awaited_once() + update_default_budget.assert_awaited_once() + sync_ui_settings.assert_awaited_once() + assert proxy_server_module._initial_prisma_connect_recovery_state.recovered_worker_heartbeat is worker_heartbeat async def _run_scheduled_background_jobs(): diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py index a97dcb41e44..f9e553dfed2 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py @@ -1,5 +1,5 @@ """Pin ProxyLogging lifecycle: ``__init__``, ``startup_event``, -``update_values``, ``_add_proxy_hooks``, ``get_proxy_hook``, and +``update_values``, ``add_missing_proxy_hooks``, ``get_proxy_hook``, and ``_init_litellm_callbacks``. Also covers ``update_request_status`` and ``_convert_user_api_key_auth_to_dict`` @@ -211,7 +211,7 @@ async def test_startup_event_skips_the_daily_report_when_it_is_not_an_alert_type # --------------------------------------------------------------------------- -# _add_proxy_hooks +# add_missing_proxy_hooks # --------------------------------------------------------------------------- @@ -243,7 +243,7 @@ def test_add_proxy_hooks_registers_callbacks(proxy_logging, monkeypatch): ) with patch("litellm.proxy.proxy_server.prisma_client", None): - proxy_logging._add_proxy_hooks(llm_router=None) + proxy_logging.add_missing_proxy_hooks(llm_router=None) keys = list(proxy_logging.proxy_hook_mapping.keys()) snapshot = { @@ -294,7 +294,7 @@ def test_add_proxy_hooks_skips_prisma_requiring_hook_when_no_db(proxy_logging, m ) with patch("litellm.proxy.proxy_server.prisma_client", None): - proxy_logging._add_proxy_hooks(llm_router=None) + proxy_logging.add_missing_proxy_hooks(llm_router=None) snapshot = { "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), @@ -326,7 +326,7 @@ def test_add_proxy_hooks_registers_prisma_requiring_hook_with_db(proxy_logging, ) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): - proxy_logging._add_proxy_hooks(llm_router=None) + proxy_logging.add_missing_proxy_hooks(llm_router=None) snapshot = { "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), @@ -342,6 +342,65 @@ def test_add_proxy_hooks_registers_prisma_requiring_hook_with_db(proxy_logging, } +def test_add_proxy_hooks_after_db_recovery_only_registers_missing_hooks(proxy_logging, monkeypatch): + hook_classes = _stub_hook_classes() + registered: List[Any] = [] + success_callbacks: List[Any] = [] + failure_callbacks: List[Any] = [] + async_success_callbacks: List[Any] = [] + async_failure_callbacks: List[Any] = [] + fake_prisma = MagicMock() + + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy import utils as utils_mod + + monkeypatch.setattr(utils_mod, "PROXY_HOOKS", list(hook_classes.keys())) + monkeypatch.setattr(utils_mod, "get_proxy_hook", hook_classes.__getitem__) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_callback", + lambda callback: registered.append(callback), + ) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_success_callback", + lambda callback: success_callbacks.append(callback), + ) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_failure_callback", + lambda callback: failure_callbacks.append(callback), + ) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_async_success_callback", + lambda callback: async_success_callbacks.append(callback), + ) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_async_failure_callback", + lambda callback: async_failure_callbacks.append(callback), + ) + + monkeypatch.setattr(proxy_server_module, "prisma_client", None) + proxy_logging.add_missing_proxy_hooks(llm_router=None) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + proxy_logging.add_missing_proxy_hooks(llm_router=None) + proxy_logging.add_missing_proxy_hooks(llm_router=None) + + assert list(proxy_logging.proxy_hook_mapping) == ["cache_control_check", "needs_db_hook", "db_only_hook"] + expected_callback_types = [ + "_PrismaFreeHook", + "_PrismaRequiringHook", + "_PrismaOnlyHook", + ] + assert [type(callback).__name__ for callback in registered] == expected_callback_types + assert [type(callback).__name__ for callback in success_callbacks] == expected_callback_types + assert [type(callback).__name__ for callback in failure_callbacks] == expected_callback_types + assert [type(callback).__name__ for callback in async_success_callbacks] == expected_callback_types + assert [type(callback).__name__ for callback in async_failure_callbacks] == expected_callback_types + + def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch): from litellm.proxy import utils as utils_mod @@ -352,7 +411,7 @@ def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch): monkeypatch.setattr(utils_mod, "get_proxy_hook", bad_resolver) with pytest.raises(KeyError): - proxy_logging._add_proxy_hooks(llm_router=None) + proxy_logging.add_missing_proxy_hooks(llm_router=None) # ---------------------------------------------------------------------------