diff --git a/litellm/litellm_core_utils/tool_fields.py b/litellm/litellm_core_utils/tool_fields.py index 5ee96ea9133..10b6f07bc4d 100644 --- a/litellm/litellm_core_utils/tool_fields.py +++ b/litellm/litellm_core_utils/tool_fields.py @@ -10,7 +10,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params ROUTES_WITH_TOOLS: Final = frozenset( { @@ -58,6 +58,19 @@ _TOOL_ENTRY: Final = TypeAdapter(_ToolEntry) _WEB_SEARCH_OPTIONS: Final = TypeAdapter(dict[str, object] | None, config=ConfigDict(strict=True)) +def declared_function_name(tool: object) -> str | None: + fields: Final = as_str_mapping(tool) + if fields is None: + return None + function: Final = as_str_mapping(fields.get("function")) + name: Final = (fields if function is None else function).get("name") + return name if isinstance(name, str) and name else None + + +def qualified_member_name(namespace: str, name: str) -> str: + return f"{namespace}__{name}" + + @dataclass(frozen=True, slots=True) class DropPolicy: drop_params: bool diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 32d7b39a20e..80244cfc3eb 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -5,9 +5,11 @@ import json import re import threading import time +from collections import Counter from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from functools import lru_cache, partial +from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, Union, cast, get_args @@ -34,7 +36,11 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.litellm_core_utils.prompt_templates.factory import encode_tool_call_id_with_signature from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning -from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp +from litellm.litellm_core_utils.tool_fields import ( + declared_function_name, + qualified_member_name, + tool_params_without_hosted_mcp, +) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -544,6 +550,138 @@ def mapped_gemini_speech_config( } +def _as_chat_tool(member: JsonValue) -> JsonValue: + if not isinstance(member, dict) or "function" in member or "name" not in member: + return member + return { + "type": member.get("type", "function"), + "function": {key: value for key, value in member.items() if key != "type"}, + } + + +def _namespace_members_or_tool(tool: JsonValue) -> tuple[JsonValue, ...]: + if not isinstance(tool, dict) or tool.get("type") != "namespace": + return (tool,) + members: Final = tool.get("tools") + return _members_named_once(tuple(map(_as_chat_tool, members))) if isinstance(members, list) else () + + +def _members_named_once( + members: tuple[JsonValue, ...], repeatable: frozenset[str] = frozenset() +) -> tuple[JsonValue, ...]: + names: Final = tuple(map(declared_function_name, members)) + first_index: Final = {name: index for index, name in reversed(tuple(enumerate(names)))} + return tuple( + member + for index, (member, name) in enumerate(zip(members, names, strict=True)) + if name is None or name in repeatable or first_index[name] == index + ) + + +def _is_namespace(tool: JsonValue) -> bool: + return isinstance(tool, dict) and tool.get("type") == "namespace" + + +def _top_level_name(tool: JsonValue) -> str | None: + name: Final = tool.get("name") if isinstance(tool, dict) else None + return name if isinstance(name, str) and name else None + + +def _listed(tools: JsonValue) -> tuple[JsonValue, ...]: + return tuple(tools) if isinstance(tools, list) else () + + +def _declarable_names(tool: JsonValue) -> tuple[str | None, ...]: + members: Final = _namespace_members_or_tool(tool) if _is_namespace(tool) else () + return (declared_function_name(tool), *map(declared_function_name, members)) + + +def _declarable_names_in(tools: tuple[JsonValue, ...]) -> tuple[str | None, ...]: + return tuple(chain.from_iterable(map(_declarable_names, tools))) + + +def _inlines_members(tool: JsonValue, declarable: Counter[str | None]) -> bool: + members: Final = _namespace_members_or_tool(tool) + return _is_namespace(tool) and all( + name is None or declarable[name] == 1 for name in map(declared_function_name, members) + ) + + +def _qualified_member(namespace: str, member: JsonValue) -> JsonValue: + function: Final = member.get("function") if isinstance(member, dict) else None + name: Final = declared_function_name(member) + if not isinstance(member, dict) or not isinstance(function, dict) or name is None: + return member + return {**member, "function": {**function, "name": qualified_member_name(namespace, name)}} + + +def _as_declared( + tool: JsonValue, declarable: Counter[str | None], top_level_names: frozenset[str] +) -> tuple[JsonValue, ...]: + if _inlines_members(tool, declarable): + return _namespace_members_or_tool(tool) + namespace: Final = _top_level_name(tool) + if not _is_namespace(tool) or namespace is None: + return (tool,) + qualified: Final = (_qualified_member(namespace, member) for member in _namespace_members_or_tool(tool)) + return tuple(member for member in qualified if declared_function_name(member) not in top_level_names) + + +def _top_level_names_in(tools: tuple[JsonValue, ...]) -> frozenset[str]: + top_level_tools: Final = (tool for tool in tools if not _is_namespace(tool)) + return frozenset(name for name in map(declared_function_name, top_level_tools) if name is not None) + + +def _with_namespace_members_inlined(tools: JsonValue, outside_declarations: JsonValue = None) -> JsonValue: + if not isinstance(tools, list) or not any(map(_is_namespace, tools)): + return tools + candidates: Final = (*tools, *_listed(outside_declarations)) + declarable: Final = Counter(_declarable_names_in(candidates)) + top_level_names: Final = _top_level_names_in(candidates) + declared: Final = chain.from_iterable(_as_declared(tool, declarable, top_level_names) for tool in tools) + return list(_members_named_once(tuple(declared), top_level_names)) + + +def _names_each_tool_declares(candidates: tuple[JsonValue, ...]) -> tuple[tuple[str, ...], ...]: + declarable: Final = Counter(_declarable_names_in(candidates)) + top_level_names: Final = _top_level_names_in(candidates) + declared: Final = tuple( + tuple(filter(None, map(declared_function_name, _as_declared(tool, declarable, top_level_names)))) + for tool in candidates + ) + owned_names: Final = chain.from_iterable(zip(repeat(index), names) for index, names in enumerate(declared)) + first_owner: Final = {name: index for index, name in reversed(tuple(owned_names))} + return tuple(tuple(name for name in names if first_owner[name] == index) for index, names in enumerate(declared)) + + +def _qualified_names_of( + member_name: str, tool: JsonValue, declarable: Counter[str | None], declared: tuple[str, ...] +) -> tuple[str, ...]: + namespace: Final = _top_level_name(tool) + if not _is_namespace(tool) or namespace is None or _inlines_members(tool, declarable): + return () + qualified: Final = qualified_member_name(namespace, member_name) + return (qualified,) if qualified in declared else () + + +def _forced_function_names(name: str, tools: JsonValue) -> tuple[str, ...]: + candidates: Final = _listed(tools) + if any(not _is_namespace(tool) and declared_function_name(tool) == name for tool in candidates): + return (name,) + declarable: Final = Counter(_declarable_names_in(candidates)) + declared_by_tool: Final = tuple(zip(candidates, _names_each_tool_declares(candidates), strict=True)) + named: Final = next( + (declared for tool, declared in declared_by_tool if _is_namespace(tool) and _top_level_name(tool) == name), + None, + ) + if named is not None: + return named + qualified: Final = chain.from_iterable( + _qualified_names_of(name, tool, declarable, declared) for tool, declared in declared_by_tool + ) + return tuple(qualified) or (name,) + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -815,7 +953,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return supported_params - def map_tool_choice_values(self, model: str, tool_choice: str | dict) -> ToolConfig | None: + def map_tool_choice_values( + self, model: str, tool_choice: str | dict, tools: JsonValue = None, drop_params: bool = False + ) -> ToolConfig | None: if tool_choice == "none": return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="NONE")) elif tool_choice == "required": @@ -825,7 +965,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif isinstance(tool_choice, dict): # only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html name: Final = tool_choice.get("function", {}).get("name", "") - return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY", allowed_function_names=[name])) + forced: Final = _forced_function_names(name, tools) + if forced: + return ToolConfig( + functionCallingConfig=FunctionCallingConfig(mode="ANY", allowed_function_names=list(forced)) + ) + if drops_unsupported_params(drop_params): + return None + raise litellm.utils.UnsupportedParamsError( + message=f"VertexAI can't force tool_choice '{name}', since none of that namespace's members can be sent to Gemini as a function. Name one of its functions instead, or set `litellm.drop_params = True` to drop tool_choice.", + status_code=400, + ) else: raise litellm.utils.UnsupportedParamsError( message=f"VertexAI doesn't support tool_choice={tool_choice}. Supported tool_choice values=['auto', 'required', json object]. To drop it from the call, set `litellm.drop_params = True.", @@ -1068,7 +1218,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return googleSearch, googleSearchRetrieval, enterpriseWebSearch, urlContext - def _map_function(self, value: list[dict], optional_params: dict, send_integer_enums: bool = False) -> list[Tools]: + def _map_function( + self, + value: list[dict], + optional_params: dict, + send_integer_enums: bool = False, + outside_declarations: JsonValue = None, + ) -> list[Tools]: """ Map OpenAI-style tools/functions to Vertex AI format. @@ -1099,7 +1255,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # remove 'strict' from tools remove_strict_from_schema(schema_value) - for tool in cast(list[dict], schema_value): # cast-ok: a deep copy keeps the caller's tool shape + inlined_tools: Final = _with_namespace_members_inlined(schema_value, outside_declarations) + for tool in cast(list[dict], inlined_tools): # cast-ok: a deep copy keeps the caller's tool shape openai_function_object: ChatCompletionToolParamFunctionChunk | None = None if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk(**tool["function"]) @@ -1734,7 +1891,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif (param == "tools" or param == "functions") and isinstance(value, list) and value: # Pass optional_params so _map_function can add toolConfig if needed mapped_tools = with_integer_enums_that_fit( - partial(self._map_function, value, optional_params), + partial( + self._map_function, + value, + optional_params, + outside_declarations=request_params.get("functions" if param == "tools" else "tools"), + ), allow_integer_enums and self._is_gemini_3_or_newer(model), ) optional_params = self._add_tools_to_optional_params(optional_params, mapped_tools) @@ -1742,6 +1904,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tool_choice_value = self.map_tool_choice_values( model=model, tool_choice=value, + tools=[*_listed(non_default_params.get("tools")), *_listed(non_default_params.get("functions"))], + drop_params=drop_params, ) if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b290adc47ce..ab85b6f8842 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -20233,7 +20233,7 @@ "type": "null" } ], - "description": "Regex pattern applied to the tool's function name", + "description": "Regex pattern applied to the tool's name, or to its type when the tool has no name", "title": "Tool Name" }, "tool_type": { diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 5f9dc4a7a31..5f99b1a4611 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -1,13 +1,16 @@ import json import re -from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence -from typing import Any, Final, Literal, TypedDict +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from functools import partial, reduce +from itertools import chain, groupby +from types import MappingProxyType +from typing import Any, Final, Literal, TypeAlias, TypedDict, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter from typing_extensions import ReadOnly, Required -from litellm import ChatCompletionToolParam from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.completion_extras.litellm_responses_transformation.transformation import ( @@ -19,6 +22,9 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none +from litellm.litellm_core_utils.tool_fields import declared_function_name, qualified_member_name +from litellm.llms.base_llm.guardrail_translation.utils import anthropic_tool_names from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -30,7 +36,10 @@ from litellm.proxy.guardrails.anthropic_sse import ( is_raw_sse_stream, is_sse_error_stream, ) +from litellm.responses.additional_tools import ADDITIONAL_TOOLS_INPUT_ITEM_TYPE from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.llms.anthropic import ANTHROPIC_HOSTED_TOOLS +from litellm.types.llms.vertex_ai import VertexToolName from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ToolPermissionRule, @@ -48,8 +57,48 @@ from litellm.types.utils import ( ) GUARDRAIL_NAME: Final = "tool_permission" +_MCP_TOOLSET_TYPE: Final = "mcp_toolset" +_HOSTED_SERVER_NAME_FIELDS: Final = MappingProxyType({"mcp": "server_label", _MCP_TOOLSET_TYPE: "mcp_server_name"}) +_TOOL_CARRYING_INPUT_TYPES: Final = frozenset({"tool_search_output", ADDITIONAL_TOOLS_INPUT_ITEM_TYPE}) +_DECLARATION_KEYS: Final = ("functionDeclarations", "function_declarations") +_NON_TOOL_PART_KEYS: Final = frozenset({*_DECLARATION_KEYS, "cachePoint", "type"}) +_TOOL_CONFIG_KEYS: Final = ("toolConfig", "tool_config") +_FUNCTION_CALLING_CONFIG_KEYS: Final = ("functionCallingConfig", "function_calling_config") +_ALLOWED_FUNCTION_NAMES_KEYS: Final = ("allowedFunctionNames", "allowed_function_names") +_FORCED_FUNCTION_CALLING_MODES: Final = frozenset({"ANY", "VALIDATED"}) +_NAME_IDENTIFIED_TYPES: Final = frozenset({"function", "custom", "namespace", *_HOSTED_SERVER_NAME_FIELDS}) +_TEXT_EDITOR_TOOL_NAMES: Final = frozenset({"str_replace_editor", "str_replace_based_edit_tool"}) +_TYPE_VERSION_SUFFIX: Final = re.compile(r"_\d{8}$") +_GEMINI_TYPED_HOSTED_TOOLS: Final = ("computer_use", "web_search", "web_search_preview") +_NAMED_CHOICE_TYPES: Final = frozenset({"tool", "function", "custom", "allowed_tools", "any", "auto", "none"}) _DECODED_ARGUMENTS: Final = TypeAdapter(object) _REQUEST_DATA: Final = TypeAdapter(object) +_ToolTarget: TypeAlias = tuple[str | None, str | None] +_JudgedCalls: TypeAlias = Mapping[str, tuple[tuple[str, str], ...]] +_GroupedT: Final = TypeVar("_GroupedT") + + +def _snake_case(name: str) -> str: + return re.sub(r"(? str: + head, *rest = name.split("_") + return head + "".join(part.title() for part in rest) + + +_GEMINI_HOSTED_TOOL_KEYS: Final = frozenset( + chain.from_iterable((_snake_case(tool.value), _camel_case(tool.value)) for tool in VertexToolName) +) + + +def _name_identifies_tool(name: str, tool_type: str) -> bool: + stem: Final = _TYPE_VERSION_SUFFIX.sub("", tool_type) + return ( + tool_type in _NAME_IDENTIFIED_TYPES + or name == stem + or (stem == ANTHROPIC_HOSTED_TOOLS.TEXT_EDITOR.value and name in _TEXT_EDITOR_TOOL_NAMES) + ) def _object_mapping(value: object) -> Mapping[str, object] | None: @@ -68,6 +117,356 @@ def _function_tool_name(tool: object) -> str | None: return name if isinstance(name, str) else None +def _non_empty_string(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _request_tool_names(tool: object) -> tuple[str, ...]: + fields: Final = _object_mapping(tool) + custom: Final = ( + _object_mapping(fields.get("custom")) if fields is not None and fields.get("type") == "custom" else None + ) + custom_name: Final = _non_empty_string(custom.get("name")) if custom is not None else None + server_field: Final = _HOSTED_SERVER_NAME_FIELDS.get(str(fields.get("type"))) if fields is not None else None + server_name: Final = ( + _non_empty_string(fields.get(server_field)) if fields is not None and server_field is not None else None + ) + extra_names: Final = tuple(name for name in (_function_name(tool), custom_name, server_name) if name is not None) + return tuple(dict.fromkeys(anthropic_tool_names(tool) + extra_names)) + + +def _function_name(tool: object) -> str | None: + fields: Final = _object_mapping(tool) + function: Final = _object_mapping(fields.get("function")) if fields is not None else None + return _non_empty_string(function.get("name")) if function is not None else None + + +def _namespace_function_targets(tool: object) -> tuple[_ToolTarget, ...]: + function_name: Final = _function_name(tool) + return ((function_name, "function"),) if function_name is not None else () + + +def _gemini_hosted_key_targets(tool: object) -> tuple[_ToolTarget, ...]: + fields: Final = _object_mapping(tool) + if fields is None or fields.get("type") in _GEMINI_TYPED_HOSTED_TOOLS: + return () + keys: Final = tuple(key for key in fields if key != "type") + return (_part_target(fields, keys[0]),) if len(keys) == 1 and keys[0] in _GEMINI_HOSTED_TOOL_KEYS else () + + +def _is_namespace(tool: object) -> bool: + fields: Final = _object_mapping(tool) + return fields is not None and fields.get("type") == "namespace" + + +def _namespace_members(tool: object) -> Sequence[object]: + fields: Final = _object_mapping(tool) + members: Final = _object_list(fields.get("tools")) if fields is not None and _is_namespace(tool) else None + return members or () + + +def _returned_namespace(tool_call: ChatCompletionMessageToolCall) -> str | None: + tags: Final = as_str_mapping(attribute_or_none(tool_call, "provider_specific_fields")) + return _non_empty_string(tags.get("namespace")) if tags is not None else None + + +def _judged_call_names( + tool_call: ChatCompletionMessageToolCall, returned_name: str, judged_names: _JudgedCalls +) -> tuple[tuple[str, str], ...]: + namespace: Final = _returned_namespace(tool_call) + if namespace is not None: + return ((qualified_member_name(namespace, returned_name), returned_name),) + return judged_names.get(returned_name, ()) + + +def _pair_namespace(qualified: str, bare: str) -> str | None: + suffix: Final = qualified_member_name("", bare) + if qualified == bare or not qualified.endswith(suffix): + return None + return _non_empty_string(qualified.removesuffix(suffix)) + + +def _qualified_target(namespace: str | None, target: _ToolTarget) -> _ToolTarget: + name, tool_type = target + return (qualified_member_name(namespace, name), tool_type) if namespace is not None and name is not None else target + + +def _member_name(member: object) -> str | None: + fields: Final = _object_mapping(member) + return _non_empty_string(fields.get("name")) if fields is not None else None + + +def _declaration_targets(declaration: object) -> tuple[_ToolTarget, ...]: + name: Final = _member_name(declaration) + return ((name, "function"),) if name is not None else () + + +@dataclass(frozen=True, slots=True) +class _MemberGroup: + key: str + members: Sequence[object] + targets_of: Callable[[object], tuple[_ToolTarget, ...]] + + +def _member_group( + tool: object, namespace_member_targets: Callable[[str | None, object], tuple[_ToolTarget, ...]] +) -> _MemberGroup | None: + fields: Final = _object_mapping(tool) + if fields is None: + return None + members: Final = _namespace_members(tool) + if members: + return _MemberGroup("tools", members, partial(namespace_member_targets, _member_name(tool))) + key: Final = next((key for key in _DECLARATION_KEYS if _object_list(fields.get(key))), None) + declarations: Final = _object_list(fields.get(key)) if key is not None else None + return _MemberGroup(key, declarations, _declaration_targets) if key is not None and declarations else None + + +def _keyed_tool_fields(tool: object) -> Mapping[str, object] | None: + fields: Final = _object_mapping(tool) + if fields is None or isinstance(fields.get("type"), str) or _request_tool_names(tool): + return None + return fields + + +def _part_keys(fields: Mapping[str, object]) -> tuple[str, ...]: + return tuple(key for key in fields if key not in _NON_TOOL_PART_KEYS) + + +def _part_target(fields: Mapping[str, object], key: str) -> _ToolTarget: + return (None if key in _GEMINI_HOSTED_TOOL_KEYS else _member_name(fields.get(key)), key) + + +def _input_item_tools(item: object) -> Sequence[object] | None: + fields: Final = _object_mapping(item) + carries_tools: Final = fields is not None and _non_empty_string(fields.get("type")) in _TOOL_CARRYING_INPUT_TYPES + return _object_list(fields.get("tools")) if fields is not None and carries_tools else None + + +def _request_list(data: object, key: str) -> Sequence[object] | None: + fields: Final = _object_mapping(data) + return _object_list(fields.get(key)) if fields is not None else None + + +def _input_items(data: object) -> Sequence[object]: + return _request_list(data, "input") or () + + +def _tools_in_items(items: Sequence[object]) -> tuple[object, ...]: + return tuple(chain.from_iterable(_input_item_tools(item) or () for item in items)) + + +def _carries_tools(items: Sequence[object]) -> bool: + return any(_input_item_tools(item) is not None for item in items) + + +def _pair_key(pair: tuple[str, object]) -> str: + return pair[0] + + +def _grouped_values(pairs: Iterable[tuple[str, _GroupedT]]) -> Mapping[str, tuple[_GroupedT, ...]]: + groups: Final = groupby(sorted(pairs, key=_pair_key), key=_pair_key) + return MappingProxyType({key: tuple(dict.fromkeys(value for _, value in group)) for key, group in groups}) + + +def _request_layers(data: object) -> tuple[Mapping[str, object], ...]: + fields: Final = _object_mapping(data) + extra_body: Final = _object_mapping(fields.get("extra_body")) if fields is not None else None + return tuple(layer for layer in (fields, extra_body) if layer is not None) + + +def _layer_sent_tools(layer: Mapping[str, object]) -> tuple[object, ...]: + return ( + *(_request_list(layer, "tools") or ()), + *(_request_list(layer, "functions") or ()), + *_tools_in_items(_input_items(layer)), + ) + + +def _sent_tools(data: object) -> tuple[object, ...]: + return tuple(chain.from_iterable(map(_layer_sent_tools, _request_layers(data)))) + + +def _sent_fields(data: object) -> Mapping[str, object]: + return dict(chain.from_iterable(layer.items() for layer in _request_layers(data))) + + +@dataclass(frozen=True, slots=True) +class _KeptLayer: + fields: Mapping[str, object] + tools: tuple[object, ...] | None + functions: tuple[object, ...] | None + mcp_servers: tuple[object, ...] | None + input: tuple[object, ...] | None + + def sent_tools(self) -> tuple[object, ...]: + kept_input: Final = _input_items(self.fields) if self.input is None else self.input + return (*(self.tools or ()), *(self.functions or ()), *_tools_in_items(kept_input)) + + def kept_lists(self) -> dict[str, object]: + lists: Final = ( + ("tools", self.tools), + ("functions", self.functions), + ("mcp_servers", self.mcp_servers), + ("input", self.input), + ) + return {key: list(kept) for key, kept in lists if kept is not None} + + def overlaid_by(self, over: "_KeptLayer") -> "_KeptLayer": + return _KeptLayer( + fields={**self.fields, **over.fields}, + tools=over.tools if "tools" in over.fields else self.tools, + functions=over.functions if "functions" in over.fields else self.functions, + mcp_servers=over.mcp_servers if "mcp_servers" in over.fields else self.mcp_servers, + input=over.input if "input" in over.fields else self.input, + ) + + +def _mcp_server_target(server: object) -> _ToolTarget: + fields: Final = _object_mapping(server) + return (_non_empty_string(fields.get("name")) if fields is not None else None, _MCP_TOOLSET_TYPE) + + +def _reference_name(part: object) -> str | None: + fields: Final = _object_mapping(part) + is_reference: Final = fields is not None and fields.get("type") == "tool_reference" + return _non_empty_string(fields.get("tool_name")) if fields is not None and is_reference else None + + +def _references_one_of(part: object, names: frozenset[str]) -> bool: + name: Final = _reference_name(part) + return name is not None and name in names + + +def _without_references_to(parts: Sequence[object], names: frozenset[str]) -> Sequence[object]: + removed: Final = [name for name in map(_reference_name, parts) if name is not None and name in names] + kept: Final = [part for part in parts if not _references_one_of(part, names)] + if kept or not removed: + return kept + return [{"type": "text", "text": f"Permission denied: tool {', '.join(removed)} is not available"}] + + +def _tool_result_without_references(fields: Mapping[str, object], names: frozenset[str]) -> object: + parts: Final = _object_list(fields.get("content")) + return {**fields, "content": _without_references_to(parts, names)} if parts else fields + + +def _search_result_without_references(fields: Mapping[str, object], names: frozenset[str]) -> object: + result: Final = _object_mapping(fields.get("content")) + references: Final = _object_list(result.get("tool_references")) if result is not None else None + if result is None or references is None: + return fields + kept: Final = [reference for reference in references if not _references_one_of(reference, names)] + return {**fields, "content": {**result, "tool_references": kept}} + + +def _without_references_in_block(block: object, names: frozenset[str]) -> object: + fields: Final = _object_mapping(block) + if fields is None: + return block + match fields.get("type"): + case "tool_result": + return _tool_result_without_references(fields, names) + case "tool_search_tool_result": + return _search_result_without_references(fields, names) + case _: + return block + + +def _without_references_in_message(message: object, names: frozenset[str]) -> object: + fields: Final = _object_mapping(message) + content: Final = _object_list(fields.get("content")) if fields is not None else None + if fields is None or content is None: + return message + if fields.get("role") == "tool": + return {**fields, "content": _without_references_to(content, names)} + return {**fields, "content": [_without_references_in_block(block, names) for block in content]} + + +def _allowed_tool_entries(tool_choice: Mapping[str, object]) -> Sequence[object] | None: + if tool_choice.get("type") != "allowed_tools": + return None + nested: Final = _object_mapping(tool_choice.get("allowed_tools")) + return _object_list((nested if nested is not None else tool_choice).get("tools")) + + +def _is_removed_entry(entry: object, removed_names: frozenset[str], removed_types: frozenset[str]) -> bool: + names: Final = _request_tool_names(entry) + if names: + return any(name in removed_names for name in names) + fields: Final = _object_mapping(entry) + entry_type: Final = _non_empty_string(fields.get("type")) if fields is not None else None + return entry_type is not None and entry_type in removed_types + + +def _with_allowed_tool_entries(tool_choice: Mapping[str, object], entries: list[object]) -> dict[str, object]: + nested: Final = _object_mapping(tool_choice.get("allowed_tools")) + if nested is None: + return {**tool_choice, "tools": entries} + return {**tool_choice, "allowed_tools": {**nested, "tools": entries}} + + +def _is_denied_target( + target: _ToolTarget, denied_targets: frozenset[_ToolTarget], denied_names: frozenset[str] +) -> bool: + return target in denied_targets or target[0] in denied_names + + +def _kept_members(group: _MemberGroup | None, is_denied: Callable[[_ToolTarget], bool]) -> tuple[object, ...]: + if group is None: + return () + return tuple(member for member in group.members if not any(map(is_denied, group.targets_of(member)))) + + +def _none_tool_choice(tool_choice: object) -> object: + choice: Final = _object_mapping(tool_choice) + return {"type": "none"} if choice is not None and choice.get("type") in ("tool", "any") else "none" + + +def _requires_a_tool(tool_choice: object) -> bool: + choice: Final = _object_mapping(tool_choice) + return tool_choice == "required" or (choice is not None and choice.get("type") == "any") + + +def _kept_allowed_function_names(config: Mapping[str, object], removed_names: frozenset[str]) -> Mapping[str, object]: + listed: Final = tuple((key, _object_list(config.get(key))) for key in _ALLOWED_FUNCTION_NAMES_KEYS) + return { + key: [name for name in names if _non_empty_string(name) not in removed_names] + for key, names in listed + if names is not None + } + + +def _rewritten_function_calling_config( + config: Mapping[str, object], removed_names: frozenset[str], functions_remain: bool +) -> Mapping[str, object]: + kept_names: Final = _kept_allowed_function_names(config, removed_names) + had_names: Final = any(_object_list(config.get(key)) for key in _ALLOWED_FUNCTION_NAMES_KEYS) + every_listed_name_removed: Final = had_names and not any(kept_names.values()) + forces_a_function: Final = str(config.get("mode", "")).upper() in _FORCED_FUNCTION_CALLING_MODES + if not every_listed_name_removed and (functions_remain or not forces_a_function): + return {**config, **kept_names} + return {**{key: value for key, value in config.items() if key not in _ALLOWED_FUNCTION_NAMES_KEYS}, "mode": "NONE"} + + +def _rewritten_tool_config_field( + key: str, value: object, removed_names: frozenset[str], functions_remain: bool +) -> object: + config: Final = _object_mapping(value) + if key not in _FUNCTION_CALLING_CONFIG_KEYS or config is None: + return value + return _rewritten_function_calling_config(config, removed_names, functions_remain) + + +def _rewritten_tool_config(tool_config: object, removed_names: frozenset[str], functions_remain: bool) -> object: + fields: Final = _object_mapping(tool_config) + if fields is None: + return tool_config + return { + key: _rewritten_tool_config_field(key, value, removed_names, functions_remain) for key, value in fields.items() + } + + class _ToolPermissionRuleFields(TypedDict, total=False): """The config-file shape a :class:`ToolPermissionRule` is built from.""" @@ -258,7 +657,7 @@ class ToolPermissionGuardrail(CustomGuardrail): name_required: Final = rule.tool_name is not None type_required: Final = rule.tool_type is not None - name_matched: Final = self._matches_regex(name_pattern, tool_name) if name_required else True + name_matched: Final = self._matches_regex(name_pattern, tool_name or tool_type) if name_required else True type_matched: Final = self._matches_regex(type_pattern, tool_type) if type_required else True overall_match: Final = name_matched and type_matched @@ -320,6 +719,72 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug(message) return is_allowed, None, message + def _is_name_rule_for(self, rule: ToolPermissionRule, tool_name: str, tool_type: str | None) -> bool: + return rule.tool_name is not None and self._rule_matches_tool(rule, tool_name=tool_name, tool_type=tool_type)[0] + + def _decides_call( + self, rule: ToolPermissionRule, tool_name: str, tool_call: ChatCompletionMessageToolCall | None + ) -> bool: + if tool_call is None or not rule.allowed_param_patterns: + return True + arguments, parse_error = self._parse_tool_call_arguments(tool_call) + if parse_error or not arguments: + return True + return self._patterns_match_for_rule(arguments=arguments, rule=rule, tool_name=tool_name)[0] + + def _first_name_rule( + self, tool_name: str, tool_type: str | None, tool_call: ChatCompletionMessageToolCall | None + ) -> ToolPermissionRule | None: + return next( + ( + rule + for rule in self.rules + if self._is_name_rule_for(rule, tool_name, tool_type) and self._decides_call(rule, tool_name, tool_call) + ), + None, + ) + + def _narrower_name( + self, + name: str, + other: str, + tool_type: str | None, + tool_call: ChatCompletionMessageToolCall | None, + ) -> str: + rule: Final = self._first_name_rule(name, tool_type, tool_call) + other_rule: Final = self._first_name_rule(other, tool_type, tool_call) + if other_rule is None or other_rule is rule: + return name + if rule is None or self._is_name_rule_for(rule, other, tool_type): + return other + if self._is_name_rule_for(other_rule, name, tool_type): + return name + return other if other_rule.decision == "deny" else name + + def _deciding_member_name( + self, + qualified: str, + bare: str, + tool_type: str | None, + tool_call: ChatCompletionMessageToolCall | None = None, + ) -> str: + member: Final = self._narrower_name(qualified, bare, tool_type, tool_call) + namespace: Final = _pair_namespace(qualified, bare) + return member if namespace is None else self._narrower_name(member, namespace, tool_type, tool_call) + + def _check_target_permission( + self, target: _ToolTarget, bare_names: Mapping[str, tuple[str, ...]] + ) -> tuple[bool, str | None, str | None]: + name, tool_type = target + candidates: Final = tuple(sorted(bare_names.get(name, ()))) if name is not None else () + if name is None or not candidates: + return self._check_tool_permission(name, tool_type) + results: Final = tuple( + self._check_tool_permission(self._deciding_member_name(name, bare_name, tool_type), tool_type) + for bare_name in candidates + ) + return next((result for result in results if not result[0]), results[0]) + def _parse_tool_call_arguments( self, tool_call: ChatCompletionMessageToolCall ) -> tuple[Mapping[str, object] | None, str | None]: @@ -411,9 +876,9 @@ class ToolPermissionGuardrail(CustomGuardrail): return True, None def _get_permission_for_tool_call( - self, tool_call: ChatCompletionMessageToolCall + self, tool_call: ChatCompletionMessageToolCall, judged_name: str | None = None ) -> tuple[bool, str | None, str | None]: - tool_name: Final = tool_call.function.name if tool_call.function else None + tool_name: Final = judged_name or (tool_call.function.name if tool_call.function else None) tool_type: Final = getattr(tool_call, "type", None) if not tool_name and not tool_type: return self.default_action == "allow", None, None @@ -556,10 +1021,46 @@ class ToolPermissionGuardrail(CustomGuardrail): tool_call for block in content if (tool_call := self._anthropic_tool_use_to_tool_call(block)) is not None ) + def _returned_call_names(self, data: Mapping[str, object]) -> _JudgedCalls: + sent_tools: Final = _sent_tools(_REQUEST_DATA.validate_python(data)) + member_names: Final = tuple(chain.from_iterable(map(self._bare_member_names, sent_tools))) + if not member_names: + return MappingProxyType({}) + top_level_tools: Final = (tool for tool in sent_tools if not _is_namespace(tool)) + top_level_names: Final = frozenset(name for name in map(declared_function_name, top_level_tools) if name) + listed: Final = frozenset(name for name, _ in self._collect_request_tools(data) if name is not None) + own_names: Final = frozenset(bare for _, bare in member_names if bare in listed) | top_level_names + returned_members: Final = _grouped_values( + chain( + ((bare, (name, bare)) for name, bare in member_names if bare not in top_level_names), + ((name, (name, bare)) for name, bare in member_names if name not in top_level_names), + ) + ) + return MappingProxyType( + { + name: (*members, *(((name, name),) if name in own_names else ())) + for name, members in returned_members.items() + } + ) + + def _judged_name(self, tool_call: ChatCompletionMessageToolCall, qualified: str, bare: str) -> str: + return qualified if qualified == bare else self._deciding_member_name(qualified, bare, "function", tool_call) + + def _judge_tool_call( + self, tool_call: ChatCompletionMessageToolCall, judged_names: _JudgedCalls + ) -> tuple[bool, str | None, str | None]: + returned_name: Final = tool_call.function.name if tool_call.function else None + names: Final = _judged_call_names(tool_call, returned_name, judged_names) if returned_name is not None else () + if returned_name is None or not names: + return self._get_permission_for_tool_call(tool_call) + judged: Final = dict.fromkeys(self._judged_name(tool_call, qualified, bare) for qualified, bare in names) + results: Final = tuple(self._get_permission_for_tool_call(tool_call, name) for name in judged) + return next((result for result in results if not result[0]), results[0]) + def _evaluate_tool_calls( - self, tool_calls: Sequence[ChatCompletionMessageToolCall] + self, tool_calls: Sequence[ChatCompletionMessageToolCall], judged_names: _JudgedCalls ) -> tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...]: - checked: Final = tuple((tool_call, *self._get_permission_for_tool_call(tool_call)) for tool_call in tool_calls) + checked: Final = tuple((tool_call, *self._judge_tool_call(tool_call, judged_names)) for tool_call in tool_calls) for _tool_call, is_allowed, _rule_id, message in checked: if not is_allowed and message is not None: @@ -619,112 +1120,311 @@ class ToolPermissionGuardrail(CustomGuardrail): if not any(_is_tool_use_block(block) for block in kept_blocks): response["stop_reason"] = "end_turn" # rebind-ok: dropping every tool_use ends the turn - def _get_request_tool_name(self, tool: object) -> tuple[str | None, str | None]: - tool_type: Final = self._get_mapping_value(tool, "type") - if tool_type != "function": - return None, tool_type + def _unqualified_member_targets(self, member: object) -> tuple[_ToolTarget, ...]: + fields: Final = _object_mapping(member) + if fields is None: + return () + targets: Final = self._get_request_tool_targets(member) + own_name: Final = _member_name(member) + judged_names: Final = frozenset(name for name, _ in targets) + own_target: Final = ( + ((own_name, _non_empty_string(fields.get("type"))),) + if own_name is not None and own_name not in judged_names + else () + ) + return targets + own_target - function: Final = self._get_mapping_value(tool, "function") - tool_name: Final = self._get_mapping_value(function, "name") - return tool_name, tool_type + def _namespace_member_targets(self, namespace: str | None, member: object) -> tuple[_ToolTarget, ...]: + return tuple(_qualified_target(namespace, target) for target in self._unqualified_member_targets(member)) - def _get_legacy_function_name(self, function: object) -> str | None: - return self._get_mapping_value(function, "name") + def _bare_member_names(self, tool: object) -> tuple[tuple[str, str], ...]: + namespace: Final = _member_name(tool) + members: Final = _namespace_members(tool) + if namespace is None or not members: + return () + targets: Final = chain.from_iterable(map(self._unqualified_member_targets, members)) + return tuple((qualified_member_name(namespace, name), name) for name, _ in targets if name is not None) - def _get_named_tool_choice(self, data: dict) -> str | None: - tool_choice: Final = data.get("tool_choice") + def _get_request_tool_targets(self, tool: object) -> tuple[_ToolTarget, ...]: + group: Final = _member_group(tool, self._namespace_member_targets) + member_targets: Final = ( + tuple(chain.from_iterable(map(group.targets_of, group.members))) if group is not None else () + ) + keyed: Final = _keyed_tool_fields(tool) + part_targets: Final = tuple(_part_target(keyed, key) for key in _part_keys(keyed)) if keyed is not None else () + return member_targets + part_targets + self._own_tool_targets(tool) + + def _own_tool_targets(self, tool: object) -> tuple[_ToolTarget, ...]: + if _keyed_tool_fields(tool) is not None: + return () + named: Final = _namespace_function_targets(tool) if _namespace_members(tool) else self._named_tool_targets(tool) + return named + _gemini_hosted_key_targets(tool) + + def _named_tool_targets(self, tool: object) -> tuple[_ToolTarget, ...]: + raw_type: Final = self._get_mapping_value(tool, "type") + tool_type: Final = raw_type if isinstance(raw_type, str) else None + anthropic_custom: Final = ( + tool_type == "custom" + and self._get_mapping_value(tool, "input_schema") is not None + and _object_mapping(self._get_mapping_value(tool, "custom")) is None + ) + normalized_type: Final = "function" if tool_type is None or anthropic_custom else tool_type + names: Final = _request_tool_names(tool) + named_targets: Final = tuple((name, normalized_type) for name in names) + if names and all(_name_identifies_tool(name, normalized_type) for name in names): + return named_targets + if names: + return (*named_targets, (None, normalized_type)) + return ((None, tool_type),) if tool_type is not None else () + + def _get_named_tool_choice(self, tool_choice: object) -> str | None: if not tool_choice or tool_choice in ("auto", "none", "required"): return None if isinstance(tool_choice, str): return tool_choice - if self._get_mapping_value(tool_choice, "type") != "function": + choice: Final = _object_mapping(tool_choice) + if choice is None: return None - return self._get_mapping_value(self._get_mapping_value(tool_choice, "function"), "name") + choice_type: Final = choice.get("type") + server_field: Final = _HOSTED_SERVER_NAME_FIELDS.get(str(choice_type)) + if server_field is not None: + return _non_empty_string(choice.get(server_field)) + if choice_type not in ("tool", "function", "custom"): + return None + nested: Final = _object_mapping(choice.get(str(choice_type))) + nested_name: Final = _non_empty_string(nested.get("name")) if nested is not None else None + return nested_name or _non_empty_string(choice.get("name")) - def _get_named_function_call(self, data: dict) -> str | None: + def _get_named_function_call(self, data: Mapping[str, object]) -> str | None: function_call: Final = data.get("function_call") if not function_call or function_call in ("auto", "none"): return None if isinstance(function_call, str): return function_call - return self._get_mapping_value(function_call, "name") + return _non_empty_string(self._get_mapping_value(function_call, "name")) - def _collect_request_tools(self, data: dict) -> list[tuple[str, str | None]]: - request_tools: Final[list[tuple[str, str | None]]] = [] + def _listed_layer_targets(self, layer: Mapping[str, object]) -> tuple[_ToolTarget, ...]: + return ( + *chain.from_iterable(map(self._get_request_tool_targets, _request_list(layer, "tools") or ())), + *map(_mcp_server_target, _request_list(layer, "mcp_servers") or ()), + *chain.from_iterable(map(self._legacy_function_targets, _request_list(layer, "functions") or ())), + ) - for tool in data.get("tools") or []: - tool_name, tool_type = self._get_request_tool_name(tool) - if tool_name is not None: - request_tools.append((tool_name, tool_type)) + def _forced_layer_names(self, layer: Mapping[str, object]) -> tuple[str | None, ...]: + return (self._get_named_tool_choice(layer.get("tool_choice")), self._get_named_function_call(layer)) - for function in data.get("functions") or []: - function_name = self._get_legacy_function_name(function) - if function_name is not None: - request_tools.append((function_name, "function")) + def _collect_request_tools(self, data: Mapping[str, object]) -> tuple[_ToolTarget, ...]: + request: Final = _REQUEST_DATA.validate_python(data) + layers: Final = _request_layers(request) + listed: Final = tuple(chain.from_iterable(map(self._listed_layer_targets, layers))) + listed_names: Final = frozenset(name for name, _ in listed) + forced: Final = chain.from_iterable(map(self._forced_layer_names, layers)) + replayed: Final = map(_function_tool_name, replayed_tool_search_tools(request)) + input_tools: Final = _tools_in_items(tuple(chain.from_iterable(map(_input_items, layers)))) + return ( + *listed, + *((name, "function") for name in forced if name is not None and name not in listed_names), + *((name, "function") for name in replayed if name is not None), + *chain.from_iterable(map(self._get_request_tool_targets, input_tools)), + ) - for forced_tool_name in ( - self._get_named_tool_choice(data), - self._get_named_function_call(data), - ): - if forced_tool_name is not None: - request_tools.append((forced_tool_name, "function")) + def _legacy_function_targets(self, function: object) -> tuple[_ToolTarget, ...]: + targets: Final = self._get_request_tool_targets(function) + as_functions: Final = tuple((name, "function") for name, _ in targets if name is not None) + return tuple(dict.fromkeys(targets + as_functions)) - replayed: Final = map(_function_tool_name, replayed_tool_search_tools(_REQUEST_DATA.validate_python(data))) - request_tools.extend((name, "function") for name in replayed if name is not None) + def _kept_layer( + self, + layer: Mapping[str, object], + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> _KeptLayer: + is_denied: Final = partial(_is_denied_target, denied_targets=denied_targets, denied_names=denied_names) + tools: Final = _request_list(layer, "tools") + functions: Final = _request_list(layer, "functions") + mcp_servers: Final = _request_list(layer, "mcp_servers") + return _KeptLayer( + fields=layer, + tools=None if tools is None else self._kept_tool_list(tools, denied_targets, denied_names), + functions=None if functions is None else self._kept_tool_list(functions, denied_targets, denied_names), + mcp_servers=( + None + if mcp_servers is None + else tuple(server for server in mcp_servers if not is_denied(_mcp_server_target(server))) + ), + input=self._kept_input_list(_input_items(layer), denied_targets, denied_names), + ) - return request_tools + def _rewritten_layer_fields( + self, + kept: _KeptLayer, + rewrite_choice: Callable[[object], object], + removed_names: frozenset[str], + functions_remain: bool, + ) -> dict[str, object]: + layer: Final = kept.fields + choice: Final = layer.get("tool_choice") + relaxes_function_call: Final = self._get_named_function_call(layer) in removed_names + return { + **kept.kept_lists(), + **({"tool_choice": rewrite_choice(choice)} if choice is not None else {}), + **({"function_call": "none"} if relaxes_function_call else {}), + **{ + key: _rewritten_tool_config(layer[key], removed_names, functions_remain) + for key in _TOOL_CONFIG_KEYS + if key in layer + }, + } + + def _tool_names_in(self, tools: Sequence[object]) -> frozenset[str]: + targets: Final = chain.from_iterable(self._get_request_tool_targets(tool) for tool in tools) + return frozenset(name for name, _ in targets if name is not None) + + def _referable_names_in(self, tools: Sequence[object]) -> frozenset[str]: + member_names: Final = chain.from_iterable(map(self._bare_member_names, tools)) + namespace_names: Final = (_member_name(tool) for tool in tools if _namespace_members(tool)) + return ( + self._tool_names_in(tools) + | frozenset(bare for _, bare in member_names) + | frozenset(name for name in namespace_names if name is not None) + ) + + def _kept_tools( + self, + tool: object, + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...]: + is_denied: Final = partial(_is_denied_target, denied_targets=denied_targets, denied_names=denied_names) + group: Final = _member_group(tool, self._namespace_member_targets) + keyed: Final = _keyed_tool_fields(tool) + fields: Final = _object_mapping(tool) + if fields is None or (group is None and keyed is None): + return () if any(map(is_denied, self._get_request_tool_targets(tool))) else (tool,) + own_targets: Final = self._own_tool_targets(tool) + if any(map(is_denied, own_targets)): + return () + kept_members: Final = _kept_members(group, is_denied) + denied_keys: Final = ( + frozenset(key for key in _part_keys(keyed) if is_denied(_part_target(keyed, key))) + if keyed is not None + else frozenset[str]() + ) + if not denied_keys and len(kept_members) == len(group.members if group is not None else ()): + return (tool,) + emptied_group: Final = frozenset({group.key}) if group is not None and not kept_members else frozenset[str]() + dropped_keys: Final = denied_keys | emptied_group + kept_fields: Final = { + key: list(kept_members) if group is not None and key == group.key else value + for key, value in fields.items() + if key not in dropped_keys + } + own_name_left: Final = any(name is not None for name, _ in own_targets) + parts_left: Final = bool(kept_members) or own_name_left or (keyed is not None and bool(_part_keys(kept_fields))) + return (kept_fields,) if parts_left else () + + def _kept_tool_list( + self, + tools: Sequence[object], + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...]: + return tuple(chain.from_iterable(self._kept_tools(tool, denied_targets, denied_names) for tool in tools)) + + def _kept_input_list( + self, + items: Sequence[object], + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...] | None: + if not _carries_tools(items): + return None + return tuple(chain.from_iterable(self._kept_input_items(item, denied_targets, denied_names) for item in items)) + + def _kept_input_items( + self, + item: object, + denied_targets: frozenset[_ToolTarget], + denied_names: frozenset[str], + ) -> tuple[object, ...]: + fields: Final = _object_mapping(item) + tools: Final = _input_item_tools(item) + if fields is None or tools is None: + return (item,) + kept: Final = self._kept_tool_list(tools, denied_targets, denied_names) + emptied: Final = bool(tools) and not kept and fields.get("type") == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE + return () if emptied else ({**fields, "tools": list(kept)},) + + def _tool_types_in(self, tools: Sequence[object]) -> frozenset[str]: + targets: Final = chain.from_iterable(self._get_request_tool_targets(tool) for tool in tools) + return frozenset(tool_type for _, tool_type in targets if tool_type is not None) + + def _rewritten_tool_choice( + self, tool_choice: object, removed_names: frozenset[str], removed_types: frozenset[str], tools_remain: bool + ) -> object: + named: Final = self._get_named_tool_choice(tool_choice) + if named is not None and named in removed_names: + return _none_tool_choice(tool_choice) + choice: Final = _object_mapping(tool_choice) + hosted_type: Final = _non_empty_string(choice.get("type")) if choice is not None else None + if hosted_type not in _NAMED_CHOICE_TYPES and hosted_type in removed_types: + return _none_tool_choice(tool_choice) + allowed: Final = _allowed_tool_entries(choice) if choice is not None else None + if choice is not None and allowed is not None: + kept: Final = [entry for entry in allowed if not _is_removed_entry(entry, removed_names, removed_types)] + return _with_allowed_tool_entries(choice, kept) if kept and tools_remain else _none_tool_choice(choice) + if not tools_remain and (named is not None or _requires_a_tool(tool_choice)): + return _none_tool_choice(tool_choice) + return tool_choice def _modify_request_with_permission_errors( self, data: dict, - denied_tool_names: list[str], + denied_targets: frozenset[_ToolTarget], ): """ - Modify the request to replace denied tool_calls blocks with error results - - Args: - data: The model request to modify - denied_tools: List of (tool_use, error) tuples for denied tools + Remove denied tools from the request, and relax any tool_choice that would still require one """ - if not denied_tool_names: + if not denied_targets: return data - verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_tool_names)) + verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_targets)) - # Create a mapping of tool_use_id to error result - error_tool_names: Final = set() - for tool_use in denied_tool_names: - error_tool_names.add(tool_use) - - tools: Final[list[ChatCompletionToolParam] | None] = data.get("tools") - if tools is not None: - new_tools: Final = [] - for tool in tools: - tool_name, tool_type = self._get_request_tool_name(tool) - if tool_type == "function" and tool_name in error_tool_names: - continue - new_tools.append(tool) - data["tools"] = new_tools - - functions: Final = data.get("functions") - if functions is not None: - data["functions"] = [ - function for function in functions if self._get_legacy_function_name(function) not in error_tool_names - ] - - named_tool_choice: Final = self._get_named_tool_choice(data) - if named_tool_choice in error_tool_names: - data["tool_choice"] = "none" - - named_function_call: Final = self._get_named_function_call(data) - if named_function_call in error_tool_names: - data["function_call"] = "none" - - messages: Final = without_tool_search_items_loading( - _REQUEST_DATA.validate_python(data), frozenset(denied_tool_names) + denied_names: Final = frozenset(name for name, _ in denied_targets if name is not None) + request: Final = _REQUEST_DATA.validate_python(data) + kept_layers: Final = tuple( + self._kept_layer(layer, denied_targets, denied_names) for layer in _request_layers(request) ) + + sent_tools: Final = _layer_sent_tools(_sent_fields(request)) + kept_sent: Final = reduce(_KeptLayer.overlaid_by, kept_layers) + kept_sent_tools: Final = kept_sent.sent_tools() + removed_names: Final = denied_names | ( + self._referable_names_in(sent_tools) - self._referable_names_in(kept_sent_tools) + ) + kept_types: Final = self._tool_types_in(kept_sent_tools) + removed_types: Final = self._tool_types_in(sent_tools) - kept_types + + tools_remain: Final = bool(kept_sent_tools) or bool(kept_sent.mcp_servers) + rewrite_choice: Final = partial( + self._rewritten_tool_choice, + removed_names=removed_names, + removed_types=removed_types, + tools_remain=tools_remain, + ) + rewritten: Final = tuple( + self._rewritten_layer_fields(kept, rewrite_choice, removed_names, "function" in kept_types) + for kept in kept_layers + ) + data.update(rewritten[0]) + if len(rewritten) > 1 and rewritten[1]: + data["extra_body"] = {**kept_layers[1].fields, **rewritten[1]} # rebind-ok: rewrite mode edits in place + + messages: Final = without_tool_search_items_loading(_REQUEST_DATA.validate_python(data), denied_names) if messages is not None: - data["messages"] = messages # rebind-ok: rewrite mode edits the request in place + data["messages"] = [ # rebind-ok: rewrite mode edits the request in place + _without_references_in_message(message, removed_names) for message in messages + ] return data @@ -843,25 +1543,29 @@ class ToolPermissionGuardrail(CustomGuardrail): ) return data - # Check permissions for each tool - denied_tool_names: Final = [] - for tool_name, tool_type in new_tools: - is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) + member_names: Final = tuple( + chain.from_iterable(map(self._bare_member_names, _sent_tools(_REQUEST_DATA.validate_python(data)))) + ) + bare_names: Final = _grouped_values(member_names) + decisions: Final = tuple((target, self._check_target_permission(target, bare_names)) for target in new_tools) + denied: Final = tuple( + (target, message) + for target, (is_allowed, _, message) in decisions + if not is_allowed and message is not None + ) + for _, message in denied: + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) + if denied and self.on_disallowed_action == "block": + raise HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "detection_message": denied[0][1], + }, + ) - if not is_allowed and message is not None: - verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) - if self.on_disallowed_action == "block": - raise HTTPException( - status_code=400, - detail={ - "error": "Violated guardrail policy", - "detection_message": message, - }, - ) - denied_tool_names.append(tool_name) - - if denied_tool_names: - data = self._modify_request_with_permission_errors(data, denied_tool_names) + if denied: + self._modify_request_with_permission_errors(data, frozenset(target for target, _ in denied)) verbose_proxy_logger.debug("Tool Permission Guardrail Pre-Call Hook: All tools allowed") @@ -908,7 +1612,7 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail: Found %s tool calls", len(tool_calls)) - denied_tools: Final = self._evaluate_tool_calls(tool_calls) + denied_tools: Final = self._evaluate_tool_calls(tool_calls, self._returned_call_names(data)) if not denied_tools: verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed") @@ -932,7 +1636,7 @@ class ToolPermissionGuardrail(CustomGuardrail): Args: user_api_key_dict: User API key information (unused but required by interface) response: The model response to check - request_data: The model request (unused but required by interface) + request_data: The model request """ # Import here to avoid circular imports @@ -949,7 +1653,9 @@ class ToolPermissionGuardrail(CustomGuardrail): stream_chunk_builder(chunks=all_chunks) if not is_raw_sse_stream(all_chunks) else None ) if isinstance(assembled_model_response, ModelResponse): - denied_tools = self._check_assembled_stream(assembled_model_response) + denied_tools = self._check_assembled_stream( + assembled_model_response, self._returned_call_names(request_data) + ) if denied_tools: self._modify_response_with_permission_errors(assembled_model_response, denied_tools) @@ -975,18 +1681,23 @@ class ToolPermissionGuardrail(CustomGuardrail): yield chunk return - for sse_chunk in self._anthropic_stream_to_deliver(anthropic_response, all_chunks): + for sse_chunk in self._anthropic_stream_to_deliver( + anthropic_response, all_chunks, self._returned_call_names(request_data) + ): yield sse_chunk def _anthropic_stream_to_deliver( - self, anthropic_response: ModelResponse, source_chunks: Sequence[object] + self, + anthropic_response: ModelResponse, + source_chunks: Sequence[object], + judged_names: _JudgedCalls, ) -> Sequence[object]: """The frames the client gets: on a failed upstream only what the check read, never the upstream's own. Assembly reads past an upstream error, and a client can read a frame of a broken stream that assembly does not, so forwarding those frames would pass a tool call nobody checked. """ - denials: Final = self._check_assembled_stream(anthropic_response) + denials: Final = self._check_assembled_stream(anthropic_response, judged_names) if denials: self._modify_response_with_permission_errors(anthropic_response, denials) elif not anthropic_sse_error_events(source_chunks): @@ -994,7 +1705,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return anthropic_sse_chunks_from_response(anthropic_response, source_chunks, keep_explanations=True) def _check_assembled_stream( - self, assembled: ModelResponse + self, assembled: ModelResponse, judged_names: _JudgedCalls ) -> tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...]: verbose_proxy_logger.debug("Tool Permission Guardrail: Checking response") tool_calls: Final = self._extract_tool_calls_from_response(assembled) @@ -1002,7 +1713,7 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail: No tool uses found") return () verbose_proxy_logger.debug("Tool Permission Guardrail: Found %s tool calls", len(tool_calls)) - denied_tools: Final = self._evaluate_tool_calls(tool_calls) + denied_tools: Final = self._evaluate_tool_calls(tool_calls, judged_names) if not denied_tools: verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed") return denied_tools diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 3f76bcf69ad..c4cd7c2d411 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2162,7 +2162,7 @@ _WS_FRAME: Final = TypeAdapter(dict[str, object]) _GUARDRAIL_INFORMATION: Final = TypeAdapter(list[dict[str, object]]) _PII_TOKENS: Final = TypeAdapter(dict[str, str]) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" -_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") +_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools", "tool_choice", "extra_body") _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") _METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config") diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py index 6f521d1fd1b..99ff7953ee2 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -16,7 +16,7 @@ class ToolPermissionRule(LiteLLMBaseModel): id: str = Field(description="Unique identifier for the rule") tool_name: str | None = Field( default=None, - description="Regex pattern applied to the tool's function name", + description="Regex pattern applied to the tool's name, or to its type when the tool has no name", ) tool_type: str | None = Field( default=None, diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index 854958fc39e..457f03406f1 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -20,8 +20,8 @@ - {id: guardrail.ibm_guardrails.post_call.blocks, module: guardrail, tier: P1, hook_point: post_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/ibm_guardrails/ibm_detector.py", rationale: "Output policy validation"} - {id: guardrail.semantic_guard.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/semantic_guard", rationale: "Semantic policy compliance"} - {id: guardrail.block_code_execution.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/block_code_execution", rationale: "Code-injection prevention"} -- {id: guardrail.tool_permission.pre_call.allows, module: guardrail, tier: P1, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_hooks/tool_permission.py", rationale: "Grant allowed tools"} -- {id: guardrail.tool_permission.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/tool_permission.py", rationale: "Block unauthorized tools"} +- {id: guardrail.tool_permission.pre_call.allows, module: guardrail, tier: P1, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/tool_permission.py", rationale: "Grant allowed tools"} +- {id: guardrail.tool_permission.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/tool_permission.py", rationale: "Block unauthorized tools"} - {id: guardrail.microsoft_purview.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/microsoft_purview/purview_dlp.py", rationale: "DLP sensitive-data disclosure"} - {id: guardrail.headroom.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/headroom/headroom.py", rationale: "Anomaly detection threshold"} - {id: guardrail.generic_guardrail_api.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py", rationale: "Vendor-agnostic custom API"} diff --git a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py index d5a6a9d9866..87545299169 100644 --- a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py @@ -19,10 +19,9 @@ verdict comes from the proxy itself. from __future__ import annotations -from typing import Final +from typing import Final, Literal import pytest - from e2e_config import unique_marker from e2e_http import StreamingResponse, UnknownApiError from e2e_metadata import Capability, Domain, Mode, Provider, Subject, meta @@ -31,14 +30,29 @@ from guardrails_client import ( ToolPermissionParamsBody, ToolPermissionRuleBody, poll_until_blocked, + poll_until_blocked_stream, poll_until_guardrail_applied, ) from lifecycle import ResourceManager -from models import ChatResponse, ChatTool, ChatToolFunction +from models import ( + AnthropicCustomTool, + AnthropicMessagesBody, + AnthropicMessagesResponse, + AnthropicToolChoice, + ChatMessage, + ChatResponse, + ChatTool, + ChatToolFunction, + JsonSchemaProperty, + ResponsesFunctionTool, + ResponsesToolBody, + ResponsesToolOutput, + ToolInputSchema, +) -pytestmark = pytest.mark.e2e +pytestmark: Final = pytest.mark.e2e -MODEL = "gemini-2.5-flash" +MODEL: Final = "gemini-2.5-flash" #: The one tool the guardrail permits, and one it does not. Both are declared by #: the caller in the request body; the guardrail reads them there. @@ -63,11 +77,15 @@ DENIED_TOOL: Final = ChatTool( TOOL_PROMPT: Final = "What is the weather in Paris right now?" +type FlatToolEndpoint = Literal["messages", "responses"] +ALLOWED_TOOL_SCHEMA: Final = ToolInputSchema(properties={"city": JsonSchemaProperty(type="string")}, required=["city"]) +DENIED_TOOL_SCHEMA: Final = ToolInputSchema() + def _register_tool_permission(client: GuardrailsClient, resources: ResourceManager, *, name: str) -> None: """Allow-list exactly one tool: everything else falls to `default_action=deny` and, with `on_disallowed_action=block`, is rejected outright.""" - guardrail_id = client.register( + guardrail_id: Final = client.register( name, ToolPermissionParamsBody( mode="pre_call", @@ -100,6 +118,57 @@ def _tool_call_names(response: ChatResponse) -> tuple[str, ...]: ) +def _send_flat_tool( + client: GuardrailsClient, + key: str, + endpoint: FlatToolEndpoint, + *, + guardrail: str, + tool: ChatTool, + schema: ToolInputSchema, + require_tool: bool = False, +) -> StreamingResponse: + name: Final = tool.function.name + match endpoint: + case "messages": + return client.proxy.transport.send( + "/v1/messages", + headers=client.proxy.transport.bearer(key), + json=AnthropicMessagesBody( + model=MODEL, + messages=[ChatMessage(role="user", content=TOOL_PROMPT)], + max_tokens=128, + tools=[ + AnthropicCustomTool(name=name, description=tool.function.description or "", input_schema=schema) + ], + tool_choice=AnthropicToolChoice(type="any") if require_tool else None, + guardrails=[guardrail], + ), + ) + case "responses": + return client.proxy.transport.send( + "/v1/responses", + headers=client.proxy.transport.bearer(key), + json=ResponsesToolBody( + model=MODEL, + input=TOOL_PROMPT, + tools=[ResponsesFunctionTool(name=name, description=tool.function.description, parameters=schema)], + tool_choice="required" if require_tool else None, + guardrails=[guardrail], + ), + ) + + +def _flat_tool_call_names(endpoint: FlatToolEndpoint, body: str) -> tuple[str, ...]: + match endpoint: + case "messages": + content: Final = AnthropicMessagesResponse.model_validate_json(body).content or [] + return tuple(block.name for block in content if block.type == "tool_use" and block.name) + case "responses": + output: Final = ResponsesToolOutput.model_validate_json(body).output + return tuple(item.name for item in output if item.type == "function_call" and item.name) + + class TestToolPermissionPreCall: @pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["chat_completions"]) @meta( @@ -118,10 +187,10 @@ class TestToolPermissionPreCall: rejected with a 400 that names the denied tool. An unauthorized tool that merely reaches the model is the whole failure mode this guardrail exists to prevent, so a 200 here is a hard failure.""" - name = f"e2e-toolperm-block-{unique_marker()}" + name: Final = f"e2e-toolperm-block-{unique_marker()}" _register_tool_permission(client, resources, name=name) - result = poll_until_blocked( + result: Final = poll_until_blocked( lambda: client.chat( scoped_key, MODEL, @@ -161,10 +230,10 @@ class TestToolPermissionPreCall: and the model calls it. Without the header check a guardrail that never attached would pass this test for the wrong reason, so the 200 alone is not the contract.""" - name = f"e2e-toolperm-allow-{unique_marker()}" + name: Final = f"e2e-toolperm-allow-{unique_marker()}" _register_tool_permission(client, resources, name=name) - outcome = poll_until_guardrail_applied( + outcome: Final = poll_until_guardrail_applied( lambda: client.chat_raw( scoped_key, MODEL, @@ -178,13 +247,117 @@ class TestToolPermissionPreCall: ) assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" - applied = _applied_guardrails(outcome) + applied: Final = _applied_guardrails(outcome) assert name in applied, ( "the allowed call must carry x-litellm-applied-guardrails naming the guardrail; " f"without it the 200 only proves the guardrail never ran. Got {applied!r}" ) - called = _tool_call_names(ChatResponse.model_validate_json(outcome.body)) + called: Final = _tool_call_names(ChatResponse.model_validate_json(outcome.body)) + assert called == (ALLOWED_TOOL.function.name,), ( + f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}" + ) + + @pytest.mark.parametrize( + "endpoint", + [ + pytest.param( + "messages", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["messages"]), + ), + pytest.param( + "responses", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["responses"]), + ), + ], + ) + @meta( + Subject( + domain=Domain.GUARDRAILS, + providers=(Provider.GEMINI,), + models=(MODEL,), + capabilities=(Capability.FUNCTION_CALLING,), + mode=Mode.NONSTREAM, + ) + ) + def test_pre_call_blocks_flat_tool_outside_the_allow_list( + self, + client: GuardrailsClient, + resources: ResourceManager, + scoped_key: str, + endpoint: FlatToolEndpoint, + ) -> None: + name: Final = f"e2e-toolperm-block-{endpoint}-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + outcome: Final = poll_until_blocked_stream( + lambda: _send_flat_tool( + client, scoped_key, endpoint, guardrail=name, tool=DENIED_TOOL, schema=DENIED_TOOL_SCHEMA + ) + ) + + assert outcome.status_code == 400, ( + f"tool_permission let a tool outside the allow-list through /{endpoint}; " + f"got {outcome.status_code}: {outcome.body[:400]}" + ) + assert DENIED_TOOL.function.name in outcome.body, ( + f"the block must name the denied tool so the caller can fix the request; got: {outcome.body[:400]}" + ) + assert "Violated guardrail policy" in outcome.body, ( + f"the block must be a guardrail verdict, got: {outcome.body[:400]}" + ) + + @pytest.mark.parametrize( + "endpoint", + [ + pytest.param( + "messages", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["messages"]), + ), + pytest.param( + "responses", + marks=pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["responses"]), + ), + ], + ) + @meta( + Subject( + domain=Domain.GUARDRAILS, + providers=(Provider.GEMINI,), + models=(MODEL,), + capabilities=(Capability.FUNCTION_CALLING,), + mode=Mode.NONSTREAM, + ) + ) + def test_pre_call_allows_permitted_flat_tool( + self, + client: GuardrailsClient, + resources: ResourceManager, + scoped_key: str, + endpoint: FlatToolEndpoint, + ) -> None: + name: Final = f"e2e-toolperm-allow-{endpoint}-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + outcome: Final = poll_until_guardrail_applied( + lambda: _send_flat_tool( + client, + scoped_key, + endpoint, + guardrail=name, + tool=ALLOWED_TOOL, + schema=ALLOWED_TOOL_SCHEMA, + require_tool=True, + ), + name, + ) + + assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" + applied: Final = _applied_guardrails(outcome) + assert name in applied, ( + f"the allowed call must carry x-litellm-applied-guardrails naming the guardrail; got {applied!r}" + ) + called: Final = _flat_tool_call_names(endpoint, outcome.body) assert called == (ALLOWED_TOOL.function.name,), ( f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}" ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 25683d64313..86af30f413a 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -653,6 +653,35 @@ class ResponsesStreamBody(BaseModel): cache: dict[str, bool] | None = {"no-cache": True} +class ResponsesFunctionTool(BaseModel): + model_config = ConfigDict(frozen=True) + type: Literal["function"] = "function" + name: str + description: str | None = None + parameters: ToolInputSchema + + +class ResponsesToolBody(BaseModel): + model_config = ConfigDict(frozen=True) + model: str + input: str + tools: list[ResponsesFunctionTool] + tool_choice: Literal["required"] | None = None + guardrails: list[str] + max_output_tokens: int = 128 + + +class ResponsesOutputItem(BaseModel): + model_config = ConfigDict(frozen=True) + type: str + name: str | None = None + + +class ResponsesToolOutput(BaseModel): + model_config = ConfigDict(frozen=True) + output: list[ResponsesOutputItem] + + class CountTokensBody(BaseModel): """POST /v1/messages/count_tokens body: the /v1/messages shape minus max_tokens (the endpoint only counts the prompt).""" diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 2eecb83e2df..03bc308a89d 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1785,6 +1785,289 @@ def test_vertex_ai_map_tools(): assert tools == new_tools +def test_vertex_ai_map_tools_declares_a_namespace_by_its_members() -> None: + namespace = { + "type": "namespace", + "name": "github", + "tools": [ + {"type": "function", "name": "create_issue", "parameters": {"type": "object", "properties": {}}}, + {"type": "function", "function": {"name": "close_issue"}}, + ], + } + + tools = VertexGeminiConfig().map_function(value=[namespace], optional_params={}) + + assert [declaration["name"] for declaration in tools[0]["function_declarations"]] == [ + "create_issue", + "close_issue", + ] + + +def test_vertex_ai_map_tools_declares_a_namespace_member_like_the_same_chat_tool() -> None: + parameters = { + "type": "object", + "properties": {"repo": {"$ref": "#/$defs/Repo"}, "label": {"type": ["string", "null"]}}, + "$defs": {"Repo": {"type": "string"}}, + } + function = {"name": "create_issue", "description": "Open an issue", "parameters": parameters} + namespace = {"type": "namespace", "name": "github", "tools": [{"type": "function", **function}]} + + from_namespace = VertexGeminiConfig().map_function(value=[namespace], optional_params={}) + from_chat_tool = VertexGeminiConfig().map_function( + value=[{"type": "function", "function": function}], optional_params={} + ) + + assert from_namespace == from_chat_tool + + +_GET_WEATHER = {"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}} + + +_WEATHER_NAMESPACE = {"type": "namespace", "name": "weather", "tools": [_GET_WEATHER]} +_TOP_LEVEL_GET_WEATHER = {"type": "function", "function": {k: v for k, v in _GET_WEATHER.items() if k != "type"}} + + +@pytest.mark.parametrize( + ("tools", "declared"), + [ + ( + [{"type": "namespace", "name": "forecast", "tools": [_GET_WEATHER]}, _WEATHER_NAMESPACE], + ["forecast__get_weather", "weather__get_weather"], + ), + ([_TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], ["get_weather", "weather__get_weather"]), + ([_WEATHER_NAMESPACE, _TOP_LEVEL_GET_WEATHER], ["weather__get_weather", "get_weather"]), + ], + ids=["namespaces_share_a_member", "top_level_first", "namespace_first"], +) +def test_vertex_ai_map_tools_declares_qualified_member_names_when_another_tool_declares_a_member_name( + tools: list[dict], declared: list[str] +) -> None: + mapped = VertexGeminiConfig().map_function(value=tools, optional_params={}) + + assert [declaration["name"] for declaration in mapped[0]["function_declarations"]] == declared + + +def test_vertex_ai_map_tools_declares_a_member_a_namespace_lists_twice_once_from_its_first_entry() -> None: + first = {**_GET_WEATHER, "description": "Current weather"} + repeat = {**_GET_WEATHER, "description": "Forecast"} + namespace = {"type": "namespace", "name": "weather", "tools": [first, repeat]} + + mapped = VertexGeminiConfig().map_function(value=[namespace], optional_params={}) + + declarations = mapped[0]["function_declarations"] + assert [(declaration["name"], declaration["description"]) for declaration in declarations] == [ + ("get_weather", "Current weather") + ] + + +@pytest.mark.parametrize( + ("tools", "declared"), + [ + ( + [ + _WEATHER_NAMESPACE, + _TOP_LEVEL_GET_WEATHER, + {"type": "function", "function": {"name": "weather__get_weather", "description": "Top level"}}, + ], + [("get_weather", None), ("weather__get_weather", "Top level")], + ), + ( + [ + {"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]}, + {"type": "namespace", "name": "mcp", "tools": [{"type": "function", "name": "ops__exec"}]}, + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "ops__exec"}}, + ], + [("mcp__ops__exec", None), ("exec", None), ("ops__exec", None)], + ), + ( + [_TOP_LEVEL_GET_WEATHER, _TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], + [("get_weather", None), ("get_weather", None), ("weather__get_weather", None)], + ), + ], + ids=["top_level_function_has_the_qualified_name", "namespaces_build_one_qualified_name", "client_repeats_a_function"], +) +def test_vertex_ai_map_tools_never_repeats_a_generated_name( + tools: list[dict], declared: list[tuple[str, str | None]] +) -> None: + mapped = VertexGeminiConfig().map_function(value=tools, optional_params={}) + + declarations = mapped[0]["function_declarations"] + assert [(declaration["name"], declaration.get("description")) for declaration in declarations] == declared + + +def test_vertex_ai_map_openai_params_qualifies_members_by_the_namespace_top_level_name() -> None: + namespace = { + "type": "namespace", + "name": "safe", + "function": {"name": "lookup"}, + "tools": [{"type": "function", "name": "lookup"}], + } + + mapped = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tools": [namespace], + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + declared = [declaration["name"] for tool in mapped["tools"] for declaration in tool["function_declarations"]] + assert declared == ["safe__lookup"] + assert mapped["tool_choice"]["functionCallingConfig"]["allowed_function_names"] == ["safe__lookup"] + + +def test_vertex_ai_map_openai_params_declares_qualified_member_names_when_functions_repeat_a_member() -> None: + mapped = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tools": [{"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]}], + "functions": [{"name": "lookup", "parameters": {"type": "object", "properties": {}}}], + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + declared = [declaration["name"] for tool in mapped["tools"] for declaration in tool["function_declarations"]] + assert declared == ["mcp__weather__lookup", "lookup"] + + +def test_vertex_ai_map_tools_keeps_repeated_chat_tools_when_no_namespace_is_sent() -> None: + function = {k: v for k, v in _GET_WEATHER.items() if k != "type"} + + mapped = VertexGeminiConfig().map_function( + value=[{"type": "function", "function": function}, {"type": "function", "function": function}], + optional_params={}, + ) + + assert [declaration["name"] for declaration in mapped[0]["function_declarations"]] == ["get_weather", "get_weather"] + + +@pytest.mark.parametrize( + ("forced_name", "other_tools", "allowed_function_names"), + [ + ("github", [], ["create_issue", "close_issue"]), + ("close_issue", [], ["close_issue"]), + ("github", [{"type": "function", "function": {"name": "github"}}], ["github"]), + ( + "github", + [{"type": "function", "function": {"name": "close_issue"}}], + ["github__create_issue", "github__close_issue"], + ), + ( + "close_issue", + [{"type": "namespace", "name": "gitlab", "tools": [{"type": "function", "name": "close_issue"}]}], + ["github__close_issue", "gitlab__close_issue"], + ), + ( + "github", + [ + {"type": "function", "function": {"name": "close_issue"}}, + {"type": "function", "function": {"name": "github__close_issue"}}, + ], + ["github__create_issue"], + ), + ( + "close_issue", + [ + {"type": "namespace", "name": "gitlab", "tools": [{"type": "function", "name": "close_issue"}]}, + {"type": "function", "function": {"name": "github__close_issue"}}, + ], + ["gitlab__close_issue"], + ), + ( + "mcp", + [ + {"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]}, + { + "type": "namespace", + "name": "mcp", + "tools": [{"type": "function", "name": "ops__exec"}, {"type": "function", "name": "status"}], + }, + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "ops__exec"}}, + ], + ["mcp__status"], + ), + ], + ids=[ + "names_the_namespace", + "names_a_member", + "names_a_function_sharing_the_namespace_name", + "names_a_namespace_whose_member_a_function_repeats", + "names_a_member_two_namespaces_share", + "names_a_namespace_whose_generated_name_a_function_holds", + "names_a_member_whose_generated_name_a_function_holds", + "names_a_namespace_whose_generated_name_another_namespace_holds", + ], +) +def test_vertex_ai_tool_choice_forces_only_functions_gemini_is_sent( + forced_name: str, other_tools: list[dict[str, object]], allowed_function_names: list[str] +) -> None: + namespace = { + "type": "namespace", + "name": "github", + "tools": [{"type": "function", "name": "create_issue"}, {"type": "function", "name": "close_issue"}], + } + + mapped = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tools": [namespace, *other_tools], + "tool_choice": {"type": "function", "function": {"name": forced_name}}, + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + declared = [declaration["name"] for declaration in mapped["tools"][0]["function_declarations"]] + assert mapped["tool_choice"]["functionCallingConfig"]["allowed_function_names"] == allowed_function_names + assert set(allowed_function_names) <= set(declared) + + +@pytest.mark.parametrize( + ("forced_name", "tools"), + [ + ( + "weather", + [ + {"type": "namespace", "name": "weather", "tools": [{"type": "function", "name": "lookup"}]}, + {"type": "function", "function": {"name": "lookup"}}, + {"type": "function", "function": {"name": "weather__lookup"}}, + ], + ), + ( + "ops__exec", + [ + {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]}, + {"type": "namespace", "name": "ops__exec", "tools": [{"type": "function", "name": "status"}]}, + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "status"}}, + {"type": "function", "function": {"name": "ops__exec__status"}}, + ], + ), + ], + ids=["its_members_lose_to_functions", "another_namespace_generates_its_name"], +) +def test_vertex_ai_tool_choice_naming_a_namespace_gemini_gets_no_member_of( + forced_name: str, tools: list[dict[str, object]] +) -> None: + params = {"tools": tools, "tool_choice": {"type": "function", "function": {"name": forced_name}}} + + with pytest.raises(litellm.UnsupportedParamsError, match=f"tool_choice '{forced_name}'"): + VertexGeminiConfig().map_openai_params( + non_default_params=params, optional_params={}, model="gemini-3.8-flash", drop_params=False + ) + dropped = VertexGeminiConfig().map_openai_params( + non_default_params=params, optional_params={}, model="gemini-3.8-flash", drop_params=True + ) + + assert "tool_choice" not in dropped + assert dropped["tools"][0]["function_declarations"] + + def test_vertex_ai_map_tool_with_anyof(): """ Related issue: https://github.com/BerriAI/litellm/issues/11164 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 91f00edfb39..46e766f753c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -2,11 +2,12 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ +import copy import json import logging import re from itertools import chain -from typing import Final +from typing import Final, Literal from unittest.mock import patch import pytest @@ -26,6 +27,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) from litellm.types.utils import ( + CallTypesLiteral, ChatCompletionCustomToolCallPayload, ChatCompletionDeltaToolCall, ChatCompletionMessageCustomToolCall, @@ -467,6 +469,346 @@ class TestToolPermissionGuardrail: ] assert kept == [[], ["Read"]] + @pytest.mark.asyncio + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + async def test_post_call_judges_a_returned_namespace_member_by_its_qualified_name(self, stream: bool) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + data = {"guardrails": ["post-mcp"], "tools": [namespace]} + function = Function(name="lookup", arguments='{"city": "Paris"}') + + if stream: + + async def _stream(): + yield ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_lookup", index=0, type="function", function=function + ) + ] + ), + ) + ] + ) + + chunks = [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), response=_stream(), request_data=data + ) + ] + deltas = [choice.delta for choice in chain.from_iterable(chunk.choices for chunk in chunks)] + called = [call.function.name for call in chain.from_iterable(delta.tool_calls or () for delta in deltas)] + content = "".join(delta.content or "" for delta in deltas) + else: + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_lookup", "type": "function", "function": function}]}, + ) + ] + ) + await guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + message = response.choices[0].message + called = [call.function.name for call in message.tool_calls or ()] + content = message.content or "" + + assert called == [] + assert "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" in content + + @pytest.mark.asyncio + async def test_post_call_judges_a_returned_namespace_member_by_the_rule_its_arguments_satisfy(self) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-exec", + rules=[ + { + "id": "allow_ls", + "tool_name": r"^exec$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^ls$"}, + }, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_ops", "tool_name": r"^ops__.*$", "decision": "allow"}, + ], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + data = {"guardrails": ["post-exec"], "tools": [namespace]} + function = Function(name="exec", arguments='{"command": "rm -rf /data"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_exec", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), response=response) + + message = response.choices[0].message + assert [call.function.name for call in message.tool_calls or ()] == [] + assert "Tool 'exec'" in (message.content or "") + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rules", "command", "kept"), + [ + ( + [ + { + "id": "allow_ls", + "tool_name": r"^exec$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^ls$"}, + }, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "ls", + ["run"], + ), + ( + [ + { + "id": "allow_ls", + "tool_name": r"^exec$", + "decision": "allow", + "allowed_param_patterns": {"command": r"^ls$"}, + }, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "pwd", + [], + ), + ( + [ + {"id": "allow_run", "tool_name": r"^run$", "decision": "allow"}, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "ls", + [], + ), + ], + ids=["namespace_rule_allows", "namespace_rule_denies", "namespace_deny_beats_member_allow"], + ) + async def test_post_call_judges_a_returned_namespace_member_by_a_rule_on_its_namespace_name( + self, rules: list[dict[str, object]], command: str, kept: list[str] + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-exec", + rules=rules, + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "exec", "tools": [{"type": "function", "name": "run"}]} + function = Function(name="run", arguments=json.dumps({"command": command})) + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_run", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-exec"], "tools": [namespace]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + message = response.choices[0].message + assert [call.function.name for call in message.tool_calls or ()] == kept + assert ("Tool 'exec'" in (message.content or "")) == (not kept) + + @pytest.mark.asyncio + @pytest.mark.parametrize("namespace_first", [True, False], ids=["namespace_first", "top_level_first"]) + async def test_post_call_judges_a_returned_name_a_top_level_function_carries_as_that_function( + self, namespace_first: bool + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + top_level = {"type": "function", "function": {"name": "lookup"}} + tools = [namespace, top_level] if namespace_first else [top_level, namespace] + function = Function(name="lookup", arguments='{"country": "France"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_lookup", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-mcp"], "tools": tools}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == ["lookup"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "same_named_tool", + [ + {"type": "mcp", "server_label": "lookup", "server_url": "https://mcp.example.com/sse"}, + {"type": "mcp_toolset", "mcp_server_name": "lookup"}, + {"ignored": {"name": "lookup"}}, + {"type": "custom", "custom": {"name": "lookup"}}, + {"functionDeclarations": [{"name": "lookup"}]}, + ], + ids=["mcp", "mcp_toolset", "unknown_keyed_block", "chat_custom", "native_declaration"], + ) + async def test_post_call_judges_a_returned_namespace_member_as_that_member_beside_a_same_named_non_function( + self, same_named_tool: dict[str, object] + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + function = Function(name="lookup", arguments='{"country": "France"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_lookup", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-mcp"], "tools": [namespace, same_named_tool]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == [] + + @pytest.mark.asyncio + async def test_post_call_judges_a_call_tagged_with_a_namespace_as_that_member_beside_a_same_named_function( + self, + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-mcp", + rules=[{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]} + top_level = {"type": "function", "function": {"name": "lookup"}} + tagged_call = { + "id": "call_lookup", + "type": "function", + "function": Function(name="lookup", arguments='{"country": "France"}'), + "provider_specific_fields": {"namespace": "mcp__weather"}, + } + response = ModelResponse( + choices=[Choices(finish_reason="tool_calls", message={"tool_calls": [tagged_call]})] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-mcp"], "tools": [namespace, top_level]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + message = response.choices[0].message + assert [call.function.name for call in message.tool_calls or ()] == [] + assert "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" in (message.content or "") + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("decision", "default_action", "kept"), + [("allow", "deny", ["weather__get_weather"]), ("deny", "allow", [])], + ids=["bare_name_allowed", "bare_name_denied"], + ) + async def test_post_call_judges_a_returned_qualified_member_name_by_the_member_names_pre_call_judged( + self, decision: Literal["allow", "deny"], default_action: Literal["allow", "deny"], kept: list[str] + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-weather", + rules=[{"id": "get_weather", "tool_name": r"^get_weather$", "decision": decision}], + default_action=default_action, + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "weather", "tools": [{"type": "function", "name": "get_weather"}]} + top_level = {"type": "function", "function": {"name": "get_weather"}} + function = Function(name="weather__get_weather", arguments='{"city": "Istanbul"}') + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_weather", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-weather"], "tools": [namespace, top_level]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == kept + + @pytest.mark.asyncio + async def test_post_call_judges_a_top_level_function_by_its_own_name_beside_a_member_qualified_the_same( + self, + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-exec", + rules=[{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + top_level = {"type": "function", "function": {"name": "ops__exec"}} + function = Function(name="ops__exec", arguments="{}") + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_ops", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-exec"], "tools": [namespace, top_level]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == ["ops__exec"] + @pytest.mark.asyncio async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule( self, @@ -1361,6 +1703,1723 @@ class TestToolPermissionGuardrailAnthropicMessages: def _tool_use(self, name, tool_id="tu_1"): return {"type": "tool_use", "id": tool_id, "name": name, "input": {"command": "ls"}} + def _always_on_pre_call(self, on_disallowed_action: Literal["block", "rewrite"]) -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name=f"anthropic-pre-call-{on_disallowed_action}", + rules=self.rules, + default_action="deny", + on_disallowed_action=on_disallowed_action, + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "denied_tool", + [ + {"name": "Read", "input_schema": {"type": "object", "properties": {}}}, + {"type": "custom", "name": "Read", "input_schema": {"type": "object", "properties": {}}}, + {"type": "function", "name": "Read", "parameters": {"type": "object", "properties": {}}}, + ], + ids=["anthropic", "anthropic_custom_type", "responses_api_flat_function"], + ) + async def test_pre_call_blocks_denied_request_tool_in_flat_format(self, denied_tool: dict[str, object]) -> None: + data = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "tools": [denied_tool]} + + with pytest.raises(HTTPException) as excinfo: + await self._always_on_pre_call("block").async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="anthropic_messages", + ) + + assert excinfo.value.status_code == 400 + assert excinfo.value.detail["detection_message"] == "Tool 'Read' denied by rule 'deny_read'" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool_fields", "rule_tool_type", "call_type"), + [ + ({"format": {"type": "text"}}, r"^custom$", "responses"), + ({"input_schema": {"type": "object", "properties": {}}}, r"^function$", "anthropic_messages"), + ({"input_schema": {"type": "object", "properties": {}}}, r"^function$", "pass_through_endpoint"), + ({"custom": {"name": "apply_patch"}, "input_schema": {"type": "object"}}, r"^custom$", "acompletion"), + ], + ids=[ + "responses_custom_stays_custom", + "anthropic_custom_reads_as_function", + "anthropic_custom_reads_as_function_on_a_pass_through_route", + "chat_custom_with_a_stray_input_schema_stays_custom", + ], + ) + async def test_pre_call_tool_type_rules_follow_the_request_format( + self, tool_fields: dict[str, object], rule_tool_type: str, call_type: CallTypesLiteral + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name=f"tool-type-{call_type}", + rules=[{"id": "deny_type", "tool_type": rule_tool_type, "decision": "deny"}], + default_action="allow", + on_disallowed_action="block", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "custom", "name": "apply_patch", **tool_fields}], + } + + with pytest.raises(HTTPException) as excinfo: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type=call_type, + ) + + assert excinfo.value.status_code == 400 + assert excinfo.value.detail["detection_message"] == "Tool 'apply_patch' denied by rule 'deny_type'" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool_shape", "tool_choice", "call_type", "expected_tool_choice"), + [ + ( + {"input_schema": {"type": "object", "properties": {}}}, + {"type": "tool", "name": "Read"}, + "anthropic_messages", + {"type": "none"}, + ), + ( + {"type": "function", "parameters": {"type": "object", "properties": {}}}, + {"type": "function", "name": "Read"}, + "responses", + "none", + ), + ], + ids=["anthropic", "responses_api"], + ) + async def test_pre_call_rewrite_strips_denied_flat_tool_and_forced_choice( + self, + tool_shape: dict[str, object], + tool_choice: dict[str, str], + call_type: CallTypesLiteral, + expected_tool_choice: dict[str, str] | str, + ) -> None: + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "Bash", **tool_shape}, {"name": "Read", **tool_shape}], + "tool_choice": tool_choice, + } + + result = await self._always_on_pre_call("rewrite").async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type=call_type, + ) + + assert [tool["name"] for tool in result["tools"]] == ["Bash"] + assert result["tool_choice"] == expected_tool_choice + + def _pre_call( + self, + rules: list[dict[str, str]], + on_disallowed_action: Literal["block", "rewrite"], + default_action: Literal["allow", "deny"] = "allow", + ) -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name=f"request-shapes-{on_disallowed_action}", + rules=rules, + default_action=default_action, + on_disallowed_action=on_disallowed_action, + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + async def _run_pre_call( + self, guardrail: ToolPermissionGuardrail, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: + return await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type=call_type, + ) + + def _namespace(self, *member_names: str) -> dict[str, object]: + return { + "type": "namespace", + "name": "mcp__everything", + "tools": [{"type": "function", "name": name, "parameters": {"type": "object"}} for name in member_names], + } + + @pytest.mark.asyncio + async def test_pre_call_blocks_a_denied_tool_inside_a_responses_namespace(self) -> None: + guardrail = self._pre_call( + [{"id": "deny_sum", "tool_name": r"^mcp__everything__get_sum$", "decision": "deny"}], "block" + ) + data = {"model": "gpt-5", "input": "add 2 and 3", "tools": [self._namespace("echo", "get_sum")]} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + + assert excinfo.value.detail["detection_message"] == "Tool 'mcp__everything__get_sum' denied by rule 'deny_sum'" + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_denied_namespace_members_and_empty_namespaces(self) -> None: + guardrail = self._pre_call([{"id": "deny_sum", "tool_name": r"^.*get_sum$", "decision": "deny"}], "rewrite") + partly_denied = {**self._namespace("echo", "get_sum"), "name": "mcp__kept"} + data = {"model": "gpt-5", "input": "add", "tools": [partly_denied, self._namespace("get_sum")]} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{**partly_denied, "tools": partly_denied["tools"][:1]}] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "malformed"), + [ + ( + "agenerate_content", + {"toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": [{"a": 1}]}}}, + ), + ("responses", {"input": [{"type": {"x": 1}, "content": "hi"}]}), + ("acompletion", {"extra_body": {"input": [{"type": [1]}]}}), + ("acompletion", {"extra_body": {"function_call": {"name": 5}}}), + ], + ids=["unhashable_allowed_function_name", "unhashable_input_type", "extra_body_input_type", "int_function_call"], + ) + async def test_pre_call_rewrite_leaves_a_malformed_value_for_the_provider_to_reject( + self, call_type: CallTypesLiteral, malformed: dict[str, object] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], "rewrite", "deny" + ) + tools = [{"type": "function", "function": {"name": name}} for name in ("get_weather", "shell")] + + result = await self._run_pre_call(guardrail, {"model": "m", "tools": tools, **malformed}, call_type) + + assert result["tools"] == tools[:1] + assert {key: result[key] for key in malformed} == malformed + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rule", "tool", "call_type", "detection_message"), + [ + ( + {"id": "deny_mcp", "tool_type": r"^mcp$", "decision": "deny"}, + {"type": "mcp", "server_label": "weather", "server_url": "https://mcp.example.com"}, + "responses", + "Tool 'weather' denied by rule 'deny_mcp'", + ), + ( + {"id": "allow_bash", "tool_name": r"^Bash$", "decision": "allow"}, + {"type": "mcp_toolset", "mcp_server_name": "weather"}, + "anthropic_messages", + "Tool 'weather' denied by default action", + ), + ( + {"id": "deny_weather", "tool_name": r"^weather$", "decision": "deny"}, + {"type": "mcp", "server_label": "weather", "server_url": "https://mcp.example.com"}, + "responses", + "Tool 'weather' denied by rule 'deny_weather'", + ), + ( + {"id": "deny_lookup", "tool_name": r"^mcp__weather__lookup$", "decision": "deny"}, + {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "custom", "name": "lookup"}]}, + "responses", + "Tool 'mcp__weather__lookup' denied by rule 'deny_lookup'", + ), + ( + {"id": "deny_lookup", "tool_name": r"^mcp__weather__lookup$", "decision": "deny"}, + {"functionDeclarations": [{"name": "mcp__weather__lookup"}, {"name": "get_weather"}]}, + "agenerate_content", + "Tool 'mcp__weather__lookup' denied by rule 'deny_lookup'", + ), + ( + {"id": "allow_bash", "tool_name": r"^Bash$", "decision": "allow"}, + {"googleSearch": {}}, + "acompletion", + "Tool 'googleSearch' denied by default action", + ), + ( + {"id": "deny_custom", "tool_type": r"^custom$", "decision": "deny"}, + {"type": "custom", "custom": {"name": "grammar_tool"}}, + "acompletion", + "Tool 'grammar_tool' denied by rule 'deny_custom'", + ), + ( + {"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}, + {"type": "custom", "custom": {"name": "Read"}}, + "acompletion", + "Tool 'Read' denied by rule 'deny_read'", + ), + ], + ids=[ + "responses_mcp_by_type", + "anthropic_mcp_toolset_by_server_name", + "responses_mcp_by_server_label", + "responses_namespace_custom_member", + "gemini_function_declarations", + "gemini_type_only_tool_by_its_key", + "chat_custom_keeps_its_type", + "chat_custom_nested_name", + ], + ) + async def test_pre_call_checks_tools_that_carry_no_top_level_name( + self, rule: dict[str, str], tool: dict[str, object], call_type: CallTypesLiteral, detection_message: str + ) -> None: + guardrail = self._pre_call([rule], "block", default_action="deny") + data = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "input": "hi", "tools": [tool]} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, call_type) + + assert excinfo.value.detail["detection_message"] == detection_message + + @pytest.mark.asyncio + async def test_pre_call_lets_a_type_rule_allow_a_gemini_tool_by_its_key(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_search", "tool_type": r"^googleSearch$", "decision": "allow"}], "block", default_action="deny" + ) + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [{"googleSearch": {}}]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [{"googleSearch": {}}] + + @pytest.mark.asyncio + async def test_pre_call_blocks_a_denied_tool_loaded_by_responses_tool_search(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "block") + loaded = {"type": "tool_search_output", "tools": [{"type": "function", "name": "mcp__weather__lookup"}]} + data = {"model": "gpt-5", "input": [loaded], "tools": []} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + + assert excinfo.value.detail["detection_message"] == "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" + + @pytest.mark.asyncio + async def test_pre_call_rewrite_keeps_the_allowed_members_of_grouped_and_loaded_tools(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + get_weather = {"type": "function", "name": "get_weather"} + lookup = {"type": "function", "name": "mcp__weather__lookup"} + data = { + "model": "m", + "input": [{"role": "user", "content": "hi"}, {"type": "tool_search_output", "tools": [get_weather, lookup]}], + "tools": [{"functionDeclarations": [{"name": "get_weather"}, {"name": "mcp__weather__lookup"}]}], + } + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{"functionDeclarations": [{"name": "get_weather"}]}] + assert result["input"] == [ + {"role": "user", "content": "hi"}, + {"type": "tool_search_output", "tools": [get_weather]}, + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "message", "rewritten"), + [ + ( + "anthropic_messages", + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [{"type": "tool_reference", "tool_name": "mcp__weather__lookup"}], + }, + {"type": "text", "text": "Tool loaded."}, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + { + "type": "text", + "text": "Permission denied: tool mcp__weather__lookup is not available", + } + ], + }, + {"type": "text", "text": "Tool loaded."}, + ], + }, + ), + ( + "anthropic_messages", + { + "role": "assistant", + "content": [ + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [ + {"type": "tool_reference", "tool_name": "mcp__weather__lookup"}, + {"type": "tool_reference", "tool_name": "get_weather"}, + ], + }, + } + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}], + }, + } + ], + }, + ), + ( + "acompletion", + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "tool_reference", "tool_name": "mcp__weather__lookup"}, + {"type": "tool_reference", "tool_name": "get_weather"}, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "tool_reference", "tool_name": "get_weather"}], + }, + ), + ], + ) + async def test_pre_call_rewrite_drops_references_to_a_removed_tool(self, call_type, message, rewritten) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + data = { + "model": "m", + "messages": [{"role": "user", "content": "weather in Istanbul"}, message], + "tools": [ + {"name": "get_weather", "input_schema": {"type": "object"}, "defer_loading": True}, + {"name": "mcp__weather__lookup", "input_schema": {"type": "object"}, "defer_loading": True}, + ], + } + + result = await self._run_pre_call(guardrail, data, call_type) + + assert [tool["name"] for tool in result["tools"]] == ["get_weather"] + assert result["messages"] == [{"role": "user", "content": "weather in Istanbul"}, rewritten] + + @pytest.mark.asyncio + @pytest.mark.parametrize("tool_choice", ["required", {"type": "function", "name": "get_weather"}]) + async def test_pre_call_rewrite_keeps_tool_choice_when_only_loaded_tools_remain(self, tool_choice) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + get_weather = {"type": "function", "name": "get_weather"} + lookup = {"type": "function", "name": "mcp__weather__lookup"} + data = { + "model": "m", + "input": [{"role": "user", "content": "hi"}, {"type": "tool_search_output", "tools": [get_weather, lookup]}], + "tools": [lookup], + "tool_choice": tool_choice, + } + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [] + assert result["input"][1] == {"type": "tool_search_output", "tools": [get_weather]} + assert result["tool_choice"] == tool_choice + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rule", "tools", "tool_choice"), + [ + ( + {"id": "deny_search", "tool_type": r"^web_search$", "decision": "deny"}, + [{"type": "web_search"}, {"type": "function", "name": "get_weather"}], + {"type": "web_search"}, + ), + ( + {"id": "deny_weather", "tool_name": r"^weather$", "decision": "deny"}, + [ + {"type": "mcp", "server_label": "weather", "server_url": "https://a.example.com"}, + {"type": "mcp", "server_label": "files", "server_url": "https://b.example.com"}, + ], + {"type": "mcp", "server_label": "weather"}, + ), + ], + ids=["hosted_tool_type", "hosted_mcp_server"], + ) + async def test_pre_call_rewrite_relaxes_a_hosted_tool_choice_it_removed( + self, rule: dict[str, str], tools: list[dict[str, str]], tool_choice: dict[str, str] + ) -> None: + guardrail = self._pre_call([rule], "rewrite") + data = {"model": "gpt-5", "input": "hi", "tools": tools, "tool_choice": tool_choice} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == tools[1:] + assert result["tool_choice"] == "none" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tools", "tool_choice", "call_type", "expected_tool_choice"), + [ + ( + [{"type": "custom", "name": "Bash"}, {"type": "custom", "name": "Read"}], + {"type": "custom", "name": "Read"}, + "responses", + "none", + ), + ( + [{"type": "function", "name": "Bash"}, {"type": "function", "name": "Read"}], + { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "function", "name": "Read"}, {"type": "function", "name": "Bash"}], + }, + "responses", + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "Bash"}]}, + ), + ( + [{"type": "function", "name": "Bash"}, {"type": "function", "name": "Read"}], + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "Read"}]}, + "responses", + "none", + ), + ( + [ + {"type": "function", "function": {"name": "Bash"}}, + {"type": "function", "function": {"name": "Read"}}, + ], + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": "Read"}}]}, + }, + "acompletion", + "none", + ), + ( + [{"name": "Read", "input_schema": {"type": "object"}}], + {"type": "any"}, + "anthropic_messages", + {"type": "none"}, + ), + ( + [ + {"name": "Bash", "input_schema": {"type": "object"}}, + {"name": "Read", "input_schema": {"type": "object"}}, + ], + {"type": "any"}, + "anthropic_messages", + {"type": "any"}, + ), + ( + [{"type": "custom", "custom": {"name": "Bash"}}, {"type": "custom", "custom": {"name": "Read"}}], + {"type": "custom", "custom": {"name": "Read"}}, + "acompletion", + "none", + ), + ], + ids=[ + "responses_custom_choice", + "responses_allowed_tools_keeps_the_rest", + "responses_allowed_tools_all_denied", + "chat_allowed_tools_all_denied", + "anthropic_any_with_no_tools_left", + "anthropic_any_with_a_tool_left", + "chat_custom_choice", + ], + ) + async def test_pre_call_rewrite_leaves_no_tool_choice_that_needs_a_removed_tool( + self, + tools: list[dict[str, object]], + tool_choice: dict[str, object], + call_type: CallTypesLiteral, + expected_tool_choice: dict[str, object] | str, + ) -> None: + guardrail = self._pre_call(self.rules, "rewrite") + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "tools": tools, + "tool_choice": tool_choice, + } + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tool_choice"] == expected_tool_choice + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool", "tool_choice", "call_type", "allowed_tool_type"), + [ + ({"type": "custom", "name": "apply_patch"}, {"type": "custom", "name": "apply_patch"}, "responses", r"^custom$"), + ( + {"type": "web_search_20250305", "name": "web_search"}, + {"type": "tool", "name": "web_search"}, + "anthropic_messages", + r"^web_search_\d+$", + ), + ], + ids=["responses_custom_tool", "anthropic_server_tool"], + ) + async def test_pre_call_checks_a_forced_choice_with_its_tool_type( + self, + tool: dict[str, str], + tool_choice: dict[str, str], + call_type: CallTypesLiteral, + allowed_tool_type: str, + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_type", "tool_type": allowed_tool_type, "decision": "allow"}], "block", default_action="deny" + ) + data = {"model": "m", "input": "hi", "tools": [tool], "tool_choice": tool_choice} + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tools"] == [tool] + assert result["tool_choice"] == tool_choice + + @pytest.mark.asyncio + async def test_pre_call_rewrite_relaxes_a_choice_for_a_tool_a_type_rule_removed(self) -> None: + guardrail = self._pre_call([{"id": "deny_custom", "tool_type": r"^custom$", "decision": "deny"}], "rewrite") + data = { + "model": "gpt-5", + "input": "patch it", + "tools": [{"type": "function", "name": "Bash"}, {"type": "custom", "name": "apply_patch"}], + "tool_choice": {"type": "custom", "name": "apply_patch"}, + } + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{"type": "function", "name": "Bash"}] + assert result["tool_choice"] == "none" + + async def _blocked_and_rewritten( + self, rules: list[dict[str, str]], request: dict[str, object], call_type: CallTypesLiteral + ) -> tuple[str, dict[str, object]]: + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(self._pre_call(rules, "block"), copy.deepcopy(request), call_type) + rewritten = await self._run_pre_call(self._pre_call(rules, "rewrite"), request, call_type) + return excinfo.value.detail["detection_message"], rewritten + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("denied_tool", "detection_message"), + [ + ( + {"type": "function", "name": "mcp__weather__lookup"}, + "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'", + ), + ( + {"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]}, + "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'", + ), + ({"type": "web_search"}, "Tool 'web_search' denied by rule 'deny_search'"), + ], + ids=["flat", "namespace_member", "hosted_type"], + ) + async def test_pre_call_checks_and_strips_tools_in_responses_additional_tools_items( + self, denied_tool: dict[str, object], detection_message: str + ) -> None: + rules = [ + {"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}, + {"id": "deny_search", "tool_type": r"^web_search$", "decision": "deny"}, + ] + get_weather = {"type": "function", "name": "get_weather"} + request = { + "model": "gpt-5", + "input": [ + {"role": "user", "content": "hi"}, + {"type": "additional_tools", "tools": [get_weather, denied_tool]}, + {"type": "additional_tools", "tools": [denied_tool]}, + ], + "tools": [], + "tool_choice": "required", + } + + message, rewritten = await self._blocked_and_rewritten(rules, request, "responses") + + assert message == detection_message + assert rewritten["input"] == [ + {"role": "user", "content": "hi"}, + {"type": "additional_tools", "tools": [get_weather]}, + ] + assert rewritten["tool_choice"] == "required" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rules", "default_action", "kept"), + [ + ([{"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}], "deny", ["run"]), + ([{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], "allow", []), + ( + [ + {"id": "allow_run", "tool_name": r"^run$", "decision": "allow"}, + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + ], + "allow", + [], + ), + ( + [ + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_all", "tool_name": r"^.*$", "decision": "allow"}, + ], + "allow", + [], + ), + ( + [ + {"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}, + {"id": "deny_all", "tool_name": r"^.*$", "decision": "deny"}, + ], + "deny", + ["run"], + ), + ], + ids=[ + "namespace_allowed", + "namespace_denied", + "namespace_deny_beats_member_allow", + "namespace_deny_beats_a_catch_all_allow", + "namespace_allow_beats_a_catch_all_deny", + ], + ) + async def test_pre_call_judges_a_namespace_member_by_a_rule_on_its_namespace_name( + self, rules: list[dict[str, str]], default_action: Literal["allow", "deny"], kept: list[str] + ) -> None: + namespace = {"type": "namespace", "name": "exec", "tools": [{"type": "function", "name": "run"}]} + request = {"model": "gpt-5", "input": "Run pwd", "tools": [namespace]} + + rewritten = await self._run_pre_call(self._pre_call(rules, "rewrite", default_action), request, "responses") + + tools = rewritten.get("tools") or [] + assert [member["name"] for member in chain.from_iterable(tool["tools"] for tool in tools)] == kept + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "bash", "read"), + [ + ( + "acompletion", + {"type": "function", "function": {"name": "Bash"}}, + {"type": "function", "function": {"name": "Read"}}, + ), + ("responses", {"type": "function", "name": "Bash"}, {"type": "function", "name": "Read"}), + ( + "anthropic_messages", + {"name": "Bash", "input_schema": {"type": "object"}}, + {"name": "Read", "input_schema": {"type": "object"}}, + ), + ], + ids=["chat", "responses", "anthropic"], + ) + async def test_pre_call_checks_and_strips_tools_in_extra_body( + self, call_type: CallTypesLiteral, bash: dict[str, object], read: dict[str, object] + ) -> None: + request = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "input": "hi", + "extra_body": {"tools": [bash, read], "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(self.rules, request, call_type) + + assert message == "Tool 'Read' denied by rule 'deny_read'" + assert rewritten["extra_body"] == {"tools": [bash], "seed": 7} + + @pytest.mark.asyncio + async def test_pre_call_checks_and_strips_tools_loaded_in_extra_body_input(self) -> None: + rules = [{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}] + get_weather = {"type": "function", "name": "get_weather"} + lookup = {"type": "function", "name": "mcp__weather__lookup"} + user = {"role": "user", "content": "hi"} + request = { + "model": "gpt-5", + "input": "hi", + "tools": [], + "extra_body": {"input": [user, {"type": "tool_search_output", "tools": [get_weather, lookup]}], "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(rules, request, "responses") + + assert message == "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'" + assert rewritten["extra_body"] == { + "input": [user, {"type": "tool_search_output", "tools": [get_weather]}], + "seed": 7, + } + + @pytest.mark.asyncio + async def test_rewrite_relaxes_an_extra_body_tool_choice_that_names_a_dropped_tool(self) -> None: + guardrail = self._pre_call([{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "function": {"name": "get_weather"}} + delete_db = {"type": "function", "function": {"name": "delete_db"}} + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "extra_body": { + "tools": [weather, delete_db], + "tool_choice": {"type": "function", "function": {"name": "delete_db"}}, + "seed": 7, + }, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["extra_body"] == {"tools": [weather], "tool_choice": "none", "seed": 7} + + @pytest.mark.asyncio + async def test_rewrite_relaxes_a_required_tool_choice_when_the_extra_body_tools_sent_lose_every_tool( + self, + ) -> None: + guardrail = self._pre_call([{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "function": {"name": "get_weather"}} + delete_db = {"type": "function", "function": {"name": "delete_db"}} + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "tools": [weather], + "extra_body": {"tools": [delete_db], "tool_choice": "required", "seed": 7}, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + assert result["extra_body"] == {"tools": [], "tool_choice": "none", "seed": 7} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("call_type", "tool_choice"), + [ + ("acompletion", {"type": "function", "function": {"name": "Read"}}), + ("responses", {"type": "function", "name": "Read"}), + ], + ids=["chat", "responses"], + ) + async def test_pre_call_checks_an_extra_body_tool_choice_that_names_an_unlisted_tool( + self, call_type: CallTypesLiteral, tool_choice: dict[str, object] + ) -> None: + request = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "input": "hi", + "tools": [{"type": "function", "name": "Bash"}], + "extra_body": {"tool_choice": tool_choice, "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(self.rules, request, call_type) + + assert message == "Tool 'Read' denied by rule 'deny_read'" + assert rewritten["extra_body"] == {"tool_choice": "none", "seed": 7} + + @pytest.mark.asyncio + async def test_pre_call_checks_and_strips_legacy_functions_in_extra_body(self) -> None: + rules = [{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}] + weather = {"name": "get_weather", "parameters": {"type": "object"}} + delete_db = {"name": "delete_db", "parameters": {"type": "object"}} + request = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "extra_body": {"functions": [weather, delete_db], "function_call": {"name": "delete_db"}, "seed": 7}, + } + + message, rewritten = await self._blocked_and_rewritten(rules, request, "acompletion") + + assert message == "Tool 'delete_db' denied by rule 'deny_delete'" + assert rewritten["extra_body"] == {"functions": [weather], "function_call": "none", "seed": 7} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "fields", + [ + {"functions": [{"name": "get_weather"}, {"name": "delete_db"}], "function_call": "delete_db"}, + { + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + "mcp_servers": [{"type": "url", "url": "https://mcp.example.com/mcp", "name": "delete_db"}], + "tool_choice": "required", + }, + { + "tools": [{"functionDeclarations": [{"name": "get_weather"}, {"name": "delete_db"}]}], + "toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["delete_db"]}}, + }, + ], + ids=["legacy_functions", "mcp_servers", "gemini_tool_config"], + ) + async def test_extra_body_fields_are_judged_and_rewritten_like_the_same_top_level_fields( + self, fields: dict[str, object] + ) -> None: + rules = [{"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}] + base = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + top_request = {**base, **copy.deepcopy(fields)} + extra_request = {**base, "extra_body": {**copy.deepcopy(fields), "seed": 7}} + + top_message, top = await self._blocked_and_rewritten(rules, top_request, "acompletion") + extra_message, extra = await self._blocked_and_rewritten(rules, extra_request, "acompletion") + + top_fields = {key: top[key] for key in fields} + assert top_fields != fields + assert extra_message == top_message + assert extra["extra_body"] == {**top_fields, "seed": 7} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "rule", + [ + {"id": "deny_deepwiki", "tool_name": r"^deepwiki$", "decision": "deny"}, + {"id": "deny_deepwiki", "tool_name": r"^deepwiki$", "tool_type": r"^mcp_toolset$", "decision": "deny"}, + ], + ids=["by_server_name", "by_server_name_and_mcp_toolset_type"], + ) + async def test_pre_call_checks_and_strips_anthropic_mcp_servers_like_an_mcp_toolset( + self, rule: dict[str, str] + ) -> None: + deepwiki = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + files = {"type": "url", "url": "https://files.example.com/mcp", "name": "files"} + request = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "mcp_servers": [deepwiki, files], + "tool_choice": {"type": "any"}, + } + + message, rewritten = await self._blocked_and_rewritten([rule], request, "anthropic_messages") + + assert message == "Tool 'deepwiki' denied by rule 'deny_deepwiki'" + assert rewritten["mcp_servers"] == [files] + assert rewritten["tool_choice"] == {"type": "any"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rule", "tool", "detection_message", "kept_tools"), + [ + ( + {"id": "deny_code", "tool_type": r"^codeExecution$", "decision": "deny"}, + {"googleSearch": {}, "codeExecution": {}}, + "Tool 'codeExecution' denied by rule 'deny_code'", + [{"googleSearch": {}}], + ), + ( + {"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}, + {"functionDeclarations": [{"name": "mcp__weather__lookup"}], "googleSearch": {}}, + "Tool 'mcp__weather__lookup' denied by rule 'deny_mcp'", + [{"googleSearch": {}}], + ), + ( + {"id": "deny_search", "tool_type": r"^googleSearch$", "decision": "deny"}, + {"functionDeclarations": [{"name": "get_weather"}], "googleSearch": {}}, + "Tool 'googleSearch' denied by rule 'deny_search'", + [{"functionDeclarations": [{"name": "get_weather"}]}], + ), + ( + {"id": "deny_builtin", "tool_type": r"^(googleSearch|codeExecution)$", "decision": "deny"}, + {"googleSearch": {}, "codeExecution": {}}, + "Tool 'googleSearch' denied by rule 'deny_builtin'", + [], + ), + ], + ids=["denied_key_beside_a_key", "every_declaration_denied", "denied_key_beside_declarations", "all_denied"], + ) + async def test_pre_call_checks_and_strips_each_part_of_a_multi_key_gemini_tool( + self, + rule: dict[str, str], + tool: dict[str, object], + detection_message: str, + kept_tools: list[dict[str, object]], + ) -> None: + request = {"model": "gemini-3", "contents": [{"role": "user", "parts": [{"text": "hi"}]}], "tools": [tool]} + + message, rewritten = await self._blocked_and_rewritten([rule], request, "agenerate_content") + + assert message == detection_message + assert rewritten["tools"] == kept_tools + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("tool", "call_type"), + [({"type": "web_search"}, "responses"), ({"googleSearch": {}}, "acompletion")], + ids=["responses_hosted_type", "gemini_tool_key"], + ) + async def test_pre_call_matches_a_name_rule_against_a_nameless_tool_by_its_type( + self, + tool: dict[str, object], + call_type: CallTypesLiteral, + on_disallowed_action: Literal["block", "rewrite"], + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_all", "tool_name": r".*", "decision": "allow"}], on_disallowed_action, default_action="deny" + ) + data = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "input": "hi", "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("rules", "default_action", "detection_message"), + [ + ( + [{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], + "allow", + "Tool 'exec' denied by rule 'deny_exec'", + ), + ([{"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}], "deny", None), + ( + [ + {"id": "allow_exec", "tool_name": r"^exec$", "decision": "allow"}, + {"id": "deny_functions", "tool_name": r"^functions__.*$", "decision": "deny"}, + ], + "allow", + "Tool 'functions__exec' denied by rule 'deny_functions'", + ), + ( + [ + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_functions", "tool_type": r"^function$", "decision": "allow"}, + ], + "allow", + "Tool 'exec' denied by rule 'deny_exec'", + ), + ( + [ + {"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}, + {"id": "allow_all", "tool_name": r"^.*$", "decision": "allow"}, + ], + "deny", + "Tool 'exec' denied by rule 'deny_exec'", + ), + ( + [ + {"id": "allow_functions", "tool_name": r"^functions__.*$", "decision": "allow"}, + {"id": "deny_all", "tool_name": r"^.*$", "decision": "deny"}, + ], + "allow", + None, + ), + ], + ids=[ + "bare_name_deny_rule", + "bare_name_allow_rule", + "qualified_rule_wins_over_an_earlier_bare_rule", + "type_rule_keeps_an_earlier_bare_name_rule_in_charge", + "bare_name_deny_beats_a_later_catch_all_allow", + "qualified_allow_beats_a_later_catch_all_deny", + ], + ) + async def test_pre_call_judges_a_namespace_member_by_its_qualified_and_bare_names( + self, rules: list[dict[str, str]], default_action: Literal["allow", "deny"], detection_message: str | None + ) -> None: + guardrail = self._pre_call(rules, "block", default_action=default_action) + namespace = {"type": "namespace", "name": "functions", "tools": [{"type": "function", "name": "exec"}]} + data = {"model": "gpt-5", "input": "run date", "tools": [namespace]} + + if detection_message is None: + assert (await self._run_pre_call(guardrail, data, "responses"))["tools"] == [namespace] + return + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + assert excinfo.value.detail["detection_message"] == detection_message + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_a_namespace_member_a_bare_name_rule_denies(self) -> None: + guardrail = self._pre_call([{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], "rewrite") + members = [{"type": "function", "name": "exec"}, {"type": "function", "name": "read_file"}] + namespace = {"type": "namespace", "name": "functions", "tools": members} + data = {"model": "gpt-5", "input": "run date", "tools": [namespace]} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [{"type": "namespace", "name": "functions", "tools": members[1:]}] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("choice", "relaxed"), + [ + ({"tool_choice": {"type": "function", "function": {"name": "exec"}}}, {"tool_choice": "none"}), + ({"tool_choice": {"type": "function", "function": {"name": "ops"}}}, {"tool_choice": "none"}), + ( + { + "tool_choice": { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [ + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "get_weather"}}, + ], + }, + } + }, + { + "tool_choice": { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + }, + } + }, + ), + ( + {"toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["exec"]}}}, + {"toolConfig": {"functionCallingConfig": {"mode": "NONE"}}}, + ), + ], + ids=["bare_member_choice", "namespace_choice", "allowed_tools_entry", "allowed_function_names"], + ) + async def test_pre_call_rewrite_relaxes_a_bare_name_a_dropped_namespace_member_was_sent_under( + self, choice: dict[str, object], relaxed: dict[str, object] + ) -> None: + guardrail = self._pre_call([{"id": "deny_exec", "tool_name": r"^ops__exec$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} + namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + data = { + "model": "gemini/gemini-2.5-flash", + "messages": [{"role": "user", "content": "hi"}], + "tools": [weather, namespace], + **choice, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + assert {key: result[key] for key in relaxed} == relaxed + + @pytest.mark.asyncio + async def test_pre_call_rewrite_keeps_a_bare_name_another_kept_namespace_still_serves(self) -> None: + guardrail = self._pre_call([{"id": "deny_exec", "tool_name": r"^ops__exec$", "decision": "deny"}], "rewrite") + dropped_namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + kept_namespace = {"type": "namespace", "name": "dev", "tools": [{"type": "function", "name": "exec"}]} + tool_choice = {"type": "function", "function": {"name": "exec"}} + data = { + "model": "gemini/gemini-2.5-flash", + "messages": [{"role": "user", "content": "hi"}], + "tools": [dropped_namespace, kept_namespace], + "tool_choice": tool_choice, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [kept_namespace] + assert result["tool_choice"] == tool_choice + + @pytest.mark.asyncio + @pytest.mark.parametrize("exec_first", [True, False], ids=["exec_namespace_first", "exec_namespace_last"]) + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_denies_a_bare_name_another_namespace_shares_a_qualified_name_with( + self, on_disallowed_action: Literal["block", "rewrite"], exec_first: bool + ) -> None: + guardrail = self._pre_call( + [{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], on_disallowed_action + ) + exec_namespace = {"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]} + twin_namespace = {"type": "namespace", "name": "mcp", "tools": [{"type": "function", "name": "ops__exec"}]} + tools = [exec_namespace, twin_namespace] if exec_first else [twin_namespace, exec_namespace] + data = {"model": "gpt-5", "input": "run date", "tools": tools} + + if on_disallowed_action == "rewrite": + assert (await self._run_pre_call(guardrail, data, "responses"))["tools"] == [] + return + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "responses") + assert excinfo.value.detail["detection_message"] == "Tool 'exec' denied by rule 'deny_exec'" + + @pytest.mark.asyncio + @pytest.mark.parametrize("placement", ["tools", "functions"]) + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("member", "detection_message"), + [ + ({"type": "function", "function": {"name": "delete_db"}}, "Tool 'delete_db' denied by rule 'deny_delete'"), + ( + {"type": "function", "name": "lookup", "function": {"name": "delete_db"}}, + "Tool 'delete_db' denied by rule 'deny_delete'", + ), + ({"googleSearch": {}}, "Tool 'googleSearch' denied by rule 'deny_search'"), + ( + { + "type": "namespace", + "name": "inner", + "tools": [{"type": "function", "name": "noop"}], + "function": {"name": "delete_db"}, + }, + "Tool 'delete_db' denied by rule 'deny_delete'", + ), + ], + ids=["chat_shaped_member", "function_object_behind_another_name", "hosted_key_member", "nested_namespace"], + ) + async def test_pre_call_judges_a_namespace_member_by_what_a_provider_reads_from_it( + self, + member: dict[str, object], + detection_message: str, + on_disallowed_action: Literal["block", "rewrite"], + placement: Literal["tools", "functions"], + ) -> None: + guardrail = self._pre_call( + [ + {"id": "deny_delete", "tool_name": r"^delete_db$", "decision": "deny"}, + {"id": "deny_search", "tool_type": r"^googleSearch$", "decision": "deny"}, + ], + on_disallowed_action, + ) + weather = {"type": "function", "name": "get_weather"} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + placement: [{"type": "namespace", "name": "ops", "tools": [weather, member]}], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == detection_message + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result[placement] == [{"type": "namespace", "name": "ops", "tools": [weather]}] + + @pytest.mark.asyncio + async def test_pre_call_allows_a_chat_shaped_namespace_member_the_allowlist_names(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_create", "tool_name": r"^github__create_issue$", "decision": "allow"}], + "block", + default_action="deny", + ) + namespace = { + "type": "namespace", + "name": "github", + "tools": [{"type": "function", "function": {"name": "create_issue"}}], + } + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [namespace]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [namespace] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("tool_config", "rewritten_tool_config"), + [ + ( + { + "functionCallingConfig": { + "mode": "ANY", + "allowedFunctionNames": ["mcp__weather__lookup", "get_weather"], + } + }, + {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["get_weather"]}}, + ), + ( + {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["mcp__weather__lookup"]}}, + {"functionCallingConfig": {"mode": "NONE"}}, + ), + ( + {"function_calling_config": {"mode": "any", "allowed_function_names": ["mcp__weather__lookup"]}}, + {"function_calling_config": {"mode": "NONE"}}, + ), + ({"functionCallingConfig": {"mode": "ANY"}}, {"functionCallingConfig": {"mode": "ANY"}}), + ( + {"functionCallingConfig": {"mode": "AUTO", "allowedFunctionNames": []}}, + {"functionCallingConfig": {"mode": "AUTO", "allowedFunctionNames": []}}, + ), + ], + ids=[ + "some_allowed_names_left", + "every_allowed_name_removed", + "snake_case_fields", + "no_allowed_names", + "empty_allowed_names_list", + ], + ) + async def test_pre_call_rewrite_drops_removed_functions_from_a_gemini_tool_config( + self, tool_config: dict[str, object], rewritten_tool_config: dict[str, object] + ) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + config_key = "tool_config" if "function_calling_config" in tool_config else "toolConfig" + data = { + "model": "gemini-3", + "contents": [{"role": "user", "parts": [{"text": "weather in Paris?"}]}], + "tools": [{"functionDeclarations": [{"name": "mcp__weather__lookup"}, {"name": "get_weather"}]}], + config_key: tool_config, + } + + result = await self._run_pre_call(guardrail, data, "agenerate_content") + + assert result["tools"] == [{"functionDeclarations": [{"name": "get_weather"}]}] + assert result[config_key] == rewritten_tool_config + + @pytest.mark.asyncio + async def test_pre_call_rewrite_turns_off_forced_gemini_calling_when_no_function_is_left(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + data = { + "model": "gemini-3", + "contents": [{"role": "user", "parts": [{"text": "weather in Paris?"}]}], + "tools": [{"functionDeclarations": [{"name": "mcp__weather__lookup"}]}, {"googleSearch": {}}], + "toolConfig": {"functionCallingConfig": {"mode": "ANY"}}, + } + + result = await self._run_pre_call(guardrail, data, "agenerate_content") + + assert result["tools"] == [{"googleSearch": {}}] + assert result["toolConfig"] == {"functionCallingConfig": {"mode": "NONE"}} + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_judges_a_named_hosted_tool_by_its_type_as_well_as_its_name( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + "tools": [weather, {"type": "web_search", "name": "get_weather"}], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'web_search' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_judges_a_gemini_hosted_key_by_the_key_whatever_name_it_carries( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + computer_use = {"computerUse": {"name": "get_weather", "environment": "browser"}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [weather, computer_use]} + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'computerUse' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "rules, tool", + [ + ( + [{"id": "allow_bash", "tool_name": r"^bash$", "decision": "allow"}], + {"type": "bash_20250124", "name": "bash"}, + ), + ( + [{"id": "allow_edit", "tool_name": r"^str_replace_based_edit_tool$", "decision": "allow"}], + {"type": "text_editor_20250728", "name": "str_replace_based_edit_tool"}, + ), + ( + [{"id": "allow_search", "tool_name": r"^web_search$", "decision": "allow"}], + {"type": "web_search_20250305", "name": "web_search"}, + ), + ( + [ + { + "id": "allow_search", + "tool_name": r"^web_search$", + "tool_type": r"^web_search_20250305$", + "decision": "allow", + } + ], + {"type": "web_search_20250305", "name": "web_search"}, + ), + ], + ids=["client_tool_by_name", "client_tool_with_its_own_name", "server_tool_by_name", "name_and_type_rule"], + ) + async def test_pre_call_allows_an_anthropic_defined_tool_that_an_allowlist_names( + self, rules: list[dict[str, object]], tool: dict[str, object] + ) -> None: + guardrail = self._pre_call(rules, "block", default_action="deny") + data = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "tool, call_type", + [ + ({"type": "web_search_20250305", "name": "get_weather"}, "anthropic_messages"), + ( + { + "type": "computer_20251124", + "name": "get_weather", + "function": {"parameters": {"display_width_px": 1024, "display_height_px": 768}}, + }, + "acompletion", + ), + ({"type": "computer_use", "name": "get_weather", "environment": "browser"}, "acompletion"), + ({"type": "text_editor_20250728", "name": "get_weather"}, "anthropic_messages"), + ], + ids=["server_tool", "anthropic_computer_tool", "gemini_computer_use", "text_editor"], + ) + async def test_pre_call_blocks_a_provider_defined_tool_whose_name_is_not_its_own( + self, tool: dict[str, object], call_type: CallTypesLiteral + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + "block", + default_action="deny", + ) + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, call_type) + + assert excinfo.value.detail["detection_message"] == f"Tool '{tool['type']}' denied by default action" + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + "tool", + [ + {"type": "function", "function": {"name": "delete_db"}, "functionDeclarations": [{"name": "get_weather"}]}, + {"type": "function", "function": {"name": "delete_db"}, "function_declarations": [{"name": "get_weather"}]}, + {"type": "custom", "name": "get_weather", "function": {"name": "delete_db"}}, + {"type": "mcp", "server_label": "get_weather", "function": {"name": "delete_db"}}, + { + "type": "namespace", + "name": "github", + "tools": [{"type": "function", "name": "create_issue"}], + "function": {"name": "delete_db"}, + }, + { + "type": "computer_20251124", + "name": "computer", + "function": {"name": "delete_db", "parameters": {"display_width_px": 1024, "display_height_px": 768}}, + }, + ], + ids=[ + "declarations_list", + "snake_case_declarations_list", + "custom_top_level_name", + "mcp_server_label", + "namespace_with_members", + "computer_top_level_name", + ], + ) + async def test_pre_call_judges_the_function_a_tool_runs_behind_an_allowed_name( + self, tool: dict[str, object], on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_known", "tool_name": r"^(get_weather|computer|github__create_issue)$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [weather, tool]} + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'delete_db' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("tool", "hosted_key"), + [ + ({"type": "get_weather", "googleSearch": {}}, "googleSearch"), + ({"type": "function", "codeExecution": {}}, "codeExecution"), + ({"type": "get_weather", "url_context": {}}, "url_context"), + ], + ids=["named_type_google_search", "function_type_code_execution", "named_type_url_context"], + ) + async def test_pre_call_judges_a_gemini_hosted_tool_key_whatever_the_tool_type( + self, tool: dict[str, object], hosted_key: str, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [ + {"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}, + {"id": "allow_functions", "tool_type": r"^function$", "decision": "allow"}, + ], + on_disallowed_action, + default_action="deny", + ) + weather = {"type": "function", "function": {"name": "get_weather"}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [weather, tool]} + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == f"Tool '{hosted_key}' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_leaves_a_gemini_hosted_key_that_sits_beside_a_function( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + tool = {"type": "function", "function": {"name": "get_weather"}, "googleSearch": {}} + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + @pytest.mark.parametrize( + ("function", "denied"), + [ + ({"type": "function", "function": {"name": "delete_db"}}, "delete_db"), + ({"name": "get_weather", "function": {"name": "delete_db"}}, "delete_db"), + ({"googleSearch": {}}, "googleSearch"), + ], + ids=["function_object", "function_object_behind_an_allowed_name", "gemini_hosted_key"], + ) + async def test_pre_call_judges_a_legacy_functions_entry_like_a_tool( + self, function: dict[str, object], denied: str, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"name": "get_weather"} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + "functions": [weather, function], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == f"Tool '{denied}' denied by default action" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["functions"] == [weather] + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_judges_a_legacy_functions_entry_as_a_function_whatever_its_type( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "deny_delete", "tool_name": r"^delete_db$", "tool_type": r"^function$", "decision": "deny"}], + on_disallowed_action, + ) + weather = {"name": "get_weather"} + data = { + "model": "gemini-3", + "messages": [{"role": "user", "content": "hi"}], + "functions": [weather, {"type": "custom", "name": "delete_db"}], + } + + if on_disallowed_action == "block": + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "acompletion") + assert excinfo.value.detail["detection_message"] == "Tool 'delete_db' denied by rule 'deny_delete'" + return + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["functions"] == [weather] + + @pytest.mark.asyncio + async def test_rewrite_keeps_an_allowed_function_whose_declaration_list_is_denied(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + "rewrite", + default_action="deny", + ) + tool = { + "type": "function", + "function": {"name": "get_weather"}, + "functionDeclarations": [{"name": "delete_db"}], + } + data = {"model": "gemini-3", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [{"type": "function", "function": {"name": "get_weather"}}] + + @pytest.mark.asyncio + async def test_pre_call_allows_a_custom_tool_by_its_function_name(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + "block", + default_action="deny", + ) + tool = {"type": "custom", "function": {"name": "get_weather"}} + data = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + async def test_pre_call_blocks_a_tool_whose_type_is_not_a_string_by_its_name(self) -> None: + guardrail = self._pre_call( + [ + {"id": "allow_search", "tool_type": r"^web_search$", "decision": "allow"}, + {"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}, + ], + "block", + ) + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": ["custom"], "name": "Read", "input_schema": {"type": "object"}}], + } + + with pytest.raises(HTTPException) as excinfo: + await self._run_pre_call(guardrail, data, "anthropic_messages") + + assert excinfo.value.status_code == 400 + assert excinfo.value.detail["detection_message"] == "Tool 'Read' denied by rule 'deny_read'" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("malformed_tool", "tool_choice", "call_type"), + [ + ({"type": ["custom"], "name": "Read", "input_schema": {"type": "object"}}, None, "anthropic_messages"), + ({"type": "function", "name": "Read"}, {"type": {"a": 1}}, "responses"), + ], + ids=["tool_type_is_a_list", "tool_choice_type_is_an_object"], + ) + async def test_pre_call_rewrite_strips_a_denied_tool_and_leaves_malformed_fields_for_the_provider( + self, malformed_tool: dict[str, object], tool_choice: dict[str, object] | None, call_type: CallTypesLiteral + ) -> None: + guardrail = self._pre_call([{"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}], "rewrite") + bash = {"type": "function", "name": "Bash"} + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "tools": [bash, malformed_tool], + **({"tool_choice": tool_choice} if tool_choice is not None else {}), + } + + result = await self._run_pre_call(guardrail, data, call_type) + + assert result["tools"] == [bash] + assert result.get("tool_choice") == tool_choice + + @pytest.mark.asyncio + async def test_pre_call_judges_a_tool_with_a_null_type_by_its_nested_function_name(self) -> None: + guardrail = self._pre_call( + [{"id": "allow_read", "tool_name": r"^Read$", "decision": "allow"}], "block", default_action="deny" + ) + tool = {"type": None, "function": {"name": "Read", "parameters": {"type": "object"}}} + data = {"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}], "tools": [tool]} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == [tool] + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_a_removed_hosted_tool_from_allowed_tools(self) -> None: + guardrail = self._pre_call([{"id": "deny_search", "tool_type": r"^web_search$", "decision": "deny"}], "rewrite") + weather = {"type": "function", "name": "get_weather"} + tool_choice = {"type": "allowed_tools", "mode": "required", "tools": [weather, {"type": "web_search"}]} + data = {"model": "gpt-5", "input": "hi", "tools": [weather, {"type": "web_search"}], "tool_choice": tool_choice} + + result = await self._run_pre_call(guardrail, data, "responses") + + assert result["tools"] == [weather] + assert result["tool_choice"] == {"type": "allowed_tools", "mode": "required", "tools": [weather]} + + @pytest.mark.asyncio + @pytest.mark.parametrize("tools", [{"not": "a list"}, '[{"type": "function"}]'], ids=["object", "json_string"]) + async def test_pre_call_rewrite_leaves_a_tools_field_that_is_not_a_list_for_the_provider_to_reject( + self, tools: object + ) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + data = { + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "tools": tools, + "tool_choice": {"type": "function", "function": {"name": "mcp__weather__lookup"}}, + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == tools + assert result["tool_choice"] == "none" + + @pytest.mark.asyncio + @pytest.mark.parametrize("on_disallowed_action", ["block", "rewrite"]) + async def test_pre_call_reads_bedrock_tool_blocks_by_name_and_skips_cache_points( + self, on_disallowed_action: Literal["block", "rewrite"] + ) -> None: + guardrail = self._pre_call( + [{"id": "allow_weather", "tool_name": r"^get_weather$", "decision": "allow"}], + on_disallowed_action, + default_action="deny", + ) + weather = {"toolSpec": {"name": "get_weather", "inputSchema": {"json": {}}}} + tools = [weather, {"cachePoint": {"type": "default"}}] + data = {"model": "bedrock/claude", "messages": [{"role": "user", "content": "hi"}], "tools": tools} + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == tools + + @pytest.mark.asyncio + async def test_pre_call_rewrite_drops_a_denied_bedrock_tool_block_by_its_name(self) -> None: + guardrail = self._pre_call([{"id": "deny_mcp", "tool_name": r"^mcp__.*$", "decision": "deny"}], "rewrite") + kept = [{"toolSpec": {"name": "get_weather"}}, {"cachePoint": {"type": "default"}}] + data = { + "model": "bedrock/claude", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"toolSpec": {"name": "mcp__weather__lookup"}}, *kept], + } + + result = await self._run_pre_call(guardrail, data, "acompletion") + + assert result["tools"] == kept + @pytest.mark.asyncio async def test_denied_anthropic_tool_use_is_blocked(self): response = self._response({"type": "text", "text": "reading"}, self._tool_use("Read")) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index ef43fe56146..b7d0d9e62cd 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2556,6 +2556,66 @@ class TestResponsesWSFirstFrameModelAuth: "response": {"input": "hi [guarded]", "tools": tools[:1]}, } + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [True, False]) + async def test_frame_guard_forwards_the_tool_choice_a_guardrail_rewrote(self, nested: bool): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {**data, "tools": [], "tool_choice": "none"} + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + body = { + "input": "hi", + "tools": [{"type": "function", "name": "a"}], + "tool_choice": {"type": "function", "name": "a"}, + } + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + guarded = json.loads((await guard(json.dumps(frame))).frame) + + sent = guarded["response"] if nested else guarded + assert (sent["tools"], sent["tool_choice"]) == ([], "none") + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [True, False]) + async def test_frame_guard_forwards_the_extra_body_tools_a_rewrite_guardrail_stripped(self, nested: bool): + from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = ToolPermissionGuardrail( + guardrail_name="ws-rewrite", + rules=[{"id": "deny_shell", "tool_name": r"^shell$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=MagicMock(), data=data, call_type=call_type + ) + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + weather = {"type": "function", "name": "get_weather"} + body = {"input": "hi", "extra_body": {"tools": [weather, {"type": "function", "name": "shell"}]}} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + guarded = json.loads((await guard(json.dumps(frame))).frame) + + sent = guarded["response"] if nested else guarded + assert sent["extra_body"] == {"tools": [weather]} + @pytest.mark.asyncio async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard @@ -2991,7 +3051,7 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) + @pytest.mark.parametrize("field", ["input", "instructions", "tools", "tool_choice", "extra_body"]) async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field: str): from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2191454a2f7..15d440c0ad4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -48150,7 +48150,7 @@ export interface components { id: string; /** * Tool Name - * @description Regex pattern applied to the tool's function name + * @description Regex pattern applied to the tool's name, or to its type when the tool has no name */ tool_name?: string | null; /**