fix(guardrails): bound tool payload scans and stop nested tool_result recursion

Scan each tool value beside its key instead of an indented dump, read
tool_result content one level deep, keep non-finite tool arguments as
raw text, and give read-only guardrails one fragment per tool payload

enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py

@@ -30,7 +30,7 @@ from litellm.integrations.custom_guardrail import (
     log_guardrail_information,
 )
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.guardrails._content_utils import iter_tool_json_text, map_tool_json_text, walk_user_text
+from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import GenericGuardrailAPIInputs
 
@@ -584,10 +584,16 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
         return pattern.sub("[REDACTED]", text)
 
     def redact_tool_payload(self, payload: JsonValue) -> JsonValue:
-        texts: Final = tuple(iter_tool_json_text(payload))
-        scan_text: Final = "\n".join((*texts, json.dumps(payload, indent=2, ensure_ascii=False)))
+        """Redact secrets in a tool payload's string values, scanning each value alone and beside its key."""
+        fields: Final = tuple(iter_tool_json_fields(payload))
+        texts: Final = tuple(text for _, text in fields)
+        keyed: Final = tuple(
+            f"{json.dumps(key, ensure_ascii=False)}: {json.dumps(text, ensure_ascii=False)}"
+            for key, text in fields
+            if key is not None
+        )
         detected: Final = tuple(
-            secret for secret in self.scan_message_for_secrets(scan_text)
+            secret for secret in self.scan_message_for_secrets("\n".join((*texts, *keyed)))
             if any(secret["value"] in text for text in texts)
         )
         if not detected:
@@ -596,6 +602,10 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
         if counts is not None:
             for secret in detected:
                 counts[secret["type"]] = counts.get(secret["type"], 0) + 1
+        verbose_proxy_logger.warning(
+            "Detected and redacted secrets in tool payload: %s",
+            sorted(dict.fromkeys(secret["type"] for secret in detected)),
+        )
         pattern: Final = re.compile("|".join(re.escape(secret["value"]) for secret in detected))
         return map_tool_json_text(payload, lambda text: pattern.sub("[REDACTED]", text))
 

litellm/proxy/guardrails/_content_utils.py

@@ -9,6 +9,7 @@ every text fragment.
 """
 
 import json
+import math
 from collections.abc import Callable, Iterator, Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any, Final
@@ -81,31 +82,64 @@ def _tool_json_object(pairs: Sequence[tuple[str, JsonValue]]) -> JsonValue:
     return value
 
 
+def _finite_number(literal: str) -> float:
+    number: Final = float(literal)
+    if not math.isfinite(number):
+        raise ValueError("Non-finite tool argument number")
+    return number
+
+
 def _tool_arguments(arguments: str) -> JsonValue:
-    return _tool_json(json.loads(arguments, object_pairs_hook=_tool_json_object))
+    return _tool_json(
+        json.loads(
+            arguments,
+            object_pairs_hook=_tool_json_object,
+            parse_float=_finite_number,
+            parse_constant=_finite_number,
+        )
+    )
 
 
-def iter_tool_json_text(value: object) -> Iterator[str]:
-    pending: Final[list[tuple[object, bool]]] = [(value, False)]  # mutable-ok: owned DFS stack avoids Python recursion
+def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:
+    """Yield ``(nearest key, text)`` for every non-empty string leaf; list items take their parent's key."""
+    pending: Final[list[tuple[object, str | None, bool]]] = [  # mutable-ok: owned DFS stack avoids Python recursion
+        (value, None, False)
+    ]
     active: Final[set[int]] = set()  # mutable-ok: track active ancestry only
     while pending:
-        current, leaving = pending.pop()
+        current, key, leaving = pending.pop()
         if leaving:
             active.remove(id(current))
             continue
         if isinstance(current, str):
             if current:
-                yield current
+                yield key, current
             continue
         if current is None or isinstance(current, (bool, int, float)):
             continue
         fields = _text_object(current)
-        children = tuple(fields.values()) if fields is not None else _object_list(current)
+        if fields is not None:
+            children = tuple(fields.items())
+        else:
+            items = _object_list(current)
+            children = None if items is None else tuple((key, child) for child in items)
         if children is None or id(current) in active:
             raise ValueError("Invalid tool JSON payload")
         active.add(id(current))
