From 334d09b3b21c728e6f7152d1804cb8b5aac643ef Mon Sep 17 00:00:00 2001 From: YutaSaito <36355491+uc4w6c@users.noreply.github.com> Date: Fri, 28 Nov 2025 14:26:27 +0900 Subject: [PATCH] feat: add regex-based tool_name/tool_type matching for tool-permission (#17164) * feat: add regex-based tool_name/tool_type matching for tool-permission * docs: update tool permission quick start for UI workflow --- .../docs/proxy/guardrails/tool_permission.md | 45 +++++- .../guardrail_hooks/tool_permission.py | 139 +++++++++++------ .../guardrail_hooks/tool_permission.py | 33 +++- .../guardrail_hooks/test_tool_permission.py | 144 ++++++++++++------ .../ToolPermissionRulesEditor.tsx | 36 ++++- 5 files changed, 286 insertions(+), 111 deletions(-) diff --git a/docs/my-website/docs/proxy/guardrails/tool_permission.md b/docs/my-website/docs/proxy/guardrails/tool_permission.md index 19b674c9e55..897c31d9dab 100644 --- a/docs/my-website/docs/proxy/guardrails/tool_permission.md +++ b/docs/my-website/docs/proxy/guardrails/tool_permission.md @@ -7,9 +7,38 @@ import TabItem from '@theme/TabItem'; LiteLLM provides the LiteLLM Tool Permission Guardrail that lets you control which **tool calls** a model is allowed to invoke, using configurable allow/deny rules. This offers fine-grained, provider-agnostic control over tool execution (e.g., OpenAI Chat Completions `tool_calls`, Anthropic Messages `tool_use`, MCP tools). ## Quick Start -### 1. Define Guardrails on your LiteLLM config.yaml -Define your guardrails under the `guardrails` section +### LiteLLM UI + +#### Step 1: Select Tool Permission Guardrail + +Open the LiteLLM Dashboard, click **Add New Guardrail**, and choose **LiteLLM Tool Permission Guardrail**. This loads the rule builder UI. + +Configure tool permission guardrail in LiteLLM UI + +#### Step 2: Define Regex Rules + +1. Click **Add Rule**. +2. Enter a unique Rule ID. +3. Provide a regex for the tool name (e.g., `^mcp__github_.*$`). +4. Optionally add a regex for tool type (e.g., `^function$`). +5. Pick **Allow** or **Deny**. + +Configure tool permission guardrail in LiteLLM UI + +#### Step 3: Restrict Tool Arguments (Optional) + +Select **+ Restrict tool arguments** to attach regex validations to nested paths (dot + `[]` notation). This enforces that sensitive parameters (such as `arguments.to[]`) conform to pre-approved formats. + +#### Step 4: Choose Defaults & Actions + +- Set the fallback decision (`default_action`) for tools that do not hit any rule. +- Decide how disallowed tools behave: **Block** halts the request, **Rewrite** strips forbidden tools and returns an error message inside the response. +- Customize `violation_message_template` if you want branded error copy. +- Save the guardrail. + +### LiteLLM Config.yaml Setup + ```yaml guardrails: - guardrail_name: "tool-permission-guardrail" @@ -21,16 +50,17 @@ guardrails: tool_name: "Bash" decision: "allow" - id: "allow_github_mcp" - tool_name: "mcp__github_*" + tool_name: "^mcp__github_.*$" decision: "allow" - id: "allow_aws_documentation" - tool_name: "mcp__aws-documentation_*_documentation" + tool_name: "^mcp__aws-documentation_.*_documentation$" decision: "allow" - id: "deny_read_commands" tool_name: "Read" - decision: "Deny" + decision: "deny" - id: "mail-domain" - tool_name: "send_email" + tool_name: "^send_email$" + tool_type: "^function$" decision: "allow" allowed_param_patterns: "to[]": "^.+@berri\\.ai$" @@ -44,7 +74,8 @@ guardrails: ```yaml - id: "unique_rule_id" # Unique identifier for the rule - tool_name: "pattern" # Tool name or pattern to match + tool_name: "^regex$" # Regex for tool name (optional, at least one of name/type required) + tool_type: "^function$" # Regex for tool type (optional) decision: "allow" # "allow" or "deny" allowed_param_patterns: # Optional - regex map for argument paths (dot + [] notation) "path.to[].field": "^regex$" diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 02e06acbda4..64753d9fa85 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -62,6 +62,7 @@ class ToolPermissionGuardrail(CustomGuardrail): 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): @@ -70,6 +71,30 @@ class ToolPermissionGuardrail(CustomGuardrail): 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(): @@ -100,59 +125,75 @@ class ToolPermissionGuardrail(CustomGuardrail): return ToolPermissionGuardrailConfigModel - def _matches_pattern(self, tool_name: str, pattern: str) -> bool: - """ - Check if a tool name matches a pattern - - Supports patterns like: - - "Bash" - exact match - - "mcp__*" - prefix pattern (matches names starting wich "mcp__") - - "*_read" - suffix wildcard (matches names ending with "_read") - - "mcp__github_*_read" - infix wildcard (matches names like "mcp__github_mark_all_notifications_read") - - Args: - tool_name: Name of the tool to check - pattern: Pattern to match against - - Returns: - True if the tool name matches the pattern - """ - # Handle exact matches - if tool_name == pattern: + def _matches_regex( + self, pattern: Optional[re.Pattern], value: Optional[str] + ) -> bool: + if pattern is None: return True + if value is None: + return False + return bool(pattern.fullmatch(value)) - if "*" in pattern: - # Escape regex special chars except '*' - escaped_pattern = re.escape(pattern) - # Turn \* into .* - regex_pattern = escaped_pattern.replace(r"\*", ".*") - return bool(re.fullmatch(regex_pattern, tool_name)) + def _rule_matches_tool( + self, + rule: ToolPermissionRule, + *, + tool_name: Optional[str], + tool_type: Optional[str] = None, + ) -> tuple[bool, bool]: + target_patterns = self._compiled_rule_targets.get(rule.id, {}) + name_pattern = target_patterns.get("tool_name") + type_pattern = target_patterns.get("tool_type") - return False + name_required = rule.tool_name is not None + type_required = rule.tool_type is not None + + name_matched = ( + self._matches_regex(name_pattern, tool_name) if name_required else True + ) + type_matched = ( + self._matches_regex(type_pattern, tool_type) if type_required else True + ) + + overall_match = name_matched and type_matched + should_check_params = name_required and name_matched + + return overall_match, should_check_params def _check_tool_permission( - self, tool_name: str + self, + tool_name: Optional[str], + tool_type: Optional[str] = None, ) -> tuple[bool, Optional[str], Optional[str]]: """ Check if a tool is allowed based on the configured rules Args: tool_name: Name of the tool to check + tool_type: Type of the tool to check Returns: Tuple of (is_allowed, rule_id, message) """ - verbose_proxy_logger.debug(f"Checking permission for tool: {tool_name}") + verbose_proxy_logger.debug( + f"Checking permission for tool: {tool_name or tool_type}" + ) # Check each rule in order for rule in self.rules: - if self._matches_pattern(tool_name, rule.tool_name): + matches, _ = self._rule_matches_tool( + rule, + tool_name=tool_name, + tool_type=tool_type, + ) + if matches: is_allowed = rule.decision == "allow" - default_message = f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" + tool_identifier = tool_name or tool_type or "unknown_tool" + default_message = f"Tool '{tool_identifier}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" message = self.render_violation_message( default=default_message, context={ - "tool_name": tool_name, + "tool_name": tool_name or tool_identifier, "rule_id": rule.id, }, ) @@ -161,11 +202,12 @@ class ToolPermissionGuardrail(CustomGuardrail): # No rule matched, use default action is_allowed = self.default_action == "allow" - default_message = f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by default action" + tool_identifier = tool_name or tool_type or "unknown_tool" + default_message = f"Tool '{tool_identifier}' {'allowed' if is_allowed else 'denied'} by default action" message = self.render_violation_message( default=default_message, context={ - "tool_name": tool_name, + "tool_name": tool_name or tool_identifier, "rule_id": None, }, ) @@ -228,7 +270,7 @@ class ToolPermissionGuardrail(CustomGuardrail): *, arguments: Dict[str, Any], rule: ToolPermissionRule, - tool_name: str, + tool_name: Optional[str], ) -> tuple[bool, Optional[str]]: compiled_patterns = self._compiled_rule_patterns.get(rule.id) if not compiled_patterns: @@ -249,7 +291,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return ( False, f"Value '{raw_value}' for path '{path}' does not match allowed pattern" - f" '{compiled_pattern.pattern}' for tool '{tool_name}'", + f" '{compiled_pattern.pattern}' for tool '{tool_name or 'unknown_tool'}'", ) return True, None @@ -258,19 +300,27 @@ class ToolPermissionGuardrail(CustomGuardrail): self, tool_call: ChatCompletionMessageToolCall ) -> tuple[bool, Optional[str], Optional[str]]: tool_name = tool_call.function.name if tool_call.function else None - if not tool_name: + tool_type = getattr(tool_call, "type", None) + if not tool_name and not tool_type: return self.default_action == "allow", None, None + tool_identifier = tool_name or tool_type or "unknown_tool" + last_pattern_failure_msg: Optional[str] = None for rule in self.rules: - if not self._matches_pattern(tool_name, rule.tool_name): + matches, should_check_params = self._rule_matches_tool( + rule, + tool_name=tool_name, + tool_type=tool_type, + ) + if not matches: continue - if rule.allowed_param_patterns: + if rule.allowed_param_patterns and should_check_params: arguments = self._parse_tool_call_arguments(tool_call) if not arguments: - last_pattern_failure_msg = f"Tool '{tool_name}' is missing arguments required by rule '{rule.id}'" + last_pattern_failure_msg = f"Tool '{tool_identifier}' is missing arguments required by rule '{rule.id}'" continue patterns_match, failure_message = self._patterns_match_for_rule( @@ -283,10 +333,10 @@ class ToolPermissionGuardrail(CustomGuardrail): continue is_allowed = rule.decision == "allow" - default_message = f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" + default_message = f"Tool '{tool_identifier}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" message = self.render_violation_message( default=default_message, - context={"tool_name": tool_name, "rule_id": rule.id}, + context={"tool_name": tool_identifier, "rule_id": rule.id}, ) return is_allowed, rule.id, message @@ -294,11 +344,11 @@ class ToolPermissionGuardrail(CustomGuardrail): default_message = ( last_pattern_failure_msg if (last_pattern_failure_msg and not is_allowed) - else f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by default action" + else f"Tool '{tool_identifier}' {'allowed' if is_allowed else 'denied'} by default action" ) message = self.render_violation_message( default=default_message, - context={"tool_name": tool_name, "rule_id": None}, + context={"tool_name": tool_identifier, "rule_id": None}, ) return is_allowed, None, message @@ -469,8 +519,9 @@ class ToolPermissionGuardrail(CustomGuardrail): if tool["type"] != "function": continue tool_name: str = tool["function"]["name"] + tool_type: Optional[str] = tool.get("type") - is_allowed, _, message = self._check_tool_permission(tool_name) + is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) if not is_allowed and message is not None: verbose_proxy_logger.warning(f"Tool Permission Guardrail: {message}") diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py index e78cfad8bdb..2ed1f3d2e3a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -1,7 +1,7 @@ # Tool Permission Guardrail Type Definitions from typing import Dict, List, Literal, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator, model_validator from .base import GuardrailConfigModel @@ -12,8 +12,13 @@ class ToolPermissionRule(BaseModel): """ id: str = Field(description="Unique identifier for the rule") - tool_name: str = Field( - description="Tool name or pattern (e.g., 'Bash', 'mcp__github_*', 'mcp__github_*_read', '*_read')" + tool_name: Optional[str] = Field( + default=None, + description="Regex pattern applied to the tool's function name", + ) + tool_type: Optional[str] = Field( + default=None, + description="Regex pattern applied to the tool type (e.g., function)", ) decision: Literal["allow", "deny"] = Field( description="Whether to allow or deny this tool usage" @@ -23,6 +28,26 @@ class ToolPermissionRule(BaseModel): description="Optional regex map enforcing nested parameter values using dot/[] paths", ) + @field_validator("tool_name", "tool_type", mode="before") + @classmethod + def _blank_to_none(cls, value: Optional[str]) -> Optional[str]: + if value is None: + return None + if isinstance(value, str): + stripped = value.strip() + if not stripped: + return None + return stripped + return value + + @model_validator(mode="after") + def _ensure_target_present(self): + if self.tool_name is None and self.tool_type is None: + raise ValueError( + "Each rule must specify at least a tool_name or tool_type regex" + ) + return self + class ToolResult(BaseModel): """ @@ -52,7 +77,7 @@ class ToolPermissionGuardrailConfigModel(GuardrailConfigModel): rules: Optional[List[ToolPermissionRule]] = Field( default=None, - description="Ordered allow/deny rules. Patterns support * wildcards and optional regex constraints on tool arguments.", + description="Ordered allow/deny rules. Patterns use regex for tool names/types and optional regex constraints on tool arguments.", ) default_action: Literal["allow", "deny"] = Field( default="deny", description="Fallback decision when no rule matches" 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 5468dcf9491..a7fd1c64955 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 @@ -3,6 +3,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ import os +import re import sys from unittest.mock import patch @@ -36,15 +37,19 @@ class TestToolPermissionGuardrail: def setup_method(self): """Set up test fixtures""" self.test_rules = [ - {"id": "allow_bash", "tool_name": "Bash", "decision": "allow"}, - {"id": "allow_github", "tool_name": "mcp__github_*", "decision": "allow"}, + {"id": "allow_bash", "tool_name": r"^Bash$", "decision": "allow"}, { - "id": "allow_documentation", - "tool_name": "mcp__aws-documentation_*_documentation", + "id": "allow_github", + "tool_name": r"^mcp__github_.*$", "decision": "allow", }, - {"id": "deny_read", "tool_name": "Read", "decision": "deny"}, - {"id": "deny_get", "tool_name": "*_get", "decision": "deny"}, + { + "id": "allow_documentation", + "tool_name": r"^mcp__aws-documentation_.*_documentation$", + "decision": "allow", + }, + {"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}, + {"id": "deny_get", "tool_name": r".*_get$", "decision": "deny"}, ] self.guardrail = ToolPermissionGuardrail( @@ -64,50 +69,91 @@ class TestToolPermissionGuardrail: self.guardrail.supported_event_hooks or [] ) - def test_pattern_matching_exact(self): - """Test exact pattern matching""" - assert self.guardrail._matches_pattern("Read", "Read") is True - assert self.guardrail._matches_pattern("Write", "Read") is False + def test_matches_regex_helper(self): + pattern = re.compile(r"^Read$") + assert self.guardrail._matches_regex(pattern, "Read") is True + assert self.guardrail._matches_regex(pattern, "Write") is False + assert self.guardrail._matches_regex(None, "Any") is True + assert self.guardrail._matches_regex(pattern, None) is False - def test_pattern_matching_wildcards(self): - """Test wildcard pattern matching""" - assert ( - self.guardrail._matches_pattern( - "mcp__github_add_issue_comment", "mcp__github_*" - ) - is True + def test_rule_matches_tool_with_type_only(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="type-only", + rules=[ + { + "id": "allow_functions", + "tool_type": r"^function$", + "decision": "allow", + } + ], + default_action="deny", + on_disallowed_action="block", ) - assert ( - self.guardrail._matches_pattern( - "mcp__github_add_issue_comment", "mcp__github_*_comment" - ) - is True + + is_allowed, rule_id, _ = guardrail._check_tool_permission("AnyTool", "function") + assert is_allowed is True + assert rule_id == "allow_functions" + + is_allowed, rule_id, _ = guardrail._check_tool_permission("AnyTool", "custom") + assert is_allowed is False + assert rule_id is None + + def test_rule_matches_tool_with_name_and_type(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="name-type", + rules=[ + { + "id": "allow_specific", + "tool_name": r"^Bash$", + "tool_type": r"^function$", + "decision": "allow", + } + ], + default_action="deny", + on_disallowed_action="block", ) - assert ( - self.guardrail._matches_pattern( - "mcp__github_add_issue_comment", "*_comment" + + is_allowed, rule_id, _ = guardrail._check_tool_permission("Bash", "function") + assert is_allowed is True + assert rule_id == "allow_specific" + + is_allowed, rule_id, _ = guardrail._check_tool_permission("Bash", "custom") + assert is_allowed is False + assert rule_id is None + + def test_rule_requires_name_or_type(self): + with pytest.raises(ValueError): + ToolPermissionGuardrail( + guardrail_name="invalid-rule", + rules=[{"id": "no_target", "decision": "allow"}], + default_action="deny", + on_disallowed_action="block", ) - is True + + def test_type_only_rule_skips_param_patterns(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="type-param", + rules=[ + { + "id": "allow_type_only", + "tool_type": r"^function$", + "decision": "allow", + "allowed_param_patterns": {"foo": r"^bar$"}, + } + ], + default_action="deny", + on_disallowed_action="block", ) - assert ( - self.guardrail._matches_pattern( - "mcp__git_add_issue_comment", "mcp__github_*" - ) - is False - ) - assert ( - self.guardrail._matches_pattern( - "mcp__github_assign_copilot_to_issue", "mcp__github_*_comment" - ) - is False - ) - assert ( - self.guardrail._matches_pattern( - "mcp__github_assign_copilot_to_issue", "*_comment" - ) - is False + + tool_call = ChatCompletionMessageToolCall( + function={"name": "AnyTool", "arguments": "{}"}, + type="function", ) + is_allowed, rule_id, _ = guardrail._get_permission_for_tool_call(tool_call) + assert is_allowed is True + assert rule_id == "allow_type_only" + def test_check_tool_permission_allow(self): is_allowed, rule_id, msg = self.guardrail._check_tool_permission("Bash") assert is_allowed is True @@ -232,7 +278,7 @@ class TestToolPermissionGuardrail: rules=[ { "id": "allow_mail", - "tool_name": "mail_mcp-send_email", + "tool_name": r"^mail_mcp-send_email$", "decision": "allow", "allowed_param_patterns": { "to[]": r"^.+@berri\.ai$", @@ -267,7 +313,7 @@ class TestToolPermissionGuardrail: rules=[ { "id": "allow_mail", - "tool_name": "mail_mcp-send_email", + "tool_name": r"^mail_mcp-send_email$", "decision": "allow", "allowed_param_patterns": {"to[]": r"^.+@berri\.ai$"}, } @@ -300,7 +346,7 @@ class TestToolPermissionGuardrail: rules=[ { "id": "allow_mail", - "tool_name": "mail_mcp-send_email", + "tool_name": r"^mail_mcp-send_email$", "decision": "allow", "allowed_param_patterns": {"to[]": r"^.+@berri\.ai$"}, } @@ -339,7 +385,7 @@ class TestToolPermissionGuardrail: rules=[ { "id": "deny_gmail", - "tool_name": "mail_mcp-send_email", + "tool_name": r"^mail_mcp-send_email$", "decision": "deny", "allowed_param_patterns": {"to[]": r"^.+@gmail\.com$"}, } @@ -486,7 +532,9 @@ class TestToolPermissionGuardrailIntegration: def test_default_action_allow(self): guardrail = ToolPermissionGuardrail( guardrail_name="test-allow-default", - rules=[{"id": "deny_read", "tool_name": "Read", "decision": "deny"}], + rules=[ + {"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"} + ], default_action="allow", ) diff --git a/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.tsx b/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.tsx index 790876ed3f0..68bbc055f45 100644 --- a/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.tsx @@ -9,7 +9,8 @@ export type ToolPermissionOnDisallowedAction = "block" | "rewrite"; export interface ToolPermissionRuleConfig { id: string; - tool_name: string; + tool_name?: string; + tool_type?: string; decision: ToolPermissionDecision; allowed_param_patterns?: Record; } @@ -67,7 +68,6 @@ const ToolPermissionRulesEditor: React.FC = ({ ...config.rules, { id: `rule_${Math.random().toString(36).slice(2, 8)}`, - tool_name: "", decision: "allow" as ToolPermissionDecision, allowed_param_patterns: undefined, }, @@ -195,8 +195,8 @@ const ToolPermissionRulesEditor: React.FC = ({
LiteLLM Tool Permission Guardrail - Use wildcards (e.g., mcp__github_*) to scope which tools can run and optionally constrain - payload fields. + Provide regex patterns (e.g., ^mcp__github_.*$) for tool names or types and optionally + constrain payload fields.
{!disabled && ( @@ -242,12 +242,32 @@ const ToolPermissionRulesEditor: React.FC = ({ />
- Tool Name / Pattern + Tool Name (optional) updateRule(index, { tool_name: e.target.value })} + placeholder="^mcp__github_.*$" + value={rule.tool_name ?? ""} + onChange={(e) => + updateRule(index, { + tool_name: e.target.value.trim() === "" ? undefined : e.target.value, + }) + } + /> +
+ + +
+
+ Tool Type (optional) + + updateRule(index, { + tool_type: e.target.value.trim() === "" ? undefined : e.target.value, + }) + } />