From 990cda80f210d24baa6f715f43530a86d13a54bb Mon Sep 17 00:00:00 2001 From: Shin Date: Sun, 22 Feb 2026 04:31:47 +0000 Subject: [PATCH] fix(websearch_interception): Convert agentic loop response to streaming format for Claude Code Fixes #20187 - When using websearch_interception in Bedrock with Claude Code: 1. Output tokens were showing as 0 because the agentic loop response wasn't being converted back to streaming format 2. The response from the agentic loop (follow-up request) was returned as a non-streaming dict, but Claude Code expects a streaming response This fix adds streaming format conversion for the agentic loop response when the original request was streaming (detected via the websearch_interception_converted_stream flag in logging_obj). The fix applies to both: - Anthropic Messages API (_call_agentic_completion_hooks) - Chat Completions API (_call_agentic_chat_completion_hooks) The fix ensures: - Output tokens are correctly included in the message_delta event - stop_reason is properly preserved - The response format matches what Claude Code expects Note: This fix was previously in PR #20631 but was merged to a staging branch (litellm_oss_staging_02_07_2026) and never made it to main. --- litellm/llms/custom_httpx/llm_http_handler.py | 55 +++++++++++ .../test_websearch_interception_handler.py | 99 +++++++++++++++++++ 2 files changed, 154 insertions(+) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 0a5364bfcfe..011d0ebe5d2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4445,6 +4445,37 @@ class BaseLLMHTTPHandler: stream=stream, kwargs=kwargs_with_provider, ) + + # Check if we need to convert agentic loop response to fake stream + # This happens when the original request was streaming but was converted + # to non-streaming for WebSearch interception, and the agentic loop ran + 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 agentic_response is not None: + 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: Agentic loop completed, converting response to fake stream" + ) + + # Convert the non-streaming response to a fake stream + if isinstance(agentic_response, dict): + fake_stream = FakeAnthropicMessagesStreamIterator( + response=cast(AnthropicMessagesResponse, agentic_response) + ) + return fake_stream + # First hook that runs agentic loop wins return agentic_response @@ -4552,6 +4583,30 @@ class BaseLLMHTTPHandler: stream=stream, kwargs=kwargs_with_provider, ) + + # Check if we need to convert agentic loop response to fake stream + # This happens when the original request was streaming but was converted + # to non-streaming for WebSearch interception, and the agentic loop ran + 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 agentic_response is not None: + from litellm.llms.base_llm.base_model_iterator import ( + convert_model_response_to_streaming, + ) + + verbose_logger.debug( + "WebSearchInterception: Chat completion agentic loop completed, converting response to fake stream" + ) + + # Convert the non-streaming ModelResponse to a fake stream + if hasattr(agentic_response, "choices"): + fake_stream = convert_model_response_to_streaming(agentic_response) + return fake_stream + # First hook that runs agentic loop wins return agentic_response diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index 85ed5dddee1..d711a65c812 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -273,3 +273,102 @@ async def test_async_pre_call_deployment_hook_provider_derived_from_model_name() # Full kwargs preserved assert result["model"] == "openai/gpt-4o-mini" assert result["api_key"] == "fake-key" + + +def test_fake_anthropic_messages_stream_iterator_includes_output_tokens(): + """Test that FakeAnthropicMessagesStreamIterator includes output_tokens in message_delta. + + Regression test for issue #20187 - Output tokens showing as 0 in Claude Code session logs. + The FakeAnthropicMessagesStreamIterator must include output_tokens in the message_delta event. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + + # Create a sample non-streaming response + response = { + "id": "msg_test123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [ + {"type": "text", "text": "Here is the response to your query."} + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 100, + "output_tokens": 50 + } + } + + # Create the fake stream iterator + iterator = FakeAnthropicMessagesStreamIterator(response=response) + + # Collect all chunks + chunks = list(iterator) + + # Find the message_delta event + message_delta_found = False + output_tokens_in_delta = None + for chunk in chunks: + chunk_str = chunk.decode("utf-8") + if "message_delta" in chunk_str: + message_delta_found = True + import json + # Parse the data line + for line in chunk_str.split("\n"): + if line.startswith("data: "): + data = json.loads(line[6:]) + output_tokens_in_delta = data.get("usage", {}).get("output_tokens") + break + + assert message_delta_found, "message_delta event not found in fake stream" + assert output_tokens_in_delta == 50, f"Expected output_tokens=50, got {output_tokens_in_delta}" + + +def test_fake_anthropic_messages_stream_iterator_preserves_stop_reason(): + """Test that FakeAnthropicMessagesStreamIterator preserves stop_reason in message_delta. + + Regression test for issue #20187 - The stop_reason must be included in the message_delta event. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + + # Create a sample non-streaming response + response = { + "id": "msg_test456", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [ + {"type": "text", "text": "Response text."} + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 50, + "output_tokens": 25 + } + } + + # Create the fake stream iterator + iterator = FakeAnthropicMessagesStreamIterator(response=response) + + # Collect all chunks + chunks = list(iterator) + + # Find the message_delta event and verify stop_reason + stop_reason_found = None + for chunk in chunks: + chunk_str = chunk.decode("utf-8") + if "message_delta" in chunk_str: + import json + for line in chunk_str.split("\n"): + if line.startswith("data: "): + data = json.loads(line[6:]) + stop_reason_found = data.get("delta", {}).get("stop_reason") + break + + assert stop_reason_found == "end_turn", f"Expected stop_reason='end_turn', got {stop_reason_found}"