feat: sglang-omni model picker on manual server launch

3 files changed, 323 insertions, 7 deletions

diff --git a/app/tests/test_backends_servers.py b/app/tests/test_backends_servers.py
index 4065a34..569c7bd 100644
--- a/app/tests/test_backends_servers.py
+++ b/app/tests/test_backends_servers.py
@@ -160,6 +160,14 @@ class StartTests(unittest.TestCase):
             'architecture")', "y"]))
         self.assertIsNone(servers._boot_hint([]))
 
+    def test_boot_hint_names_missing_companion_packages(self):
+        # A model whose companion pip packages are absent from the backend
+        # venv (e.g. sglang-omni's Qwen3-TTS models need qwen_tts) dies on
+        # the import; the hint says what to do about it.
+        self.assertIn("companion", servers._boot_hint(
+            ['File "...", in resolve_checkpoint',
+             "ModuleNotFoundError: No module named 'qwen_tts'"]))
+
     def test_console_progress_prints_the_hint(self):
         out = io.StringIO()
         with redirect_stdout(out):
diff --git a/app/tests/test_backends_sglomni.py b/app/tests/test_backends_sglomni.py
index a11aa92..1fd5f5a 100644
--- a/app/tests/test_backends_sglomni.py
+++ b/app/tests/test_backends_sglomni.py
@@ -11,7 +11,7 @@ from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
-from backends import envs, probe, servers
+from backends import common, envs, probe, servers
 from backends.sglomni import catalog, constants, models, pythonenv, status
 from backends.sglomni.catalog import CAPABILITY_CLONE, CAPABILITY_DESIGN, \
     CAPABILITY_SPEAKER, ENTRIES, config_path, entry_by_key, entry_by_repo, \
@@ -184,6 +184,82 @@ class ModelInstallStateTests(unittest.TestCase):
         self.assertIn("--model", str(ctx.exception))
 
 
+class CompanionPackageTests(unittest.TestCase):
+    """The venv probe and start-time heal for a model's companion packages.
+
+    Weights can reach the shared HF cache by another route (another
+    backend's install) or outlive a failed companion pip run — which
+    install_model only warns about — so the probe is what keeps a
+    "installed" model from booting into a ModuleNotFoundError.
+    """
+
+    def test_extra_import_names(self):
+        self.assertEqual(catalog.extra_import_name("sox"), "sox")
+        self.assertEqual(catalog.extra_import_name("einops"), "einops")
+        self.assertEqual(catalog.extra_import_name("qwen-tts==0.1.1"),
+                         "qwen_tts")
+        self.assertEqual(
+            catalog.extra_import_name("descript-audiotools==0.7.2"),
+            "audiotools")
+        self.assertEqual(
+            catalog.extra_import_name("descript-audio-codec==1.0.0"), "dac")
+
+    def test_missing_companions_reports_absent_modules(self):
+        entry = entry_by_key("qwen3_tts_1_7b_base")
+        with patch.object(envs, "env_exists", return_value=True), \
+                patch.object(envs, "module_available",
+                             side_effect=lambda name, env_dir:
+                             name != "qwen_tts"):
+            missing = models.missing_companions(entry)
+        self.assertEqual([spec for spec, _no_deps in missing],
+                         ["qwen-tts==0.1.1"])
+
+    def test_missing_companions_skips_every_extra_that_imports(self):
+        entry = entry_by_key("qwen3_tts_1_7b_base")
+        with patch.object(envs, "env_exists", return_value=True), \
+                patch.object(envs, "module_available", return_value=True):
+            self.assertEqual(models.missing_companions(entry), [])
+
+    def test_missing_companions_no_verdict_without_a_venv(self):
+        # A missing venv cannot be healed here: the start flow fails on
+        # the missing sgl-omni executable instead.
+        entry = entry_by_key("qwen3_tts_1_7b_base")
+        with patch.object(envs, "env_exists", return_value=False):
+            self.assertEqual(models.missing_companions(entry), [])
+
+    def test_install_companions_installs_only_missing_with_no_deps(self):
+        entry = entry_by_key("qwen3_tts_1_7b_base")
+        calls = []
+
+        def fake_pip_install(specs, **kwargs):
+            calls.append((list(specs), kwargs.get("extra_args")))
+            return 0
+
+        with patch.object(envs, "env_exists", return_value=True), \
+                patch.object(envs, "module_available",
+                             side_effect=lambda name, env_dir:
+                             name not in ("sox", "qwen_tts")), \
+                patch.object(common, "pip_install",
+                             side_effect=fake_pip_install):
+            self.assertEqual(models.install_companions(entry), 0)
+        self.assertEqual([specs for specs, _args in calls],
+                         [["sox"], ["qwen-tts==0.1.1"]])
+        self.assertTrue(all(args == ["--no-deps"] for _specs, args in calls))
+
+    def test_install_companions_stops_at_the_first_failure(self):
+        entry = entry_by_key("qwen3_tts_1_7b_base")
+        with patch.object(envs, "env_exists", return_value=True), \
+                patch.object(envs, "module_available",
+                             return_value=False), \
+                patch.object(common, "pip_install", return_value=23):
+            self.assertEqual(models.install_companions(entry), 23)
+
+    def test_models_with_no_extras_need_nothing(self):
+        entry = entry_by_key("higgs_audio_v3_tts")
+        self.assertEqual(entry.extras, ())
+        self.assertEqual(models.missing_companions(entry), [])
+
+
 class PythonEnvTests(unittest.TestCase):
     """Interpreter selection for the version-pinned venv."""
 
