feat: sglang-omni model picker on manual server launch

1 files changed, 238 insertions, 6 deletions

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], [], [])