Skip to content

feat(guardrails): streaming text transformation in generic_guardrail_api - #32915

Closed
yucheng-berri wants to merge 12 commits into
litellm_internal_stagingfrom
litellm_guardrail_streaming_text_transform
Closed

feat(guardrails): streaming text transformation in generic_guardrail_api#32915
yucheng-berri wants to merge 12 commits into
litellm_internal_stagingfrom
litellm_guardrail_streaming_text_transform

Conversation

@yucheng-berri

@yucheng-berri yucheng-berri commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Title

feat(guardrails): streaming text transformation in generic_guardrail_api

Relevant issues

Linear ticket

Pre-Submission checklist

  • I have added meaningful tests
  • My PR passes all CI/CD checks
  • My PR's scope is as isolated as possible
  • I have received a Greptile Confidence Score of at least 4/5

Type

New Feature

Changes

Adds an opt-in streaming_transform_mode: incremental_diff config on generic_guardrail_api so an HTTP guardrail can rewrite streamed response text (PII masking, pseudonym reversal, redaction) instead of only blocking on it.

The default block_only preserves historical behavior: the guardrail can terminate a stream via a BLOCKED action, but GUARDRAIL_INTERVENED text rewrites were silently dropped on streams. incremental_diff withholds the raw upstream chunks, runs the guardrail every N chunks over the accumulated text, and emits the guardrail's rewritten text as synthetic deltas by diffing the mutated accumulated text against what has already been streamed.

Companion stream_holdback_chars on the guardrail response contract lets a guardrail server request the framework withhold N trailing chars per choice (word-boundary safety across rounds, e.g. pause on "Thomas W" until the next chunk arrives so the guardrail can decide whether to redact "Thomas Weber").

v1 scope:

  • Only the OpenAI chat completions streaming path supports incremental_diff; other routes silently fall back to block_only.
  • Only delta.content text is transformed. delta.tool_calls chunks pass through with content stripped and are inspected for a block decision at end of stream (matching block_only), but never text-rewritten.
  • A rewrite that is not a forward extension of what has already been streamed (would retract emitted bytes) is failed closed with HTTP 400 stream_transform_underflow; a guardrail that needs to rewrite recent output must withhold it first via stream_holdback_chars.

