diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6cdf49818e6..8861a9e3bbc 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -34,6 +34,12 @@ def is_text_content_call_type(call_type: str) -> bool: TEXT_PART_TYPES: FrozenSet[str] = frozenset({"text", "input_text", "output_text"}) +# Responses-API item types whose ``output`` field carries user/tool text +# that guardrails should inspect. ``function_call_output`` is the +# built-in shape; ``custom_tool_call_output`` is the custom-tool +# counterpart (see ``ChatCompletionCustomToolCallOutput``). +_OUTPUT_ITEM_TYPES: frozenset[str] = frozenset({"function_call_output", "custom_tool_call_output"}) + def _iter_text_parts_in_content(content: Any) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or @@ -72,7 +78,7 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]: messages.append({"role": item.get("role") or "user", "content": [item]}) elif "content" in item: messages.append({"role": item.get("role") or "user", "content": item["content"]}) - elif item.get("type") == "function_call_output" and "output" in item: + elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: messages.append({"role": item.get("role") or "tool", "content": item["output"]}) return messages @@ -157,7 +163,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: input_value[idx] = {**item, "text": visit(item["text"])} elif "content" in item: item["content"] = _rewrite_content(item["content"]) - elif item.get("type") == "function_call_output" and "output" in item: + elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: item["output"] = _rewrite_content(item["output"]) return visited diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 34d92505359..3dfb98c12ea 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -524,3 +524,42 @@ def test_apply_redacted_messages_back_skips_input_when_not_string(): data = {"input": [{"type": "text", "text": "leak"}]} apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) assert data["input"] == [{"type": "text", "text": "leak"}] + + +# ------------------------------------------------------------------- +# LIT-4302: custom_tool_call_output walking +# ------------------------------------------------------------------- + +def test_iter_message_text_walks_custom_tool_call_output(): + """custom_tool_call_output items should yield their output text.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "tool-secret"}, + ] + } + from litellm.proxy.guardrails._content_utils import iter_message_text + texts = list(iter_message_text(data)) + assert "tool-secret" in texts + + +def test_walk_user_text_redacts_custom_tool_call_output(): + """walk_user_text should rewrite text inside custom_tool_call_output.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "PII-data"}, + ] + } + count = walk_user_text(data, lambda t: t.replace("PII-data", "[MASKED]")) + assert count >= 1 + assert data["input"][0]["output"] == "[MASKED]" + + +def test_build_inspection_messages_custom_tool_call_output(): + """build_inspection_messages should include custom_tool_call_output text.""" + data = { + "input": [ + {"type": "custom_tool_call_output", "output": "custom-tool-leak"}, + ] + } + msgs = build_inspection_messages(data) + assert any("custom-tool-leak" in m["content"] for m in msgs)