Fix runtime policy attachment initialization

Mark runtime-created policies and attachments initialized so global policy attachments created from the policy builder apply immediately without requiring a restart.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
shivam 2026-05-01 17:11:57 -07:00
parent 0b87bdff3d
commit 536a24c5ca
No known key found for this signature in database
3 changed files with 55 additions and 0 deletions

View file

@ -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:

View file

@ -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:

View file

@ -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():
"""