From 2a679d2d151e9bc756c67e3f319cc8d1e44aa58c Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Sun, 1 Feb 2026 01:46:15 +0000 Subject: [PATCH] fix(websearch): convert agentic response to fake stream when stream was converted Fixes #20187 When using websearch_interception with Claude Code and stream=True: 1. WebSearch interception converts stream=True to stream=False for the initial request 2. The agentic loop executes the search and makes a follow-up request 3. The follow-up response has the correct output_tokens Previously, the agentic response was returned as-is (non-streaming), but Claude Code expected a streaming response. This caused the usage to show 0 output tokens because the fake stream conversion was only applied when no agentic loop ran. This fix checks if websearch_interception_converted_stream is True after the agentic loop runs, and if so, converts the agentic response to a FakeAnthropicMessagesStreamIterator that preserves the correct output tokens from the follow-up request. --- litellm/llms/custom_httpx/llm_http_handler.py | 29 +++ .../test_websearch_stream_conversion.py | 228 ++++++++++++++++++ 2 files changed, 257 insertions(+) create mode 100644 tests/test_litellm/integrations/websearch_interception/test_websearch_stream_conversion.py diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d2ea7e872a2..89ba6baea03 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4408,6 +4408,35 @@ class BaseLLMHTTPHandler: stream=stream, kwargs=kwargs_with_provider, ) + + # Check if we need to convert agentic response to fake stream + # This happens when stream was converted from True to False for interception + # Fixes https://github.com/BerriAI/litellm/issues/20187 + websearch_converted_stream = ( + logging_obj.model_call_details.get("websearch_interception_converted_stream", False) + if logging_obj is not None + else False + ) + + if websearch_converted_stream and isinstance(agentic_response, dict): + from typing import cast + + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + ) + + verbose_logger.debug( + "WebSearchInterception: Converting agentic response to fake stream" + ) + + # Convert the non-streaming agentic response to a fake stream + return FakeAnthropicMessagesStreamIterator( + response=cast(AnthropicMessagesResponse, agentic_response) + ) + # First hook that runs agentic loop wins return agentic_response diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_stream_conversion.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_stream_conversion.py new file mode 100644 index 00000000000..558e70acdd6 --- /dev/null +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_stream_conversion.py @@ -0,0 +1,228 @@ +""" +Test that websearch interception properly converts agentic response to fake stream. + +Fixes https://github.com/BerriAI/litellm/issues/20187 + +The issue was that when: +1. Client requested stream=True +2. WebSearch interception converted stream to False +3. Agentic loop ran and returned a non-streaming response with correct usage +4. The response was returned without conversion to streaming format + +This caused Claude Code to receive a response with 0 output tokens instead of +the actual output tokens from the follow-up request. +""" + +import json +import pytest +from typing import cast +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, +) +from litellm.integrations.custom_logger import CustomLogger + + +class MockCustomCallback(CustomLogger): + """Mock callback that inherits from CustomLogger for proper isinstance checks""" + + def __init__(self, agentic_response): + super().__init__() + self.agentic_response = agentic_response + + async def async_should_run_agentic_loop(self, response, model, messages, tools, stream, custom_llm_provider, kwargs): + return True, {"tool_calls": [{"id": "tc_123"}]} + + async def async_run_agentic_loop(self, tools, model, messages, response, anthropic_messages_provider_config, anthropic_messages_optional_request_params, logging_obj, stream, kwargs): + return self.agentic_response + + +class TestWebsearchStreamConversion: + """Tests for websearch interception stream conversion""" + + @pytest.mark.asyncio + async def test_agentic_response_converted_to_fake_stream_when_stream_was_converted(self): + """ + Test that when websearch_interception_converted_stream is True, + the agentic response is converted to a fake stream. + """ + # Create a mock agentic response with actual output tokens + mock_agentic_response = { + "id": "msg_test123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [ + { + "type": "text", + "text": "This is the response from the follow-up request with search results." + } + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 38196, + "output_tokens": 395 # This should be preserved in the fake stream + } + } + + # Create a mock CustomLogger that simulates websearch interception + mock_callback = MockCustomCallback(mock_agentic_response) + + # Create a mock logging object with websearch_interception_converted_stream=True + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = { + "websearch_interception_converted_stream": True + } + mock_logging_obj.dynamic_success_callbacks = [] + + # Create the handler + handler = BaseLLMHTTPHandler() + + # Patch litellm.callbacks to include our mock callback + with patch("litellm.callbacks", [mock_callback]): + result = await handler._call_agentic_completion_hooks( + response={"content": [{"type": "tool_use", "name": "WebSearch"}]}, # Initial response with tool_use + model="claude-sonnet-4-5-20250929", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=mock_logging_obj, + stream=False, + custom_llm_provider="bedrock", + kwargs={}, + ) + + # The result should be a FakeAnthropicMessagesStreamIterator + assert isinstance(result, FakeAnthropicMessagesStreamIterator) + + # Verify the fake stream contains the correct output tokens + chunks = list(result) + + # Find the message_delta chunk which contains output_tokens + message_delta_chunk = None + for chunk in chunks: + chunk_str = chunk.decode() if isinstance(chunk, bytes) else chunk + if "message_delta" in chunk_str: + # Parse the event data + for line in chunk_str.split('\n'): + if line.startswith('data: '): + data = json.loads(line[6:]) + if data.get("type") == "message_delta": + message_delta_chunk = data + break + + assert message_delta_chunk is not None + assert message_delta_chunk["usage"]["output_tokens"] == 395 + + @pytest.mark.asyncio + async def test_agentic_response_not_converted_when_stream_was_not_converted(self): + """ + Test that when websearch_interception_converted_stream is False, + the agentic response is returned as-is (not converted to fake stream). + """ + mock_agentic_response = { + "id": "msg_test123", + "content": [{"type": "text", "text": "Response"}], + "usage": {"output_tokens": 100} + } + + mock_callback = MockCustomCallback(mock_agentic_response) + + # websearch_interception_converted_stream is False + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = { + "websearch_interception_converted_stream": False + } + mock_logging_obj.dynamic_success_callbacks = [] + + handler = BaseLLMHTTPHandler() + + with patch("litellm.callbacks", [mock_callback]): + result = await handler._call_agentic_completion_hooks( + response={}, + model="test-model", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=mock_logging_obj, + stream=False, + custom_llm_provider="bedrock", + kwargs={}, + ) + + # The result should be the dict response, not a FakeAnthropicMessagesStreamIterator + assert result == mock_agentic_response + assert not isinstance(result, FakeAnthropicMessagesStreamIterator) + + +class TestFakeStreamIteratorUsage: + """Tests for FakeAnthropicMessagesStreamIterator usage handling""" + + def test_fake_stream_preserves_output_tokens(self): + """Test that FakeAnthropicMessagesStreamIterator preserves output_tokens in the stream""" + response = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-test", + "content": [{"type": "text", "text": "Hello world"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 100, + "output_tokens": 50 + } + } + + iterator = FakeAnthropicMessagesStreamIterator(response=response) + chunks = list(iterator) + + # Find message_delta chunk + output_tokens_found = None + for chunk in chunks: + chunk_str = chunk.decode() if isinstance(chunk, bytes) else chunk + if "message_delta" in chunk_str: + for line in chunk_str.split('\n'): + if line.startswith('data: '): + data = json.loads(line[6:]) + if data.get("type") == "message_delta": + output_tokens_found = data["usage"]["output_tokens"] + break + + assert output_tokens_found == 50 + + def test_fake_stream_preserves_input_tokens_in_message_start(self): + """Test that FakeAnthropicMessagesStreamIterator preserves input_tokens in message_start""" + response = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-test", + "content": [{"type": "text", "text": "Hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 200, + "output_tokens": 30 + } + } + + iterator = FakeAnthropicMessagesStreamIterator(response=response) + chunks = list(iterator) + + # Find message_start chunk + input_tokens_found = None + for chunk in chunks: + chunk_str = chunk.decode() if isinstance(chunk, bytes) else chunk + if "message_start" in chunk_str: + for line in chunk_str.split('\n'): + if line.startswith('data: '): + data = json.loads(line[6:]) + if data.get("type") == "message_start": + input_tokens_found = data["message"]["usage"]["input_tokens"] + break + + assert input_tokens_found == 200