From 311d9bba37b44099461e6e699232b0903474fc62 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:19:20 -0700 Subject: [PATCH] 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 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/policy_engine/attachment_registry.py | 6 +++++- .../policy_engine/test_attachment_registry.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index a8ead86ac36..76b2291774e 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/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) ) diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 87b1bc56659..fa37a02a37c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/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()