diff --git a/app/tests/test_hub.py b/app/tests/test_hub.py
index 8d5edd9..8e71104 100644
--- a/app/tests/test_hub.py
+++ b/app/tests/test_hub.py
@@ -805,6 +805,176 @@ class SubmenuStatusTableTests(unittest.TestCase):
         self.assertEqual(len(flashed), 1)
         self.assertIn("No backend is installed", flashed[0])
 
+    def test_server_menu_flashes_when_sglang_omni_has_no_models(self):
+        # An installed-but-model-less SGLang-Omni backend stays listed so
+        # selecting it can explain itself instead of being absent.
+        captured = {}
+        flashed = []
+        st = BackendStatus(hub.BACKEND_SGLOMNI, "SGLang-Omni", installed=True,
+                           configured=False)
+
+        def fake_menu(stdscr, title, options, **kwargs):
+            captured["title"] = title
+            captured["options"] = options
+            captured.update(kwargs)
+            return options[0][1]  # select the SGLang-Omni entry
+
+        with patch.object(hub.tui, "menu", fake_menu), \
+                patch.object(hub.tui, "flash",
+                             lambda *a, **k: flashed.append(a[1])), \
+                patch.object(hub, "detect_all", return_value=[st]), \
+                patch.object(hub.shutil, "which", return_value="/x"):
+            handler = hub._Hub(None).screen_server()
+            # The no-models handler flashes and lands back on the list.
+            self.assertIs(handler(), tui.Wizard.BACK)
+        self.assertEqual(captured["table_rows"],
+                         [("SGLang-Omni", "no models", "err", "body")])
+        self.assertIn("No SGLang-Omni models", flashed[-1])
+        self.assertIn("Configure Backends", flashed[-1])
+
+    def test_server_menu_asks_which_sglang_omni_model_to_load(self):
+        entries = [hub.sglomni_backend.ENTRIES[0],
+                   hub.sglomni_backend.ENTRIES[4]]
+        spec = hub.sglomni_backend.build_spec(entries[0])
+        st = BackendStatus(hub.BACKEND_SGLOMNI, "SGLang-Omni", installed=True,
+                           configured=True, servers=[spec])
+        menus = []
+        started = []
+        answers = [(hub.BACKEND_SGLOMNI, spec), entries[1]]
+
+        def fake_menu(stdscr, title, options, **kwargs):
+            menus.append((title, options))
+            return answers.pop(0)
+
+        def fake_run_steps(scr, title, steps, **kwargs):
+            for step in steps:
+                step.work(None, None)
+            return 0
+
+        with tempfile.TemporaryDirectory() as td, \
+                patch.object(hub.tui, "menu", fake_menu), \
+                patch.object(hub.tui, "flash", lambda *a, **k: None), \
+                patch.object(hub.common, "server_running",
+                             return_value=False), \
+                patch.object(hub.servers, "LOG_DIR", Path(td)), \
+                patch.object(hub.servers, "start",
+                             side_effect=lambda spec, **k:
+                             started.append(spec) or True), \
+                patch.object(hub.taskview, "run_steps", fake_run_steps), \
+                patch.object(hub.sglomni_backend, "installed_entries",
+                             return_value=entries), \
+                patch.object(hub.sglomni_backend, "missing_companions",
+                             return_value=[]), \
+                patch.object(hub, "detect_all", return_value=[st]), \
+                patch.object(hub.shutil, "which", return_value="/x"):
+            toggle = hub._Hub(None).screen_server()
+            toggle()  # the wizard runs the returned screen
+        self.assertEqual([title for title, _options in menus],
+                         ["Start / Stop A Server",
+                          "Load Which SGLang-Omni Model?"])
+        self.assertEqual([opt[0] for opt in menus[1][1]],
+                         [entry.label for entry in entries])
+        # The server boots the picked model's spec, not the menu's default
+        # (the first installed catalog model).
+        self.assertEqual(len(started), 1)
+        self.assertIn(entries[1].repo, started[0].argv)
+
+    def test_server_menu_starts_a_single_sglang_omni_model_without_asking(self):
+        entry = hub.sglomni_backend.ENTRIES[4]
+        spec = hub.sglomni_backend.build_spec(entry)
+        st = BackendStatus(hub.BACKEND_SGLOMNI, "SGLang-Omni", installed=True,
+                           configured=True, servers=[spec])
+        menus = []
+        started = []
+        answer = (hub.BACKEND_SGLOMNI, spec)
+
+        def fake_menu(stdscr, title, options, **kwargs):
+            menus.append((title, options))
+            return answer
+
+        def fake_run_steps(scr, title, steps, **kwargs):
+            for step in steps:
+                step.work(None, None)
+            return 0
+
+        with tempfile.TemporaryDirectory() as td, \
+                patch.object(hub.tui, "menu", fake_menu), \
+                patch.object(hub.tui, "flash", lambda *a, **k: None), \
+                patch.object(hub.common, "server_running",
+                             return_value=False), \
+                patch.object(hub.servers, "LOG_DIR", Path(td)), \
+                patch.object(hub.servers, "start",
+                             side_effect=lambda spec, **k:
+                             started.append(spec) or True), \
+                patch.object(hub.taskview, "run_steps", fake_run_steps), \
+                patch.object(hub.sglomni_backend, "installed_entries",
+                             return_value=[entry]), \
+                patch.object(hub, "detect_all", return_value=[st]), \
+                patch.object(hub.shutil, "which", return_value="/x"):
+            toggle = hub._Hub(None).screen_server()
+            toggle()
+        self.assertEqual(len(menus), 1)  # no model picker for one model
+        self.assertEqual(len(started), 1)
+        self.assertIs(started[0], spec)
+
+    def test_server_menu_asks_which_qwen_model_to_load(self):
+        spec = ServerSpec("qwen", "http://127.0.0.1:8300", ["demo"])
+        st = BackendStatus("qwen", "qwen-tts", installed=True,
+                           configured=True, servers=[spec])
+        menus = []
+        started = []
+        answers = [("qwen", spec), "Base"]
+
+        def fake_menu(stdscr, title, options, **kwargs):
+            menus.append((title, options))
+            return answers.pop(0)
+
+        def fake_run_steps(scr, title, steps, **kwargs):
+            for step in steps:
+                step.work(None, None)
+            return 0
+
+        with tempfile.TemporaryDirectory() as td, \
+                patch.object(hub.tui, "menu", fake_menu), \
+                patch.object(hub.tui, "flash", lambda *a, **k: None), \
+                patch.object(hub.common, "server_running",
+                             return_value=False), \
+                patch.object(hub.servers, "LOG_DIR", Path(td)), \
+                patch.object(hub.servers, "start",
+                             side_effect=lambda spec, **k:
+                             started.append(spec) or True), \
+                patch.object(hub.taskview, "run_steps", fake_run_steps), \
+                patch.object(hub, "detect_all", return_value=[st]), \
+                patch.object(hub.shutil, "which", return_value="/x"):
+            toggle = hub._Hub(None).screen_server()
+            toggle()
+        self.assertEqual(menus[1][0], "Load Which qwen-tts Model?")
+        self.assertEqual([opt[0] for opt in menus[1][1]],
+                         list(hub.qwen_backend.MODEL_REPOS))
+        self.assertEqual(len(started), 1)
+        self.assertEqual(started[0].argv[1],
+                         hub.qwen_backend.MODEL_REPOS["Base"])
+        self.assertEqual(started[0].identity,
+                         hub.backend_probe.IDENTITY_QWEN_CLONE)
+
+    def test_start_model_picker_cancel_keeps_the_server_stopped(self):
+        entry = hub.sglomni_backend.ENTRIES[4]
+        spec = hub.sglomni_backend.build_spec(entry)
+        answers = [tui.Wizard.BACK]  # Esc in the model picker
+
+        def fake_menu(stdscr, title, options, **kwargs):
+            return answers.pop(0)
+
+        with patch.object(hub.tui, "menu", fake_menu), \
+                patch.object(hub.sglomni_backend, "installed_entries",
+                             return_value=[entry, entry]):
+            result = hub._pick_start_model(None, hub.BACKEND_SGLOMNI, spec)
+        self.assertIsNone(result)
+
+    def test_pick_start_model_passes_other_backends_through(self):
+        spec = ServerSpec("audiocpp", "http://127.0.0.1:8080", ["server"])
+        self.assertIs(hub._pick_start_model(None, "audiocpp", spec), spec)
+
     def test_submenu_repeats_ffmpeg_warning(self):
         captured = {}
         infos = [BackendInfo("qwen", "qwen-tts", lambda: None, lambda: 0)]
