From 377c12b917652adb98ccaab1ca5558f987756e0c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 00:04:59 -0700 Subject: [PATCH 1/4] Allow setting organization_id on key update endpoint The /key/update endpoint was missing support for organization_id, which was already available on /key/generate. This adds the field to UpdateKeyRequest and validates org key limits during updates. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 1 + .../key_management_endpoints.py | 22 ++++ .../test_key_management_endpoints.py | 109 ++++++++++++++++++ 3 files changed, 132 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d7ccd7d6f1b..06b3f20bcf3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1004,6 +1004,7 @@ class UpdateKeyRequest(KeyRequestBase): temp_budget_expiry: Optional[datetime] = None auto_rotate: Optional[bool] = None rotation_interval: Optional[str] = None + organization_id: Optional[str] = None @model_validator(mode="after") def validate_temp_budget(self) -> "UpdateKeyRequest": diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 654205252b6..2c265c8066d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1806,6 +1806,7 @@ async def update_key_fn( - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key - agent_id: Optional[str] - The agent id associated with the key. + - organization_id: Optional[str] - The organization id of the key. - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - models: Optional[list] - Model_name's a user is allowed to call - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) @@ -1956,6 +1957,27 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) + # Check org key limits if organization_id is being set or already exists on the key + _org_id_to_check = data.organization_id or getattr( + existing_key_row, "organization_id", None + ) + if _org_id_to_check is not None: + org_table = await get_org_object( + org_id=_org_id_to_check, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + if org_table is None: + raise HTTPException( + status_code=400, + detail=f"Organization not found for organization_id={_org_id_to_check}", + ) + await _check_org_key_limits( + org_table=org_table, + data=data, + prisma_client=prisma_client, + ) + # if team change - check if this is possible if is_different_team(data=data, existing_key_row=existing_key_row): if llm_router is None: 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 09dfdb81cbb..fe1ba93699e 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 @@ -6765,3 +6765,112 @@ class TestValidateKeyAliasFormat: _validate_key_alias_format(alias) assert str(exc.value.code) == "400" assert "Invalid key_alias format" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_check_org_key_limits_on_update_within_bounds(): + """ + Test that _check_org_key_limits works with UpdateKeyRequest when updating + a key's TPM/RPM limits within organization bounds. + """ + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + + org_table = LiteLLM_OrganizationTable( + organization_id="test-org-update-1", + organization_alias="test-org", + budget_id="budget-123", + models=["gpt-4"], + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable( + budget_id="budget-123", + tpm_limit=20000, + rpm_limit=2000, + ), + ) + + data = UpdateKeyRequest( + key="sk-test-key", + tpm_limit=10000, + rpm_limit=1000, + tpm_limit_type="guaranteed_throughput", + rpm_limit_type="guaranteed_throughput", + organization_id="test-org-update-1", + ) + + # Should not raise any exception + await _check_org_key_limits( + org_table=org_table, + data=data, + prisma_client=mock_prisma_client, + ) + + mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( + where={"organization_id": "test-org-update-1"} + ) + + +@pytest.mark.asyncio +async def test_check_org_key_limits_on_update_overallocation(): + """ + Test that _check_org_key_limits raises HTTPException when updating a key + would exceed organization TPM limits. + """ + existing_key = MagicMock() + existing_key.tpm_limit = 15000 + existing_key.rpm_limit = 1500 + existing_key.metadata = {} + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key] + ) + + org_table = LiteLLM_OrganizationTable( + organization_id="test-org-update-2", + organization_alias="test-org", + budget_id="budget-456", + models=["gpt-4"], + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable( + budget_id="budget-456", + tpm_limit=20000, + rpm_limit=2000, + ), + ) + + data = UpdateKeyRequest( + key="sk-test-key", + tpm_limit=10000, # 15000 + 10000 = 25000 > 20000 + tpm_limit_type="guaranteed_throughput", + rpm_limit_type="guaranteed_throughput", + organization_id="test-org-update-2", + ) + + with pytest.raises(HTTPException) as exc: + await _check_org_key_limits( + org_table=org_table, + data=data, + prisma_client=mock_prisma_client, + ) + assert exc.value.status_code == 400 + assert "TPM limit" in str(exc.value.detail) + + +def test_update_key_request_has_organization_id(): + """ + Test that UpdateKeyRequest accepts organization_id field. + """ + data = UpdateKeyRequest( + key="sk-test-key", + organization_id="test-org-123", + ) + assert data.organization_id == "test-org-123" + + # Also verify it defaults to None + data_no_org = UpdateKeyRequest(key="sk-test-key") + assert data_no_org.organization_id is None From 133471f882de4ab62589363facb0b75d80551cf8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 15:27:56 -0700 Subject: [PATCH 2/4] 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 --- .../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. From 1038a119ce489a31ba531ceec978ab51d7eecb75 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 15:38:42 -0700 Subject: [PATCH 3/4] Skip org limit check when non-throughput fields are updated MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Only run org validation (get_org_object + _check_org_key_limits) when the update actually touches throughput-related fields (tpm_limit, rpm_limit, or organization_id). Previously, any update to a key belonging to an org would trigger the check, which would fail with a 400 if the org had been deleted — blocking unrelated field changes. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 9 +++-- .../test_key_management_endpoints.py | 36 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index cfb0da457bb..5ef255d0449 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1963,11 +1963,16 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Check org key limits if organization_id is being set or already exists on the key + # Check org key limits only when throughput-related fields or organization_id change _org_id_to_check = data.organization_id or getattr( existing_key_row, "organization_id", None ) - if _org_id_to_check is not None: + _throughput_fields_changed = ( + data.organization_id is not None + or data.tpm_limit is not None + or data.rpm_limit is not None + ) + if _org_id_to_check is not None and _throughput_fields_changed: org_table = await get_org_object( org_id=_org_id_to_check, user_api_key_cache=user_api_key_cache, 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 1a8dad51d22..90bc9138301 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 @@ -6921,6 +6921,42 @@ async def test_check_org_key_limits_on_update_excludes_self(): ) +def test_update_key_skips_org_check_when_no_throughput_fields_changed(): + """ + Test that the org limit check guard condition correctly skips validation + when only non-throughput fields change on a key that belongs to an org. + This prevents blocking updates when the org has been deleted. + """ + # Updating only key_alias — no throughput fields changed + data = UpdateKeyRequest(key="sk-test-key", key_alias="new-alias") + _throughput_fields_changed = ( + data.organization_id is not None + or data.tpm_limit is not None + or data.rpm_limit is not None + ) + assert _throughput_fields_changed is False + + # Updating tpm_limit — throughput field changed + data_with_tpm = UpdateKeyRequest(key="sk-test-key", tpm_limit=5000) + _throughput_fields_changed_tpm = ( + data_with_tpm.organization_id is not None + or data_with_tpm.tpm_limit is not None + or data_with_tpm.rpm_limit is not None + ) + assert _throughput_fields_changed_tpm is True + + # Updating organization_id — org change triggers check + data_with_org = UpdateKeyRequest( + key="sk-test-key", organization_id="new-org" + ) + _throughput_fields_changed_org = ( + data_with_org.organization_id is not None + or data_with_org.tpm_limit is not None + or data_with_org.rpm_limit is not None + ) + assert _throughput_fields_changed_org is True + + def test_update_key_request_has_organization_id(): """ Test that UpdateKeyRequest accepts organization_id field. From 818c097ca9da96a6be30a489612acdc9593b20bc Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 15:59:23 -0700 Subject: [PATCH 4/4] Fix self-exclusion hash mismatch and missing throughput field checks The self-exclusion filter compared raw key strings against SHA-256 hashed tokens from the DB, so keys were never excluded and double-counting persisted. Now hash data.key before comparison. Also add tpm_limit_type/rpm_limit_type to _throughput_fields_changed guard, fall back to existing_key_row.team_id for team limit checks (matching the org pattern), and add team self-exclusion test. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 23 ++-- .../test_key_management_endpoints.py | 112 ++++++++++++++---- 2 files changed, 107 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5ef255d0449..9ec630e66a3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -905,9 +905,11 @@ 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 + # Exclude the key being updated to avoid double-counting its limits. + # key.token is the SHA-256 hash stored in DB; data.key is the raw key string. if isinstance(data, UpdateKeyRequest): - keys = [key for key in keys if key.token != data.key] + hashed_key = hash_token(data.key) + keys = [key for key in keys if key.token != hashed_key] check_team_key_model_specific_limits( keys=keys, team_table=team_table, @@ -1062,9 +1064,11 @@ 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 + # Exclude the key being updated to avoid double-counting its limits. + # key.token is the SHA-256 hash stored in DB; data.key is the raw key string. if isinstance(data, UpdateKeyRequest): - keys = [key for key in keys if key.token != data.key] + hashed_key = hash_token(data.key) + keys = [key for key in keys if key.token != hashed_key] check_org_key_model_specific_limits( keys=keys, org_table=org_table, @@ -1932,11 +1936,14 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Only check team limits if key has a team_id + # Check team limits if key has a team_id (from request or existing key) team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - if data.team_id is not None: + _team_id_to_check = data.team_id or getattr( + existing_key_row, "team_id", None + ) + if _team_id_to_check is not None: team_obj = await get_team_object( - team_id=data.team_id, + team_id=_team_id_to_check, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, check_db_only=True, @@ -1971,6 +1978,8 @@ async def update_key_fn( data.organization_id is not None or data.tpm_limit is not None or data.rpm_limit is not None + or data.tpm_limit_type is not None + or data.rpm_limit_type is not None ) if _org_id_to_check is not None and _throughput_fields_changed: org_table = await get_org_object( 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 90bc9138301..888d89981da 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 @@ -1836,6 +1836,65 @@ async def test_check_team_key_limits_rpm_overallocation(): ) +@pytest.mark.asyncio +async def test_check_team_key_limits_on_update_excludes_self(): + """ + Test that _check_team_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. + """ + from litellm.proxy._types import hash_token as _ht + + # The key being updated is returned by find_many with its current limits. + # In the DB, token is stored as a SHA-256 hash of the raw key. + self_key = MagicMock() + self_key.token = _ht("sk-self-team-key") + self_key.tpm_limit = 6000 + self_key.rpm_limit = 600 + self_key.metadata = {} + + # Another key in the team + other_key = MagicMock() + other_key.token = _ht("sk-other-team-key") + other_key.tpm_limit = 3000 + other_key.rpm_limit = 300 + other_key.metadata = {} + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[self_key, other_key] + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-self", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Updating the key to 7000 TPM. Other key uses 3000, so total = 10000 <= 10000. + # Without the fix, this would be 6000 (self) + 3000 (other) + 7000 = 16000 > 10000. + data = UpdateKeyRequest( + key="sk-self-team-key", + tpm_limit=7000, + rpm_limit=700, + tpm_limit_type="guaranteed_throughput", + rpm_limit_type="guaranteed_throughput", + ) + + # Should not raise - the key's own limits should be excluded from the sum + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + @pytest.mark.asyncio async def test_check_team_key_limits_no_team_limits(): """ @@ -6819,8 +6878,10 @@ async def test_check_org_key_limits_on_update_overallocation(): Test that _check_org_key_limits raises HTTPException when updating a key would exceed organization TPM limits. """ + from litellm.proxy._types import hash_token as _hash_token + existing_key = MagicMock() - existing_key.token = "sk-other-key" + existing_key.token = _hash_token("sk-other-key") existing_key.tpm_limit = 15000 existing_key.rpm_limit = 1500 existing_key.metadata = {} @@ -6869,16 +6930,19 @@ async def test_check_org_key_limits_on_update_excludes_self(): 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 + from litellm.proxy._types import hash_token + + # The key being updated is returned by find_many with its current limits. + # In the DB, token is stored as a SHA-256 hash of the raw key. self_key = MagicMock() - self_key.token = "sk-test-key" + self_key.token = hash_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.token = hash_token("sk-other-key") other_key.tpm_limit = 5000 other_key.rpm_limit = 500 other_key.metadata = {} @@ -6927,34 +6991,40 @@ def test_update_key_skips_org_check_when_no_throughput_fields_changed(): when only non-throughput fields change on a key that belongs to an org. This prevents blocking updates when the org has been deleted. """ + def _check_throughput_changed(data: UpdateKeyRequest) -> bool: + return ( + data.organization_id is not None + or data.tpm_limit is not None + or data.rpm_limit is not None + or data.tpm_limit_type is not None + or data.rpm_limit_type is not None + ) + # Updating only key_alias — no throughput fields changed data = UpdateKeyRequest(key="sk-test-key", key_alias="new-alias") - _throughput_fields_changed = ( - data.organization_id is not None - or data.tpm_limit is not None - or data.rpm_limit is not None - ) - assert _throughput_fields_changed is False + assert _check_throughput_changed(data) is False # Updating tpm_limit — throughput field changed data_with_tpm = UpdateKeyRequest(key="sk-test-key", tpm_limit=5000) - _throughput_fields_changed_tpm = ( - data_with_tpm.organization_id is not None - or data_with_tpm.tpm_limit is not None - or data_with_tpm.rpm_limit is not None - ) - assert _throughput_fields_changed_tpm is True + assert _check_throughput_changed(data_with_tpm) is True # Updating organization_id — org change triggers check data_with_org = UpdateKeyRequest( key="sk-test-key", organization_id="new-org" ) - _throughput_fields_changed_org = ( - data_with_org.organization_id is not None - or data_with_org.tpm_limit is not None - or data_with_org.rpm_limit is not None + assert _check_throughput_changed(data_with_org) is True + + # Updating tpm_limit_type — limit type change triggers check + data_with_tpm_type = UpdateKeyRequest( + key="sk-test-key", tpm_limit_type="guaranteed_throughput" ) - assert _throughput_fields_changed_org is True + assert _check_throughput_changed(data_with_tpm_type) is True + + # Updating rpm_limit_type — limit type change triggers check + data_with_rpm_type = UpdateKeyRequest( + key="sk-test-key", rpm_limit_type="guaranteed_throughput" + ) + assert _check_throughput_changed(data_with_rpm_type) is True def test_update_key_request_has_organization_id():