Skip to content
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,16 @@ def _build_request_headers(self) -> dict:
headers.update(self.headers)
return headers

@staticmethod
def _echo_sent_tool_calls(
guardrailed: GenericGuardrailAPIInputs, *, sent: GenericGuardrailAPIInputs
) -> GenericGuardrailAPIInputs:
"""Keep the tool calls the guardrail did not return, so the handlers do not report them missing."""
sent_tool_calls: Final = sent.get("tool_calls")
if not sent_tool_calls or "tool_calls" in guardrailed:
return guardrailed
return {**guardrailed, "tool_calls": sent_tool_calls}

def _build_guardrail_return_inputs(
self,
*,
Expand Down Expand Up @@ -365,6 +375,8 @@ def _build_guardrail_return_inputs(
)
if rows_to_write_back is not None:
return_inputs["structured_messages"] = list(rows_to_write_back)
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved
if guardrail_response.tool_calls:
return_inputs["tool_calls"] = list(guardrail_response.tool_calls)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Medium: Multi-choice streams bypass text redaction

In a post-call policy pipeline using this guardrail, a caller can request stream: true, n: 2 and receive original, unredacted assistant text when the guardrail also rewrites a tool call: _write_ended_stream_tool_call_rewrites rejects multi-choice streams, and PipelineExecutor._run_streaming_step catches that rejection and restores all original chunks. The PR changes this integration from omitting tool-call rewrites to forwarding them, introducing a rollback path that also discards text redactions.

Before forwarding these rewrites on this path, implement per-choice tool-call delivery or make unsupported rewrites fail closed rather than releasing the original stream.

@ninadphalak ninadphalak Oct 6, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Leaving this: outside pipelines it fails closed, and inside them the executor rolls back every guardrail's undeliverable rewrite. Happy to change the executor separately

if guardrail_response.stream_holdback_chars is not None:
return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
return return_inputs
Expand Down Expand Up @@ -499,13 +511,16 @@ async def apply_guardrail(
blocked_content=True,
)

return self._build_guardrail_return_inputs(
texts=texts,
images=images,
tools=tools,
structured_messages=structured_messages,
shown_messages=guardrail_request.structured_messages,
guardrail_response=guardrail_response,
return self._echo_sent_tool_calls(
self._build_guardrail_return_inputs(
texts=texts,
images=images,
tools=tools,
structured_messages=structured_messages,
shown_messages=guardrail_request.structured_messages,
guardrail_response=guardrail_response,
),
sent=inputs,
)

except GuardrailRaisedException:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from pydantic import ConfigDict, Field
from typing_extensions import TypedDict

from litellm._logging import verbose_proxy_logger
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import (
AllMessageValues,
Expand Down Expand Up @@ -168,6 +169,43 @@ def structured_messages_from_response(value: object) -> Sequence[AllMessageValue
return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get


def _json_object(value: object) -> Mapping[str, object] | None:
"""`value` as a JSON object, or None; the keys of a JSON object are strings."""
if not isinstance(value, Mapping):
return None
return cast("Mapping[str, object]", value) # cast-ok: isinstance leaves a Mapping's item types unknown


def _json_field(container: object, key: str) -> object:
fields: Final = _json_object(container)
return fields.get(key) if fields is not None else None


def _is_usable_tool_call(row: object) -> bool:
"""Whether a returned tool call carries what the chat handlers index: a function name and string arguments."""
function: Final = _json_field(row, "function")
return isinstance(_json_field(function, "name"), str) and isinstance(_json_field(function, "arguments"), str)


def _json_array(value: object) -> Sequence[object] | None:
if not isinstance(value, list):
return None
return cast("Sequence[object]", value) # cast-ok: isinstance leaves a list's item type unknown


def tool_calls_from_response(value: object) -> Sequence[ChatCompletionToolCallChunk] | None:
"""The tool calls a guardrail handed back, or None when absent or when any row is unusable."""
rows: Final = _json_array(value)
if rows is None:
return None
if not all(_is_usable_tool_call(row) for row in rows):
verbose_proxy_logger.warning(
"Generic Guardrail API returned tool calls missing a function name or arguments; keeping the ones already sent",
)
return None
return cast("Sequence[ChatCompletionToolCallChunk]", rows) # cast-ok: every row checked above


class GenericGuardrailAPIResponse:
"""Response model for the Generic Guardrail API"""

Expand All @@ -178,6 +216,7 @@ class GenericGuardrailAPIResponse:
action: str
blocked_reason: str | None
stream_holdback_chars: list[int] | None
tool_calls: Sequence[ChatCompletionToolCallChunk] | None

def __init__(
self,
Expand All @@ -188,6 +227,7 @@ def __init__(
tools: list[GuardrailToolParam] | None = None,
stream_holdback_chars: list[int] | None = None,
structured_messages: Sequence[AllMessageValues] | None = None,
tool_calls: Sequence[ChatCompletionToolCallChunk] | None = None,
) -> None:
self.action = action
self.blocked_reason = blocked_reason
Expand All @@ -199,9 +239,11 @@ def __init__(
# framework must withhold from streaming emission until the next
# processing round (word-boundary safety for text transformations).
self.stream_holdback_chars = stream_holdback_chars
self.tool_calls = tool_calls

@classmethod
def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse":
fields: Final = _json_object(data) or {}
raw_holdback: Final = data.get("stream_holdback_chars")
stream_holdback_chars: Final = (
[coerce_stream_holdback_value(value) for value in raw_holdback] if isinstance(raw_holdback, list) else None
Expand All @@ -213,5 +255,6 @@ def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse":
images=data.get("images"),
tools=data.get("tools"),
stream_holdback_chars=stream_holdback_chars,
structured_messages=structured_messages_from_response(data.get("structured_messages")),
structured_messages=structured_messages_from_response(fields.get("structured_messages")),
tool_calls=tool_calls_from_response(fields.get("tool_calls")),
)
Original file line number Diff line number Diff line change
Expand Up @@ -1519,6 +1519,35 @@ def test_from_dict_malformed_holdback_degrades_to_zero(self):

assert response.stream_holdback_chars == [3, 0, 0, 0]

def test_from_dict_parses_tool_calls(self):
"""``tool_calls`` goes out on the request, so the response has to carry it back."""
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)

tool_calls = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]

response = GenericGuardrailAPIResponse.from_dict(
{"action": "GUARDRAIL_INTERVENED", "texts": ["Done"], "tool_calls": tool_calls}
)

assert response.tool_calls == tool_calls

def test_from_dict_tool_calls_absent_is_none(self):
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)

response = GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["hi"]})

assert response.tool_calls is None

@pytest.mark.asyncio
async def test_apply_guardrail_flows_holdback_back_to_inputs(self, generic_guardrail):
"""A GUARDRAIL_INTERVENED response with stream_holdback_chars is surfaced on
Expand Down Expand Up @@ -2015,6 +2044,158 @@ async def test_builtin_tools_without_function_block_do_not_crash(

assert forwarded_tools == tools

@pytest.mark.asyncio
async def test_guardrail_returned_tool_calls_flow_back_to_inputs(self, generic_guardrail):
"""A guardrail that rewrites tool arguments has to be able to return them.

``tool_calls`` is sent on the request, so the matching response field is what closes
the loop. Without it the calls were dropped on the floor and the caller silently
kept the values it sent in -- for a rewriting guardrail (reversible redaction, for
instance) the rewritten arguments never reached the client.
"""
guardrailed_tool_calls = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]
mock_response = MagicMock()
mock_response.json.return_value = {
"action": "GUARDRAIL_INTERVENED",
"texts": ["Done"],
"tool_calls": guardrailed_tool_calls,
}
mock_response.raise_for_status = MagicMock()

with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response):
result = await generic_guardrail.apply_guardrail(
inputs={"texts": ["Done"]},
request_data={},
input_type="response",
)

assert result["tool_calls"] == guardrailed_tool_calls

def test_guardrail_tool_calls_take_precedence_over_the_sent_ones(self, generic_guardrail):
"""The returned calls win over the ones that went out.

Built through ``_build_guardrail_return_inputs`` rather than ``apply_guardrail`` so
the assertion is about this mapping, not about request-model coercion of the calls
the caller supplied.
"""
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)

sent = [{"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": "{}"}}]
returned = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]

result = generic_guardrail._build_guardrail_return_inputs(
texts=["Done"],
images=None,
tools=None,
structured_messages=None,
shown_messages=None,
guardrail_response=GenericGuardrailAPIResponse.from_dict(
{"action": "GUARDRAIL_INTERVENED", "texts": ["Done"], "tool_calls": returned}
),
)

assert result["tool_calls"] == returned
assert result["tool_calls"] != sent

@pytest.mark.asyncio
async def test_untouched_tool_calls_are_echoed_not_dropped(self, generic_guardrail):
"""A guardrail that ignores tool calls must not blank them out.

The handlers compare what comes back against what they sent and fall back to their
own copy when the lengths differ, and the Responses handler logs a warning when no
list comes back. Echoing the unchanged list keeps both quiet.
"""
sent = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]
mock_response = MagicMock()
mock_response.json.return_value = {"action": "NONE", "texts": ["Done"]}
mock_response.raise_for_status = MagicMock()

with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response):
result = await generic_guardrail.apply_guardrail(
inputs={"texts": ["Done"], "tool_calls": sent},
request_data={},
input_type="response",
)

assert result["tool_calls"] is sent

@pytest.mark.asyncio
async def test_malformed_returned_tool_calls_are_discarded(self, generic_guardrail):
"""A same-length list with no function block must not reach chat handling.

The chat handlers index ``function.name`` and ``function.arguments`` on whatever the
guardrail returns, so forwarding an unusable list raises there instead of returning a
response. The calls that were sent are kept instead.
"""
sent = [{"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": "{}"}}]
mock_response = MagicMock()
mock_response.json.return_value = {
"action": "GUARDRAIL_INTERVENED",
"texts": ["Done"],
"tool_calls": [{"id": "call_1", "type": "function"}],
}
mock_response.raise_for_status = MagicMock()

with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response):
result = await generic_guardrail.apply_guardrail(
inputs={"texts": ["Done"], "tool_calls": sent},
request_data={},
input_type="response",
)

assert result["tool_calls"] is sent

def test_from_dict_rejects_tool_calls_without_string_arguments(self):
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)

response = GenericGuardrailAPIResponse.from_dict(
{
"action": "GUARDRAIL_INTERVENED",
"tool_calls": [{"id": "call_1", "function": {"name": "send_email", "arguments": {"to": "x"}}}],
}
)

assert response.tool_calls is None

def test_no_tool_calls_on_either_side_stay_absent(self, generic_guardrail):
"""An empty list is not the same as no list: the key must stay out of the result."""
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)

result = generic_guardrail._build_guardrail_return_inputs(
texts=["Done"],
images=None,
tools=None,
structured_messages=None,
shown_messages=None,
guardrail_response=GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["Done"]}),
)

assert "tool_calls" not in result


class TestFailOnError:
"""Test fail_on_error: complete fail-open on any guardrail error"""
Expand Down
Loading