feat: sglang-omni model picker on manual server launch

1 files changed, 135 insertions, 25 deletions

diff --git a/app/ui/hub.py b/app/ui/hub.py
index f9151f8..faf2291 100644
--- a/app/ui/hub.py
+++ b/app/ui/hub.py
@@ -534,9 +534,12 @@ class _Hub:
         toggles it directly (starts a stopped server, stops a running one)
         without an extra action menu. The state lives in the table, not on
         the entries, because the menu's selection bar would cover inline
-        colors. Each server is labelled by its backend's name; qwen hosts
-        one model at a time (its default start runs CustomVoice — a
-        Generate-audiobooks run needing another model restarts it).
+        colors. Each server is labelled by its backend's name; the
+        one-model-per-port backends (qwen-tts, SGLang-Omni) ask which model
+        to load on a fresh start (a Generate-audiobooks run needing another
+        model restarts the server with its own pick), and an SGLang-Omni
+        backend without downloaded models stays listed so selecting it can
+        say so instead of the backend silently missing from the menu.
         """
         statuses = detect_all()
         candidates = [st for st in statuses if st.installed]
@@ -548,11 +551,15 @@ class _Hub:
         rows = []
         for st in candidates:
             specs = st.servers
+            if st.key == BACKEND_SGLOMNI and not specs and not st.remote:
+                options.append((st.label, (st.key, _SGLOMNI_NO_MODELS)))
+                rows.append((st.label, "no models", "err", "body"))
+                continue
             for spec in specs:
                 running = common.server_running(spec.url)
                 label = st.label if len(specs) == 1 \
                     else f"{st.label} — {spec.name}"
-                options.append((label, spec))
+                options.append((label, (st.key, spec)))
                 rows.append((label,
                              "running" if running else "stopped",
                              "ok" if running else "err", "body"))
@@ -560,31 +567,48 @@ class _Hub:
             tui.flash(self.stdscr, "No backend server is configured yet — "
                       "use 'Configure Backends' first.")
             return tui.Wizard.BACK
-        spec = tui.menu(self.stdscr, "Start / Stop A Server", options,
-                        back_value=tui.Wizard.BACK,
-                        help_lines=["Start/stop local servers manually.",
-                                    "'Generate Audiobooks' handles this "
-                                    "automatically."],
-                        table_rows=rows,
-                        notice_lines=_notice_lines())
-        if spec is tui.Wizard.BACK:
+        chosen = tui.menu(self.stdscr, "Start / Stop A Server", options,
+                          back_value=tui.Wizard.BACK,
+                          help_lines=["Start/stop local servers manually.",
+                                      "'Generate Audiobooks' handles this "
+                                      "automatically."],
+                          table_rows=rows,
+                          notice_lines=_notice_lines())
+        if chosen is tui.Wizard.BACK:
             return tui.Wizard.BACK
-        return functools.partial(self._server_toggle, spec)
+        status_key, target = chosen
+        if target is _SGLOMNI_NO_MODELS:
+            def no_models():
+                tui.flash(self.stdscr, "No SGLang-Omni models are "
+                          "downloaded — install one via Configure "
+                          "Backends → SGLang-Omni (Configure).", "err")
+                return tui.Wizard.BACK
+            return no_models
+        return functools.partial(self._server_toggle, status_key, target)
 
-    def _server_toggle(self, spec):
+    def _server_toggle(self, status_key, spec):
         """Start SPEC's server when stopped, stop it when running.
 
-        Runs inside the task view (no console drop); the server module's
-        plain-console output is tee'd to a log file under ``app/logs`` so
-        nothing is lost, and on failure a flash points the user at that
-        file. Returns BACK so the stack lands back on the server list,
-        which re-reads each server's live state.
+        Starting one of the one-model-per-port backends (qwen-tts,
+        SGLang-Omni) asks which model to load first — their detect() specs
+        aim at a default the menu must not silently boot (see
+        _pick_start_model). Runs inside the task view (no console drop);
+        the server module's plain-console output is tee'd to a log file
+        under ``app/logs`` so nothing is lost, and on failure a flash
+        points the user at that file. Returns BACK so the stack lands back
+        on the server list, which re-reads each server's live state.
         """
         running = common.server_running(spec.url)
         action = "stop" if running else "start"
-        step, log_path = _server_action_step(spec, action)
+        if action == "start":
+            picked = _pick_start_model(self.stdscr, status_key, spec)
+            if picked is None:
+                return tui.Wizard.BACK
+            spec = picked
+        step, log_path = _server_action_step(
+            spec, action, prep=_companion_prep(spec, action))
         taskview.run_steps(self.stdscr, f"{action.capitalize()} "
-                                       f"{spec.name} server", [step],
+                                        f"{spec.name} server", [step],
                            wait_on_finish=False)
         # Re-check the server instead of trusting the step's exit code
         # (cancel and failure both come back non-zero): did the toggle take?
@@ -599,15 +623,18 @@ class _Hub:
         return tui.Wizard.BACK
 
 
-def _server_action_step(spec, action: str):
+def _server_action_step(spec, action: str, prep=None):
     """Build a task step that starts/stops SPEC's server, logged to a file.
 
     ACTION is "start" or "stop". The step runs inside the task view (no
     console drop): the server module's output is tee'd to a timestamped
     ``<name>_<action>_*.log`` artifact under ``servers.LOG_DIR`` (see
-    ``logging_kit.run_artifact``) and to the view's log tail. Returns
-    ``(TaskStep, log_path)`` so the caller can point the user at the file
-    on failure.
+    ``logging_kit.run_artifact``) and to the view's log tail. PREP, when
+    given, runs inside a start step before the spawn (the companion-package
+    heal) — the task view keeps going past a failed step, so the gate has
+    to live in the step's own work: a non-zero prep result aborts the start.
+    Returns ``(TaskStep, log_path)`` so the caller can point the user at
+    the file on failure.
     """
     title = (f"Start {spec.name} server" if action == "start"
              else f"Stop {spec.name} server")
@@ -620,6 +647,9 @@ def _server_action_step(spec, action: str):
             with contextlib.redirect_stdout(
                     logging_kit.TeeWriter(logf, inner)):
                 if action == "start":
+                    if prep is not None and prep(emit, cancel) != 0:
+                        print("[ERROR] the server was not started")
+                        return 1
                     ok = servers.start(spec, cancel=cancel)
                 else:
                     ok = servers.stop(spec.name)
@@ -630,6 +660,86 @@ def _server_action_step(spec, action: str):
     return taskview.TaskStep(title, work), log_path
 
 
+# Selecting the SGLang-Omni entry while it has no downloaded models: the
+# handler flashes the remediation instead of toggling anything.
+_SGLOMNI_NO_MODELS = object()
+
+
+def _pick_start_model(stdscr, status_key: str, spec):
+    """The spec a manual start from the Start/Stop menu should boot.
+
+    One process hosts one model on the qwen-tts and SGLang-Omni backends,
+    and their detect() specs aim at a backend default — so a fresh start
+    asks which model to load instead of silently booting that default:
+    SGLang-Omni offers its downloaded catalog models (a model without
+    weights cannot boot), qwen-tts its three demos (a first boot downloads
+    the weights, like the Generate form's voice picker). SPEC returns
+    unchanged for every other backend, a running server (the toggle is a
+    stop), and a single installed SGLang-Omni model (it is the default
+    already); None means the user cancelled the picker.
+    """
+    if status_key == BACKEND_SGLOMNI:
+        entries = sglomni_backend.installed_entries()
+        if len(entries) < 2:
+            return spec
+        chosen = tui.menu(stdscr, "Load Which SGLang-Omni Model?",
+                          [(entry.label, entry) for entry in entries],
+                          back_value=tui.Wizard.BACK,
+                          help_lines=["One server process hosts one model;",
+                                      "stop it from this menu to load a",
+                                      "different one."])
+        if chosen is tui.Wizard.BACK:
+            return None
+        return sglomni_backend.build_spec(chosen)
+    if status_key == BACKEND_QWEN:
+        chosen = tui.menu(stdscr, "Load Which qwen-tts Model?",
+                          [(model, model) for model in
+                           qwen_backend.MODEL_REPOS],
+                          back_value=tui.Wizard.BACK,
+                          help_lines=["One demo process hosts one model;",
+                                      "stop it from this menu to load a",
+                                      "different one."])
+        if chosen is tui.Wizard.BACK:
+            return None
+        return qwen_backend.build_spec(chosen)
+    return spec
+
+
+def _sglomni_spec_entry(spec):
+    """The catalog entry SPEC's --model-path hosts (None when not sglomni's)."""
+    if spec.name != sglomni_backend.SERVER_NAME:
+        return None
+    argv = list(spec.argv)
+    try:
+        repo = argv[argv.index("--model-path") + 1]
+    except (ValueError, IndexError):
+        return None
+    return sglomni_backend.entry_by_repo(repo)
+
+
+def _companion_prep(spec, action: str):
+    """A start-step callable healing the spec's venv, or None.
+
+    An sglang-omni model whose companion packages are absent from the
+    backend venv (weights present via the shared HuggingFace cache, or a
+    failed install-time pip run — which only warns) boots into a
+    ``ModuleNotFoundError``. When the probe finds any missing, the returned
+    callable pip-installs exactly those first (the model's own recipe) so
+    the start aborts the boot rather than spawning a server that cannot
+    load its model. None when this spec is not sglang-omni's, hosts no
+    catalog model, or its venv already holds every companion.
+    """
+    if action != "start":
+        return None
+    entry = _sglomni_spec_entry(spec)
+    if entry is None or not sglomni_backend.missing_companions(entry):
+        return None
+    def prep(emit, cancel):
+        return sglomni_backend.install_companions(entry, emit=emit,
+                                                  cancel=cancel)
+    return prep
+
+
 def _configurable(info) -> bool:
     """True when INFO has a configure screen worth running from the hub.