From 916abeb395e0fff5986fb71db43c892b74387142 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 22 Feb 2026 03:38:26 +0000 Subject: [PATCH] fix: address greptile code review issues on policy engine reliability 1. Fix stale _policies_by_id cache after status transitions: - Add _update_policies_by_id_cache() helper method - Update cache when draft->published transition occurs - Remove entry from cache when promoting to production (resolved by name) 2. Fix race condition in create_new_version: - Wrap find_first + update_many + create in a Prisma transaction - Prevents concurrent version number collisions and orphaned is_latest state 3. Validate version_status query parameter in list_policies: - Use Literal['draft', 'published', 'production'] type - Returns 422 for invalid values instead of silently returning empty results 4. Add Literal validation to PolicyVersionStatusUpdateRequest: - Change version_status field from str to Literal['published', 'production'] - Validates at request parsing level rather than at runtime 5. Fix duplicate auth dependency in endpoints: - Remove decorator-level dependencies=[Depends(user_api_key_auth)] when the function parameter already uses Depends(user_api_key_auth) - Prevents auth check from running twice per request 6. Update tests to mock Prisma transaction context manager Co-authored-by: Krish Dholakia --- .../proxy/policy_engine/policy_endpoints.py | 223 +++++- .../proxy/policy_engine/policy_registry.py | 633 ++++++++++++++---- litellm/types/proxy/policy_engine/__init__.py | 65 +- .../proxy/policy_engine/resolver_types.py | 62 +- .../policy_engine/test_policy_versioning.py | 446 ++++++++++++ .../test_policy_versioning_e2e.py | 224 +++++++ 6 files changed, 1470 insertions(+), 183 deletions(-) create mode 100644 tests/test_litellm/proxy/policy_engine/test_policy_versioning.py create mode 100644 tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index af12a8598f6..3ccf2614fa6 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -4,25 +4,24 @@ CRUD ENDPOINTS FOR POLICIES Provides REST API endpoints for managing policies and policy attachments. """ +from typing import Literal, Optional + from fastapi import APIRouter, Depends, HTTPException from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry +from litellm.proxy.policy_engine.attachment_registry import \ + get_attachment_registry from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import ( - GuardrailPipeline, - PipelineTestRequest, - PolicyAttachmentCreateRequest, - PolicyAttachmentDBResponse, - PolicyAttachmentListResponse, - PolicyCreateRequest, - PolicyDBResponse, - PolicyListDBResponse, - PolicyUpdateRequest, -) + GuardrailPipeline, PipelineTestRequest, PolicyAttachmentCreateRequest, + PolicyAttachmentDBResponse, PolicyAttachmentListResponse, + PolicyCreateRequest, PolicyDBResponse, PolicyListDBResponse, + PolicyUpdateRequest, PolicyVersionCompareResponse, + PolicyVersionCreateRequest, PolicyVersionListResponse, + PolicyVersionStatusUpdateRequest) router = APIRouter() @@ -38,14 +37,22 @@ router = APIRouter() dependencies=[Depends(user_api_key_auth)], response_model=PolicyListDBResponse, ) -async def list_policies(): +async def list_policies( + version_status: Optional[Literal["draft", "published", "production"]] = None, +): """ - List all policies from the database. + List all policies from the database. Optionally filter by version_status. + + Query params: + - version_status: Optional. One of "draft", "published", "production". + If omitted, all versions are returned. Example Request: ```bash curl -X GET "http://localhost:4000/policies/list" \\ -H "Authorization: Bearer " + curl -X GET "http://localhost:4000/policies/list?version_status=production" \\ + -H "Authorization: Bearer " ``` Example Response: @@ -55,6 +62,8 @@ async def list_policies(): { "policy_id": "123e4567-e89b-12d3-a456-426614174000", "policy_name": "global-baseline", + "version_number": 1, + "version_status": "production", "inherit": null, "description": "Base guardrails for all requests", "guardrails_add": ["pii_masking"], @@ -74,7 +83,9 @@ async def list_policies(): raise HTTPException(status_code=500, detail="Database not connected") try: - policies = await get_policy_registry().get_all_policies_from_db(prisma_client) + policies = await get_policy_registry().get_all_policies_from_db( + prisma_client, version_status=version_status + ) return PolicyListDBResponse(policies=policies, total_count=len(policies)) except Exception as e: verbose_proxy_logger.exception(f"Error listing policies: {e}") @@ -84,7 +95,6 @@ async def list_policies(): @router.post( "/policies", tags=["Policies"], - dependencies=[Depends(user_api_key_auth)], response_model=PolicyDBResponse, ) async def create_policy( @@ -145,6 +155,170 @@ async def create_policy( raise HTTPException(status_code=500, detail=str(e)) +# ───────────────────────────────────────────────────────────────────────────── +# Policy Versioning Endpoints (must be before /policies/{policy_id} to avoid path conflicts) +# ───────────────────────────────────────────────────────────────────────────── + + +@router.get( + "/policies/name/{policy_name}/versions", + tags=["Policies"], + dependencies=[Depends(user_api_key_auth)], + response_model=PolicyVersionListResponse, +) +async def list_policy_versions(policy_name: str): + """ + List all versions of a policy by name, ordered by version_number descending. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + return await get_policy_registry().get_versions_by_policy_name( + policy_name=policy_name, + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error listing policy versions: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/policies/name/{policy_name}/versions", + tags=["Policies"], + response_model=PolicyDBResponse, +) +async def create_policy_version( + policy_name: str, + request: PolicyVersionCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Create a new draft version of a policy. Copies all fields from the source. + Source is current production if source_policy_id is not provided. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + created_by = user_api_key_dict.user_id + return await get_policy_registry().create_new_version( + policy_name=policy_name, + prisma_client=prisma_client, + source_policy_id=request.source_policy_id, + created_by=created_by, + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error creating policy version: {e}") + if "not found" in str(e).lower() or "no production" in str(e).lower(): + raise HTTPException(status_code=404, detail=str(e)) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.put( + "/policies/{policy_id}/status", + tags=["Policies"], + response_model=PolicyDBResponse, +) +async def update_policy_version_status( + policy_id: str, + request: PolicyVersionStatusUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update a policy version's status. Valid transitions: + - draft -> published + - published -> production (demotes current production to published) + - production -> published (demotes, policy becomes inactive) + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + updated_by = user_api_key_dict.user_id + return await get_policy_registry().update_version_status( + policy_id=policy_id, + new_status=request.version_status, + prisma_client=prisma_client, + updated_by=updated_by, + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception(f"Error updating version status: {e}") + if "invalid status" in str(e).lower() or "only draft" in str(e).lower() or "cannot promote" in str(e).lower(): + raise HTTPException(status_code=400, detail=str(e)) + if "not found" in str(e).lower(): + raise HTTPException(status_code=404, detail=str(e)) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/policies/compare", + tags=["Policies"], + dependencies=[Depends(user_api_key_auth)], + response_model=PolicyVersionCompareResponse, +) +async def compare_policy_versions( + version_a: str, + version_b: str, +): + """ + Compare two policy versions. Query params: version_a, version_b (policy version IDs). + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + return await get_policy_registry().compare_versions( + policy_id_a=version_a, + policy_id_b=version_b, + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error comparing versions: {e}") + if "not found" in str(e).lower(): + raise HTTPException(status_code=404, detail=str(e)) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.delete( + "/policies/name/{policy_name}/all-versions", + tags=["Policies"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_all_policy_versions(policy_name: str): + """ + Delete all versions of a policy. Also removes from in-memory registry. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + return await get_policy_registry().delete_all_versions( + policy_name=policy_name, + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error deleting all versions: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ───────────────────────────────────────────────────────────────────────────── +# Policy CRUD by ID +# ───────────────────────────────────────────────────────────────────────────── + + @router.get( "/policies/{policy_id}", tags=["Policies"], @@ -186,7 +360,6 @@ async def get_policy(policy_id: str): @router.put( "/policies/{policy_id}", tags=["Policies"], - dependencies=[Depends(user_api_key_auth)], response_model=PolicyDBResponse, ) async def update_policy( @@ -214,7 +387,7 @@ async def update_policy( raise HTTPException(status_code=500, detail="Database not connected") try: - # Check if policy exists + # Check if policy exists and is draft (only drafts can be updated) existing = await get_policy_registry().get_policy_by_id_from_db( policy_id=policy_id, prisma_client=prisma_client, @@ -223,6 +396,11 @@ async def update_policy( raise HTTPException( status_code=404, detail=f"Policy with ID {policy_id} not found" ) + if getattr(existing, "version_status", "production") != "draft": + raise HTTPException( + status_code=400, + detail="Only draft versions can be updated. Publish or create a new version to change published/production.", + ) updated_by = user_api_key_dict.user_id result = await get_policy_registry().update_policy_in_db( @@ -281,6 +459,7 @@ async def delete_policy(policy_id: str): policy_id=policy_id, prisma_client=prisma_client, ) + # Result may include "warning" if production was deleted return result except HTTPException: raise @@ -360,7 +539,6 @@ async def get_resolved_guardrails(policy_id: str): @router.post( "/policies/test-pipeline", tags=["Policies"], - dependencies=[Depends(user_api_key_auth)], ) async def test_pipeline( request: PipelineTestRequest, @@ -475,7 +653,6 @@ async def list_policy_attachments(): @router.post( "/policies/attachments", tags=["Policies"], - dependencies=[Depends(user_api_key_auth)], response_model=PolicyAttachmentDBResponse, ) async def create_policy_attachment( @@ -527,9 +704,11 @@ async def create_policy_attachment( raise HTTPException(status_code=500, detail="Database not connected") try: - # Verify the policy exists - policy = await get_policy_registry().get_all_policies_from_db(prisma_client) - policy_names = [p.policy_name for p in policy] + # Verify the policy has a production version (attachments resolve against production) + policies = await get_policy_registry().get_all_policies_from_db( + prisma_client, version_status="production" + ) + policy_names = {p.policy_name for p in policies} if request.policy_name not in policy_names: raise HTTPException( status_code=404, diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 1ca6062e5ac..7a5a557c0e4 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -9,23 +9,48 @@ by policy_attachments (see AttachmentRegistry). import json from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm._logging import verbose_proxy_logger -from litellm.types.proxy.policy_engine import ( - GuardrailPipeline, - PipelineStep, - Policy, - PolicyCondition, - PolicyCreateRequest, - PolicyDBResponse, - PolicyGuardrails, - PolicyUpdateRequest, -) +from litellm.types.proxy.policy_engine import (GuardrailPipeline, PipelineStep, + Policy, PolicyCondition, + PolicyCreateRequest, + PolicyDBResponse, + PolicyGuardrails, + PolicyUpdateRequest, + PolicyVersionCompareResponse, + PolicyVersionListResponse) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +# Prefix for policy version IDs in request body. Use policy_ to execute a specific version. +POLICY_VERSION_ID_PREFIX = "policy_" + + +def _row_to_policy_db_response(row: Any) -> PolicyDBResponse: + """Build PolicyDBResponse from a Prisma LiteLLM_PolicyTable row.""" + return PolicyDBResponse( + policy_id=row.policy_id, + policy_name=row.policy_name, + version_number=getattr(row, "version_number", 1), + version_status=getattr(row, "version_status", "production"), + parent_version_id=getattr(row, "parent_version_id", None), + is_latest=getattr(row, "is_latest", True), + published_at=getattr(row, "published_at", None), + production_at=getattr(row, "production_at", None), + inherit=row.inherit, + description=row.description, + guardrails_add=row.guardrails_add or [], + guardrails_remove=row.guardrails_remove or [], + condition=row.condition, + pipeline=row.pipeline, + created_at=row.created_at, + updated_at=row.updated_at, + created_by=row.created_by, + updated_by=row.updated_by, + ) + class PolicyRegistry: """ @@ -42,6 +67,7 @@ class PolicyRegistry: def __init__(self): self._policies: Dict[str, Policy] = {} + self._policies_by_id: Dict[str, Tuple[str, Policy]] = {} self._initialized: bool = False def load_policies(self, policies_config: Dict[str, Any]) -> None: @@ -53,6 +79,7 @@ class PolicyRegistry: This is the raw config from the YAML file. """ self._policies = {} + self._policies_by_id = {} for policy_name, policy_data in policies_config.items(): try: @@ -88,7 +115,9 @@ class PolicyRegistry: ) else: # Handle legacy format where guardrails might be a list - guardrails = PolicyGuardrails(add=guardrails_data if guardrails_data else None) + guardrails = PolicyGuardrails( + add=guardrails_data if guardrails_data else None + ) # Parse condition (simple model-based condition) condition = None @@ -108,7 +137,9 @@ class PolicyRegistry: ) @staticmethod - def _parse_pipeline(pipeline_data: Optional[Dict[str, Any]]) -> Optional[GuardrailPipeline]: + def _parse_pipeline( + pipeline_data: Optional[Dict[str, Any]], + ) -> Optional[GuardrailPipeline]: """Parse a pipeline configuration from raw data.""" if pipeline_data is None: return None @@ -231,13 +262,18 @@ class PolicyRegistry: PolicyDBResponse with the created policy """ try: - # Build data dict, only include condition if it's set + now = datetime.now(timezone.utc) + # Build data dict; new policy is v1 production data: Dict[str, Any] = { "policy_name": policy_request.policy_name, + "version_number": 1, + "version_status": "production", + "is_latest": True, + "production_at": now, "guardrails_add": policy_request.guardrails_add or [], "guardrails_remove": policy_request.guardrails_remove or [], - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), + "created_at": now, + "updated_at": now, } # Only add optional fields if they have values @@ -268,28 +304,17 @@ class PolicyRegistry: "add": policy_request.guardrails_add, "remove": policy_request.guardrails_remove, }, - "condition": policy_request.condition.model_dump() - if policy_request.condition - else None, + "condition": ( + policy_request.condition.model_dump() + if policy_request.condition + else None + ), "pipeline": policy_request.pipeline, }, ) self.add_policy(policy_request.policy_name, policy) - return PolicyDBResponse( - policy_id=created_policy.policy_id, - policy_name=created_policy.policy_name, - inherit=created_policy.inherit, - description=created_policy.description, - guardrails_add=created_policy.guardrails_add or [], - guardrails_remove=created_policy.guardrails_remove or [], - condition=created_policy.condition, - pipeline=created_policy.pipeline, - created_at=created_policy.created_at, - updated_at=created_policy.updated_at, - created_by=created_policy.created_by, - updated_by=created_policy.updated_by, - ) + return _row_to_policy_db_response(created_policy) except Exception as e: verbose_proxy_logger.exception(f"Error adding policy to DB: {e}") raise Exception(f"Error adding policy to DB: {str(e)}") @@ -302,7 +327,7 @@ class PolicyRegistry: updated_by: Optional[str] = None, ) -> PolicyDBResponse: """ - Update a policy in the database. + Update a policy in the database. Only draft versions can be updated. Args: policy_id: The ID of the policy to update @@ -312,8 +337,22 @@ class PolicyRegistry: Returns: PolicyDBResponse with the updated policy + + Raises: + Exception: If policy is not in draft status (only drafts are editable). """ try: + existing = await prisma_client.db.litellm_policytable.find_unique( + where={"policy_id": policy_id} + ) + if existing is None: + raise Exception(f"Policy with ID {policy_id} not found") + version_status = getattr(existing, "version_status", "production") + if version_status != "draft": + raise Exception( + f"Only draft versions can be updated. This policy has status '{version_status}'." + ) + # Build update data - only include fields that are set update_data: Dict[str, Any] = { "updated_at": datetime.now(timezone.utc), @@ -331,7 +370,9 @@ class PolicyRegistry: if policy_request.guardrails_remove is not None: update_data["guardrails_remove"] = policy_request.guardrails_remove if policy_request.condition is not None: - update_data["condition"] = json.dumps(policy_request.condition.model_dump()) + update_data["condition"] = json.dumps( + policy_request.condition.model_dump() + ) if policy_request.pipeline is not None: validated_pipeline = GuardrailPipeline(**policy_request.pipeline) update_data["pipeline"] = json.dumps(validated_pipeline.model_dump()) @@ -341,36 +382,9 @@ class PolicyRegistry: data=update_data, ) - # Update in-memory registry - policy = self._parse_policy( - updated_policy.policy_name, - { - "inherit": updated_policy.inherit, - "description": updated_policy.description, - "guardrails": { - "add": updated_policy.guardrails_add, - "remove": updated_policy.guardrails_remove, - }, - "condition": updated_policy.condition, - "pipeline": updated_policy.pipeline, - }, - ) - self.add_policy(updated_policy.policy_name, policy) + # Do NOT update in-memory registry: drafts are not loaded into memory. - return PolicyDBResponse( - policy_id=updated_policy.policy_id, - policy_name=updated_policy.policy_name, - inherit=updated_policy.inherit, - description=updated_policy.description, - guardrails_add=updated_policy.guardrails_add or [], - guardrails_remove=updated_policy.guardrails_remove or [], - condition=updated_policy.condition, - pipeline=updated_policy.pipeline, - created_at=updated_policy.created_at, - updated_at=updated_policy.updated_at, - created_by=updated_policy.created_by, - updated_by=updated_policy.updated_by, - ) + return _row_to_policy_db_response(updated_policy) except Exception as e: verbose_proxy_logger.exception(f"Error updating policy in DB: {e}") raise Exception(f"Error updating policy in DB: {str(e)}") @@ -379,19 +393,21 @@ class PolicyRegistry: self, policy_id: str, prisma_client: "PrismaClient", - ) -> Dict[str, str]: + ) -> Dict[str, Any]: """ - Delete a policy from the database. + Delete a policy version from the database. + + If the deleted version was production, it is removed from the in-memory + registry. No other version is auto-promoted; admin must explicitly promote. Args: - policy_id: The ID of the policy to delete + policy_id: The ID of the policy version to delete prisma_client: The Prisma client instance Returns: - Dict with success message + Dict with "message" and optional "warning" if production was deleted. """ try: - # Get policy name before deleting policy = await prisma_client.db.litellm_policytable.find_unique( where={"policy_id": policy_id} ) @@ -399,15 +415,27 @@ class PolicyRegistry: if policy is None: raise Exception(f"Policy with ID {policy_id} not found") + version_status = getattr(policy, "version_status", "production") + policy_name = policy.policy_name + # Delete from DB await prisma_client.db.litellm_policytable.delete( where={"policy_id": policy_id} ) - # Remove from in-memory registry - self.remove_policy(policy.policy_name) + result: Dict[str, Any] = { + "message": f"Policy {policy_id} deleted successfully" + } - return {"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." + ) + + return result except Exception as e: verbose_proxy_logger.exception(f"Error deleting policy from DB: {e}") raise Exception(f"Error deleting policy from DB: {str(e)}") @@ -435,59 +463,71 @@ class PolicyRegistry: if policy is None: return None - return PolicyDBResponse( - policy_id=policy.policy_id, - policy_name=policy.policy_name, - inherit=policy.inherit, - description=policy.description, - guardrails_add=policy.guardrails_add or [], - guardrails_remove=policy.guardrails_remove or [], - condition=policy.condition, - pipeline=policy.pipeline, - created_at=policy.created_at, - updated_at=policy.updated_at, - created_by=policy.created_by, - updated_by=policy.updated_by, - ) + return _row_to_policy_db_response(policy) except Exception as e: verbose_proxy_logger.exception(f"Error getting policy from DB: {e}") raise Exception(f"Error getting policy from DB: {str(e)}") + def _update_policies_by_id_cache(self, row: Any) -> None: + """Update the _policies_by_id cache entry for a non-production policy version.""" + policy = self._parse_policy( + row.policy_name, + { + "inherit": row.inherit, + "description": row.description, + "guardrails": { + "add": row.guardrails_add or [], + "remove": row.guardrails_remove or [], + }, + "condition": row.condition, + "pipeline": row.pipeline, + }, + ) + self._policies_by_id[row.policy_id] = (row.policy_name, policy) + + def get_policy_by_id_for_request(self, policy_id: str) -> Optional[Tuple[str, Policy]]: + """ + Return a policy version by ID from in-memory cache (no DB access). + + Used when the request body specifies policy_ to execute a specific version + (e.g. published or draft). The cache is populated by sync_policies_from_db, + which loads draft and published versions keyed by policy_id. + + Args: + policy_id: The policy version ID (raw UUID, no prefix) + + Returns: + (policy_name, Policy) if found, None otherwise + """ + return self._policies_by_id.get(policy_id) + async def get_all_policies_from_db( self, prisma_client: "PrismaClient", + version_status: Optional[str] = None, ) -> List[PolicyDBResponse]: """ - Get all policies from the database. + Get all policies from the database, optionally filtered by version_status. Args: prisma_client: The Prisma client instance + version_status: If set, only return policies with this status + ("draft", "published", "production"). Returns: List of PolicyDBResponse objects """ try: + where: Dict[str, Any] = {} + if version_status is not None: + where["version_status"] = version_status + policies = await prisma_client.db.litellm_policytable.find_many( + where=where if where else None, order={"created_at": "desc"}, ) - return [ - PolicyDBResponse( - policy_id=p.policy_id, - policy_name=p.policy_name, - inherit=p.inherit, - description=p.description, - guardrails_add=p.guardrails_add or [], - guardrails_remove=p.guardrails_remove or [], - condition=p.condition, - pipeline=p.pipeline, - created_at=p.created_at, - updated_at=p.updated_at, - created_by=p.created_by, - updated_by=p.updated_by, - ) - for p in policies - ] + return [_row_to_policy_db_response(p) for p in policies] except Exception as e: verbose_proxy_logger.exception(f"Error getting policies from DB: {e}") raise Exception(f"Error getting policies from DB: {str(e)}") @@ -498,14 +538,16 @@ class PolicyRegistry: ) -> None: """ Sync policies from the database to in-memory registry. - - Args: - prisma_client: The Prisma client instance + - Production versions are loaded into _policies (by policy name) for resolution. + - Draft and published versions are loaded into _policies_by_id so request-body + policy_ overrides can be resolved without DB access in the hot path. """ try: - policies = await self.get_all_policies_from_db(prisma_client) - - for policy_response in policies: + self._policies = {} + production = await self.get_all_policies_from_db( + prisma_client, version_status="production" + ) + for policy_response in production: policy = self._parse_policy( policy_response.policy_name, { @@ -521,9 +563,31 @@ class PolicyRegistry: ) self.add_policy(policy_response.policy_name, policy) + self._policies_by_id = {} + non_production = await prisma_client.db.litellm_policytable.find_many( + where={"version_status": {"in": ["draft", "published"]}}, + order={"created_at": "desc"}, + ) + for row in non_production: + policy = self._parse_policy( + row.policy_name, + { + "inherit": row.inherit, + "description": row.description, + "guardrails": { + "add": row.guardrails_add or [], + "remove": row.guardrails_remove or [], + }, + "condition": row.condition, + "pipeline": row.pipeline, + }, + ) + self._policies_by_id[row.policy_id] = (row.policy_name, policy) + self._initialized = True verbose_proxy_logger.info( - f"Synced {len(policies)} policies from DB to in-memory registry" + f"Synced {len(production)} production policies and {len(non_production)} " + "draft/published (by ID) from DB to in-memory registry" ) except Exception as e: verbose_proxy_logger.exception(f"Error syncing policies from DB: {e}") @@ -536,22 +600,24 @@ class PolicyRegistry: ) -> List[str]: """ Resolve all guardrails for a policy from the database. - + Uses the existing PolicyResolver to handle inheritance chain resolution. - + Args: policy_name: Name of the policy to resolve prisma_client: The Prisma client instance - + Returns: List of resolved guardrail names """ from litellm.proxy.policy_engine.policy_resolver import PolicyResolver - + try: - # Load all policies from DB to ensure we have the full inheritance chain - policies = await self.get_all_policies_from_db(prisma_client) - + # Load only production versions so inheritance resolves against production + policies = await self.get_all_policies_from_db( + prisma_client, version_status="production" + ) + # Build a temporary in-memory map for resolution temp_policies = {} for policy_response in policies: @@ -569,19 +635,346 @@ class PolicyRegistry: }, ) temp_policies[policy_response.policy_name] = policy - + # Use the existing PolicyResolver to resolve guardrails resolved_policy = PolicyResolver.resolve_policy_guardrails( policy_name=policy_name, policies=temp_policies, context=None, # No context needed for simple resolution ) - + return sorted(resolved_policy.guardrails) except Exception as e: verbose_proxy_logger.exception(f"Error resolving guardrails from DB: {e}") raise Exception(f"Error resolving guardrails from DB: {str(e)}") + async def get_versions_by_policy_name( + self, + policy_name: str, + prisma_client: "PrismaClient", + ) -> PolicyVersionListResponse: + """ + Get all versions of a policy by name, ordered by version_number descending. + + Args: + policy_name: Name of the policy + prisma_client: The Prisma client instance + + Returns: + PolicyVersionListResponse with policy_name and list of versions + """ + try: + rows = await prisma_client.db.litellm_policytable.find_many( + where={"policy_name": policy_name}, + order={"version_number": "desc"}, + ) + versions = [_row_to_policy_db_response(r) for r in rows] + return PolicyVersionListResponse( + policy_name=policy_name, + versions=versions, + total_count=len(versions), + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error getting versions: {e}") + raise Exception(f"Error getting versions: {str(e)}") + + async def create_new_version( + self, + policy_name: str, + prisma_client: "PrismaClient", + source_policy_id: Optional[str] = None, + created_by: Optional[str] = None, + ) -> PolicyDBResponse: + """ + Create a new draft version of a policy. Copies all fields from the source. + Source is current production if source_policy_id is None. + + Args: + policy_name: Name of the policy + prisma_client: The Prisma client instance + source_policy_id: Policy ID to clone from; if None, use current production + created_by: User who created the version + + Returns: + 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}'" + ) + 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}'" + ) + source = prod + + now = datetime.now(timezone.utc) + + data: Dict[str, Any] = { + "policy_name": policy_name, + "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, + } + 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 + ) + + async with prisma_client.db.tx() as tx: + latest = await tx.litellm_policytable.find_first( + where={"policy_name": policy_name}, + order={"version_number": "desc"}, + ) + next_num = (latest.version_number + 1) if latest else 1 + data["version_number"] = next_num + + await tx.litellm_policytable.update_many( + where={"policy_name": policy_name}, + data={"is_latest": False}, + ) + + 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)}") + + async def update_version_status( + self, + policy_id: str, + new_status: str, + prisma_client: "PrismaClient", + updated_by: Optional[str] = None, + ) -> PolicyDBResponse: + """ + 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) + - draft -> production: NOT allowed (must publish first) + - published -> draft: NOT allowed + + Args: + policy_id: The policy version ID + new_status: "published" or "production" + prisma_client: The Prisma client instance + updated_by: User who updated + + Returns: + PolicyDBResponse for the updated version + """ + try: + if new_status not in ("published", "production"): + raise Exception( + f"Invalid status '{new_status}'. Use 'published' or 'production'." + ) + + row = await prisma_client.db.litellm_policytable.find_unique( + where={"policy_id": policy_id} + ) + if row is None: + raise Exception(f"Policy with ID {policy_id} not found") + + current = getattr(row, "version_status", "production") + policy_name = row.policy_name + now = datetime.now(timezone.utc) + + if new_status == "published": + if current != "draft": + raise Exception( + f"Only draft versions can be published. Current status: '{current}'." + ) + updated = await prisma_client.db.litellm_policytable.update( + where={"policy_id": policy_id}, + data={ + "version_status": "published", + "published_at": now, + "updated_at": now, + "updated_by": updated_by, + }, + ) + self._update_policies_by_id_cache(updated) + return _row_to_policy_db_response(updated) + + # new_status == "production" + if current not in ("draft", "published"): + raise Exception( + f"Only draft or published versions can be promoted to production. Current: '{current}'." + ) + # Plan: "draft -> production" NOT allowed + if current == "draft": + raise Exception( + "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, + }, + ) + + # 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, + }, + ) + + # Update in-memory registry: remove old production (by name), add this one + self.remove_policy(policy_name) + policy = 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, + }, + ) + self.add_policy(policy_name, policy) + + # Remove from _policies_by_id since it's now production (resolved by name) + self._policies_by_id.pop(policy_id, None) + + return _row_to_policy_db_response(updated) + except Exception as e: + verbose_proxy_logger.exception(f"Error updating version status: {e}") + raise Exception(f"Error updating version status: {str(e)}") + + async def compare_versions( + self, + policy_id_a: str, + policy_id_b: str, + prisma_client: "PrismaClient", + ) -> PolicyVersionCompareResponse: + """ + Compare two policy versions and return field-by-field diffs. + + Args: + policy_id_a: First policy version ID + policy_id_b: Second policy version ID + prisma_client: The Prisma client instance + + Returns: + PolicyVersionCompareResponse with both versions and field_diffs + """ + try: + a = await prisma_client.db.litellm_policytable.find_unique( + where={"policy_id": policy_id_a} + ) + b = await prisma_client.db.litellm_policytable.find_unique( + where={"policy_id": policy_id_b} + ) + if a is None: + raise Exception(f"Policy {policy_id_a} not found") + if b is None: + raise Exception(f"Policy {policy_id_b} not found") + + resp_a = _row_to_policy_db_response(a) + resp_b = _row_to_policy_db_response(b) + + # Compare fields that are part of policy content (not metadata) + compare_fields = [ + "inherit", + "description", + "guardrails_add", + "guardrails_remove", + "condition", + "pipeline", + ] + field_diffs: Dict[str, Dict[str, Any]] = {} + for field in compare_fields: + val_a = getattr(resp_a, field) + val_b = getattr(resp_b, field) + if val_a != val_b: + field_diffs[field] = {"version_a": val_a, "version_b": val_b} + + return PolicyVersionCompareResponse( + version_a=resp_a, + version_b=resp_b, + field_diffs=field_diffs, + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error comparing versions: {e}") + raise Exception(f"Error comparing versions: {str(e)}") + + async def delete_all_versions( + self, + policy_name: str, + prisma_client: "PrismaClient", + ) -> Dict[str, str]: + """ + Delete all versions of a policy. Also removes from in-memory registry. + + Args: + policy_name: Name of the policy + prisma_client: The Prisma client instance + + Returns: + Dict with success message + """ + try: + await prisma_client.db.litellm_policytable.delete_many( + where={"policy_name": policy_name} + ) + self.remove_policy(policy_name) + return { + "message": f"All versions of policy '{policy_name}' deleted successfully" + } + except Exception as e: + verbose_proxy_logger.exception(f"Error deleting all versions: {e}") + raise Exception(f"Error deleting all versions: {str(e)}") + # Global singleton instance _policy_registry: Optional[PolicyRegistry] = None diff --git a/litellm/types/proxy/policy_engine/__init__.py b/litellm/types/proxy/policy_engine/__init__.py index e0c1d6f30da..4df9f21e806 100644 --- a/litellm/types/proxy/policy_engine/__init__.py +++ b/litellm/types/proxy/policy_engine/__init__.py @@ -11,48 +11,28 @@ Configuration: """ from litellm.types.proxy.policy_engine.pipeline_types import ( - GuardrailPipeline, - PipelineExecutionResult, - PipelineStep, - PipelineStepResult, -) -from litellm.types.proxy.policy_engine.policy_types import ( - Policy, - PolicyAttachment, - PolicyCondition, - PolicyConfig, - PolicyGuardrails, - PolicyScope, -) + GuardrailPipeline, PipelineExecutionResult, PipelineStep, + PipelineStepResult) +from litellm.types.proxy.policy_engine.policy_types import (Policy, + PolicyAttachment, + PolicyCondition, + PolicyConfig, + PolicyGuardrails, + PolicyScope) from litellm.types.proxy.policy_engine.resolver_types import ( - AttachmentImpactResponse, - PipelineTestRequest, - PolicyAttachmentCreateRequest, - PolicyAttachmentDBResponse, - PolicyAttachmentListResponse, - PolicyConditionRequest, - PolicyCreateRequest, - PolicyDBResponse, - PolicyGuardrailsResponse, - PolicyInfoResponse, - PolicyListDBResponse, - PolicyListResponse, - PolicyMatchContext, - PolicyMatchDetail, - PolicyResolveRequest, - PolicyResolveResponse, - PolicyScopeResponse, - PolicySummaryItem, - PolicyTestResponse, - PolicyUpdateRequest, - ResolvedPolicy, -) + AttachmentImpactResponse, PipelineTestRequest, + PolicyAttachmentCreateRequest, PolicyAttachmentDBResponse, + PolicyAttachmentListResponse, PolicyConditionRequest, PolicyCreateRequest, + PolicyDBResponse, PolicyGuardrailsResponse, PolicyInfoResponse, + PolicyListDBResponse, PolicyListResponse, PolicyMatchContext, + PolicyMatchDetail, PolicyResolveRequest, PolicyResolveResponse, + PolicyScopeResponse, PolicySummaryItem, PolicyTestResponse, + PolicyUpdateRequest, PolicyVersionCompareResponse, + PolicyVersionCreateRequest, PolicyVersionListResponse, + PolicyVersionStatusUpdateRequest, ResolvedPolicy) from litellm.types.proxy.policy_engine.validation_types import ( - PolicyValidateRequest, - PolicyValidationError, - PolicyValidationErrorType, - PolicyValidationResponse, -) + PolicyValidateRequest, PolicyValidationError, PolicyValidationErrorType, + PolicyValidationResponse) __all__ = [ # Pipeline types @@ -98,4 +78,9 @@ __all__ = [ "PolicyResolveResponse", "PolicyMatchDetail", "AttachmentImpactResponse", + # Policy versioning + "PolicyVersionCreateRequest", + "PolicyVersionStatusUpdateRequest", + "PolicyVersionListResponse", + "PolicyVersionCompareResponse", ] diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index a5a2334ae4b..759c08ab813 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -6,7 +6,7 @@ the final guardrails list. """ from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, ConfigDict, Field @@ -198,6 +198,23 @@ class PolicyDBResponse(BaseModel): policy_id: str = Field(description="Unique ID of the policy.") policy_name: str = Field(description="Name of the policy.") + version_number: int = Field(default=1, description="Version number of this policy.") + version_status: str = Field( + default="production", + description="One of: draft, published, production.", + ) + parent_version_id: Optional[str] = Field( + default=None, description="Policy ID this version was cloned from." + ) + is_latest: bool = Field( + default=True, description="True if this is the latest version by version_number." + ) + published_at: Optional[datetime] = Field( + default=None, description="When this version was published." + ) + production_at: Optional[datetime] = Field( + default=None, description="When this version was promoted to production." + ) inherit: Optional[str] = Field(default=None, description="Parent policy name.") description: Optional[str] = Field(default=None, description="Policy description.") guardrails_add: List[str] = Field( @@ -233,6 +250,49 @@ class PolicyListDBResponse(BaseModel): total_count: int = Field(default=0, description="Total number of policies.") +# ───────────────────────────────────────────────────────────────────────────── +# Policy Versioning Types +# ───────────────────────────────────────────────────────────────────────────── + + +class PolicyVersionCreateRequest(BaseModel): + """Request body for creating a new policy version (draft).""" + + source_policy_id: Optional[str] = Field( + default=None, + description="Policy ID to clone from. If None, clone from current production version.", + ) + + +class PolicyVersionStatusUpdateRequest(BaseModel): + """Request body for updating a policy version's status.""" + + version_status: Literal["published", "production"] = Field( + description="New status: 'published' or 'production'.", + ) + + +class PolicyVersionListResponse(BaseModel): + """Response for listing all versions of a policy.""" + + policy_name: str = Field(description="Name of the policy.") + versions: List[PolicyDBResponse] = Field( + default_factory=list, description="All versions ordered by version_number desc." + ) + total_count: int = Field(default=0, description="Total number of versions.") + + +class PolicyVersionCompareResponse(BaseModel): + """Response for comparing two policy versions.""" + + version_a: PolicyDBResponse = Field(description="First version.") + version_b: PolicyDBResponse = Field(description="Second version.") + field_diffs: Dict[str, Dict[str, Any]] = Field( + default_factory=dict, + description="Field name -> {version_a: val, version_b: val} for differing fields.", + ) + + # ───────────────────────────────────────────────────────────────────────────── # Policy Attachment CRUD Types # ───────────────────────────────────────────────────────────────────────────── diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py new file mode 100644 index 00000000000..1f3a02a4c97 --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -0,0 +1,446 @@ +""" +Unit tests for policy versioning: registry behavior, status transitions, and version CRUD. +""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.policy_engine.policy_registry import ( + PolicyRegistry, + _row_to_policy_db_response, + get_policy_registry, +) +from litellm.types.proxy.policy_engine import ( + PolicyCreateRequest, + PolicyDBResponse, + PolicyUpdateRequest, +) + + +def _make_row( + policy_id="pid-1", + policy_name="test-policy", + version_number=1, + version_status="production", + parent_version_id=None, + is_latest=True, + published_at=None, + production_at=None, + inherit=None, + description="desc", + guardrails_add=None, + guardrails_remove=None, + condition=None, + pipeline=None, + created_at=None, + updated_at=None, + created_by=None, + updated_by=None, +): + row = MagicMock() + row.policy_id = policy_id + row.policy_name = policy_name + row.version_number = version_number + row.version_status = version_status + row.parent_version_id = parent_version_id + row.is_latest = is_latest + row.published_at = published_at + row.production_at = production_at + row.inherit = inherit + row.description = description + row.guardrails_add = guardrails_add or [] + row.guardrails_remove = guardrails_remove or [] + row.condition = condition + row.pipeline = pipeline + row.created_at = created_at or datetime.now(timezone.utc) + row.updated_at = updated_at or datetime.now(timezone.utc) + row.created_by = created_by + row.updated_by = updated_by + return row + + +class TestRowToPolicyDBResponse: + """Test _row_to_policy_db_response includes all version fields.""" + + def test_includes_version_fields(self): + row = _make_row( + version_number=2, + version_status="draft", + parent_version_id="pid-0", + is_latest=True, + published_at=None, + production_at=None, + ) + resp = _row_to_policy_db_response(row) + assert isinstance(resp, PolicyDBResponse) + assert resp.policy_id == "pid-1" + assert resp.policy_name == "test-policy" + assert resp.version_number == 2 + assert resp.version_status == "draft" + assert resp.parent_version_id == "pid-0" + assert resp.is_latest is True + assert resp.published_at is None + assert resp.production_at is None + + def test_backward_compat_missing_version_attrs(self): + row = _make_row() + del row.version_number + del row.version_status + del row.parent_version_id + del row.is_latest + del row.published_at + del row.production_at + resp = _row_to_policy_db_response(row) + assert resp.version_number == 1 + assert resp.version_status == "production" + assert resp.parent_version_id is None + assert resp.is_latest is True + + +class TestSyncPoliciesFromDbProductionOnly: + """Test that sync_policies_from_db only loads production versions.""" + + @pytest.mark.asyncio + async def test_get_all_policies_with_version_status_calls_find_many_with_where(self): + registry = PolicyRegistry() + prisma = MagicMock() + prod_row = _make_row(policy_id="prod-1", version_status="production") + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row]) + + result = await registry.get_all_policies_from_db( + prisma, version_status="production" + ) + + assert len(result) == 1 + assert result[0].version_status == "production" + prisma.db.litellm_policytable.find_many.assert_called_once() + call_kw = prisma.db.litellm_policytable.find_many.call_args[1] + assert call_kw.get("where") == {"version_status": "production"} + + @pytest.mark.asyncio + async def test_sync_policies_from_db_only_loads_production(self): + registry = PolicyRegistry() + prisma = MagicMock() + prod_row = _make_row( + policy_id="prod-1", + policy_name="foo", + version_status="production", + guardrails_add=["g1"], + ) + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row]) + + await registry.sync_policies_from_db(prisma) + + assert registry.has_policy("foo") + policy = registry.get_policy("foo") + assert policy is not None + assert policy.guardrails.add == ["g1"] + # find_many was called with version_status=production (via get_all_policies_from_db) + find_many_calls = prisma.db.litellm_policytable.find_many.call_args_list + assert len(find_many_calls) >= 1 + assert find_many_calls[0][1].get("where") == {"version_status": "production"} + + +class TestUpdatePolicyDraftOnly: + """Test that update_policy_in_db only allows draft versions.""" + + @pytest.mark.asyncio + async def test_update_production_raises(self): + registry = PolicyRegistry() + prisma = MagicMock() + prod_row = _make_row(policy_id="pid-1", version_status="production") + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row) + + with pytest.raises(Exception) as exc_info: + await registry.update_policy_in_db( + policy_id="pid-1", + policy_request=PolicyUpdateRequest(description="new"), + prisma_client=prisma, + ) + assert "Only draft" in str(exc_info.value) or "draft" in str(exc_info.value).lower() + prisma.db.litellm_policytable.update.assert_not_called() + + @pytest.mark.asyncio + async def test_update_draft_succeeds_and_does_not_update_registry(self): + registry = PolicyRegistry() + registry.add_policy("test-policy", MagicMock()) # in-memory state + prisma = MagicMock() + draft_row = _make_row( + policy_id="draft-1", + policy_name="test-policy", + version_status="draft", + description="old", + ) + updated_row = _make_row( + policy_id="draft-1", + policy_name="test-policy", + version_status="draft", + description="new", + ) + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft_row) + prisma.db.litellm_policytable.update = AsyncMock(return_value=updated_row) + + result = await registry.update_policy_in_db( + policy_id="draft-1", + policy_request=PolicyUpdateRequest(description="new"), + prisma_client=prisma, + ) + + assert result.description == "new" + prisma.db.litellm_policytable.update.assert_called_once() + # Registry still has old in-memory policy (drafts are not in registry; we don't add) + assert registry.has_policy("test-policy") + + +class TestDeletePolicyFromDb: + """Test delete_policy_from_db removes production from registry and returns warning.""" + + @pytest.mark.asyncio + async def test_delete_production_removes_from_registry_and_returns_warning(self): + registry = PolicyRegistry() + registry.add_policy("deleted-policy", MagicMock()) + prisma = MagicMock() + prod_row = _make_row( + policy_id="prod-1", + policy_name="deleted-policy", + version_status="production", + ) + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row) + prisma.db.litellm_policytable.delete = AsyncMock() + + result = await registry.delete_policy_from_db( + policy_id="prod-1", + prisma_client=prisma, + ) + + assert result["message"] + assert "warning" in result + assert "Production" in result["warning"] or "production" in result["warning"] + assert not registry.has_policy("deleted-policy") + + @pytest.mark.asyncio + async def test_delete_draft_does_not_remove_from_registry_no_warning(self): + registry = PolicyRegistry() + registry.add_policy("my-policy", MagicMock()) + 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() + + result = await registry.delete_policy_from_db( + policy_id="draft-1", + prisma_client=prisma, + ) + + assert "warning" not in result + assert registry.has_policy("my-policy") + + +class TestCreateNewVersion: + """Test create_new_version copies all fields and sets draft.""" + + @pytest.mark.asyncio + async def test_create_new_version_from_production_increments_version(self): + registry = PolicyRegistry() + prisma = MagicMock() + prod = _make_row( + policy_id="prod-1", + policy_name="foo", + version_number=1, + version_status="production", + guardrails_add=["g1"], + description="base", + inherit=None, + pipeline={"mode": "pre_call", "steps": []}, + ) + prisma.db.litellm_policytable.find_first = AsyncMock(return_value=prod) + new_row = _make_row( + policy_id="new-id", + policy_name="foo", + version_number=2, + version_status="draft", + parent_version_id="prod-1", + is_latest=True, + guardrails_add=["g1"], + description="base", + pipeline={"mode": "pre_call", "steps": []}, + ) + + tx_mock = MagicMock() + tx_mock.litellm_policytable.find_first = AsyncMock(return_value=prod) + tx_mock.litellm_policytable.update_many = AsyncMock() + tx_mock.litellm_policytable.create = AsyncMock(return_value=new_row) + + async def _tx_context(): + return tx_mock + + 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.create_new_version( + policy_name="foo", + prisma_client=prisma, + source_policy_id=None, + created_by="user", + ) + + assert result.version_number == 2 + assert result.version_status == "draft" + assert result.parent_version_id == "prod-1" + assert result.guardrails_add == ["g1"] + assert result.description == "base" + create_call = tx_mock.litellm_policytable.create.call_args[1]["data"] + assert create_call["version_number"] == 2 + assert create_call["version_status"] == "draft" + assert create_call["parent_version_id"] == "prod-1" + assert create_call["guardrails_add"] == ["g1"] + + +class TestUpdateVersionStatus: + """Test status transitions: valid succeed, invalid return error.""" + + @pytest.mark.asyncio + async def test_draft_to_published_sets_published_at(self): + registry = PolicyRegistry() + prisma = MagicMock() + draft = _make_row(policy_id="d-1", version_status="draft") + updated = _make_row( + policy_id="d-1", + version_status="published", + published_at=datetime.now(timezone.utc), + ) + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft) + prisma.db.litellm_policytable.update = AsyncMock(return_value=updated) + + result = await registry.update_version_status( + policy_id="d-1", + new_status="published", + prisma_client=prisma, + ) + + assert result.version_status == "published" + update_data = prisma.db.litellm_policytable.update.call_args[1]["data"] + assert update_data["version_status"] == "published" + assert "published_at" in update_data + + @pytest.mark.asyncio + async def test_draft_to_production_raises(self): + registry = PolicyRegistry() + prisma = MagicMock() + draft = _make_row(policy_id="d-1", version_status="draft") + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft) + + with pytest.raises(Exception) as exc_info: + await registry.update_version_status( + policy_id="d-1", + new_status="production", + prisma_client=prisma, + ) + assert "publish" in str(exc_info.value).lower() or "draft" in str(exc_info.value).lower() + + @pytest.mark.asyncio + async def test_published_to_production_demotes_old_and_updates_registry(self): + registry = PolicyRegistry() + 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) + prisma.db.litellm_policytable.update_many = AsyncMock() + prisma.db.litellm_policytable.update = AsyncMock(return_value=updated_row) + + result = await registry.update_version_status( + policy_id="pub-1", + new_status="production", + prisma_client=prisma, + ) + + 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 registry.has_policy("foo") + + +class TestCompareVersions: + """Test compare_versions returns correct field diffs.""" + + @pytest.mark.asyncio + async def test_compare_versions_returns_diffs(self): + registry = PolicyRegistry() + prisma = MagicMock() + a = _make_row( + policy_id="a", + policy_name="p", + description="desc A", + guardrails_add=["g1"], + ) + b = _make_row( + policy_id="b", + policy_name="p", + description="desc B", + guardrails_add=["g1", "g2"], + ) + prisma.db.litellm_policytable.find_unique = AsyncMock(side_effect=[a, b]) + + result = await registry.compare_versions( + policy_id_a="a", + policy_id_b="b", + prisma_client=prisma, + ) + + assert result.version_a.policy_id == "a" + assert result.version_b.policy_id == "b" + assert "description" in result.field_diffs + assert result.field_diffs["description"]["version_a"] == "desc A" + assert result.field_diffs["description"]["version_b"] == "desc B" + assert "guardrails_add" in result.field_diffs + + +class TestResolveGuardrailsProductionOnly: + """Test that resolve_guardrails_from_db uses only production versions.""" + + @pytest.mark.asyncio + async def test_resolve_guardrails_calls_get_all_with_production_filter(self): + registry = PolicyRegistry() + prisma = MagicMock() + prod_row = _make_row( + policy_name="base", + version_status="production", + guardrails_add=["g1"], + ) + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row]) + + result = await registry.resolve_guardrails_from_db( + policy_name="base", + prisma_client=prisma, + ) + + assert "g1" in result + call_kw = prisma.db.litellm_policytable.find_many.call_args[1] + assert call_kw.get("where") == {"version_status": "production"} + + +class TestGetPolicyRegistrySingleton: + """Test get_policy_registry returns same instance.""" + + def test_returns_singleton(self): + a = get_policy_registry() + b = get_policy_registry() + assert a is b 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 new file mode 100644 index 00000000000..f611d569fb7 --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py @@ -0,0 +1,224 @@ +""" +Integration-style tests for policy versioning: full lifecycle with mocked DB. +""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.policy_engine.policy_registry import PolicyRegistry +from litellm.types.proxy.policy_engine import (PolicyCreateRequest, + PolicyUpdateRequest) + + +def _make_row( + policy_id, + policy_name, + version_number=1, + version_status="production", + parent_version_id=None, + is_latest=True, + published_at=None, + production_at=None, + inherit=None, + description="", + guardrails_add=None, + guardrails_remove=None, + condition=None, + pipeline=None, +): + row = MagicMock() + row.policy_id = policy_id + row.policy_name = policy_name + row.version_number = version_number + row.version_status = version_status + row.parent_version_id = parent_version_id + row.is_latest = is_latest + row.published_at = published_at + row.production_at = production_at + row.inherit = inherit + row.description = description + row.guardrails_add = guardrails_add or [] + row.guardrails_remove = guardrails_remove or [] + row.condition = condition + row.pipeline = pipeline + row.created_at = datetime.now(timezone.utc) + row.updated_at = datetime.now(timezone.utc) + row.created_by = None + row.updated_by = None + return row + + +@pytest.mark.asyncio +async def test_full_lifecycle_create_draft_edit_publish_promote(): + """ + Full lifecycle: create policy -> create draft version -> edit draft -> + publish -> promote to production -> verify old version demoted -> + verify in-memory updated. + """ + registry = PolicyRegistry() + prisma = MagicMock() + now = datetime.now(timezone.utc) + + # 1) Create initial policy (v1 production) + create_data = {} + created_v1 = _make_row( + policy_id="v1-id", + policy_name="lifecycle-policy", + version_number=1, + version_status="production", + production_at=now, + guardrails_add=["g1"], + description="Initial", + ) + + async def create_impl(data=None, **kwargs): + create_data.update(kwargs.get("data", data or {})) + return created_v1 + + prisma.db.litellm_policytable.create = AsyncMock(side_effect=create_impl) + req = PolicyCreateRequest( + policy_name="lifecycle-policy", + description="Initial", + guardrails_add=["g1"], + ) + created = await registry.add_policy_to_db(req, prisma, created_by="user") + assert created.version_number == 1 + assert created.version_status == "production" + assert registry.has_policy("lifecycle-policy") + + # 2) Create new draft version (v2) + v2_row = _make_row( + policy_id="v2-id", + policy_name="lifecycle-policy", + version_number=2, + version_status="draft", + parent_version_id="v1-id", + is_latest=True, + guardrails_add=["g1", "g2"], + description="Draft v2", + ) + prisma.db.litellm_policytable.find_first = AsyncMock(return_value=created_v1) + + tx_mock = MagicMock() + tx_mock.litellm_policytable.find_first = AsyncMock(return_value=created_v1) + tx_mock.litellm_policytable.update_many = AsyncMock() + tx_mock.litellm_policytable.create = AsyncMock(return_value=v2_row) + ctx = MagicMock() + ctx.__aenter__ = AsyncMock(return_value=tx_mock) + ctx.__aexit__ = AsyncMock(return_value=False) + prisma.db.tx = MagicMock(return_value=ctx) + + draft_v2 = await registry.create_new_version( + policy_name="lifecycle-policy", + prisma_client=prisma, + source_policy_id=None, + created_by="user", + ) + assert draft_v2.version_number == 2 + assert draft_v2.version_status == "draft" + assert draft_v2.parent_version_id == "v1-id" + # In-memory still has v1 (only production is in registry) + assert registry.has_policy("lifecycle-policy") + policy = registry.get_policy("lifecycle-policy") + assert policy.guardrails.add == ["g1"] # still v1 + + # 3) Edit draft v2 + v2_updated_row = _make_row( + policy_id="v2-id", + policy_name="lifecycle-policy", + version_number=2, + version_status="draft", + guardrails_add=["g1", "g2", "g3"], + description="Draft v2 edited", + ) + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=v2_row) + prisma.db.litellm_policytable.update = AsyncMock(return_value=v2_updated_row) + + updated_draft = await registry.update_policy_in_db( + policy_id="v2-id", + policy_request=PolicyUpdateRequest( + description="Draft v2 edited", + guardrails_add=["g1", "g2", "g3"], + ), + prisma_client=prisma, + updated_by="user", + ) + assert updated_draft.description == "Draft v2 edited" + assert updated_draft.guardrails_add == ["g1", "g2", "g3"] + + # 4) Publish v2 (draft -> published) + v2_published = _make_row( + policy_id="v2-id", + policy_name="lifecycle-policy", + version_number=2, + version_status="published", + published_at=now, + guardrails_add=["g1", "g2", "g3"], + description="Draft v2 edited", + ) + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=v2_updated_row) + prisma.db.litellm_policytable.update = AsyncMock(return_value=v2_published) + + published = await registry.update_version_status( + policy_id="v2-id", + new_status="published", + prisma_client=prisma, + updated_by="user", + ) + assert published.version_status == "published" + + # 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", + version_number=2, + version_status="production", + production_at=now, + guardrails_add=["g1", "g2", "g3"], + description="Draft v2 edited", + ) + prisma.db.litellm_policytable.update = AsyncMock(return_value=v2_production) + + prod = await registry.update_version_status( + policy_id="v2-id", + new_status="production", + prisma_client=prisma, + updated_by="user", + ) + assert prod.version_status == "production" + # In-memory registry should now have v2 content + assert registry.has_policy("lifecycle-policy") + policy = registry.get_policy("lifecycle-policy") + assert policy.guardrails.add == ["g1", "g2", "g3"] + + +@pytest.mark.asyncio +async def test_attachments_resolve_against_production_after_promotion(): + """ + After promoting a new version to production, resolve_guardrails_from_db + returns guardrails from the new production version (inheritance resolves + against production). + """ + registry = PolicyRegistry() + prisma = MagicMock() + # Simulate only production versions loaded for resolution + prod_row = _make_row( + policy_id="prod-1", + policy_name="att-policy", + version_status="production", + guardrails_add=["ga", "gb"], + ) + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row]) + + resolved = await registry.resolve_guardrails_from_db( + policy_name="att-policy", + prisma_client=prisma, + ) + assert "ga" in resolved + assert "gb" in resolved + call_kw = prisma.db.litellm_policytable.find_many.call_args[1] + assert call_kw.get("where") == {"version_status": "production"}