fix(health): list a wildcard deployment only for models the router sends through it, degrade a failed project read like auth

This commit is contained in:
mateo-berri 2026-09-05 20:22:30 -07:00
parent 3eda47b2b2
commit bd95fd3399
2 changed files with 133 additions and 16 deletions

View file

@ -1046,15 +1046,18 @@ async def _health_user_scope(
async def _health_project_scope(
caller: UserAPIKeyAuth, llm_router: Router | None, stores: _AuthStores
) -> _ModelScope | None:
"""Auth's project layer: the project's own allowlist."""
"""Auth's project layer: the project's own allowlist; a row that cannot be read leaves the layer off, as auth's ``_safe_fetch`` does."""
if caller.project_id is None or stores.prisma_client is None:
return None
project_object: Final = await get_project_object(
project_id=caller.project_id,
prisma_client=stores.prisma_client,
user_api_key_cache=stores.user_api_key_cache,
proxy_logging_obj=stores.proxy_logging_obj,
)
try:
project_object: Final = await get_project_object(
project_id=caller.project_id,
prisma_client=stores.prisma_client,
user_api_key_cache=stores.user_api_key_cache,
proxy_logging_obj=stores.proxy_logging_obj,
)
except Exception: # noqa: BLE001 # auth's _safe_fetch swallows every failure here and lets the request run on the token's own fields
return None
return _model_scope(tuple(project_object.models) if project_object is not None else (), llm_router, None, None)
@ -1153,26 +1156,34 @@ 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_candidates(
deployment: Mapping[str, object],
own_names: Sequence[str],
model_scopes: Sequence[_ModelScope],
requested_model: str | None,
team_id: str | None,
llm_router: Router | None,
) -> tuple[str, ...]:
"""The concrete models a caller could send through a wildcard-named deployment: the one asked for, else every allowlist entry its pattern serves."""
patterns: Final = tuple(name for name in own_names if _is_wildcard_pattern(name))
if not patterns:
"""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,) if any(pattern_serves_model(p, requested_model) for p in patterns) else ()
access_groups: Final = llm_router.get_model_access_groups() if llm_router is not None else _NO_ENTRIES
return (requested_model,)
access_groups: Final = llm_router.get_model_access_groups()
deployment_id: Final = _deployment_id(deployment)
return tuple(
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 any(pattern_serves_model(p, entry) for p in patterns)
and _router_serves_through(llm_router, entry, team_id, deployment_id)
)
@ -1191,7 +1202,7 @@ def _caller_may_probe_deployment(
candidates: Final = (
own_names
+ _alias_names_for(own_names, caller, llm_router)
+ _wildcard_candidates(own_names, model_scopes, requested_model, llm_router)
+ _wildcard_candidates(deployment, own_names, model_scopes, requested_model, caller.team_id, llm_router)
)
return any(
all(

View file

@ -1733,6 +1733,7 @@ def _proxy_health_globals(
use_background_health_checks: bool = False,
health_check_results: Mapping[str, object] | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
prisma_client: object | None = None,
) -> Iterator[None]:
with (
patch( # test-quality-ok: proxy module global, no injection seam
@ -1742,7 +1743,8 @@ def _proxy_health_globals(
"litellm.proxy.proxy_server.llm_router", llm_router
),
patch( # test-quality-ok: proxy module global, no injection seam
"litellm.proxy.proxy_server.prisma_client", None if user_api_key_cache is None else MagicMock()
"litellm.proxy.proxy_server.prisma_client",
prisma_client or (None if user_api_key_cache is None else MagicMock()),
),
patch( # test-quality-ok: proxy module global, no injection seam
"litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache or UserApiKeyCache()
@ -2824,6 +2826,7 @@ async def _live_probed_model_ids(
model_id: str | None = None,
router: Router | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
prisma_client: object | None = None,
) -> set[str]:
from fastapi import Response
@ -2838,7 +2841,10 @@ async def _live_probed_model_ids(
with (
_proxy_health_globals(
model_list, router if router is not None else _router_for(model_list), user_api_key_cache=user_api_key_cache
model_list,
router if router is not None else _router_for(model_list),
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
),
patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
@ -3153,6 +3159,48 @@ async def test_health_endpoint_shows_a_wildcard_deployment_that_serves_a_key_mod
assert probed == {"id-bedrock-wildcard"}
_WILDCARD_BESIDE_EXACT_NAME_MODEL_LIST = [
{"model_name": "openai/*", "litellm_params": {"model": "openai/*"}, "model_info": {"id": "id-openai-wildcard"}},
_ACCESS_GROUP_MODEL_LIST[1],
]
_WILDCARD_BESIDE_ANOTHER_TEAMS_EXACT_NAME_MODEL_LIST = [
_WILDCARD_BESIDE_EXACT_NAME_MODEL_LIST[0],
{
"model_name": "gpt-5.4-mini_team-b_9f2c",
"litellm_params": {"model": "openai/gpt-5.4-mini"},
"model_info": {"id": "id-team-b-mini", "team_id": "team-b", "team_public_model_name": "gpt-5.4-mini"},
},
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("model_list", "team_models", "expected_ids"),
[
(_WILDCARD_BESIDE_EXACT_NAME_MODEL_LIST, ["gpt-5.4-mini"], {"id-openai"}),
(
_WILDCARD_BESIDE_EXACT_NAME_MODEL_LIST,
["gpt-5.4-mini", "openai/gpt-5.4-nano"],
{"id-openai", "id-openai-wildcard"},
),
(_WILDCARD_BESIDE_ANOTHER_TEAMS_EXACT_NAME_MODEL_LIST, ["gpt-5.4-mini"], {"id-openai-wildcard"}),
],
)
async def test_health_endpoint_shows_a_wildcard_deployment_only_for_a_model_the_router_sends_through_it(
model_list, team_models, expected_ids
):
"""
The router sends a name to the deployment carrying it exactly and reaches ``openai/*`` only for a
name no deployment this team can use carries, so /health lists the wildcard deployment on the same terms.
"""
probed = await _live_probed_model_ids(
model_list,
UserAPIKeyAuth(api_key="hashed-test-key", models=[], team_id="team-a", team_models=team_models),
)
assert probed == expected_ids
@pytest.mark.asyncio
async def test_health_endpoint_hides_a_wildcard_deployment_when_no_single_model_passes_both_the_key_and_the_team_allowlist():
"""Auth checks the one requested model against both allowlists, so a key and a team allowing different models call nothing through ``bedrock/*``."""
@ -3648,6 +3696,64 @@ async def test_health_endpoint_applies_the_projects_allowlist():
assert probed == {"id-bedrock"}
@pytest.mark.asyncio
async def test_health_endpoint_leaves_the_project_layer_off_when_the_project_row_cannot_be_read():
"""Auth's ``_safe_fetch`` swallows a failed project read and lets the request run on the token, so /health must not 500 there."""
db_down = SimpleNamespace(
db=SimpleNamespace(litellm_projecttable=SimpleNamespace(find_unique=AsyncMock(side_effect=ClientNotConnectedError())))
)
probed = await _live_probed_model_ids(
_ACCESS_GROUP_MODEL_LIST,
UserAPIKeyAuth(api_key="hashed-test-key", models=[], project_id="p1"),
user_api_key_cache=UserApiKeyCache(),
prisma_client=db_down,
)
assert probed == {"id-bedrock", "id-openai"}
@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."""
probed = await _live_probed_model_ids(
_ACCESS_GROUP_MODEL_LIST,
UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
user_api_key_cache=await _auth_cache_with(("admin", LiteLLM_UserTable(user_id="admin", models=["bedrock-nova"]))),
)
assert probed == {"id-bedrock", "id-openai"}
@pytest.mark.asyncio
async def test_health_endpoint_scopes_a_proxy_admins_key_by_its_team_allowlist():
"""Auth holds an admin's key on a restricted team to that team's allowlist, so /health must narrow it the same way."""
probed = await _live_probed_model_ids(
_ACCESS_GROUP_MODEL_LIST,
UserAPIKeyAuth(
api_key="hashed-test-key",
models=[],
team_id="team-a",
team_models=["gpt-5.4-mini"],
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
assert probed == {"id-openai"}
@pytest.mark.asyncio
async def test_health_endpoint_reads_the_team_allowlist_from_the_team_row_auth_loads():
"""Auth checks the team row it loads, not the allowlist copied onto the token when the key was minted."""
probed = await _live_probed_model_ids(
_ACCESS_GROUP_MODEL_LIST,
UserAPIKeyAuth(api_key="hashed-test-key", models=[], team_id="team-a", team_models=["gpt-5.4-mini"]),
user_api_key_cache=await _auth_cache_with(_team_a(["bedrock-nova"])),
)
assert probed == {"id-bedrock"}
@pytest.mark.asyncio
async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team_on_background_cache_path():
from fastapi import Response