feat(policy_engine): explicit priority for policy attachment execution order

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-proxy-extras/litellm_proxy_extras/migrations/20260917055603_add_policy_attachment_priority/migration.sql

@@ -0,0 +1 @@
+ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN "priority" INTEGER;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1378,6 +1378,7 @@ model LiteLLM_PolicyAttachmentTable {
   keys             String[]  @default([])  // Key aliases or patterns
   models           String[]  @default([])  // Model names or patterns
   tags             String[]  @default([])  // Tag patterns (e.g., ["healthcare", "prod-*"])
+  priority         Int?                    // Explicit execution order
   created_at       DateTime  @default(now())
   created_by       String?
   updated_at       DateTime  @default(now()) @updatedAt

litellm/proxy/policy_engine/attachment_registry.py

@@ -48,6 +48,13 @@ def _attachment_specificity(attachment: PolicyAttachment) -> tuple[int, int]:
     return (max(dims, default=0), len(dims))
 
 
+def _attachment_sort_key(attachment: PolicyAttachment) -> tuple[int, int, int, int]:
+    specificity: Final = _attachment_specificity(attachment)
+    if attachment.priority is not None:
+        return (0, attachment.priority, *specificity)
+    return (1, 0, *specificity)
+
+
 class AttachmentRegistry:
     """
     In-memory registry for storing and managing policy attachments.
@@ -111,6 +118,7 @@ class AttachmentRegistry:
             keys=attachment_data.get("keys"),
             models=attachment_data.get("models"),
             tags=attachment_data.get("tags"),
+            priority=attachment_data.get("priority"),
         )
 
     def get_attached_policies(self, context: PolicyMatchContext) -> list[str]:
@@ -140,7 +148,7 @@ class AttachmentRegistry:
                 for attachment in self._attachments
                 if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
             ),
-            key=_attachment_specificity,
+            key=_attachment_sort_key,
         )
         broadest_attachment_by_policy: Final = MappingProxyType(
             {attachment.policy: attachment for attachment in reversed(matching_attachments)}
@@ -315,6 +323,7 @@ class AttachmentRegistry:
                     "keys": attachment_request.keys or [],
                     "models": attachment_request.models or [],
                     "tags": attachment_request.tags or [],
+                    "priority": attachment_request.priority,
                     "created_at": datetime.now(timezone.utc),
                     "updated_at": datetime.now(timezone.utc),
                     "created_by": created_by,
@@ -330,6 +339,7 @@ class AttachmentRegistry:
                 keys=attachment_request.keys,
                 models=attachment_request.models,
                 tags=attachment_request.tags,
+                priority=attachment_request.priority,
             )
             self.add_attachment(attachment)
 
@@ -341,6 +351,7 @@ class AttachmentRegistry:
                 keys=created_attachment.keys or [],
                 models=created_attachment.models or [],
                 tags=created_attachment.tags or [],
+                priority=created_attachment.priority,
                 created_at=created_attachment.created_at,
                 updated_at=created_attachment.updated_at,
                 created_by=created_attachment.created_by,
@@ -417,6 +428,7 @@ class AttachmentRegistry:
                 keys=attachment.keys or [],
                 models=attachment.models or [],
                 tags=attachment.tags or [],
+                priority=attachment.priority,
                 created_at=attachment.created_at,
                 updated_at=attachment.updated_at,
                 created_by=attachment.created_by,
@@ -455,6 +467,7 @@ class AttachmentRegistry:
                     keys=a.keys or [],
                     models=a.models or [],
                     tags=a.tags or [],
+                    priority=a.priority,
                     created_at=a.created_at,
                     updated_at=a.updated_at,
                     created_by=a.created_by,
@@ -488,6 +501,7 @@ class AttachmentRegistry:
                     keys=attachment_response.keys if attachment_response.keys else None,
                     models=(attachment_response.models if attachment_response.models else None),
                     tags=attachment_response.tags if attachment_response.tags else None,
+                    priority=attachment_response.priority,
                 )
                 for attachment_response in attachments
             ]

litellm/proxy/policy_engine/policy_endpoints.py

@@ -60,6 +60,7 @@ def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment)
         keys=attachment.keys or [],
         models=attachment.models or [],
         tags=attachment.tags or [],
+        priority=attachment.priority,
         definition_location="config",
     )
 

litellm/proxy/schema.prisma

@@ -1378,6 +1378,7 @@ model LiteLLM_PolicyAttachmentTable {
   keys             String[]  @default([])  // Key aliases or patterns
   models           String[]  @default([])  // Model names or patterns
   tags             String[]  @default([])  // Tag patterns (e.g., ["healthcare", "prod-*"])
+  priority         Int?                    // Explicit execution order
   created_at       DateTime  @default(now())
   created_by       String?
   updated_at       DateTime  @default(now()) @updatedAt

litellm/types/proxy/policy_engine/policy_types.py

@@ -288,6 +288,10 @@ class PolicyAttachment(BaseModel):
         default=None,
         description="Tag patterns this attachment applies to. Supports wildcards (e.g., health-*).",
     )
+    priority: int | None = Field(
+        default=None,
+        description="Explicit execution order. Attachments with a priority run before those without, lower first; ties fall back to scope specificity.",
+    )
 
     model_config = ConfigDict(extra="forbid")
 

litellm/types/proxy/policy_engine/resolver_types.py

@@ -305,6 +305,10 @@ class PolicyAttachmentCreateRequest(BaseModel):
         default=None,
         description="Tag patterns this attachment applies to. Supports wildcards (e.g., health-*).",
     )
+    priority: int | None = Field(
+        default=None,
+        description="Explicit execution order. Attachments with a priority run before those without, lower first; ties fall back to scope specificity.",
+    )
 
 
 class PolicyAttachmentDBResponse(BaseModel):
@@ -317,6 +321,10 @@ class PolicyAttachmentDBResponse(BaseModel):
     keys: list[str] = Field(default_factory=list, description="Key patterns.")
     models: list[str] = Field(default_factory=list, description="Model patterns.")
     tags: list[str] = Field(default_factory=list, description="Tag patterns.")
+    priority: int | None = Field(
+        default=None,
+        description="Explicit execution order. Attachments with a priority run before those without, lower first; ties fall back to scope specificity.",
+    )
     created_at: datetime | None = Field(default=None, description="When the attachment was created.")
     updated_at: datetime | None = Field(default=None, description="When the attachment was last updated.")
     created_by: str | None = Field(default=None, description="Who created the attachment.")

schema.prisma

@@ -1378,6 +1378,7 @@ model LiteLLM_PolicyAttachmentTable {
   keys             String[]  @default([])  // Key aliases or patterns
   models           String[]  @default([])  // Model names or patterns
   tags             String[]  @default([])  // Tag patterns (e.g., ["healthcare", "prod-*"])
+  priority         Int?                    // Explicit execution order
   created_at       DateTime  @default(now())
   created_by       String?
   updated_at       DateTime  @default(now()) @updatedAt

tests/test_litellm/proxy/policy_engine/test_attachment_registry.py

@@ -158,6 +158,37 @@ class TestGetAttachedPolicies:
             "model-policy",
         ]
 
+    def test_prioritized_attachments_run_before_unprioritized_attachments(self):
+        registry = AttachmentRegistry()
+        registry.load_attachments(
+            [
+                {"policy": "unprioritized-tag", "tags": ["prod"]},
+                {"policy": "prioritized-tag", "tags": ["prod"], "priority": 5},
+                {"policy": "prioritized-model", "models": ["gpt-4"], "priority": 0},
+            ]
+        )
+
+        context = PolicyMatchContext(model="gpt-4", tags=["prod"])
+
+        assert registry.get_attached_policies(context) == [
+            "prioritized-model",
+            "prioritized-tag",
+            "unprioritized-tag",
+        ]
+
+    def test_prioritized_attachments_order_by_priority_across_scope_tiers(self):
+        registry = AttachmentRegistry()
+        registry.load_attachments(
+            [
+                {"policy": "team-policy", "teams": ["team-a"], "priority": 2},
+                {"policy": "model-policy", "models": ["gpt-4"], "priority": 1},
+            ]
+        )
+
+        context = PolicyMatchContext(team_alias="team-a", model="gpt-4")
+
+        assert registry.get_attached_policies(context) == ["model-policy", "team-policy"]
+
     def test_combined_team_and_model_attachment_uses_model_specificity(self):
         registry = AttachmentRegistry()
         registry.load_attachments(
@@ -474,8 +505,28 @@ class TestAttachmentRegistrySingleton:
         registry2 = get_attachment_registry()
         assert registry1 is registry2
 
+    def test_parse_attachment_reads_priority(self):
+        registry = AttachmentRegistry()
+        registry.load_attachments(
+            [
+                {"policy": "prioritized", "priority": 4},
+                {"policy": "unprioritized"},
+            ]
+        )
 
-def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scope=None, teams=None):
+        attachments = registry.get_all_attachments()
+
+        assert attachments[0].priority == 4
+        assert attachments[1].priority is None
+
+
+def _make_db_attachment_row(
+    attachment_id: str = "att-1",
+    policy_name: str = "db-policy",
+    scope: str | None = None,
+    teams: list[str] | None = None,
+    priority: int | None = None,
+) -> MagicMock:
     row = MagicMock()
     row.attachment_id = attachment_id
     row.policy_name = policy_name
@@ -484,6 +535,7 @@ def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scop
     row.keys = []
     row.models = []
     row.tags = []
+    row.priority = priority
     row.created_at = datetime.now(timezone.utc)
     row.updated_at = datetime.now(timezone.utc)
     row.created_by = None
@@ -491,9 +543,11 @@ def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scop
     return row
 
 
-def _prisma_with_attachment_rows(rows):
+def _prisma_with_attachment_rows(rows: list[MagicMock]) -> MagicMock:
     prisma = MagicMock()
-    prisma.db.litellm_policyattachmenttable.find_many = AsyncMock(return_value=rows)
+    prisma.configure_mock(
+        **{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)}
+    )
     return prisma
 
 
@@ -535,6 +589,15 @@ class TestConfigAttachmentsPreservedAcrossDbSync:
 
         assert len(registry.get_all_attachments()) == 1
 
+    @pytest.mark.asyncio
+    async def test_sync_round_trips_db_attachment_priority(self):
+        registry = AttachmentRegistry()
+        db_row = _make_db_attachment_row(priority=7)
+
+        await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row]))
+
+        assert registry.get_all_attachments()[0].priority == 7
+
     @pytest.mark.asyncio
     async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
         registry = AttachmentRegistry()