fix(streaming): do not raise when a collected stream carries raw SSE frames

The guardrail post-call hooks collect a stream by appending whatever the
response iterator yields, so `all_chunks` can hold raw SSE frames (bytes or
str) next to parsed chunks. `stream_chunk_builder` indexes every entry, so
the first raw frame raises and the hook turns a completion the client had
already received into a 500.

Six of the nine guardrails that reassemble a collected stream call
`stream_chunk_builder` with no guard for this: noma, model_armor, repelloai,
panw_prisma_airs, cisco_ai_defense and microsoft_purview. Only the bedrock
and tool_permission hooks check `is_raw_sse_stream` first, and presidio
filters by type while collecting.

Fix it once at the assembler rather than six times at the callers. Assembling
only the parsed remainder would hand back a response that silently omits
content, which a guardrail would then scan and pass, so treat the whole
stream as unassemblable instead, matching the `any(...)` stance
`is_raw_sse_stream` already takes. Returning None meets the empty-case
contract immediately above it, and every one of those callers already routes
a non-ModelResponse result to its pass-through path.

Relates to #37873
This commit is contained in:
Vineeth Sai 2026-08-22 09:54:21 -07:00
parent 6a0d03914c
commit 462199de08
2 changed files with 64 additions and 0 deletions

View file

@ -8548,6 +8548,19 @@ def stream_chunk_builder(
if not chunks:
return None
# A collected stream can carry raw SSE frames (bytes/str) alongside parsed
# chunks, because the guardrail post-call hooks append whatever the response
# iterator yields. Those have no chunk shape, so the first one reached raises
# and the caller turns a completion the client already received into a 500.
# Assembling only the parsed remainder would instead hand back a response that
# silently omits content, so treat the whole stream as unassemblable, matching
# the `any(...)` stance `is_raw_sse_stream` already takes for this shape.
# Returning None meets the empty-case contract just above, which every caller
# handles by passing the stream through unscanned.
if any(isinstance(chunk, (str, bytes, bytearray)) for chunk in chunks):
verbose_logger.debug("stream_chunk_builder: raw SSE frames in the collected stream, not assembling")
return None
processor: Final = ChunkProcessor(chunks, messages)
chunks = processor.chunks

View file

@ -1337,3 +1337,54 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate(
assert usage.completion_tokens == 100
assert usage.completion_tokens_details.reasoning_tokens == expected_reasoning_tokens
assert usage.completion_tokens_details.text_tokens == expected_text_tokens
class TestRawSSEFramesInACollectedStream:
"""A collected stream can carry raw SSE frames next to parsed chunks.
The guardrail post-call hooks append whatever the response iterator yields, so
`stream_chunk_builder` receives bytes/str entries it cannot index. Reaching one
used to raise, and the caller turned a completion the client had already
received into a 500 (issue #37873). Only two of the nine guardrails that
reassemble a stream carry the `is_raw_sse_stream` guard that avoids this call.
"""
@staticmethod
def _chunk(content):
return ModelResponseStream(
id="chatcmpl-a83572f7",
created=1,
model="nvidia_nim/meta/llama-3.1-8b-instruct",
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(role="assistant", content=content),
)
],
)
@pytest.mark.parametrize(
"frame",
[b'data: {"id":"x"}\n\n', 'data: {"id":"x"}\n\n', bytearray(b"data: x\n\n")],
)
def test_a_raw_frame_does_not_raise(self, frame):
chunks = [self._chunk("Hello"), frame, self._chunk(" world")]
# Must not raise: an APIError here becomes a 500 for an already-delivered completion.
assert stream_chunk_builder(chunks=chunks) is None
def test_an_all_raw_stream_does_not_raise(self):
assert stream_chunk_builder(chunks=[b"data: a\n\n", b"data: b\n\n"]) is None
def test_a_stream_of_real_chunks_is_still_assembled(self):
"""Control: the ordinary path is untouched, so this passes with or without the fix."""
built = stream_chunk_builder(chunks=[self._chunk("Hello"), self._chunk(" world")])
assert built is not None
assert built.choices[0].message.content == "Hello world"
def test_empty_and_none_contracts_are_unchanged(self):
"""Control: the existing empty-case contract still holds."""
assert stream_chunk_builder(chunks=[]) is None