fix(policy_engine): resolve policies once and apply fallback semantics in get_matching_policies

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-21 17:36:03 +00:00
parent 99e14efb62
commit e647255909
2 changed files with 60 additions and 9 deletions

View file

@ -114,7 +114,7 @@ class PolicyMatcher:
verbose_proxy_logger.debug("AttachmentRegistry not initialized, returning empty list")
return []
return registry.get_attached_policies(context)
return registry.get_attached_policies(context, PolicyMatcher.policy_applies(context))
@staticmethod
def get_matching_policies_from_registry(
@ -137,14 +137,22 @@ class PolicyMatcher:
policies: dict[str, Policy] | None = None,
) -> Callable[[str], bool]:
"""Predicate telling whether a policy exists and its condition matches the context."""
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
return lambda policy_name: bool(
PolicyMatcher.get_policies_with_matching_conditions(
policy_names=(policy_name,),
context=context,
policies=policies,
policies=resolved,
)
)
@staticmethod
def _registry_policies() -> dict[str, Policy]:
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
registry: Final = get_policy_registry()
return registry.get_all_policies() if registry.is_initialized() else {}
@staticmethod
def get_policies_with_matching_conditions(
policy_names: Sequence[str],
@ -167,17 +175,12 @@ class PolicyMatcher:
List of policy names whose conditions match the context
"""
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
if policies is None:
registry: Final = get_policy_registry()
if not registry.is_initialized():
return []
policies = registry.get_all_policies()
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
matching_policies: Final = []
for policy_name in policy_names:
policy = policies.get(policy_name)
policy = resolved.get(policy_name)
if policy is None:
continue
# Policy matches if it has no condition OR condition evaluates to True

View file

@ -8,8 +8,11 @@ Tests:
import pytest
import litellm.proxy.policy_engine.attachment_registry as attachment_registry_module
import litellm.proxy.policy_engine.policy_registry as policy_registry_module
from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
from litellm.proxy.policy_engine.policy_registry import PolicyRegistry
from litellm.types.proxy.policy_engine import (
PolicyMatchContext,
PolicyScope,
@ -196,3 +199,48 @@ class TestPolicyMatcherWithAttachments:
attached = registry.get_attached_policies(context)
assert "healthcare-policy" not in attached
def _global_registries(monkeypatch):
policies = PolicyRegistry()
policies.load_policies(
{
"guardrail-y": {"guardrails": {"add": ["y"]}},
"guardrail-x": {"guardrails": {"add": ["x"]}, "condition": {"model": "claude.*"}},
}
)
attachments = AttachmentRegistry()
attachments.load_attachments(
[
{"policy": "guardrail-x", "tags": ["opt-in"]},
{"policy": "guardrail-y", "scope": "*", "default": True},
]
)
monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: policies)
monkeypatch.setattr(attachment_registry_module, "get_attachment_registry", lambda: attachments)
return policies
class TestGetMatchingPoliciesFallback:
def test_condition_failing_opt_in_falls_back_to_default(self, monkeypatch):
_global_registries(monkeypatch)
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"])
assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-y"]
def test_condition_passing_opt_in_suppresses_default(self, monkeypatch):
_global_registries(monkeypatch)
context = PolicyMatchContext(team_alias="t", key_alias="k", model="claude-haiku", tags=["opt-in"])
assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-x"]
def test_policy_applies_reads_registry_once(self, monkeypatch):
policies = _global_registries(monkeypatch)
calls = []
original = policies.get_all_policies
monkeypatch.setattr(policies, "get_all_policies", lambda: calls.append(1) or original())
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"])
PolicyMatcher.get_matching_policies(context=context)
assert len(calls) == 1