From d2ab01a88bba980da53fbda20d084845a5172d9a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 1 Oct 2026 18:11:49 -0700
Subject: [PATCH] fix(router): keep content blocks on a trailing assistant turn
 when extending it with the partial output

A trailing assistant message whose content is a list of blocks was stringified before the partial output was appended, so the fallback received the Python repr of the blocks as its prefill. The partial now lands as one more text block, and a content shape that cannot be extended surfaces the original stream error instead of a garbled continuation.
---
 litellm/router.py                     | 32 +++++++++++++++++----------
 tests/unit/test_router/test_router.py | 15 +++++++++++++
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index 2d3d6971e3b..bb8485e0ef5 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -475,6 +475,7 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType
 _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"})
 _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
+_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
 
 
 def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
@@ -2997,12 +2998,13 @@ class Router:
                             getattr(complete_response_object_usage, "completion_tokens", 0) or 0
                         )
                         reduced_ceilings: Final = self._continuation_output_ceilings(initial_kwargs, emitted_tokens)
-                        if reduced_ceilings is None:
-                            self._raise_original_mid_stream_error(e)
-                        initial_kwargs.update(reduced_ceilings)
-                        initial_kwargs["messages"] = self._build_completion_continuation_input(
+                        continuation_messages: Final = self._build_completion_continuation_input(
                             messages, e.generated_content
                         )
+                        if reduced_ceilings is None or continuation_messages is None:
+                            self._raise_original_mid_stream_error(e)
+                        initial_kwargs.update(reduced_ceilings)
+                        initial_kwargs["messages"] = continuation_messages
                         initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER
                     else:
                         initial_kwargs["messages"] = messages
@@ -3243,17 +3245,23 @@ class Router:
 
     @staticmethod
     def _build_completion_continuation_input(
-        messages: Sequence[Mapping[str, str]],
+        messages: Sequence[Mapping[str, object]],
         generated_content: str,
-    ) -> Sequence[Mapping[str, object]]:
+    ) -> Sequence[Mapping[str, object]] | None:
         """Append the partial output as an assistant prefill, or extend a trailing assistant turn in place so the
-        request never ends in two assistant messages."""
+        request never ends in two assistant messages. None when that turn's content has a shape this cannot extend."""
         last: Final = messages[-1] if messages else None
-        if last is not None and last.get("role") == "assistant":
-            merged: Final = {**last, "content": str(last.get("content") or "") + generated_content, "prefix": True}
-            return [*messages[:-1], merged]
-        prefill: Final = {"role": "assistant", "content": generated_content, "prefix": True}
-        return [*messages, prefill]
+        if last is None or last.get("role") != "assistant":
+            return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}]
+        content: Final = last.get("content")
+        if content is None or isinstance(content, str):
+            return [*messages[:-1], {**last, "content": f"{content or ''}{generated_content}", "prefix": True}]
+        try:
+            blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content)
+        except ValidationError:
+            return None
+        extended: Final = [*blocks, {"type": "text", "text": generated_content}]
+        return [*messages[:-1], {**last, "content": extended, "prefix": True}]
 
     @staticmethod
     def _build_responses_continuation_input(
diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py
index d59674f1c57..3d592737e1e 100644
--- a/tests/unit/test_router/test_router.py
+++ b/tests/unit/test_router/test_router.py
@@ -2901,6 +2901,21 @@ def test_build_completion_continuation_input_folds_into_trailing_plain_assistant
     assert built[-1] == {"role": "assistant", "content": "Sure, here it is", "prefix": True}
 
 
+def test_build_completion_continuation_input_keeps_structured_assistant_content():
+    """Content blocks on a trailing assistant turn stay blocks: the partial lands as one more text
+    block instead of a stringified list, and a content shape that cannot be extended declines."""
+    messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}]}]
+    built = litellm.Router._build_completion_continuation_input(messages, "here it is")
+    assert built is not None
+    assert [m["role"] for m in built] == ["user", "assistant"]
+    assert built[-1] == {
+        "role": "assistant",
+        "content": [{"type": "text", "text": "Sure, "}, {"type": "text", "text": "here it is"}],
+        "prefix": True,
+    }
+    assert litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42}], "x") is None
+
+
 def test_continuation_output_ceilings_reduces_by_emitted_tokens():
     """A continuation must complete within the caller's original allowance, so each
     output ceiling is reduced by the tokens already emitted."""