From cab0d82be1e6ceae2e79520be669619d7bf8245c Mon Sep 17 00:00:00 2001 From: lawrence Date: Thu, 5 Nov 2026 15:48:24 +0800 Subject: [PATCH] fix(presidio): reject duplicate JSON tool argument keys --- litellm/proxy/guardrails/_content_utils.py | 10 +++-- .../guardrails/guardrail_hooks/presidio.py | 21 ++++++--- .../guardrail_hooks/test_presidio.py | 45 ++++++++++++++++++- 3 files changed, 65 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 51934c99db6..8034ca74f56 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -97,10 +97,14 @@ def _text_object(value: object) -> Mapping[str, object] | None: return None +class DuplicateToolArgumentKeyError(ValueError): + pass + + def _tool_json_object(pairs: Sequence[tuple[str, JsonValue]]) -> JsonValue: value: Final = dict(pairs) if len(value) != len(pairs): - raise ValueError("Duplicate tool argument keys") + raise DuplicateToolArgumentKeyError("Duplicate tool argument keys") return value @@ -111,7 +115,7 @@ def _finite_number(literal: str) -> float: return number -def _tool_arguments(arguments: str) -> JsonValue: +def parse_tool_arguments(arguments: str) -> JsonValue: parsed: Final = _tool_json( json.loads( arguments, @@ -363,7 +367,7 @@ def walk_user_text( if not isinstance(arguments, str) or not arguments: return arguments try: - parsed: Final = _tool_arguments(arguments) + parsed: Final = parse_tool_arguments(arguments) except (ValueError, RecursionError): visited += 1 # rebind-ok: malformed arguments still undergo the text callback return visit(arguments) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index c082989f237..3e00f1dd905 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -30,7 +30,7 @@ from operator import itemgetter from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias, cast import aiohttp -from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config +from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -55,6 +55,7 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames +from litellm.proxy.guardrails._content_utils import DuplicateToolArgumentKeyError, parse_tool_arguments from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, assemble_anthropic_sse_stream, @@ -303,7 +304,15 @@ def _set_tool_call_arguments(tool_call: _ToolCall, arguments: str) -> None: function_chunk["arguments"] = arguments -_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +def _parsed_arguments(arguments: str) -> JsonValue: + try: + return parse_tool_arguments(arguments) + except DuplicateToolArgumentKeyError as error: + raise GuardrailRaisedException( + message="Presidio cannot safely scan tool arguments with duplicate JSON object keys", + should_wrap_with_default_message=False, + blocked_content=True, + ) from error def _json_literal_tokens(value: JsonValue) -> frozenset[str]: @@ -323,8 +332,8 @@ def _segment_literal_tokens(segment: tuple[_TextSlot | _ToolCallSlot, str]) -> f if isinstance(slot, _TextSlot): return _literal_tokens(source) try: - parsed: Final = _JSON_VALUE.validate_json(source) - except ValidationError: + parsed: Final = _parsed_arguments(source) + except ValueError: return _literal_tokens(source) return _json_literal_tokens(parsed) @@ -353,8 +362,8 @@ async def _rewritten_arguments(arguments: str, rewrite: Callable[[str, bool], Aw quote or a newline has to be escaped for the arguments to stay valid JSON. """ try: - parsed: Final = _JSON_VALUE.validate_json(arguments) - except ValidationError: + parsed: Final = _parsed_arguments(arguments) + except ValueError: return await rewrite(arguments, True) rewritten: Final = await _rewritten_json_strings(parsed, rewrite) return arguments if rewritten == parsed else json.dumps(rewritten, ensure_ascii=False) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 90e52a925bc..95bcda488ab 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -5592,8 +5592,49 @@ async def test_escaped_tool_argument_tokens_never_restore_to_unrelated_pii(argum await guardrail._close_http_session() -async def test_safe_tool_argument_keys_and_scalar_types_keep_original_bytes() -> None: - arguments: Final = '{ "count": 42, "fraction": 1.25, "enabled": true, "absent": null, "items": [false, 0] }' +@pytest.mark.parametrize( + ("input_type", "action"), + [("request", PiiAction.MASK), ("request", PiiAction.BLOCK), ("response", PiiAction.MASK)], +) +@pytest.mark.parametrize( + "arguments", + [ + '{"guest":"test@example.com","guest":"safe"}', + '{"items":[{"guest":"test@example.com","guest":"safe"}]}', + r'{"guest":"test@example.com","\u0067uest":"safe"}', + '{"guest":"","guest":"safe"}', + ], +) +async def test_duplicate_tool_argument_keys_cannot_hide_values_from_presidio( + arguments: str, input_type: Literal["request", "response"], action: PiiAction +) -> None: + analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue() + guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True, action=action) + tool_call: Final = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="lookup", arguments=arguments) + ) + try: + with pytest.raises(GuardrailRaisedException, match="duplicate JSON object keys") as excinfo: + await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [tool_call]}, + request_data={"metadata": {"pii_tokens": {"": _EMAIL}}}, + input_type=input_type, + ) + assert excinfo.value.status_code == 400 + assert excinfo.value.blocked_content is True + assert _EMAIL not in str(excinfo.value) + finally: + await guardrail._close_http_session() + + +@pytest.mark.parametrize( + "arguments", + [ + '{ "count": 42, "fraction": 1.25, "enabled": true, "absent": null, "items": [false, 0] }', + '{ "left": {"guest": "safe"}, "right": {"guest": "safe"} }', + ], +) +async def test_safe_tool_argument_keys_and_scalar_types_keep_original_bytes(arguments: str) -> None: analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue() guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, output_parse_pii=True) tool_call: Final = ChatCompletionMessageToolCall(