fix(presidio): reject duplicate JSON tool argument keys

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)

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)

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":"<EMAIL_ADDRESS_1>","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_ADDRESS_1>": _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(