From 809eda2d9973bb59d82423d7823784361f8152b7 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Mon, 22 Jun 2026 17:15:52 -0700
Subject: [PATCH] fix(batches): return original JSONL on unparseable row to
 avoid silent batch truncation

The streaming rewrite of replace_model_in_jsonl accumulated physical lines and
skipped a row on JSONDecodeError to support multi-line objects, but a genuinely
malformed or truncated row never completes: it poisons the buffer, swallows every
following row, and the function still returned the partial rewrite (the rows before
the bad one, already model-rewritten) as if the batch were complete. That turned the
pre-rewrite behavior of returning the original file unchanged (so the provider rejects
the bad batch loudly) into a silent partial submission.

Restore the original-content fallback: when an unparseable remainder is left after the
loop, return the original file_content (rewinding a consumed seekable source) instead of
the truncated output. The multi-line happy path is unchanged.
---
 litellm/router_utils/batch_utils.py           | 10 ++++
 .../test_router_batch_utils.py                | 57 +++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py
index 6e6a17d9e0b..ddec753d362 100644
--- a/litellm/router_utils/batch_utils.py
+++ b/litellm/router_utils/batch_utils.py
@@ -139,9 +139,19 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File
             buffer = ""
 
         if buffer.strip():
+            # A row never parsed (truncated/malformed, or it swallowed the rows
+            # that followed it). Returning the partial `output` would silently
+            # drop those rows; return the unchanged original so the provider
+            # rejects the batch loudly instead of accepting a truncated one.
             verbose_logger.error(
                 f"error parsing trailing batch content: {buffer[:100]}..."
             )
+            if hasattr(source, "seek"):
+                try:
+                    source.seek(0)  # type: ignore[attr-defined]
+                except (OSError, ValueError):
+                    pass
+            return file_content
 
         # If no valid JSON objects were found, return the original content
         if not wrote_any:
diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py
index cba176bd290..b8760906645 100644
--- a/tests/router_unit_tests/test_router_batch_utils.py
+++ b/tests/router_unit_tests/test_router_batch_utils.py
@@ -233,6 +233,63 @@ def test_parse_jsonl_with_embedded_newlines_whitespace_only():
     assert len(result) == 0
 
 
+def test_replace_model_in_jsonl_malformed_middle_row_returns_original():
+    """Regression: a malformed/truncated middle row must not silently drop the
+    rows that follow it. The streaming rewrite accumulates physical lines into a
+    buffer; a row that never parses poisons the buffer so every later valid row
+    is concatenated into it and dropped. Returning that partial rewrite would
+    ship a truncated batch with no error to the caller. Instead the original
+    content is returned unchanged so the provider rejects the bad batch loudly."""
+    content = (
+        b'{"custom_id":"a","body":{"model":"x"}}\n'
+        b'{"custom_id":"b","body":{"model":\n'  # truncated, never completes
+        b'{"custom_id":"c","body":{"model":"x"}}\n'
+    )
+
+    result = replace_model_in_jsonl(content, "new-model")
+
+    assert (
+        result == content
+    ), "must return the original unchanged, not a partial rewrite"
+
+
+def test_replace_model_in_jsonl_malformed_row_seekable_handle_rewound():
+    """When the source is a seekable handle that gets consumed during the failed
+    rewrite, it must be rewound to 0 so the caller can re-read the full original."""
+    content = (
+        b'{"custom_id":"a","body":{"model":"x"}}\n'
+        b'{"custom_id":"b","body":{"model":\n'
+        b'{"custom_id":"c","body":{"model":"x"}}\n'
+    )
+    handle = BytesIO(content)
+
+    result = replace_model_in_jsonl(handle, "new-model")
+
+    assert result is handle
+    assert handle.read() == content, "handle must be rewound for the caller to re-read"
+
+
+def test_replace_model_in_jsonl_multi_row_rewrites_every_model():
+    """Happy path: a well-formed multi-row file gets every row's model rewritten
+    and no row is dropped."""
+    content = (
+        b'{"custom_id":"a","body":{"model":"old1"}}\n'
+        b'{"custom_id":"b","body":{"model":"old2"}}\n'
+        b'{"custom_id":"c","body":{"model":"old3"}}\n'
+    )
+
+    result = replace_model_in_jsonl(content, "new-model")
+
+    assert isinstance(result, InMemoryFile)
+    rows = [
+        json.loads(line)
+        for line in result.getvalue().decode("utf-8").splitlines()
+        if line.strip()
+    ]
+    assert [row["custom_id"] for row in rows] == ["a", "b", "c"]
+    assert all(row["body"]["model"] == "new-model" for row in rows)
+
+
 def test_replace_model_in_jsonl_with_embedded_newlines():
     """Test that replace_model_in_jsonl works correctly with embedded newlines in content"""
     # Create a JSONL with embedded newlines in the message content