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
This commit is contained in:
mateo-berri 2026-09-11 19:49:37 -07:00
parent 787f2dee0c
commit 6209f0694b
6 changed files with 231 additions and 46 deletions

View file

@ -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(

View file

@ -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.

View file

@ -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)

View file

@ -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)

View file

@ -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

View file

@ -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):