From 58753a268f88ca24bedb58269f75c218e92f5256 Mon Sep 17 00:00:00 2001
From: Rudra Dudhat <contact.rdudhat@gmail.com>
Date: Tue, 15 Sep 2026 23:15:41 +0530
Subject: [PATCH] Fix LineIterator hanging forever when the stream ends without
 a trailing newline

When the underlying event stream is exhausted (StopIteration) while the
buffer still has an unread trailing partial line (no "\n" at the end,
a normal case since the model's last streamed token has no reason to
be a newline), __next__ just did `continue`. That re-reads the exact
same unterminated line, fails the same check, and calls next() on an
already-exhausted iterator again, forever. Nothing in that branch ever
changes state, so it spins indefinitely instead of returning or
raising.

Return the leftover partial data once instead, so the following call
correctly hits the read_pos == buffer size case and raises
StopIteration normally.

Added test_line_iterator_no_trailing_newline_at_end since the existing
test_line_iterator_incomplete_line_at_end doesn't actually cover this,
it only tests a stream that ends cleanly with a trailing newline.
---
 .../src/sagemaker/core/iterators.py           | 10 +++++++-
 sagemaker-core/tests/unit/test_iterators.py   | 23 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/sagemaker-core/src/sagemaker/core/iterators.py b/sagemaker-core/src/sagemaker/core/iterators.py
index 60914cbdd0..37f39fce6f 100644
--- a/sagemaker-core/src/sagemaker/core/iterators.py
+++ b/sagemaker-core/src/sagemaker/core/iterators.py
@@ -175,7 +175,15 @@ def __next__(self):
                 chunk = next(self.byte_iterator)
             except StopIteration:
                 if self.read_pos < self.buffer.getbuffer().nbytes:
-                    continue
+                    # Stream ended with a trailing partial line (no "\n").
+                    # `continue` alone would spin forever here: byte_iterator
+                    # is already exhausted, so it keeps raising StopIteration
+                    # and read_pos/buffer never change. Return the remainder
+                    # once, so the next call correctly raises StopIteration.
+                    self.buffer.seek(self.read_pos)
+                    remainder = self.buffer.read()
+                    self.read_pos += len(remainder)
+                    return remainder
                 raise
             if "PayloadPart" not in chunk:
                 # handle API response errors and force terminate.
diff --git a/sagemaker-core/tests/unit/test_iterators.py b/sagemaker-core/tests/unit/test_iterators.py
index 02ed29e2be..3bbefda59e 100644
--- a/sagemaker-core/tests/unit/test_iterators.py
+++ b/sagemaker-core/tests/unit/test_iterators.py
@@ -232,3 +232,26 @@ def test_line_iterator_incomplete_line_at_end():
     # After consuming all complete lines, should raise StopIteration
     with pytest.raises(StopIteration):
         next(iterator)
+
+
+def test_line_iterator_no_trailing_newline_at_end():
+    """Test LineIterator returns the final chunk even if it has no trailing "\n".
+
+    Regression test: previously this hung forever instead of returning or
+    raising, because the StopIteration handler for a leftover partial line
+    just did `continue`, re-reading the same unterminated line and calling
+    next() on an already-exhausted iterator every time.
+    """
+    mock_stream = [
+        {"PayloadPart": {"Bytes": b'{"outputs": [" first"]}\n'}},
+        {"PayloadPart": {"Bytes": b'{"outputs": [" second"]}'}},  # no trailing \n
+    ]
+    iterator = LineIterator(mock_stream)
+
+    line1 = next(iterator)
+    line2 = next(iterator)
+    assert line1 == b'{"outputs": [" first"]}'
+    assert line2 == b'{"outputs": [" second"]}'
+
+    with pytest.raises(StopIteration):
+        next(iterator)