Included:

  • StreamTransformSink on BaseTranslation.process_output_streaming_response for handlers that support the streaming text-diff protocol
  • Diff engine in the OpenAI chat handler that accumulates by StreamingChoices.index (not enumerate position) so n>1 streams do not collapse
  • coerce_stream_holdback_value for defensive parsing of malformed guardrail responses
  • End-of-stream block inspection over the full assembled response including tool calls (so incremental_diff matches block_only's guardrail coverage)
  • Regression tests for the mixed-content chunk case (delta.content + delta.tool_calls in the same item, and n>1 with text on one choice and a tool call on another)

Credit

Adopted from #32084 by @schneidermr (PalenaAI). Mirrored onto a litellm_ branch so CircleCI and the internal lint workflow run against the current HEAD. Original PR: #32084

Follow-ups (not in this PR)

  • Docs page at docs.litellm.ai/docs/adding_provider/generic_guardrail_api needs a "Streaming text transformations" section covering the diff protocol, stream_holdback_chars, the fail-closed rule, and the v1 scope. Source lives outside this monorepo.
  • Add build_block_sse_chunks override to the OpenAI chat translation handler so a mid-stream stream_transform_underflow HTTPException surfaces as a well-formed SSE error frame + data: [DONE] instead of relying on the outer async_data_generator fallback. The current behavior mirrors the pre-existing pattern used by 9 other streaming guardrails on OpenAI-format streams, so this is not a regression, but under incremental_diff the fail-closed path fires far more routinely than under block_only and would benefit from a cleaner client signal.

Note

High Risk
Large changes to proxy streaming guardrail behavior on the hot path for chat completions, including fail-closed suppression and tool-call/mixed-chunk edge cases where incorrect ordering could leak unredacted text.

Overview
Adds opt-in streaming_transform_mode: incremental_diff on generic_guardrail_api so HTTP guardrails can rewrite streamed chat text (masking, redaction) instead of only blocking. Default block_only is unchanged: streams still pass raw deltas and only BLOCK stops the stream; GUARDRAIL_INTERVENED text edits were previously dropped on streams.

With incremental_diff, the unified streaming hook withholds raw chunks, runs the guardrail on accumulated string delta.content (sampled or end-of-stream), and emits synthetic deltas by diffing guardrailed text against what was already sent. Guardrail responses can include stream_holdback_chars per choice for word-boundary withholding; rewrites that would retract already-sent bytes fail closed with stream_transform_underflow.

Implementation spans StreamTransformSink on BaseTranslation.process_output_streaming_response, a non-mutating transform path in the OpenAI chat handler (accumulation by StreamingChoices.index), wiring in UnifiedLLMGuardrails (tool-call passthrough with content stripped, end-of-stream block inspection, mixed text+tool SSE ordering), and typed/config plumbing plus extensive regression tests.

Reviewed by Cursor Bugbot for commit c47afcd. Bugbot is set up for automated code reviews on this repo. Configure here.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@CLAassistant

CLAassistant commented Jul 11, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@codecov

codecov Bot commented Jul 11, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 89.84962% with 27 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...drail_hooks/unified_guardrail/unified_guardrail.py 89.44% 19 Missing ⚠️
.../llms/openai/chat/guardrail_translation/handler.py 87.69% 8 Missing ⚠️

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds opt-in streaming_transform_mode: incremental_diff to generic_guardrail_api, enabling HTTP guardrails to rewrite streamed text (PII masking, redaction) instead of only blocking. The default block_only mode is unchanged.

  • Core engine (unified_guardrail.py): _run_incremental_transform_stream drives the new path with pre-tool-call text flushing, deferred finish_reason for mixed content+tool chunks, n>1 index-sorted guardrail inputs, and a terminator chunk when text is fully suppressed.
  • Handler (handler.py): _process_streaming_transform accumulates raw text by StreamingChoices.index, runs the guardrail without mutating responses_so_far, and reports results on a StreamTransformSink out-parameter.
  • Tests: 783 lines of new mock-only tests covering holdback, underflow fail-closed, per-choice finish_reason, tool-call passthrough, mixed-content chunks, and raw-accumulator integrity.

Confidence Score: 4/5

The incremental_diff path is opt-in and the default block_only behavior is unchanged; the new findings are edge cases that do not affect well-formed streaming responses.

The implementation addresses prior review concerns but two edge-case bugs remain: raw-accumulator corruption for malformed streams without finish_reason, and missing saw_text_content tracking for mixed-chunk-only streams. Open issues from prior threads (double guardrail call, n>1 finish_reason drop) are also unresolved.

unified_guardrail.py - the _inspect_full_response_for_block interaction and saw_text_content tracking for mixed chunks deserve a second look.

Important Files Changed

Filename Overview
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py Large new incremental_diff stream implementation. Two new edge-case findings: raw-accumulator corruption via _inspect_full_response_for_block on malformed streams, and saw_text_content never set for mixed-only streams.
litellm/llms/openai/chat/guardrail_translation/handler.py Extracts _process_streaming_block_only and adds _process_streaming_transform. Core logic is correct; has_stream_ended check still inspects only choices[0] (pre-existing).
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py Adds streaming_transform_mode config param and propagates stream_holdback_chars from guardrail response.
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py Adds streaming_transform_mode optional param, stream_holdback_chars response field, coerce_stream_holdback_value helper. All new fields are additive and backward-compatible.
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py 783 lines of new streaming-transform tests covering core scenarios. All tests use sampling_rate=1, which masks behaviors that differ at the default sampling_rate=5.

Reviews (15): Last reviewed commit: "fix(guardrails): four correctness fixes ..." | Re-trigger Greptile

@greptile-apps

greptile-apps Bot commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds an opt-in streaming_transform_mode: incremental_diff to generic_guardrail_api that withholds raw upstream chunks, runs the guardrail every N chunks over the accumulated text, and emits only the guardrail-rewritten text as synthetic SSE deltas (diffed against what has already been sent). The default block_only mode is unchanged.

  • StreamTransformSink is introduced as an out-parameter on BaseTranslation.process_output_streaming_response; the OpenAI chat handler populates it while leaving responses_so_far untouched so it stays a correct raw accumulator across rounds.
  • UnifiedLLMGuardrails._run_incremental_transform_stream orchestrates the new path: it accumulates chunks, runs sampled transform rounds, passes tool-call chunks through raw (content stripped), performs an end-of-stream block inspection over tool calls, and flushes any remaining held-back text on the final round.
  • coerce_stream_holdback_value and the stream_holdback_chars response field let a guardrail server request word-boundary holdback per choice; malformed values degrade defensively to 0.

Confidence Score: 3/5

The incremental_diff path is safe for text-only and tool-call-only streams, but a protocol violation in the mixed content+tool_calls edge case means the guardrail transformed text is silently lost for SSE-compliant clients when finish_reason rides on the tool passthrough chunk before the synthetic text is emitted.

The incremental_diff path works correctly for the common cases but contains a protocol violation in the mixed content+tool_calls edge case that would cause silent data loss for conformant SSE clients. The existing test exercises this code path but only checks the security property (no raw text leak), not the ordering violation.

unified_guardrail.py (lines 580-587, tool-call passthrough branch) and test_unified_guardrail.py (mixed-chunk test should assert output ordering relative to finish_reason).

Important Files Changed

Filename Overview
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py Core orchestration for incremental_diff streaming — protocol violation in mixed content+tool_calls edge case; also makes a redundant guardrail API call for tool-call-containing streams.
litellm/llms/openai/chat/guardrail_translation/handler.py Cleanly refactored into block-only and transform paths; accumulates by StreamingChoices.index correctly; does not mutate responses_so_far.
litellm/llms/base_llm/guardrail_translation/base_translation.py Adds StreamTransformSink and stream_transform_sink parameter; backward-compatible, no issues.
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py Adds streaming_transform_mode, stream_holdback_chars, and coerce_stream_holdback_value; defensive coercion degrades invalid values to 0.
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py Adds streaming_transform_mode param and surfaces stream_holdback_chars through apply_guardrail return; no issues.
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/init.py One-line change forwarding streaming_transform_mode to initialize_guardrail; no issues.
litellm/types/utils.py Adds stream_holdback_chars to GenericGuardrailAPIInputs TypedDict with total=False; backward-compatible.
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py Well-structured mock-only unit tests for response parsing, holdback coercion, and config defaults.
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py Comprehensive transform-path tests; mixed-chunk test verifies no raw leak but does not assert output ordering relative to finish_reason.
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/example_config.yaml Documents new streaming options clearly; no issues.

Reviews (2): Last reviewed commit: "docs(generic_guardrail_api): document st..." | Re-trigger Greptile

Comment thread litellm/llms/openai/chat/guardrail_translation/handler.py
@veria-ai

veria-ai Bot commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

PR overview

This PR adds streaming text transformation support to the generic guardrail API, updating the unified guardrail hook path to process and emit transformed text during streamed responses.

One issue remains open around how transformed streaming output is released. With the current behavior, text can be sent to the client before later guardrail processing determines it should be redacted, allowing sensitive content split across chunks to leak during streaming. One earlier issue has already been addressed, but this remaining streaming-buffering problem leaves a meaningful confidentiality risk until fixed.

Open issues (1)

Fixed/addressed: 1 · PR risk: 7/10

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri
yucheng-berri force-pushed the litellm_guardrail_streaming_text_transform branch from 93acde1 to 42227d0 Compare July 11, 2026 17:51
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

Comment thread litellm/llms/openai/chat/guardrail_translation/handler.py
mappings=endpoint_guardrail_translation_mappings,
):
yield transformed_item
return

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.

