scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(health): resolve aliases per allowlist layer and scope stored health rows

Browse source 
/health decided which models a caller could probe with its own alias handling, so a
model reachable only because one allowlist layer names the alias and another names
the target came back filtered out. The per-layer expansion `_can_object_call_model`
already does is now a shared helper both call sites use, which makes /health agree
with what a chat request through the same key would be allowed to do.

Wildcard allowlist entries used to be re-resolved against the router once per
deployment, which is quadratic on a large model list. Each entry now resolves once
per call into the set of deployment ids it serves.

/health/latest and /health/history returned every stored row regardless of who
asked, so a scoped key could read health for deployments it cannot call. Both now
filter to the caller's readable deployments, matching what /health itself returns.

The team-member layer only applies when the team row was actually read, the same
condition `common_checks` uses, so a failed team lookup no longer widens the check.
This commit is contained in:
mateo-berri 2026-09-05 22:18:44 -07:00
parent bd95fd3399
commit 126d75e6f2
3 changed files with 274 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/auth/auth_checks.py
View file

@ -3962,6 +3962,18 @@ def _check_model_access_helper(
return True
def _potential_models_for_auth_check(model: str, llm_router: Router | None) -> tuple[str, ...]:
"""The names one allowlist layer tries for a requested model: the name itself, then the alias it resolves to."""
if model in litellm.model_alias_map:
return (model, litellm.model_alias_map[model])
aliased: Final = (
llm_router._get_model_from_alias(model)
if llm_router is not None and model in llm_router.model_group_alias
else None
)
return (model, aliased) if aliased else (model,)
def _can_object_call_model(
model: str | list[str],
llm_router: Router | None,
@ -4002,13 +4014,7 @@ def _can_object_call_model(
)
return True
potential_models: Final = [model]
if model in litellm.model_alias_map:
potential_models.append(litellm.model_alias_map[model])
elif llm_router and model in llm_router.model_group_alias:
_model: Final = llm_router._get_model_from_alias(model)
if _model:
potential_models.append(_model)
potential_models: Final = _potential_models_for_auth_check(model=model, llm_router=llm_router)
## check model access for alias + underlying model - allow if either is in allowed models
for m in potential_models:

178
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -45,6 +45,7 @@ from litellm.proxy.auth.auth_checks import (
_check_model_access_helper, # pyright: ignore[reportPrivateUsage] # the auth layer's model access predicate, reused so /health scopes exactly like a request
_get_models_from_access_groups, # pyright: ignore[reportPrivateUsage] # the auth layer's access-group expansion, reused so /health scopes exactly like a request
_is_wildcard_pattern, # pyright: ignore[reportPrivateUsage] # the auth layer's pattern test, reused so /health scopes exactly like a request
_potential_models_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's per-layer alias expansion, reused so /health scopes exactly like a request
_resolve_all_team_model_sentinel_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's sentinel resolution, reused so /health scopes exactly like a request
_resolve_key_models_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's sentinel resolution, reused so /health scopes exactly like a request
get_authorized_resources_from_key_access_groups,
@ -1003,10 +1004,13 @@ async def _health_team_scope(
async def _health_team_member_scope(
caller: UserAPIKeyAuth, llm_router: Router | None, stores: _AuthStores
caller: UserAPIKeyAuth,
team_object: LiteLLM_TeamTableCachedObj | None,
llm_router: Router | None,
stores: _AuthStores,
) -> _ModelScope | None:
"""Auth's team-member layer: the member's own allowlist, when the membership's budget names one."""
if caller.team_id is None or caller.user_id is None or stores.prisma_client is None:
"""Auth's team-member layer: the member's own allowlist, when the membership's budget names one and the team row loaded, as ``common_checks`` gates it."""
if team_object is None or caller.team_id is None or caller.user_id is None or stores.prisma_client is None:
return None
membership: Final = await get_team_membership(
user_id=caller.user_id,
@ -1069,7 +1073,7 @@ async def _health_caller_model_scopes(
scopes: Final = (
await _health_key_scope(caller, llm_router, stores),
await _health_team_scope(caller, team_object, llm_router, stores),
await _health_team_member_scope(caller, llm_router, stores),
await _health_team_member_scope(caller, team_object, llm_router, stores),
await _health_user_scope(caller, llm_router, stores),
await _health_project_scope(caller, llm_router, stores),
)
@ -1156,34 +1160,88 @@ def _alias_names_for(own_names: Sequence[str], caller: UserAPIKeyAuth, llm_route
)
def _router_serves_through(llm_router: Router, model: str, team_id: str | None, deployment_id: str | None) -> bool:
"""Whether a request for ``model`` from this team lands on the deployment: ``Router.get_model_list`` resolves a name the way a request does and reaches a wildcard deployment only when none carries the exact name."""
served_by: Final = llm_router.get_model_list(model_name=model, team_id=team_id) or ()
return deployment_id is not None and any(_deployment_id(row) == deployment_id for row in served_by)
def _wildcard_names(deployment: Mapping[str, object]) -> tuple[str, ...]:
"""The deployment's own names that are wildcard patterns rather than a model a request can name outright."""
model_info: Final = deployment.get("model_info")
info: Final = model_info if isinstance(model_info, Mapping) else _NO_ENTRIES
return tuple(
name
for name in (deployment.get("model_name"), info.get("team_public_model_name"))
if isinstance(name, str) and _is_wildcard_pattern(name)
)
def _wildcard_candidates(
deployment: Mapping[str, object],
own_names: Sequence[str],
_NO_WILDCARD_ROUTES: Final[Mapping[str, frozenset[str]]] = MappingProxyType({})
def _wildcard_route_targets(
model_scopes: Sequence[_ModelScope],
requested_model: str | None,
team_id: str | None,
llm_router: Router | None,
) -> tuple[str, ...]:
"""The models a caller sends through a wildcard-named deployment: the one asked for, else every allowlist entry the router resolves to this deployment, which a deployment carrying the entry's exact name pre-empts as it does for a request."""
if llm_router is None or not any(_is_wildcard_pattern(name) for name in own_names):
return ()
if requested_model is not None:
return (requested_model,)
model_list: Sequence[Mapping[str, object]],
) -> Mapping[str, frozenset[str]]:
"""The deployments each allowlist entry a wildcard could serve actually routes to, resolved once per call.
``Router.get_model_list`` walks every deployment, so resolving an entry per wildcard row costs the
whole list each time. It also reaches a wildcard deployment only when none carries the entry's exact
name, which is the pre-emption a request gets, so the ids it returns are what the caller would hit.
"""
if llm_router is None:
return _NO_WILDCARD_ROUTES
patterns: Final = tuple(name for row in model_list for name in _wildcard_names(row))
if not patterns:
return _NO_WILDCARD_ROUTES
access_groups: Final = llm_router.get_model_access_groups()
deployment_id: Final = _deployment_id(deployment)
return tuple(
entries: Final = frozenset(
entry
for scope in model_scopes
for entry in scope.models
if entry not in access_groups
and entry != SpecialModelNames.all_team_models.value
and _router_serves_through(llm_router, entry, team_id, deployment_id)
and any(pattern_serves_model(pattern, entry) for pattern in patterns)
)
return MappingProxyType(
{
entry: frozenset(
ident
for row in (llm_router.get_model_list(model_name=entry, team_id=team_id) or ())
if (ident := _deployment_id(row)) is not None
)
for entry in entries
}
)
def _wildcard_candidates(
deployment: Mapping[str, object],
own_names: Sequence[str],
requested_model: str | None,
wildcard_routes: Mapping[str, frozenset[str]],
) -> tuple[str, ...]:
"""The models a caller sends through a wildcard-named deployment: the one asked for, else every allowlist entry the router resolves to this deployment."""
if not any(_is_wildcard_pattern(name) for name in own_names):
return ()
if requested_model is not None:
return (requested_model,)
deployment_id: Final = _deployment_id(deployment)
return tuple(entry for entry, ids in wildcard_routes.items() if deployment_id in ids)
def _passes_every_layer(candidate: str, model_scopes: Sequence[_ModelScope], llm_router: Router | None) -> bool:
"""Whether one model name clears every allowlist layer, each layer also trying the alias it resolves to, as ``_can_object_call_model`` does per layer."""
names: Final = _potential_models_for_auth_check(model=candidate, llm_router=llm_router)
return all(
any(
_check_model_access_helper(
model=name,
llm_router=llm_router,
models=scope.models,
team_model_aliases=scope.team_model_aliases,
team_id=scope.team_id,
)
for name in names
)
for scope in model_scopes
)
@ -1194,6 +1252,7 @@ def _caller_may_probe_deployment(
requested_model: str | None,
llm_router: Router | None,
caller_is_admin: bool,
wildcard_routes: Mapping[str, frozenset[str]] = _NO_WILDCARD_ROUTES,
) -> bool:
"""Mirror auth: the deployment must be reachable by the caller's team, and one model it serves for the caller must pass every allowlist layer."""
if not caller_is_admin and not Router._deployment_usable_by_team(deployment, caller.team_id):
@ -1202,20 +1261,31 @@ def _caller_may_probe_deployment(
candidates: Final = (
own_names
+ _alias_names_for(own_names, caller, llm_router)
+ _wildcard_candidates(deployment, own_names, model_scopes, requested_model, caller.team_id, llm_router)
+ _wildcard_candidates(deployment, own_names, requested_model, wildcard_routes)
)
return any(
all(
_check_model_access_helper(
model=candidate,
llm_router=llm_router,
models=scope.models,
team_model_aliases=scope.team_model_aliases,
team_id=scope.team_id,
)
for scope in model_scopes
)
for candidate in candidates
return any(_passes_every_layer(candidate, model_scopes, llm_router) for candidate in candidates)
async def _health_readable_deployment_ids(
caller: UserAPIKeyAuth,
llm_router: Router | None,
llm_model_list: Sequence[Mapping[str, object]] | None,
stores: _AuthStores,
) -> frozenset[str] | None:
"""The deployments whose stored health rows this caller may read, scoped exactly as ``GET /health`` scopes live results; ``None`` leaves them unscoped."""
model_scopes: Final = await _health_caller_model_scopes(caller, llm_router, stores)
is_admin: Final = _is_proxy_admin(caller)
if is_admin and not model_scopes:
return None
deployments: Final = tuple(llm_model_list or ())
alias_copies: Final = _router_alias_copies(llm_router)
wildcard_routes: Final = _wildcard_route_targets(model_scopes, caller.team_id, llm_router, deployments)
return frozenset(
ident
for row in deployments
if not _is_router_alias_copy(row, alias_copies)
and _caller_may_probe_deployment(row, caller, model_scopes, None, llm_router, is_admin, wildcard_routes)
and (ident := _deployment_id(row)) is not None
)
@ -1452,6 +1522,11 @@ async def health_endpoint(
)
restrict_to_allowed_models: Final = not is_admin or bool(model_scopes)
alias_copies: Final = _router_alias_copies(llm_router)
wildcard_routes: Final = (
_wildcard_route_targets(model_scopes, user_api_key_dict.team_id, llm_router, llm_model_list)
if restrict_to_allowed_models
else _NO_WILDCARD_ROUTES
)
_llm_model_list: Final = [
m
for m in copy.deepcopy(llm_model_list)
@ -1459,7 +1534,7 @@ async def health_endpoint(
and (
not restrict_to_allowed_models
or _caller_may_probe_deployment(
m, user_api_key_dict, model_scopes, requested_model, llm_router, is_admin
m, user_api_key_dict, model_scopes, requested_model, llm_router, is_admin, wildcard_routes
)
)
]
@ -1552,9 +1627,22 @@ async def health_check_history_endpoint(
Returns historical health check data with optional filtering.
"""
from litellm.proxy.proxy_server import (
llm_model_list,
llm_router,
proxy_logging_obj,
user_api_key_cache,
)
prisma_client: Final = _check_prisma_client()
try:
readable_ids: Final = await _health_readable_deployment_ids(
user_api_key_dict,
llm_router,
llm_model_list,
_AuthStores(prisma_client, user_api_key_cache, proxy_logging_obj),
)
history: Final = await prisma_client.get_health_check_history(
model_name=model,
limit=limit,
@ -1563,7 +1651,11 @@ async def health_check_history_endpoint(
)
# Convert to dict format for JSON response using helper function
history_data: Final = [_convert_health_check_to_dict(check) for check in history]
history_data: Final = [
_convert_health_check_to_dict(check)
for check in history
if readable_ids is None or check.model_id in readable_ids
]
return {
"health_checks": history_data,
@ -1588,15 +1680,29 @@ async def latest_health_checks_endpoint(
Returns the most recent health check result for each model.
"""
from litellm.proxy.proxy_server import (
llm_model_list,
llm_router,
proxy_logging_obj,
user_api_key_cache,
)
prisma_client: Final = _check_prisma_client()
try:
readable_ids: Final = await _health_readable_deployment_ids(
user_api_key_dict,
llm_router,
llm_model_list,
_AuthStores(prisma_client, user_api_key_cache, proxy_logging_obj),
)
latest_checks: Final = await prisma_client.get_all_latest_health_checks()
# Convert to dict format for JSON response using helper function
checks_data: Final = {
(check.model_id if check.model_id else check.model_name): _convert_health_check_to_dict(check)
for check in latest_checks
if readable_ids is None or check.model_id in readable_ids
}
return {

119
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -3713,6 +3713,125 @@ async def test_health_endpoint_leaves_the_project_layer_off_when_the_project_row
assert probed == {"id-bedrock", "id-openai"}
@pytest.mark.asyncio
@pytest.mark.parametrize("model", [None, "nova-router-alias"])
async def test_health_endpoint_keeps_a_model_whose_alias_and_target_sit_on_different_allowlist_layers(model):
"""Auth checks each layer against both the alias and its target, so a key allowed the alias on a team allowed the target may health-check the deployment it can call."""
router = Router(
model_list=copy.deepcopy(_ACCESS_GROUP_MODEL_LIST),
model_group_alias={"nova-router-alias": "bedrock-nova"},
)
probed = await _live_probed_model_ids(
_ACCESS_GROUP_MODEL_LIST,
UserAPIKeyAuth(api_key="hashed-test-key", models=["nova-router-alias"], team_id="team-a"),
model=model,
router=router,
user_api_key_cache=await _auth_cache_with(_team_a(["bedrock-nova"])),
)
assert probed == {"id-bedrock"}
@pytest.mark.asyncio
async def test_health_endpoint_leaves_the_team_member_layer_off_when_the_team_row_cannot_be_read():
"""``common_checks`` applies a membership's own allowlist only once the team row loads, so /health must not be stricter than chat when it does not."""
membership = LiteLLM_TeamMembership(
user_id="u1", team_id="team-a", litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["bedrock-nova"])
)
probed = await _live_probed_model_ids(
_ACCESS_GROUP_MODEL_LIST,
UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_id="u1", team_id="team-a", team_models=[]),
user_api_key_cache=await _auth_cache_with(
(team_membership_reservation_cache_key(user_id="u1", team_id="team-a"), membership)
),
prisma_client=SimpleNamespace(
db=SimpleNamespace(litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=None)))
),
)
assert probed == {"id-bedrock", "id-openai"}
def _stored_health_row(model_id: str, model_name: str):
return SimpleNamespace(
health_check_id=f"hc-{model_id}",
model_name=model_name,
model_id=model_id,
status="healthy",
healthy_count=1,
unhealthy_count=0,
error_message=None,
response_time_ms=12,
details=None,
checked_by="u1",
checked_at=None,
created_at=None,
)
_STORED_HEALTH_ROWS = (
_stored_health_row("id-bedrock", "bedrock-nova"),
_stored_health_row("id-openai", "gpt-5.4-mini"),
)
def _stored_health_prisma():
return SimpleNamespace(
get_health_check_history=AsyncMock(return_value=list(_STORED_HEALTH_ROWS)),
get_all_latest_health_checks=AsyncMock(return_value=list(_STORED_HEALTH_ROWS)),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("caller", "expected_ids"),
[
(UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"]), {"id-bedrock"}),
(
UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN),
{"id-bedrock", "id-openai"},
),
],
)
async def test_health_latest_endpoint_scopes_stored_rows_to_the_callers_own_deployments(caller, expected_ids):
"""A stored health row names a deployment, so /health/latest must hand a key only the rows /health would probe for it."""
from litellm.proxy.health_endpoints._health_endpoints import latest_health_checks_endpoint
with _proxy_health_globals(
_ACCESS_GROUP_MODEL_LIST, _ACCESS_GROUP_ROUTER, prisma_client=_stored_health_prisma()
):
result = await latest_health_checks_endpoint(user_api_key_dict=caller)
assert set(result["latest_health_checks"]) == expected_ids
assert result["total_models"] == len(expected_ids)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("caller", "expected_ids"),
[
(UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"]), ["id-bedrock"]),
(
UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN),
["id-bedrock", "id-openai"],
),
],
)
async def test_health_history_endpoint_scopes_stored_rows_to_the_callers_own_deployments(caller, expected_ids):
"""/health/history reads the same stored rows, so it must not hand a key the history of a deployment it cannot probe."""
from litellm.proxy.health_endpoints._health_endpoints import health_check_history_endpoint
with _proxy_health_globals(
_ACCESS_GROUP_MODEL_LIST, _ACCESS_GROUP_ROUTER, prisma_client=_stored_health_prisma()
):
result = await health_check_history_endpoint(user_api_key_dict=caller)
assert [row["model_id"] for row in result["health_checks"]] == expected_ids
assert result["total_records"] == len(expected_ids)
@pytest.mark.asyncio
async def test_health_endpoint_ignores_the_users_allowlist_for_a_proxy_admin():
"""Auth swaps an unrestricted user object in for a proxy admin, so the admin's own user row must not narrow /health."""