diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ef5b7b30d..1856028c6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -593,7 +593,7 @@ async def update_key_fn( await _enforce_unique_key_alias( key_alias=non_default_values.get("key_alias", None), prisma_client=prisma_client, - existing_key_id=existing_key_row.token, + existing_key_token=existing_key_row.token, ) response = await prisma_client.update_data( @@ -1899,7 +1899,7 @@ async def test_key_logging( async def _enforce_unique_key_alias( key_alias: Optional[str], prisma_client: Any, - existing_key_id: Optional[str] = None, + existing_key_token: Optional[str] = None, ) -> None: """ Helper to enforce unique key aliases across all keys. @@ -1907,16 +1907,17 @@ async def _enforce_unique_key_alias( Args: key_alias (Optional[str]): The key alias to check prisma_client (Any): Prisma client instance - existing_key_id (Optional[str]): ID of existing key being updated, to exclude from uniqueness check + existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check + (The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating) Raises: HTTPException: If key alias already exists on a different key """ if key_alias is not None and prisma_client is not None: where_clause: dict[str, Any] = {"key_alias": key_alias} - if existing_key_id: + if existing_key_token: # Exclude the current key from the uniqueness check - where_clause["NOT"] = {"token": existing_key_id} + where_clause["NOT"] = {"token": existing_key_token} existing_key = await prisma_client.db.litellm_verificationtoken.find_first( where=where_clause diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 935e4c352..87bbe76d3 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3623,3 +3623,79 @@ async def test_key_alias_uniqueness(prisma_client): print("got exceptions, e=", e) print("vars(e)=", vars(e)) pytest.fail(f"An unexpected error occurred: {str(e)}") + + +@pytest.mark.asyncio +async def test_enforce_unique_key_alias(prisma_client): + """ + Unit test the _enforce_unique_key_alias function: + 1. Test it allows unique aliases + 2. Test it blocks duplicate aliases for new keys + 3. Test it allows updating a key with its own existing alias + 4. Test it blocks updating a key with another key's alias + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _enforce_unique_key_alias, + ) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + await litellm.proxy.proxy_server.prisma_client.connect() + + try: + # Test 1: Allow unique alias + unique_alias = f"test-alias-{uuid.uuid4()}" + await _enforce_unique_key_alias( + key_alias=unique_alias, + prisma_client=prisma_client, + ) # Should pass + + # Create a key with this alias in the database + key1 = await generate_key_fn( + data=GenerateKeyRequest(key_alias=unique_alias), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Test 2: Block duplicate alias for new key + try: + await _enforce_unique_key_alias( + key_alias=unique_alias, + prisma_client=prisma_client, + ) + pytest.fail("Should not allow duplicate alias") + except Exception as e: + assert "Unique key aliases across all keys are required" in str(e.message) + + # Test 3: Allow updating key with its own alias + await _enforce_unique_key_alias( + key_alias=unique_alias, + existing_key_token=hash_token(key1.key), + prisma_client=prisma_client, + ) # Should pass + + # Test 4: Block updating with another key's alias + another_key = await generate_key_fn( + data=GenerateKeyRequest(key_alias=f"test-alias-{uuid.uuid4()}"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + try: + await _enforce_unique_key_alias( + key_alias=unique_alias, + existing_key_token=another_key.key, + prisma_client=prisma_client, + ) + pytest.fail("Should not allow using another key's alias") + except Exception as e: + assert "Unique key aliases across all keys are required" in str(e.message) + + except Exception as e: + print("Unexpected error:", e) + pytest.fail(f"An unexpected error occurred: {str(e)}")