From f38bbad108e7af453f1bbc0242b61dca7ba31fcd Mon Sep 17 00:00:00 2001
From: Moe Khalil <moe@berri.ai>
Date: Tue, 22 Sep 2026 03:29:18 +0000
Subject: [PATCH] fix(router): initialize shared Fusion replay stream state

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/fusion_router.py                 | 36 +++++++++---------------
 tests/test_litellm/test_fusion_router.py |  6 ++--
 2 files changed, 18 insertions(+), 24 deletions(-)

diff --git a/litellm/fusion_router.py b/litellm/fusion_router.py
index af7638b0032..14b03dc5b74 100644
--- a/litellm/fusion_router.py
+++ b/litellm/fusion_router.py
@@ -829,26 +829,22 @@ class FusionReplayStream(CustomStreamWrapper):
         chunks: Sequence[ModelResponseStream],
         fusion_metadata: Mapping[str, object],
     ) -> None:
-        # Deliberately do not call CustomStreamWrapper.__init__. The source
-        # wrapper already normalized and logged these chunks while Fusion
-        # buffered them to determine whether its private tool was invoked.
-        source_model: Final = getattr(source, "model", "")
-        self.model = source_model if isinstance(source_model, str) else ""
-        self.custom_llm_provider = source.custom_llm_provider
-        self.logging_obj = source.logging_obj
-        self._hidden_params = dict(  # mutable-ok: local provider payload
-            getattr(
-                source,
-                "_hidden_params",
-                {},  # mutable-ok: stream metadata defaults to a native mapping
-            )  # mutable-ok: local provider payload
-        )  # mutable-ok: local provider payload
-        self._hidden_params["fusion"] = dict(  # mutable-ok: local provider payload
-            fusion_metadata
-        )  # mutable-ok: local provider payload
-        self._source = source
+        super().__init__(
+            completion_stream=source,
+            model=source.model,
+            logging_obj=source.logging_obj,
+            custom_llm_provider=source.custom_llm_provider,
+            stream_options=source.stream_options,
+        )
+        self._hidden_params = {  # mutable-ok: stream consumers attach response metadata
+            **source._hidden_params,
+            "fusion": dict(fusion_metadata),
+        }
         self._iterator = iter(chunks)
 
+    def __next__(self) -> ModelResponseStream:
+        return next(self._iterator)
+
     def __aiter__(self) -> FusionReplayStream:
         return self
 
@@ -858,10 +854,6 @@ class FusionReplayStream(CustomStreamWrapper):
         except StopIteration as exc:
             raise StopAsyncIteration from exc
 
-    async def aclose(self) -> None:
-        if hasattr(self._source, "aclose"):
-            await self._source.aclose()
-
 
 class FusionRouter:
     def __init__(
diff --git a/tests/test_litellm/test_fusion_router.py b/tests/test_litellm/test_fusion_router.py
index 45e121d4077..bccc4427ea6 100644
--- a/tests/test_litellm/test_fusion_router.py
+++ b/tests/test_litellm/test_fusion_router.py
@@ -1118,7 +1118,8 @@ async def test_proxy_fusion_fails_closed_without_authorization_context() -> None
 
 
 @pytest.mark.asyncio
-async def test_router_replays_direct_outer_response_as_an_async_stream() -> None:
+@pytest.mark.parametrize("consume_sync", [True, False])
+async def test_router_replays_direct_outer_response_without_reprocessing(consume_sync: bool) -> None:
     router = Router(model_list=_router_model_list())
     response = await router.acompletion(
         model="fusion/test",
@@ -1127,11 +1128,12 @@ async def test_router_replays_direct_outer_response_as_an_async_stream() -> None
     )
 
     assert isinstance(response, CustomStreamWrapper)
-    chunks = [chunk async for chunk in response]
+    chunks = list(response) if consume_sync else [chunk async for chunk in response]
     rebuilt = litellm.stream_chunk_builder(chunks=chunks)
     assert isinstance(rebuilt, ModelResponse)
     assert rebuilt.choices[0].message.content == "Final"
     assert response._hidden_params["fusion"]["invoked"] is False
+    await response.aclose()
 
 
 @pytest.mark.asyncio