mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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:
parent
3eda47b2b2
commit
bd95fd3399
2 changed files with 133 additions and 16 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue