From 12005c4a02a4c022429c9219c722eac352126724 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 6 Mar 2026 18:20:10 -0800 Subject: [PATCH] Fix /key/aliases auth for internal users and scope results by role Internal users were blocked from accessing /key/aliases because the route was missing from key_management_routes. Added the route and scoped query results so non-admin users only see aliases for their own keys and their teams' keys, matching /key/list behavior. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 2 + .../key_management_endpoints.py | 39 +++++++ .../test_key_management_endpoints.py | 100 +++++++++++++++++- 3 files changed, 137 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 95dabd8bfd..b64192c47d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -238,6 +238,7 @@ class KeyManagementRoutes(str, enum.Enum): # list routes KEY_LIST = "/key/list" + KEY_ALIASES = "/key/aliases" # team usage routes TEAM_DAILY_ACTIVITY = "/team/daily/activity" @@ -514,6 +515,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.KEY_BULK_UPDATE.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value, KeyManagementRoutes.KEY_RESET_SPEND.value, + KeyManagementRoutes.KEY_ALIASES.value, ] management_routes = [ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5d189e0db4..d3090a55a5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4136,6 +4136,7 @@ async def list_keys( ) @management_endpoint_wrapper async def key_aliases( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), page: int = Query(1, ge=1, description="Page number"), size: int = Query(50, ge=1, le=100, description="Page size"), search: Optional[str] = Query( @@ -4145,6 +4146,9 @@ async def key_aliases( """ Lists key aliases with pagination and optional search. + Non-admin users only see aliases for keys they own or keys belonging to + their teams. + Returns: { "aliases": List[str], @@ -4174,6 +4178,41 @@ async def key_aliases( "key_alias != ''", "(team_id IS NULL OR team_id != $1)", ] + + # Scope results for non-admin users: only show aliases for keys the + # user owns or keys belonging to teams they are a member of. + is_proxy_admin = user_api_key_dict.user_role in [ + LitellmUserRoles.PROXY_ADMIN.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ] + if not is_proxy_admin: + scope_conditions: List[str] = [] + if user_api_key_dict.user_id: + query_params.append(user_api_key_dict.user_id) + scope_conditions.append(f"user_id = ${len(query_params)}") + + # Look up the user's teams from the user table + user_teams: List[str] = [] + if user_api_key_dict.user_id: + user_row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id} + ) + if user_row is not None: + user_teams = getattr(user_row, "teams", []) or [] + + if user_teams: + team_placeholders = ", ".join( + f"${len(query_params) + i + 1}" for i in range(len(user_teams)) + ) + query_params.extend(user_teams) + scope_conditions.append(f"team_id IN ({team_placeholders})") + + if scope_conditions: + where_parts.append(f"({' OR '.join(scope_conditions)})") + else: + # No user_id and no teams — return nothing + where_parts.append("FALSE") + if search: query_params.append(f"%{search}%") where_parts.append(f"key_alias ILIKE ${len(query_params)}") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index bc5c692510..4613b1e277 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6356,6 +6356,13 @@ async def test_generate_key_helper_fn_agent_id(): ) +def _make_admin_key_dict() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_id="admin-user", + ) + + @pytest.mark.asyncio async def test_key_aliases_response_shape(): """Test that key_aliases returns the correct paginated response shape.""" @@ -6368,7 +6375,10 @@ async def test_key_aliases_response_shape(): ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): - result = await key_aliases(page=1, size=50, search=None) + result = await key_aliases( + user_api_key_dict=_make_admin_key_dict(), + page=1, size=50, search=None, + ) assert result["aliases"] == ["alias-alpha", "alias-beta"] assert result["total_count"] == 2 @@ -6395,7 +6405,10 @@ async def test_key_aliases_pagination_skip_take(): ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): - result = await key_aliases(page=3, size=25, search=None) + result = await key_aliases( + user_api_key_dict=_make_admin_key_dict(), + page=3, size=25, search=None, + ) assert result["current_page"] == 3 assert result["size"] == 25 @@ -6420,7 +6433,10 @@ async def test_key_aliases_search_filter(): ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): - await key_aliases(page=1, size=50, search="my-key") + await key_aliases( + user_api_key_dict=_make_admin_key_dict(), + page=1, size=50, search="my-key", + ) count_call = mock_prisma_client.db.query_raw.call_args_list[0] count_sql = count_call.args[0] @@ -6442,12 +6458,88 @@ async def test_key_aliases_no_search_omits_ilike_filter(): ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): - await key_aliases(page=1, size=50, search=None) + await key_aliases( + user_api_key_dict=_make_admin_key_dict(), + page=1, size=50, search=None, + ) count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] assert "ILIKE" not in count_sql +@pytest.mark.asyncio +async def test_key_aliases_internal_user_scoped_to_own_keys_and_teams(): + """Test that internal users only see aliases for their own keys and team keys.""" + mock_prisma_client = AsyncMock() + + # Mock user table lookup to return teams + mock_user_row = MagicMock() + mock_user_row.teams = ["team-abc", "team-xyz"] + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mock_user_row + ) + + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 1}], + [{"key_alias": "my-alias"}], + ] + ) + + internal_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + user_id="user-123", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases( + user_api_key_dict=internal_user, + page=1, size=50, search=None, + ) + + assert result["aliases"] == ["my-alias"] + + # Verify the SQL includes user_id and team_id scoping + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + assert "user_id = " in count_sql + assert "team_id IN " in count_sql + + # Verify the params include user_id and team IDs + count_params = mock_prisma_client.db.query_raw.call_args_list[0].args[1:] + assert "user-123" in count_params + assert "team-abc" in count_params + assert "team-xyz" in count_params + + +@pytest.mark.asyncio +async def test_key_aliases_admin_sees_all(): + """Test that proxy admins see all aliases without user/team scoping.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 3}], + [ + {"key_alias": "alias-a"}, + {"key_alias": "alias-b"}, + {"key_alias": "alias-c"}, + ], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases( + user_api_key_dict=_make_admin_key_dict(), + page=1, size=50, search=None, + ) + + assert len(result["aliases"]) == 3 + + # Admin queries should NOT have user_id or team_id scoping + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + assert "user_id = " not in count_sql + assert "team_id IN " not in count_sql + + class TestValidateKeyAliasFormat: def test_validate_key_alias_format_valid(self):