feat: sglang backend support

1 files changed, 828 insertions, 0 deletions

diff --git a/app/tests/test_backends_sglomni.py b/app/tests/test_backends_sglomni.py
new file mode 100644
index 0000000..a11aa92
--- /dev/null
+++ b/app/tests/test_backends_sglomni.py
@@ -0,0 +1,828 @@
+"""Tests for the SGLang-Omni backend package (backends/sglomni)."""
+
+import io
+import json
+import urllib.error
+import sys
+import tempfile
+import unittest
+from contextlib import redirect_stdout
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from backends import 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, \
+    install_tree_families
+from backends.sglomni.pythonenv import SGLOMNI_ENV
+
+
+class CatalogTests(unittest.TestCase):
+    """The model catalog is the single source of hosting/voice facts."""
+
+    def test_unique_keys_and_repos(self):
+        keys = [entry.key for entry in ENTRIES]
+        repos = [entry.repo for entry in ENTRIES]
+        self.assertEqual(len(keys), len(set(keys)))
+        self.assertEqual(len(repos), len(set(repos)))
+
+    def test_capabilities_are_known(self):
+        for entry in ENTRIES:
+            self.assertIn(entry.capability,
+                          (CAPABILITY_SPEAKER, CAPABILITY_CLONE,
+                           CAPABILITY_DESIGN))
+
+    def test_vendored_config_files_exist(self):
+        for entry in ENTRIES:
+            path = config_path(entry)
+            if entry.config is None:
+                self.assertIsNone(path)
+            else:
+                self.assertTrue(path.is_file(), f"missing {path}")
+
+    def test_config_declares_the_entry_repo(self):
+        # The vendored yaml pins model_path — it must match the entry's
+        # repo, or the server would host something else than the run
+        # selected (the client's connect check would refuse it). The
+        # files are flat `key: value` documents, parsed by hand here.
+        for entry in ENTRIES:
+            path = config_path(entry)
+            if path is None:
+                continue
+            data = {}
+            for line in path.read_text(encoding="utf-8").splitlines():
+                line = line.strip()
+                if not line or line.startswith("#") or ":" not in line:
+                    continue
+                key, _, value = line.partition(":")
+                data[key.strip()] = value.strip()
+            self.assertEqual(data.get("model_path"), entry.repo,
+                             f"stale config for {entry.key}")
+
+    def test_clone_capability_matches_reference_requirement(self):
+        # Only clone models carry a reference requirement.
+        for entry in ENTRIES:
+            if entry.requires_reference:
+                self.assertEqual(entry.capability, CAPABILITY_CLONE,
+                                 entry.key)
+
+    def test_entry_lookup_by_key_and_repo(self):
+        entry = ENTRIES[0]
+        self.assertIs(entry_by_key(entry.key), entry)
+        self.assertIs(entry_by_repo(entry.repo), entry)
+        self.assertIsNone(entry_by_key("nope"))
+        self.assertIsNone(entry_by_repo("nope"))
+
+    def test_install_tree_covers_every_entry(self):
+        families = install_tree_families(list(ENTRIES))
+        covered = [option["key"] for family in families
+                   for option in family["options"]]
+        self.assertEqual(sorted(covered),
+                         sorted(entry.key for entry in ENTRIES))
+
+    def test_install_tree_has_no_detail_line(self):
+        # The install screen's status line under the Confirm/Back buttons
+        # would only repeat the catalog keys under the cursor — the tree
+        # carries no detail at all, so no status line is drawn.
+        families = install_tree_families(list(ENTRIES))
+        self.assertTrue(families)
+        for family in families:
+            self.assertNotIn("detail", family)
+
+    def test_install_tree_filters_unavailable(self):
+        some = [ENTRIES[0]]
+        families = install_tree_families(some)
+        covered = [option["key"] for family in families
+                   for option in family["options"]]
+        self.assertEqual(covered, [ENTRIES[0].key])
+
+
+class ModelInstallStateTests(unittest.TestCase):
+    """Install state reads the shared HuggingFace hub cache layout."""
+
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.cache = Path(self._tmp.name)
+        patcher = patch.object(models, "_hf_cache_dir", return_value=self.cache)
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        self.addCleanup(self._tmp.cleanup)
+
+    def _seed_repo(self, repo):
+        directory = self.cache / ("models--" + repo.replace("/", "--"))
+        (directory / "refs").mkdir(parents=True)
+        (directory / "refs" / "main").write_text("hash\n")
+        (directory / "snapshots" / "abc").mkdir(parents=True)
+        (directory / "snapshots" / "abc" / "weights.safetensors") \
+            .write_bytes(b"x")
+        return directory
+
+    def test_model_installed_needs_refs_and_snapshots(self):
+        entry = entry_by_key("higgs_audio_v3_tts")
+        self.assertFalse(models.model_installed(entry))
+        self._seed_repo(entry.repo)
+        self.assertTrue(models.model_installed(entry))
+
+    def test_installed_entries_in_catalog_order(self):
+        first, second = ENTRIES[0], ENTRIES[4]
+        self._seed_repo(second.repo)
+        self._seed_repo(first.repo)
+        keys = models.installed_keys()
+        self.assertEqual(keys, [first.key, second.key])
+
+    def test_delete_model_weights_removes_only_targeted_repos(self):
+        entry = ENTRIES[0]
+        other = ENTRIES[1]
+        self._seed_repo(entry.repo)
+        self._seed_repo(other.repo)
+        removed = models.delete_model_weights([entry])
+        self.assertEqual(removed, 1)
+        self.assertFalse(models.model_installed(entry))
+        self.assertTrue(models.model_installed(other))
+
+    def test_preset_voices_from_voice_embedding(self):
+        # Voxtral-style: preset voices ship as voice_embedding/*.pt in the
+        # downloaded snapshot.
+        entry = entry_by_key("voxtral_tts")
+        directory = self._seed_repo(entry.repo)
+        (directory / "snapshots" / "abc" / "voice_embedding").mkdir()
+        (directory / "snapshots" / "abc" / "voice_embedding" / "casual_male.pt") \
+            .write_bytes(b"x")
+        (directory / "snapshots" / "abc" / "voice_embedding" / "default.pt") \
+            .write_bytes(b"x")
+        self.assertEqual(models.preset_voices(entry),
+                         ["casual_male", "default"])
+
+    def test_preset_voices_from_catalog_table(self):
+        entry = entry_by_key("qwen3_tts_0_6b_customvoice")
+        self.assertTrue(entry.speakers)
+        self.assertEqual(models.preset_voices(entry), list(entry.speakers))
+
+    def test_resolve_model_unknown_key_raises(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            models.resolve_model("nope")
+        self.assertIn("Unknown sglang-omni model", str(ctx.exception))
+
+    def test_resolve_model_requires_installed_weights(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            models.resolve_model("higgs_audio_v3_tts")
+        self.assertIn("not downloaded", str(ctx.exception))
+
+    def test_resolve_model_auto_selects_the_single_install(self):
+        entry = ENTRIES[0]
+        self._seed_repo(entry.repo)
+        self.assertIs(models.resolve_model(None), entry)
+        self.assertIs(models.resolve_model(entry.key), entry)
+
+    def test_resolve_model_needs_a_pick_with_several_installs(self):
+        self._seed_repo(ENTRIES[0].repo)
+        self._seed_repo(ENTRIES[1].repo)
+        with self.assertRaises(RuntimeError) as ctx:
+            models.resolve_model(None)
+        self.assertIn("--model", str(ctx.exception))
+
+
+class PythonEnvTests(unittest.TestCase):
+    """Interpreter selection for the version-pinned venv."""
+
+    def test_env_compatible_needs_310_to_312(self):
+        with patch.object(pythonenv, "env_version", return_value=(3, 12)):
+            self.assertTrue(pythonenv.env_compatible())
+        with patch.object(pythonenv, "env_version", return_value=(3, 13)):
+            self.assertFalse(pythonenv.env_compatible())
+        with patch.object(pythonenv, "env_version", return_value=None):
+            self.assertFalse(pythonenv.env_compatible())
+
+    def test_prepare_env_noop_on_compatible_venv(self):
+        with patch.object(pythonenv, "env_compatible", return_value=True), \
+                patch.object(envs, "create_env") as mock_create, \
+                patch.object(envs, "provision_env_with_uv") as mock_uv:
+            self.assertEqual(pythonenv.prepare_env(), 0)
+        mock_create.assert_not_called()
+        mock_uv.assert_not_called()
+
+    def test_prepare_env_uses_compatible_system_interpreter(self):
+        interpreter = Path("/usr/bin/python3.12")
+        with patch.object(pythonenv, "env_compatible", return_value=False), \
+                patch.object(envs, "env_exists", return_value=False), \
+                patch.object(envs, "compatible_interpreter",
+                             return_value=interpreter) as mock_find, \
+                patch.object(envs, "create_env",
+                             return_value=0) as mock_create:
+            self.assertEqual(pythonenv.prepare_env(), 0)
+        mock_find.assert_called_once()
+        mock_create.assert_called_once_with(SGLOMNI_ENV, interpreter)
+
+    def test_prepare_env_falls_back_to_uv(self):
+        with patch.object(pythonenv, "env_compatible", return_value=False), \
+                patch.object(envs, "env_exists", return_value=False), \
+                patch.object(envs, "compatible_interpreter",
+                             return_value=None), \
+                patch.object(envs, "ensure_uv", return_value=0) as mock_uv_install, \
+                patch.object(envs, "provision_env_with_uv",
+                             return_value=0) as mock_uv:
+            self.assertEqual(pythonenv.prepare_env(), 0)
+        mock_uv_install.assert_called_once()
+        mock_uv.assert_called_once()
+
+    def test_prepare_env_reports_uv_failure(self):
+        with patch.object(pythonenv, "env_compatible", return_value=False), \
+                patch.object(envs, "env_exists", return_value=False), \
+                patch.object(envs, "compatible_interpreter",
+                             return_value=None), \
+                patch.object(envs, "ensure_uv", return_value=1):
+            self.assertEqual(pythonenv.prepare_env(), 1)
+
+
+class BuildSpecTests(unittest.TestCase):
+    """The managed ServerSpec: sgl-omni serve with model/config/port."""
+
+    def test_spec_hosts_the_entry_repo_with_config(self):
+        entry = entry_by_key("qwen3_tts_0_6b_customvoice")
+        spec = status.build_spec(entry)
+        self.assertEqual(spec.name, constants.SERVER_NAME)
+        self.assertEqual(spec.identity, probe.IDENTITY_SGLOMNI)
+        self.assertEqual(spec.start_timeout, constants.SERVER_START_TIMEOUT)
+        script, serve, flag, repo, cfg_flag, cfg, port_flag, port = spec.argv
+        self.assertEqual(serve, "serve")
+        self.assertEqual(flag, "--model-path")
+        self.assertEqual(repo, entry.repo)
+        self.assertEqual(cfg_flag, "--config")
+        self.assertEqual(Path(cfg), config_path(entry))
+        self.assertEqual(port_flag, "--port")
+        self.assertIn(port, str(spec.url))
+
+    def test_spec_omits_config_when_none_needed(self):
+        entry = entry_by_key("higgs_audio_v3_tts")
+        spec = status.build_spec(entry)
+        self.assertNotIn("--config", spec.argv)
+
+
+class GpuCapabilityTests(unittest.TestCase):
+    """The nvidia-smi-backed GPU facts are best-effort and cached."""
+
+    def setUp(self):
+        gpu_module = status.gpu
+        gpu_module._query.cache_clear()
+        self.addCleanup(gpu_module._query.cache_clear)
+
+    def _nvidia_smi(self, *, stdout="", returncode=0, installed=True):
+        def fake_run(argv, **_kwargs):
+            if not installed:
+                raise FileNotFoundError("nvidia-smi")
+            return SimpleNamespace(returncode=returncode, stdout=stdout)
+        return fake_run
+
+    def test_parses_name_and_compute_capability(self):
+        with patch.object(status.gpu.shutil, "which", return_value="/x"), \
+                patch.object(status.gpu.subprocess, "run",
+                             side_effect=self._nvidia_smi(
+                                 stdout="NVIDIA GeForce RTX 3090, 8.6\n")):
+            self.assertEqual(status.gpu.compute_capability(), (8, 6))
+            self.assertEqual(status.gpu.describe(),
+                             "NVIDIA GeForce RTX 3090 (compute capability 8.6)")
+
+    def test_none_when_nvidia_smi_missing(self):
+        with patch.object(status.gpu.shutil, "which", return_value=None):
+            self.assertIsNone(status.gpu.compute_capability())
+            self.assertIsNone(status.gpu.describe())
+
+    def test_none_when_the_query_fails_or_is_garbage(self):
+        for kwargs in (dict(returncode=1),
+                       dict(stdout=""),
+                       dict(stdout="name only\n")):
+            with self.subTest(stdout=kwargs.get("stdout")):
+                with patch.object(status.gpu.shutil, "which",
+                                  return_value="/x"), \
+                        patch.object(status.gpu.subprocess, "run",
+                                     side_effect=self._nvidia_smi(**kwargs)):
+                    self.assertIsNone(status.gpu.compute_capability())
+
+
+class Fp8FallbackTests(unittest.TestCase):
+    """FP8-only pipelines fall back to a vendored bf16 config on old GPUs."""
+
+    ZONOS2 = "zonos2"
+
+    def _fallback(self, capability):
+        return patch("backends.sglomni.gpu.compute_capability",
+                     return_value=capability)
+
+    def test_fallback_config_declares_the_repo_and_disables_fp8(self):
+        entry = entry_by_key("zonos2")
+        path = catalog.fallback_config_path(entry)
+        self.assertIsNotNone(path)
+        self.assertTrue(path.is_file(), f"missing {path}")
+        text = path.read_text(encoding="utf-8")
+        self.assertIn(f"model_path: {entry.repo}", text)
+        self.assertRegex(text, r"fp8:\s*false")
+        # bf16 weights need a bigger static pool than the builder's 0.5
+        # default (24 GB card: >=0.64 for any KV cache at all).
+        self.assertRegex(text, r"mem_fraction_static:\s*0\.70")
+
+    def test_only_fp8_models_carry_a_fallback(self):
+        for entry in ENTRIES:
+            if entry.fp8_moe:
+                self.assertIsNotNone(entry.fp8_min_compute_capability,
+                                     entry.key)
+                self.assertIsNotNone(entry.bf16_config, entry.key)
+            else:
+                self.assertIsNone(catalog.fallback_config_path(entry))
+
+    def test_fallback_needed_below_the_capability_floor(self):
+        entry = entry_by_key("zonos2")
+        with patch("backends.sglomni.gpu.compute_capability",
+                   return_value=(8, 6)):
+            self.assertTrue(status.needs_fp8_fallback(entry))
+            self.assertEqual(status.launch_config_path(entry),
+                             catalog.fallback_config_path(entry))
+            note = status.gpu_fallback_note(entry)
+        self.assertIn("bf16", note)
+        self.assertIn("8.9", note)
+
+    def test_no_fallback_at_or_above_the_capability(self):
+        entry = entry_by_key("zonos2")
+        for capability in ((8, 9), (9, 0), (10, 0)):
+            with self.subTest(capability=capability):
+                with patch("backends.sglomni.gpu.compute_capability",
+                           return_value=capability):
+                    self.assertFalse(status.needs_fp8_fallback(entry))
+                    self.assertIsNone(status.gpu_fallback_note(entry))
+                    self.assertEqual(status.launch_config_path(entry),
+                                     config_path(entry))
+
+    def test_no_fallback_without_an_answerable_gpu(self):
+        # A GPU this tool cannot read keeps upstream defaults instead of
+        # second-guessing the host.
+        entry = entry_by_key("zonos2")
+        with patch("backends.sglomni.gpu.compute_capability",
+                   return_value=None):
+            self.assertFalse(status.needs_fp8_fallback(entry))
+            self.assertIsNone(status.gpu_fallback_note(entry))
+            self.assertEqual(status.launch_config_path(entry),
+                             config_path(entry))
+
+    def test_non_fp8_models_never_fall_back(self):
+        for entry in ENTRIES:
+            if entry.key == "zonos2":
+                continue
+            with patch("backends.sglomni.gpu.compute_capability",
+                       return_value=(1, 0)):
+                self.assertFalse(status.needs_fp8_fallback(entry))
+
+    def test_spec_launches_the_bf16_config_on_an_old_gpu(self):
+        entry = entry_by_key("zonos2")
+        with patch("backends.sglomni.gpu.compute_capability",
+                   return_value=(8, 6)):
+            spec = status.build_spec(entry)
+        self.assertIn("--config", spec.argv)
+        self.assertEqual(Path(spec.argv[spec.argv.index("--config") + 1]),
+                         catalog.fallback_config_path(entry))
+
+    def test_spec_keeps_the_default_pipeline_on_modern_gpus(self):
+        entry = entry_by_key("zonos2")
+        with patch("backends.sglomni.gpu.compute_capability",
+                   return_value=(9, 0)):
+            spec = status.build_spec(entry)
+        self.assertNotIn("--config", spec.argv)
+
+    def test_detect_tags_a_fallback_model(self):
+        entry = entry_by_key("zonos2")
+        with patch("backends.sglomni.status._is_installed",
+                   return_value=True), \
+                patch("backends.sglomni.status.installed_entries",
+                      return_value=[entry]), \
+                patch("backends.sglomni.gpu.compute_capability",
+                      return_value=(8, 6)), \
+                patch.object(servers, "manages", return_value=False), \
+                patch.object(status, "_detect_remote",
+                             return_value=([], {})):
+            st = status.detect()
+        self.assertIn("zonos2 (bf16 fallback)",
+                      next(line for line in st.details
+                           if line.startswith("models: ")))
+
+    def test_install_prints_the_fallback_note(self):
+        out = io.StringIO()
+        with patch("backends.sglomni.models.prepare_env", return_value=0), \
+                patch("backends.common.pip_install", return_value=0), \
+                patch("backends.sglomni.models._hf_download_prefix",
+                      return_value=["hf"]), \
+                patch("backends.common.run_console_subprocess",
+                      return_value=0), \
+                patch("backends.sglomni.gpu.compute_capability",
+                      return_value=(8, 6)), \
+                redirect_stdout(out):
+            rc = models.install_model("zonos2")
+        self.assertEqual(rc, 0)
+        self.assertIn("bf16", out.getvalue())
+
+
+class DetectTests(unittest.TestCase):
+    """detect() reports install state, models, and the running model."""
+
+    def _detect(self, *, script=False, module=False, entries=()):
+        entry = entries[0] if entries else None
+        spec = [status.build_spec(entry)] if entry else []
+        with patch("backends.sglomni.status._is_installed",
+                   return_value=script or module), \
+                patch("backends.sglomni.status.installed_entries",
+                      return_value=list(entries)), \
+                patch.object(servers, "manages", return_value=False), \
+                patch.object(status, "_detect_remote",
+                             return_value=([], {})):
+            return status.detect(), spec
+
+    def test_not_installed(self):
+        st, _spec = self._detect()
+        self.assertFalse(st.installed)
+        self.assertFalse(st.configured)
+        self.assertFalse(st.running)
+        self.assertEqual(st.servers, [])
+        self.assertEqual(st.partial, "")
+
+    def test_installed_without_models_is_partial(self):
+        st, _spec = self._detect(script=True)
+        self.assertTrue(st.installed)
+        self.assertFalse(st.configured)
+        self.assertEqual(st.partial, "installed (no models)")
+
+    def test_configured_reports_a_spec_and_ready(self):
+        entry = entry_by_key("higgs_audio_v3_tts")
+        st, _spec = self._detect(script=True, entries=[entry])
+        self.assertTrue(st.configured)
+        self.assertTrue(st.ready)
+        self.assertEqual(len(st.servers), 1)
+        self.assertIn(entry.repo, st.servers[0].argv)
+
+
+class ProbeIdentityTests(unittest.TestCase):
+    """A healthy sglang-omni /health identifies the sglomni backend."""
+
+    def _urlopen_returning(self, payloads):
+        calls = {"index": 0}
+
+        def fake_urlopen(url, timeout=3.0):
+            if calls["index"] >= len(payloads):
+                # Past the scripted payloads (e.g. the gradio fallback
+                # probe): behave like a 404 — urlopen raises, _get_json
+                # maps that to None.
+                calls["index"] += 1
+                raise urllib.error.URLError("HTTP 404")
+            payload = payloads[calls["index"]]
+            calls["index"] += 1
+            response = MagicMock()
+            response.__enter__.return_value = response
+            response.read.return_value = json.dumps(payload).encode("utf-8")
+            return response
+
+        return fake_urlopen, calls
+
+    def test_healthy_server_with_stages_identifies_sglomni(self):
+        health = {"status": "healthy", "running": True,
+                  "stages": ["preprocessing", "tts_generation", "vocoder"]}
+        fake_urlopen, _calls = self._urlopen_returning([health])
+        with patch.object(probe.common, "server_running", return_value=True), \
+                patch("backends.probe.urllib.request.urlopen", fake_urlopen):
+            self.assertEqual(probe.identify_server("http://127.0.0.1:8100"),
+                             probe.IDENTITY_SGLOMNI)
+
+    def test_unhealthy_server_is_not_sglomni(self):
+        health = {"status": "unhealthy", "running": False, "stages": []}
+        fake_urlopen, _calls = self._urlopen_returning([health])
+        with patch.object(probe.common, "server_running", return_value=True), \
+                patch("backends.probe.urllib.request.urlopen", fake_urlopen):
+            self.assertIsNone(probe.identify_server("http://127.0.0.1:8100"))
+
+    def test_served_model_read_from_v1_models(self):
+        payload = {"object": "list", "data": [
+            {"id": "bosonai/higgs-audio-v3-tts-4b", "root":
+             "bosonai/higgs-audio-v3-tts-4b"}]}
+        fake_urlopen, _calls = self._urlopen_returning([payload])
+        with patch("backends.probe.urllib.request.urlopen", fake_urlopen):
+            self.assertEqual(
+                probe.sglomni_served_model("http://127.0.0.1:8100"),
+                "bosonai/higgs-audio-v3-tts-4b")
+
+    def test_served_model_none_on_garbage(self):
+        fake_urlopen, _calls = self._urlopen_returning([{"data": []}])
+        with patch("backends.probe.urllib.request.urlopen", fake_urlopen):
+            self.assertIsNone(
+                probe.sglomni_served_model("http://127.0.0.1:8100"))
+
+    def test_voice_names_read_from_uploaded_voices(self):
+        payload = {"uploaded_voice_names": ["narrator", "second narrator"]}
+        fake_urlopen, _calls = self._urlopen_returning([payload])
+        with patch("backends.probe.urllib.request.urlopen", fake_urlopen):
+            self.assertEqual(
+                probe.sglomni_voice_names("http://127.0.0.1:8100"),
+                ["narrator", "second narrator"])
+
+
+class ModelsScreenTests(unittest.TestCase):
+    """models_screen: the audio.cpp-style checkbox tree driving steps.
+
+    The tree is scripted (like the qwen models_screen tests): each
+    fake render records its arguments and returns the next scripted
+    answer; the task-view run executes its steps inline so delegation to
+    install/uninstall_model is observable.
+    """
+
+    FAMILY_INDEX = {option["key"]: index
+                    for index, family in
+                    enumerate(install_tree_families(list(ENTRIES)))
+                    for option in family["options"]}
+
+    def _screen(self, answers, *, installed=(), package=True, confirm=True,
+                extra=()):
+        """Run models_screen with scripted tree answers; record calls.
+
+        Returns ``(rc, trees, confirms, flashes, runs)``: trees holds one
+        (title, kwargs) per render, confirms every uninstall question,
+        flashes every (text, kind), and runs each (title, step titles)
+        while executing its steps' work inline.
+        """
+        import contextlib
+
+        from backends.sglomni import wizard
+        choices = list(answers)
+        trees, confirms, flashes, runs = [], [], [], []
+
+        def fake_tree(stdscr, title, families, **kwargs):
+            trees.append((title, kwargs))
+            return choices.pop(0)
+
+        def fake_confirm(scr, question, **kwargs):
+            confirms.append((question, kwargs))
+            return confirm
+
+        def fake_flash(scr, text, kind="warn"):
+            flashes.append((text, kind))
+
+        def fake_run(scr, title, steps, **kwargs):
+            runs.append((title, [step.title for step in steps]))
+            for step in steps:
+                step.work(None, None)
+            return 0
+
+        patches = [
+            patch.object(wizard, "_is_installed", return_value=package),
+            patch.object(models, "installed_keys",
+                         return_value=list(installed)),
+            patch.object(wizard.tui, "checkbox_tree", fake_tree),
+            patch.object(wizard.tui, "confirm", fake_confirm),
+            patch.object(wizard.tui, "flash", fake_flash),
+            patch.object(wizard.taskview, "run_steps", fake_run),
+            *extra,
+        ]
+        with contextlib.ExitStack() as stack:
+            for ctx in patches:
+                stack.enter_context(ctx)
+            rc = wizard.models_screen(None)
+        return rc, trees, confirms, flashes, runs
+
+    def test_tree_is_the_audio_cpp_modify_flow(self):
+        from backends.sglomni import wizard
+        first = ENTRIES[0]
+        rc, trees, _confirms, _flashes, _runs = self._screen(
+            [wizard._GO_BACK], installed=(first.key,))
+        self.assertEqual(rc, 0)
+        title, kwargs = trees[0]
+        self.assertEqual(title, "Select SGLang-Omni Models")
+        # The installed model starts checked (a modify list), Confirm is
+        # pre-focused, and an empty selection is a valid answer.
+        self.assertEqual(kwargs["checked"],
+                         {(self.FAMILY_INDEX[first.key], first.key)})
+        self.assertTrue(kwargs["start_on_buttons"])
+        self.assertTrue(kwargs["allow_empty"])
+        self.assertIs(kwargs["back_value"], wizard._GO_BACK)
+
+    def test_checking_a_model_runs_one_install_step(self):
+        from backends.sglomni import wizard
+        entry = ENTRIES[0]
+        requested = []
+
+        def capture(key, *, emit=None, cancel=None):
+            requested.append(key)
+            return 0
+
+        picked = [(self.FAMILY_INDEX[entry.key], entry.key)]
+        rc, _trees, confirms, flashes, runs = self._screen(
+            [picked, wizard._GO_BACK],
+            extra=[patch.object(models, "install_model",
+                                side_effect=capture)])
+        self.assertEqual(rc, 0)
+        self.assertEqual(requested, [entry.key])
+        self.assertEqual(runs, [("Configure SGLang-Omni",
+                                 [f"Install {entry.label}"])])
+        self.assertEqual(confirms, [])
+        self.assertEqual(flashes[-1],
+                         ("SGLang-Omni models updated: 1 installed.", "ok"))
+
+    def test_unchecking_confirms_then_deletes_the_weights(self):
+        from backends.sglomni import wizard
+        entry = ENTRIES[0]
+        rc, _trees, confirms, flashes, runs = self._screen(
+            [[], wizard._GO_BACK], installed=(entry.key,))
+        self.assertEqual(rc, 0)
+        # An empty selection is accepted (allow_empty) and uninstalls
+        # everything installed: one confirm, one removal step.
+        question, kwargs = confirms[0]
+        self.assertEqual(question, "Remove cached weights for 1 model?")
+        self.assertIn(entry.label, kwargs["body"])
+        self.assertEqual(runs, [("Configure SGLang-Omni",
+                                 [f"Delete {entry.label} weights"])])
+        self.assertEqual(flashes[-1],
+                         ("SGLang-Omni models updated: 1 removed.", "ok"))
+
+    def test_uninstall_step_stops_nothing_and_deletes_real_weights(self):
+        # The removal step is uninstall_model itself: with a redirected
+        # HF cache the seeded weight directory is really deleted.
+        from backends.sglomni import wizard
+        entry = ENTRIES[0]
+        with tempfile.TemporaryDirectory() as td:
+            directory = Path(td) / ("models--"
+                                    + entry.repo.replace("/", "--"))
+            (directory / "refs").mkdir(parents=True)
+            (directory / "refs" / "main").write_text("hash\n")
+            (directory / "snapshots" / "abc").mkdir(parents=True)
+            (directory / "snapshots" / "abc" / "weights.safetensors") \
+                .write_bytes(b"x")
+            rc, _trees, _confirms, _flashes, _runs = self._screen(
+                [[], wizard._GO_BACK], installed=(entry.key,),
+                extra=[
+                    patch.object(models, "_hf_cache_dir",
+                                 return_value=Path(td)),
+                    patch.object(models, "_managed_running_repo",
+                                 return_value=None),
+                ])
+            self.assertEqual(rc, 0)
+            self.assertFalse(directory.exists())
+
+    def test_declining_the_uninstall_confirm_runs_nothing(self):
+        from backends.sglomni import wizard
+        entry = ENTRIES[0]
+        rc, trees, confirms, flashes, runs = self._screen(
+            [[], wizard._GO_BACK], installed=(entry.key,), confirm=False)
+        self.assertEqual(rc, 0)
+        # The decline re-opens the tree (second render), no work happens.
+        self.assertEqual(len(trees), 2)
+        self.assertEqual(len(confirms), 1)
+        self.assertEqual(runs, [])
+        self.assertEqual(flashes, [])
+
+    def test_install_without_package_flashes_guidance_instead(self):
+        from backends.sglomni import wizard
+        entry = ENTRIES[0]
+        picked = [(self.FAMILY_INDEX[entry.key], entry.key)]
+        rc, _trees, confirms, flashes, runs = self._screen(
+            [picked, wizard._GO_BACK], package=False)
+        self.assertEqual(rc, 0)
+        self.assertEqual(runs, [])
+        self.assertEqual(confirms, [])
+        self.assertEqual(flashes, [(
+            "Install the SGLang-Omni backend first "
+            "(Configure Backends > Install Backend).", "warn")])
+
+    def test_unchanged_selection_re_opens_the_tree(self):
+        from backends.sglomni import wizard
+        entry = ENTRIES[0]
+        picked = [(self.FAMILY_INDEX[entry.key], entry.key)]
+        rc, trees, _confirms, flashes, runs = self._screen(
+            [picked, wizard._GO_BACK], installed=(entry.key,))
+        self.assertEqual(rc, 0)
+        self.assertEqual(len(trees), 2)
+        self.assertEqual(runs, [])
+        self.assertEqual(flashes, [])
+
+    def test_mixed_selection_removes_before_downloading(self):
+        # Unchecking the installed model and checking another in one
+        # confirm: a single run whose removal step precedes the download.
+        from backends.sglomni import wizard
+        gone, added = ENTRIES[0], ENTRIES[1]
+        picked = [(self.FAMILY_INDEX[added.key], added.key)]
+        rc, _trees, confirms, flashes, runs = self._screen(
+            [picked, wizard._GO_BACK], installed=(gone.key,))
+        self.assertEqual(rc, 0)
+        self.assertEqual(len(confirms), 1)
+        self.assertEqual(runs, [(
+            "Configure SGLang-Omni",
+            [f"Delete {gone.label} weights", f"Install {added.label}"])])
+        self.assertEqual(
+            flashes[-1],
+            ("SGLang-Omni models updated: 1 installed, 1 removed.", "ok"))
+
+
+class SetupWizardTests(unittest.TestCase):
+    """_wizard: the setup tree reconciles models like the Configure screen."""
+
+    FAMILY_INDEX = {option["key"]: index
+                    for index, family in
+                    enumerate(install_tree_families(list(ENTRIES)))
+                    for option in family["options"]}
+
+    def _wizard(self, answers, *, installed=(), package=True, confirm=True):
+        """Run _wizard with scripted tree answers; return (settings, trees,
+        confirms)."""
+        from backends.sglomni import wizard
+        trees, confirms = [], []
+
+        def fake_tree(stdscr, title, families, **kwargs):
+            trees.append((title, kwargs))
+            return answers.pop(0)
+
+        def fake_confirm(scr, question, **kwargs):
+            confirms.append(question)
+            return confirm
+
+        args = wizard.build_parser().parse_args([])
+        with patch.object(wizard, "_preflight", return_value=[]), \
+                patch.object(wizard, "_gpu_warning", return_value=None), \
+                patch.object(wizard, "_is_installed",
+                             return_value=package), \
+                patch.object(models, "installed_keys",
+                             return_value=list(installed)), \
+                patch.object(wizard.tui, "checkbox_tree", fake_tree), \
+                patch.object(wizard.tui, "confirm", fake_confirm), \
+                patch.object(wizard.tui, "flash", lambda *a, **k: None):
+            settings = wizard._wizard(None, args)
+        return settings, trees, confirms
+
+    def test_esc_aborts(self):
+        from backends.sglomni import wizard
+        settings, _trees, confirms = self._wizard([wizard._GO_BACK])
+        self.assertIsNone(settings)
+        self.assertEqual(confirms, [])
+
+    def test_modify_flow_installs_new_and_keeps_installed(self):
+        from backends.sglomni import wizard
+        first, second = ENTRIES[0], ENTRIES[1]
+        # The installed model stays checked (kept as-is); the new one is
+        # added — the diff installs the new one only.
+        picked = [(self.FAMILY_INDEX[first.key], first.key),
+                  (self.FAMILY_INDEX[second.key], second.key)]
+        settings, _trees, confirms = self._wizard(
+            [picked], installed=(first.key,))
+        self.assertEqual(settings["keys"], [second.key])
+        self.assertEqual(settings["uninstall_keys"], [])
+        self.assertEqual(confirms, [])
+
+    def test_unchecking_requires_a_confirm_then_uninstalls(self):
+        from backends.sglomni import wizard
+        first = ENTRIES[0]
+        settings, _trees, confirms = self._wizard([[]],
+                                                  installed=(first.key,))
+        self.assertEqual(settings["keys"], [])
+        self.assertEqual(settings["uninstall_keys"], [first.key])
+        self.assertEqual(confirms, ["Remove cached weights for 1 model?"])
+
+    def test_declined_confirm_re_opens_the_tree(self):
+        from backends.sglomni import wizard
+        first = ENTRIES[0]
+        settings, trees, confirms = self._wizard(
+            [[], wizard._GO_BACK], installed=(first.key,), confirm=False)
+        self.assertIsNone(settings)
+        self.assertEqual(len(trees), 2)
+        self.assertEqual(len(confirms), 1)
+
+    def test_empty_tree_installs_the_package_only(self):
+        from backends.sglomni import wizard
+        settings, _trees, confirms = self._wizard([[]])
+        self.assertEqual(settings["keys"], [])
+        self.assertEqual(settings["uninstall_keys"], [])
+        self.assertEqual(confirms, [])
+
+    def test_steps_remove_before_downloading(self):
+        from backends.sglomni import wizard
+        gone, added = ENTRIES[0], ENTRIES[1]
+        steps = wizard._execute_steps({
+            "do_python": False, "do_install": False,
+            "uninstall_keys": [gone.key], "keys": [added.key]})
+        self.assertEqual([step.title for step in steps],
+                         [f"Delete {gone.label} weights",
+                          f"Install {added.label}"])
+
+    def test_run_tui_drives_the_wizard_in_a_curses_session(self):
+        # The standalone CLI wraps the wizard in its own curses session
+        # (passing a real screen through) and runs the model work as a
+        # console tail afterwards.
+        from backends.sglomni import wizard
+        settings = {"keys": [], "uninstall_keys": [], "do_python": False,
+                    "do_install": False}
+        screens = []
+        with patch.object(wizard, "_wizard",
+                          side_effect=lambda scr, args:
+                              screens.append(scr) or settings), \
+                patch.object(wizard, "_execute", return_value=7) as execute, \
+                patch("curses.wrapper",
+                      side_effect=lambda fn: fn("SCREEN")):
+            rc = wizard.run_tui(wizard.build_parser().parse_args([]))
+        self.assertEqual(rc, 7)
+        self.assertEqual(screens, ["SCREEN"])
+        execute.assert_called_once_with(settings)
+
+
+if __name__ == "__main__":
+    unittest.main()