fix(simple-engine): stamp finish_reason="stop" on natural-stop streaming epilogue

[mvmories]

Fixes #628

Problem

mlx-lm's stream_generate chunks never claim finished on natural EOS - the generator just exhausts. In _stream_generate_impl, the if not finished: epilogue then yields the close-out chunk with finished=True, finish_reason=None, so streamed chat completions that end naturally never carry "finish_reason":"stop".
Length-capped stops are unaffected (stamped by the in-loop branch), which produced the asymmetry documented in #628.

Strict OpenAI-compatible streaming clients reject such streams as truncated (Stream ended without finish_reason).

Fix

One line: the epilogue is only reachable when the generator exhausted normally after producing tokens - which is a natural stop in OpenAI semantics - so stamp finish_reason="stop".
Exceptions bypass the epilogue (they propagate to finally:), and length-caps exit via the in-loop branch, so the stamp cannot mislabel either case.

Verification

The curl repro from #628 against mlx-community/Qwen3.6-35B-A3B-4bit (SimpleEngine, M4 Max 64GB):

• Before: grep -o '"finish_reason":"[a-z]*"' on the stream → empty
• After: → "finish_reason":"stop"
• Length-capped requests still return "finish_reason":"length"
• A strict streaming client (pi coding agent) now completes multi-turn tool-using sessions without stream errors

[Thump604]

The one-line epilogue correction matches the reported natural-exhaustion path, but this protocol fix needs a regression test before merge.

Please add a focused streaming test that proves:

• generator exhaustion emits one final chunk with finish_reason="stop" before [DONE];
• a max_tokens stop remains finish_reason="length" and is not overwritten by the epilogue;
• an exception does not get mislabeled as a clean stop.

The PR currently uses Fixes #628, but without that executable coverage the issue acceptance criteria are not yet pinned. Use Refs #628 until the regression test is present and passing, or retain Fixes once the test proves the complete reported contract.

[brandy975]

@Thump604 the regression test you asked for is ready — I don't have push access to this branch (fork owned by @mvmories), so posting it here as a patch rather than a commit.

Adds TestSimpleEngineStreamNaturalStopFinishReason to tests/test_simple_engine.py, covering exactly the three cases from your review:

• test_natural_eos_stop_emits_one_final_stop_chunk — generator exhaustion (EOS before max_tokens) yields exactly one finished chunk, stamped finish_reason="stop".
• test_max_tokens_stop_keeps_length_reason — a max_tokens stop keeps finish_reason="length" and isn't overwritten by the epilogue.
• test_exception_mid_stream_is_not_mislabeled_as_stop — a mid-stream exception propagates as-is instead of being reported as a clean stop.

Verified the first test actually catches the regression: reverting just the one-line fix in this PR (finish_reason="stop" → None in the epilogue) makes it fail with AssertionError: assert None == 'stop'; the other two are unaffected by that line, as expected — so it's not a placebo test.

