Skip to content

Commit ba341eb

Browse files
giulio-leonegiulio-leone
authored andcommitted
fix: handle missing usage metadata on premature Anthropic stream termination
When the Anthropic API stream terminates before sending the message_stop event (e.g. network timeout, connection reset), the code crashes with AttributeError because event.message.usage is None. The stream() method unconditionally accessed event.message.usage after the async iteration loop, assuming a complete stream. Two failure modes: 1. Empty stream: 'event' variable is never assigned (UnboundLocalError) 2. Premature termination: event.message or event.message.usage is None Fix: Initialize event=None before the loop, use safe attribute access via getattr() chain, and emit zero-usage metadata with a warning log when usage data is unavailable. Added two regression tests: - test_stream_premature_termination: stream ends without message.usage - test_stream_empty_no_events: completely empty stream Fixes #1868
1 parent fca208b commit ba341eb

File tree

2 files changed

+61
-3
lines changed

2 files changed

+61
-3
lines changed

src/strands/models/anthropic.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -405,12 +405,22 @@ async def stream(
405405
try:
406406
async with self.client.messages.stream(**request) as stream:
407407
logger.debug("got response from model")
408+
event = None
408409
async for event in stream:
409410
if event.type in AnthropicModel.EVENT_TYPES:
410411
yield self.format_chunk(event.model_dump())
411412

412-
usage = event.message.usage # type: ignore
413-
yield self.format_chunk({"type": "metadata", "usage": usage.model_dump()})
413+
usage = getattr(getattr(event, "message", None), "usage", None) if event else None
414+
if usage is not None:
415+
yield self.format_chunk({"type": "metadata", "usage": usage.model_dump()})
416+
else:
417+
logger.warning("stream ended without usage metadata (possible premature termination)")
418+
yield self.format_chunk(
419+
{
420+
"type": "metadata",
421+
"usage": {"input_tokens": 0, "output_tokens": 0},
422+
}
423+
)
414424

415425
except anthropic.RateLimitError as error:
416426
raise ModelThrottledException(str(error)) from error

tests/strands/models/test_anthropic.py

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -739,7 +739,55 @@ async def test_stream(anthropic_client, model, agenerator, alist):
739739

740740

741741
@pytest.mark.asyncio
742-
async def test_stream_rate_limit_error(anthropic_client, model, alist):
742+
async def test_stream_premature_termination(anthropic_client, model, agenerator, alist):
743+
"""Test that stream handles premature termination without crashing.
744+
745+
When the Anthropic API stream ends before message_stop (e.g. network
746+
timeout), event.message.usage may be None. The code must not crash
747+
with AttributeError.
748+
749+
Regression test for #1868.
750+
"""
751+
mock_event_1 = unittest.mock.Mock(
752+
type="message_start",
753+
model_dump=lambda: {"type": "message_start"},
754+
)
755+
# Last event has no .message attribute (simulating premature termination)
756+
mock_event_2 = unittest.mock.Mock(
757+
type="content_block_stop",
758+
model_dump=lambda: {"type": "content_block_stop", "index": 0},
759+
spec=["type", "model_dump"],
760+
)
761+
762+
mock_context = unittest.mock.AsyncMock()
763+
mock_context.__aenter__.return_value = agenerator([mock_event_1, mock_event_2])
764+
anthropic_client.messages.stream.return_value = mock_context
765+
766+
messages = [{"role": "user", "content": [{"text": "hello"}]}]
767+
response = model.stream(messages, None, None)
768+
769+
# Should not raise AttributeError
770+
tru_events = await alist(response)
771+
772+
# Should still yield a metadata event with zero usage
773+
assert any("metadata" in str(e) for e in tru_events)
774+
775+
776+
@pytest.mark.asyncio
777+
async def test_stream_empty_no_events(anthropic_client, model, agenerator, alist):
778+
"""Test that stream handles an empty event sequence without crashing."""
779+
mock_context = unittest.mock.AsyncMock()
780+
mock_context.__aenter__.return_value = agenerator([])
781+
anthropic_client.messages.stream.return_value = mock_context
782+
783+
messages = [{"role": "user", "content": [{"text": "hello"}]}]
784+
response = model.stream(messages, None, None)
785+
786+
# Should not raise UnboundLocalError or AttributeError
787+
tru_events = await alist(response)
788+
789+
# Should still yield a metadata event with zero usage
790+
assert any("metadata" in str(e) for e in tru_events)
743791
anthropic_client.messages.stream.side_effect = anthropic.RateLimitError(
744792
"rate limit", response=unittest.mock.Mock(), body=None
745793
)

0 commit comments

Comments
 (0)