-        pending.append((current, True))
-        pending.extend((child, False) for child in reversed(children))
+        pending.append((current, key, True))
+        pending.extend((child, child_key, False) for child_key, child in reversed(children))
+
+
+def iter_tool_json_text(value: object) -> Iterator[str]:
+    for _, text in iter_tool_json_fields(value):
+        yield text
+
+
+def _tool_payload_text(value: object) -> Iterator[str]:
+    """Yield a tool payload's string values as one newline-joined fragment, like one text part."""
+    text: Final = "\n".join(iter_tool_json_text(value))
+    if text:
+        yield text
 
 
 @dataclass(frozen=True, slots=True)
@@ -172,9 +206,9 @@ def _iter_native_content_text(content: object) -> Iterator[str]:
         yield from _iter_text_parts_in_content([part])
         fields = _text_object(part)
         if fields is not None and fields.get("type") == "tool_result":
-            yield from _iter_native_content_text(fields.get("content"))
+            yield from _iter_text_parts_in_content(fields.get("content"))
         elif fields is not None and fields.get("type") == "tool_use":
-            yield from iter_tool_json_text(fields.get("input"))
+            yield from _tool_payload_text(fields.get("input"))
 
 
 # Responses-API item types whose ``output`` field carries user/tool text
@@ -283,9 +317,9 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
                         except (ValueError, RecursionError):
                             yield arguments
                         else:
-                            yield from iter_tool_json_text(parsed)
+                            yield from _tool_payload_text(parsed)
                     elif isinstance(arguments, dict):
