From 030eca6e17a6bfa8b23b3dc0d0ccc16bcd1393a6 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 06:30:17 -0700 Subject: [PATCH] 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_callbacks/secret_detection.py | 18 +++- litellm/proxy/guardrails/_content_utils.py | 72 +++++++++++---- .../test_banned_keywords_tool_traffic.py | 87 +++++++++++++++++++ .../test_hide_secrets_tool_traffic.py | 40 +++++++++ .../test_secret_detection.py | 12 +++ .../proxy/guardrails/test_content_utils.py | 33 +++++++ .../guardrails/test_guardrail_coverage.py | 30 ++++++- 7 files changed, 267 insertions(+), 25 deletions(-) create mode 100644 tests/integration/security/test_banned_keywords_tool_traffic.py diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 3f7d05bfcc1..183d2342e79 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/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)) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 35ebe598769..75beef19691 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/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 diff --git a/tests/integration/security/test_banned_keywords_tool_traffic.py b/tests/integration/security/test_banned_keywords_tool_traffic.py new file mode 100644 index 00000000000..e611ef2d896 --- /dev/null +++ b/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() == () diff --git a/tests/integration/security/test_hide_secrets_tool_traffic.py b/tests/integration/security/test_hide_secrets_tool_traffic.py index e86bd213578..e35c82eba76 100644 --- a/tests/integration/security/test_hide_secrets_tool_traffic.py +++ b/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) diff --git a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py index 9f1f8f6e375..8e0b47e67fd 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/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 diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index 8418e825869..87dff2d7feb 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/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",) diff --git a/tests/unit/proxy/guardrails/test_guardrail_coverage.py b/tests/unit/proxy/guardrails/test_guardrail_coverage.py index b28486ea4c1..c02d9f24a54 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_coverage.py +++ b/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", ), ),