fix: sglomni retry contract, booting detection, CLI keys, companion refresh, ffmpeg depHEADmain

2 files changed, 387 insertions, 3 deletions

diff --git a/app/tests/test_backends_sglomni.py b/app/tests/test_backends_sglomni.py
index 54659d6..d4ec12d 100644
--- a/app/tests/test_backends_sglomni.py
+++ b/app/tests/test_backends_sglomni.py
@@ -2,6 +2,7 @@
 
 import io
 import json
+import shutil
 import urllib.error
 import sys
 import tempfile
@@ -105,7 +106,8 @@ class ModelInstallStateTests(unittest.TestCase):
     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 = patch("backends.common.hf_cache_dir",
+                        return_value=self.cache)
         patcher.start()
         self.addCleanup(patcher.stop)
         self.addCleanup(self._tmp.cleanup)
@@ -850,8 +852,8 @@ class ModelsScreenTests(unittest.TestCase):
             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("backends.common.hf_cache_dir",
+                          return_value=Path(td)),
                     patch.object(models, "_managed_running_repo",
                                  return_value=None),
                 ])
@@ -1028,5 +1030,288 @@ class SetupWizardTests(unittest.TestCase):
         execute.assert_called_once_with(settings)
 
 
+class SystemDepTests(unittest.TestCase):
+    """Docs-required system binaries are checked at install time."""
+
+    def test_dac_models_require_ffmpeg(self):
+        # The Fish Audio and ZONOS2 pipelines shell out to ffmpeg; the
+        # install flow must warn like it does for the Qwen entries' sox.
+        for key in ("fish_s2_pro", "zonos2"):
+            with self.subTest(key=key):
+                entry = entry_by_key(key)
+                self.assertEqual(entry.system_dep, "ffmpeg")
+                self.assertIn("ffmpeg", entry.system_hint)
+
+    def test_missing_system_dep_warns(self):
+        with patch("shutil.which", return_value=None):
+            message = models.system_dep_missing(entry_by_key("fish_s2_pro"))
+        self.assertIsNotNone(message)
+        self.assertIn("ffmpeg", message)
+
+    def test_present_system_dep_is_silent(self):
+        with patch("shutil.which", return_value="/usr/bin/ffmpeg"):
+            self.assertIsNone(
+                models.system_dep_missing(entry_by_key("fish_s2_pro")))
+
+    def test_missing_system_dep_warns_during_install(self):
+        out = io.StringIO()
+        with patch("shutil.which", return_value=None), \
+                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), \
+                redirect_stdout(out):
+            rc = models.install_model("fish_s2_pro")
+        self.assertEqual(rc, 0)
+        self.assertIn("ffmpeg", out.getvalue())
+
+
+class UninstallModelServerStopTests(unittest.TestCase):
+    """uninstall_model stops a managed server hosting the model."""
+
+    def test_stops_a_server_hosting_the_model(self):
+        entry = ENTRIES[0]
+        with patch.object(models, "_managed_running_repo",
+                          return_value=entry.repo), \
+                patch.object(models, "delete_model_weights"), \
+                patch.object(servers, "stop") as stop:
+            rc = models.uninstall_model(entry.key)
+        self.assertEqual(rc, 0)
+        stop.assert_called_once_with(constants.SERVER_NAME)
+
+    def test_leaves_a_server_hosting_something_else_alone(self):
+        entry = ENTRIES[0]
+        with patch.object(models, "_managed_running_repo",
+                          return_value=ENTRIES[1].repo), \
+                patch.object(servers, "stop") as stop, \
+                patch.object(models, "delete_model_weights") as weights:
+            rc = models.uninstall_model(entry.key)
+        self.assertEqual(rc, 0)
+        stop.assert_not_called()
+        weights.assert_called_once_with([entry])
+
+
+class WizardUninstallTests(unittest.TestCase):
+    """wizard.uninstall: stop the server, delete weights, remove the venv.
+
+    The venv removal IS the cleanup (the heavyweight CUDA stack is the
+    install): there is deliberately no pip-uninstall phase ahead of it.
+    """
+
+    def _run(self, *, pid=None, cancel_after=None, venv=True, pythons=True,
+             rmtree_removes=True):
+        """Run uninstall with a temp venv tree; return (rc, output,
+        stop_calls, weights, pip_uninstall).
+
+        CANCEL_AFTER N lets the first N cancel checks pass (None = none
+        of them do).
+        """
+        from backends.sglomni import wizard
+        out = io.StringIO()
+        with tempfile.TemporaryDirectory() as td:
+            venv_dir = Path(td) / "sglomni"
+            pythons_dir = Path(td) / "pythons"
+            for make, directory in ((venv, venv_dir),
+                                    (pythons, pythons_dir)):
+                if make:
+                    directory.mkdir()
+            state = {"stops": [], "cancel": 0}
+
+            def fake_stop(name):
+                state["stops"].append(name)
+
+            def fake_cancel(_cancel):
+                state["cancel"] += 1
+                return cancel_after is not None \
+                    and state["cancel"] > cancel_after
+
+            real_rmtree = shutil.rmtree
+
+            def fake_rmtree(path, ignore_errors=False):
+                if rmtree_removes:
+                    real_rmtree(path, ignore_errors=True)
+
+            with patch.object(servers, "pid_for", return_value=pid), \
+                    patch.object(servers, "stop", side_effect=fake_stop), \
+                    patch.object(models, "delete_model_weights") as weights, \
+                    patch.object(common, "pip_uninstall") as pip_uninstall, \
+                    patch.object(wizard, "SGLOMNI_ENV", venv_dir), \
+                    patch.object(envs, "PYTHON_INSTALL_DIR", pythons_dir), \
+                    patch.object(common, "cancel_requested",
+                                 side_effect=fake_cancel), \
+                    patch("shutil.rmtree", side_effect=fake_rmtree), \
+                    redirect_stdout(out):
+                rc = wizard.uninstall()
+                existed = venv_dir.exists() or pythons_dir.exists()
+        return rc, out.getvalue(), state["stops"], weights, \
+            pip_uninstall, existed
+
+    def test_removes_the_venv_tree_without_a_pip_phase(self):
+        rc, out, _stops, weights, pip_uninstall, existed = self._run()
+        self.assertEqual(rc, 0)
+        weights.assert_called_once_with()
+        pip_uninstall.assert_not_called()
+        self.assertFalse(existed)
+
+    def test_stops_only_a_running_managed_server(self):
+        _rc, _out, stops, _w, _p, _e = self._run(pid=123)
+        self.assertEqual(stops, ["sglomni"])
+        _rc, _out, stops, _w, _p, _e = self._run(pid=None)
+        self.assertEqual(stops, [])
+
+    def test_cancel_before_anything_removes_nothing(self):
+        rc, _out, _stops, weights, _p, _e = self._run(cancel_after=0)
+        self.assertEqual(rc, 130)
+        weights.assert_not_called()
+
+    def test_cancel_before_the_venv_removal_keeps_the_dirs(self):
+        rc, _out, _stops, weights, _p, existed = self._run(cancel_after=1)
+        self.assertEqual(rc, 130)
+        weights.assert_called_once_with()
+        self.assertTrue(existed)
+
+    def test_a_stuck_venv_warns_but_the_uninstall_succeeds(self):
+        rc, out, _stops, _w, _p, existed = self._run(rmtree_removes=False)
+        self.assertEqual(rc, 0)
+        self.assertTrue(existed)
+        self.assertIn("Could not fully remove", out)
+
+    def test_missing_venv_dirs_are_fine(self):
+        rc, _out, _stops, _w, _p, _e = self._run(venv=False, pythons=False)
+        self.assertEqual(rc, 0)
+
+
+class WizardUpdateTests(unittest.TestCase):
+    """wizard.update: package upgrade, then a companion refresh.
+
+    The refresh re-runs every installed model's extras (a satisfied pin
+    is a pip no-op), healing version drift the import probe cannot see.
+    """
+
+    def _run(self, *, pid=None, venv=True, pip_rc=0, companions_rc=0,
+             installed=(0, 1), cancel_after=None):
+        from backends.sglomni import wizard
+        out = io.StringIO()
+        entries = [ENTRIES[i] for i in installed]
+        state = {"pip": [], "companions": [], "stops": [], "cancel": 0}
+
+        def fake_stop(name):
+            state["stops"].append(name)
+
+        def fake_cancel(_cancel):
+            state["cancel"] += 1
+            return cancel_after is not None \
+                and state["cancel"] > cancel_after
+
+        def fake_pip(specs, **kwargs):
+            state["pip"].append((list(specs), kwargs))
+            return pip_rc
+
+        def fake_companions(entry, **kwargs):
+            state["companions"].append((entry.key, kwargs.get("force")))
+            return companions_rc
+
+        with patch.object(servers, "pid_for", return_value=pid), \
+                patch.object(servers, "stop", side_effect=fake_stop), \
+                patch.object(envs, "env_exists", return_value=venv), \
+                patch.object(common, "pip_install", side_effect=fake_pip), \
+                patch.object(models, "installed_entries",
+                             return_value=entries), \
+                patch.object(models, "install_companions",
+                             side_effect=fake_companions), \
+                patch.object(common, "cancel_requested",
+                             side_effect=fake_cancel), \
+                redirect_stdout(out):
+            rc = wizard.update()
+        return rc, out.getvalue(), state
+
+    def test_no_venv_is_a_no_op(self):
+        rc, _out, state = self._run(venv=False)
+        self.assertEqual(rc, 0)
+        self.assertEqual(state["pip"], [])
+        self.assertEqual(state["companions"], [])
+
+    def test_stops_the_server_then_upgrades_the_package(self):
+        rc, _out, state = self._run(pid=123)
+        self.assertEqual(rc, 0)
+        self.assertEqual(state["stops"], ["sglomni"])
+        specs, kwargs = state["pip"][0]
+        self.assertEqual(specs, ["sglang-omni"])
+        self.assertTrue(kwargs["upgrade"])
+        self.assertEqual(kwargs["extra_args"], ["--pre"])
+
+    def test_upgrade_failure_skips_the_companion_refresh(self):
+        rc, _out, state = self._run(pip_rc=23)
+        self.assertEqual(rc, 23)
+        self.assertEqual(state["companions"], [])
+
+    def test_refreshes_every_installed_models_companions(self):
+        rc, _out, state = self._run(installed=(0, 4))
+        self.assertEqual(rc, 0)
+        self.assertEqual(state["companions"],
+                         [(ENTRIES[0].key, True), (ENTRIES[4].key, True)])
+
+    def test_companion_failure_warns_but_the_update_succeeds(self):
+        rc, out, _state = self._run(companions_rc=23)
+        self.assertEqual(rc, 0)
+        self.assertIn("companion packages", out)
+
+    def test_cancel_before_pip_removes_nothing(self):
+        rc, _out, state = self._run(cancel_after=0)
+        self.assertEqual(rc, 130)
+        self.assertEqual(state["pip"], [])
+
+
+class NonInteractiveCliTests(unittest.TestCase):
+    """_collect_from_flags: the flag-driven (non-TUI) setup path.
+
+    Positional keys are the space-separated twin of --models: both feed
+    the same deduplicated install list.
+    """
+
+    def _collect(self, argv, *, installed=True):
+        from backends.sglomni import wizard
+        parser = wizard.build_parser()
+        args = parser.parse_args(argv)
+        with patch.object(wizard, "_preflight", return_value=[]), \
+                patch.object(wizard, "_gpu_warning", return_value=None), \
+                patch.object(wizard, "_is_installed",
+                             return_value=installed):
+            return wizard._collect_from_flags(args, parser)
+
+    def test_positional_keys_are_installed(self):
+        settings = self._collect(["higgs_audio_v3_tts"])
+        self.assertEqual(settings["keys"], ["higgs_audio_v3_tts"])
+        self.assertTrue(settings["do_python"])
+        self.assertFalse(settings["do_install"])
+
+    def test_positional_and_models_flags_merge_in_order(self):
+        settings = self._collect(["moss_tts", "--models",
+                                  "higgs_audio_v3_tts,moss_tts"])
+        self.assertEqual(settings["keys"],
+                         ["moss_tts", "higgs_audio_v3_tts"])
+
+    def test_all_installs_every_catalog_model(self):
+        settings = self._collect(["--all"])
+        self.assertEqual(settings["keys"], [entry.key for entry in ENTRIES])
+
+    def test_unknown_key_stops_the_run(self):
+        with self.assertRaises(SystemExit):
+            self._collect(["nope"])
+
+    def test_no_keys_installs_the_package_only(self):
+        settings = self._collect([], installed=False)
+        self.assertEqual(settings["keys"], [])
+        self.assertTrue(settings["do_install"])
+
+    def test_skip_flags_are_honored(self):
+        settings = self._collect(["--skip-python", "--skip-install",
+                                  "higgs_audio_v3_tts"])
+        self.assertFalse(settings["do_python"])
+        self.assertFalse(settings["do_install"])
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/app/tests/test_tts_sglomni.py b/app/tests/test_tts_sglomni.py
index 0e3a7de..c4cdfaf 100644
--- a/app/tests/test_tts_sglomni.py
+++ b/app/tests/test_tts_sglomni.py
@@ -146,6 +146,29 @@ class ConnectHealthTests(unittest.TestCase):
         self.assertIn("not reachable", str(ctx.exception))
         self.assertIn("sgl-omni", str(ctx.exception))
 
+    def test_booting_503_tells_the_user_to_wait(self):
+        # A booting sgl-omni answers /health with 503 + an "unhealthy"
+        # body (urlopen surfaces that as an HTTPError before any JSON
+        # could be inspected) — the message must say wait, not start.
+        with patch("converter.clients.sglomni.urllib.request.urlopen",
+                   side_effect=_http_error(503, '{"status": "unhealthy"}')):
+            with self.assertRaises(RuntimeError) as ctx:
+                SgOmniTTSClient(_DUMMY_CHUNKS, model="higgs_audio_v3_tts")
+        message = str(ctx.exception)
+        self.assertIn("not healthy yet", message)
+        self.assertIn("HTTP 503", message)
+        self.assertIn("booting", message)
+        self.assertNotIn("not reachable", message)
+
+    def test_other_health_errors_name_the_code(self):
+        with patch("converter.clients.sglomni.urllib.request.urlopen",
+                   side_effect=_http_error(404, "<html>nope</html>")):
+            with self.assertRaises(RuntimeError) as ctx:
+                SgOmniTTSClient(_DUMMY_CHUNKS, model="higgs_audio_v3_tts")
+        message = str(ctx.exception)
+        self.assertIn("HTTP 404", message)
+        self.assertIn("Is this an sgl-omni server?", message)
+
     def test_booting_server_raises(self):
         with self.assertRaises(RuntimeError) as ctx:
             self._connect([{"status": "unhealthy"}])
@@ -192,6 +215,7 @@ class PayloadTests(unittest.TestCase):
         client.language = "English"
         client._seed = None
         client._kv_fit = None
+        client._ref_audio_cached = None
         return client
 
     def test_speaker_payload_sends_the_preset_name(self):
@@ -237,6 +261,20 @@ class PayloadTests(unittest.TestCase):
             b"abc")
         self.assertNotIn("ref_text", payload)
 
+    def test_reference_audio_is_encoded_once_per_run(self):
+        # The clip cannot change mid-run: the data URL (or resolved path)
+        # is computed on the first sub-request and reused verbatim.
+        from converter.clients.sglomni import _data_url as real_data_url
+        client = self._make_client("higgs_audio_v3_tts",
+                                   ref_audio=str(self.ref))
+        client.api_url = "http://10.20.30.40:8100"
+        with patch("converter.clients.sglomni._data_url",
+                   wraps=real_data_url) as encode:
+            first = client._request_payload("Hello.")
+            second = client._request_payload("Hello again.")
+        self.assertEqual(encode.call_count, 1)
+        self.assertEqual(first["ref_audio"], second["ref_audio"])
+
     def test_clone_without_reference_sends_no_reference_fields(self):
         client = self._make_client("higgs_audio_v3_tts")
         payload = client._request_payload("Hello.")
@@ -423,6 +461,33 @@ class KvAdmissionTests(unittest.TestCase):
         self.assertEqual(second["max_new_tokens"], 2531)
 
 
+class ErrorClassificationTests(unittest.TestCase):
+    """HTTP status → retry decision: every 4xx envelope is deterministic."""
+
+    def _request_error(self, status, detail):
+        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
+        return client._request_error(status, detail)
+
+    def test_every_4xx_envelope_is_non_retryable(self):
+        # Including types outside the OpenAI-style names: the identical
+        # request fails identically on every attempt.
+        exception = self._request_error(
+            401, json.dumps({"error": {"message": "bad key",
+                                       "type": "AuthenticationError"}}))
+        self.assertIsInstance(exception, NonRetryableTTSError)
+        self.assertIn("bad key", str(exception))
+
+    def test_non_json_4xx_bodies_are_non_retryable(self):
+        exception = self._request_error(400, "plain text refusal")
+        self.assertIsInstance(exception, NonRetryableTTSError)
+        self.assertIn("plain text refusal", str(exception))
+
+    def test_5xx_stays_retryable(self):
+        exception = self._request_error(500, "CUDA out of memory")
+        self.assertNotIsInstance(exception, NonRetryableTTSError)
+        self.assertIn("CUDA out of memory", str(exception))
+
+
 class GenerateChunkTests(unittest.TestCase):
     """Chunk generation: WAV output, sub-chunking, bookkeeping."""
 
