fix: address remaining policy engine reliability issues

1. Fix race condition in update_version_status promote-to-production:
   - Wrap demote (update_many) + promote (update) in a Prisma transaction
   - Prevents concurrent promote calls from leaving two versions as production

2. Fix stale _policies_by_id cache in delete methods:
   - delete_policy_from_db: remove entry from _policies_by_id when deleting
     a draft/published version
   - delete_all_versions: scan and remove all _policies_by_id entries
     matching the deleted policy_name

3. Fix docstring for update_version_status:
   - Remove documented-but-unimplemented 'production -> published' transition
   - Clarify which transitions are valid vs invalid

4. Add tests for new behavior:
   - test_published_to_production_removes_from_policies_by_id
   - test_delete_draft_removes_from_policies_by_id_cache
   - test_delete_all_versions_cleans_policies_by_id_cache
   - Update existing promote-to-production tests to mock transaction

Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>

litellm/proxy/policy_engine/policy_registry.py

@@ -427,13 +427,14 @@ class PolicyRegistry:
                 "message": f"Policy {policy_id} deleted successfully"
             }
 
-            # Remove from in-memory registry only if this was the production version
             if version_status == "production":
                 self.remove_policy(policy_name)
                 result["warning"] = (
                     "Production version was deleted. No other version was promoted. "
                     "Promote another version to production if this policy should remain active."
                 )
+            else:
+                self._policies_by_id.pop(policy_id, None)
 
             return result
         except Exception as e:
@@ -784,10 +785,13 @@ class PolicyRegistry:
         """
         Update a policy version's status. Valid transitions:
         - draft -> published (sets published_at)
-        - published -> production (sets production_at, demotes current production to published, updates in-memory)
-        - production -> published (demotes, removes from in-memory)
+        - published -> production (sets production_at, demotes current production
+          to published, updates in-memory registry)
+
+        Invalid transitions:
         - draft -> production: NOT allowed (must publish first)
         - published -> draft: NOT allowed
+        - production -> published: NOT allowed (use delete or create a new version)
 
         Args:
             policy_id: The policy version ID
@@ -842,29 +846,30 @@ class PolicyRegistry:
                     "Cannot promote draft directly to production. Publish the version first."
                 )
 
-            # Demote current production to published
-            await prisma_client.db.litellm_policytable.update_many(
-                where={
-                    "policy_name": policy_name,
-                    "version_status": "production",
-                },
-                data={
-                    "version_status": "published",
-                    "updated_at": now,
-                    "updated_by": updated_by,
-                },
-            )
+            async with prisma_client.db.tx() as tx:
+                # Demote current production to published
+                await tx.litellm_policytable.update_many(
+                    where={
+                        "policy_name": policy_name,
+                        "version_status": "production",
+                    },
+                    data={
+                        "version_status": "published",
+                        "updated_at": now,
+                        "updated_by": updated_by,
+                    },
+                )
 
-            # Promote this version to production
-            updated = await prisma_client.db.litellm_policytable.update(
-                where={"policy_id": policy_id},
-                data={
-                    "version_status": "production",
-                    "production_at": now,
-                    "updated_at": now,
-                    "updated_by": updated_by,
-                },
-            )
+                # Promote this version to production
+                updated = await tx.litellm_policytable.update(
+                    where={"policy_id": policy_id},
+                    data={
+                        "version_status": "production",
+                        "production_at": now,
+                        "updated_at": now,
+                        "updated_by": updated_by,
+                    },
+                )
 
             # Update in-memory registry: remove old production (by name), add this one
             self.remove_policy(policy_name)
@@ -968,6 +973,13 @@ class PolicyRegistry:
                 where={"policy_name": policy_name}
             )
             self.remove_policy(policy_name)
+            stale_ids = [
+                pid
+                for pid, (pname, _) in self._policies_by_id.items()
+                if pname == policy_name
+            ]
+            for pid in stale_ids:
+                del self._policies_by_id[pid]
             return {
                 "message": f"All versions of policy '{policy_name}' deleted successfully"
             }

tests/test_litellm/proxy/policy_engine/test_policy_versioning.py

@@ -241,6 +241,27 @@ class TestDeletePolicyFromDb:
         assert "warning" not in result
         assert registry.has_policy("my-policy")
 
+    @pytest.mark.asyncio
+    async def test_delete_draft_removes_from_policies_by_id_cache(self):
+        registry = PolicyRegistry()
+        dummy_policy = MagicMock()
+        registry._policies_by_id["draft-1"] = ("my-policy", dummy_policy)
+        prisma = MagicMock()
+        draft_row = _make_row(
+            policy_id="draft-1",
+            policy_name="my-policy",
+            version_status="draft",
+        )
+        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft_row)
+        prisma.db.litellm_policytable.delete = AsyncMock()
+
+        await registry.delete_policy_from_db(
+            policy_id="draft-1",
+            prisma_client=prisma,
+        )
+
+        assert "draft-1" not in registry._policies_by_id
+
 
 class TestCreateNewVersion:
     """Test create_new_version copies all fields and sets draft."""
@@ -362,8 +383,14 @@ class TestUpdateVersionStatus:
             production_at=datetime.now(timezone.utc),
         )
         prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=published_row)
-        prisma.db.litellm_policytable.update_many = AsyncMock()
-        prisma.db.litellm_policytable.update = AsyncMock(return_value=updated_row)
+
+        tx_mock = MagicMock()
+        tx_mock.litellm_policytable.update_many = AsyncMock()
+        tx_mock.litellm_policytable.update = AsyncMock(return_value=updated_row)
+        ctx = MagicMock()
+        ctx.__aenter__ = AsyncMock(return_value=tx_mock)
+        ctx.__aexit__ = AsyncMock(return_value=False)
+        prisma.db.tx = MagicMock(return_value=ctx)
 
         result = await registry.update_version_status(
             policy_id="pub-1",
@@ -372,11 +399,72 @@ class TestUpdateVersionStatus:
         )
 
         assert result.version_status == "production"
-        # update_many should have been called to demote current production
-        assert prisma.db.litellm_policytable.update_many.called
-        # Registry should have been updated with new production
+        assert tx_mock.litellm_policytable.update_many.called
         assert registry.has_policy("foo")
 
+    @pytest.mark.asyncio
+    async def test_published_to_production_removes_from_policies_by_id(self):
+        registry = PolicyRegistry()
+        dummy_policy = MagicMock()
+        registry._policies_by_id["pub-1"] = ("foo", dummy_policy)
+
+        prisma = MagicMock()
+        published_row = _make_row(
+            policy_id="pub-1",
+            policy_name="foo",
+            version_status="published",
+        )
+        updated_row = _make_row(
+            policy_id="pub-1",
+            policy_name="foo",
+            version_status="production",
+            production_at=datetime.now(timezone.utc),
+        )
+        prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=published_row)
+
+        tx_mock = MagicMock()
+        tx_mock.litellm_policytable.update_many = AsyncMock()
+        tx_mock.litellm_policytable.update = AsyncMock(return_value=updated_row)
+        ctx = MagicMock()
+        ctx.__aenter__ = AsyncMock(return_value=tx_mock)
+        ctx.__aexit__ = AsyncMock(return_value=False)
+        prisma.db.tx = MagicMock(return_value=ctx)
+
+        await registry.update_version_status(
+            policy_id="pub-1",
+            new_status="production",
+            prisma_client=prisma,
+        )
+
+        assert "pub-1" not in registry._policies_by_id
+
+
+class TestDeleteAllVersions:
+    """Test delete_all_versions cleans up all caches."""
+
+    @pytest.mark.asyncio
+    async def test_delete_all_versions_cleans_policies_by_id_cache(self):
+        registry = PolicyRegistry()
+        registry.add_policy("my-policy", MagicMock())
+        dummy = MagicMock()
+        registry._policies_by_id["draft-1"] = ("my-policy", dummy)
+        registry._policies_by_id["pub-1"] = ("my-policy", dummy)
+        registry._policies_by_id["other-1"] = ("other-policy", dummy)
+
+        prisma = MagicMock()
+        prisma.db.litellm_policytable.delete_many = AsyncMock()
+
+        result = await registry.delete_all_versions(
+            policy_name="my-policy",
+            prisma_client=prisma,
+        )
+
+        assert "deleted successfully" in result["message"]
+        assert not registry.has_policy("my-policy")
+        assert "draft-1" not in registry._policies_by_id
+        assert "pub-1" not in registry._policies_by_id
+        assert "other-1" in registry._policies_by_id
+
 
 class TestCompareVersions:
     """Test compare_versions returns correct field diffs."""

tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py

@@ -171,7 +171,6 @@ async def test_full_lifecycle_create_draft_edit_publish_promote():
 
     # 5) Promote v2 to production (demote v1 to published, update registry)
     prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=v2_published)
-    prisma.db.litellm_policytable.update_many = AsyncMock()
     v2_production = _make_row(
         policy_id="v2-id",
         policy_name="lifecycle-policy",
@@ -181,7 +180,14 @@ async def test_full_lifecycle_create_draft_edit_publish_promote():
         guardrails_add=["g1", "g2", "g3"],
         description="Draft v2 edited",
     )
-    prisma.db.litellm_policytable.update = AsyncMock(return_value=v2_production)
+
+    promote_tx = MagicMock()
+    promote_tx.litellm_policytable.update_many = AsyncMock()
+    promote_tx.litellm_policytable.update = AsyncMock(return_value=v2_production)
+    promote_ctx = MagicMock()
+    promote_ctx.__aenter__ = AsyncMock(return_value=promote_tx)
+    promote_ctx.__aexit__ = AsyncMock(return_value=False)
+    prisma.db.tx = MagicMock(return_value=promote_ctx)
 
     prod = await registry.update_version_status(
         policy_id="v2-id",