fix(anthropic): keep pinging while the held-back follow-up stream is in flight

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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)
 

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."""