fix(anthropic): buffer streamed responses carrying server-fulfilled tools so retrieval tool calls never reach the client

litellm/integrations/compression_interception/handler.py

@@ -7,7 +7,7 @@
 
 import time
 import uuid
-from typing import Any, Final, cast
+from typing import Any, ClassVar, Final, cast
 
 from litellm._logging import verbose_logger
 from litellm.compression import compress
@@ -72,6 +72,8 @@ class CompressionInterceptionLogger(CustomLogger):
     4. Build typed rerun plan with tool_result blocks from the compressed cache.
     """
 
+    server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({LITELLM_CONTENT_RETRIEVE_TOOL_NAME})
+
     def __init__(
         self,
         enabled: bool = True,

litellm/integrations/custom_logger.py

@@ -3,7 +3,7 @@
 import re
 import traceback
 from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Any, Final, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
 
 from pydantic import BaseModel
 
@@ -60,6 +60,8 @@
 
 class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callback#callback-class
     # Class variables or attributes
+    server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset()
+
     def __init__(
         self,
         turn_off_message_logging: bool = False,

litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py

@@ -6,14 +6,38 @@
 all bytes, and on stream exhaustion rebuilds the full Anthropic response
 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 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.
 """
 
+import asyncio
+import contextlib
 import json
 from collections.abc import AsyncIterator
 from typing import Any, Final, cast
 
 from litellm._logging import verbose_logger
 
+PING_SSE_BYTES: Final = b'event: ping\ndata: {"type": "ping"}\n\n'
+HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0
+SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = (
+    b"event: error\n"
+    b'data: {"type": "error", "error": {"type": "api_error", "message": '
+    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)
 # ---------------------------------------------------------------------------
@@ -156,6 +180,9 @@ def __init__(
         logging_obj: Any,
         custom_llm_provider: str,
         kwargs: dict,
+        hold_back: bool = False,
+        server_fulfilled_tool_names: frozenset[str] = frozenset(),
+        ping_interval_seconds: float = HOLD_BACK_PING_INTERVAL_SECONDS,
     ):
         self._inner = completion_stream.__aiter__()
         self._http_handler = http_handler
@@ -166,16 +193,27 @@ def __init__(
         self._logging_obj = logging_obj
         self._custom_llm_provider = custom_llm_provider
         self._kwargs = kwargs
+        self._hold_back = hold_back
+        self._server_fulfilled_tool_names = server_fulfilled_tool_names
+        self._ping_interval_seconds = ping_interval_seconds
 
         self._collected_bytes: list[bytes] = []
         self._stream_exhausted = False
         self._hook_processing_done = False
         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
 
     def __aiter__(self):
         return self
 
     async def __anext__(self) -> bytes:
+        if self._hold_back:
+            return await self._anext_held_back()
+
         # Phase 1: yield from upstream, collect bytes
         if not self._stream_exhausted:
             try:
@@ -194,11 +232,102 @@ async def __anext__(self) -> bytes:
 
         raise StopAsyncIteration
 
+    async def _drain_upstream(self) -> None:
+        try:
+            while True:
+                self._collected_bytes.append(await self._inner.__anext__())
+        except StopAsyncIteration:
+            return
+
+    async def _completed_within_ping_interval(self, task: asyncio.Task) -> bool:
+        try:
+            await asyncio.wait_for(asyncio.shield(task), timeout=self._ping_interval_seconds)
+        except asyncio.TimeoutError:
+            return False
+        return True
+
+    async def _anext_held_back(self) -> bytes:
+        if self._drain_task is None:
+            self._drain_task = asyncio.create_task(self._drain_upstream())
+            return PING_SSE_BYTES
+
+        if not self._stream_exhausted:
+            if not await self._completed_within_ping_interval(self._drain_task):
+                return PING_SSE_BYTES
+            self._stream_exhausted = True
+
+        if self._hook_task is None:
+            self._hook_task = asyncio.create_task(self._process_agentic_hooks())
+        if not await self._completed_within_ping_interval(self._hook_task):
+            return PING_SSE_BYTES
+
+        if self._follow_up_iterator is not None:
+            return await self._next_follow_up_chunk(self._follow_up_iterator)
+
+        if self._buffer_holds_server_fulfilled_tool_use():
+            if self._error_emitted:
+                raise StopAsyncIteration
+            self._error_emitted = True
+            verbose_logger.error(
+                "AgenticStreamingIterator: hooks did not replace a message containing a server-fulfilled "
+                "tool_use [model=%s]; emitting an SSE error instead of leaking the tool call to the client",
+                self._model,
+            )
+            return SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES
+
+        if self._replay_index < len(self._collected_bytes):
+            chunk: Final = self._collected_bytes[self._replay_index]
+            self._replay_index += 1
+            return chunk
+
+        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
+        started_blocks: Final = (
+            data.get("content_block")
+            for event_type, data in _parse_sse_events(b"".join(self._collected_bytes))
+            if event_type == "content_block_start"
+        )
+        return any(
+            isinstance(block, dict)
+            and block.get("type") == "tool_use"
+            and block.get("name") in self._server_fulfilled_tool_names
+            for block in started_blocks
+        )
+
+    @staticmethod
+    async def _settle_task(task: asyncio.Task | None) -> None:
+        if task is None:
+            return
+        if task.done():
+            if not task.cancelled():
+                task.exception()
+            return
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+
     async def aclose(self) -> None:
         from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
             aclose_if_supported,
         )
 