Buffer flag ignored on transform

Low Severity

When streaming_transform_mode is incremental_diff and the route resolves, the hook enters _run_incremental_transform_stream and returns before the legacy loop. streaming_buffer_until_moderated is never applied there, so configured end-of-stream buffering is silently ignored and transformed deltas can reach the client before moderation would have completed on the block-only path.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 42227d0. Configure here.

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.

streaming_buffer_until_moderated is an all-or-nothing buffer flag: it withholds every chunk until end-of-stream moderation passes, then releases the (clean) raw chunks or emits the block message. That model is inherently incompatible with incremental_diff, which is a stream-and-transform mode — buffering everything until the guardrail says OK would defeat the whole purpose of emitting synthetic deltas as text arrives.

The two config knobs already model this decision:

  • streaming_buffer_until_moderated: for block_only, delay every chunk until end of stream.
  • streaming_transform_mode: incremental_diff: stream transformed text as it arrives, no buffering.

Under block_only (the default) both knobs interact as documented. Under incremental_diff the transform flag takes precedence, which the _resolve_transform_call_type guard at line 306-320 makes explicit — we enter _run_incremental_transform_stream and return before the buffer logic. Operators who want moderation-first behavior should stick with block_only; operators who want in-flight rewrites use incremental_diff and accept that the buffer flag is a no-op.

A doc note on this interaction is worth adding — will include it in the follow-up docs.litellm.ai section already tracked in the PR body. Not fixing the code here.

@codspeed-hq

