From 0729a7cac44aa694a2ee2e64e21ee97e6f6d5c5d Mon Sep 17 00:00:00 2001 From: DineshGirbide Date: Thu, 4 Jun 2026 13:10:13 +0530 Subject: [PATCH 1/3] fix(guardrails): recompile ToolPermissionGuardrail rules on update_in_memory_litellm_params ToolPermissionGuardrail builds self.rules and the compiled target/pattern maps only in __init__. The base update_in_memory_litellm_params re-sets raw attributes via setattr but never rebuilds those maps, so a guardrail updated in place (PUT /guardrails, or the immediate in-memory sync) keeps enforcing the construction-time rules until it is reinitialized (PATCH path, periodic DB poll, or restart). Extract the compile step into _load_rules and override update_in_memory_litellm_params to rebuild from it (dict- and model-safe), re-normalizing default_action / on_disallowed_action. Mirrors the existing PresidioGuardrail override of the same method. Adds regression tests. Fixes #29592. --- .../guardrail_hooks/tool_permission.py | 129 +++++++++++------- .../guardrail_hooks/test_tool_permission.py | 81 ++++++++++- 2 files changed, 161 insertions(+), 49 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 37be832d350..b72da2b2887 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -16,7 +16,7 @@ from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ToolPermissionRule, @@ -60,53 +60,7 @@ def __init__( super().__init__(**kwargs) - self.rules: List[ToolPermissionRule] = [] - self._compiled_rule_patterns: Dict[str, Dict[str, re.Pattern]] = {} - self._compiled_rule_targets: Dict[str, Dict[str, Optional[re.Pattern]]] = {} - if rules: - for rule_item in rules: - if isinstance(rule_item, ToolPermissionRule): - rule = rule_item - else: - rule = ToolPermissionRule(**rule_item) - self.rules.append(rule) - - compiled_target_patterns: Dict[str, Optional[re.Pattern]] = { - "tool_name": None, - "tool_type": None, - } - if rule.tool_name is not None: - try: - compiled_target_patterns["tool_name"] = re.compile( - rule.tool_name - ) - except re.error as exc: - raise ValueError( - f"Invalid regex for tool_name in rule '{rule.id}': {exc}" - ) from exc - if rule.tool_type is not None: - try: - compiled_target_patterns["tool_type"] = re.compile( - rule.tool_type - ) - except re.error as exc: - raise ValueError( - f"Invalid regex for tool_type in rule '{rule.id}': {exc}" - ) from exc - self._compiled_rule_targets[rule.id] = compiled_target_patterns - - if rule.allowed_param_patterns: - compiled_patterns: Dict[str, re.Pattern] = {} - for path, pattern in rule.allowed_param_patterns.items(): - try: - compiled_patterns[path] = re.compile(pattern) - except re.error as exc: - raise ValueError( - f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}" - ) from exc - - if compiled_patterns: - self._compiled_rule_patterns[rule.id] = compiled_patterns + self._load_rules(rules) # Normalize to lowercase for case-insensitive handling self.default_action = ( @@ -126,6 +80,85 @@ def __init__( self.default_action, ) + def _load_rules(self, rules: Optional[List[Any]]) -> None: + """Parse ``rules`` and (re)build the compiled target/pattern lookups. + + ``self.rules`` plus ``_compiled_rule_targets`` / ``_compiled_rule_patterns`` + are the state every matching path reads. Centralizing the build here lets + both ``__init__`` and ``update_in_memory_litellm_params`` recompile from a + single source of truth, so an in-place update (PUT /guardrails, immediate + sync) reflects rule changes instead of keeping the construction-time maps. + """ + self.rules = [] + self._compiled_rule_patterns = {} + self._compiled_rule_targets = {} + if not rules: + return + + for rule_item in rules: + rule = ( + rule_item + if isinstance(rule_item, ToolPermissionRule) + else ToolPermissionRule(**rule_item) + ) + self.rules.append(rule) + + compiled_target_patterns: Dict[str, Optional[re.Pattern]] = { + "tool_name": None, + "tool_type": None, + } + if rule.tool_name is not None: + try: + compiled_target_patterns["tool_name"] = re.compile(rule.tool_name) + except re.error as exc: + raise ValueError( + f"Invalid regex for tool_name in rule '{rule.id}': {exc}" + ) from exc + if rule.tool_type is not None: + try: + compiled_target_patterns["tool_type"] = re.compile(rule.tool_type) + except re.error as exc: + raise ValueError( + f"Invalid regex for tool_type in rule '{rule.id}': {exc}" + ) from exc + self._compiled_rule_targets[rule.id] = compiled_target_patterns + + if rule.allowed_param_patterns: + compiled_patterns: Dict[str, re.Pattern] = {} + for path, pattern in rule.allowed_param_patterns.items(): + try: + compiled_patterns[path] = re.compile(pattern) + except re.error as exc: + raise ValueError( + f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}" + ) from exc + if compiled_patterns: + self._compiled_rule_patterns[rule.id] = compiled_patterns + + def update_in_memory_litellm_params( + self, litellm_params: Union[LitellmParams, dict] + ) -> None: + """Apply updated params in place, rebuilding the compiled rule state. + + The base implementation only ``setattr``s raw fields, which would leave + ``_compiled_rule_targets`` / ``_compiled_rule_patterns`` (built in + ``__init__``) stale, so a guardrail updated without reinitialization would + keep enforcing the old ruleset. Recompile here so PUT /guardrails and the + immediate in-memory sync take effect, mirroring the PresidioGuardrail + override of this method. + """ + super().update_in_memory_litellm_params(litellm_params) + params = ( + litellm_params if isinstance(litellm_params, dict) else vars(litellm_params) + ) + self._load_rules(params.get("rules")) + default_action = params.get("default_action") + if isinstance(default_action, str): + self.default_action = default_action.lower() + on_disallowed_action = params.get("on_disallowed_action") + if isinstance(on_disallowed_action, str): + self.on_disallowed_action = on_disallowed_action.lower() + @staticmethod def get_config_model(): from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 716b4470d25..68c0a3260c3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -2,6 +2,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ +import json import os import re import sys @@ -20,7 +21,7 @@ from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ( ToolPermissionGuardrail, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) @@ -850,3 +851,81 @@ def test_case_insensitive_decision_in_rules(self): is_allowed, rule_id, _ = guardrail._check_tool_permission("Read") assert is_allowed is False assert rule_id == "deny_read" + + +class TestToolPermissionGuardrailInMemoryUpdate: + """Regression: an in-memory params update (PUT /guardrails path) must rebuild + the compiled rule maps, not just self.rules, so the new rules are enforced + without reinitializing the guardrail.""" + + def _bash(self, command): + return ChatCompletionMessageToolCall( + function={"name": "Bash", "arguments": json.dumps({"command": command})}, + type="function", + ) + + def test_update_in_memory_recompiles_added_param_pattern(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="tp", + rules=[{"id": "native-bash", "tool_name": r"^Bash$", "decision": "allow"}], + default_action="deny", + on_disallowed_action="block", + ) + # No pattern yet: any Bash command is allowed. + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0] + is True + ) + + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="tool_permission", + mode=["pre_call", "post_call"], + default_action="deny", + on_disallowed_action="block", + rules=[ + { + "id": "native-bash", + "tool_name": r"^Bash$", + "decision": "allow", + "allowed_param_patterns": { + "command": r"^(?!(echo blockme)$).*$" + }, + } + ], + ) + ) + + # The compiled map must be rebuilt, and enforcement must reflect it. + assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {}) + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0] + is False + ) + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo hello"))[0] is True + ) + + def test_update_in_memory_recompiles_tool_name_target(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="tp", + rules=[], + default_action="allow", + on_disallowed_action="block", + ) + # No rules: default_action allow lets Bash through. + assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is True + + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="tool_permission", + mode=["pre_call", "post_call"], + default_action="allow", + on_disallowed_action="block", + rules=[{"id": "deny-bash", "tool_name": r"^Bash$", "decision": "deny"}], + ) + ) + + # A newly added deny rule (new id) must match -> its compiled target was rebuilt. + assert "deny-bash" in guardrail._compiled_rule_targets + assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is False From d22634595887c3983820d9d9c56c8b455b8072f3 Mon Sep 17 00:00:00 2001 From: DineshGirbide Date: Thu, 4 Jun 2026 13:22:30 +0530 Subject: [PATCH 2/3] fix(guardrails): handle dict params in ToolPermissionGuardrail in-memory update Delegate to super() only for LitellmParams input (the base setattr loop is model-only); apply the raw-dict case inline. Fixes the mypy arg-type error and makes the recompile work when the proxy passes the raw DB dict. --- .../guardrails/guardrail_hooks/tool_permission.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index b72da2b2887..fb39a366735 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -147,10 +147,17 @@ def update_in_memory_litellm_params( immediate in-memory sync take effect, mirroring the PresidioGuardrail override of this method. """ - super().update_in_memory_litellm_params(litellm_params) - params = ( - litellm_params if isinstance(litellm_params, dict) else vars(litellm_params) - ) + # ``litellm_params`` may arrive as the raw DB dict (the proxy ``cast()``s + # it to ``LitellmParams`` without converting), so handle both shapes. The + # base ``setattr`` loop is model-only, so apply the dict case here. + if isinstance(litellm_params, dict): + params = litellm_params + for key, value in params.items(): + setattr(self, key, value) + else: + super().update_in_memory_litellm_params(litellm_params) + params = vars(litellm_params) + self._load_rules(params.get("rules")) default_action = params.get("default_action") if isinstance(default_action, str): From a1fdbe1217a8e02b131145e10be702a1aff41902 Mon Sep 17 00:00:00 2001 From: DineshGirbide Date: Thu, 4 Jun 2026 14:44:28 +0530 Subject: [PATCH 3/3] fix(guardrails): preserve tool-permission rules on a partial in-memory update A partial update (e.g. a LitellmParams whose rules field is None) ran through the generic setattr, which set self.rules to None, and the recompile was skipped, leaving the guardrail with no rules. Snapshot the previous rules and restore them when the update carries no rules; an explicit empty list still clears them. Adds a regression test for the rules-absent case. Addresses the Greptile review note on #29655. --- .../guardrail_hooks/tool_permission.py | 12 ++++++- .../guardrail_hooks/test_tool_permission.py | 34 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index fb39a366735..90152f86911 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -150,6 +150,7 @@ def update_in_memory_litellm_params( # ``litellm_params`` may arrive as the raw DB dict (the proxy ``cast()``s # it to ``LitellmParams`` without converting), so handle both shapes. The # base ``setattr`` loop is model-only, so apply the dict case here. + previous_rules = self.rules if isinstance(litellm_params, dict): params = litellm_params for key, value in params.items(): @@ -158,7 +159,16 @@ def update_in_memory_litellm_params( super().update_in_memory_litellm_params(litellm_params) params = vars(litellm_params) - self._load_rules(params.get("rules")) + # The generic update above sets ``self.rules`` from the incoming value + # (None on a partial update that omits rules), but never rebuilds the + # compiled maps. Rebuild them when rules are provided; otherwise restore + # the previous ruleset so a partial update doesn't silently wipe it. An + # explicit empty list still clears the rules. + rules = params.get("rules") + if rules is not None: + self._load_rules(rules) + else: + self.rules = previous_rules default_action = params.get("default_action") if isinstance(default_action, str): self.default_action = default_action.lower() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 68c0a3260c3..48516666120 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -929,3 +929,37 @@ def test_update_in_memory_recompiles_tool_name_target(self): # A newly added deny rule (new id) must match -> its compiled target was rebuilt. assert "deny-bash" in guardrail._compiled_rule_targets assert guardrail._get_permission_for_tool_call(self._bash("echo x"))[0] is False + + def test_update_in_memory_preserves_rules_when_rules_absent(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="tp", + rules=[ + { + "id": "native-bash", + "tool_name": r"^Bash$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^(?!(echo blockme)$).*$"}, + } + ], + default_action="deny", + on_disallowed_action="block", + ) + assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {}) + + # A partial update that does not carry `rules` must NOT wipe the existing + # ruleset / compiled maps. + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="tool_permission", + mode=["pre_call", "post_call"], + default_action="deny", + on_disallowed_action="block", + ) + ) + + assert len(guardrail.rules) == 1 + assert "command" in guardrail._compiled_rule_patterns.get("native-bash", {}) + assert ( + guardrail._get_permission_for_tool_call(self._bash("echo blockme"))[0] + is False + )