mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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:
parent
377c12b917
commit
133471f882
2 changed files with 66 additions and 0 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue