diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index 50c05daee11..f5e08e7b603 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -15,7 +15,11 @@ from litellm.types.utils import Choices, ModelResponse def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool: - return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks) + # A genuine SSE stream carries only str/bytes frames. A stream that mixes + # parsed chunk objects with a stray str/bytes chunk (e.g. nvidia_nim, + # #37873) is not SSE: it must go to stream_chunk_builder, not the + # fail-closed Anthropic-SSE path. + return len(all_chunks) > 0 and all(isinstance(chunk, (str, bytes)) for chunk in all_chunks) def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 30526d30dc5..e9ad0ec5ac7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -2751,11 +2751,23 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # /v1/messages arrives as SSE frames, which stream_chunk_builder cannot assemble raw_sse: Final = is_raw_sse_stream(all_chunks) - assembled_model_response: ModelResponse | TextCompletionResponse | None = ( - assemble_anthropic_sse_stream(all_chunks, restore_identity=True) - if raw_sse - else stream_chunk_builder(chunks=all_chunks) - ) + assembled_model_response: ModelResponse | TextCompletionResponse | None + if raw_sse: + assembled_model_response = assemble_anthropic_sse_stream(all_chunks, restore_identity=True) + else: + # Some providers yield stray str/bytes chunks among the parsed + # chunks (e.g. nvidia_nim, #37873). They carry no model content + # and crash stream_chunk_builder, so exclude them from assembly; + # the OUTPUT scan still runs on every real chunk. + assemblable_chunks: Final = [ # mutable-ok: filtered copy passed to stream_chunk_builder + c for c in all_chunks if not isinstance(c, (str, bytes)) + ] + if len(assemblable_chunks) != len(all_chunks): + verbose_proxy_logger.warning( + "BedrockGuardrail: ignoring %d non-frame chunk(s) before stream assembly", + len(all_chunks) - len(assemblable_chunks), + ) + assembled_model_response = stream_chunk_builder(chunks=assemblable_chunks) if isinstance(assembled_model_response, ModelResponse): pre_guardrail_text: Final = model_response_text(assembled_model_response) _pre_block_response: Final = assembled_model_response diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index a8b33109900..025a9bfe421 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -916,9 +916,23 @@ class ToolPermissionGuardrail(CustomGuardrail): async for chunk in response: all_chunks.append(chunk) - assembled_model_response: Final[ModelResponse | TextCompletionResponse | None] = ( - stream_chunk_builder(chunks=all_chunks) if not is_raw_sse_stream(all_chunks) else None - ) + assembled_model_response: ModelResponse | TextCompletionResponse | None + if is_raw_sse_stream(all_chunks): + assembled_model_response = None + else: + # Some providers yield stray str/bytes chunks among the parsed + # chunks (e.g. nvidia_nim, #37873). They carry no model content + # and crash stream_chunk_builder, so exclude them from assembly; + # permission checks still run on every real chunk. + assemblable_chunks: Final = [ # mutable-ok: filtered copy passed to stream_chunk_builder + c for c in all_chunks if not isinstance(c, (str, bytes)) + ] + if len(assemblable_chunks) != len(all_chunks): + verbose_proxy_logger.warning( + "ToolPermissionGuardrail: ignoring %d non-frame chunk(s) before stream assembly", + len(all_chunks) - len(assemblable_chunks), + ) + assembled_model_response = stream_chunk_builder(chunks=assemblable_chunks) if isinstance(assembled_model_response, ModelResponse): denied_tools = self._check_assembled_stream(assembled_model_response) if denied_tools: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 953e3de1519..0bd7a672433 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2251,6 +2251,59 @@ async def test_streaming_post_call_only_runs_output_scan(): assert len(input_calls) == 0 +@pytest.mark.asyncio +async def test_post_call_streaming_iterator_hook_ignores_stray_bytes_chunk(): + """ + Regression test for #37873: a provider (nvidia_nim) yielding a stray raw + bytes chunk made stream_chunk_builder raise TypeError on c["choices"], + which turned a successful streamed completion into a 500. The hook must + ignore non-frame chunks for assembly and still run the OUTPUT scan on + the real chunks, so enforcement is never skipped. + """ + guardrail = BedrockGuardrail( + guardrail_name="bedrock-stream-bytes-chunk", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + mock_chunks = [ + litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="Hi", role="assistant"), + finish_reason="stop", + index=0, + ) + ], + created=1, + model="nvidia_nim/meta/llama-3.1-8b-instruct", + object="chat.completion.chunk", + ), + b"raw bytes chunk that stream_chunk_builder cannot index", + ] + + async def mock_stream(): + for c in mock_chunks: + yield c + + minimal = {"action": "NONE", "assessments": [], "outputs": []} + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: + out = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data={"model": "nvidia_nim/meta/llama-3.1-8b-instruct", "messages": []}, + ): + out.append(chunk) + + output_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT"] + assert len(output_calls) == 1 + assert out, "streamed content must reach the client" + + + @pytest.mark.asyncio async def test_streaming_post_call_output_only_path_passes_request_data_to_make_bedrock(): """When INPUT validation is skipped (pre/during already ran), OUTPUT still gets request_data.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 0dbd4591ac9..310707ea989 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -25,8 +25,10 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( from litellm.types.utils import ( ChatCompletionMessageToolCall, Choices, + Delta, ModelResponse, ModelResponseStream, + StreamingChoices, ) @@ -1199,6 +1201,46 @@ class TestToolPermissionGuardrailAnthropicMessages: assert out == chunks, "an allowed stream must not be re-serialized" + @pytest.mark.asyncio + async def test_denied_tool_call_in_mixed_stream_with_stray_bytes_is_blocked(self): + # Regression test for #37873: a provider yielding a stray bytes chunk + # must not crash stream assembly nor disable permission enforcement. + tool_call_chunk = ModelResponseStream( + id="chatcmpl-tool", + created=1700000000, + model="nvidia_nim/meta/llama-3.1-8b-instruct", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "Read", "arguments": '{"file_path": "/etc/passwd"}'}, + } + ], + ), + ) + ], + ) + finish_chunk = ModelResponseStream( + id="chatcmpl-tool", + created=1700000000, + model="nvidia_nim/meta/llama-3.1-8b-instruct", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="tool_calls")], + ) + chunks = [tool_call_chunk, b"stray provider bytes", finish_chunk] + + with patch.object(self.blocking, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self._drain(self.blocking, chunks) + + @pytest.mark.asyncio async def test_rewrite_mode_removes_denied_tool_use_from_anthropic_sse_stream(self): with patch.object(self.rewriting, "should_run_guardrail", return_value=True): diff --git a/tests/test_litellm/proxy/guardrails/test_anthropic_sse.py b/tests/test_litellm/proxy/guardrails/test_anthropic_sse.py new file mode 100644 index 00000000000..b1c981c4e8e --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_anthropic_sse.py @@ -0,0 +1,39 @@ +""" +Unit tests for litellm/proxy/guardrails/anthropic_sse.py +""" + +import litellm +from litellm.proxy.guardrails.anthropic_sse import is_raw_sse_stream + + +def _chunk() -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="Hi", role="assistant"), + finish_reason=None, + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + +def test_is_raw_sse_stream_all_str_or_bytes(): + assert is_raw_sse_stream([b"event: message_start\ndata: {}"]) is True + assert is_raw_sse_stream(["event: message_start", b"data: {}"]) is True + + +def test_is_raw_sse_stream_parsed_chunks(): + assert is_raw_sse_stream([_chunk(), _chunk()]) is False + assert is_raw_sse_stream([]) is False + + +def test_is_raw_sse_stream_mixed_stream_is_not_sse(): + # #37873: a provider yielding one stray bytes chunk among parsed chunks + # must classify as a normal (non-SSE) stream so it reaches + # stream_chunk_builder, not the fail-closed Anthropic-SSE path. + assert is_raw_sse_stream([_chunk(), b"stray bytes"]) is False