From f6cd0a827ae84cffa838eac859eaea504bd6464a Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Wed, 18 Mar 2026 16:53:29 -0700
Subject: [PATCH] fix: /key/update returns 404 (not 401) for nonexistent body
 key

The /key/update endpoint's get_data() call raises a 401 when the body
`key` field doesn't exist in the DB, because get_data() treats the
token as an auth credential. This caused the auth layer to resolve the
body key instead of the Authorization header bearer token.

Replace prisma_client.get_data() with direct Prisma find_unique() in
both _get_and_validate_existing_key() and update_key_fn(), matching
the pattern used in the /key/block and /key/unblock fix (PR #23977).
Also fix the incorrect "Team not found" error message in update_key_fn.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../key_management_endpoints.py               | 40 +++++++++++++------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index ec33816b78..2884c8d374 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -1662,16 +1662,23 @@ async def _get_and_validate_existing_key(
             detail={"error": "Database not connected"},
         )
 
-    existing_key_row = await prisma_client.get_data(
-        token=token,
-        table_name="key",
-        query_type="find_unique",
+    from litellm.proxy.proxy_server import hash_token
+
+    if token.startswith("sk-"):
+        hashed_token = hash_token(token=token)
+    else:
+        hashed_token = token
+
+    existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
+        where={"token": hashed_token}
     )
 
     if existing_key_row is None:
-        raise HTTPException(
-            status_code=404,
-            detail={"error": f"Key not found: {token}"},
+        raise ProxyException(
+            message=f"Key not found. Passed key={token}",
+            type=ProxyErrorTypes.not_found_error,
+            param="key",
+            code=status.HTTP_404_NOT_FOUND,
         )
 
     return existing_key_row
@@ -2112,14 +2119,23 @@ async def update_key_fn(
         if prisma_client is None:
             raise Exception("Not connected to DB!")
 
-        existing_key_row = await prisma_client.get_data(
-            token=data.key, table_name="key", query_type="find_unique"
+        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}
         )
 
         if existing_key_row is None:
-            raise HTTPException(
-                status_code=404,
-                detail={"error": f"Team not found, passed team_id={data.team_id}"},
+            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(