diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15-py3-none-any.whl deleted file mode 100644 index ba2e5e5fce..0000000000 Binary files a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15-py3-none-any.whl and /dev/null differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15.tar.gz deleted file mode 100644 index 7d01b3de6f..0000000000 Binary files a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15.tar.gz and /dev/null differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22-py3-none-any.whl deleted file mode 100644 index 1e2f6967dc..0000000000 Binary files a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22-py3-none-any.whl and /dev/null differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22.tar.gz deleted file mode 100644 index 1864c77ddd..0000000000 Binary files a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22.tar.gz and /dev/null differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql deleted file mode 100644 index 6ca66ddaad..0000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql +++ /dev/null @@ -1,117 +0,0 @@ --- CreateTable -CREATE TABLE "LiteLLM_DeletedTeamTable" ( - "id" TEXT NOT NULL, - "team_id" TEXT NOT NULL, - "team_alias" TEXT, - "organization_id" TEXT, - "object_permission_id" TEXT, - "admins" TEXT[], - "members" TEXT[], - "members_with_roles" JSONB NOT NULL DEFAULT '{}', - "metadata" JSONB NOT NULL DEFAULT '{}', - "max_budget" DOUBLE PRECISION, - "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, - "models" TEXT[], - "max_parallel_requests" INTEGER, - "tpm_limit" BIGINT, - "rpm_limit" BIGINT, - "budget_duration" TEXT, - "budget_reset_at" TIMESTAMP(3), - "blocked" BOOLEAN NOT NULL DEFAULT false, - "model_spend" JSONB NOT NULL DEFAULT '{}', - "model_max_budget" JSONB NOT NULL DEFAULT '{}', - "team_member_permissions" TEXT[] DEFAULT ARRAY[]::TEXT[], - "model_id" INTEGER, - "created_at" TIMESTAMP(3), - "updated_at" TIMESTAMP(3), - "deleted_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, - "deleted_by" TEXT, - "deleted_by_api_key" TEXT, - "litellm_changed_by" TEXT, - - CONSTRAINT "LiteLLM_DeletedTeamTable_pkey" PRIMARY KEY ("id") -); - --- CreateTable -CREATE TABLE "LiteLLM_DeletedVerificationToken" ( - "id" TEXT NOT NULL, - "token" TEXT NOT NULL, - "key_name" TEXT, - "key_alias" TEXT, - "soft_budget_cooldown" BOOLEAN NOT NULL DEFAULT false, - "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, - "expires" TIMESTAMP(3), - "models" TEXT[], - "aliases" JSONB NOT NULL DEFAULT '{}', - "config" JSONB NOT NULL DEFAULT '{}', - "user_id" TEXT, - "team_id" TEXT, - "permissions" JSONB NOT NULL DEFAULT '{}', - "max_parallel_requests" INTEGER, - "metadata" JSONB NOT NULL DEFAULT '{}', - "blocked" BOOLEAN, - "tpm_limit" BIGINT, - "rpm_limit" BIGINT, - "max_budget" DOUBLE PRECISION, - "budget_duration" TEXT, - "budget_reset_at" TIMESTAMP(3), - "allowed_cache_controls" TEXT[] DEFAULT ARRAY[]::TEXT[], - "allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[], - "model_spend" JSONB NOT NULL DEFAULT '{}', - "model_max_budget" JSONB NOT NULL DEFAULT '{}', - "budget_id" TEXT, - "organization_id" TEXT, - "object_permission_id" TEXT, - "created_at" TIMESTAMP(3), - "created_by" TEXT, - "updated_at" TIMESTAMP(3), - "updated_by" TEXT, - "rotation_count" INTEGER DEFAULT 0, - "auto_rotate" BOOLEAN DEFAULT false, - "rotation_interval" TEXT, - "last_rotation_at" TIMESTAMP(3), - "key_rotation_at" TIMESTAMP(3), - "deleted_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, - "deleted_by" TEXT, - "deleted_by_api_key" TEXT, - "litellm_changed_by" TEXT, - - CONSTRAINT "LiteLLM_DeletedVerificationToken_pkey" PRIMARY KEY ("id") -); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_team_id_idx" ON "LiteLLM_DeletedTeamTable"("team_id"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_deleted_at_idx" ON "LiteLLM_DeletedTeamTable"("deleted_at"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_organization_id_idx" ON "LiteLLM_DeletedTeamTable"("organization_id"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_team_alias_idx" ON "LiteLLM_DeletedTeamTable"("team_alias"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedTeamTable_created_at_idx" ON "LiteLLM_DeletedTeamTable"("created_at"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_token_idx" ON "LiteLLM_DeletedVerificationToken"("token"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_deleted_at_idx" ON "LiteLLM_DeletedVerificationToken"("deleted_at"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_user_id_idx" ON "LiteLLM_DeletedVerificationToken"("user_id"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_team_id_idx" ON "LiteLLM_DeletedVerificationToken"("team_id"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_organization_id_idx" ON "LiteLLM_DeletedVerificationToken"("organization_id"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_key_alias_idx" ON "LiteLLM_DeletedVerificationToken"("key_alias"); - --- CreateIndex -CREATE INDEX "LiteLLM_DeletedVerificationToken_created_at_idx" ON "LiteLLM_DeletedVerificationToken"("created_at"); - diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 71b398c59a..56fe093a8b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -132,49 +132,6 @@ model LiteLLM_TeamTable { object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } -// Audit table for deleted teams - preserves spend and team information for historical tracking -model LiteLLM_DeletedTeamTable { - id String @id @default(uuid()) - team_id String // Original team_id - team_alias String? - organization_id String? - object_permission_id String? - admins String[] - members String[] - members_with_roles Json @default("{}") - metadata Json @default("{}") - max_budget Float? - spend Float @default(0.0) - models String[] - max_parallel_requests Int? - tpm_limit BigInt? - rpm_limit BigInt? - budget_duration String? - budget_reset_at DateTime? - blocked Boolean @default(false) - model_spend Json @default("{}") - model_max_budget Json @default("{}") - router_settings Json? @default("{}") - team_member_permissions String[] @default([]) - model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases - - // Original timestamps from team creation/updates - created_at DateTime? @map("created_at") - updated_at DateTime? @map("updated_at") - - // Deletion metadata - deleted_at DateTime @default(now()) @map("deleted_at") - deleted_by String? @map("deleted_by") // User who deleted the team - deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion - litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided - - @@index([team_id]) - @@index([deleted_at]) - @@index([organization_id]) - @@index([team_alias]) - @@index([created_at]) -} - // Track spend, rate limit, budget Users model LiteLLM_UserTable { user_id String @id @@ -302,62 +259,6 @@ model LiteLLM_VerificationToken { object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } -// Audit table for deleted keys - preserves spend and key information for historical tracking -model LiteLLM_DeletedVerificationToken { - id String @id @default(uuid()) - token String // Original token (hashed) - key_name String? - key_alias String? - soft_budget_cooldown Boolean @default(false) - spend Float @default(0.0) - expires DateTime? - models String[] - aliases Json @default("{}") - config Json @default("{}") - user_id String? - team_id String? - permissions Json @default("{}") - max_parallel_requests Int? - metadata Json @default("{}") - blocked Boolean? - tpm_limit BigInt? - rpm_limit BigInt? - max_budget Float? - budget_duration String? - budget_reset_at DateTime? - allowed_cache_controls String[] @default([]) - allowed_routes String[] @default([]) - model_spend Json @default("{}") - model_max_budget Json @default("{}") - router_settings Json? @default("{}") - budget_id String? - organization_id String? - object_permission_id String? - created_at DateTime? // Original creation timestamp - created_by String? // Original creator - updated_at DateTime? // Last update timestamp before deletion - updated_by String? // Last user who updated before deletion - rotation_count Int? @default(0) - auto_rotate Boolean? @default(false) - rotation_interval String? - last_rotation_at DateTime? - key_rotation_at DateTime? - - // Deletion metadata - deleted_at DateTime @default(now()) @map("deleted_at") - deleted_by String? @map("deleted_by") // User who deleted the key - deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion - litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided - - @@index([token]) - @@index([deleted_at]) - @@index([user_id]) - @@index([team_id]) - @@index([organization_id]) - @@index([key_alias]) - @@index([created_at]) -} - model LiteLLM_EndUserTable { user_id String @id alias String? // admin-facing alias diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0f2cd4b3a7..3c6e210526 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1723,21 +1723,6 @@ class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): last_refreshed_at: Optional[float] = None -class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): - """ - Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - class TeamRequest(LiteLLMPydanticObjectBase): teams: List[str] @@ -2132,21 +2117,6 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) -class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): - """ - Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): """ Combined view of litellm verification token + litellm team table (select values) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 2672c41893..1850ffa256 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -412,19 +412,6 @@ async def new_user( status_code=403, detail="License is over limit. Please contact support@berri.ai to upgrade your license.", ) - - # Only proxy admins can create administrative users - # Check if user_api_key_dict is actually a UserAPIKeyAuth instance (not a Depends object) - # This can happen when the function is called directly in tests - if ( - data.user_role in [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] - and isinstance(user_api_key_dict, UserAPIKeyAuth) - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN - ): - raise HTTPException( - status_code=403, - detail=f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {data.user_role}. Your role: {user_api_key_dict.user_role}" - ) data_json = data.json() # type: ignore data_json = _update_internal_new_user_params(data_json, data) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3c1053c7b0..39b6774a61 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -16,7 +16,7 @@ import secrets import traceback import yaml from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Literal, Optional, Tuple, cast +from typing import List, Literal, Optional, Tuple, cast from litellm.litellm_core_utils.safe_json_dumps import safe_dumps import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status @@ -1791,10 +1791,6 @@ async def delete_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") - # Normalize litellm_changed_by: if it's a Header object or not a string, convert to None - if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): - litellm_changed_by = None - ## only allow user to delete keys they own verbose_proxy_logger.debug( f"user_api_key_dict.user_role: {user_api_key_dict.user_role}" @@ -1807,7 +1803,6 @@ async def delete_key_fn( tokens=data.keys, user_api_key_cache=user_api_key_cache, user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, ) num_keys_to_be_deleted = len(data.keys) deleted_keys = data.keys @@ -1817,7 +1812,6 @@ async def delete_key_fn( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, ) num_keys_to_be_deleted = len(data.key_aliases) deleted_keys = data.key_aliases @@ -2439,7 +2433,6 @@ async def delete_verification_tokens( tokens: List, user_api_key_cache: DualCache, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: """ Helper that deletes the list of tokens from the database @@ -2476,43 +2469,38 @@ async def delete_verification_tokens( detail={"error": "No keys found"}, ) - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: - authorized_keys = _keys_being_deleted - else: - authorized_keys = [] - for key in _keys_being_deleted: - if await can_modify_verification_token( - key_info=key, - user_api_key_cache=user_api_key_cache, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ): - authorized_keys.append(key) - else: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "You are not authorized to delete this key" - }, - ) - await _persist_deleted_verification_tokens( - keys=authorized_keys, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - + # Assuming 'db' is your Prisma Client instance + # check if admin making request - don't filter by user-id if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: deleted_tokens = await prisma_client.delete_data(tokens=tokens) + # else else: - deletion_tasks = [ - prisma_client.delete_data(tokens=[key.token]) - for key in authorized_keys - ] - await asyncio.gather(*deletion_tasks) + tasks = [] + deleted_tokens = [] + for key in _keys_being_deleted: - deleted_tokens = [key.token for key in authorized_keys] - if len(deleted_tokens) != len(tokens): + async def _delete_key(key: LiteLLM_VerificationToken): + if await can_modify_verification_token( + key_info=key, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ): + await prisma_client.delete_data(tokens=[key.token]) + deleted_tokens.append(key.token) + else: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "You are not authorized to delete this key" + }, + ) + + tasks.append(_delete_key(key)) + await asyncio.gather(*tasks) + + _num_deleted_tokens = len(deleted_tokens) + if _num_deleted_tokens != len(tokens): failed_tokens = [ token for token in tokens if token not in deleted_tokens ] @@ -2540,81 +2528,11 @@ async def delete_verification_tokens( return {"deleted_keys": deleted_tokens}, _keys_being_deleted -def _transform_verification_tokens_to_deleted_records( - keys: List[LiteLLM_VerificationToken], - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> List[Dict[str, Any]]: - """Transform verification tokens into deleted token records ready for persistence.""" - if not keys: - return [] - - deleted_at = datetime.now(timezone.utc) - records = [] - for key in keys: - key_payload = key.model_dump() - deleted_record = LiteLLM_DeletedVerificationToken( - **key_payload, - deleted_at=deleted_at, - deleted_by=user_api_key_dict.user_id, - deleted_by_api_key=user_api_key_dict.api_key, - litellm_changed_by=litellm_changed_by, - ) - record = deleted_record.model_dump() - - # Map org_id to organization_id (model uses org_id, but schema expects organization_id) - org_id_value = record.pop("org_id", None) - if org_id_value is not None: - record["organization_id"] = org_id_value - - for json_field in ["aliases", "config", "permissions", "metadata", "model_spend", "model_max_budget", "router_settings"]: - if json_field in record and record[json_field] is not None: - record[json_field] = json.dumps(record[json_field]) - - for rel_key in ("litellm_budget_table", "litellm_organization_table", "object_permission", "id"): - record.pop(rel_key, None) - - records.append(record) - - return records - - -async def _save_deleted_verification_token_records( - records: List[Dict[str, Any]], - prisma_client: PrismaClient, -) -> None: - """Save deleted verification token records to the database.""" - if not records: - return - await prisma_client.db.litellm_deletedverificationtoken.create_many( - data=records - ) - - -async def _persist_deleted_verification_tokens( - keys: List[LiteLLM_VerificationToken], - prisma_client: PrismaClient, - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> None: - """Persist deleted verification token records by transforming and saving them.""" - records = _transform_verification_tokens_to_deleted_records( - keys=keys, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - await _save_deleted_verification_token_records( - records=records, - prisma_client=prisma_client, - ) - - async def delete_key_aliases( key_aliases: List[str], user_api_key_cache: DualCache, prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: _keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many( where={"key_alias": {"in": key_aliases}} @@ -2625,7 +2543,6 @@ async def delete_key_aliases( tokens=tokens, user_api_key_cache=user_api_key_cache, user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c606420cc0..d1549b5116 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -34,10 +34,8 @@ from litellm.proxy._types import ( LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamMembership, LiteLLM_TeamTable, - LiteLLM_DeletedTeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, - LiteLLM_VerificationToken, LitellmTableNames, LitellmUserRoles, Member, @@ -2020,28 +2018,6 @@ async def team_member_delete( ## DELETE KEYS CREATED BY USER FOR THIS TEAM if user_ids_to_delete: - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _persist_deleted_verification_tokens, - ) - - # Fetch keys before deletion to persist them - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } - ) - ) - - if keys_to_delete: - await _persist_deleted_verification_tokens( - keys=keys_to_delete, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=None, - ) - await prisma_client.db.litellm_verificationtoken.delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, @@ -2427,13 +2403,6 @@ async def delete_team( team_row_pydantic = LiteLLM_TeamTable(**team_row_base.model_dump()) team_rows.append(team_row_pydantic) - await _persist_deleted_team_records( - teams=team_rows, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True # we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes if litellm.store_audit_logs is True: @@ -2469,25 +2438,6 @@ async def delete_team( # End of Audit logging ## DELETE ASSOCIATED KEYS - # Fetch keys before deletion to persist them - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _persist_deleted_verification_tokens, - ) - - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} - ) - ) - - if keys_to_delete: - await _persist_deleted_verification_tokens( - keys=keys_to_delete, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key") # ## DELETE TEAM MEMBERSHIPS @@ -2516,70 +2466,6 @@ async def delete_team( return deleted_teams - -def _transform_teams_to_deleted_records( - teams: List[LiteLLM_TeamTable], - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> List[Dict[str, Any]]: - """Transform teams into deleted team records ready for persistence.""" - if not teams: - return [] - - deleted_at = datetime.now(timezone.utc) - records = [] - for team in teams: - team_payload = team.model_dump() - deleted_record = LiteLLM_DeletedTeamTable( - **team_payload, - deleted_at=deleted_at, - deleted_by=user_api_key_dict.user_id, - deleted_by_api_key=user_api_key_dict.api_key, - litellm_changed_by=litellm_changed_by, - ) - record = deleted_record.model_dump() - - for json_field in ["members_with_roles", "metadata", "model_spend", "model_max_budget", "router_settings"]: - if json_field in record and record[json_field] is not None: - record[json_field] = json.dumps(record[json_field]) - - for rel_key in ("litellm_model_table", "object_permission", "id"): - record.pop(rel_key, None) - - records.append(record) - - return records - - -async def _save_deleted_team_records( - records: List[Dict[str, Any]], - prisma_client: PrismaClient, -) -> None: - """Save deleted team records to the database.""" - if not records: - return - await prisma_client.db.litellm_deletedteamtable.create_many( - data=records - ) - - -async def _persist_deleted_team_records( - teams: List[LiteLLM_TeamTable], - prisma_client: PrismaClient, - user_api_key_dict: UserAPIKeyAuth, - litellm_changed_by: Optional[str] = None, -) -> None: - """Persist deleted team records by transforming and saving them.""" - records = _transform_teams_to_deleted_records( - teams=teams, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - await _save_deleted_team_records( - records=records, - prisma_client=prisma_client, - ) - def validate_membership( user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable ): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 71b398c59a..56fe093a8b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -132,49 +132,6 @@ model LiteLLM_TeamTable { object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } -// Audit table for deleted teams - preserves spend and team information for historical tracking -model LiteLLM_DeletedTeamTable { - id String @id @default(uuid()) - team_id String // Original team_id - team_alias String? - organization_id String? - object_permission_id String? - admins String[] - members String[] - members_with_roles Json @default("{}") - metadata Json @default("{}") - max_budget Float? - spend Float @default(0.0) - models String[] - max_parallel_requests Int? - tpm_limit BigInt? - rpm_limit BigInt? - budget_duration String? - budget_reset_at DateTime? - blocked Boolean @default(false) - model_spend Json @default("{}") - model_max_budget Json @default("{}") - router_settings Json? @default("{}") - team_member_permissions String[] @default([]) - model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases - - // Original timestamps from team creation/updates - created_at DateTime? @map("created_at") - updated_at DateTime? @map("updated_at") - - // Deletion metadata - deleted_at DateTime @default(now()) @map("deleted_at") - deleted_by String? @map("deleted_by") // User who deleted the team - deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion - litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided - - @@index([team_id]) - @@index([deleted_at]) - @@index([organization_id]) - @@index([team_alias]) - @@index([created_at]) -} - // Track spend, rate limit, budget Users model LiteLLM_UserTable { user_id String @id @@ -302,62 +259,6 @@ model LiteLLM_VerificationToken { object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } -// Audit table for deleted keys - preserves spend and key information for historical tracking -model LiteLLM_DeletedVerificationToken { - id String @id @default(uuid()) - token String // Original token (hashed) - key_name String? - key_alias String? - soft_budget_cooldown Boolean @default(false) - spend Float @default(0.0) - expires DateTime? - models String[] - aliases Json @default("{}") - config Json @default("{}") - user_id String? - team_id String? - permissions Json @default("{}") - max_parallel_requests Int? - metadata Json @default("{}") - blocked Boolean? - tpm_limit BigInt? - rpm_limit BigInt? - max_budget Float? - budget_duration String? - budget_reset_at DateTime? - allowed_cache_controls String[] @default([]) - allowed_routes String[] @default([]) - model_spend Json @default("{}") - model_max_budget Json @default("{}") - router_settings Json? @default("{}") - budget_id String? - organization_id String? - object_permission_id String? - created_at DateTime? // Original creation timestamp - created_by String? // Original creator - updated_at DateTime? // Last update timestamp before deletion - updated_by String? // Last user who updated before deletion - rotation_count Int? @default(0) - auto_rotate Boolean? @default(false) - rotation_interval String? - last_rotation_at DateTime? - key_rotation_at DateTime? - - // Deletion metadata - deleted_at DateTime @default(now()) @map("deleted_at") - deleted_by String? @map("deleted_by") // User who deleted the key - deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion - litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided - - @@index([token]) - @@index([deleted_at]) - @@index([user_id]) - @@index([team_id]) - @@index([organization_id]) - @@index([key_alias]) - @@index([created_at]) -} - model LiteLLM_EndUserTable { user_id String @id alias String? // admin-facing alias diff --git a/schema.prisma b/schema.prisma index 52170f2f3e..a16380fb5f 100644 --- a/schema.prisma +++ b/schema.prisma @@ -132,49 +132,6 @@ model LiteLLM_TeamTable { object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } -// Audit table for deleted teams - preserves spend and team information for historical tracking -model LiteLLM_DeletedTeamTable { - id String @id @default(uuid()) - team_id String // Original team_id - team_alias String? - organization_id String? - object_permission_id String? - admins String[] - members String[] - members_with_roles Json @default("{}") - metadata Json @default("{}") - max_budget Float? - spend Float @default(0.0) - models String[] - max_parallel_requests Int? - tpm_limit BigInt? - rpm_limit BigInt? - budget_duration String? - budget_reset_at DateTime? - blocked Boolean @default(false) - model_spend Json @default("{}") - model_max_budget Json @default("{}") - router_settings Json? @default("{}") - team_member_permissions String[] @default([]) - model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases - - // Original timestamps from team creation/updates - created_at DateTime? @map("created_at") - updated_at DateTime? @map("updated_at") - - // Deletion metadata - deleted_at DateTime @default(now()) @map("deleted_at") - deleted_by String? @map("deleted_by") // User who deleted the team - deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion - litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided - - @@index([team_id]) - @@index([deleted_at]) - @@index([organization_id]) - @@index([team_alias]) - @@index([created_at]) -} - // Track spend, rate limit, budget Users model LiteLLM_UserTable { user_id String @id @@ -302,62 +259,6 @@ model LiteLLM_VerificationToken { object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } -// Audit table for deleted keys - preserves spend and key information for historical tracking -model LiteLLM_DeletedVerificationToken { - id String @id @default(uuid()) - token String // Original token (hashed) - key_name String? - key_alias String? - soft_budget_cooldown Boolean @default(false) - spend Float @default(0.0) - expires DateTime? - models String[] - aliases Json @default("{}") - config Json @default("{}") - user_id String? - team_id String? - permissions Json @default("{}") - max_parallel_requests Int? - metadata Json @default("{}") - blocked Boolean? - tpm_limit BigInt? - rpm_limit BigInt? - max_budget Float? - budget_duration String? - budget_reset_at DateTime? - allowed_cache_controls String[] @default([]) - allowed_routes String[] @default([]) - model_spend Json @default("{}") - model_max_budget Json @default("{}") - router_settings Json? @default("{}") - budget_id String? - organization_id String? - object_permission_id String? - created_at DateTime? // Original creation timestamp - created_by String? // Original creator - updated_at DateTime? // Last update timestamp before deletion - updated_by String? // Last user who updated before deletion - rotation_count Int? @default(0) - auto_rotate Boolean? @default(false) - rotation_interval String? - last_rotation_at DateTime? - key_rotation_at DateTime? - - // Deletion metadata - deleted_at DateTime @default(now()) @map("deleted_at") - deleted_by String? @map("deleted_by") // User who deleted the key - deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion - litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided - - @@index([token]) - @@index([deleted_at]) - @@index([user_id]) - @@index([team_id]) - @@index([organization_id]) - @@index([key_alias]) - @@index([created_at]) -} - model LiteLLM_EndUserTable { user_id String @id alias String? // admin-facing alias diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index a196080ead..126718af84 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -1061,7 +1061,6 @@ async def test_list_key_helper(prisma_client): api_key="sk-1234", user_id="admin", ), - litellm_changed_by=None, ) @@ -1182,7 +1181,6 @@ async def test_list_key_helper_team_filtering(prisma_client): api_key="sk-1234", user_id="admin", ), - litellm_changed_by=None, ) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 1a613a3db5..e0d6b7e81b 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -1166,10 +1166,8 @@ def test_delete_key_auth(prisma_client): asyncio.run(test()) except Exception as e: print("Got Exception", e) - # Handle different exception types - ProxyException has .message, others might have .detail or str(e) - error_message = getattr(e, "message", None) or getattr(e, "detail", None) or str(e) - print(f"Error message: {error_message}") - assert "Authentication Error" in error_message or "Invalid proxy server token" in error_message or "not found in db" in error_message + print(e.message) + assert "Authentication Error" in e.message pass @@ -2710,12 +2708,7 @@ async def test_reset_spend_authentication(prisma_client): _response = await new_user( data=NewUserRequest( tpm_limit=20, - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=master_key, - user_id="1234", - ), + ) ) generate_key = "Bearer " + _response.key @@ -2735,12 +2728,7 @@ async def test_reset_spend_authentication(prisma_client): data=NewUserRequest( user_role=LitellmUserRoles.PROXY_ADMIN, tpm_limit=20, - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=master_key, - user_id="1234", - ), + ) ) generate_key = "Bearer " + _response.key 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 47395a1f32..c9a10e3c4d 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 @@ -31,13 +31,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_team_key_limits, _common_key_generation_helper, _list_key_helper, - _persist_deleted_verification_tokens, - _save_deleted_verification_token_records, - _transform_verification_tokens_to_deleted_records, can_modify_verification_token, check_org_key_model_specific_limits, check_team_key_model_specific_limits, - delete_verification_tokens, generate_key_helper_fn, prepare_key_update_data, validate_key_team_change, @@ -2732,364 +2728,64 @@ def test_check_org_key_model_specific_limits_org_model_tpm_overallocation(): ) -def test_transform_verification_tokens_to_deleted_records(): - from datetime import datetime, timezone - - user_api_key_dict = UserAPIKeyAuth( - user_id="user-123", - api_key="sk-test", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - key1 = LiteLLM_VerificationToken( - token="hashed-token-1", - user_id="user-123", - team_id="team-456", - key_alias="test-key-1", - spend=100.0, - max_budget=1000.0, - models=["gpt-4"], - aliases={}, - config={}, - permissions={}, - metadata={"test": "value"}, - model_max_budget={}, - model_spend={}, - soft_budget_cooldown=False, - allowed_routes=[], - ) - - key2 = LiteLLM_VerificationToken( - token="hashed-token-2", - user_id="user-789", - team_id=None, - key_alias="test-key-2", - spend=50.0, - max_budget=500.0, - models=["gpt-3.5-turbo"], - aliases={"alias": "model"}, - config={"config": "value"}, - permissions={"permission": True}, - metadata={}, - model_max_budget={"gpt-4": {"budget_limit": 100.0}}, - model_spend={}, - soft_budget_cooldown=False, - allowed_routes=[], - ) - - records = _transform_verification_tokens_to_deleted_records( - keys=[key1, key2], - user_api_key_dict=user_api_key_dict, - litellm_changed_by="admin-user", - ) - - assert len(records) == 2 - assert all("deleted_at" in record for record in records) - assert all("deleted_by" in record for record in records) - assert all("deleted_by_api_key" in record for record in records) - assert all("litellm_changed_by" in record for record in records) - assert all(record["deleted_by"] == "user-123" for record in records) - assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records) - assert all(record["litellm_changed_by"] == "admin-user" for record in records) - - record1 = records[0] - assert record1["token"] == "hashed-token-1" - assert record1["user_id"] == "user-123" - assert record1["team_id"] == "team-456" - assert isinstance(record1["aliases"], str) - assert isinstance(record1["config"], str) - assert isinstance(record1["permissions"], str) - assert isinstance(record1["metadata"], str) - assert "litellm_budget_table" not in record1 - assert "litellm_organization_table" not in record1 - assert "object_permission" not in record1 - assert "id" not in record1 - - record2 = records[1] - assert record2["token"] == "hashed-token-2" - assert isinstance(record2["model_max_budget"], str) - - -def test_transform_verification_tokens_to_deleted_records_empty_list(): - user_api_key_dict = UserAPIKeyAuth( - user_id="user-123", - api_key="sk-test", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - records = _transform_verification_tokens_to_deleted_records( - keys=[], - user_api_key_dict=user_api_key_dict, - ) - - assert records == [] - - -@pytest.mark.asyncio -async def test_save_deleted_verification_token_records(): - mock_prisma_client = AsyncMock() - mock_create_many = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many - ) - - records = [ - { - "token": "hashed-token-1", - "user_id": "user-123", - "deleted_at": "2024-01-01T00:00:00Z", - "deleted_by": "admin", - }, - { - "token": "hashed-token-2", - "user_id": "user-456", - "deleted_at": "2024-01-01T00:00:00Z", - "deleted_by": "admin", - }, - ] - - await _save_deleted_verification_token_records( - records=records, prisma_client=mock_prisma_client - ) - - mock_create_many.assert_called_once_with(data=records) - - -@pytest.mark.asyncio -async def test_save_deleted_verification_token_records_empty_list(): - mock_prisma_client = AsyncMock() - mock_create_many = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many - ) - - await _save_deleted_verification_token_records( - records=[], prisma_client=mock_prisma_client - ) - - mock_create_many.assert_not_called() - - -@pytest.mark.asyncio -async def test_persist_deleted_verification_tokens(): - mock_prisma_client = AsyncMock() - mock_create_many = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many - ) - - user_api_key_dict = UserAPIKeyAuth( - user_id="user-123", - api_key="sk-test", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - key = LiteLLM_VerificationToken( - token="hashed-token-1", - user_id="user-123", - team_id="team-456", - key_alias="test-key", - spend=100.0, - max_budget=1000.0, - models=["gpt-4"], - aliases={}, - config={}, - permissions={}, - metadata={}, - model_max_budget={}, - model_spend={}, - soft_budget_cooldown=False, - allowed_routes=[], - ) - - await _persist_deleted_verification_tokens( - keys=[key], - prisma_client=mock_prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by="admin-user", - ) - - mock_create_many.assert_called_once() - call_args = mock_create_many.call_args - assert "data" in call_args.kwargs - records = call_args.kwargs["data"] - assert len(records) == 1 - assert records[0]["token"] == "hashed-token-1" - assert records[0]["deleted_by"] == "user-123" - assert records[0]["litellm_changed_by"] == "admin-user" - - -@pytest.mark.asyncio -async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch): - mock_prisma_client = AsyncMock() - mock_user_api_key_cache = MagicMock() - - user_api_key_dict = UserAPIKeyAuth( - user_id="admin-user", - api_key="sk-admin", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - key1 = LiteLLM_VerificationToken( - token="hashed-token-1", - user_id="user-123", - team_id="team-456", - key_alias="test-key-1", - spend=100.0, - max_budget=1000.0, - models=["gpt-4"], - aliases={}, - config={}, - permissions={}, - metadata={}, - model_max_budget={}, - model_spend={}, - soft_budget_cooldown=False, - allowed_routes=[], - ) - - key2 = LiteLLM_VerificationToken( - token="hashed-token-2", - user_id="user-789", - team_id=None, - key_alias="test-key-2", - spend=50.0, - max_budget=500.0, - models=["gpt-3.5-turbo"], - aliases={}, - config={}, - permissions={}, - metadata={}, - model_max_budget={}, - model_spend={}, - soft_budget_cooldown=False, - allowed_routes=[], - ) - - mock_find_many = AsyncMock(return_value=[key1, key2]) - mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many - - # delete_data returns {"deleted_keys": ...} from utils.py line 3049 - # The function at line 2410 assigns it to deleted_tokens - # Then at line 2444 returns {"deleted_keys": deleted_tokens} - # So if delete_data returns {"deleted_keys": list}, then result would be nested - # But looking at the error, it seems like delete_data might return just the list - # Or the code extracts it. Let's return the list directly since that's what the test expects - mock_delete_data = AsyncMock(return_value=["hashed-token-1", "hashed-token-2"]) - mock_prisma_client.delete_data = mock_delete_data - - # Mock cache delete_cache method - mock_user_api_key_cache.delete_cache = MagicMock() - - mock_create_many = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many - ) - - def mock_hash_token(token): - return token if not token.startswith("sk-") else f"hashed-{token}" - - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - mock_hash_token, - ) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.hash_token", - mock_hash_token, - ) - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", - mock_prisma_client, - ) - - result, deleted_keys = await delete_verification_tokens( - tokens=["sk-token-1", "sk-token-2"], - user_api_key_cache=mock_user_api_key_cache, - user_api_key_dict=user_api_key_dict, - litellm_changed_by="admin-user", - ) - - mock_create_many.assert_called_once() - call_args = mock_create_many.call_args - assert "data" in call_args.kwargs - records = call_args.kwargs["data"] - assert len(records) == 2 - assert all(record["deleted_by"] == "admin-user" for record in records) - assert all(record["litellm_changed_by"] == "admin-user" for record in records) - # delete_data returns the list directly, which gets wrapped in {"deleted_keys": ...} - assert isinstance(result["deleted_keys"], list) - assert set(result["deleted_keys"]) == {"hashed-token-1", "hashed-token-2"} - assert len(deleted_keys) == 2 - - -@pytest.mark.asyncio -async def test_delete_key_fn_persists_deleted_keys(monkeypatch): - from litellm.proxy._types import KeyRequest - from litellm.proxy.management_endpoints.key_management_endpoints import ( - delete_key_fn, - delete_verification_tokens, - ) - - mock_prisma_client = AsyncMock() - mock_user_api_key_cache = MagicMock() - - user_api_key_dict = UserAPIKeyAuth( - user_id="admin-user", - api_key="sk-admin", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - key1 = LiteLLM_VerificationToken( - token="hashed-token-1", - user_id="user-123", - team_id="team-456", - key_alias="test-key-1", - spend=100.0, - max_budget=1000.0, - models=["gpt-4"], - aliases={}, - config={}, - permissions={}, - metadata={}, - model_max_budget={}, - model_spend={}, - soft_budget_cooldown=False, - allowed_routes=[], - ) - - async def mock_delete_verification_tokens(*args, **kwargs): - return ({"deleted_keys": ["sk-token-1"]}, [key1]) - - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.delete_verification_tokens", - mock_delete_verification_tokens, - ) - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", - mock_prisma_client, - ) - monkeypatch.setattr( - "litellm.proxy.proxy_server.user_api_key_cache", - mock_user_api_key_cache, - ) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_deleted_hook", - AsyncMock(), - ) - - data = KeyRequest(keys=["sk-token-1"]) - - result = await delete_key_fn( - data=data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by="admin-user", - ) - - assert result["deleted_keys"] == ["sk-token-1"] - - @pytest.mark.asyncio async def test_can_delete_verification_token_proxy_admin_team_key(monkeypatch): + """Test that proxy admin can delete any team key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_proxy_admin_personal_key(monkeypatch): + """Test that proxy admin can delete any personal key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_modify_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_team_admin_own_team(monkeypatch): """Test that team admin can delete team keys from their own team.""" key_info = LiteLLM_VerificationToken( token="test-token", 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..bbff7448e1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -33,13 +33,8 @@ from litellm.proxy.management_endpoints.team_endpoints import ( from litellm.proxy.management_endpoints.team_endpoints import ( GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, - _persist_deleted_team_records, - _save_deleted_team_records, - _transform_teams_to_deleted_records, - delete_team, router, team_member_add_duplication_check, - team_member_delete, validate_team_org_change, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( @@ -2265,7 +2260,6 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a # Verification token deletion should be called mock_db_client.db.litellm_verificationtoken = MagicMock() - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) # Execute @@ -2313,7 +2307,6 @@ async def test_team_member_delete_cleans_verification_tokens(mock_db_client, moc mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock()) mock_db_client.db.litellm_verificationtoken = MagicMock() - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) await team_member_delete( @@ -4332,348 +4325,6 @@ async def test_update_team_guardrails_with_org_id(): assert first_call_kwargs["include"]["teams"] is True -def test_transform_teams_to_deleted_records(): - from datetime import datetime, timezone - - user_api_key_dict = UserAPIKeyAuth( - user_id="user-123", - api_key="sk-test", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - team1 = LiteLLM_TeamTable( - team_id="team-1", - team_alias="test-team-1", - members_with_roles=[ - Member(user_id="user-1", role="admin"), - Member(user_id="user-2", role="user"), - ], - metadata={"test": "value"}, - model_max_budget={}, - model_spend={}, - ) - - team2 = LiteLLM_TeamTable( - team_id="team-2", - team_alias="test-team-2", - members_with_roles=[], - metadata=None, - model_max_budget={"gpt-4": {"budget_limit": 100.0}}, - model_spend={}, - ) - - records = _transform_teams_to_deleted_records( - teams=[team1, team2], - user_api_key_dict=user_api_key_dict, - litellm_changed_by="admin-user", - ) - - assert len(records) == 2 - assert all("deleted_at" in record for record in records) - assert all("deleted_by" in record for record in records) - assert all("deleted_by_api_key" in record for record in records) - assert all("litellm_changed_by" in record for record in records) - assert all(record["deleted_by"] == "user-123" for record in records) - # UserAPIKeyAuth hashes the api_key, so we check against the hashed value - assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records) - assert all(record["litellm_changed_by"] == "admin-user" for record in records) - - record1 = records[0] - assert record1["team_id"] == "team-1" - assert isinstance(record1["members_with_roles"], str) - assert isinstance(record1["metadata"], str) - assert "litellm_model_table" not in record1 - assert "object_permission" not in record1 - assert "id" not in record1 - - record2 = records[1] - assert record2["team_id"] == "team-2" - # model_max_budget should be converted to JSON string if it exists - if "model_max_budget" in record2: - assert isinstance(record2["model_max_budget"], str) - - -def test_transform_teams_to_deleted_records_empty_list(): - user_api_key_dict = UserAPIKeyAuth( - user_id="user-123", - api_key="sk-test", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - records = _transform_teams_to_deleted_records( - teams=[], - user_api_key_dict=user_api_key_dict, - ) - - assert records == [] - - -@pytest.mark.asyncio -async def test_save_deleted_team_records(): - mock_prisma_client = AsyncMock() - mock_create_many = AsyncMock() - mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many - - records = [ - { - "team_id": "team-1", - "team_alias": "test-team-1", - "deleted_at": "2024-01-01T00:00:00Z", - "deleted_by": "admin", - }, - { - "team_id": "team-2", - "team_alias": "test-team-2", - "deleted_at": "2024-01-01T00:00:00Z", - "deleted_by": "admin", - }, - ] - - await _save_deleted_team_records(records=records, prisma_client=mock_prisma_client) - - mock_create_many.assert_called_once_with(data=records) - - -@pytest.mark.asyncio -async def test_save_deleted_team_records_empty_list(): - mock_prisma_client = AsyncMock() - mock_create_many = AsyncMock() - mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many - - await _save_deleted_team_records(records=[], prisma_client=mock_prisma_client) - - mock_create_many.assert_not_called() - - -@pytest.mark.asyncio -async def test_persist_deleted_team_records(): - mock_prisma_client = AsyncMock() - mock_create_many = AsyncMock() - mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many - - user_api_key_dict = UserAPIKeyAuth( - user_id="user-123", - api_key="sk-test", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - team = LiteLLM_TeamTable( - team_id="team-1", - team_alias="test-team", - members_with_roles=[ - Member(user_id="user-1", role="admin"), - ], - metadata={}, - model_max_budget={}, - model_spend={}, - ) - - await _persist_deleted_team_records( - teams=[team], - prisma_client=mock_prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by="admin-user", - ) - - mock_create_many.assert_called_once() - call_args = mock_create_many.call_args - assert "data" in call_args.kwargs - records = call_args.kwargs["data"] - assert len(records) == 1 - assert records[0]["team_id"] == "team-1" - assert records[0]["deleted_by"] == "user-123" - assert records[0]["litellm_changed_by"] == "admin-user" - - -@pytest.mark.asyncio -async def test_delete_team_persists_deleted_teams(monkeypatch): - from litellm.proxy._types import DeleteTeamRequest - - mock_prisma_client = AsyncMock() - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="admin-user", - api_key="sk-admin", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - team1 = LiteLLM_TeamTable( - team_id="team-1", - team_alias="test-team-1", - members_with_roles=[ - Member(user_id="user-1", role="admin"), - ], - metadata={}, - model_max_budget={}, - model_spend={}, - ) - - mock_find_unique = AsyncMock(return_value=team1) - mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique - - mock_delete_data = AsyncMock(return_value={"deleted_teams": ["team-1"]}) - mock_prisma_client.delete_data = mock_delete_data - - mock_create_many_teams = AsyncMock() - mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many_teams - - mock_create_many_keys = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many_keys - ) - - mock_find_many_keys = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys - - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", - mock_prisma_client, - ) - monkeypatch.setattr( - "litellm.proxy.proxy_server.create_audit_log_for_update", - AsyncMock(), - ) - monkeypatch.setattr( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", - "admin", - ) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.team_endpoints.team_member_delete", - AsyncMock(return_value=team1), - ) - - data = DeleteTeamRequest(team_ids=["team-1"]) - - result = await delete_team( - data=data, - http_request=MagicMock(), - user_api_key_dict=mock_user_api_key_dict, - litellm_changed_by="admin-user", - ) - - mock_create_many_teams.assert_called_once() - call_args = mock_create_many_teams.call_args - assert "data" in call_args.kwargs - records = call_args.kwargs["data"] - assert len(records) == 1 - assert records[0]["team_id"] == "team-1" - assert records[0]["deleted_by"] == "admin-user" - assert records[0]["litellm_changed_by"] == "admin-user" - - -@pytest.mark.asyncio -async def test_team_member_delete_persists_deleted_keys(monkeypatch): - from litellm.proxy._types import TeamMemberDeleteRequest - from litellm.proxy.management_endpoints.key_management_endpoints import ( - LiteLLM_VerificationToken, - ) - - mock_prisma_client = AsyncMock() - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="admin-user", - api_key="sk-admin", - user_role=LitellmUserRoles.PROXY_ADMIN.value, - ) - - team = LiteLLM_TeamTable( - team_id="team-1", - team_alias="test-team", - members_with_roles=[ - Member(user_id="user-123", role="admin"), - ], - metadata={}, - model_max_budget={}, - model_spend={}, - ) - - key1 = LiteLLM_VerificationToken( - token="hashed-token-1", - user_id="user-123", - team_id="team-1", - key_alias="test-key-1", - spend=100.0, - max_budget=1000.0, - models=["gpt-4"], - aliases={}, - config={}, - permissions={}, - metadata={}, - model_max_budget={}, - ) - - key2 = LiteLLM_VerificationToken( - token="hashed-token-2", - user_id="user-123", - team_id="team-1", - key_alias="test-key-2", - spend=50.0, - max_budget=500.0, - models=["gpt-3.5-turbo"], - aliases={}, - config={}, - permissions={}, - metadata={}, - model_max_budget={}, - ) - - mock_find_unique_team = AsyncMock(return_value=team) - mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team - - mock_find_many_user = AsyncMock( - return_value=[ - MagicMock( - user_id="user-123", - teams=["team-1"], - model_dump=lambda: {"user_id": "user-123", "teams": ["team-1"]}, - ) - ] - ) - mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_user - - mock_update_team = AsyncMock() - mock_prisma_client.db.litellm_teamtable.update = mock_update_team - - mock_update_user = AsyncMock() - mock_prisma_client.db.litellm_usertable.update = mock_update_user - - mock_delete_membership = AsyncMock() - mock_prisma_client.db.litellm_teammembership.delete_many = mock_delete_membership - - mock_find_many_keys = AsyncMock(return_value=[key1, key2]) - mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys - - mock_delete_keys = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.delete_many = mock_delete_keys - - mock_create_many_keys = AsyncMock() - mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( - mock_create_many_keys - ) - - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", - mock_prisma_client, - ) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", - lambda **kwargs: True, - ) - - data = TeamMemberDeleteRequest(team_id="team-1", user_id="user-123") - - result = await team_member_delete( - data=data, - user_api_key_dict=mock_user_api_key_dict, - ) - - mock_create_many_keys.assert_called_once() - call_args = mock_create_many_keys.call_args - assert "data" in call_args.kwargs - records = call_args.kwargs["data"] - assert len(records) == 2 - assert all(record["deleted_by"] == "admin-user" for record in records) - assert all(record["team_id"] == "team-1" for record in records) - assert all(record["user_id"] == "user-123" for record in records) - mock_delete_keys.assert_called_once() @pytest.mark.asyncio async def test_new_team_negative_max_budget(): """