Fix local CLI streamed generation error handling (#7135)

studio/backend/core/inference/orchestrator.py

@@ -54,9 +54,8 @@ class GenStreamError(str):
     """A stream chunk carrying a real backend/generation error, not model text.
 
     Subclasses str so existing display/logging consumers are unaffected, while
-    callers that must abort a distributed run on error (raise_on_streamed_error)
-    can distinguish a real error from model output whose visible text starts with
-    "Error:" by checking isinstance(chunk, GenStreamError).
+    callers can distinguish a real error from model output whose visible text
+    starts with "Error:" by checking isinstance(chunk, GenStreamError).
     """
 
     __slots__ = ("public",)

unsloth_cli/_inference.py

@@ -235,7 +235,7 @@ def collect_stream(stream, show_thinking: bool) -> str:
 def raise_on_streamed_error(stream):
     # Match real backend errors by type (GenStreamError), not the "Error:" text
     # prefix, so a completion whose text opens with "Error:" is not misread as a
-    # failure that aborts a distributed run.
+    # backend failure.
     try:
         ensure_studio_backend_path()
         from core.inference.orchestrator import GenStreamError

unsloth_cli/commands/chat.py

@@ -331,7 +331,7 @@ def chat(
             enable_thinking = show_thinking,
             use_adapter = use_adapter,
         )
-        return raise_on_streamed_error(stream) if is_mlx_distributed else stream
+        return raise_on_streamed_error(stream)
 
     if should_print:
         console.print()

unsloth_cli/commands/inference.py

@@ -111,15 +111,12 @@ def inference(
             repetition_penalty = repetition_penalty,
             enable_thinking = think,
         )
-        if is_mlx_distributed:
-            stream = raise_on_streamed_error(stream)
+        stream = raise_on_streamed_error(stream)
         if rank == 0:
             typer.echo("Assistant:")
             try:
                 stream_to_stdout(stream, show_thinking = think)
             except RuntimeError as exc:
-                if not is_mlx_distributed:
-                    raise
                 typer.echo(f"Error: {exc}", err = True)
                 raise typer.Exit(code = 1)
         else:

unsloth_cli/tests/test_inference_chat.py

@@ -873,6 +873,104 @@ def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch):
     assert set(closed) == {"tuned", "base"}
 
 
+@pytest.mark.parametrize(
+    ("chunk_kind", "expected_exit"),
+    [
+        ("answer", 0),
+        ("model_text_error", 0),
+        ("real_error", 1),
+    ],
+)
+def test_inference_local_handles_stream(monkeypatch, chunk_kind, expected_exit):
+    from unsloth_cli.commands import inference as infermod
+    from unsloth_cli._inference import ensure_studio_backend_path
+
+    ensure_studio_backend_path()
+    from core.inference.orchestrator import GenStreamError
+
+    chunks = {
+        "answer": ["answer"],
+        "model_text_error": ["Error: printed by the model, not a backend failure"],
+        "real_error": [GenStreamError("Error: generation failed")],
+    }[chunk_kind]
+    closed = []
+
+    class _FakeBackend:
+        def stream(self, messages, **kwargs):
+            return iter(chunks)
+
+        def close(self):
+            closed.append(True)
+
+    monkeypatch.setattr(
+        infermod,
+        "connect_studio_server",
+        lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
+    )
+    monkeypatch.setattr(infermod, "load_chat_backend", lambda *a, **k: _FakeBackend())
+
+    result = CliRunner().invoke(
+        _inference_app(),
+        ["fake-model", "hello", "--no-server"],
+    )
+
+    assert result.exit_code == expected_exit, result.output
+    assert closed == [True]
+    if chunk_kind == "real_error":
+        assert result.stdout == "Assistant:\n"
+        assert result.stderr == "Error: generation failed\n"
+    else:
+        assert chunks[0] in result.output
+
+
+@pytest.mark.parametrize("chunk_kind", ["answer", "model_text_error", "real_error"])
+def test_chat_local_handles_stream(monkeypatch, chunk_kind):
+    from unsloth_cli._inference import ensure_studio_backend_path
+
+    ensure_studio_backend_path()
+    from core.inference.orchestrator import GenStreamError
+
+    first_chunk = {
+        "answer": "answer",
+        "model_text_error": "Error: printed by the model, not a backend failure",
+        "real_error": GenStreamError("Error: generation failed"),
+    }[chunk_kind]
+    calls, closed = [], []
+
+    class _FakeChatBackend:
+        def stream(self, messages, **kwargs):
+            calls.append([dict(message) for message in messages])
+            return iter([first_chunk if len(calls) == 1 else "second answer"])
+
+        def close(self):
+            closed.append(True)
+
+    monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
+    monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
+    monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
+    monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
+
+    result = CliRunner().invoke(
+        _chat_app(),
+        ["fake-model"],
+        input = "first\nsecond\n/exit\n",
+    )
+
+    assert result.exit_code == 0, result.output
+    assert closed == [True]
+    if chunk_kind == "real_error":
+        assert calls[1] == [{"role": "user", "content": "second"}]
+        assert "(error: generation failed)" in result.output
+        assert "Error: generation failed" not in result.output
+    else:
+        assert calls[1] == [
+            {"role": "user", "content": "first"},
+            {"role": "assistant", "content": first_chunk},
+            {"role": "user", "content": "second"},
+        ]
+        assert first_chunk in result.output
+
+
 @pytest.mark.parametrize(
     ("chunk_kind", "expected_exit"),
     [