Merge pull request #38878 from BerriAI/litellm_fix_master_key_rotation_blocked

fix(proxy): preserve model table columns on master key rotation

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -19,13 +19,13 @@ import re
 import secrets
 import traceback
 from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
-from contextlib import AbstractAsyncContextManager
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast
 
 import fastapi
 import yaml
 from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
+from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -155,6 +155,7 @@ from litellm.types.utils import (
 )
 
 if TYPE_CHECKING:
+    import prisma
     from prisma import Prisma
     from prisma import models as prisma_models
 
@@ -182,6 +183,14 @@ class _TxTables(Protocol):
     litellm_proxymodeltable: TableActions[object]
 
 
+class _ModelParamsUpdate(TypedDict):
+    litellm_params: ReadOnly["prisma.Json"]
+
+
+class _ModelRowWhere(TypedDict):
+    model_id: ReadOnly[str]
+
+
 class _ConfigTableActions(Protocol):
     """Config table surface this module needs; the shared repository seam exposes no ``update``."""
 
@@ -273,12 +282,6 @@ def _env_vars_param_value(param: _EnvVarsParam) -> Mapping[str, str] | None:
     return param.param_value
 
 
-def _tx_tables_context(
-    open_tx: Callable[[], AbstractAsyncContextManager[_TxTables]],
-) -> AbstractAsyncContextManager[_TxTables]:
-    return open_tx()
-
-
 async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
     """Raise 403 if custom API keys are disabled and a custom key was provided."""
     if custom_key_value is None:
@@ -4484,27 +4487,29 @@ async def _rotate_master_key(
     if models:
         decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models)
         verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models))
-        new_models: Final[list[dict[str, object]]] = []
-        for model in decrypted_models:
-            new_model = await _add_model_to_db(
-                model_params=Deployment(**model),
-                user_api_key_dict=user_api_key_dict,
-                prisma_client=prisma_client,
-                new_encryption_key=new_master_key,
-                should_create_model_in_db=False,
-            )
-            if new_model:
-                _dumped = dict[str, object](_as_object_dict(new_model.model_dump(exclude_none=True)))
-                _dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"])
-                _dumped["model_info"] = prisma.Json(_dumped["model_info"])
-                new_models.append(_dumped)
-        verbose_proxy_logger.debug("Resetting proxy model table")
-        async with _tx_tables_context(prisma_client.db.tx) as tx:
-            await tx.litellm_proxymodeltable.delete_many()
-            verbose_proxy_logger.debug("Creating %s models", len(new_models))
-            await tx.litellm_proxymodeltable.create_many(
-                data=new_models,
-            )
+        reencrypted_models: Final = tuple(
+            [
+                reencrypted
+                for model in decrypted_models
+                if (
+                    reencrypted := await _add_model_to_db(
+                        model_params=Deployment(**model),
+                        user_api_key_dict=user_api_key_dict,
+                        prisma_client=prisma_client,
+                        new_encryption_key=new_master_key,
+                        should_create_model_in_db=False,
+                    )
+                )
+            ]
+        )
+        verbose_proxy_logger.debug("Re-encrypting litellm_params on %s model rows", len(reencrypted_models))
+        async with prisma_client.db.tx(timeout=timedelta(minutes=2)) as tx_ctx:
+            tx: Final[_TxTables] = tx_ctx
+            for reencrypted_model in reencrypted_models:
+                await tx.litellm_proxymodeltable.update_many(
+                    data=_ModelParamsUpdate(litellm_params=prisma.Json(reencrypted_model.litellm_params)),
+                    where=_ModelRowWhere(model_id=reencrypted_model.model_id),
+                )
         await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
     # 3. process config table
     try:

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -8312,15 +8312,17 @@ async def test_key_does_not_override_explicit_budget_duration():
 @patch(
     "litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key"
 )
