From 074cf0dfb31a2dc2587567be8a9252be616ce6f2 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 6 May 2026 23:30:51 +0000
Subject: [PATCH] Avoid redundant bulk team key lookups

---
 .../management_endpoints/key_management_endpoints.py | 12 ++++++++----
 .../test_key_management_endpoints.py                 | 11 +++++++----
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index d6702cfbdf1..4439a55c1c6 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -1889,6 +1889,7 @@ async def _process_single_key_update(
     proxy_logging_obj: Any,
     llm_router: Optional[Router],
     user_custom_key_update: Optional[Callable] = None,
+    existing_key_row: Optional[LiteLLM_VerificationToken] = None,
 ) -> Dict[str, Any]:
     """
     Process a single key update with all validations and checks.
@@ -1904,6 +1905,7 @@ async def _process_single_key_update(
         user_api_key_cache: User API key cache
         proxy_logging_obj: Proxy logging object
         llm_router: LLM router instance
+        existing_key_row: Optional pre-fetched key row to avoid redundant lookups
 
     Returns:
         Dict containing the updated key information
@@ -1915,10 +1917,11 @@ async def _process_single_key_update(
     _validate_max_budget(update_key_request.max_budget)
 
     # Get and validate existing key
-    existing_key_row = await _get_and_validate_existing_key(
-        token=update_key_request.key,
-        prisma_client=prisma_client,
-    )
+    if existing_key_row is None:
+        existing_key_row = await _get_and_validate_existing_key(
+            token=update_key_request.key,
+            prisma_client=prisma_client,
+        )
 
     # Check team member permissions
     if prisma_client is not None:
@@ -2852,6 +2855,7 @@ async def bulk_update_team_keys(
                 proxy_logging_obj=proxy_logging_obj,
                 llm_router=llm_router,
                 user_custom_key_update=user_custom_key_update,
+                existing_key_row=existing_by_token[db_token],
             )
 
             successful_updates.append(
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 653cf161e06..66716400c4f 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
@@ -10117,10 +10117,11 @@ async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch):
     )
 
     keys = [_make_team_key("tok-a"), _make_team_key("tok-b")]
+    find_unique = AsyncMock(side_effect=keys)
     mock = _setup_team_keys_mocks(
         monkeypatch,
         find_many=keys,
-        find_unique=AsyncMock(side_effect=keys),
+        find_unique=find_unique,
         update_data=AsyncMock(
             side_effect=[{"data": _updated({"max_budget": 50.0})}] * 2
         ),
@@ -10139,6 +10140,7 @@ async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch):
     where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
     assert where["team_id"] == "team-abc"
     assert where["token"] == {"in": ["tok-a", "tok-b"]}
+    find_unique.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -10149,10 +10151,11 @@ async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
     )
 
     keys = [_make_team_key(f"tok-{i}") for i in range(3)]
+    find_unique = AsyncMock(side_effect=keys)
     mock = _setup_team_keys_mocks(
         monkeypatch,
         find_many=keys,
-        find_unique=AsyncMock(side_effect=keys),
+        find_unique=find_unique,
         update_data=AsyncMock(
             side_effect=[{"data": _updated({"max_budget": 50.0})}] * 3
         ),
@@ -10178,6 +10181,7 @@ async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
         for c in expires_or
         if isinstance(c.get("expires"), dict)
     )
+    find_unique.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -10479,8 +10483,7 @@ async def test_bulk_update_team_keys_does_not_log_raw_sk_token_on_failure(
     _setup_team_keys_mocks(
         monkeypatch,
         find_many=[row],
-        # Force the per-key fetch to raise so the exception logger fires.
-        find_unique=AsyncMock(side_effect=RuntimeError("boom")),
+        update_data=AsyncMock(side_effect=RuntimeError("boom")),
         hash_identity=False,
     )