@@ -820,6 +990,65 @@ class SubmenuStatusTableTests(unittest.TestCase):
                          [("Warning: ffmpeg not installed!", "err")])
 
 
+class CompanionPrepTests(unittest.TestCase):
+    """The start step's companion-package heal for sglang-omni specs."""
+
+    def test_prep_none_for_other_backends_and_for_stops(self):
+        spec = ServerSpec("qwen", "http://127.0.0.1:8300", ["demo"])
+        self.assertIsNone(hub._companion_prep(spec, "start"))
+        sgl = hub.sglomni_backend.build_spec(
+            hub.sglomni_backend.ENTRIES[4])
+        self.assertIsNone(hub._companion_prep(sgl, "stop"))
+
+    def test_prep_none_when_every_companion_imports(self):
+        sgl = hub.sglomni_backend.build_spec(
+            hub.sglomni_backend.ENTRIES[0])
+        with patch.object(hub.sglomni_backend, "missing_companions",
+                          return_value=[]):
+            self.assertIsNone(hub._companion_prep(sgl, "start"))
+
+    def test_prep_installs_the_missing_companions(self):
+        entry = hub.sglomni_backend.ENTRIES[0]
+        sgl = hub.sglomni_backend.build_spec(entry)
+        with patch.object(hub.sglomni_backend, "missing_companions",
+                          return_value=[("qwen-tts==0.1.1", True)]), \
+                patch.object(hub.sglomni_backend, "install_companions",
+                             return_value=0) as install:
+            prep = hub._companion_prep(sgl, "start")
+            self.assertIsNotNone(prep)
+            self.assertEqual(prep(None, None), 0)
+        install.assert_called_once()
+
+    def test_start_step_runs_prep_before_the_spawn(self):
+        spec = ServerSpec("qwen", "http://127.0.0.1:8300", ["demo"])
+        order = []
+
+        def prep(emit, cancel):
+            order.append("prep")
+            return 0
+
+        with tempfile.TemporaryDirectory() as td, \
+                patch.object(hub.servers, "LOG_DIR", Path(td)), \
+                patch.object(hub.servers, "start",
+                             side_effect=lambda *a, **k:
+                             order.append("start") or True):
+            step, _log = hub._server_action_step(spec, "start", prep=prep)
+            self.assertEqual(step.work(None, None), 0)
+        self.assertEqual(order, ["prep", "start"])
+
+    def test_start_step_aborts_when_prep_fails(self):
+        # The task view keeps going past a failed step, so the gate has to
+        # live in the start step's own work: no spawn after a failed heal.
+        spec = ServerSpec("qwen", "http://127.0.0.1:8300", ["demo"])
+        with tempfile.TemporaryDirectory() as td, \
+                patch.object(hub.servers, "LOG_DIR", Path(td)), \
+                patch.object(hub.servers, "start") as start:
+            step, _log = hub._server_action_step(
+                spec, "start", prep=lambda emit, cancel: 23)
+            self.assertEqual(step.work(None, None), 1)
+        start.assert_not_called()
+
+
 class ConvertFlowTests(unittest.TestCase):
     """_convert_form / screen_convert: one form whose first field is the
     Backend picker, followed by that backend's options (local config or live
@@ -4444,7 +4673,8 @@ class HubNavigationTests(unittest.TestCase):
                                configured=True, servers=specs)
         registry = [self._info("qwen", "qwen-tts")]
         titles = []
-        script = ["configure_backends", "server", specs[0],
+        script = ["configure_backends", "server", ("qwen", specs[0]),
+                  "CustomVoice",
                   tui.Wizard.BACK, tui.Wizard.BACK, tui.Wizard.BACK]
 
         def menu(stdscr, title, options, **kwargs):
@@ -4461,14 +4691,16 @@ class HubNavigationTests(unittest.TestCase):
                 patch.object(hub.audiocpp_backend, "find_local_checkout",
                              return_value=None):
             hub._Hub(None).run()
-        # Selecting a server toggles it directly (no action sub-menu), then
-        # Esc steps back one screen at a time: server list → Configure
-        # Backends → main menu.
+        # Selecting a server asks which model to load (one process hosts
+        # one), toggles it directly (no action sub-menu), then Esc steps
+        # back one screen at a time: model picker → server list →
+        # Configure Backends → main menu.
         self.assertEqual(
             titles,
             ["tts-audiobook-generator", "Configure Backends",
-             "Start / Stop A Server", "Start / Stop A Server",
-             "Configure Backends", "tts-audiobook-generator"])
+             "Start / Stop A Server", "Load Which qwen-tts Model?",
+             "Start / Stop A Server", "Configure Backends",
+             "tts-audiobook-generator"])
 
     def test_esc_on_main_menu_quits(self):
         titles = self._drive([tui.Wizard.BACK], [], [])