-async def test_rotate_master_key_model_data_valid_for_prisma(
+async def test_rotate_master_key_reencrypts_model_params_in_place(
     mock_rotate_mcp,
 ):
     """
-    Test that _rotate_master_key produces valid data for Prisma create_many().
-
-    Regression test for: master key rotation fails with Prisma validation error
-    because created_at/updated_at are None (non-nullable DateTime) and
-    litellm_params/model_info are JSON strings (create_many expects dicts).
+    Regression test for: master key rotation wipes every non-credential column
+    on LiteLLM_ProxyModelTable. Rotation used to rebuild the table via
+    delete_many + create_many from Deployment objects, which carry no
+    blocked/created_at/created_by/updated_at/updated_by, so every rotation
+    reset blocked to False (silently unblocking blocked models) and rewrote the
+    audit columns. Rotation must instead update only litellm_params (the sole
+    encrypted column) on each existing row, keyed by model_id.
     """
     from unittest.mock import AsyncMock, MagicMock
 
@@ -8352,6 +8354,7 @@ async def test_rotate_master_key_model_data_valid_for_prisma(
     mock_tx.litellm_proxymodeltable = MagicMock()
     mock_tx.litellm_proxymodeltable.delete_many = AsyncMock()
     mock_tx.litellm_proxymodeltable.create_many = AsyncMock()
+    mock_tx.litellm_proxymodeltable.update_many = AsyncMock()
     mock_prisma_client.db.tx = MagicMock(
         return_value=AsyncMock(
             __aenter__=AsyncMock(return_value=mock_tx),
@@ -8400,36 +8403,33 @@ async def test_rotate_master_key_model_data_valid_for_prisma(
             new_master_key="sk-new-master-key",
         )
 
-    # Verify create_many was called
-    mock_tx.litellm_proxymodeltable.create_many.assert_called_once()
+    # Rotation must never rewrite whole rows: no delete + recreate
+    mock_tx.litellm_proxymodeltable.delete_many.assert_not_called()
+    mock_tx.litellm_proxymodeltable.create_many.assert_not_called()
 
-    # Get the data passed to create_many
-    call_args = mock_tx.litellm_proxymodeltable.create_many.call_args
-    created_models = call_args.kwargs.get("data") or call_args[1].get("data")
+    mock_tx.litellm_proxymodeltable.update_many.assert_called_once()
+    call_args = mock_tx.litellm_proxymodeltable.update_many.call_args
 
-    assert len(created_models) == 1
-    model_data = created_models[0]
+    assert call_args.kwargs["where"] == {
+        "model_id": "model-1"
+    }, "the re-encrypted params must land on the same row, keyed by model_id"
 
-    # Verify timestamps are NOT present (Prisma @default(now()) should apply)
-    assert (
-        "created_at" not in model_data
-    ), "created_at should be excluded so Prisma @default(now()) applies"
-    assert (
-        "updated_at" not in model_data
-    ), "updated_at should be excluded so Prisma @default(now()) applies"
+    update_data = call_args.kwargs["data"]
+    assert set(update_data.keys()) == {"litellm_params"}, (
+        "rotation must touch only the encrypted litellm_params column; writing any "
+        f"other column wipes it (blocked, audit columns), got {sorted(update_data.keys())}"
+    )
 
-    # Verify litellm_params and model_info are prisma.Json wrappers, NOT JSON strings
     import prisma
 
     assert isinstance(
-        model_data["litellm_params"], prisma.Json
-    ), f"litellm_params should be prisma.Json for create_many(), got {type(model_data['litellm_params'])}"
-    assert isinstance(
-        model_data["model_info"], prisma.Json
-    ), f"model_info should be prisma.Json for create_many(), got {type(model_data['model_info'])}"
-
-    # Verify delete_many was called inside the transaction (before create_many)
-    mock_tx.litellm_proxymodeltable.delete_many.assert_called_once()
+        update_data["litellm_params"], prisma.Json
+    ), f"litellm_params should be prisma.Json for update_many(), got {type(update_data['litellm_params'])}"
+    reencrypted_params = update_data["litellm_params"].data
+    assert set(reencrypted_params.keys()) >= {"model", "api_key"}
+    assert (
+        reencrypted_params["api_key"] != "sk-decrypted-key"
+    ), "api_key must be stored re-encrypted under the new master key, not in plaintext"
 
 
 async def test_default_key_generate_params_duration(monkeypatch):