diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1eda01e5c6..35c9edb937 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -776,6 +776,45 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict: return out +def _resolve_targeted_model_ids( + model_list: list, model: Optional[str], model_id: Optional[str] +) -> Optional[set]: + """ + Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of + deployment IDs the response should be scoped to. + + Mirrors the live-path semantics in ``perform_health_check()``: ``model`` + matches either the deployment's ``model_name`` alias or its + ``litellm_params.model`` provider string. ``model_id`` matches + ``model_info.id``. + + Both query params are validated against the supplied ``model_list``. + Callers pass an already-scoped list (filtered to the caller's allowed + models for non-admins, full list for admins), so a ``model_id`` that + isn't present resolves to an empty set rather than a single-element + set — preventing a non-admin from reading another deployment's cached + health entry by guessing its ID. + + Returns ``None`` when no targeting is requested — callers should treat + that as "no filter." + """ + if not model and not model_id: + return None + target_ids: set = set() + for m in model_list: + deployment_id = (m.get("model_info") or {}).get("id") + if not deployment_id: + continue + if model_id and deployment_id == model_id: + target_ids.add(deployment_id) + continue + if model: + litellm_model = (m.get("litellm_params") or {}).get("model") + if m.get("model_name") == model or litellm_model == model: + target_ids.add(deployment_id) + return target_ids + + def _filter_health_check_results_by_model_ids( results: dict, allowed_model_ids: set ) -> dict: @@ -926,6 +965,7 @@ async def health_endpoint( ) is_admin = _is_proxy_admin(user_api_key_dict) + model_specific_request = bool(model or model_id) def _post_process(result: dict) -> dict: # api_base / api_version reveal which provider/region/internal host the @@ -933,6 +973,12 @@ async def health_endpoint( # still see model/model_id and the healthy/unhealthy status. We also # set a header so non-admin clients that previously parsed those # fields can detect the change programmatically. + # When a caller asked about a specific model/model_id and zero + # endpoints came back healthy, surface that as a 503 so monitoring + # systems can rely on the HTTP status instead of having to parse the + # body. The body shape is unchanged. + if model_specific_request and result.get("healthy_count", 0) == 0: + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE if is_admin: return result response.headers["Litellm-Health-Field-Notice"] = ( @@ -975,16 +1021,29 @@ async def health_endpoint( m for m in _llm_model_list if m.get("model_name") in allowed_models ] if use_background_health_checks: + # The cached background result covers every model. When the + # caller targets a specific model/model_id we have to narrow the + # cache to that deployment before _post_process evaluates + # healthy_count, otherwise an unhealthy "foo" combined with any + # other healthy model would still report healthy_count > 0 and + # the targeted-503 path would never fire. + targeted_ids = _resolve_targeted_model_ids(_llm_model_list, model, model_id) if len(user_api_key_dict.models) > 0: allowed_model_ids = { (m.get("model_info") or {}).get("id") for m in _llm_model_list if (m.get("model_info") or {}).get("id") } - filtered = _filter_health_check_results_by_model_ids( - health_check_results, allowed_model_ids + # _llm_model_list is already scoped to the caller's allowed + # model_names above, so targeted_ids is implicitly the + # intersection of "targeted" and "allowed." + filter_ids = ( + targeted_ids if targeted_ids is not None else allowed_model_ids ) - if not allowed_model_ids: + filtered = _filter_health_check_results_by_model_ids( + health_check_results, filter_ids + ) + if targeted_ids is None and not allowed_model_ids: # Caller has accessible model_names but none of the # matching deployments expose a model_info.id, so the # cache filter (which keys on model_id) drops every @@ -1005,6 +1064,15 @@ async def health_endpoint( "to populate model_info.id for these models." ] return _post_process(filtered) + if targeted_ids is not None: + # Admin caller targeting a specific model: filter the cache + # so the response (and the targeted-503 check) reflects only + # that deployment, not the global aggregate. + return _post_process( + _filter_health_check_results_by_model_ids( + health_check_results, targeted_ids + ) + ) return _post_process(health_check_results) else: router_result = await _perform_health_check_and_save( @@ -1384,7 +1452,7 @@ def callback_name(callback): tags=["health"], dependencies=[Depends(user_api_key_auth)], ) -async def health_readiness(): +async def health_readiness(response: Response): """ Unprotected endpoint for checking if worker can receive requests """ @@ -1417,8 +1485,8 @@ async def health_readiness(): try: index_info = await litellm.cache.cache._index_info() except Exception as e: - index_info = "index does not exist - error: " + str(e) - cache_type = {"type": cache_type, "index_info": index_info} + index_info = "index does not exist - error: " + str(e) # type: ignore[assignment] + cache_type = {"type": cache_type, "index_info": index_info} # type: ignore[assignment] # check log level log_level_name = logging.getLevelName(verbose_logger.getEffectiveLevel()) @@ -1427,6 +1495,12 @@ async def health_readiness(): # check DB if prisma_client is not None: # if db passed in, check if it's connected db_health_status = 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 db_health_status["status"] != "connected": + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", "db": db_health_status["status"], 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 d59682c2d5..353e67c9f7 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -927,8 +927,14 @@ async def test_health_endpoint_filters_background_cache_by_user_access(): ): from fastapi import Response + # Pass model=None, model_id=None explicitly: direct calls to the + # handler skip FastAPI's Query() resolution, so unspecified params + # would otherwise carry the Query() sentinel (which is truthy). result = await health_endpoint( - response=Response(), user_api_key_dict=user_api_key_dict + response=Response(), + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, ) # Sanity: the source cache had two entries before scoping; the scoping @@ -1016,10 +1022,16 @@ async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not(): admin_response = Response() non_admin_response = Response() admin_result = await health_endpoint( - response=admin_response, user_api_key_dict=admin_key + response=admin_response, + user_api_key_dict=admin_key, + model=None, + model_id=None, ) non_admin_result = await health_endpoint( - response=non_admin_response, user_api_key_dict=non_admin_key + response=non_admin_response, + user_api_key_dict=non_admin_key, + model=None, + model_id=None, ) finally: for p in common_patches: @@ -1113,7 +1125,10 @@ async def test_health_endpoint_warns_when_scoped_models_lack_model_id(): patch("litellm.proxy.proxy_server.health_check_concurrency", 1), ): result = await health_endpoint( - response=Response(), user_api_key_dict=user_api_key_dict + response=Response(), + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, ) assert result["healthy_count"] == 0 @@ -1125,6 +1140,415 @@ async def test_health_endpoint_warns_when_scoped_models_lack_model_id(): assert any("model_info.id" in w for w in result["warnings"]) +@pytest.mark.asyncio +async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cache(): + """ + A non-admin scoped to model-a must not be able to read model-b's cached + health entry by guessing its model_id. Before the fix, + _resolve_targeted_model_ids returned {model_id} unconditionally, so the + cache filter was driven by an unvalidated ID and the global cache + leaked id-b's entry to the caller. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", # caller has no access + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-b", + "api_base": "https://leaky-internal.test", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-scoped", + models=["model-a"], + ) + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + # llm_router None here means the model_id 404 lookup short-circuits; + # we patch _llm_model_list directly instead to drive the cache path. + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + # Calling with model="model-b" rather than model_id="id-b" because + # the model_id branch raises 404 when llm_router is None. The bug + # being verified is the same: targeted resolver must drop entries + # not in the caller's scoped model_list. With the fix, the result + # has no leaked endpoints and the targeted-503 path fires. + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-b", + model_id=None, + ) + + leaked_ids = {ep.get("model_id") for ep in result.get("healthy_endpoints", [])} + leaked_ids |= {ep.get("model_id") for ep in result.get("unhealthy_endpoints", [])} + assert ( + "id-b" not in leaked_ids + ), "background cache leaked an out-of-scope deployment to a scoped caller" + assert result["healthy_count"] == 0 + assert response.status_code == 503 + + +@pytest.mark.asyncio +async def test_health_endpoint_503_for_targeted_unhealthy_model_under_background_cache_admin(): + """ + With background_health_checks enabled, an admin calling /health?model=foo + must get 503 when foo specifically has zero healthy endpoints — even if + other unrelated models in the cache are healthy. Without the cache-path + filter, the global healthy_count would mask the targeted failure. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", # the unhealthy target + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", # an unrelated healthy model + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-b"}, + ], + "unhealthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"}, + ], + "healthy_count": 1, + "unhealthy_count": 1, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + model_id=None, + ) + + assert response.status_code == 503 + # Body must be scoped to the targeted model — not the global cache. + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 1 + returned_ids = {ep["model_id"] for ep in result.get("unhealthy_endpoints", [])} + assert returned_ids == {"id-a"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_503_when_requested_model_has_no_healthy_endpoints(): + """ + /health?model=foo must return 503 when the targeted model resolves but + has zero healthy endpoints. Body shape stays the same so existing + parsers still work; only the HTTP status changes. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "error": "boom", + } + ], + "healthy_count": 0, + "unhealthy_count": 1, + } + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + ) + + assert response.status_code == 503 + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 1 + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_200_when_requested_model_has_healthy_endpoints(): + """ + /health?model=foo with a healthy endpoint must keep returning the + default 200. Verifies the 503 path doesn't fire when healthy_count > 0. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [{"model": "openai/gpt-4o", "model_id": "id-a"}], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + response = Response() + # Default Response() exposes status_code as None; the endpoint should + # leave it alone for the healthy path. + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + ) + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_health_endpoint_no_model_param_returns_200_even_when_zero_healthy(): + """ + The non-targeted /health (no model / model_id query) preserves the + legacy 200 behavior even when healthy_count == 0. Existing K8s probes + and dashboards depend on this; only the targeted call became 5xx-aware. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"} + ], + "healthy_count": 0, + "unhealthy_count": 1, + } + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + # Pass model=None, model_id=None explicitly: when invoked through + # FastAPI, the Query(None) defaults resolve to None, but direct + # function calls in unit tests receive Query() sentinel objects + # (which are truthy). The explicit None mirrors production routing. + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, + ) + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_health_readiness_returns_503_when_db_disconnected(): + """ + When a Prisma client is configured but its health_check fails, the + readiness probe should mark the worker as unhealthy via the HTTP + status — not just a body field — so K8s removes the pod from the + Service endpoints. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope")) + mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope")) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result["db"] == "disconnected" + assert result["status"] == "healthy" # body shape unchanged for back-compat + + +@pytest.mark.asyncio +async def test_health_readiness_returns_200_when_db_connected(): + """Happy path: connected DB keeps the legacy 200.""" + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock() + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result["db"] == "connected" + + +@pytest.mark.asyncio +async def test_health_readiness_returns_200_when_no_db_configured(): + """ + `prisma_client is None` means the operator chose not to use a DB. That + is a valid configuration — the worker should still report ready. We + only flip to 503 when a DB *was* configured but is unreachable. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", None): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result["db"] == "Not connected" + + def test_clean_endpoint_data_strips_credentials_keeps_routing_fields(): """ _clean_endpoint_data() drops credentials but leaves api_base /