codspeed-hq Bot commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_guardrail_streaming_text_transform (527f110) with litellm_internal_staging (aa9dcb4)

Open in CodSpeed

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

Live e2e — 6/6 scenarios green

Real proxy at 127.0.0.1:4141 running commit 6b05b08, real Bedrock LLM (claude-haiku-4-5), real HTTP guardrail server at 127.0.0.1:8765 speaking the generic guardrail API contract. Zero mocks.

Scenarios

# What it exercises Outcome
A action: NONE, non-triggering prompt streams unchanged PASS — 2 chunks, hello world, finish=stop
B SSN pattern in response, guardrail wants to shorten → stream_transform_underflow fails closed PASS — client sees ChunkedEncodingError (framework re-raise), no full SSN reaches client
C Alice → PERSON_A rewrite with stream_holdback_chars=[3] for word-boundary safety PASS — 3 chunks, 'PERSON_A went to the market and bought apples.', guardrail called 3×
D Guardrail returns texts: [] mid-stream — fails closed (underflow), only partial-leak-before-detection survives PASS — 3 chars leaked before guardrail identified trigger, no full response emitted
E action: BLOCKED — stream terminates via error frame PASS — framework's async_data_generator emitted well-formed JSON error frame {"message": "Blocked: BLOCK_TOKEN present", "code": "400"}
F Non-triggering prompt with sampling every chunk — transform runs, emits identity PASS — 6 chunks, 'The sky looks cloudy today.', guardrail called 6× per sampled round

What this proves

  • Streaming text rewrites reach the client end-to-end (C).
  • stream_transform_underflow fires when the guardrail tries to retract already-sent bytes (B, D). Framework fails closed; the SSE fail-open manifestation on the OpenAI chat path is the pre-existing framework-wide pattern, follow-up tracked to add build_block_sse_chunks override on the OpenAI handler.
  • BLOCKED action still surfaces cleanly as an SSE error frame via async_data_generator (E).
  • Default (block_only) is unchanged; scenario A/F still stream normally with the guardrail running post-chunk.
  • Guardrail is invoked per sampled round (streaming_sampling_rate: 1 in test config), text stays consistent (raw accumulator preserved), no double invocations at end of stream for tool-call-free streams.

Test infrastructure (for reproduction)

  • Guardrail server: /tmp/oss-adopt/e2e/guardrail_server.py — 100 lines, stdlib http.server, no dependencies
  • Proxy config: /tmp/oss-adopt/e2e/proxy_config.yaml — routes haiku-4-5 model through Bedrock, wires test-transform guardrail with streaming_transform_mode: incremental_diff and streaming_sampling_rate: 1
  • Test harness: /tmp/oss-adopt/e2e/run_e2e.py — 6 scenarios, requests library streams SSE, asserts on chunk sequence + finish_reason + guardrail invocation count

All 6 scenarios ran end-to-end against a real LLM producing real chunks — no mocks, no synthetic streams.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri
yucheng-berri force-pushed the litellm_guardrail_streaming_text_transform branch from 6b05b08 to 06ff6b9 Compare July 11, 2026 18:57
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri
yucheng-berri force-pushed the litellm_guardrail_streaming_text_transform branch from 06ff6b9 to c6af8f5 Compare July 11, 2026 18:58
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

2 similar comments
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

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.

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit c47afcd. Configure here.

# semantic on the streaming_transform_mode.
if self.streaming_transform_mode == "incremental_diff":
if guardrail_response.texts is not None:
return_inputs["texts"] = guardrail_response.texts

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.

Non-streaming ignores empty texts

Medium Severity

With streaming_transform_mode set to incremental_diff, _build_guardrail_return_inputs forwards an explicit empty texts list from the guardrail as full suppression. Non-streaming process_output_response still skips applying guardrail text when guardrailed_texts is empty, so the model’s original message content is returned unchanged on post-call paths.

Additional Locations (1)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit c47afcd. Configure here.

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.

This finding is technically accurate but the fix would introduce a different regression. Reasoning:

  • streaming_transform_mode names its scope: it controls streaming behavior. On the non-streaming path (process_output_response), "empty texts" retains the historical block_only contract of "no change" — that's the shape thousands of deployed guardrails were built against.
  • The streaming path is where incremental_diff genuinely needs the is-not-None semantic, because a rewrite-mode guardrail's only signal for "suppress the whole response" is texts: [] (a BLOCK action terminates the stream entirely; GUARDRAIL_INTERVENED with empty texts is the only way to say "emit nothing").
  • Non-streaming callers have a proper channel for suppression: BLOCKED action with a blocked_reason. Overloading empty texts to mean "suppress" on non-streaming would silently change behavior for every deployed guardrail returning texts: [] as a no-op — which is exactly the P1 regression I was called out on in the previous Greptile round.

