From 0ea2557ffe0871d25835fc44d4df075a331dceb0 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:00:11 -0700 Subject: [PATCH] fix(agentcore): end a stream on messageStop when the runtime sends no metadata or final message --- .../bedrock/chat/agentcore/transformation.py | 27 +++++++ .../test_agentcore_transformation.py | 80 ++++++++++++++++++- 2 files changed, 106 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 48b3795e274..b3706b6c4e8 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -540,6 +540,17 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return AgentCoreParsedResponse(content=content, usage=usage_data, final_message=final_message) + @staticmethod + def _message_stop_chunk(model: str) -> ModelResponseStream: + chunk: Final = ModelResponseStream( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=model, + object="chat.completion.chunk", + ) + chunk.choices = [StreamingChoices(finish_reason="stop", index=0, delta=Delta())] + return chunk + def _stream_agentcore_response_sync( self, response: httpx.Response, @@ -549,6 +560,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): Internal sync generator that parses SSE and yields ModelResponse chunks. """ buffer = "" + sent_finish = False + message_stopped = False # runtimes that send no metadata or final message end on messageStop for text_chunk in response.iter_text(): buffer += text_chunk @@ -620,7 +633,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): total_tokens=usage_data.get("totalTokens", 0), ), ) + sent_finish = True yield chunk + message_stopped |= "messageStop" in event_payload # Process final message if "message" in data_obj and isinstance(data_obj["message"], dict): @@ -637,12 +652,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): delta=Delta(), ) ] + sent_finish = True yield chunk except json.JSONDecodeError: verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100]) continue + if message_stopped and not sent_finish: + yield self._message_stop_chunk(model) + def get_sync_custom_stream_wrapper( self, model: str, @@ -771,6 +790,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): Internal async generator that parses SSE and yields ModelResponse chunks. """ buffer = "" + sent_finish = False + message_stopped = False # runtimes that send no metadata or final message end on messageStop async for text_chunk in response.aiter_text(): buffer += text_chunk @@ -842,7 +863,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): total_tokens=usage_data.get("totalTokens", 0), ), ) + sent_finish = True yield chunk + message_stopped |= "messageStop" in event_payload # Process final message if "message" in data_obj and isinstance(data_obj["message"], dict): @@ -859,12 +882,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): delta=Delta(), ) ] + sent_finish = True yield chunk except json.JSONDecodeError: verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100]) continue + if message_stopped and not sent_finish: + yield self._message_stop_chunk(model) + async def get_async_custom_stream_wrapper( self, model: str, diff --git a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index 053ba477212..8ad70d120ac 100644 --- a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -19,8 +19,9 @@ from typing import Final from unittest.mock import MagicMock, Mock, patch import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler class TestAgentCoreAcceptHeader: @@ -1269,3 +1270,80 @@ def test_agentcore_runtime_response_becomes_assistant_message( assert json.loads(route.calls[0].request.content)["prompt"] == "Explain machine learning in simple terms" assert response.choices[0].message.content == "Machine learning lets computers learn from data." assert response.choices[0].finish_reason == "stop" + + +_AGENTCORE_RUNTIME_URL: Final = ( + "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/" + "arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A888602223428%3Aruntime%2Fhosted_agent_13sf6-cALnp38iZD/invocations" +) +_AGENTCORE_DELTAS: Final = ( + 'data: {"event":{"contentBlockDelta":{"delta":{"text":"Machine learning lets "}}}}\n\n' + 'data: {"event":{"contentBlockDelta":{"delta":{"text":"computers learn from data."}}}}\n\n' +) + + +async def _stream_agentcore(body: str, sync_mode: bool) -> list[tuple[str | None, str | None]]: + """Stream `body` from the runtime as SSE and return each chunk's (content, finish_reason).""" + kwargs: Final = { + "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/" + "hosted_agent_13sf6-cALnp38iZD", + "messages": [{"role": "user", "content": "Explain machine learning in simple terms"}], + "aws_access_key_id": "AKIAAGENTCOREKEY", + "aws_secret_access_key": "agentcore-secret", + "stream": True, + } + if sync_mode: + return [ + (chunk.choices[0].delta.content, chunk.choices[0].finish_reason) + for chunk in litellm.completion(**kwargs, client=HTTPHandler()) + ] + return [ + (chunk.choices[0].delta.content, chunk.choices[0].finish_reason) + async for chunk in await litellm.acompletion(**kwargs, client=AsyncHTTPHandler()) + ] + + +@pytest.mark.parametrize( + "ending", + [ + pytest.param('data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n', id="message_stop"), + pytest.param( + 'data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n' + 'data: {"event":{"metadata":{"usage":{"inputTokens":9,"outputTokens":7,"totalTokens":16}}}}\n\n' + 'data: {"message":{"role":"assistant","content":[{"text":"Machine learning lets computers learn from data."}]}}\n\n', + id="strands_full", + ), + ], +) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.respx(assert_all_called=True) +async def test_agentcore_stream_that_signals_its_end_completes( + ending: str, sync_mode: bool, respx_mock: respx.Router, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(_AGENTCORE_RUNTIME_URL).mock( + return_value=httpx.Response( + 200, content=(_AGENTCORE_DELTAS + ending).encode(), headers={"content-type": "text/event-stream"} + ) + ) + + chunks: Final = await _stream_agentcore(_AGENTCORE_DELTAS + ending, sync_mode) + + assert "".join(content for content, _ in chunks if content) == "Machine learning lets computers learn from data." + assert [finish for _, finish in chunks if finish] == ["stop"] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.respx(assert_all_called=True) +async def test_agentcore_stream_cut_before_it_signals_its_end_raises( + sync_mode: bool, respx_mock: respx.Router, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(_AGENTCORE_RUNTIME_URL).mock( + return_value=httpx.Response( + 200, content=_AGENTCORE_DELTAS.encode(), headers={"content-type": "text/event-stream"} + ) + ) + + with pytest.raises(MidStreamFallbackError, match="Stream ended before a finish_reason was received"): + await _stream_agentcore(_AGENTCORE_DELTAS, sync_mode)