-                        yield from iter_tool_json_text(arguments)
+                        yield from _tool_payload_text(arguments)
 
 
 def walk_user_text(
@@ -295,7 +329,7 @@ def walk_user_text(
 ) -> int:
     """Rewrite every text fragment in place via ``visit``.
 
-    Mutates ``data["messages"]`` and ``data["input"]``. Returns the number
+    Mutates ``data["messages"]`` (including tool payloads), ``data["system"]`` and ``data["input"]``. Returns the number
     of fragments visited so callers can short-circuit when nothing was
     inspected.
     """
@@ -326,19 +360,19 @@ def walk_user_text(
         rewritten: Final = _rewrite_tool(parsed)
         return arguments if rewritten == parsed else json.dumps(rewritten, ensure_ascii=False)
 
-    def _rewrite_part(part: Mapping[str, object]) -> Mapping[str, object]:
+    def _rewrite_part(part: Mapping[str, object], native: bool) -> Mapping[str, object]:
         nonlocal visited
         text: Final = _part_text(part)
         rewritten: Final = {**part, **({"text": visit(text)} if text is not None else {})}
         if text is not None:
             visited += 1  # rebind-ok: tally each existing text field alongside the native payload
-        if part.get("type") == "tool_result" and "content" in part:
-            return {**rewritten, "content": _rewrite_content(part["content"])}
-        if part.get("type") == "tool_use" and "input" in part:
+        if native and part.get("type") == "tool_result" and "content" in part:
+            return {**rewritten, "content": _rewrite_content(part["content"], native=False)}
+        if native and part.get("type") == "tool_use" and "input" in part:
             return {**rewritten, "input": _rewrite_tool(part["input"])}
         return rewritten
 
-    def _rewrite_content(content: object) -> object:
+    def _rewrite_content(content: object, native: bool = True) -> object:
         nonlocal visited
         if isinstance(content, str):
             if content:
@@ -352,7 +386,7 @@ def walk_user_text(
                     visited += 1
                     new_parts.append(visit(part))
                 elif isinstance(part, dict):
-                    new_parts.append(_rewrite_part(part))
+                    new_parts.append(_rewrite_part(part, native))
                 else:
                     new_parts.append(part)
             return new_parts

tests/integration/security/test_banned_keywords_tool_traffic.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from functools import reduce
+from typing import Final
+
+import pytest
+from integration.security.test_hide_secrets_tool_traffic import Rig, tool_rig
+from pydantic import JsonValue
+
+KEYWORD: Final = "zebracornflux"
+
+
+def _ban_keyword(configuration: dict[str, JsonValue]) -> None:
+    settings: Final = configuration["litellm_settings"]
+    assert isinstance(settings, dict)
+    settings["banned_keywords_list"] = [KEYWORD]
+    settings["callbacks"] = ["banned_keywords"]
+    del configuration["guardrails"]
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+    with tool_rig(tmp_path_factory.mktemp("banned-keywords-tools"), configure=_ban_keyword) as value:
+        yield value
+
+
+def _chat(arguments: str, *, stream: bool = False) -> dict[str, JsonValue]:
+    return {
+        "model": "gpt",
+        "stream": stream,
+        "cache": {"no-cache": True},
+        "messages": [
+            {"role": "user", "content": "read .env"},
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {"id": "call_1", "type": "function", "function": {"name": "read_file", "arguments": arguments}}
+                ],
+            },
+            {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+        ],
+    }
+
+
+@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
+def test_keyword_in_tool_call_arguments_is_blocked(rig: Rig, stream: bool) -> None:
+    rig.wire.drain()
+    response: Final = rig.proxy.request(
+        "POST", "/v1/chat/completions", _chat(json.dumps({"path": f"notes/{KEYWORD}.txt"}), stream=stream)
+    )
+    assert response.status_code == 400, response.text
+    assert f"Keyword={KEYWORD}" in response.text
+    assert rig.wire.drain() == ()
+
+
+def test_tool_call_arguments_without_keyword_reach_provider(rig: Rig) -> None:
+    rig.wire.drain()
+    arguments: Final = json.dumps({"path": "notes/plain.txt"})
+    response: Final = rig.proxy.request("POST", "/v1/chat/completions", _chat(arguments))
+    assert response.status_code == 200, response.text
+    sent: Final = rig.wire.drain()
+    assert len(sent) == 1
+    assert json.loads(sent[0].body)["messages"][1]["tool_calls"][0]["function"]["arguments"] == arguments
+
+
+def test_deep_tool_result_chain_cannot_hide_a_later_keyword(rig: Rig) -> None:
+    rig.wire.drain()
+    chain: Final = reduce(
+        lambda content, _: [{"type": "tool_result", "tool_use_id": "tool_1", "content": content}],
+        range(1000),
+        [{"type": "text", "text": "inner"}],
+    )
+    response: Final = rig.proxy.request(
+        "POST",
+        "/v1/chat/completions",
+        {
+            "model": "gpt",
+            "cache": {"no-cache": True},
+            "messages": [{"role": "assistant", "content": chain}, {"role": "user", "content": f"say {KEYWORD}"}],
+        },
+    )
+    assert response.status_code == 400, response.text
+    assert f"Keyword={KEYWORD}" in response.text
+    assert rig.wire.drain() == ()

tests/integration/security/test_hide_secrets_tool_traffic.py

@@ -15,6 +15,7 @@ import yaml
 from integration._support.client import Gateway
 from integration._support.database import scratch_database
 from integration._support.process import OwnedProxy, owned_proxy_process
+from integration._support.timing import assert_not_delayed_by, timed
 from integration._support.wire import Reply, Request, Wire, wire_server
 from integration.security._license_edge import license_edge
 from pydantic import JsonValue
@@ -317,3 +318,42 @@ def test_chat_tool_arguments_remain_valid_json(rig: Rig) -> None:
     assert arguments["path"] == _input()["path"]
     assert arguments["enabled"] is False
     assert arguments["optional"] is None
+
+
+def test_deeply_nested_tool_arguments_do_not_stall_the_worker(rig: Rig) -> None:
+    def call(depth: int) -> float:
+        rig.wire.drain()
+        arguments: Final = '{"k":' * depth + json.dumps(TEXT) + "}" * depth
+        sent: Final = timed(
+            lambda: rig.proxy.request(
+                "POST",
+                "/v1/chat/completions",
+                {
+                    "model": "gpt",
+                    "cache": {"no-cache": True},
+                    "messages": [
+                        {"role": "user", "content": "read .env"},
+                        {
+                            "role": "assistant",
+                            "content": None,
+                            "tool_calls": [
+                                {
+                                    "id": "call_1",
+                                    "type": "function",
+                                    "function": {"name": "read", "arguments": arguments},
+                                }
+                            ],
+                        },
+                        {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+                    ],
+                },
+            )
+        )
+        assert sent.value.status_code == 200, sent.value.text
+        upstream: Final = rig.wire.drain()
+        assert len(upstream) == 1
+        assert SECRET.encode() not in upstream[0].body
+        return sent.seconds
+
+    control: Final = tuple(call(1) for _ in range(3))
+    assert_not_delayed_by((call(300),), control=control, injected_delay=5.0)

tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py

@@ -81,6 +81,18 @@ async def test_tool_arguments_native_hook_cannot_forward_unredacted_credentials(
     assert "[REDACTED]" in rewritten
 
 
+@pytest.mark.asyncio
+async def test_deeply_nested_tool_arguments_are_scanned_in_bounded_time():
+    arguments = '{"k":' * 200 + json.dumps(f"AWS_ACCESS_KEY_ID={AWS_KEY}") + "}" * 200
+    data = {"messages": [{"role": "assistant", "tool_calls": [{"id": "call_1", "function": {
+        "name": "save_config", "arguments": arguments}}]}]}
+    started = time.perf_counter()
+    await _guardrail().async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion")
+    assert time.perf_counter() - started < 3.0
+    rewritten = data["messages"][0]["tool_calls"][0]["function"]["arguments"]
+    assert json.loads(rewritten) == json.loads(arguments.replace(AWS_KEY, "[REDACTED]"))
+
+
 def _recorded(request_data: dict) -> dict:
     entries = request_data["metadata"]["standard_logging_guardrail_information"]
     assert len(entries) == 1

tests/unit/proxy/guardrails/test_content_utils.py

@@ -985,3 +985,36 @@ def test_tool_json_graph_unsupported_values_fail_without_echoing_values(value: o
     with pytest.raises(ValueError, match=r"^Invalid tool JSON payload$") as caught:
         _exercise_tool_json(value, operation)
     assert "secret" not in str(caught.value)
+
+
+def test_nested_tool_result_chain_leaves_later_messages_inspected() -> None:
+    chain = reduce(
+        lambda content, _: [{"type": "tool_result", "tool_use_id": "tool", "content": content}],
+        range(1500),
+        [{"type": "text", "text": "inner"}],
+    )
+    data = {"messages": [{"role": "assistant", "content": chain}, {"role": "user", "content": "secret"}]}
+    assert "secret" in tuple(iter_message_text(data))
+    walk_user_text(data, lambda text: text.replace("secret", "[MASKED]"))
+    assert data["messages"][1]["content"] == "[MASKED]"
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    ('{"token":"secret","limit":1e400}', '{"token":"secret","limit":-Infinity}', '{"token":"secret","ratio":NaN}'),
+)
+def test_tool_arguments_with_non_finite_numbers_keep_their_original_bytes(arguments: str) -> None:
+    data = {"messages": [{"tool_calls": [{"function": {"arguments": arguments}}]}]}
+    assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1
+    assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == arguments.replace("secret", "[MASKED]")
+
+
+@pytest.mark.parametrize(
+    "data",
+    (
+        {"messages": [{"content": [{"type": "tool_use", "input": {"a": "first", "b": [7, {"c": "second"}, "third"]}}]}]},
+        {"messages": [{"tool_calls": [{"function": {"arguments": '{"a":"first","b":[7,{"c":"second"},"third"]}'}}]}]},
+    ),
+)
+def test_native_reader_yields_one_fragment_per_tool_payload(data) -> None:
+    assert tuple(iter_message_text(data)) == ("first\nsecond\nthird",)

tests/unit/proxy/guardrails/test_guardrail_coverage.py

@@ -15,6 +15,7 @@ inspection by the relevant guardrail hook:
 import json
 import sys
 from collections.abc import Mapping
+from functools import reduce
 from types import SimpleNamespace
 from typing import Any, Dict, Final, NamedTuple
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -501,6 +502,31 @@ def test_banned_keywords_blocks_responses_api_input(monkeypatch):
         asyncio.run(_run())
 
 
+def test_banned_keywords_still_checks_messages_after_a_deep_tool_result_chain(monkeypatch):
+    monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False)
+    from enterprise.enterprise_hooks.banned_keywords import ENTERPRISE_BannedKeywords
+
+    chain: Final = reduce(
+        lambda content, _: [{"type": "tool_result", "tool_use_id": "tool", "content": content}],
+        range(1500),
+        [{"type": "text", "text": "inner"}],
+    )
+    data: Final = {"messages": [{"role": "assistant", "content": chain}, {"role": "user", "content": "forbidden"}]}
+
+    import asyncio
+
+    with pytest.raises(HTTPException) as exc:
+        asyncio.run(
+            ENTERPRISE_BannedKeywords().async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"),
+                cache=DualCache(),
+                data=data,
+                call_type="acompletion",
+            )
+        )
+    assert "forbidden" in str(exc.value.detail).lower()
+
+
 @pytest.mark.parametrize("call_type", ["completion", "acompletion", "aresponses"])
 def test_banned_keywords_fires_on_text_content_call_types(monkeypatch, call_type):
     """Locks the call-type gate to the runtime ``route_type`` values the
@@ -961,7 +987,7 @@ class _GoogleRequest(NamedTuple):
                     },
                 ]
             },
-            "ordinary controlinput first input blocked fragment",
+            "ordinary controlinput first \ninput blocked fragment",
             id="native-tool-input",
         ),
         pytest.param(
@@ -1011,7 +1037,7 @@ class _GoogleRequest(NamedTuple):
                     },
                 ]
             },
-            "ordinary controlargument first argument blocked fragment",
+            "ordinary controlargument first \nargument blocked fragment",
             id="chat-tool-arguments",
         ),
     ),