From 787f2dee0c3d9e895c14adfccd21ca7e090fea24 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:52:59 -0700 Subject: [PATCH] fix(health): match team public model names when targeting /health by model --- litellm/proxy/health_check.py | 9 ++- .../health_endpoints/_health_endpoints.py | 9 +-- .../health_endpoints/test_health_endpoints.py | 67 ++++++++++++++++++- .../proxy/test_health_check_functions.py | 29 ++++++++ 4 files changed, 107 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 78bb01aa6ac..8cda318b937 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -258,6 +258,13 @@ def _deployment_model(deployment: Mapping[str, object]) -> str | None: return params.get("model") if isinstance(params, Mapping) else None +def deployment_answers_to(deployment: Mapping[str, object], model_name: str) -> bool: + """True when `model_name` is the deployment's model_name or the public name a team key reaches it by.""" + info: Final = deployment.get("model_info") + public_name: Final = info.get("team_public_model_name") if isinstance(info, Mapping) else None + return model_name in (deployment.get("model_name"), public_name) + + def _narrow_to_target( model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None ) -> tuple[Mapping[str, object], ...]: @@ -268,7 +275,7 @@ def _narrow_to_target( if model is None: return tuple(model_list) by_param: Final = tuple(x for x in model_list if _deployment_model(x) == model) - return by_param or tuple(x for x in model_list if x.get("model_name") == model) + return by_param or tuple(x for x in model_list if deployment_answers_to(x, model)) def _is_strategy_router_deployment(litellm_params: Mapping[str, object]) -> bool: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c66eafc724f..267b094526b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, _update_litellm_params_for_health_check, + deployment_answers_to, health_check_filter_kwargs_from_general_settings, perform_health_check, run_with_timeout, @@ -930,9 +931,9 @@ def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: s 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``. + matches the deployment's ``model_name`` alias, its ``litellm_params.model`` + provider string, or the ``model_info.team_public_model_name`` a team key + reaches it by. ``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 @@ -956,7 +957,7 @@ def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: s continue if model: litellm_model = (m.get("litellm_params") or {}).get("model") - if m.get("model_name") == model or litellm_model == model: + if litellm_model == model or deployment_answers_to(m, model): target_ids.add(deployment_id) return target_ids 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 022cf1e5505..e0cf6d97141 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2802,7 +2802,7 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, async def _live_probed_model_ids( - model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth + model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth, model: str | None = None ) -> set[str]: from fastapi import Response @@ -2821,7 +2821,7 @@ async def _live_probed_model_ids( side_effect=fake_perform, ), ): - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=model, model_id=None) return {m["model_info"]["id"] for m in captured["model_list"]} @@ -3025,6 +3025,69 @@ async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team_on assert result["healthy_count"] == 1 +_TEAM_ONLY_MODEL_LIST = [_TEAM_MODEL_LIST[1]] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_a_team_only_deployment_by_its_public_name_on_live_path(): + """ + A team key targets its deployment by ``team_public_model_name``; when that + name resolves to nothing but the team deployment, the probe must run rather + than 403 as if the key were out of scope. + """ + probed = await _live_probed_model_ids( + _TEAM_ONLY_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + ) + + assert probed == {"id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_a_team_only_deployment_by_its_public_name_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_ONLY_MODEL_LIST, + _router_for(_TEAM_ONLY_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-team-b"] + + +@pytest.mark.asyncio +async def test_health_endpoint_targets_both_deployments_behind_a_shared_public_name_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock", "id-team-b"] + + def test_health_test_connection_keeps_error_and_raw_request_through_the_allowlist(monkeypatch): """ The dashboard's Test Connect button reads ``result.error`` and diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index fdae11d517a..137eda223d1 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -623,6 +623,35 @@ async def test_perform_health_check_and_save_forwards_skip_disabled_background_f assert call_kwargs["health_check_skip_disabled_background_models"] is True +@pytest.mark.asyncio +async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_name(): + """``/health?model=`` must probe the team deployment, not an empty list.""" + from litellm.proxy.health_check import perform_health_check + + team_deployment = { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-team-b", "team_id": "team-b", "team_public_model_name": "bedrock-nova"}, + } + other_deployment = { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-openai"}, + } + probe = AsyncMock(return_value=([{"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-team-b"}], [], {})) + + with patch( # test-quality-ok: the deployments handed to the probe are the assertion; no injection seam + "litellm.proxy.health_check._perform_health_check", probe + ): + healthy, unhealthy, _ = await perform_health_check( + model_list=[team_deployment, other_deployment], model="bedrock-nova" + ) + + assert [m["model_info"]["id"] for m in probe.call_args.args[0]] == ["id-team-b"] + assert [ep["model_id"] for ep in healthy] == ["id-team-b"] + assert unhealthy == [] + + def test_parse_background_health_check_model_groups_unset_returns_none(): from litellm.proxy.health_check import parse_background_health_check_model_groups