diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index da7e1f5a04..e6b7b7d285 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -239,6 +239,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" @@ -515,6 +516,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 dd7cc0ec9c..072d4181c6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4140,6 +4140,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( @@ -4149,6 +4150,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], @@ -4178,6 +4182,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 c09840561a..2533f0d66d 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 @@ -6486,6 +6486,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.""" @@ -6498,7 +6505,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 @@ -6525,7 +6535,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 @@ -6550,7 +6563,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] @@ -6572,12 +6588,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):