diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 7a5a557c0e4..16cc8bd102b 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/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" } diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index 1f3a02a4c97..b1b70059feb 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/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.""" diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py index f611d569fb7..5e473d87314 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py +++ b/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",