mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge 157733dc9b into b03e913ccf
This commit is contained in:
mphilippnv
GitHub
commit
1ac795bf5b
12 changed files with 956 additions and 120 deletions
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <migration_name>'")
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue