From 20e24bfd1a76bfb182357a2e7e608aba32172bce Mon Sep 17 00:00:00 2001 From: Kunal Kumar Date: Mon, 19 Oct 2026 09:40:41 +0530 Subject: [PATCH] fix(guardrails): block a Highflame redaction on pass-through endpoints The pass-through handler forwards the original payload and drops the texts a guardrail returns, so a Shield modify there would let the unredacted content through. Treat pass-through text as not rewritable, the same as a streamed response, so a redaction blocks instead. --- .../guardrail_hooks/highflame/highflame.py | 20 +++++-- .../guardrail_hooks/test_highflame.py | 60 ++++++++++++++++++- 2 files changed, 74 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/highflame/highflame.py b/litellm/proxy/guardrails/guardrail_hooks/highflame/highflame.py index a0dfb719ed6..7d58f63494a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/highflame/highflame.py +++ b/litellm/proxy/guardrails/guardrail_hooks/highflame/highflame.py @@ -23,7 +23,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -177,7 +177,11 @@ class HighflameGuardrail(CustomGuardrail): logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: evaluations: Final = _plan_evaluations( - inputs, request_data, input_type, streamed=request_data.get("stream") is True + inputs, + request_data, + input_type, + streamed=request_data.get("stream") is True, + pass_through=_is_pass_through(logging_obj), ) if not evaluations: return inputs @@ -326,17 +330,23 @@ class HighflameGuardrail(CustomGuardrail): return HighflameGuardrailConfigModel +def _is_pass_through(logging_obj: object) -> bool: + call_type: Final[object] = getattr(logging_obj, "call_type", None) + return call_type == CallTypes.pass_through.value + + def _plan_evaluations( inputs: GenericGuardrailAPIInputs, request_data: Mapping[str, object], input_type: Literal["request", "response"], *, streamed: bool, + pass_through: bool, ) -> tuple[_Evaluation, ...]: is_request: Final = input_type == "request" - # LiteLLM's streaming path drops text rewrites and replays the original chunks, so a redaction - # of a streamed response cannot be applied and is refused instead - rewritable: Final = is_request or not streamed + # LiteLLM's streaming path replays the original chunks and the pass-through handler forwards the + # original payload, so a redaction cannot be applied there and is refused instead + rewritable: Final = not pass_through and (is_request or not streamed) texts: Final = tuple( _Evaluation( content=text, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_highflame.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_highflame.py index f2318581c29..17f760ec28c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_highflame.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_highflame.py @@ -7,6 +7,7 @@ path under test — notably that it raises on every non-2xx status, which is how (401) and defer (425) verdicts arrive. """ +import datetime import json from collections.abc import Callable, Mapping, Sequence @@ -18,6 +19,8 @@ import pytest import litellm from litellm.caching.caching import DualCache from litellm.exceptions import GuardrailRaisedException +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.highflame.highflame import ( @@ -31,7 +34,7 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai from litellm.proxy.guardrails.guardrail_endpoints import get_guardrail_ui_settings from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import ChatCompletionMessageToolCall, Function +from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, Function from litellm.types.utils import ( Choices, Delta, @@ -631,6 +634,61 @@ class TestThroughTheProxyHooks: assert fake.guard_requests[0]["action"] == "process_response" +def _logging_obj(call_type: str) -> LiteLLMLogging: + return LiteLLMLogging( + model="gpt-4o", + messages=[], + stream=False, + call_type=call_type, + start_time=datetime.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + + +class TestPassThrough: + """The pass-through handler forwards the original payload, so a redaction there must block.""" + + @pytest.mark.asyncio + async def test_redaction_of_a_pass_through_request_blocks(self): + fake = FakeHighflame({"decision": "modify", "redacted_content": "ssn [SSN]"}) + + with pytest.raises(GuardrailRaisedException) as exc: + await PassThroughEndpointHandler().process_input_messages( + data={"prompt": "ssn 123-45-6789"}, + guardrail_to_apply=_guardrail(fake), + litellm_logging_obj=_logging_obj(CallTypes.pass_through.value), + ) + + assert exc.value.blocked_content is True + + @pytest.mark.asyncio + async def test_redaction_of_a_pass_through_response_blocks(self): + fake = FakeHighflame({"decision": "modify", "redacted_content": "ssn [SSN]"}) + + with pytest.raises(GuardrailRaisedException) as exc: + await PassThroughEndpointHandler().process_output_response( + response={"output": "ssn 123-45-6789"}, + guardrail_to_apply=_guardrail(fake), + litellm_logging_obj=_logging_obj(CallTypes.pass_through.value), + ) + + assert exc.value.blocked_content is True + + @pytest.mark.asyncio + async def test_redaction_still_applies_outside_pass_through(self): + fake = FakeHighflame({"decision": "modify", "redacted_content": "ssn [SSN]"}) + + result = await _guardrail(fake).apply_guardrail( + inputs={"texts": ["ssn 123-45-6789"]}, + request_data={}, + input_type="request", + logging_obj=_logging_obj(CallTypes.completion.value), + ) + + assert result["texts"] == ["ssn [SSN]"] + + class TestStreaming: @pytest.mark.asyncio async def test_redaction_of_a_streamed_response_blocks_since_the_stream_cannot_be_rewritten(self):