perf(policy_engine): dedup attachments in one pass after sorting (#40883)

get_attached_policies_with_reasons rescanned the sorted matches with next() once
per distinct policy, which is quadratic and misses the one second budget past a
few thousand global attachments. Build a policy to broadest attachment map in one
pass instead, keeping the specificity sort and result order.

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/policy_engine/attachment_registry.py

@@ -6,6 +6,7 @@ This allows the same policy to be attached to multiple scopes.
 """
 
 from datetime import datetime, timezone
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, TypedDict
 
 from litellm._logging import verbose_proxy_logger
@@ -141,8 +142,11 @@ class AttachmentRegistry:
             ),
             key=_attachment_specificity,
         )
+        broadest_attachment_by_policy: Final = MappingProxyType(
+            {attachment.policy: attachment for attachment in reversed(matching_attachments)}
+        )
         unique_attachments: Final = tuple(
-            next(attachment for attachment in matching_attachments if attachment.policy == policy_name)
+            broadest_attachment_by_policy[policy_name]
             for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments)
         )
 

tests/test_litellm/proxy/policy_engine/test_attachment_registry.py

@@ -4,6 +4,7 @@ Unit tests for AttachmentRegistry - tests policy attachment matching.
 Tests the main entry point: get_attached_policies()
 """
 
+import time
 from datetime import datetime, timezone
 from unittest.mock import AsyncMock, MagicMock
 
@@ -222,6 +223,21 @@ class TestGetAttachedPolicies:
         # Should only appear once
         assert attached.count("multi-policy") == 1
 
+    def test_many_distinct_policies_resolve_in_linear_time(self):
+        policy_count = 20_000
+        registry = AttachmentRegistry()
+        registry.load_attachments(
+            [{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)]
+        )
+        context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
+
+        started = time.perf_counter()
+        attached = registry.get_attached_policies(context)
+        elapsed = time.perf_counter() - started
+
+        assert attached == [f"policy-{index}" for index in range(policy_count)]
+        assert elapsed < 1.0, f"{policy_count} attachments took {elapsed:.2f}s, dedup is no longer one pass"
+
     def test_no_attachments_returns_empty(self):
         """Test empty attachments returns empty list."""
         registry = AttachmentRegistry()