Full suite on this branch (rebased cleanly onto current main, no conflicts): 2250 passed, 23 skipped. Same 3 known-unrelated failures present on main independent of this change (SpecPrefill/residency cancellation tests, root-caused and fixed separately in #643, not yet merged).

diff --git a/tests/test_simple_engine.py b/tests/test_simple_engine.py
index ae5d49e..d771c13 100644
--- a/tests/test_simple_engine.py
+++ b/tests/test_simple_engine.py
@@ -2617,6 +2617,105 @@ class TestSimpleEngineConcurrency:
         assert await asyncio.to_thread(prefill_cancelled.wait, 1.0)
 
 
+class TestSimpleEngineStreamNaturalStopFinishReason:
+    """Regression coverage for the natural-stop streaming epilogue (#628).
+
+    mlx_lm's own GenerationResponse has no ``finished`` field -- only
+    ``finish_reason`` ("stop"/"length"/None) on its one closing chunk.
+    SimpleEngine's loop instead derives ``finished`` from
+    ``completion_tokens >= max_tokens``, which only lines up with mlx_lm's
+    closing chunk when generation actually runs to the token budget. When
+    the model stops early (EOS before max_tokens), that closing chunk is
+    seen as unfinished, the underlying generator is then exhausted, and
+    control falls through to the post-loop epilogue in
+    ``_stream_generate_impl``. Before the fix that epilogue always stamped
+    ``finish_reason=None``, so ordinary EOS-terminated generations reported
+    no finish reason at all.
+    """
+
+    def _make_engine(self, stream_generate_side_effect):
+        from vllm_mlx.engine.simple import SimpleEngine
+
+        with patch("vllm_mlx.engine.simple.is_mllm_model", return_value=False):
+            engine = SimpleEngine("test-model")
+        engine._loaded = True
+        engine._model = MagicMock()
+        engine._model.stream_generate = MagicMock(
+            side_effect=stream_generate_side_effect
+        )
+        return engine
+
+    @pytest.mark.anyio
+    async def test_natural_eos_stop_emits_one_final_stop_chunk(self):
+        """Generator exhaustion (EOS before max_tokens) must yield exactly
+        one finished chunk, stamped finish_reason='stop'."""
+
+        def fake_stream_generate(**kwargs):
+            # Mirrors real mlx_lm chunks: no `finished` attribute at all,
+            # and `finish_reason` only set (non-None) on the closing chunk.
+            yield SimpleNamespace(text="Hel", prompt_tokens=3, finish_reason=None)
+            yield SimpleNamespace(text="lo", prompt_tokens=3, finish_reason="stop")
+
+        engine = self._make_engine(fake_stream_generate)
+
+        outputs = [
+            chunk async for chunk in engine.stream_generate(prompt="hi", max_tokens=50)
+        ]
+
+        finished_chunks = [o for o in outputs if o.finished]
+        assert len(finished_chunks) == 1, (
+            f"expected exactly one finished chunk before [DONE], got "
+            f"{len(finished_chunks)}: {finished_chunks}"
+        )
+        assert finished_chunks[0].finish_reason == "stop"
+
+    @pytest.mark.anyio
+    async def test_max_tokens_stop_keeps_length_reason(self):
+        """Hitting the token budget must report finish_reason='length' and
+        must not be overwritten by the natural-stop epilogue afterwards."""
+
+        def fake_stream_generate(**kwargs):
+            yield SimpleNamespace(text="a", prompt_tokens=3, finish_reason=None)
+            yield SimpleNamespace(text="b", prompt_tokens=3, finish_reason=None)
+            # Closing chunk from mlx_lm when max_tokens is exhausted without
+            # hitting EOS: generation_tokens has now reached max_tokens.
+            yield SimpleNamespace(text="c", prompt_tokens=3, finish_reason="length")
+
+        engine = self._make_engine(fake_stream_generate)
+
+        outputs = [
+            chunk async for chunk in engine.stream_generate(prompt="hi", max_tokens=3)
+        ]
+
+        finished_chunks = [o for o in outputs if o.finished]
+        assert len(finished_chunks) == 1, (
+            "epilogue must not append a second finished chunk once the "
+            f"token budget itself already finished the stream: {finished_chunks}"
+        )
+        assert finished_chunks[0].finish_reason == "length"
+
+    @pytest.mark.anyio
+    async def test_exception_mid_stream_is_not_mislabeled_as_stop(self):
+        """A genuine failure while streaming must propagate as-is, not be
+        swallowed and reported as a clean finish_reason='stop' epilogue."""
+
+        def fake_stream_generate(**kwargs):
+            yield SimpleNamespace(text="partial", prompt_tokens=3, finish_reason=None)
+            raise RuntimeError("backend exploded mid-generation")
+
+        engine = self._make_engine(fake_stream_generate)
+
+        outputs = []
+        with pytest.raises(RuntimeError, match="backend exploded mid-generation"):
+            async for chunk in engine.stream_generate(prompt="hi", max_tokens=50):
+                outputs.append(chunk)
+
+        assert not any(o.finished for o in outputs), (
+            "no chunk should have been reported as a clean finish before "
+            f"the exception propagated: {outputs}"
+        )
+
+
 class TestSimpleEngineClearRuntimeCaches:
     """Operational reset (DELETE /v1/cache) must actually release the
     multi-slot system-prompt KV cache state introduced in the LRU patch —

@mvmories feel free to apply this patch directly (or let me know if you'd rather I open a PR against your branch instead).

[waybarrios]

Added the requested regression coverage in ceee2d2.

It covers natural exhaustion returning stop, max-token termination keeping length, and exceptions propagating without a clean stop. tests/test_simple_engine.py is already included in the Apple Silicon CI jobs.

The three focused tests, Ruff, Black, compilation, and diff checks pass locally.

[waybarrios]

Resolved the conflict with current main in 41260b5. I kept the broader termination coverage from #681 and preserved the final-chunk assertions. All four focused tests, Ruff, Black, and compilation pass locally.