scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

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

Browse source
This commit is contained in:
mateo-berri 2026-08-07 19:23:34 -07:00
parent ecb0ea2f1c
commit e4c2ad4627
8 changed files with 345 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/integrations/compression_interception/handler.py
View file

@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
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,

4
litellm/integrations/custom_logger.py
View file

@ -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 @@ _BASE64_INLINE_PATTERN: Final = re.compile(
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,

59
litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py
View file

@ -6,14 +6,27 @@ yields every chunk to the caller (preserving real streaming), collects
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, with SSE ping events emitted while the upstream message is
in flight. On exhaustion the hooks run first: if a follow-up response
replaces the message, only the follow-up is yielded and the buffered
message is dropped; otherwise the buffer is replayed verbatim. This is
required for server-fulfilled tools (e.g. ``headroom_retrieve``), whose
tool_use blocks must never reach a client that cannot execute them.
"""
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
# ---------------------------------------------------------------------------
# SSE parsing helpers (module-level to keep the class lean)
# ---------------------------------------------------------------------------
@ -156,6 +169,8 @@ class AgenticAnthropicStreamingIterator:
logging_obj: Any,
custom_llm_provider: str,
kwargs: dict,
hold_back: bool = False,
ping_interval_seconds: float = HOLD_BACK_PING_INTERVAL_SECONDS,
):
self._inner = completion_stream.__aiter__()
self._http_handler = http_handler
@ -166,16 +181,23 @@ class AgenticAnthropicStreamingIterator:
self._logging_obj = logging_obj
self._custom_llm_provider = custom_llm_provider
self._kwargs = kwargs
self._hold_back = hold_back
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._replay_index = 0
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 +216,48 @@ class AgenticAnthropicStreamingIterator:
raise StopAsyncIteration
async def _drain_upstream(self) -> None:
try:
while True:
self._collected_bytes.append(await self._inner.__anext__())
except StopAsyncIteration:
return
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
while not self._stream_exhausted:
try:
await asyncio.wait_for(asyncio.shield(self._drain_task), timeout=self._ping_interval_seconds)
except asyncio.TimeoutError:
return PING_SSE_BYTES
self._stream_exhausted = True
await self._process_agentic_hooks()
if self._follow_up_iterator is not None:
return await self._follow_up_iterator.__anext__()
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 aclose(self) -> None:
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
aclose_if_supported,
)
if self._drain_task is not None and self._drain_task.done():
if not self._drain_task.cancelled():
self._drain_task.exception()
elif self._drain_task is not None:
self._drain_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._drain_task
await aclose_if_supported(self._inner)
await aclose_if_supported(self._follow_up_iterator)

23
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -2189,6 +2189,10 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
hold_back=self._should_hold_back_stream(
logging_obj=logging_obj,
tools=anthropic_messages_optional_request_params.get("tools"),
),
)
return AnthropicMessagesStreamingResponse(
completion_stream=initial_response,
@ -5033,6 +5037,25 @@ class BaseLLMHTTPHandler:
return True
return False
@staticmethod
def _should_hold_back_stream(logging_obj: LiteLLMLoggingObj, tools: object) -> bool:
"""
True when the request carries a tool that a registered callback fulfills
server-side (e.g. ``headroom_retrieve``). The model's tool_use for such a
tool must never reach the client, which cannot execute it: the agentic
loop replaces the whole message with a follow-up response, so the stream
is buffered (with ping keepalives) instead of forwarded live.
"""
if not isinstance(tools, list) or not tools:
return False
from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name
return any(
has_tool_with_name(tools, name)
for cb in _custom_logger_callbacks(logging_obj)
for name in getattr(cb, "server_fulfilled_tool_names", frozenset())
)
@staticmethod
def _check_agentic_loop_safety(
tool_calls: object,

1
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -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]:

178
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py
View file

@ -2,6 +2,7 @@
Tests for AgenticAnthropicStreamingIterator and SSE rebuild helpers.
"""
import asyncio
import json
import os
import sys
@ -13,6 +14,7 @@ import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
PING_SSE_BYTES,
AgenticAnthropicStreamingIterator,
_handle_content_block_delta,
_handle_content_block_start,
@ -230,6 +232,51 @@ class MockAsyncStream:
return chunk
class MockSlowAsyncStream(MockAsyncStream):
"""Async iterator that sleeps before every chunk."""
def __init__(self, chunks: List[bytes], delay_seconds: float):
super().__init__(chunks)
self._delay_seconds = delay_seconds
async def __anext__(self) -> bytes:
await asyncio.sleep(self._delay_seconds)
return await super().__anext__()
class MockFailingAsyncStream(MockAsyncStream):
"""Async iterator that raises after yielding its chunks."""
def __init__(self, chunks: List[bytes], error: Exception):
super().__init__(chunks)
self._error = error
async def __anext__(self) -> bytes:
if self._idx >= len(self._chunks):
raise self._error
return await super().__anext__()
def _build_hold_back_iterator(
stream: MockAsyncStream,
mock_handler: MagicMock,
ping_interval_seconds: float = 15.0,
) -> AgenticAnthropicStreamingIterator:
return AgenticAnthropicStreamingIterator(
completion_stream=stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
hold_back=True,
ping_interval_seconds=ping_interval_seconds,
)
# ---------------------------------------------------------------------------
# Tests for _parse_sse_events
# ---------------------------------------------------------------------------
@ -790,3 +837,134 @@ class TestAgenticStreamingIteratorErrorHandling:
call_kwargs = mock_handler._call_agentic_completion_hooks.call_args
assert call_kwargs.kwargs["stream"] is True
class TestAgenticStreamingIteratorHoldBack:
@pytest.mark.asyncio
async def test_should_not_leak_intercepted_message_when_follow_up_fires(self):
"""The buffered tool_use message must be dropped: only pings and follow-up bytes reach the client."""
phase1_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=MockAsyncStream(phase2_chunks))
iterator = _build_hold_back_iterator(MockAsyncStream(phase1_chunks), mock_handler)
collected = []
async for chunk in iterator:
collected.append(chunk)
non_ping = [c for c in collected if c != PING_SSE_BYTES]
assert non_ping == phase2_chunks
assert b"litellm_content_retrieve" not in b"".join(collected)
assert collected[0] == PING_SSE_BYTES
@pytest.mark.asyncio
async def test_should_replay_buffer_verbatim_when_no_hook_fires(self):
"""Without interception the buffered message is replayed byte-identical after the pings."""
chunks = _build_simple_text_stream()
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert [c for c in collected if c != PING_SSE_BYTES] == chunks
mock_handler._call_agentic_completion_hooks.assert_awaited_once()
@pytest.mark.asyncio
async def test_should_emit_pings_while_upstream_is_slow(self):
"""Pings keep the client connection alive while the upstream message is buffered."""
chunks = _build_simple_text_stream()
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = _build_hold_back_iterator(
MockSlowAsyncStream(chunks, delay_seconds=0.05),
mock_handler,
ping_interval_seconds=0.02,
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert collected.count(PING_SSE_BYTES) >= 2
assert [c for c in collected if c != PING_SSE_BYTES] == chunks
@pytest.mark.asyncio
async def test_should_propagate_upstream_error_instead_of_partial_message(self):
"""An upstream failure surfaces as an error; the client never receives a truncated message."""
chunks = _build_simple_text_stream()[:2]
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = _build_hold_back_iterator(
MockFailingAsyncStream(chunks, RuntimeError("upstream died")),
mock_handler,
)
collected = []
with pytest.raises(RuntimeError, match="upstream died"):
async for chunk in iterator:
collected.append(chunk)
assert all(c == PING_SSE_BYTES for c in collected)
mock_handler._call_agentic_completion_hooks.assert_not_awaited()
@pytest.mark.asyncio
async def test_should_replay_buffer_when_hook_processing_errors(self):
"""A hook crash degrades to replaying the original message rather than dropping it."""
chunks = _build_tool_use_stream()
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded"))
mock_logging = MagicMock()
mock_logging.litellm_call_id = "test_call_holdback"
iterator = AgenticAnthropicStreamingIterator(
completion_stream=MockAsyncStream(chunks),
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=mock_logging,
custom_llm_provider="anthropic",
kwargs={},
hold_back=True,
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert [c for c in collected if c != PING_SSE_BYTES] == chunks
@pytest.mark.asyncio
async def test_aclose_cancels_drain_task(self):
"""Closing the iterator mid-buffer must cancel the background drain task."""
chunks = _build_simple_text_stream()
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = _build_hold_back_iterator(
MockSlowAsyncStream(chunks, delay_seconds=5.0),
mock_handler,
)
first = await iterator.__anext__()
assert first == PING_SSE_BYTES
assert iterator._drain_task is not None
await iterator.aclose()
assert iterator._drain_task.cancelled()

71
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -2071,3 +2071,74 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques
retry_authorization = posts[1]["headers"]["Authorization"]
assert retry_authorization.startswith("AWS4-HMAC-SHA256")
assert retry_authorization != first_attempt_headers["Authorization"]
class TestShouldHoldBackStream:
"""_should_hold_back_stream gates the buffered (non-leaking) streaming mode
for server-fulfilled tools like headroom_retrieve."""
@staticmethod
def _logging_obj_with(callbacks):
logging_obj = Mock()
logging_obj.dynamic_success_callbacks = callbacks
return logging_obj
def test_should_hold_back_when_callback_owns_tool_in_request(self):
from litellm.integrations.custom_logger import CustomLogger
class RetrievalCallback(CustomLogger):
server_fulfilled_tool_names = frozenset({"headroom_retrieve"})
tools = [
{"name": "Bash", "input_schema": {"type": "object"}},
{"name": "headroom_retrieve", "input_schema": {"type": "object"}},
]
assert (
BaseLLMHTTPHandler._should_hold_back_stream(
logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools
)
is True
)
def test_should_stream_live_when_tool_absent_from_request(self):
from litellm.integrations.custom_logger import CustomLogger
class RetrievalCallback(CustomLogger):
server_fulfilled_tool_names = frozenset({"headroom_retrieve"})
tools = [{"name": "Bash", "input_schema": {"type": "object"}}]
assert (
BaseLLMHTTPHandler._should_hold_back_stream(
logging_obj=self._logging_obj_with([RetrievalCallback()]), tools=tools
)
is False
)
def test_should_stream_live_when_no_callback_declares_tool_names(self):
from litellm.integrations.custom_logger import CustomLogger
tools = [{"name": "headroom_retrieve", "input_schema": {"type": "object"}}]
assert (
BaseLLMHTTPHandler._should_hold_back_stream(
logging_obj=self._logging_obj_with([CustomLogger()]), tools=tools
)
is False
)
def test_should_stream_live_without_tools(self):
assert BaseLLMHTTPHandler._should_hold_back_stream(logging_obj=self._logging_obj_with([]), tools=None) is False
def test_interception_callbacks_declare_their_retrieval_tools(self):
from litellm.integrations.compression_interception.handler import (
LITELLM_CONTENT_RETRIEVE_TOOL_NAME,
CompressionInterceptionLogger,
)
from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import (
HEADROOM_RETRIEVE_TOOL_NAME,
HeadroomGuardrail,
)
assert HeadroomGuardrail.server_fulfilled_tool_names == frozenset({HEADROOM_RETRIEVE_TOOL_NAME})
assert CompressionInterceptionLogger.server_fulfilled_tool_names == frozenset(
{LITELLM_CONTENT_RETRIEVE_TOOL_NAME}
)

7
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -21391,6 +21391,13 @@ export interface components {
* @description What the routed traffic actually cost
*/
spend: number;
/**
* Tier Turns
* @description Turns per tier, keyed by the tier name the routing decision recorded at request time (never re-derived at read time, since the tier-to-model mapping is mutable config). Tier names are scoped to this group's router_type and are not comparable across types: a complexity router reports 'simple'/'medium'/'complex'/'reasoning', a quality router reports its numeric quality tier, and an adaptive router records no tier at all. Turns no tier served (the classifier fell back to default_model) are absent rather than pooled under a sentinel key, so the values may sum to less than turns
*/
tier_turns?: {
[key: string]: number;
};
/** Turns */
turns: number;
};