fix(autotag): keep each download route's error shape and harden the fake download

py/api/autotag_routes.py

@@ -257,13 +257,20 @@ class AutotagRoutesMixin:
         if model_type not in MODELS:
             invalid = {"type": "error", "message": f"Invalid model type: {model_type}"}
             return await self._stream_sse(request, _single(invalid))
-        job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND)
-        if job is not None and job.scope.get("model_type") != model_type:
-            return model_busy_response()
-        if job is None:
-            job, refusal = self._start_model_download(model_type, get_autotag_service())
-            if refusal is not None:
-                return refusal
+        try:
+            job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND)
+            if job is not None and job.scope.get("model_type") != model_type:
+                return model_busy_response()
+            if job is None:
+                job, refusal = self._start_model_download(
+                    model_type, get_autotag_service()
+                )
+                if refusal is not None:
+                    return refusal
+        except Exception:
+            self.logger.exception("Download model error")
+            failed = {"type": "error", "message": _INTERNAL_ERROR}
+            return await self._stream_sse(request, _single(failed))
         return await self._stream_sse(request, self._download_frames(job))
 
     @staticmethod
@@ -300,7 +307,13 @@ class AutotagRoutesMixin:
             return self._bad_request(
                 "model_type must be one of " + ", ".join(sorted(MODELS))
             )
-        job, refusal = self._start_model_download(model_type, get_autotag_service())
+        try:
+            job, refusal = self._start_model_download(model_type, get_autotag_service())
+        except Exception as e:
+            self.logger.exception("Model download start error")
+            return web.json_response(
+                {"success": False, "error": self._public_error(e)}, status=500
+            )
         if refusal is not None:
             return refusal
         return web.json_response({"success": True, "job": job.summary()}, status=202)
@@ -340,10 +353,20 @@ class AutotagRoutesMixin:
                     self._run_model_download, service, model_type, owner
                 ),
             )
-        except JobBusy as refused:
+        except BaseException as refused:
+            # JobBusy, or anything else the registry raised: give the gate back.
             model_gate.release(owner)
-            return None, busy_response(refused)
-        job.task.add_done_callback(lambda _task: model_gate.release(owner))
+            _download_titles.pop(owner, None)
+            if isinstance(refused, JobBusy):
+                return None, busy_response(refused)
+            raise
+
+        def ended(_task):
+            model_gate.release(owner)
+            if model_gate.holder() != owner:  # not while a thread still runs
+                _download_titles.pop(owner, None)
+
+        job.task.add_done_callback(ended)
         return job, None
 
     async def _run_model_download(self, service, model_type, owner, job):

tests/browser/workspace_tagging_fixture.py

@@ -185,14 +185,20 @@ def register_tagging_commands(
     # The fake download (``/__test__/tagging/download``): waits on ``event``,
     # then answers ``result``. Every reset sets the event and puts back the
     # model's raising default, so no test can leave a worker hanging.
-    download = {"event": threading.Event(), "result": True}
+    # ``generation`` grows on every reset: a download a reset released (or one
+    # that timed out) fails without touching the next test's model state.
+    download = {"event": threading.Event(), "result": True, "generation": 0}
 
     def restore_download() -> None:
+        download["generation"] += 1
         download["event"].set()
         model.__dict__.pop("download_model", None)
 
     def fake_download(model_type: str, progress_callback=None) -> bool:
-        download["event"].wait(DOWNLOAD_HOLD_LIMIT_S)
+        event, generation = download["event"], download["generation"]
+        released = event.wait(DOWNLOAD_HOLD_LIMIT_S)
+        if not released or download["generation"] != generation:
+            return False
         if download["result"]:
             model.unavailable = model.unavailable - {model_type}
         return download["result"]

tests/test_api_autotag.py

@@ -1476,6 +1476,13 @@ class DownloadJobTestCase(AutotagAPITestCase):
             "/prompt_manager/jobs/model-download", {"model_type": model_type}
         )
 
+    async def _until(self, condition, limit=10):
+        """Yield to the loop until ``condition()`` holds (fails after ``limit`` s)."""
+        deadline = asyncio.get_running_loop().time() + limit
+        while not condition():
+            self.assertLess(asyncio.get_running_loop().time(), deadline)
+            await asyncio.sleep(0)
+
     async def _ended(self, job_id):
         job = self.api._jobs.get(job_id)
         self.proceed.set()
@@ -1633,6 +1640,8 @@ class TestLegacyDownloadAttaches(DownloadJobTestCase):
         resp = await self.client.post("/prompt_manager/autotag/download/wd14-vit")
         self.assertEqual(resp.status, 200)
         self.assertEqual(self._sse_events(await resp.text()), self.GOLDEN_OK)
+        # The gate is free by the time the stream says the download ended.
+        self.assertIsNone(model_gate.holder())
         jobs = self.api._jobs.list()
         self.assertEqual(
             [(j.kind, j.title, j.state) for j in jobs],
@@ -1651,13 +1660,8 @@ class TestLegacyDownloadAttaches(DownloadJobTestCase):
             self.client.post("/prompt_manager/autotag/download/wd14-vit")
         )
         try:
-            # Wait until the legacy stream is subscribed to the running job.
             job = self.api._jobs.get(summary["id"])
-            for _ in range(300):
-                if job.subscriber_count:
-                    break
-                await asyncio.sleep(0.01)
-            self.assertEqual(job.subscriber_count, 1)
+            await self._until(lambda: job.subscriber_count == 1)
         finally:
             await self._ended(summary["id"])
         resp = await legacy
@@ -1687,11 +1691,20 @@ class TestLegacyDownloadAttaches(DownloadJobTestCase):
         await self._started()
         resp.close()  # the page closed the modal
         (job,) = self.api._jobs.list()
-        for _ in range(300):
-            if not job.subscriber_count:
-                break
-            await asyncio.sleep(0.01)
+        # The closed stream lets go of the job; the download goes on.
+        await self._until(lambda: job.subscriber_count == 0)
         self.assertTrue(job.running)
+        self.assertTrue(model_gate.holder().startswith("model-download:"))
         job = await self._ended(job.id)
         self.assertEqual(job.state, "finished")
         self.assertIsNone(model_gate.holder())
+
+    async def test_unexpected_start_error_keeps_each_route_shape(self):
+        with patch("py.autotag.get_autotag_service", side_effect=RuntimeError("x")):
+            legacy = await self.client.post("/prompt_manager/autotag/download/wd14-vit")
+            events = self._sse_events(await legacy.text())
+            native = await self._start()
+        self.assertEqual(events[-1]["type"], "error")
+        self.assertEqual(native.status, 500)
+        self.assertFalse((await native.json())["success"])
+        self.assertIsNone(model_gate.holder())

tests/test_job_admission.py

@@ -460,7 +460,7 @@ class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase):
     async def test_second_legacy_autotag_batch_keeps_its_json_409(self):
         from py.api.autotag_routes import _BUSY_ERROR
 
-        self._run_job("autotag", "Auto-tag images (legacy page)")
+        self._run_job("autotag", "Tag output images (legacy page)")
         resp = await self.client.request(
             "POST", "/prompt_manager/autotag/start", json={}
         )