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>
This commit is contained in:
yuneng-jiang 2026-03-13 15:27:56 -07:00
parent 377c12b917
commit 133471f882
2 changed files with 66 additions and 0 deletions

View file

@ -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,

View file

@ -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.