From 175f2e46e2ab2380c975d5cb769d389df6a4bfea Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 23:10:46 +0000 Subject: [PATCH] fix(bedrock): raise on ConverseStream ending without messageStop instead of silently completing --- litellm/llms/bedrock/chat/invoke_handler.py | 26 ++++++ .../llms/bedrock/chat/test_invoke_handler.py | 93 ++++++++++++++++++- 2 files changed, 118 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 4c256be1ab8..2da6abf0ab7 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1288,6 +1288,8 @@ class AWSEventStreamDecoder: self.response_id: Optional[str] = None self.json_mode = json_mode self._current_tool_name: Optional[str] = None + self.converse_stream_started: bool = False + self.converse_stop_reason_seen: bool = False def check_empty_tool_call_args(self) -> bool: """ @@ -1485,6 +1487,7 @@ class AWSEventStreamDecoder: # Capture the conversationId from the first messageStart event # and use it as the consistent ID for all subsequent chunks. self._initialize_converse_response_id(chunk_data) + self.converse_stream_started = True verbose_logger.debug("\n\nRaw Chunk: {}\n\n".format(chunk_data)) text = "" @@ -1517,6 +1520,7 @@ class AWSEventStreamDecoder: elif "contentBlockIndex" in chunk_data: # stop block, no 'start' or 'delta' object tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: + self.converse_stop_reason_seen = True finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) elif "usage" in chunk_data: usage = converse_config._transform_usage(chunk_data.get("usage", {})) @@ -1605,6 +1609,26 @@ class AWSEventStreamDecoder: tool_use=None, ) + def _raise_if_converse_stream_incomplete(self) -> None: + """ + A Bedrock ConverseStream must end with a terminal 'messageStop' event + (carrying stopReason). AWS acknowledges mid-stream faults on cross-region + inference profiles where the HTTP stream terminates cleanly after partial + content but before messageStop/metadata. Surfacing this as an error lets + callers retry a fresh request instead of trusting a truncated response + (e.g. a tool call with unparseable arguments) delivered as a success. + """ + if self.converse_stream_started and not self.converse_stop_reason_seen: + raise BedrockError( + status_code=500, + message=( + "Bedrock ConverseStream ended without a terminal 'messageStop' event; " + "the response is incomplete and any tool-call arguments may be truncated. " + "Treating as an error so the request can be retried instead of returning a " + "silently-completed, partial response" + ), + ) + def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[Union[GChunk, ModelResponseStream, dict]]: """Given an iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer @@ -1618,6 +1642,7 @@ class AWSEventStreamDecoder: # sse_event = ServerSentEvent(data=message, event="completion") _data = json.loads(message) yield self._chunk_parser(chunk_data=_data) + self._raise_if_converse_stream_incomplete() async def aiter_bytes( self, iterator: AsyncIterator[bytes] @@ -1633,6 +1658,7 @@ class AWSEventStreamDecoder: if message: _data = json.loads(message) yield self._chunk_parser(chunk_data=_data) + self._raise_if_converse_stream_incomplete() def _parse_message_from_event(self, event) -> Optional[str]: response_stream_shape = get_bedrock_response_stream_shape() diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index 61987d25d9c..c199a8ef641 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -1,6 +1,7 @@ +import json import os import sys -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -11,6 +12,7 @@ sys.path.insert( import litellm from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, + BedrockError, BedrockLLM, make_call, make_sync_call, @@ -18,6 +20,95 @@ from litellm.llms.bedrock.chat.invoke_handler import ( from litellm.llms.custom_httpx.http_handler import HTTPHandler +def _mock_event(body: dict) -> MagicMock: + """Build a botocore-style event whose 200 body is the given ConverseStream JSON.""" + event = MagicMock() + event.to_response_dict.return_value = { + "status_code": 200, + "headers": {}, + "body": json.dumps(body).encode(), + } + return event + + +def _decode_events_sync(decoder: AWSEventStreamDecoder, bodies: list[dict]): + mock_buffer = MagicMock() + mock_buffer.__iter__.return_value = [_mock_event(b) for b in bodies] + with patch("botocore.eventstream.EventStreamBuffer", return_value=mock_buffer): + return list(decoder.iter_bytes(iter([b"raw"]))) + + +async def _decode_events_async(decoder: AWSEventStreamDecoder, bodies: list[dict]): + mock_buffer = MagicMock() + mock_buffer.__iter__.return_value = [_mock_event(b) for b in bodies] + + async def _aiter(): + yield b"raw" + + with patch("botocore.eventstream.EventStreamBuffer", return_value=mock_buffer): + return [chunk async for chunk in decoder.aiter_bytes(_aiter())] + + +_TRUNCATED_TOOL_CALL_STREAM = [ + {"messageStart": {"role": "assistant"}}, + {"start": {"toolUse": {"toolUseId": "tooluse_abc", "name": "shell"}}, "contentBlockIndex": 0}, + {"delta": {"toolUse": {"input": '{"command": ["cat",".glia/project.md"]'}}, "contentBlockIndex": 0}, +] + +_COMPLETE_TOOL_CALL_STREAM = _TRUNCATED_TOOL_CALL_STREAM + [ + {"delta": {"toolUse": {"input": "}"}}, "contentBlockIndex": 0}, + {"contentBlockIndex": 0}, + {"stopReason": "tool_use"}, + {"usage": {"inputTokens": 10641, "outputTokens": 48, "totalTokens": 10689}, "metrics": {"latencyMs": 100}}, +] + + +def test_converse_stream_without_message_stop_raises_sync(): + """A ConverseStream that ends mid-tool-call without a terminal messageStop + (stopReason) event must surface as an error, not be silently completed with a + fabricated finish_reason and a synthesized usage chunk. Regression for #32686.""" + decoder = AWSEventStreamDecoder(model="anthropic.claude-opus-4-1-20250805-v1:0") + with pytest.raises(BedrockError) as exc_info: + _decode_events_sync(decoder, _TRUNCATED_TOOL_CALL_STREAM) + assert exc_info.value.status_code == 500 + assert "messageStop" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_converse_stream_without_message_stop_raises_async(): + decoder = AWSEventStreamDecoder(model="anthropic.claude-opus-4-1-20250805-v1:0") + with pytest.raises(BedrockError) as exc_info: + await _decode_events_async(decoder, _TRUNCATED_TOOL_CALL_STREAM) + assert exc_info.value.status_code == 500 + assert "messageStop" in exc_info.value.message + + +def test_converse_stream_with_message_stop_does_not_raise_sync(): + """A well-formed ConverseStream ending in messageStop + metadata must not raise.""" + decoder = AWSEventStreamDecoder(model="anthropic.claude-opus-4-1-20250805-v1:0") + chunks = _decode_events_sync(decoder, _COMPLETE_TOOL_CALL_STREAM) + finish_reasons = [ + c.choices[0].finish_reason for c in chunks if hasattr(c, "choices") and c.choices and c.choices[0].finish_reason + ] + assert "tool_calls" in finish_reasons + + +@pytest.mark.asyncio +async def test_converse_stream_with_message_stop_does_not_raise_async(): + decoder = AWSEventStreamDecoder(model="anthropic.claude-opus-4-1-20250805-v1:0") + chunks = await _decode_events_async(decoder, _COMPLETE_TOOL_CALL_STREAM) + assert any(hasattr(c, "choices") and c.choices and c.choices[0].finish_reason == "tool_calls" for c in chunks) + + +def test_non_converse_invoke_stream_end_without_stop_reason_does_not_raise(): + """The messageStop guard is Converse-only; a non-Converse invoke text stream + (e.g. cohere) that never routes through converse_chunk_parser must not be + affected by the incomplete-stream check.""" + decoder = AWSEventStreamDecoder(model="cohere.command-text-v14") + chunks = _decode_events_sync(decoder, [{"text": "hello"}, {"text": " world"}]) + assert [c["text"] for c in chunks] == ["hello", " world"] + + def test_transform_thinking_blocks_with_redacted_content(): thinking_block = {"redactedContent": "This is a redacted content"} decoder = AWSEventStreamDecoder(model="test")