From 133471f882de4ab62589363facb0b75d80551cf8 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 13 Mar 2026 15:27:56 -0700
Subject: [PATCH] Fix double-counting bug in org/team key limit checks on
 update

When updating a key, _check_org_key_limits and _check_team_key_limits
would include the key being updated in the find_many results, causing
its current limits to be counted twice (once from the DB query, once
from the new requested limits). This caused false 400 errors on valid
limit adjustments.

Fix: exclude the key being updated (by matching token) from the
allocated totals before checking limits.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../key_management_endpoints.py               |  6 ++
 .../test_key_management_endpoints.py          | 60 +++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 2c265c8066d..cfb0da457bb 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -905,6 +905,9 @@ async def _check_team_key_limits(
     keys = await prisma_client.db.litellm_verificationtoken.find_many(
         where={"team_id": team_table.team_id},
     )
+    # Exclude the key being updated to avoid double-counting its limits
+    if isinstance(data, UpdateKeyRequest):
+        keys = [key for key in keys if key.token != data.key]
     check_team_key_model_specific_limits(
         keys=keys,
         team_table=team_table,
@@ -1059,6 +1062,9 @@ async def _check_org_key_limits(
     keys = await prisma_client.db.litellm_verificationtoken.find_many(
         where={"organization_id": org_table.organization_id},
     )
+    # Exclude the key being updated to avoid double-counting its limits
+    if isinstance(data, UpdateKeyRequest):
+        keys = [key for key in keys if key.token != data.key]
     check_org_key_model_specific_limits(
         keys=keys,
         org_table=org_table,
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 fe1ba93699e..1a8dad51d22 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
@@ -6820,6 +6820,7 @@ async def test_check_org_key_limits_on_update_overallocation():
     would exceed organization TPM limits.
     """
     existing_key = MagicMock()
+    existing_key.token = "sk-other-key"
     existing_key.tpm_limit = 15000
     existing_key.rpm_limit = 1500
     existing_key.metadata = {}
@@ -6861,6 +6862,65 @@ async def test_check_org_key_limits_on_update_overallocation():
     assert "TPM limit" in str(exc.value.detail)
 
 
+@pytest.mark.asyncio
+async def test_check_org_key_limits_on_update_excludes_self():
+    """
+    Test that _check_org_key_limits excludes the key being updated from the
+    allocated totals. Without this, the key's current limits would be
+    double-counted: once from find_many and once from data.tpm_limit/rpm_limit.
+    """
+    # The key being updated is returned by find_many with its current limits
+    self_key = MagicMock()
+    self_key.token = "sk-test-key"
+    self_key.tpm_limit = 10000
+    self_key.rpm_limit = 1000
+    self_key.metadata = {}
+
+    # Another key in the org
+    other_key = MagicMock()
+    other_key.token = "sk-other-key"
+    other_key.tpm_limit = 5000
+    other_key.rpm_limit = 500
+    other_key.metadata = {}
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
+        return_value=[self_key, other_key]
+    )
+
+    org_table = LiteLLM_OrganizationTable(
+        organization_id="test-org-self",
+        organization_alias="test-org",
+        budget_id="budget-789",
+        models=["gpt-4"],
+        created_by="admin",
+        updated_by="admin",
+        litellm_budget_table=LiteLLM_BudgetTable(
+            budget_id="budget-789",
+            tpm_limit=20000,
+            rpm_limit=2000,
+        ),
+    )
+
+    # Updating the key to 12000 TPM. Other key uses 5000, so total = 17000 < 20000.
+    # Without the fix, this would be 10000 (self) + 5000 (other) + 12000 = 27000 > 20000.
+    data = UpdateKeyRequest(
+        key="sk-test-key",
+        tpm_limit=12000,
+        rpm_limit=1200,
+        tpm_limit_type="guaranteed_throughput",
+        rpm_limit_type="guaranteed_throughput",
+        organization_id="test-org-self",
+    )
+
+    # Should not raise - the key's own limits should be excluded from the sum
+    await _check_org_key_limits(
+        org_table=org_table,
+        data=data,
+        prisma_client=mock_prisma_client,
+    )
+
+
 def test_update_key_request_has_organization_id():
     """
     Test that UpdateKeyRequest accepts organization_id field.