The scope-limited gate (if self.streaming_transform_mode == 'incremental_diff' inside _build_guardrail_return_inputs) is intentional. The asymmetry between streaming and non-streaming here mirrors the asymmetry of the modes themselves: incremental_diff is opt-in, streaming-only, and adds new semantics; non-streaming preserves the historical contract regardless of the flag.

Not fixing on this PR. If we ever want a true response-suppression signal on non-streaming, that belongs in a separate design conversation (probably a new action code, e.g. SUPPRESSED, rather than overloading empty texts).

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.

Update on this thread: rather than gate the is-not-None semantic on the mode flag (which still crossed into non-streaming callers when the operator sets incremental_diff), reverted _build_guardrail_return_inputs back to its original if guardrail_response.texts: form in 6388c3ee52. Documented as a known limitation on the incremental_diff path.

This PR now touches only the incremental_diff streaming code path — no changes to block_only, non-streaming, or pre_call semantics for any guardrail. A dedicated suppression primitive (probably a new action code, since overloading empty texts is a scope-leak trap) is deferred to a future PR.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai please rescore on the latest HEAD

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

Ready-for-review summary

HEAD: 6a950ff (after final merge of litellm_internal_staging).

What this PR does

Opt-in streaming_transform_mode: incremental_diff on generic_guardrail_api — HTTP guardrail servers can now rewrite streamed response text (PII masking, pseudonym reversal, redaction) on the fly instead of only being able to terminate the stream. Default block_only preserves historical behavior for every deployed guardrail.

Adopted from #32084 by @schneidermr. Mirrored on a litellm_ branch so CircleCI + internal lint run; then driven through the review loop to close every legit bot finding.

Bot findings resolved

Reviewer Finding Resolution
Greptile P1 Mixed content+tool_call finish_reason ordering Fixed: _tool_call_passthrough_chunk defers finish_reason for mixed choices; final synthetic text chunk owns it
Greptile P2 Double guardrail HTTP call at end of stream for tool-call streams Fixed: gated final _round on saw_text_content/deferred_finish_reason
Greptile P2 _emit_streaming_http_error non-A2A raise Rebutted on-thread: framework-wide OpenAI SSE pattern, same as 9 sibling guardrails; proper fix (build_block_sse_chunks override) is a separately-scoped follow-up
Greptile P2 has_stream_ended inspects only choices[0] Rebutted: pre-existing block_only limitation, not introduced by this PR
Greptile P1 texts is not None change breaks non-streaming empty-texts contract Fixed: gated is-not-None semantic on streaming_transform_mode == 'incremental_diff' so block_only and non-streaming callers retain the falsy check
Greptile P1 Transformed text dropped when sampling_rate > 1 and text precedes tool_call Fixed: flush accumulated text via _round(is_final=False) BEFORE yielding tool-call passthrough
Veria Medium Empty texts: [] restores raw input (leak) Fixed within the gated incremental_diff path
Bugbot High Deferred finish_reason never flushed when guardrail suppresses text Fixed: _build_transform_chunk builds a terminator chunk on is_final=True even with empty mutated_text_per_choice
Bugbot Medium _process_streaming_transform texts not sorted by choice index Fixed: sort raw_by_index keys before building indices/texts_to_check
Bugbot Medium Usage-only chunks (stream_options.include_usage) dropped on incremental_diff Fixed: new _is_usage_only_chunk helper, pass through raw
Bugbot Medium Non-streaming ignores empty texts under incremental_diff Rebutted on-thread: fixing this would silently break every deployed guardrail returning texts=[] as a no-op signal on non-streaming; incremental_diff intentionally scoped to streaming
Bugbot Low streaming_buffer_until_moderated ignored under incremental_diff Rebutted on-thread: buffer-until-moderated is inherently incompatible with in-flight text rewrites; the flags are mutually exclusive by design

Live e2e proof (posted earlier: comment)

6/6 scenarios green on a real proxy + real Bedrock LLM + real HTTP guardrail server. No mocks.

  • A: NONE action streams unchanged
  • B: SSN redaction → stream_transform_underflow fails closed (framework-wide SSE fallback via async_data_generator)
  • C: Alice → PERSON_A with stream_holdback_chars=[3] holdback across word boundary
  • D: Guardrail texts=[] mid-stream → fail-closed with partial-leak-before-detection
  • E: BLOCKED action → framework emits well-formed SSE error frame
  • F: Sampling every chunk with identity transform

