fix(guardrails): recompile ToolPermissionGuardrail rules on update_in_memory_litellm_params

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,102 @@ 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.
+        """
+        # ``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():
+                setattr(self, key, value)
+        else:
+            super().update_in_memory_litellm_params(litellm_params)
+            params = vars(litellm_params)
+
+        # 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()
+        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 (

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,115 @@ 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
+
+    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
+        )