From 48ee3eddf3e012d52ecf3fa2e2ad7ed58f057a5c Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Fri, 11 Sep 2026 00:51:10 +0000
Subject: [PATCH] refactor(proxy): extract key update caller access check

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../key_management_endpoints.py               | 68 ++++++++++++-------
 1 file changed, 45 insertions(+), 23 deletions(-)

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 8fb040edd19..6ddb51490c3 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -2680,6 +2680,42 @@ def _require_prisma_client(prisma_client: PrismaClient | None) -> PrismaClient:
     return prisma_client
 
 
+async def _validate_update_key_caller_access(
+    data: UpdateKeyRequest,
+    existing_key_row: LiteLLM_VerificationToken,
+    user_api_key_dict: UserAPIKeyAuth,
+    llm_router: Router | None,
+    premium_user: bool,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    is_proxy_admin: bool,
+) -> None:
+    _caller_is_key_owner: Final = (
+        existing_key_row.user_id is None or existing_key_row.user_id == user_api_key_dict.user_id
+    )
+    if is_proxy_admin or _caller_is_key_owner or existing_key_row.team_id is None:
+        common_key_access_checks(
+            user_api_key_dict=user_api_key_dict,
+            data=data,
+            user_id=existing_key_row.user_id,
+            llm_router=llm_router,
+            premium_user=premium_user,
+        )
+        return
+    await _check_key_admin_access(
+        user_api_key_dict=user_api_key_dict,
+        hashed_token=existing_key_row.token,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        route="/key/update",
+    )
+    _check_model_access_group(
+        models=data.models,
+        llm_router=llm_router,
+        premium_user=premium_user,
+    )
+
+
 async def _validate_update_key_data(
     data: UpdateKeyRequest,
     existing_key_row: LiteLLM_VerificationToken,
@@ -2718,30 +2754,16 @@ async def _validate_update_key_data(
         user_api_key_dict=user_api_key_dict,
     )
 
-    _caller_is_key_owner: Final = (
-        existing_key_row.user_id is None or existing_key_row.user_id == user_api_key_dict.user_id
+    await _validate_update_key_caller_access(
+        data=data,
+        existing_key_row=existing_key_row,
+        user_api_key_dict=user_api_key_dict,
+        llm_router=llm_router,
+        premium_user=premium_user,
+        prisma_client=checked_prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        is_proxy_admin=_is_proxy_admin,
     )
-    if _is_proxy_admin or _caller_is_key_owner or existing_key_row.team_id is None:
-        common_key_access_checks(
-            user_api_key_dict=user_api_key_dict,
-            data=data,
-            user_id=existing_key_row.user_id,
-            llm_router=llm_router,
-            premium_user=premium_user,
-        )
-    else:
-        await _check_key_admin_access(
-            user_api_key_dict=user_api_key_dict,
-            hashed_token=existing_key_row.token,
-            prisma_client=checked_prisma_client,
-            user_api_key_cache=user_api_key_cache,
-            route="/key/update",
-        )
-        _check_model_access_group(
-            models=data.models,
-            llm_router=llm_router,
-            premium_user=premium_user,
-        )
 
     await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
         user_api_key_dict=user_api_key_dict,