From bcf5f2002c64085798ee53513147c9de6918ccd1 Mon Sep 17 00:00:00 2001
From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 15:42:51 -0700
Subject: [PATCH] fix(guardrails): write a tool visitor's whole payload when it
 changes more than string values

---
 litellm/proxy/guardrails/_content_utils.py    | 16 +++++++--------
 .../proxy/guardrails/test_content_utils.py    | 20 +++++++++++++++----
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py
index ab4e158b1d2..af7d8f0e977 100644
--- a/litellm/proxy/guardrails/_content_utils.py
+++ b/litellm/proxy/guardrails/_content_utils.py
@@ -116,17 +116,17 @@ def _value_string_spans(arguments: str) -> Iterator[tuple[int, int]]:
 
 
 def _respliced(arguments: str, parsed: JsonValue, rewritten: JsonValue) -> str:
-    """Write the string values that differ between ``parsed`` and ``rewritten`` into ``arguments``, keeping its other text."""
-    spans: Final = tuple(_value_string_spans(arguments))
-    before: Final = tuple(iter_tool_json_text(parsed))
-    after: Final = tuple(iter_tool_json_text(rewritten))
-    if not len(spans) == len(before) == len(after):
-        return json.dumps(rewritten, ensure_ascii=False)
-    changed: Final = tuple((span, new) for span, old, new in zip(spans, before, after) if new != old)
+    """Serialize ``rewritten`` by replacing only its changed string values in ``arguments``, when nothing else changed."""
+    serialized: Final = json.dumps(rewritten, ensure_ascii=False)
+    strings: Final = iter(tuple(iter_tool_json_text(rewritten)))
+    if json.dumps(map_tool_json_text(parsed, lambda text: next(strings, text)), ensure_ascii=False) != serialized:
+        return serialized
+    values: Final = zip(_value_string_spans(arguments), iter_tool_json_text(parsed), iter_tool_json_text(rewritten))
+    changed: Final = tuple((span, new) for span, old, new in values if new != old)
     starts: Final = (0, *(end for (_, end), _ in changed))
     ends: Final = (*(start for (start, _), _ in changed), len(arguments))
     replacements: Final = (*(json.dumps(new, ensure_ascii=False) for _, new in changed), "")
-    return "".join(arguments[start:end] + replacement for start, end, replacement in zip(starts, ends, replacements))
+    return "".join(arguments[start:end] + text for start, end, text in zip(starts, ends, replacements))
 
 
 def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:
diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py
index bca50550b34..a6c366759ad 100644
--- a/tests/unit/proxy/guardrails/test_content_utils.py
+++ b/tests/unit/proxy/guardrails/test_content_utils.py
@@ -1044,10 +1044,22 @@ def test_masked_tool_arguments_change_only_the_masked_strings(arguments: str, ma
     assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == masked
 
 
-def test_tool_visitor_that_changes_the_payload_shape_writes_its_payload_as_arguments() -> None:
-    data = {"messages": [{"tool_calls": [{"function": {"arguments": '{ "token": "secret", "note": "keep" }'}}]}]}
-    assert walk_user_text(data, lambda text: text, lambda payload: {"token": "[MASKED]"}) == 1
-    assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == '{"token": "[MASKED]"}'
+@pytest.mark.parametrize(
+    "rewritten",
+    (
+        {"token": "[MASKED]"},
+        {"token": "[MASKED]", "renamed": "keep", "count": 7},
+        {"token": "[MASKED]", "note": "keep", "count": 8},
+        {"token": "[MASKED]", "note": "keep", "count": True},
+    ),
+    ids=("key-dropped", "key-renamed", "number-changed", "number-retyped"),
+)
+def test_tool_visitor_that_changes_more_than_strings_writes_its_payload_as_arguments(rewritten: JsonValue) -> None:
+    data = {
+        "messages": [{"tool_calls": [{"function": {"arguments": '{ "token": "secret", "note": "keep", "count": 1 }'}}]}]
+    }
+    assert walk_user_text(data, lambda text: text, lambda payload: rewritten) == 1
+    assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == json.dumps(rewritten)
 
 
 @pytest.mark.parametrize(