From 8477fe4108b742acbf7e40aabb79e3147fb9d329 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 06:41:13 -0700 Subject: [PATCH] fix(proxy): list key and team model aliases in GET /v1/models (#42908) * fix(proxy): list key and team model aliases in GET /v1/models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep alias listing helpers within the type discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover alias rows on GET /v1/models and /v1/models/{id} Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply team then key aliases like chat completions and keep the alias as the retrieved id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply key aliases twice like chat completions and skip only malformed alias entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply the global model_alias_map between the key alias passes like chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): list only the caller's own aliases and never rewrite a listed model id on retrieval Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format model_info alias lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): hide undiscoverable names from model retrieval so an alias named like one resolves to its target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep undiscoverable models retrievable by id while excluding them from the alias guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): pass an immutable name sequence into the model_info alias guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type the model list alias test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): annotate the new alias listing test fixtures and helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/model_listing_utils.py | 80 +++++++- litellm/proxy/proxy_server.py | 39 +++- tests/e2e/coverage_registry/other.yaml | 1 + tests/e2e/models.py | 1 + tests/e2e/other/other_client.py | 16 +- tests/e2e/other/test_jwt_auth_e2e.py | 47 ++++- .../common_utils/test_model_listing_utils.py | 177 +++++++++++++++--- .../proxy/test_model_list_aliases.py | 151 +++++++++++++++ 8 files changed, 477 insertions(+), 35 deletions(-) create mode 100644 tests/test_litellm/proxy/test_model_list_aliases.py diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 4b3ff2711b3..3c6555662e6 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -13,9 +13,12 @@ from __future__ import annotations import re from collections.abc import Container, Mapping, Sequence from dataclasses import dataclass +from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +from pydantic import TypeAdapter, ValidationError + import litellm if TYPE_CHECKING: @@ -28,6 +31,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code" _CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-" _ONE_MILLION_SUFFIX: Final = "[1m]" _ONE_MILLION_TOKENS: Final = 1_000_000 +_ALIAS_ENTRIES: Final = TypeAdapter(Mapping[object, object]) +_NO_ALIASES: Final[Mapping[str, str]] = MappingProxyType({}) def configured_display_names( @@ -152,6 +157,77 @@ class ClaudeCodeRoutingNames: ) +@dataclass(frozen=True, slots=True) +class CallerAliases: + """`own` are the caller's key and team alias maps, the names `/v1/models` lists for it. + `rewrite` are the maps `/chat/completions` rewrites its model through, in the order it + applies them: the team's, the key's in `add_litellm_data_to_request`, then the global + `model_alias_map` and the key's again in `common_processing_pre_call_logic`.""" + + own: tuple[object, ...] + rewrite: tuple[object, ...] + + +def caller_alias_maps( + key_aliases: object, + team_aliases: object, + key_team_id: str | None, + listed_team_id: str | None, +) -> CallerAliases: + """Team aliases count only when listing the team the key authenticated as.""" + if listed_team_id is not None and listed_team_id != key_team_id: + return CallerAliases((key_aliases,), (key_aliases, litellm.model_alias_map, key_aliases)) + return CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, litellm.model_alias_map, key_aliases)) + + +def _alias_map(aliases: object) -> Mapping[str, str]: + try: + entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True) + except ValidationError: + return _NO_ALIASES + return MappingProxyType( + {alias: target for alias, target in entries.items() if isinstance(alias, str) and isinstance(target, str)} + ) + + +def _alias_names(alias_maps: Sequence[Mapping[str, str]]) -> tuple[str, ...]: + return tuple(dict.fromkeys(alias for aliases in alias_maps for alias in aliases)) + + +def _rewrite(model_id: str, alias_maps: Sequence[Mapping[str, str]]) -> str | None: + target: Final = reduce(lambda name, aliases: aliases.get(name, name), alias_maps, model_id) + return None if target == model_id else target + + +def alias_target(model_id: str, aliases: CallerAliases, listed: Container[str] = frozenset()) -> str | None: + """The model group `/chat/completions` rewrites `model_id` to, else None. A `model_id` + already `listed` keeps its own row, so it is never rewritten.""" + if model_id in listed: + return None + return _rewrite(model_id, tuple(_alias_map(alias_map) for alias_map in aliases.rewrite)) + + +def alias_listing_entries( + entries: Sequence[tuple[str, str]], + aliases: CallerAliases, +) -> tuple[tuple[str, str], ...]: + """`entries` plus one `(alias, lookup_id)` row per key or team alias whose target is + listed. An alias colliding with a listed id keeps the listed entry.""" + maps: Final = tuple(_alias_map(alias_map) for alias_map in aliases.rewrite) + own: Final = tuple(_alias_map(alias_map) for alias_map in aliases.own) + lookup_by_response: Final = MappingProxyType(dict(entries)) + lookup_ids: Final = frozenset(lookup_by_response.values()) + targets: Final = MappingProxyType( + {alias: _rewrite(alias, maps) for alias in _alias_names(own) if alias not in lookup_by_response} + ) + added: Final = tuple( + (alias, lookup_by_response.get(target, target)) + for alias, target in targets.items() + if target is not None and (target in lookup_by_response or target in lookup_ids) + ) + return (*entries, *added) + + def claude_code_requested_group( requested: str, llm_router: Router, @@ -218,7 +294,7 @@ class TeamModelNameTranslator: @staticmethod def _response_to_lookup_map( - model_names: list[str], + model_names: Sequence[str], internal_to_public: dict[str, str], ) -> dict[str, str]: """Map each public response id to the first internal lookup id seen in @@ -235,7 +311,7 @@ class TeamModelNameTranslator: @staticmethod def listing_entries( - model_names: list[str], + model_names: Sequence[str], llm_router: Router | None, general_settings: Mapping[str, object], ) -> list[tuple[str, str]]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7e18db742cd..9fb72d2d1b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -422,6 +422,9 @@ from litellm.proxy.common_utils.model_deprecation import collect_model_deprecati from litellm.proxy.common_utils.model_listing_utils import ( ClaudeCodeRoutingNames, TeamModelNameTranslator, + alias_listing_entries, + alias_target, + caller_alias_maps, claude_code_view_ids, configured_display_names, is_claude_code_client, @@ -11172,14 +11175,13 @@ async def model_list( view_aliases: Final = ( view_router_settings.get("model_group_alias") if isinstance(view_router_settings, Mapping) else None ) + caller_aliases: Final = caller_alias_maps( + user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id + ) routing_names: Final = ClaudeCodeRoutingNames( llm_router, team_id or user_api_key_dict.team_id, - ( - user_api_key_dict.aliases, - user_api_key_dict.team_model_aliases, - view_aliases, - ), + (*caller_aliases.rewrite, view_aliases), ) # Validate scope parameter if provided @@ -11307,7 +11309,9 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + entries: Final = alias_listing_entries( + TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), caller_aliases + ) for response_id, lookup_id in entries: model_info = create_model_info_response( model_id=lookup_id, @@ -11391,7 +11395,8 @@ async def model_info( ) # Mirror /v1/models' visibility filter so first-occurrence resolution - # cannot land on a deployment the listing had hidden. + # cannot land on a deployment the listing had hidden. Undiscoverable + # models stay retrievable by id, they only drop out of the alias guard. blocked_names: Final = llm_router.get_fully_blocked_model_names() if llm_router is not None else set() unhealthy_names: Final = await get_hidden_unhealthy_model_names( healthy_only=healthy_only, @@ -11401,10 +11406,25 @@ async def model_info( hidden_names: Final = blocked_names | unhealthy_names if hidden_names: all_models = [m for m in all_models if m not in hidden_names] + undiscoverable_names: Final = undiscoverable_model_names( + all_models, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id + ) internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) + aliased_model_id: Final = alias_target( + model_id, + caller_alias_maps( + user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id + ), + frozenset( + response_id + for response_id, _ in TeamModelNameTranslator.listing_entries( + tuple(m for m in all_models if m not in undiscoverable_names), llm_router, settings + ) + ), + ) resolved_model_id: Final = TeamModelNameTranslator.resolve_public_name( - model_id=model_id, + model_id=aliased_model_id or model_id, available_models=all_models, llm_router=llm_router, general_settings=settings, @@ -11434,7 +11454,8 @@ async def model_info( fallback_type=None, llm_router=llm_router, ) - return {**response, "id": internal_to_public.get(resolved_model_id, model_id)} # mutable-ok: response id differs + response_id: Final = model_id if aliased_model_id else internal_to_public.get(resolved_model_id, model_id) + return {**response, "id": response_id} # mutable-ok: response id differs def _blocked_response_usage(original_response: object | None) -> "litellm.Usage": diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index f695af3cb11..3bd98ff5b0b 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -17,6 +17,7 @@ - {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} - {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"} - {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"} +- {id: other.auth.jwt.team_model_alias_listed_and_routes, module: other, tier: P1, area: auth, assertions: [team_model_alias_listed_and_routes], source: "proxy_server.py model_list / common_utils/model_listing_utils.py alias_listing_entries / LIT-8515", fail_before_fix: proven, rationale: "A team model_aliases name the JWT caller can complete on is also listed by GET /v1/models for that caller, in the OpenAI and the Anthropic (Claude Code) shapes, next to its target, so a managed client can discover the alias it is meant to send"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index bd7f5171172..c96f4b0bef1 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1426,6 +1426,7 @@ class TeamNewBody(BaseModel): team_id: str | None = None organization_id: str | None = None metadata: TeamMetadata | None = None + model_aliases: dict[str, str] | None = None class TeamNewResponse(BaseModel): diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 057a735e5e5..93c198586f6 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -14,12 +14,15 @@ own endpoints, so no test ever holds a signing key. from __future__ import annotations from dataclasses import dataclass +from typing import Final -from e2e_http import AuthHeaders, NoBody, ProbeResult, Result +from e2e_http import AnthropicHeaders, AuthHeaders, NoBody, ProbeResult, Result from idp import Keycloak, keycloak_from_env from models import ( ChatBody, ChatResponse, + ModelsListParams, + ModelsListResponse, ReadinessDetailsResponse, ReadinessResponse, UserListParams, @@ -88,6 +91,17 @@ class OtherClient: response_type=ChatResponse, ) + def list_models_as(self, token: str, *, anthropic: bool = False) -> Result[ModelsListResponse]: + """GET /v1/models under `token`, in the OpenAI shape or, with `anthropic`, the + Anthropic Models API shape Claude Code reads. Both carry `data[].id`.""" + bearer: Final = self.proxy.transport.bearer(token) + return self.proxy.transport.get( + "/v1/models", + headers=AnthropicHeaders(authorization=bearer.authorization) if anthropic else bearer, + params=ModelsListParams(return_wildcard_routes=False), + response_type=ModelsListResponse, + ) + def list_users_as(self, key: str) -> Result[UserListResponse]: """GET /user/list under `key`. Admin-only, so it doubles as the master key's authorization proof: the master key (proxy admin) reads it, a diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index de06c376f76..a8aa474f04f 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -78,9 +78,35 @@ def bound_team(client: OtherClient, resources: ResourceManager) -> BoundTeam: return BoundTeam(identity=provisioned, team_id=provisioned.group, team_alias=team_alias) -def _ping() -> ChatBody: +@dataclass(frozen=True, slots=True) +class AliasedTeam: + identity: Identity + alias: str + target: str + + +@pytest.fixture +def aliased_team(client: OtherClient, resources: ResourceManager) -> AliasedTeam: + """An identity whose team carries a model_aliases entry, the name a managed + client such as Claude Code sends and the team rewrites to a real model group.""" + marker: Final = unique_marker() + provisioned: Final = _provision(client, resources, marker=marker) + alias: Final = f"e2e-jwt-model-alias-{marker}" + team_id: Final = client.proxy.create_team( + TeamNewBody( + team_alias=f"e2e-jwt-aliased-{marker}", + team_id=provisioned.group, + models=[CHEAP_OPENAI_MODEL], + model_aliases={alias: CHEAP_OPENAI_MODEL}, + ) + ) + resources.defer(lambda: client.proxy.delete_team(team_id)) + return AliasedTeam(identity=provisioned, alias=alias, target=CHEAP_OPENAI_MODEL) + + +def _ping(model: str = CHEAP_OPENAI_MODEL) -> ChatBody: return ChatBody( - model=CHEAP_OPENAI_MODEL, + model=model, messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], max_tokens=16, ) @@ -221,6 +247,23 @@ class TestJwtTeamHeader: f"{bound_team.team_id!r}, got {by_alias!r}" ) + @pytest.mark.covers("other.auth.jwt.team_model_alias_listed_and_routes") + @pytest.mark.parametrize("anthropic", [False, True], ids=["openai_shape", "anthropic_shape"]) + def test_team_model_alias_is_listed_by_v1_models_under_the_same_token_that_routes_it( + self, client: OtherClient, aliased_team: AliasedTeam, anthropic: bool + ) -> None: + token: Final = client.idp.access_token(aliased_team.identity) + + routed: Final = unwrap(client.proxy.chat(token, _ping(model=aliased_team.alias))) + assert routed.choices, f"precondition: /chat/completions must route the team alias, got {routed}" + + listed: Final = tuple(entry.id for entry in unwrap(client.list_models_as(token, anthropic=anthropic)).data) + assert aliased_team.alias in listed, ( + f"/v1/models must list team alias {aliased_team.alias!r} that the same token routes on " + f"/chat/completions, got {listed}" + ) + assert aliased_team.target in listed, f"the alias target {aliased_team.target!r} must stay listed, got {listed}" + @pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied") def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value( self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam diff --git a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py index 7ef03140093..97e38d5e17a 100644 --- a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py @@ -4,9 +4,14 @@ from itertools import combinations import pytest +import litellm from litellm import Router from litellm.proxy.common_utils.model_listing_utils import ( + CallerAliases, ClaudeCodeRoutingNames, + alias_listing_entries, + alias_target, + caller_alias_maps, claude_code_group_name, claude_code_model_id, claude_code_requested_group, @@ -22,6 +27,10 @@ def _marked(name): return f"{_encoded(name)}[1m]" +def _caller(*maps: object) -> CallerAliases: + return CallerAliases(maps, maps) + + def _row(name, limit=1000000): return {"id": name, "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": limit} @@ -29,15 +38,16 @@ def _row(name, limit=1000000): def _router(*names, aliases=None): return Router( model_list=[ - {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} - for name in names + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names ], model_group_alias=aliases, ) @pytest.mark.parametrize("limit", [None, 999999, 1000000]) -@pytest.mark.parametrize("name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"]) +@pytest.mark.parametrize( + "name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"] +) def test_listing_round_trips_entire_source_name(name, limit): names = frozenset({name}) view = claude_code_model_id(name, limit, names) @@ -48,7 +58,15 @@ def test_listing_round_trips_entire_source_name(name, limit): def test_collision_matrix_round_trips_without_duplicate_ids(): - universe = ("foo", "foo[1m]", "claude-router-foo", _encoded("foo"), _encoded("foo") + "[1m]", "claude-opus-5", "claude-opus-5[1m]") + universe = ( + "foo", + "foo[1m]", + "claude-router-foo", + _encoded("foo"), + _encoded("foo") + "[1m]", + "claude-opus-5", + "claude-opus-5[1m]", + ) for pair in combinations(universe, 2): for visible in (pair, pair[:1], pair[1:]): names = frozenset(pair) @@ -57,18 +75,31 @@ def test_collision_matrix_round_trips_without_duplicate_ids(): assert all((claude_code_group_name(shown, names) or shown) == source for source, shown in view.items()) -@pytest.mark.parametrize("spelling", ["claude-router-foo", "claude-router-ff", "claude-router-66 6f6f", "claude-router-666F6F", "claude-router-", _encoded("missing")]) +@pytest.mark.parametrize( + "spelling", + [ + "claude-router-foo", + "claude-router-ff", + "claude-router-66 6f6f", + "claude-router-666F6F", + "claude-router-", + _encoded("missing"), + ], +) def test_unknown_or_noncanonical_ids_are_never_guessed(spelling): assert claude_code_group_name(spelling, frozenset({"foo"})) is None -@pytest.mark.parametrize("headers,enabled", [ - ({"user-agent": "claude-code/2.1.267"}, True), - ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), - ({"x-gateway-client": "Claude-Code"}, True), - ({"user-agent": "anthropic-sdk-python/0.40"}, False), - ({}, False), -]) +@pytest.mark.parametrize( + "headers,enabled", + [ + ({"user-agent": "claude-code/2.1.267"}, True), + ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), + ({"x-gateway-client": "Claude-Code"}, True), + ({"user-agent": "anthropic-sdk-python/0.40"}, False), + ({}, False), + ], +) def test_only_claude_code_gets_the_view(headers, enabled): rows = (_row("foo"), _row("claude-opus-5")) view = claude_code_view_ids(rows, headers, frozenset(row["id"] for row in rows)) @@ -77,12 +108,15 @@ def test_only_claude_code_gets_the_view(headers, enabled): @pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "team", "wildcard"]) def test_configured_names_outrank_generated_ids_even_when_hidden_from_listing(monkeypatch, layer): - import litellm - encoded = _encoded("foo") alias = {encoded: "other"} monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) - router = _router("foo", "other", *( (encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), aliases=alias if layer == "router" else None) + router = _router( + "foo", + "other", + *((encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), + aliases=alias if layer == "router" else None, + ) maps = (alias,) if layer in ("key", "team") else () names = ClaudeCodeRoutingNames(router, None, maps) assert claude_code_requested_group(encoded, router, None, maps) is None @@ -98,12 +132,113 @@ def test_mutation_breaking_the_hex_name_cannot_route_to_the_source(source): assert claude_code_requested_group(_marked(source), router, None) == source +def test_team_alias_is_listed_under_its_target_metadata_and_only_when_the_target_is_accessible() -> None: + entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("team-public", "model_name_team_1_abc")] + aliases = ( + {"gpt-4.1-mini": "team-public"}, + None, + {"claude-sonnet-4-5": "gpt-4.1-mini", "via-public": "team-public", "not-granted": "gpt-4.1"}, + ) + assert alias_listing_entries(entries, _caller(*aliases)) == ( + *entries, + ("claude-sonnet-4-5", "gpt-4.1-mini"), + ("via-public", "model_name_team_1_abc"), + ) + assert alias_listing_entries(entries, _caller(None, {})) == tuple(entries) + + +def test_alias_target_resolves_the_requested_alias_across_key_and_team_maps() -> None: + maps = _caller({"o": "gpt-4.1"}, {"claude-sonnet-4-5": "gpt-4.1-mini"}) + assert alias_target("claude-sonnet-4-5", maps) == "gpt-4.1-mini" + assert alias_target("gpt-4.1-mini", _caller(None, {"claude-sonnet-4-5": "gpt-4.1-mini"})) is None + + +def test_alias_colliding_with_a_listed_id_keeps_the_listed_model_at_list_and_retrieval() -> None: + entries = [("fast", "fast"), ("gpt-4.1-mini", "gpt-4.1-mini")] + maps = _caller({"fast": "gpt-4.1-mini"}) + listed = frozenset(response_id for response_id, _ in entries) + + assert alias_listing_entries(entries, maps) == tuple(entries) + assert alias_target("fast", maps, listed) is None + assert alias_target("fast", maps) == "gpt-4.1-mini" + + +def test_alias_maps_apply_in_the_order_chat_completions_applies_them() -> None: + team_then_key = _caller({"fast": "gpt-4.1-mini", "hop": "mid"}, {"fast": "gpt-4.1", "mid": "gpt-4.1"}) + entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("gpt-4.1", "gpt-4.1")] + + assert alias_target("fast", team_then_key) == "gpt-4.1-mini" + assert alias_target("hop", team_then_key) == "gpt-4.1" + assert alias_listing_entries(entries, team_then_key) == ( + *entries, + ("fast", "gpt-4.1-mini"), + ("hop", "gpt-4.1"), + ("mid", "gpt-4.1"), + ) + + +def test_one_bad_alias_entry_hides_only_itself() -> None: + aliases = {"fast": "gpt-4.1-mini", "broken": 5, 7: "gpt-4.1-mini"} + entries = [("gpt-4.1-mini", "gpt-4.1-mini")] + + assert alias_listing_entries(entries, _caller(aliases)) == (*entries, ("fast", "gpt-4.1-mini")) + assert alias_target("fast", _caller(aliases)) == "gpt-4.1-mini" + + +def test_chained_key_alias_is_listed_only_when_its_final_target_is_listable() -> None: + key_aliases = {"a": "b", "b": "hidden"} + entries = [("b", "b")] + + assert alias_listing_entries(entries, caller_alias_maps(key_aliases, None, "team-a", None)) == (*entries,) + assert alias_target("a", caller_alias_maps(key_aliases, None, "team-a", None)) == "hidden" + + +def test_team_aliases_only_apply_when_listing_the_team_the_key_authenticated_as( + monkeypatch: pytest.MonkeyPatch, +) -> None: + key_aliases, team_aliases, global_aliases = {"k": "gpt-4.1"}, {"t": "gpt-4.1-mini"}, {"g": "gpt-4.1"} + monkeypatch.setattr(litellm, "model_alias_map", global_aliases) + own_team = CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, global_aliases, key_aliases)) + assert caller_alias_maps(key_aliases, team_aliases, "team-a", None) == own_team + assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-a") == own_team + assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-b") == CallerAliases( + (key_aliases,), (key_aliases, global_aliases, key_aliases) + ) + + +def test_global_alias_rewrites_between_the_two_key_passes_like_chat_completions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"b": "d"}) + key_aliases = {"a": "b", "b": "c"} + entries = [("c", "c"), ("d", "d")] + maps = caller_alias_maps(key_aliases, None, "team-a", None) + + assert alias_target("a", maps) == "d" + assert alias_listing_entries(entries, maps) == (*entries, ("a", "d"), ("b", "c")) + + +def test_global_aliases_rewrite_but_are_not_listed_as_caller_rows(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"g": "gpt-4.1-mini"}) + entries = [("gpt-4.1-mini", "gpt-4.1-mini")] + maps = caller_alias_maps({"k": "g"}, None, "team-a", None) + + assert alias_listing_entries(entries, maps) == (*entries, ("k", "gpt-4.1-mini")) + assert alias_target("g", maps) == "gpt-4.1-mini" + + def test_team_public_name_uses_the_same_scope_at_list_and_request(): - router = Router(model_list=[{ - "model_name": "model_name_team-a_id", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, - "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, - }]) - shown = claude_code_view_ids((_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a"))["shared"] + router = Router( + model_list=[ + { + "model_name": "model_name_team-a_id", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, + } + ] + ) + shown = claude_code_view_ids( + (_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a") + )["shared"] assert claude_code_requested_group(shown, router, "team-a") == "shared" assert claude_code_requested_group(shown, router, "team-b") is None diff --git a/tests/test_litellm/proxy/test_model_list_aliases.py b/tests/test_litellm/proxy/test_model_list_aliases.py new file mode 100644 index 00000000000..25a941bf2ae --- /dev/null +++ b/tests/test_litellm/proxy/test_model_list_aliases.py @@ -0,0 +1,151 @@ +""" +Tests for key and team `model_aliases` on the model listing endpoints: GET /v1/models +(`model_list`, OpenAI and Anthropic shapes) and GET /v1/models/{id} (`model_info`). +An alias the caller can complete on is listed next to its target and resolves by name. +""" + +import pytest +from starlette.requests import Request + +from litellm import Router +from litellm.proxy import proxy_server +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _deployment(model_name: str, model: str = "openai/gpt-4.1-mini", **model_info: str | bool) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": f"{model_name}-id", **model_info}, + } + + +@pytest.fixture +def router(monkeypatch: pytest.MonkeyPatch) -> Router: + router = Router( + model_list=[ + _deployment("gpt-4.1-mini"), + _deployment("gpt-4.1", model="openai/gpt-4.1"), + _deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-chat"), + _deployment("hidden", model="anthropic/claude-sonnet-4-5", discoverable=False), + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.model_list) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + return router + + +def _team_member( + team_id: str = "team1", models: list[str] | None = None, **aliases: dict[str, str] | None +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=team_id, + team_models=["gpt-4.1-mini", "model_name_team1_abc"], + models=models or ["gpt-4.1-mini", "model_name_team1_abc"], + **aliases, + ) + + +def _anthropic_request(*extra_headers: tuple[bytes, bytes]) -> Request: + return Request( + scope={ + "type": "http", + "method": "GET", + "path": "/v1/models", + "query_string": b"", + "headers": [(b"anthropic-version", b"2023-06-01"), *extra_headers], + } + ) + + +def _claude_code_request() -> Request: + return _anthropic_request((b"user-agent", b"claude-cli/2.1.267 (external, cli)")) + + +async def _v1_models(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> list[str]: + response = await proxy_server.model_list(user_api_key_dict=user_api_key_dict, request=request) + return [m["id"] for m in response["data"]] + + +@pytest.mark.asyncio +async def test_v1_models_lists_team_alias_next_to_its_target_in_both_shapes(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"] + assert await _v1_models(caller, request=_anthropic_request()) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"] + + +@pytest.mark.asyncio +async def test_claude_code_picker_lists_the_alias_under_its_own_name(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + picker_ids = await _v1_models(caller, request=_claude_code_request()) + assert any(picker_id.startswith("claude-sonnet-4-5") for picker_id in picker_ids), picker_ids + + +@pytest.mark.asyncio +async def test_v1_models_lists_key_alias_and_hides_alias_to_a_model_the_caller_cannot_list(router: Router) -> None: + caller = _team_member(aliases={"mini": "gpt-4.1-mini", "big": "gpt-4.1"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "mini"] + + +@pytest.mark.asyncio +async def test_v1_models_resolves_a_team_alias_through_the_key_alias_like_chat_completions_does(router: Router) -> None: + caller = _team_member(team_model_aliases={"fast": "mid"}, aliases={"fast": "gpt-4.1", "mid": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast", "mid"] + response = await proxy_server.model_info(model_id="fast", user_api_key_dict=caller) + assert response["id"] == "fast" + + +@pytest.mark.asyncio +async def test_v1_models_skips_only_the_malformed_alias_entries(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": 5, "fast": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast"] + + +@pytest.mark.asyncio +async def test_v1_models_by_id_resolves_a_team_alias_to_its_target_metadata(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + response = await proxy_server.model_info(model_id="claude-sonnet-4-5", user_api_key_dict=caller) + assert response["id"] == "claude-sonnet-4-5" + assert response["owned_by"] == "openai" + + +@pytest.mark.asyncio +async def test_v1_models_by_id_retrieves_the_listed_model_when_an_alias_collides_with_its_id(router: Router) -> None: + caller = _team_member(aliases={"team-chat": "gpt-4.1"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat"] + response = await proxy_server.model_info(model_id="team-chat", user_api_key_dict=caller) + assert response["id"] == "team-chat" + + +@pytest.mark.asyncio +async def test_v1_models_by_id_resolves_an_alias_named_like_an_undiscoverable_model_to_the_alias_target( + router: Router, +) -> None: + caller = _team_member(aliases={"hidden": "gpt-4.1-mini"}, models=["gpt-4.1-mini", "model_name_team1_abc", "hidden"]) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "hidden"] + target = await proxy_server.model_info(model_id="gpt-4.1-mini", user_api_key_dict=caller) + response = await proxy_server.model_info(model_id="hidden", user_api_key_dict=caller) + assert response == {**target, "id": "hidden"} + + +@pytest.mark.asyncio +async def test_v1_models_by_id_keeps_the_alias_as_id_when_it_targets_a_team_scoped_model(router: Router) -> None: + caller = _team_member(team_model_aliases={"chat": "team-chat"}) + + assert "chat" in await _v1_models(caller) + response = await proxy_server.model_info(model_id="chat", user_api_key_dict=caller) + assert response["id"] == "chat"