fix: address review feedback - dedup logic, use module-level helper, add test

- Deduplicate: update_key_fn now delegates to _get_and_validate_existing_key()
  instead of inlining its own copy of the lookup logic
- Use _hash_token_if_needed (already imported at module level) instead of
  inline `from proxy_server import hash_token` + manual conditional
- Fix stale docstring: _get_and_validate_existing_key raises ProxyException,
  not HTTPException
- Add unit test: test_update_key_nonexistent_key_returns_404

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -1654,7 +1654,7 @@ async def _get_and_validate_existing_key(
         LiteLLM_VerificationToken: The existing key row
 
     Raises:
-        HTTPException: If key is not found
+        ProxyException: 404 if key is not found
     """
     if prisma_client is None:
         raise HTTPException(
@@ -1662,12 +1662,7 @@ async def _get_and_validate_existing_key(
             detail={"error": "Database not connected"},
         )
 
-    from litellm.proxy.proxy_server import hash_token
-
-    if token.startswith("sk-"):
-        hashed_token = hash_token(token=token)
-    else:
-        hashed_token = token
+    hashed_token = _hash_token_if_needed(token=token)
 
     existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
         where={"token": hashed_token}
@@ -2116,28 +2111,11 @@ async def update_key_fn(
         key = data_json.pop("key")
 
         # get the row from db
-        if prisma_client is None:
-            raise Exception("Not connected to DB!")
-
-        from litellm.proxy.proxy_server import hash_token
-
-        if data.key.startswith("sk-"):
-            hashed_token = hash_token(token=data.key)
-        else:
-            hashed_token = data.key
-
-        existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
-            where={"token": hashed_token}
+        existing_key_row = await _get_and_validate_existing_key(
+            token=data.key,
+            prisma_client=prisma_client,
         )
 
-        if existing_key_row is None:
-            raise ProxyException(
-                message=f"Key not found. Passed key={data.key}",
-                type=ProxyErrorTypes.not_found_error,
-                param="key",
-                code=status.HTTP_404_NOT_FOUND,
-            )
-
         await _validate_update_key_data(
             data=data,
             existing_key_row=existing_key_row,

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -1678,6 +1678,56 @@ async def test_unblock_key_nonexistent_key_returns_404(monkeypatch):
     mock_prisma_client.db.litellm_verificationtoken.update.assert_not_called()
 
 
+@pytest.mark.asyncio
+async def test_update_key_nonexistent_key_returns_404(monkeypatch):
+    """
+    Test that update_key_fn returns 404 (not misleading 401) when the body
+    key doesn't exist in the database, even when the caller is authenticated
+    as a proxy admin via the Authorization header.
+    """
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        update_key_fn,
+    )
+
+    mock_prisma_client = AsyncMock()
+    mock_user_api_key_cache = MagicMock()
+    mock_proxy_logging_obj = MagicMock()
+
+    # find_unique returns None → key does not exist
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=None
+    )
+
+    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.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
+
+    mock_request = MagicMock()
+    user_api_key_dict = UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
+    )
+
+    data = UpdateKeyRequest(key="sk-does-not-exist-key")
+
+    with pytest.raises(ProxyException) as exc_info:
+        await update_key_fn(
+            request=mock_request,
+            data=data,
+            user_api_key_dict=user_api_key_dict,
+            litellm_changed_by=None,
+        )
+
+    assert exc_info.value.code == "404"
+    assert "not found" in str(exc_info.value.message).lower()
+    assert "Authentication Error" not in str(exc_info.value.message)
+
+
 @pytest.mark.asyncio
 async def test_block_key_existing_key_succeeds(monkeypatch):
     """