From 6209f0694bfd7285deea4095faea4e850a7fe007 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 19:49:37 -0700 Subject: [PATCH] fix(health): keep team public names to the owning team and let model_id win over model A public name a team publishes its own deployment copy under now targets that copy only for a caller from that team, so an admin or another team probing the shared name gets the global deployment alone model_id wins when paired with model: a foreign id still gets the 403, and an id no deployment carries gets the 404 of the lone-id path, before any probe runs or a result is stored under it cache_health_check_results accepts the Mapping sequences perform_health_check returns --- litellm/proxy/health_check.py | 19 ++- .../shared_health_check_manager.py | 5 +- .../health_endpoints/_health_endpoints.py | 61 +++---- .../health_endpoints/test_health_endpoints.py | 159 ++++++++++++++++++ .../proxy/test_health_check_functions.py | 27 ++- .../proxy/test_health_check_max_tokens.py | 6 +- 6 files changed, 231 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 8cda318b937..9ca31c8b626 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -258,15 +258,18 @@ 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.""" +def deployment_answers_to(deployment: Mapping[str, object], model_name: str, team_id: str | None) -> bool: + """True when `model_name` is the deployment's model_name or the public name the caller's own team 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) + owner_team_id: Final = info.get("team_id") if isinstance(info, Mapping) else None + return model_name == deployment.get("model_name") or ( + model_name == public_name and team_id is not None and owner_team_id == team_id + ) def _narrow_to_target( - model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None + model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None, team_id: str | None ) -> tuple[Mapping[str, object], ...]: """Narrow to the requested deployment. An id matching nothing keeps the whole list.""" if model_id is not None: @@ -275,7 +278,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 deployment_answers_to(x, model)) + return by_param or tuple(x for x in model_list if deployment_answers_to(x, model, team_id)) def _is_strategy_router_deployment(litellm_params: Mapping[str, object]) -> bool: @@ -820,13 +823,15 @@ async def perform_health_check( instrumentation_context: dict | None = None, health_check_skip_disabled_background_models: bool = False, router: "Router | None" = None, + team_id: str | None = None, ): """ Perform a health check on the system. When model_id is provided, only the deployment with that id is checked (so models that share the same name but have different ids are checked separately). - When model (name) is provided, all deployments matching that name are checked. + When model (name) is provided, all deployments matching that name are checked; + a team's public model name only matches for a caller from that team (``team_id``). When ``health_check_skip_disabled_background_models`` is True (via ``general_settings.health_check_skip_disabled_background_models``), deployments @@ -857,7 +862,7 @@ async def perform_health_check( cycle_start_time: Final = time.monotonic() requested_model_count: Final = len(model_list) skip_disabled: Final = health_check_skip_disabled_background_models - narrowed: Final = _health_check_eligible(_narrow_to_target(model_list, model, model_id), skip_disabled) + narrowed: Final = _health_check_eligible(_narrow_to_target(model_list, model, model_id, team_id), skip_disabled) if not narrowed: if instrumentation_enabled: logger.debug( diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index f12cee4b636..7dbbd55d87f 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -1,6 +1,7 @@ import asyncio import json import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger @@ -143,8 +144,8 @@ class SharedHealthCheckManager: async def cache_health_check_results( self, - healthy_endpoints: list[dict[str, Any]], - unhealthy_endpoints: list[dict[str, Any]], + healthy_endpoints: Sequence[Mapping[str, object]], + unhealthy_endpoints: Sequence[Mapping[str, object]], ) -> None: """ Cache health check results in Redis. diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 267b094526b..dcef0fcd58a 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -925,41 +925,31 @@ def _caller_may_probe_deployment( ) -def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: str | None) -> set | None: +def _resolve_targeted_model_ids( + model_list: list, model: str | None, model_id: str | None, team_id: str | None +) -> set | None: """ Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of - deployment IDs the response should be scoped to. + deployment IDs the response should be scoped to, mirroring the live-path + narrowing in ``perform_health_check()``: ``model_id`` wins when given and + matches ``model_info.id`` only; ``model`` matches the deployment's + ``model_name`` alias, its ``litellm_params.model`` provider string, or the + ``model_info.team_public_model_name`` the caller's own team reaches it by. - Mirrors the live-path semantics in ``perform_health_check()``: ``model`` - 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 - 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." + Callers pass an already-scoped list, so a ``model_id`` outside the + caller's scope resolves to an empty set and never to the unvalidated id. + Returns ``None`` when no targeting is requested. """ - if not model and not model_id: + if model_id: + return {i for m in model_list if (i := (m.get("model_info") or {}).get("id")) == model_id} + if not model: return None - target_ids: Final[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 litellm_model == model or deployment_answers_to(m, model): - target_ids.add(deployment_id) - return target_ids + return { + i + for m in model_list + if (i := (m.get("model_info") or {}).get("id")) + and ((m.get("litellm_params") or {}).get("model") == model or deployment_answers_to(m, model, team_id)) + } def _filter_health_check_results_by_model_ids(results: dict, allowed_model_ids: set) -> dict: @@ -1040,8 +1030,12 @@ def _health_endpoint_resolve_target_model_name( model_id: str | None, llm_router, ) -> str | None: - """Map ``model_id`` (without ``model``) to ``model_name`` for live health checks.""" - if not model_id or model: + """Map ``model_id`` to its deployment's ``model_name`` for live health checks. + + ``model_id`` wins over ``model``, so an id no deployment carries is a 404 even + when it is paired with a known name. + """ + if not model_id: return model if llm_router is None: raise HTTPException( @@ -1161,7 +1155,7 @@ async def health_endpoint( if not restrict_to_allowed_models or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id, is_admin) ] - targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id) + targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id, user_api_key_dict.team_id) if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -1228,6 +1222,7 @@ async def health_endpoint( model_id=model_id, max_concurrency=health_check_concurrency, router=llm_router, + team_id=user_api_key_dict.team_id, **_hc_filter, ) return _post_process(router_result) 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 e0cf6d97141..c5950f6f30e 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2357,6 +2357,7 @@ async def test_health_endpoint_returns_503_when_requested_model_has_no_healthy_e response=response, user_api_key_dict=user_api_key_dict, model="model-a", + model_id=None, ) assert response.status_code == 503 @@ -2417,6 +2418,7 @@ async def test_health_endpoint_returns_200_when_requested_model_has_healthy_endp response=response, user_api_key_dict=user_api_key_dict, model="model-a", + model_id=None, ) assert response.status_code == 200 @@ -3088,6 +3090,161 @@ async def test_health_endpoint_targets_both_deployments_behind_a_shared_public_n assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock", "id-team-b"] +async def _live_narrowed_model_ids( + model_list: Sequence[Mapping[str, object]], + user_api_key_dict: UserAPIKeyAuth, + model: str | None = None, + model_id: str | None = None, +) -> set[str]: + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None): + probed = [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list] + return probed, [], {} + + with ( + _proxy_health_globals(model_list, _router_for(model_list)), + patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion + "litellm.proxy.health_check._perform_health_check", side_effect=fake_probe + ), + ): + result = await health_endpoint( + response=Response(), user_api_key_dict=user_api_key_dict, model=model, model_id=model_id + ) + + return {ep["model_id"] for ep in result["healthy_endpoints"]} + + +_ADMIN_OUTSIDE_TEAM_B = UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN) + + +@pytest.mark.asyncio +async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_public_copies(): + """ + An admin outside team-b asks for ``bedrock-nova``. Team-b's copy answers to + that name only for team-b (routing keys public names by team), so probing + it too would spend team-b's credentials and let a healthy team copy mask a + down global deployment as 200. + """ + probed = await _live_narrowed_model_ids(_TEAM_MODEL_LIST, _ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova") + + assert probed == {"id-bedrock"} + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_both_deployments_behind_a_shared_public_name_for_the_owning_team(): + probed = await _live_narrowed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + ) + + assert probed == {"id-bedrock", "id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_public_copies_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=_ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + + +def test_resolve_targeted_model_ids_lets_model_id_win_over_model(): + resolve = _health_endpoints_module._resolve_targeted_model_ids + + assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", "id-team-b", None) == {"id-team-b"} + assert resolve([_TEAM_MODEL_LIST[0]], "bedrock-nova", "id-team-b", None) == set() + assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", None, None) == {"id-bedrock"} + assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", None, "team-b") == {"id-bedrock", "id-team-b"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_rejects_an_in_scope_model_paired_with_a_foreign_model_id(use_background_health_checks): + """ + A key scoped to ``bedrock-nova`` pairs that name with another team's + deployment id. The in-scope name must not carry the foreign id past the + 403: the live path narrows by id first, so the caller's own deployment + would be probed and its result stored under the foreign id. + """ + from fastapi import HTTPException, 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=use_background_health_checks, + health_check_results=_TEAM_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"]), + model="bedrock-nova", + model_id="id-team-b", + ) + + assert refused.value.status_code == 403 + assert "id-team-b" in str(refused.value.detail) + probe.assert_not_awaited() + + +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +@pytest.mark.asyncio +async def test_health_endpoint_returns_404_for_a_model_paired_with_an_unknown_model_id(use_background_health_checks): + """ + ``model_id`` wins over ``model``: pairing a known name with an id no + deployment carries gets the same 404 as the lone unknown id, before any + probe runs or a result is stored under the unknown id. + """ + from fastapi import HTTPException, 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=use_background_health_checks, + health_check_results=_TEAM_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=_ADMIN_OUTSIDE_TEAM_B, + model="bedrock-nova", + model_id="id-nobody-has", + ) + + assert refused.value.status_code == 404 + assert "id-nobody-has" in str(refused.value.detail) + probe.assert_not_awaited() + + def test_health_test_connection_keeps_error_and_raw_request_through_the_allowlist(monkeypatch): """ The dashboard's Test Connect button reads ``result.error`` and @@ -3198,6 +3355,8 @@ async def test_health_endpoint_result_survives_non_json_safe_deployment_params() result = await health_endpoint( response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="hashed-admin-key", user_role=LitellmUserRoles.PROXY_ADMIN), + model=None, + model_id=None, ) encoded = jsonable_encoder(result) diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index 137eda223d1..af20b80a718 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -644,7 +644,7 @@ async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_n "litellm.proxy.health_check._perform_health_check", probe ): healthy, unhealthy, _ = await perform_health_check( - model_list=[team_deployment, other_deployment], model="bedrock-nova" + model_list=[team_deployment, other_deployment], model="bedrock-nova", team_id="team-b" ) assert [m["model_info"]["id"] for m in probe.call_args.args[0]] == ["id-team-b"] @@ -652,6 +652,31 @@ async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_n assert unhealthy == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("team_id", [None, "team-a"]) +async def test_perform_health_check_keeps_a_public_name_to_the_owning_team(team_id): + """A team's public model name is not a global alias: callers outside the team must not probe its deployment.""" + 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"}, + } + 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], model="bedrock-nova", team_id=team_id + ) + + probe.assert_not_awaited() + assert healthy == [] + 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 diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index 97e308d7c3c..dd3669644af 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -872,9 +872,9 @@ def test_narrowing_by_an_id_that_matches_nothing_keeps_the_whole_list(): """Pinned because the disabled-dependency fix moved this filter into its own helper.""" deployments = [{"model_name": "a", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a-1"}}] - assert hc_module._narrow_to_target(deployments, None, "no-such-id") == tuple(deployments) - assert hc_module._narrow_to_target(deployments, None, "a-1") == tuple(deployments) - assert hc_module._narrow_to_target(deployments, "a", None) == tuple(deployments) + assert hc_module._narrow_to_target(deployments, None, "no-such-id", None) == tuple(deployments) + assert hc_module._narrow_to_target(deployments, None, "a-1", None) == tuple(deployments) + assert hc_module._narrow_to_target(deployments, "a", None, None) == tuple(deployments) def _nested_router_fixture(parent_tier: str):