diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index a88b148e92c..3aaba0b139c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -8,8 +8,8 @@ to run through agentic completion hooks. If an agentic hook fires, the follow-up response is chained as Phase 2 of the same iterator. In hold-back mode (``hold_back=True``) chunks are buffered instead of yielded -live, keepalive pings run until the hooks finish, and then either the follow-up -replaces the message or the buffer replays, except that a buffered tool_use for +live, keepalive pings run whenever no other byte is ready, and then either the +follow-up replaces the message or the buffer replays, except that a tool_use for a server-fulfilled tool fails the turn rather than reaching a client that cannot execute it. """ @@ -30,6 +30,14 @@ SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n' ) + +async def _anext_or_none(iterator: AsyncIterator) -> bytes | None: + try: + return await iterator.__anext__() + except StopAsyncIteration: + return None + + # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) # --------------------------------------------------------------------------- @@ -195,6 +203,7 @@ class AgenticAnthropicStreamingIterator: self._follow_up_iterator: AsyncIterator | None = None self._drain_task: asyncio.Task | None = None self._hook_task: asyncio.Task | None = None + self._follow_up_chunk_task: asyncio.Task | None = None self._replay_index = 0 self._error_emitted = False @@ -253,7 +262,7 @@ class AgenticAnthropicStreamingIterator: return PING_SSE_BYTES if self._follow_up_iterator is not None: - return await self._follow_up_iterator.__anext__() + return await self._next_follow_up_chunk(self._follow_up_iterator) if self._buffer_holds_server_fulfilled_tool_use(): if self._error_emitted: @@ -273,6 +282,17 @@ class AgenticAnthropicStreamingIterator: raise StopAsyncIteration + async def _next_follow_up_chunk(self, follow_up_iterator: AsyncIterator) -> bytes: + if self._follow_up_chunk_task is None: + self._follow_up_chunk_task = asyncio.create_task(_anext_or_none(follow_up_iterator)) + if not await self._completed_within_ping_interval(self._follow_up_chunk_task): + return PING_SSE_BYTES + chunk: Final = self._follow_up_chunk_task.result() + self._follow_up_chunk_task = None + if chunk is None: + raise StopAsyncIteration + return chunk + def _buffer_holds_server_fulfilled_tool_use(self) -> bool: if not self._server_fulfilled_tool_names: return False @@ -307,6 +327,7 @@ class AgenticAnthropicStreamingIterator: await self._settle_task(self._drain_task) await self._settle_task(self._hook_task) + await self._settle_task(self._follow_up_chunk_task) await aclose_if_supported(self._inner) await aclose_if_supported(self._follow_up_iterator) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py index d59f76232cd..d4aebf099d1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -1001,6 +1001,65 @@ class TestAgenticStreamingIteratorHoldBack: assert [c for c in collected if c != PING_SSE_BYTES] == chunks + @pytest.mark.asyncio + async def test_should_emit_pings_while_the_follow_up_stream_is_slow(self): + """The corrected answer can be slow to generate, so the follow-up stream gets keepalives too.""" + chunks = _build_tool_use_stream() + phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"] + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockSlowAsyncStream(phase2_chunks, delay_seconds=0.06) + ) + + iterator = _build_hold_back_iterator( + MockAsyncStream(chunks), + mock_handler, + ping_interval_seconds=0.02, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + first_follow_up_index = collected.index(phase2_chunks[0]) + assert collected[first_follow_up_index + 1] == PING_SSE_BYTES + assert [c for c in collected if c != PING_SSE_BYTES] == phase2_chunks + + @pytest.mark.asyncio + async def test_should_propagate_follow_up_stream_error(self): + """A failing follow-up stream surfaces its error instead of hanging on pings forever.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockFailingAsyncStream([b"follow-up-chunk"], RuntimeError("follow-up died")) + ) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler, ping_interval_seconds=0.02) + + with pytest.raises(RuntimeError, match="follow-up died"): + async for _ in iterator: + pass + + @pytest.mark.asyncio + async def test_aclose_cancels_in_flight_follow_up_chunk_task(self): + """Closing while a follow-up chunk is pending must not orphan that task.""" + chunks = _build_tool_use_stream() + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=MockSlowAsyncStream([b"follow-up-chunk"], delay_seconds=5.0) + ) + + iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler, ping_interval_seconds=0.02) + + while iterator._follow_up_chunk_task is None: + await iterator.__anext__() + + await iterator.aclose() + assert iterator._follow_up_chunk_task.cancelled() + @pytest.mark.asyncio async def test_aclose_cancels_drain_task(self): """Closing the iterator mid-buffer must cancel the background drain task."""