From 023d437685e020e0fec81ffa3c679fcee847cef6 Mon Sep 17 00:00:00 2001
From: Clawd <fffff.c@gmail.com>
Date: Fri, 8 May 2026 15:05:29 +0800
Subject: [PATCH] fix: also wrap plan.response_override and plan.terminate
 paths, add unit tests

Address review feedback:
- Wrap plan.response_override and plan.terminate early-return paths
  with _maybe_wrap_in_fake_stream() to ensure all exit points handle
  the streaming conversion correctly
- Add unit tests for _maybe_wrap_in_fake_stream() covering:
  - Dict response with flag=True -> wrapped in FakeAnthropicMessagesStreamIterator
  - Flag=False -> response returned unchanged
  - Non-dict response -> returned unchanged
  - logging_obj=None -> response returned unchanged
---
 litellm/llms/custom_httpx/llm_http_handler.py |  8 ++-
 .../test_websearch_streaming_wrap.py          | 71 +++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py

diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index 790eedaef2f..aaba41da978 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -4885,14 +4885,18 @@ class BaseLLMHTTPHandler:
                 )
 
                 if plan.response_override is not None:
-                    return plan.response_override
+                    return self._maybe_wrap_in_fake_stream(
+                        plan.response_override, logging_obj
+                    )
                 if plan.terminate:
                     verbose_logger.debug(
                         "Agentic loop terminated by callback=%s reason=%s",
                         callback.__class__.__name__,
                         plan.stop_reason,
                     )
-                    return response
+                    return self._maybe_wrap_in_fake_stream(
+                        response, logging_obj
+                    )
                 if not plan.run_agentic_loop:
                     continue
 
diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py
new file mode 100644
index 00000000000..b250e852a47
--- /dev/null
+++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py
@@ -0,0 +1,71 @@
+"""
+Unit tests for _maybe_wrap_in_fake_stream in BaseLLMHTTPHandler.
+
+Tests that agentic loop responses are correctly wrapped in
+FakeAnthropicMessagesStreamIterator when the original request was streaming.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
+    FakeAnthropicMessagesStreamIterator,
+)
+
+
+class TestMaybeWrapInFakeStream:
+    def setup_method(self):
+        self.handler = BaseLLMHTTPHandler()
+
+    def test_wraps_dict_when_converted_stream_flag_is_true(self):
+        """When websearch_interception_converted_stream is True and response is dict, wrap it."""
+        logging_obj = MagicMock()
+        logging_obj.model_call_details = {
+            "websearch_interception_converted_stream": True
+        }
+        response = {
+            "id": "msg_123",
+            "type": "message",
+            "role": "assistant",
+            "content": [{"type": "text", "text": "hello"}],
+            "stop_reason": "end_turn",
+            "usage": {"input_tokens": 10, "output_tokens": 5},
+        }
+
+        result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)
+
+        assert isinstance(result, FakeAnthropicMessagesStreamIterator)
+
+    def test_returns_response_unchanged_when_flag_is_false(self):
+        """When flag is False, return response as-is."""
+        logging_obj = MagicMock()
+        logging_obj.model_call_details = {
+            "websearch_interception_converted_stream": False
+        }
+        response = {"id": "msg_123", "content": []}
+
+        result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)
+
+        assert result is response
+
+    def test_returns_response_unchanged_when_not_dict(self):
+        """When response is not a dict (e.g., already a stream), return as-is."""
+        logging_obj = MagicMock()
+        logging_obj.model_call_details = {
+            "websearch_interception_converted_stream": True
+        }
+        response = MagicMock()  # Not a dict
+
+        result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)
+
+        assert result is response
+
+    def test_returns_response_when_logging_obj_is_none(self):
+        """When logging_obj is None, return response as-is."""
+        response = {"id": "msg_123", "content": []}
+
+        result = self.handler._maybe_wrap_in_fake_stream(response, None)
+
+        assert result is response