diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3c1053c7b0..c9578555ba 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3150,12 +3150,14 @@ async def list_keys( ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), + status: Optional[str] = Query(None, description="Filter by status (e.g. 'deleted')"), ) -> KeyListResponseObject: """ List all keys for a given user / team / organization. Parameters: expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) + status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys. Returns: { @@ -3177,6 +3179,15 @@ async def list_keys( verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") + # Validate status parameter + if status is not None and status != "deleted": + raise HTTPException( + status_code=400, + detail={ + "error": f"Invalid status value. Currently only 'deleted' is supported." + }, + ) + complete_user_info = await validate_key_list_check( user_api_key_dict=user_api_key_dict, user_id=user_id, @@ -3217,6 +3228,7 @@ async def list_keys( sort_by=sort_by, sort_order=sort_order, expand=expand, + status=status, ) verbose_proxy_logger.debug("Successfully prepared response") @@ -3424,6 +3436,7 @@ async def _list_key_helper( sort_by: Optional[str] = None, sort_order: str = "desc", expand: Optional[List[str]] = None, + status: Optional[str] = None, ) -> KeyListResponseObject: """ Helper function to list keys @@ -3468,28 +3481,51 @@ async def _list_key_helper( else None ) + # Determine which table to query based on status + use_deleted_table = status == "deleted" + # Fetch keys with pagination - keys = await prisma_client.db.litellm_verificationtoken.find_many( - where=where, # type: ignore - skip=skip, # type: ignore - take=size, # type: ignore - order=( - order_by - if order_by - else [ - {"created_at": "desc"}, - {"token": "desc"}, # fallback sort - ] - ), - include={"object_permission": True}, - ) + if use_deleted_table: + keys = await prisma_client.db.litellm_deletedverificationtoken.find_many( + where=where, # type: ignore + skip=skip, # type: ignore + take=size, # type: ignore + order=( + order_by + if order_by + else [ + {"created_at": "desc"}, + {"token": "desc"}, # fallback sort + ] + ), + ) + else: + keys = await prisma_client.db.litellm_verificationtoken.find_many( + where=where, # type: ignore + skip=skip, # type: ignore + take=size, # type: ignore + order=( + order_by + if order_by + else [ + {"created_at": "desc"}, + {"token": "desc"}, # fallback sort + ] + ), + include={"object_permission": True}, + ) verbose_proxy_logger.debug(f"Fetched {len(keys)} keys") # Get total count of keys - total_count = await prisma_client.db.litellm_verificationtoken.count( - where=where # type: ignore - ) + if use_deleted_table: + total_count = await prisma_client.db.litellm_deletedverificationtoken.count( + where=where # type: ignore + ) + else: + total_count = await prisma_client.db.litellm_verificationtoken.count( + where=where # type: ignore + ) verbose_proxy_logger.debug(f"Total count of keys: {total_count}") @@ -3510,8 +3546,9 @@ async def _list_key_helper( key_list: List[Union[str, UserAPIKeyAuth]] = [] for key in keys: key_dict = key.dict() - # Attach object_permission if object_permission_id is set - key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) + # Attach object_permission if object_permission_id is set (only for non-deleted keys) + if not use_deleted_table: + key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) # Include user information if expand includes "user" if expand and "user" in expand and key.user_id and key.user_id in user_map: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c606420cc0..381e057da1 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2991,6 +2991,9 @@ async def list_team_v2( sort_order: str = fastapi.Query( default="asc", description="Sort order ('asc' or 'desc')" ), + status: Optional[str] = fastapi.Query( + default=None, description="Filter by status (e.g. 'deleted')" + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -3013,6 +3016,8 @@ async def list_team_v2( Column to sort by (e.g. 'team_id', 'team_alias', 'created_at') sort_order: str Sort order ('asc' or 'desc') + status: Optional[str] + Filter by status. Currently supports "deleted" to query deleted teams. """ from litellm.proxy.proxy_server import prisma_client @@ -3037,6 +3042,16 @@ async def list_team_v2( if user_id is None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: user_id = user_api_key_dict.user_id + if status is not None and status != "deleted": + raise HTTPException( + status_code=400, + detail={ + "error": f"Invalid status value. Currently only 'deleted' is supported." + }, + ) + + use_deleted_table = status == "deleted" + # Calculate skip and take for pagination skip = (page - 1) * page_size @@ -3071,16 +3086,19 @@ async def list_team_v2( detail={"error": f"User not found, passed user_id={user_id}"}, ) user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) - # Find teams where this user is a member by checking members_with_roles array - if team_id is None: - where_conditions["team_id"] = {"in": user_object_correct_type.teams} - elif team_id in user_object_correct_type.teams: - where_conditions["team_id"] = team_id + + if use_deleted_table: + where_conditions["members"] = {"has": user_id} else: - raise HTTPException( - status_code=404, - detail={"error": f"User is not a member of team_id={team_id}"}, - ) + if team_id is None: + where_conditions["team_id"] = {"in": user_object_correct_type.teams} + elif team_id in user_object_correct_type.teams: + where_conditions["team_id"] = team_id + else: + raise HTTPException( + status_code=404, + detail={"error": f"User is not a member of team_id={team_id}"}, + ) # Build order_by conditions valid_sort_columns = ["team_id", "team_alias", "created_at"] @@ -3091,14 +3109,26 @@ async def list_team_v2( order_by = {sort_by: sort_order.lower()} # Get teams with pagination - teams = await prisma_client.db.litellm_teamtable.find_many( - where=where_conditions, - skip=skip, - take=page_size, - order=order_by if order_by else {"created_at": "desc"}, # Default sort - ) - # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count(where=where_conditions) + if use_deleted_table: + teams = await prisma_client.db.litellm_deletedteamtable.find_many( + where=where_conditions, + skip=skip, + take=page_size, + order=order_by if order_by else {"created_at": "desc"}, # Default sort + ) + # Get total count for pagination + total_count = await prisma_client.db.litellm_deletedteamtable.count( + where=where_conditions + ) + else: + teams = await prisma_client.db.litellm_teamtable.find_many( + where=where_conditions, + skip=skip, + take=page_size, + order=order_by if order_by else {"created_at": "desc"}, # Default sort + ) + # Get total count for pagination + total_count = await prisma_client.db.litellm_teamtable.count(where=where_conditions) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a1e8efdbb4..771a0707b7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2068,6 +2068,7 @@ async def test_list_team_v2_security_check_non_admin_user(): http_request=mock_request, user_id=None, # Non-admin trying to query all teams user_api_key_dict=mock_user_api_key_dict_non_admin, + status=None, ) assert exc_info.value.status_code == 401 @@ -2108,6 +2109,7 @@ async def test_list_team_v2_security_check_non_admin_user_other_user(): http_request=mock_request, user_id="other_user_456", # Non-admin trying to query other user's teams user_api_key_dict=mock_user_api_key_dict_non_admin, + status=None, ) assert exc_info.value.status_code == 401 @@ -2166,6 +2168,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): team_id=None, page=1, page_size=10, + status=None, ) # Should return results without error @@ -2215,6 +2218,7 @@ async def test_list_team_v2_security_check_admin_user(): user_api_key_dict=mock_user_api_key_dict_admin, page=1, page_size=10, + status=None, ) # Should return results without error