From c1e1773a27abb4535ce78a9bf6e9dc04d1a46e92 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 22 Feb 2026 03:35:34 +0000 Subject: [PATCH] fix(policy-registry): keep version cache consistent + txn create Co-authored-by: Krish Dholakia --- .../proxy/policy_engine/policy_registry.py | 179 +++++++++++------- 1 file changed, 114 insertions(+), 65 deletions(-) diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index f8e1ebd7ba1..839c2e52fee 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -682,77 +682,78 @@ class PolicyRegistry: PolicyDBResponse for the new draft version """ try: - if source_policy_id is not None: - source = await prisma_client.db.litellm_policytable.find_unique( - where={"policy_id": source_policy_id} - ) - if source is None: - raise Exception(f"Source policy {source_policy_id} not found") - if source.policy_name != policy_name: - raise Exception( - f"Source policy name '{source.policy_name}' does not match '{policy_name}'" + async with prisma_client.db.tx() as tx: + if source_policy_id is not None: + source = await tx.litellm_policytable.find_unique( + where={"policy_id": source_policy_id} ) - else: - # Find current production version for this policy_name - prod = await prisma_client.db.litellm_policytable.find_first( - where={ - "policy_name": policy_name, - "version_status": "production", - } - ) - if prod is None: - raise Exception( - f"No production version found for policy '{policy_name}'" + if source is None: + raise Exception(f"Source policy {source_policy_id} not found") + if source.policy_name != policy_name: + raise Exception( + f"Source policy name '{source.policy_name}' does not match '{policy_name}'" + ) + else: + # Find current production version for this policy_name + prod = await tx.litellm_policytable.find_first( + where={ + "policy_name": policy_name, + "version_status": "production", + } ) - source = prod + if prod is None: + raise Exception( + f"No production version found for policy '{policy_name}'" + ) + source = prod - # Next version number - latest = await prisma_client.db.litellm_policytable.find_first( - where={"policy_name": policy_name}, - order={"version_number": "desc"}, - ) - next_num = (latest.version_number + 1) if latest else 1 - - now = datetime.now(timezone.utc) - # Set is_latest=False on all existing versions for this policy_name - await prisma_client.db.litellm_policytable.update_many( - where={"policy_name": policy_name}, - data={"is_latest": False}, - ) - - data: Dict[str, Any] = { - "policy_name": policy_name, - "version_number": next_num, - "version_status": "draft", - "parent_version_id": source.policy_id, - "is_latest": True, - "published_at": None, - "production_at": None, - "inherit": source.inherit, - "description": source.description, - "guardrails_add": source.guardrails_add or [], - "guardrails_remove": source.guardrails_remove or [], - "created_at": now, - "updated_at": now, - "created_by": created_by, - "updated_by": created_by, - } - # Prisma expects Json fields as JSON strings on create (same as add_policy_to_db) - if source.condition is not None: - data["condition"] = ( - json.dumps(source.condition) - if isinstance(source.condition, dict) - else source.condition + # Next version number + latest = await tx.litellm_policytable.find_first( + where={"policy_name": policy_name}, + order={"version_number": "desc"}, ) - if source.pipeline is not None: - data["pipeline"] = ( - json.dumps(source.pipeline) - if isinstance(source.pipeline, dict) - else source.pipeline + next_num = (latest.version_number + 1) if latest else 1 + + now = datetime.now(timezone.utc) + # Set is_latest=False on all existing versions for this policy_name + await tx.litellm_policytable.update_many( + where={"policy_name": policy_name}, + data={"is_latest": False}, ) - created = await prisma_client.db.litellm_policytable.create(data=data) - return _row_to_policy_db_response(created) + data: Dict[str, Any] = { + "policy_name": policy_name, + "version_number": next_num, + "version_status": "draft", + "parent_version_id": source.policy_id, + "is_latest": True, + "published_at": None, + "production_at": None, + "inherit": source.inherit, + "description": source.description, + "guardrails_add": source.guardrails_add or [], + "guardrails_remove": source.guardrails_remove or [], + "created_at": now, + "updated_at": now, + "created_by": created_by, + "updated_by": created_by, + } + # Prisma expects Json fields as JSON strings on create (same as add_policy_to_db) + if source.condition is not None: + data["condition"] = ( + json.dumps(source.condition) + if isinstance(source.condition, dict) + else source.condition + ) + if source.pipeline is not None: + data["pipeline"] = ( + json.dumps(source.pipeline) + if isinstance(source.pipeline, dict) + else source.pipeline + ) + + created = await tx.litellm_policytable.create(data=data) + return _row_to_policy_db_response(created) except Exception as e: verbose_proxy_logger.exception(f"Error creating new version: {e}") raise Exception(f"Error creating new version: {str(e)}") @@ -811,6 +812,23 @@ class PolicyRegistry: "updated_by": updated_by, }, ) + # Keep request-body policy_ cache fresh after transition. + self._policies_by_id[policy_id] = ( + policy_name, + self._parse_policy( + policy_name, + { + "inherit": updated.inherit, + "description": updated.description, + "guardrails": { + "add": updated.guardrails_add or [], + "remove": updated.guardrails_remove or [], + }, + "condition": updated.condition, + "pipeline": updated.pipeline, + }, + ), + ) return _row_to_policy_db_response(updated) # new_status == "production" @@ -824,6 +842,15 @@ class PolicyRegistry: "Cannot promote draft directly to production. Publish the version first." ) + # Capture current production row (if any) so we can add it to the non-production cache + # after demotion. + old_prod = await prisma_client.db.litellm_policytable.find_first( + where={ + "policy_name": policy_name, + "version_status": "production", + } + ) + # Demote current production to published await prisma_client.db.litellm_policytable.update_many( where={ @@ -848,6 +875,28 @@ class PolicyRegistry: }, ) + # Update request-body policy_ cache: + # - Old production is now published -> add to _policies_by_id + # - Newly promoted production should not be served from _policies_by_id + if old_prod is not None and old_prod.policy_id != policy_id: + self._policies_by_id[old_prod.policy_id] = ( + old_prod.policy_name, + self._parse_policy( + old_prod.policy_name, + { + "inherit": old_prod.inherit, + "description": old_prod.description, + "guardrails": { + "add": old_prod.guardrails_add or [], + "remove": old_prod.guardrails_remove or [], + }, + "condition": old_prod.condition, + "pipeline": old_prod.pipeline, + }, + ), + ) + self._policies_by_id.pop(policy_id, None) + # Update in-memory registry: remove old production (by name), add this one self.remove_policy(policy_name) policy = self._parse_policy(