+        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)
 
@@ -217,11 +346,6 @@ async def _process_agentic_hooks(self) -> None:
                 verbose_logger.debug("AgenticStreamingIterator: Could not rebuild response from SSE bytes")
                 return
 
-            [
-                (f"{b.get('type')}({b.get('name', '')})" if b.get("type") == "tool_use" else b.get("type"))
-                for b in rebuilt.get("content", [])
-            ]
-
             result: Final = await self._http_handler._call_agentic_completion_hooks(
                 response=rebuilt,
                 model=self._model,

litellm/llms/custom_httpx/llm_http_handler.py

@@ -2179,6 +2179,10 @@ async def async_anthropic_messages_handler(
                 AgenticAnthropicStreamingIterator,
             )
 
+            held_back_tool_names: Final = self._server_fulfilled_tools_in_request(
+                logging_obj=logging_obj,
+                tools=anthropic_messages_optional_request_params.get("tools"),
+            )
             initial_response = AgenticAnthropicStreamingIterator(
                 completion_stream=completion_stream,
                 http_handler=self,
@@ -2189,6 +2193,8 @@ async def async_anthropic_messages_handler(
                 logging_obj=logging_obj,
                 custom_llm_provider=custom_llm_provider,
                 kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
+                hold_back=bool(held_back_tool_names),
+                server_fulfilled_tool_names=held_back_tool_names,
             )
             return AnthropicMessagesStreamingResponse(
                 completion_stream=initial_response,
@@ -5033,6 +5039,20 @@ def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool:
                 return True
         return False
 
+    @staticmethod
+    def _server_fulfilled_tools_in_request(logging_obj: LiteLLMLoggingObj, tools: object) -> frozenset[str]:
+        """The request's tools that a registered callback fulfills server-side (e.g. ``headroom_retrieve``)."""
+        if not isinstance(tools, list) or not tools:
+            return frozenset()
+        from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name
+
+        return frozenset(
+            name
+            for cb in _custom_logger_callbacks(logging_obj)
+            for name in getattr(cb, "server_fulfilled_tool_names", frozenset())
+            if has_tool_with_name(tools, name)
+        )
+
     @staticmethod
     def _check_agentic_loop_safety(
         tool_calls: object,

litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py

@@ -339,6 +339,7 @@ def _build_responses_followup_items(
 
 class HeadroomGuardrail(CustomGuardrail):
     records_own_guardrail_information: ClassVar[bool] = True
+    server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({HEADROOM_RETRIEVE_TOOL_NAME})
 
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: