diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7877d887409..ab8e51c7c0c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4641,6 +4641,7 @@ async def info_key_fn( ), ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + include_effective_models: bool = False, ): """ Retrieve information about a key. @@ -4650,6 +4651,12 @@ async def info_key_fn( prefer the hash, since a query parameter is recorded verbatim by any HTTP access log in front of the proxy. Defaults to the key in the Authorization header. + Set include_effective_models=true to add info.effective_models: known, registered + model names resolved from the target key's grants and native model restrictions. + Stored info.models is unchanged. Inactive keys return an empty list. This is + model eligibility, not a check of budgets, provider health or request-context policies. + Only explicitly registered routes are listed; wildcard routes are not enumerated. + Returns: - key: str - The key that was looked up, echoed back as it was passed in - info: dict - The key's row, minus the hashed token. Deleted keys are served from the @@ -4695,8 +4702,12 @@ async def info_key_fn( ``` """ from litellm.proxy.proxy_server import ( + general_settings, + llm_router, model_max_budget_limiter, prisma_client, + proxy_logging_obj, + user_api_key_cache, ) try: @@ -4745,6 +4756,22 @@ async def info_key_fn( "deleted" if live_key_info is None else _derive_key_status(key_info_dict, now=datetime.now(timezone.utc)) ) + if include_effective_models: + from litellm.proxy.management_endpoints.key_model_access import effective_key_models + + key_info_dict["effective_models"] = list( # mutable-ok: preserve the JSON models-list response shape + await effective_key_models( + key=UserAPIKeyAuth.model_validate(key_info.model_dump(exclude_none=True)), + router=llm_router, + prisma_client=prisma_client, + cache=user_api_key_cache, + logging=proxy_logging_obj, + settings=TypeAdapter(Mapping[str, object]).validate_python(general_settings), + ) + if key_info_dict["status"] == "active" + else () + ) + model_max_budget = key_info_dict.get("model_max_budget") or {} budget_table: Final = key_info_dict.get("litellm_budget_table") or {} if not model_max_budget and isinstance(budget_table, dict): diff --git a/litellm/proxy/management_endpoints/key_model_access.py b/litellm/proxy/management_endpoints/key_model_access.py new file mode 100644 index 00000000000..c66f5eb7b56 --- /dev/null +++ b/litellm/proxy/management_endpoints/key_model_access.py @@ -0,0 +1,107 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.proxy._types import ModelAccessDeniedProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _check_agent_access_group_model_access, # pyright: ignore[reportPrivateUsage] # Reuse native model policy checks without duplicating them + _check_agent_caller_model_access, # pyright: ignore[reportPrivateUsage] # Reuse native model policy checks without duplicating them + _user_object_or_none, # pyright: ignore[reportPrivateUsage] # Reuse native model policy checks without duplicating them + can_key_call_resolved_model, + can_user_call_model, + get_project_object, + get_team_object, + key_model_aliases_for_auth_check, +) +from litellm.proxy.auth.team_grants import team_model_aliases +from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router import Router +from litellm.types.router import ModelInfo + + +async def effective_key_models( + key: UserAPIKeyAuth, + router: Router | None, + prisma_client: PrismaClient, + cache: UserApiKeyCache, + logging: ProxyLogging, + settings: Mapping[str, object] = MappingProxyType({}), +) -> tuple[str, ...]: + if router is None: + return () + team: Final = ( + await get_team_object( + team_id=key.team_id, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=logging, + ) + if key.team_id + else None + ) + project: Final = ( + await get_project_object( + project_id=key.project_id, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=logging, + ) + if key.project_id + else None + ) + if (team is not None and team.blocked) or (project is not None and project.blocked): + return () + target: Final = UserAPIKeyAuth.model_validate( + MappingProxyType( + { + **key.model_dump(), + "team_models": team.models if team else (), + "team_model_aliases": team_model_aliases(team), + "user_role": None, + } + ) + ) + owner: Final = ( + await _user_object_or_none(target, prisma_client, cache, logging) + if target.team_id is None and target.user_id is not None + else None + ) + registrations: Final = router.get_model_list() or () + candidates: Final = frozenset(row["model_name"] for row in registrations) + + async def allowed(model: str) -> bool: + if "*" in model: + return False + deployments: Final = router.get_model_list(model_name=model, team_id=target.team_id) or () + if not any( + ModelInfo.model_validate(deployment.get("model_info") or MappingProxyType({})).team_id # pyright: ignore[reportUnknownMemberType] # Validate the router’s legacy untyped model_info here + in (None, target.team_id) + for deployment in deployments + ): + return False + try: + await can_key_call_resolved_model(model, None, target, router) + await _check_agent_access_group_model_access(model=model, valid_token=target, llm_router=router) + await _check_agent_caller_model_access( + model=model, + valid_token=target, + llm_router=router, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=logging, + ) + await can_user_call_model( + model=model, + llm_router=router, + user_object=owner, + key_model_aliases=key_model_aliases_for_auth_check(target), + ) + except ModelAccessDeniedProxyException: + return False + return True + + eligible: Final = tuple(frozenset([model for model in candidates if await allowed(model)])) + entries: Final = TeamModelNameTranslator.listing_entries(eligible, router, settings) + return tuple(sorted(public_name for public_name, _ in entries)) diff --git a/tests/unit/proxy/management_endpoints/test_key_model_access.py b/tests/unit/proxy/management_endpoints/test_key_model_access.py new file mode 100644 index 00000000000..a939820c256 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_key_model_access.py @@ -0,0 +1,270 @@ +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm import Router +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + LiteLLM_VerificationToken, + LitellmUserRoles, + ProxyException, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + +HASH: Final = "a" * 64 + + +@dataclass(frozen=True) +class KeyInfoContext: + cache: UserApiKeyCache + database: MagicMock + router: Router + key: LiteLLM_VerificationToken + + +@pytest.fixture +def context(monkeypatch: pytest.MonkeyPatch) -> KeyInfoContext: + router: Final = Router( + model_list=[ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, + "model_info": {"access_groups": ["big", "small"]}, + }, + { + "model_name": "model-b", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, + "model_info": {"access_groups": ["big"]}, + }, + { + "model_name": "other-team-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, + "model_info": {"access_groups": ["big"], "team_id": "other-team"}, + }, + ] + ) + cache: Final = UserApiKeyCache() + cache.set_cache("owner", LiteLLM_UserTable(user_id="owner", models=[])) + cache.set_cache("team_id:team", LiteLLM_TeamTableCachedObj(team_id="team", models=["small"])) + key: Final = LiteLLM_VerificationToken(token=HASH, user_id="owner", models=["big", "small"]) + database: Final = MagicMock() + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key) + database.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + return KeyInfoContext(cache, database, router, key) + + +async def test_hash_lookup_expands_groups_without_changing_stored_grants(context: KeyInfoContext) -> None: + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["models"] == ["big", "small"] + assert result["info"]["effective_models"] == ["model-a", "model-b"] + assert context.key.models == ["big", "small"] + context.database.db.litellm_verificationtoken.find_unique.assert_awaited_once_with( + where={"token": HASH}, + include={"litellm_budget_table": True}, + ) + + +async def test_omitting_flag_preserves_existing_response(context: KeyInfoContext) -> None: + result: Final = await info_key_fn( + key=HASH, user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + ) + assert "effective_models" not in result["info"] + assert result["info"]["models"] == ["big", "small"] + + +async def test_admin_lookup_uses_target_team_ceiling_not_admin_access(context: KeyInfoContext) -> None: + context.database.db.litellm_verificationtoken.find_unique.return_value = context.key.model_copy( + update={"team_id": "team"} + ) + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, models=["all-proxy-models"]), + ) + assert result["info"]["models"] == ["big", "small"] + assert result["info"]["effective_models"] == ["model-a"] + + +@pytest.mark.parametrize("grants", [[], ["all-proxy-models"], ["model-*"], ["model-a", "big"]]) +async def test_legacy_grants_list_only_eligible_registered_models(context: KeyInfoContext, grants: list[str]) -> None: + context.database.db.litellm_verificationtoken.find_unique.return_value = context.key.model_copy( + update={"models": grants} + ) + result: Final = await info_key_fn( + key=HASH, include_effective_models=True, user_api_key_dict=UserAPIKeyAuth(user_id="owner") + ) + assert result["info"]["models"] == grants + assert result["info"]["effective_models"] == ["model-a", "model-b"] + + +async def test_another_user_cannot_expand_target_key(context: KeyInfoContext) -> None: + with pytest.raises(ProxyException) as error: + await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_id="stranger", user_role=LitellmUserRoles.INTERNAL_USER), + ) + assert str(error.value.code) == "403" + + +@pytest.mark.parametrize("state", ["blocked", "expired", "deleted"]) +async def test_inactive_key_has_no_effective_models(context: KeyInfoContext, state: str) -> None: + row: Final = context.key.model_copy( + update={ + "blocked": state == "blocked", + "expires": datetime.now(timezone.utc) - timedelta(days=1) if state == "expired" else None, + } + ) + if state == "deleted": + context.database.db.litellm_verificationtoken.find_unique.return_value = None + context.database.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=row) + else: + context.database.db.litellm_verificationtoken.find_unique.return_value = row + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["effective_models"] == [] + + +async def test_group_change_is_visible_without_rewriting_key(context: KeyInfoContext) -> None: + caller: Final = UserAPIKeyAuth(user_id="owner") + before: Final = await info_key_fn(key=HASH, include_effective_models=True, user_api_key_dict=caller) + context.router.add_deployment( + Deployment( + model_name="model-c", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", api_key="fake"), + model_info=ModelInfo(access_groups=["big"]), + ) + ) + after: Final = await info_key_fn(key=HASH, include_effective_models=True, user_api_key_dict=caller) + assert before["info"]["effective_models"] == ["model-a", "model-b"] + assert after["info"]["effective_models"] == ["model-a", "model-b", "model-c"] + assert before["info"]["models"] == after["info"]["models"] == ["big", "small"] + + +async def test_personal_key_obeys_owner_models(context: KeyInfoContext) -> None: + context.cache.set_cache("owner", LiteLLM_UserTable(user_id="owner", models=["small"])) + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["effective_models"] == ["model-a"] + + +async def test_blocked_team_has_no_effective_models(context: KeyInfoContext) -> None: + context.cache.set_cache("team_id:team", LiteLLM_TeamTableCachedObj(team_id="team", models=["big"], blocked=True)) + context.database.db.litellm_verificationtoken.find_unique.return_value = context.key.model_copy( + update={"team_id": "team"} + ) + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["effective_models"] == [] + + +async def test_http_query_flag_adds_effective_models(context: KeyInfoContext) -> None: + import httpx + from fastapi import FastAPI + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.management_endpoints.key_management_endpoints import router + + app: Final = FastAPI() + app.include_router(router) + + async def caller() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + + app.dependency_overrides[user_api_key_auth] = caller + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response: Final = await client.get(f"/key/info?key={HASH}&include_effective_models=true") + assert response.status_code == 200, response.text + assert response.json()["info"]["models"] == ["big", "small"] + assert response.json()["info"]["effective_models"] == ["model-a", "model-b"] + + +async def test_key_group_is_a_ceiling_even_when_other_routes_exist(context: KeyInfoContext) -> None: + context.database.db.litellm_verificationtoken.find_unique.return_value = context.key.model_copy( + update={"models": ["small"]} + ) + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["effective_models"] == ["model-a"] + + +async def test_project_models_narrow_key_groups(context: KeyInfoContext) -> None: + from litellm.proxy._types import LiteLLM_ProjectTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import project_cache_key + + context.cache.set_cache( + project_cache_key("project"), LiteLLM_ProjectTableCachedObj(project_id="project", models=["small"]) + ) + context.database.db.litellm_verificationtoken.find_unique.return_value = context.key.model_copy( + update={"project_id": "project"} + ) + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["effective_models"] == ["model-a"] + + +async def test_wildcard_routes_are_not_enumerated(context: KeyInfoContext) -> None: + context.router.add_deployment( + Deployment(model_name="openai/*", litellm_params=LiteLLM_Params(model="openai/*", api_key="fake")) + ) + context.database.db.litellm_verificationtoken.find_unique.return_value = context.key.model_copy( + update={"models": ["all-proxy-models"]} + ) + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["effective_models"] == ["model-a", "model-b"] + + +@pytest.mark.parametrize("public_name", [None, "project-model"]) +async def test_matching_team_can_list_its_routes(context: KeyInfoContext, public_name: str | None) -> None: + route_name: Final = "model_name_team_example" if public_name else "team-route" + context.router.add_deployment( + Deployment( + model_name=route_name, + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", api_key="fake"), + model_info=ModelInfo(team_id="team", team_public_model_name=public_name, access_groups=["big"]), + ) + ) + context.cache.set_cache("team_id:team", LiteLLM_TeamTableCachedObj(team_id="team", models=["big"])) + context.database.db.litellm_verificationtoken.find_unique.return_value = context.key.model_copy( + update={"team_id": "team"} + ) + result: Final = await info_key_fn( + key=HASH, + include_effective_models=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["info"]["effective_models"] == sorted(["model-a", "model-b", public_name or route_name])