diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql new file mode 100644 index 00000000000..a6c45448d03 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d2032cec0d0..2d7e557a9d1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable { models String[] @default([]) // Model names or patterns tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) priority Int? // Explicit execution order + is_default Boolean @default(false) // Applied only when no non-default attachment matches created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 06e157498aa..6f9a2d8c96d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -34982,6 +34982,12 @@ "PolicyAttachmentCreateRequest": { "description": "Request body for creating a policy attachment.", "properties": { + "default": { + "default": false, + "description": "Apply this attachment only when no non-default attachment matches the request.", + "title": "Default", + "type": "boolean" + }, "keys": { "anyOf": [ { @@ -35113,6 +35119,12 @@ "description": "Who created the attachment.", "title": "Created By" }, + "default": { + "default": false, + "description": "Apply this attachment only when no non-default attachment matches the request.", + "title": "Default", + "type": "boolean" + }, "definition_location": { "default": "db", "description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).", @@ -37141,6 +37153,12 @@ "PolicyAttachmentCreateRequest": { "description": "Request body for creating a policy attachment.", "properties": { + "default": { + "default": false, + "description": "Apply this attachment only when no non-default attachment matches the request.", + "title": "Default", + "type": "boolean" + }, "keys": { "anyOf": [ { diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e415a78f412..44d45dcd687 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3216,7 +3216,9 @@ def _match_and_track_policies( attachment_registry: Final = ( attachment_registry_override if attachment_registry_override is not None else get_attachment_registry() ) - matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(context) + matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons( + context, PolicyMatcher.policy_applies(context, policies_override) + ) matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons] policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 3735c335bd4..d81471b3c1a 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -5,6 +5,7 @@ Attachments define WHERE policies apply, separate from the policy definitions. This allows the same policy to be attached to multiple scopes. """ +from collections.abc import Callable from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypedDict @@ -119,35 +120,49 @@ class AttachmentRegistry: models=attachment_data.get("models"), tags=attachment_data.get("tags"), priority=attachment_data.get("priority"), + default=attachment_data.get("default", False), ) - def get_attached_policies(self, context: PolicyMatchContext) -> list[str]: + def get_attached_policies( + self, + context: PolicyMatchContext, + policy_applies: Callable[[str], bool] | None = None, + ) -> list[str]: """ Get list of policy names attached to the given context. Args: context: The request context to match against + policy_applies: Optional predicate; attachments whose policy does not apply are ignored Returns: List of policy names that are attached to matching scopes """ - return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)] + return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, policy_applies)] - def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[PolicyAttachmentMatch]: + def get_attached_policies_with_reasons( + self, + context: PolicyMatchContext, + policy_applies: Callable[[str], bool] | None = None, + ) -> list[PolicyAttachmentMatch]: """ Get list of policy names and match reasons for the given context. Returns a list of dicts with 'policy_name' and 'matched_via' keys. The 'matched_via' describes which dimension caused the match. + Attachments whose policy fails `policy_applies` are dropped before defaults are considered. """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher + in_scope: Final = tuple( + attachment + for attachment in self._attachments + if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) + and (policy_applies is None or policy_applies(attachment.policy)) + ) + non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default) matching_attachments: Final = sorted( - ( - attachment - for attachment in self._attachments - if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) - ), + non_default or tuple(attachment for attachment in in_scope if attachment.default), key=_attachment_sort_key, ) broadest_attachment_by_policy: Final = MappingProxyType( @@ -169,6 +184,11 @@ class AttachmentRegistry: @staticmethod def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: """Describe why an attachment matched the context.""" + reason: Final = AttachmentRegistry._describe_scope_match(attachment, context) + return f"default:{reason}" if attachment.default else reason + + @staticmethod + def _describe_scope_match(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher if attachment.is_global(): @@ -324,6 +344,7 @@ class AttachmentRegistry: "models": attachment_request.models or [], "tags": attachment_request.tags or [], "priority": attachment_request.priority, + "is_default": attachment_request.default, "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), "created_by": created_by, @@ -340,6 +361,7 @@ class AttachmentRegistry: models=attachment_request.models, tags=attachment_request.tags, priority=attachment_request.priority, + default=attachment_request.default, ) self.add_attachment(attachment) @@ -352,6 +374,7 @@ class AttachmentRegistry: models=created_attachment.models or [], tags=created_attachment.tags or [], priority=created_attachment.priority, + default=created_attachment.is_default, created_at=created_attachment.created_at, updated_at=created_attachment.updated_at, created_by=created_attachment.created_by, @@ -429,6 +452,7 @@ class AttachmentRegistry: models=attachment.models or [], tags=attachment.tags or [], priority=attachment.priority, + default=attachment.is_default, created_at=attachment.created_at, updated_at=attachment.updated_at, created_by=attachment.created_by, @@ -468,6 +492,7 @@ class AttachmentRegistry: models=a.models or [], tags=a.tags or [], priority=a.priority, + default=a.is_default, created_at=a.created_at, updated_at=a.updated_at, created_by=a.created_by, @@ -502,6 +527,7 @@ class AttachmentRegistry: models=(attachment_response.models if attachment_response.models else None), tags=attachment_response.tags if attachment_response.tags else None, priority=attachment_response.priority, + default=attachment_response.default, ) for attachment_response in attachments ] diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 1e30238c8b4..f4b38bea14e 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -61,6 +61,7 @@ def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment) models=attachment.models or [], tags=attachment.tags or [], priority=attachment.priority, + default=attachment.default, definition_location="config", ) diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index 001e4115374..e0f558b5085 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -7,6 +7,7 @@ apply to a given request based on team alias, key alias, and model. Policies are matched via policy_attachments which define WHERE each policy applies. """ +from collections.abc import Callable, Sequence from typing import Final from litellm._logging import verbose_proxy_logger @@ -113,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( @@ -130,9 +131,31 @@ class PolicyMatcher: """ return PolicyMatcher.get_matching_policies(context=context) + @staticmethod + def policy_applies( + context: PolicyMatchContext, + 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=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: list[str], + policy_names: Sequence[str], context: PolicyMatchContext, policies: dict[str, Policy] | None = None, ) -> list[str]: @@ -152,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 diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index a8a9856b833..898e42635c5 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -265,7 +265,9 @@ async def resolve_policies_for_context( ) # Get matching policies with reasons - match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(context=context) + match_results: Final = get_attachment_registry().get_attached_policies_with_reasons( + context=context, policy_applies=PolicyMatcher.policy_applies(context) + ) if not match_results: return PolicyResolveResponse( diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index d284c44397e..0f373b08056 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -84,7 +84,9 @@ def _retrieval_context( def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: - matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context) + matches: Final = get_attachment_registry().get_attached_policies_with_reasons( + context, PolicyMatcher.policy_applies(context) + ) if not matches: return (), MappingProxyType({}) applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d2032cec0d0..2d7e557a9d1 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable { models String[] @default([]) // Model names or patterns tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) priority Int? // Explicit execution order + is_default Boolean @default(false) // Applied only when no non-default attachment matches created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt diff --git a/litellm/types/proxy/policy_engine/policy_types.py b/litellm/types/proxy/policy_engine/policy_types.py index 66e5fbb4b49..73eeffa3585 100644 --- a/litellm/types/proxy/policy_engine/policy_types.py +++ b/litellm/types/proxy/policy_engine/policy_types.py @@ -294,6 +294,10 @@ class PolicyAttachment(BaseModel): le=2147483647, description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.", ) + default: bool = Field( + default=False, + description="Apply this attachment only when no non-default attachment matches the request.", + ) model_config = ConfigDict(extra="forbid") diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index e6f501ed4b5..ebdedb98b12 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -311,6 +311,10 @@ class PolicyAttachmentCreateRequest(BaseModel): le=2147483647, description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.", ) + default: bool = Field( + default=False, + description="Apply this attachment only when no non-default attachment matches the request.", + ) class PolicyAttachmentDBResponse(BaseModel): @@ -327,6 +331,10 @@ class PolicyAttachmentDBResponse(BaseModel): default=None, description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.", ) + default: bool = Field( + default=False, + description="Apply this attachment only when no non-default attachment matches the request.", + ) 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.") diff --git a/schema.prisma b/schema.prisma index d2032cec0d0..2d7e557a9d1 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable { models String[] @default([]) // Model names or patterns tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"]) priority Int? // Explicit execution order + is_default Boolean @default(false) // Applied only when no non-default attachment matches created_at DateTime @default(now()) created_by String? updated_at DateTime @default(now()) @updatedAt 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 089bec59583..faa8d67fe3a 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -14,7 +14,8 @@ from litellm.proxy.policy_engine.attachment_registry import ( AttachmentRegistry, get_attachment_registry, ) -from litellm.types.proxy.policy_engine import PolicyMatchContext +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails, PolicyMatchContext class TestGetAttachedPolicies: @@ -30,9 +31,7 @@ class TestGetAttachedPolicies: ) # Should match any context - context = PolicyMatchContext( - team_alias="any-team", key_alias="any-key", model="any-model" - ) + context = PolicyMatchContext(team_alias="any-team", key_alias="any-key", model="any-model") attached = registry.get_attached_policies(context) assert "global-baseline" in attached @@ -46,15 +45,11 @@ class TestGetAttachedPolicies: ) # Match - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") assert "healthcare-policy" in registry.get_attached_policies(context) # No match - different team - context_other = PolicyMatchContext( - team_alias="finance-team", key_alias="key", model="gpt-4" - ) + context_other = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4") assert "healthcare-policy" not in registry.get_attached_policies(context_other) def test_key_wildcard_pattern_attachment(self): @@ -67,15 +62,11 @@ class TestGetAttachedPolicies: ) # Match - key starts with dev-key- - context = PolicyMatchContext( - team_alias="team", key_alias="dev-key-123", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="team", key_alias="dev-key-123", model="gpt-4") assert "dev-policy" in registry.get_attached_policies(context) # No match - different prefix - context_prod = PolicyMatchContext( - team_alias="team", key_alias="prod-key-123", model="gpt-4" - ) + context_prod = PolicyMatchContext(team_alias="team", key_alias="prod-key-123", model="gpt-4") assert "dev-policy" not in registry.get_attached_policies(context_prod) def test_model_specific_attachment(self): @@ -92,9 +83,7 @@ class TestGetAttachedPolicies: assert "gpt4-policy" in registry.get_attached_policies(context) # No match - context_other = PolicyMatchContext( - team_alias="team", key_alias="key", model="gpt-3.5" - ) + context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-3.5") assert "gpt4-policy" not in registry.get_attached_policies(context_other) def test_model_wildcard_pattern(self): @@ -107,15 +96,11 @@ class TestGetAttachedPolicies: ) # Match - context = PolicyMatchContext( - team_alias="team", key_alias="key", model="bedrock/claude-3" - ) + context = PolicyMatchContext(team_alias="team", key_alias="key", model="bedrock/claude-3") assert "bedrock-policy" in registry.get_attached_policies(context) # No match - context_other = PolicyMatchContext( - team_alias="team", key_alias="key", model="openai/gpt-4" - ) + context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="openai/gpt-4") assert "bedrock-policy" not in registry.get_attached_policies(context_other) def test_multiple_attachments_match_same_context(self): @@ -129,9 +114,7 @@ class TestGetAttachedPolicies: ] ) - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") attached = registry.get_attached_policies(context) # All three should match @@ -277,9 +260,7 @@ class TestGetAttachedPolicies: ] ) - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") attached = registry.get_attached_policies(context) # Should only appear once @@ -288,9 +269,7 @@ class TestGetAttachedPolicies: 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)] - ) + 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() @@ -318,9 +297,7 @@ class TestGetAttachedPolicies: ] ) - context = PolicyMatchContext( - team_alias="finance-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4") attached = registry.get_attached_policies(context) assert attached == [] @@ -338,23 +315,15 @@ class TestGetAttachedPolicies: ) # Match - both team and model match - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-4" - ) + context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4") assert "strict-policy" in registry.get_attached_policies(context) # No match - team matches but model doesn't - context_wrong_model = PolicyMatchContext( - team_alias="healthcare-team", key_alias="key", model="gpt-3.5" - ) - assert "strict-policy" not in registry.get_attached_policies( - context_wrong_model - ) + context_wrong_model = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-3.5") + assert "strict-policy" not in registry.get_attached_policies(context_wrong_model) # No match - model matches but team doesn't - context_wrong_team = PolicyMatchContext( - team_alias="finance-team", key_alias="key", model="gpt-4" - ) + context_wrong_team = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4") assert "strict-policy" not in registry.get_attached_policies(context_wrong_team) @@ -527,6 +496,111 @@ class TestMatchAttribution: assert "catch-all" in attached +class TestDefaultAttachments: + """`default: true` attachments apply only when no non-default attachment matches.""" + + @staticmethod + def _registry() -> AttachmentRegistry: + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "guardrail-y", "scope": "*", "default": True}, + {"policy": "guardrail-x", "tags": ["opt-in"]}, + ] + ) + return registry + + def test_opted_in_request_gets_only_the_opt_in_policy(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + + assert self._registry().get_attached_policies(context) == ["guardrail-x"] + + def test_request_without_opt_in_falls_back_to_default_policy(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2") + + assert self._registry().get_attached_policies(context) == ["guardrail-y"] + + def test_default_attachment_still_honors_its_own_scope(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "team-default", "teams": ["team-a"], "default": True}]) + + assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-a", key_alias="k", model="m")) == [ + "team-default" + ] + assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-b", key_alias="k", model="m")) == [] + + def test_all_matching_defaults_apply_when_nothing_else_matches(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "default-a", "scope": "*", "default": True}, + {"policy": "default-b", "teams": ["team-a"], "default": True}, + {"policy": "opt-in", "tags": ["opt-in"]}, + ] + ) + context = PolicyMatchContext(team_alias="team-a", key_alias="k", model="m") + + assert registry.get_attached_policies(context) == ["default-a", "default-b"] + + def test_non_default_attachments_remain_additive(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "baseline", "scope": "*"}, + {"policy": "opt-in", "tags": ["opt-in"]}, + {"policy": "fallback", "scope": "*", "default": True}, + ] + ) + context = PolicyMatchContext(team_alias="t", key_alias="k", model="m", tags=["opt-in"]) + + assert registry.get_attached_policies(context) == ["baseline", "opt-in"] + + def test_default_match_reason_is_labelled(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="m") + + results = self._registry().get_attached_policies_with_reasons(context) + + assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}] + + def test_inapplicable_opt_in_policy_does_not_suppress_default(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + policies = { + "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])), + "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="claude.*")), + } + + results = self._registry().get_attached_policies_with_reasons( + context, PolicyMatcher.policy_applies(context, policies) + ) + + assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}] + + def test_attachment_to_missing_policy_does_not_suppress_default(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + policies = {"guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"]))} + + assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [ + "guardrail-y" + ] + + def test_applicable_opt_in_policy_still_wins_with_predicate(self): + context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) + policies = { + "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])), + "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="gpt.*")), + } + + assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [ + "guardrail-x" + ] + + def test_default_defaults_to_false_when_omitted(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "p"}]) + + assert registry.get_all_attachments()[0].default is False + + class TestAttachmentRegistrySingleton: """Test global singleton behavior.""" @@ -557,6 +631,7 @@ def _make_db_attachment_row( scope: str | None = None, teams: list[str] | None = None, priority: int | None = None, + is_default: bool = False, ) -> MagicMock: row = MagicMock() row.attachment_id = attachment_id @@ -567,6 +642,7 @@ def _make_db_attachment_row( row.models = [] row.tags = [] row.priority = priority + row.is_default = is_default row.created_at = datetime.now(timezone.utc) row.updated_at = datetime.now(timezone.utc) row.created_by = None @@ -576,9 +652,7 @@ def _make_db_attachment_row( def _prisma_with_attachment_rows(rows: list[MagicMock]) -> MagicMock: prisma = MagicMock() - prisma.configure_mock( - **{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)} - ) + prisma.configure_mock(**{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)}) return prisma @@ -629,6 +703,15 @@ class TestConfigAttachmentsPreservedAcrossDbSync: assert registry.get_all_attachments()[0].priority == 7 + @pytest.mark.asyncio + async def test_sync_round_trips_db_attachment_default_flag(self): + registry = AttachmentRegistry() + db_row = _make_db_attachment_row(is_default=True) + + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row])) + + assert registry.get_all_attachments()[0].default is True + @pytest.mark.asyncio async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self): registry = AttachmentRegistry() diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py index 6143898ccbe..b07137893ec 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_matcher.py @@ -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 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx index 43ad6a7cc9e..be83f73bb2e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx @@ -65,6 +65,19 @@ describe("AttachmentTable", () => { ); }); + it("should show a Default badge only for default attachments", () => { + const attachments = [ + makeAttachment({ attachment_id: "att-def00001", policy_name: "fallback", default: true }), + makeAttachment({ attachment_id: "att-def00002", policy_name: "regular" }), + ]; + renderWithProviders(); + const rows = screen.getAllByRole("row").slice(1); + const fallbackRow = rows.find((row) => within(row).queryByText("fallback")); + const regularRow = rows.find((row) => within(row).queryByText("regular")); + expect(within(fallbackRow!).getByText("Default")).toBeInTheDocument(); + expect(within(regularRow!).queryByText("Default")).not.toBeInTheDocument(); + }); + it("should show skeleton rows when isLoading is true", () => { renderWithProviders(); expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx index 9a190401d08..3265b9db834 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx @@ -181,6 +181,20 @@ export const getAttachmentTableColumns = ({ {row.original.priority} ), }, + { + id: "default", + accessorFn: (row) => (row.default ? 1 : 0), + meta: { title: "Default" }, + header: ({ column }) => , + size: 100, + enableSorting: true, + cell: ({ row }) => + row.original.default ? ( + + ) : ( + - + ), + }, { id: "created_at", accessorFn: (row) => row.created_at ?? "", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx index dfc023d428e..14af4a2b8f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx @@ -237,6 +237,21 @@ describe("AddAttachmentForm", () => { expect(createAttachment).toHaveBeenCalledWith("test-token", { policy_name: "policy-alpha", scope: "*" }); }); + it("sends default: true when the Default switch is turned on", async () => { + const user = userEvent.setup(); + const createAttachment = vi.fn().mockResolvedValue({}); + renderWithProviders(); + await selectPolicy(user, "policy-alpha"); + await user.click(screen.getByRole("switch", { name: /default/i })); + await submit(user); + await waitFor(() => expect(createAttachment).toHaveBeenCalledTimes(1)); + expect(createAttachment).toHaveBeenCalledWith("test-token", { + policy_name: "policy-alpha", + scope: "*", + default: true, + }); + }); + it.each([ ["2147483648", /at most 2147483647/i], ["-2147483649", /at least -2147483648/i], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx index 02463a89139..5cd240a0838 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx @@ -11,6 +11,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { Separator } from "@/components/ui/separator"; +import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; @@ -38,6 +39,7 @@ interface AttachmentFormValues { models: string[]; tags: string[]; priority: number | null; + default: boolean; } const EMPTY_VALUES: AttachmentFormValues = { @@ -47,6 +49,7 @@ const EMPTY_VALUES: AttachmentFormValues = { models: [], tags: [], priority: null, + default: false, }; const INT32_MIN = -2147483648; @@ -64,6 +67,7 @@ const attachmentShape = { .min(INT32_MIN, `Priority must be at least ${INT32_MIN}`) .max(INT32_MAX, `Priority must be at most ${INT32_MAX}`) .nullable(), + default: z.boolean(), }; const buildAttachmentSchema = (scopeType: ScopeType, teamsLoaded: boolean, availableTeams: string[]) => @@ -453,9 +457,23 @@ const AddAttachmentForm: React.FC = ({ /> )} + + + {({ value, onChange, ref, ...field }) => ( + + )} + - {impactResult && } + {impactResult && }
); -const ImpactPreviewAlert: React.FC = ({ impactResult }) => { +const ImpactPreviewAlert: React.FC = ({ impactResult, isDefault = false }) => { const isGlobal = impactResult.affected_keys_count === -1; + const qualifier = isDefault ? "up to " : ""; return ( @@ -47,7 +49,7 @@ const ImpactPreviewAlert: React.FC = ({ impactResult }) ) : (
- This attachment would affect{" "} + This attachment would affect {qualifier} {impactResult.affected_keys_count} key{impactResult.affected_keys_count !== 1 ? "s" : ""} {" "} @@ -57,6 +59,11 @@ const ImpactPreviewAlert: React.FC = ({ impactResult }) . + {isDefault && ( +
+ Default attachments only apply to requests no non-default attachment matches, so fewer may be affected. +
+ )} {impactResult.sample_keys.length > 0 && (