fix(proxy): resolve rotation and permission fields before the key policy and pin the effective-row contract

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -359,7 +359,9 @@ def _effective_key_after_update(
     overlay: Final = MappingProxyType(
         {column: _decode_json_string_column(column, value) for column, value in non_default_values.items()}
     )
-    return _verification_token_from_row(MappingProxyType({**existing_key_row.model_dump(), **overlay}))
+    return _verification_token_from_row(
+        MappingProxyType({**existing_key_row.model_dump(), **overlay, "object_permission": None})
+    )
 
 
 def _update_policy_request(
@@ -381,7 +383,7 @@ def _update_policy_request(
 def _generate_budget_windows(
     budget_limits: Sequence[BudgetLimitEntry] | None,
 ) -> tuple[Mapping[str, object], ...] | None:
-    if budget_limits is None:
+    if not budget_limits:
         return None
     return tuple(
         MappingProxyType(
@@ -402,15 +404,20 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite
     column_fields: Final = MappingProxyType(
         {field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS}
     )
-    request_metadata: Final = data.metadata or MappingProxyType({})
-    folded_metadata: Final = {**request_metadata, **metadata_fields}  # mutable-ok: encrypt_callback_vars needs a dict
-    columns: Final = handle_key_type(data, {**column_fields})  # mutable-ok: handle_key_type mutates data_json in place
+    metadata: Final = data.metadata or MappingProxyType({})
+    folded_metadata: Final = {**metadata, **metadata_fields}  # mutable-ok: encrypt_callback_vars needs a dict
+    columns: Final = handle_key_type(data, {**column_fields})  # mutable-ok: handle_key_type mutates in place
     expires: Final = (
         now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None
     )
     budget_reset_at: Final = (
         get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None
     )
+    key_rotation_at: Final = (
+        now + timedelta(seconds=duration_in_seconds(duration=data.rotation_interval))
+        if data.auto_rotate and data.rotation_interval
+        else None
+    )
     return _verification_token_from_row(
         MappingProxyType(
             {
@@ -418,6 +425,7 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite
                 "metadata": encrypt_callback_vars(folded_metadata),
                 "expires": expires,
                 "budget_reset_at": budget_reset_at,
+                "key_rotation_at": key_rotation_at,
                 "budget_limits": _generate_budget_windows(data.budget_limits),
                 "object_permission": None,
             }
@@ -3249,16 +3257,6 @@ async def update_key_fn(
         _enforce_upperbound_key_params(data, fill_defaults=False)
         non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row)
 
-        await _enforce_custom_key_policy(
-            hook=_custom_key_policy_hook(proxy_server),
-            build_policy_request=lambda: _update_policy_request(
-                operation="update",
-                existing_key_row=existing_key_row,
-                non_default_values=non_default_values,
-                request=data,
-            ),
-        )
-
         # Only validate key_alias format if it's actually being changed
         new_key_alias: Final = non_default_values.get("key_alias", None)
         if new_key_alias != existing_key_row.key_alias:
@@ -3278,6 +3276,16 @@ async def update_key_fn(
             existing_key_alias=existing_key_row.key_alias,
         )
 
+        await _enforce_custom_key_policy(
+            hook=_custom_key_policy_hook(proxy_server),
+            build_policy_request=lambda: _update_policy_request(
+                operation="update",
+                existing_key_row=existing_key_row,
+                non_default_values=non_default_values,
+                request=data,
+            ),
+        )
+
         if prisma_client is None:
             raise Exception("Not connected to DB!")
 

litellm/types/proxy/management_endpoints/key_management_endpoints.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any, Final, Literal
+from typing import Any, Final, Literal, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, model_validator
 from typing_extensions import ReadOnly, TypedDict
@@ -128,11 +128,18 @@ class BulkUpdateTeamKeysRequest(BaseModel):
         return self
 
 
-CustomKeyPolicyOperation = Literal["generate", "update", "regenerate"]
+CustomKeyPolicyOperation: TypeAlias = Literal["generate", "update", "regenerate"]
 
 
 class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase):
-    """What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request."""
+    """What `general_settings.custom_key_policy` receives.
+
+    `effective_key` is the verification token row as it will be written: the existing row overlaid with the
+    requested changes, with `duration` resolved to `expires` and `budget_duration` to `budget_reset_at`. Values the
+    proxy fills in after the policy stay at their defaults: `token`, `key_name`, `created_by`, `updated_by` and the
+    soft-budget `budget_id` on generate, the rotated token on regenerate, and the `object_permission` relation on
+    every operation (`object_permission_id` is set; read `request.object_permission` for the requested change).
+    """
 
     model_config = ConfigDict(protected_namespaces=(), frozen=True)
 

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -31,6 +31,7 @@ from litellm.proxy._types import (
     ResetSpendRequest,
     UpdateKeyRequest,
 )
+from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
 from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@@ -12362,7 +12363,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp
     mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key())
     received: list[CustomKeyPolicyRequest] = []
     policy = _seven_day_policy(received)
-    data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0)
+    data = UpdateKeyRequest(
+        key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0, auto_rotate=True, rotation_interval="30d"
+    )
 
     with (
         patch(  # test-quality-ok: cache eviction is outside the policy path
@@ -12381,6 +12384,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp
     mock_prisma_client.update_data.assert_awaited_once()
     assert len(received) == 1
     _assert_update_policy_request(received[0], data)
+    key_rotation_at = received[0].effective_key.key_rotation_at
+    assert key_rotation_at is not None
+    assert abs(key_rotation_at - (datetime.now(timezone.utc) + timedelta(days=30))) < timedelta(seconds=60)
 
 
 @pytest.mark.asyncio
@@ -12695,6 +12701,23 @@ async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration
     assert effective_key.expires is None
 
 
+def test_effective_key_after_update_swaps_the_object_permission_id_and_drops_the_stale_relation():
+    existing_key = LiteLLM_VerificationToken(
+        token="tok",
+        object_permission_id="op-old",
+        object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-old", mcp_servers=["old"]),
+    )
+
+    effective_key = _effective_key_after_update(
+        existing_key_row=existing_key, non_default_values={"object_permission_id": "op-new"}
+    )
+
+    assert effective_key.object_permission_id == "op-new"
+    assert effective_key.object_permission is None
+    assert existing_key.object_permission is not None
+    assert existing_key.object_permission.mcp_servers == ["old"]
+
+
 def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it():
     now = datetime(2026, 1, 1, tzinfo=timezone.utc)
     data = GenerateKeyRequest(
@@ -12705,12 +12728,21 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat
         tags=["t1"],
         budget_duration="1d",
         max_budget=3.0,
+        budget_limits=[{"budget_duration": "1d", "max_budget": 5.0}],
+        auto_rotate=True,
+        rotation_interval="30d",
+        object_permission={"mcp_servers": ["srv"]},
         key_type=LiteLLMKeyType.LLM_API,
     )
 
     effective_key = _effective_key_for_generate(data=data, now=now)
 
     assert effective_key.expires == now + timedelta(days=5)
+    assert effective_key.key_rotation_at == now + timedelta(days=30)
+    assert effective_key.budget_limits is not None
+    assert effective_key.budget_limits[0]["max_budget"] == 5.0
+    assert effective_key.budget_limits[0]["reset_at"] is not None
+    assert effective_key.object_permission is None
     assert effective_key.org_id == "org-1"
     assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]}
     assert effective_key.max_budget == 3.0
@@ -12722,6 +12754,18 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat
     assert data.guardrails == ["g1"]
     assert data.tags == ["t1"]
     assert data.duration == "5d"
+    assert data.budget_limits is not None
+    assert data.budget_limits[0].reset_at is None
+    assert data.object_permission is not None
+    assert data.object_permission.mcp_servers == ["srv"]
+
+
+def test_effective_key_for_generate_stores_no_budget_windows_for_an_empty_list():
+    effective_key = _effective_key_for_generate(
+        data=GenerateKeyRequest(budget_limits=[]), now=datetime(2026, 1, 1, tzinfo=timezone.utc)
+    )
+
+    assert effective_key.budget_limits is None
 
 
 def test_effective_key_for_generate_without_duration_never_expires():
@@ -12731,6 +12775,7 @@ def test_effective_key_for_generate_without_duration_never_expires():
 
     assert effective_key.expires is None
     assert effective_key.budget_reset_at is None
+    assert effective_key.key_rotation_at is None
     assert effective_key.key_type == "default"
 
 
@@ -12764,6 +12809,27 @@ async def test_enforce_custom_key_policy_uses_the_default_denial_message():
     assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule"
 
 
+@pytest.mark.asyncio
+async def test_enforce_custom_key_policy_allows_when_the_decision_is_missing():
+    received: list[CustomKeyPolicyRequest] = []
+
+    async def no_decision(policy_request: CustomKeyPolicyRequest) -> dict[str, object]:
+        received.append(policy_request)
+        return {}
+
+    await _enforce_custom_key_policy(hook=no_decision, build_policy_request=_policy_request_for_generate)
+
+    assert len(received) == 1
+    assert received[0].operation == "generate"
+
+
+@pytest.mark.asyncio
+async def test_enforce_custom_key_policy_never_builds_the_request_without_a_hook():
+    await _enforce_custom_key_policy(
+        hook=None, build_policy_request=lambda: pytest.fail("policy request built without a hook")
+    )
+
+
 @pytest.mark.asyncio
 async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token():
     """