diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 37be832d350..90152f86911 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,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 ( 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..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 @@ -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 + )