@@ -511,6 +576,40 @@ class GenerateChunkTests(unittest.TestCase):
             self.assertIsNone(client.generate_chunk("Hello.", 1))
         self.assertEqual(mock_wav.call_count, 1)
 
+    def test_non_retryable_errors_propagate(self):
+        # Deterministic server errors must reach the retry loop directly
+        # (which skips its remaining attempts and re-raises with the
+        # actionable message), not come back as a generic failed attempt.
+        client = self._make_client()
+        with patch.object(client, "_request_wav",
+                          side_effect=NonRetryableTTSError(
+                              "unknown voice")):
+            with self.assertRaises(NonRetryableTTSError):
+                client.generate_chunk("Hello.", 1)
+
+    def test_retry_loop_skips_remaining_attempts(self):
+        client = self._make_client()
+        with patch.object(client, "_request_wav",
+                          side_effect=NonRetryableTTSError(
+                              "unknown voice")) as mock_wav:
+            with self.assertRaises(NonRetryableTTSError):
+                client.process_chunk_with_retry(1, "Hello.")
+        self.assertEqual(mock_wav.call_count, 1)
+
+    def test_a_non_wav_200_body_fails_the_request(self):
+        # A JSON error body served with HTTP 200 must not be written as
+        # chunk bytes (it would only fail later, confusingly, in the
+        # concat step).
+        client = self._make_client()
+        response = MagicMock()
+        response.__enter__.return_value = response
+        response.read.return_value = b'{"error": {"message": "nope"}}'
+        with patch("converter.clients.sglomni.urllib.request.urlopen",
+                   return_value=response):
+            with self.assertRaises(RuntimeError) as ctx:
+                client._request_wav("Hello.")
+        self.assertIn("not a WAV file", str(ctx.exception))
+
     def test_stale_chunk_files_are_removed(self):
         stale = Path(self._tmp.name) / "chunk_0001.mp3"
         stale.write_bytes(b"old")