CI

  • 114 checks passing at HEAD (before this final staging merge; will re-run).
  • Only batches_testing fails — pre-existing flake in tests/batches_tests/, unrelated to guardrail changes. Historical flake rate on this job: high.

Docs

  • In-tree: example_config.yaml and GenericGuardrailAPI class docstring updated with all four streaming knobs and the stream_holdback_chars response contract.
  • External docs at docs.litellm.ai/docs/adding_provider/generic_guardrail_api: tracked as a follow-up (source lives outside this monorepo).

Deferred to follow-up (with team sign-off)

  • Add build_block_sse_chunks override to the OpenAI chat translation handler so mid-stream fail-closed emits a well-formed SSE error frame + data: [DONE] on the OpenAI-chat path — same fix as Anthropic already has, fixes an entire class of framework-wide fail-open across all guardrails (bedrock, model_armor, microsoft_purview, noma, repelloai, tool_permission, litellm_content_filter, aim, cato_networks). Scoped for a separate PR because it also changes block_only behavior for existing OpenAI moderation guardrails.
  • docs.litellm.ai streaming text transformation section (external repo).
  • Design conversation on non-streaming response suppression signal (SUPPRESSED action code, or similar) if we want to consolidate the streaming/non-streaming empty-texts semantics.

Note: Greptile's last score was 3/5 on bef9527c from ~3h ago, before the last three fix commits addressing its own P1s. @greptileai mentions since haven't triggered a fresh review — likely a quota or queue lag. Live e2e + the resolution table above stand in for the score.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

Scope attestation

Audited the PR diff for scope leakage. Found one — the _build_guardrail_return_inputs gate on streaming_transform_mode still crossed into non-streaming callers (pre_call, moderation, non-streaming post_call) whenever an operator configured incremental_diff, because that method is called from apply_guardrail across every hook regardless of whether the caller is streaming.

Reverted the change in 6388c3ee52. _build_guardrail_return_inputs is now byte-identical to its pre-PR form.

Known limitation (documented in class docstring and example_config.yaml)

A guardrail returning action: GUARDRAIL_INTERVENED with texts: [] is treated as a no-op (input passes through) across all modes, including incremental_diff. A dedicated suppression primitive on the streaming transform path is deferred to a future PR — likely as a new action code (e.g. SUPPRESSED) rather than overloading empty texts, which was the scope-leak trap.

What this PR touches, verified

All incremental-diff-specific additions live behind if streaming_transform_mode == "incremental_diff" in UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook:

  • _run_incremental_transform_stream and its 6 helpers (_handle_tool_call_chunk, _tool_call_passthrough_chunk, _build_transform_chunk, _emit_transform_round, _is_usage_only_chunk, _chunk_carries_text, _inspect_text_on_chunk) — new methods, only reachable behind the gate.
  • _process_streaming_transform in OpenAIChatCompletionsHandler — only called when stream_transform_sink is not None, which only my code constructs.
  • StreamTransformSink — new type, only threaded through the incremental_diff path.
  • _resolve_transform_call_type, _emit_streaming_http_error, _inspect_full_response_for_block — new helpers, only reachable behind the gate.
  • streaming_transform_mode, stream_holdback_chars, coerce_stream_holdback_value — new fields/functions; existing paths ignore them.
  • example_config.yaml and class docstring — additive documentation.

No changes to block_only, non-streaming, pre_call, moderation, or any sibling guardrail.

@mateo-berri mateo-berri left a comment

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.

The main concern is a performance issue in _handle_tool_call_chunk: a guardrail API call fires before every tool-call streaming chunk once text is in flight, meaning a streaming function call with many argument chunks generates proportionally many redundant HTTP round-trips even though subsequent calls produce no new delta

Is this greptile comment legit?

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

