diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 6d1096d5ee9..8d5d8116919 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -220,6 +220,7 @@ class AttachmentRegistry: attachment: PolicyAttachment object to add """ self._attachments.append(attachment) + self._initialized = True verbose_proxy_logger.debug(f"Added attachment for policy: {attachment.policy}") def remove_attachments_for_policy(self, policy_name: str) -> int: diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index d3df16afde6..75017c46603 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -226,6 +226,7 @@ class PolicyRegistry: policy: Policy object to add """ self._policies[policy_name] = policy + self._initialized = True verbose_proxy_logger.debug(f"Added/updated policy: {policy_name}") def remove_policy(self, policy_name: str) -> bool: diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6be1a9ecef0..72d04e8e871 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2846,6 +2846,59 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po assert "metadata" in data +@pytest.mark.asyncio +async def test_api_created_global_policy_applies_to_new_key_without_restart(): + """ + Regression: policies created at runtime via policy builder must apply + immediately when attached globally, even if the server started with no + initialized policy config. + """ + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import ( + Policy, + PolicyAttachment, + PolicyGuardrails, + ) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + } + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + policy_registry = get_policy_registry() + attachment_registry = get_attachment_registry() + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + try: + policy_registry.add_policy( + "runtime-global-policy", + Policy(guardrails=PolicyGuardrails(add=["runtime-guardrail"])), + ) + attachment_registry.add_attachment( + PolicyAttachment(policy="runtime-global-policy", scope="*") + ) + + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=user_api_key_dict, + ) + + assert "runtime-guardrail" in data["metadata"]["guardrails"] + assert "runtime-global-policy" in data["metadata"]["applied_policies"] + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_policy_version_by_id(): """