Comment on lines +689 to +693
should_run_final_round = (
state["text_pending_flush"]
or state["deferred_finish_reason_for_text"]
or (state["saw_text_content"] and not saw_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.

P1 finish_reason silently dropped for text choices in n>1 streams when they finish before separate tool-call chunks

When n>1 and one choice streams text (finishing with finish_reason="stop" in its own terminal chunk) while a separate subsequent chunk carries another choice's tool_calls, the text choice's finish_reason is permanently lost. The sequence is:

  1. Choice 0 terminal text chunk → _record_finish_reasons writes finish_reason_per_choice[0]="stop", text_pending_flush stays True
  2. Choice 1 tool-call chunk → _handle_tool_call_chunk pre-flush fires (is_final=False), emits choice 0's text but not its finish_reason, then clears text_pending_flush
  3. End of loop → should_run_final_round evaluates text_pending_flush=False OR deferred_finish_reason_for_text=False OR (saw_text_content AND NOT saw_tool_calls)=FalseFalse, so the final round never runs and finish_reason_per_choice[0] is never emitted

The existing test suite covers the case where text and tool-call choices arrive in the same chunk (test_n_gt_1_text_and_tool_call_in_same_chunk_no_text_leak) and text-only n>1 (test_per_choice_finish_reason_when_choices_finish_in_different_chunks), but not the scenario where they arrive in separate chunks.

A minimal fix is to add a fourth condition to should_run_final_round: or bool(finish_reason_per_choice). finish_reason_per_choice has entries only for deferred/unreported text-choice finish reasons (tool-only choices get their finish_reason on the passthrough and are not written there), so this triggers the final round only when needed. _build_transform_chunk with is_final=True would then emit an empty-delta chunk carrying finish_reason="stop" for choice 0.

)
holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0))
end = max(len(already), len(text) - holdback)
deltas[choice_idx] = text[len(already) : end]

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: Late guardrail rewrites can expose already streamed text

With the default zero holdback, this emits the entire currently accepted prefix. An attacker can prompt the model to split sensitive content across processing rounds: the guardrail can pass an incomplete identifier initially and redact it after later characters arrive, but the subsequent startswith failure only aborts after the earlier bytes have reached the client. Retain a safe uncommitted suffix by default or buffer until end of stream for transformations that are not guaranteed to be prefix-stable, rather than relying solely on an optional guardrail-provided holdback.

…g path

Four bug fixes on top of the OSS PR's incremental_diff streaming text
transformation, all inside the incremental_diff code paths only. No
existing block_only, non-streaming, or pre_call behavior is touched.

Fix #1 — Mixed content+tool_call finish_reason ordering
  _tool_call_passthrough_chunk now takes an optional finish_reason_per_choice
  map. For a choice carrying both delta.content and delta.tool_calls,
  finish_reason is stripped from the passthrough and recorded on the map so
  the final synthetic text chunk delivers it. Without this, SSE-compliant
  clients stopping at finish_reason drop the guardrailed text — defeating
  the redaction the whole feature exists for. (Greptile P1 twice, Veria.)

Fix #2 — Choice index sort in _process_streaming_transform
  indices/texts_to_check were derived from dict insertion order. For n>1
  streams where choice 1 emits before choice 0, guardrail-returned texts
  aligned to the input order mapped back to the wrong choice indices on
  write-back — wrong text goes to wrong choice. Sort raw_by_index.keys()
  up front so realignment is deterministic. (Bugbot Medium.)

Fix #3 — Cross-chunk pre-tool-call text flush
  With default streaming_sampling_rate=5, text chunks followed by a pure
  tool-call chunk carrying finish_reason='tool_calls' would emit the
  passthrough with finish_reason before any transformed text delta had
  fired. Same failure mode as fix #1 but cross-chunk. Now we flush any
  accumulated text via _round(is_final=False) BEFORE yielding the
  tool-call passthrough. (Greptile P1.)

Fix #4 — Terminator chunk for deferred finish_reason on empty mutated_text
  _build_transform_chunk returned None early when mutated_text_per_choice
  was empty. If a mixed content+tool_call chunk had deferred its
  finish_reason (via fix #1) and the guardrail then suppressed the text
  (empty return), the deferred finish_reason was never delivered. Now on
  is_final=True with empty mutated_text_per_choice, we emit a terminator
  carrying finish_reason per choice from finish_reason_per_choice.
  (Bugbot High.)

Also normalized Optional[X] → X | None across the OSS PR's added surface
via ruff UP045 autofix to keep the strict-rule gate within budget. Pure
mechanical typing style change, no semantic effect.

Regression tests for all four fixes:
- test_mixed_chunk_finish_reason_arrives_after_transformed_text (#1)
- test_text_flush_precedes_tool_call_passthrough (#3)
- test_final_finish_reason_flushed_when_guardrail_suppresses_text (#4)
- test_transform_sends_texts_sorted_by_choice_index (#2)

All fixes reachable only when streaming_transform_mode == 'incremental_diff'
is configured (via _run_incremental_transform_stream) or when a
StreamTransformSink is present (via _process_streaming_transform). Verified
scope-clean: no changes to block_only, non-streaming, pre_call, moderation,
or sibling guardrails.
@yucheng-berri
yucheng-berri force-pushed the litellm_guardrail_streaming_text_transform branch from 2ba08f2 to 527f110 Compare July 13, 2026 19:42
@yucheng-berri

Copy link
Copy Markdown
Contributor Author

Rescoped: only four correctness fixes on top of the OSS PR

Force-pushed a clean history — single commit 527f110a6a on top of Marton's OSS PR head. The prior polish/perf commits have been dropped and rolled up into the follow-up list below.

What this PR fixes (4 correctness bugs)

# Fix Bug class Bot
1 Mixed content+tool_call finish_reason deferred to terminator Data loss — SSE clients drop redacted text Greptile P1 (2x), Veria
2 Choice-index sort in _process_streaming_transform Data corruption — wrong text to wrong choice on n>1 Bugbot Medium
3 Cross-chunk pre-tool-call text flush Data loss — same class as #1, cross-chunk Greptile P1
4 Terminator chunk on empty mutated_text for deferred finish_reason Data loss — deferred finish_reason never delivered Bugbot High

Follow-ups (opt-in path only, no correctness impact)

# Limitation Where Follow-up approach
5 Terminal chunk guardrailed twice (sampled round + EOS flush) when sampling boundary lands on the terminal chunk _run_incremental_transform_stream end-of-loop Add not _chunk_has_finish_reason(item) gate on the sampled round
6 Every tool-call chunk after text fires a redundant guardrail HTTP call producing no delta (perf, sticky-flag bug) _handle_tool_call_chunk-equivalent inline branch Dirty-flag pattern (text_pending_flush set on text, cleared after any _round)
7 Usage-only chunks (stream_options.include_usage) dropped on incremental_diff Text-branch fallthrough Detect via _is_usage_only_chunk helper, pass through raw
8 n>1 streams where a text choice finishes in a separate chunk before/after a pure tool-call chunk: the text choice's finish_reason can be dropped _run_incremental_transform_stream end-of-loop guard Extend final-round guard to fire when finish_reason_per_choice has entries
9 SSE mid-stream raise exc on non-A2A routes: _emit_streaming_http_error re-raises rather than emitting a clean SSE error frame _emit_streaming_http_error:349 Add build_block_sse_chunks override on OpenAIChatCompletionsHandler — fixes framework-wide, not just this PR
10 Empty texts: [] in GUARDRAIL_INTERVENED treated as no-op across all modes (incl. incremental_diff) _build_guardrail_return_inputs Introduce distinct SUPPRESSED action code rather than overloading empty texts
11 Tool-call arguments not text-rewritten in v1 — only inspected at EOS v1 scope v2: add tool_call_arguments_texts to inputs; per-chunk redaction over accumulated function.arguments; diff-emit like text deltas
12 streaming_buffer_until_moderated ignored under incremental_diff Design-incompatible (buffer-until-moderated vs in-flight rewrites) Add runtime warning at config-parse time when both flags set
13 has_stream_ended inspects only choices[0] Pre-existing block_only limitation (predates this PR) any(c.finish_reason is not None for c in chunk.choices)
14 docs.litellm.ai streaming section missing External docs repo Add "Streaming text transformations" section mirroring example_config.yaml block
15 stream_transform_underflow manifests as ChunkedEncodingError on SSE clients Tied to #9 Same fix as #9

What was included in earlier revisions but dropped

  • UP045 typing autofix (Optional[X]X | None): kept in the fix commit because it was required to pass the strict-rule budget gate. Pure mechanical typing style change, no semantic effect.
  • In-tree docs (example_config.yaml "Streaming controls" section, class docstring "Streaming behavior" section): dropped. Follow-up Included PR 12 + Poetry + OpenAI fixes #14 covers them.

Live e2e proof (still applies)

The 6/6 e2e scenarios from comment above were run against a HEAD that included fixes 1-4 (semantically identical to 527f110a6a); the results still hold.

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@yucheng-berri

Copy link
Copy Markdown
Contributor Author

Superseded by #33110 — fresh branch with clean history (single fix commit on top of Marton's OSS PR head). Same content, no accumulated force-push noise.

The four correctness fixes and the follow-up limitations list carry over unchanged to #33110.

@yucheng-berri
yucheng-berri deleted the litellm_guardrail_streaming_text_transform branch July 15, 2026 00:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants