feat(autotag): scored suggestions with a configurable count behind a shared model gate

- py/model_gate.py: one process-wide, non-blocking gate for loading,
  swapping, running, downloading and unloading a tagging model. Worker
  calls run shielded; a release requested while one still runs takes
  effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
  gate and answer 409 {code: "model-busy"} while another operation holds
  it; their bodies are otherwise unchanged. /autotag/models gains
  suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
  first, GGUF/HF (name, None) in output order, first `limit` distinct
  case-insensitively; generate_tags keeps its return; the image file is
  closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
  pinned capture, with strict per-engine options, never downloads
  (409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
  10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
  metadata over the inspection budgets no longer blocks it.

py/api/annotation_suggest.py

@@ -0,0 +1,121 @@
+"""Request validation and the worker of the Tag dialog's Generate (#249 B2).
+
+``POST /prompt_manager/image-annotations/suggest`` asks one of legacy's four
+engines for the top N tags of a captured image. Options are validated
+strictly, per engine: WD14 classifiers take two thresholds, the language
+models (GGUF and HuggingFace JoyCaption) take the instructions; every engine
+takes the suggestion count and the GPU switch. Nothing here touches the
+filesystem or the database.
+"""
+
+import math
+
+WD14_MODELS = frozenset({"wd14-swinv2", "wd14-vit"})
+LLM_MODELS = frozenset({"gguf", "hf"})
+MODEL_TYPES = WD14_MODELS | LLM_MODELS
+
+MAX_INSTRUCTION_CHARS = 4000
+MAX_INSTRUCTION_BYTES = 16 * 1024
+
+SUGGEST_FIELDS = frozenset(
+    {"target", "capture_token", "expected_epoch", "model_type", "options"}
+)
+_COMMON = frozenset({"suggestion_limit", "use_gpu"})
+_WD14_REQUIRED = frozenset({"general_threshold", "character_threshold"})
+_LLM_REQUIRED = frozenset({"instructions"})
+
+UNAVAILABLE_HINT = "Download it from the legacy Auto-tag page."
+
+
+class OptionError(ValueError):
+    """A suggest request whose model type or options are not acceptable."""
+
+
+def _threshold(name: str, value) -> float:
+    if type(value) not in (int, float) or not math.isfinite(value):
+        raise OptionError(f"{name} must be a finite number")
+    if not 0 <= value <= 1:
+        raise OptionError(f"{name} must be between 0 and 1")
+    return float(value)
+
+
+def _instructions(value) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise OptionError("instructions must be non-empty text")
+    if len(value) > MAX_INSTRUCTION_CHARS:
+        raise OptionError(
+            f"instructions must be at most {MAX_INSTRUCTION_CHARS} characters"
+        )
+    try:
+        encoded = value.encode("utf-8")
+    except UnicodeEncodeError:
+        raise OptionError("instructions must be valid Unicode text") from None
+    if len(encoded) > MAX_INSTRUCTION_BYTES:
+        raise OptionError("instructions are too long")
+    return value
+
+
+def validate_options(model_type, options, *, default_limit: int, max_limit: int):
+    """Normalised ``(model_type, options)`` or :class:`OptionError`.
+
+    The answer echoes every option the request will run with, the defaults
+    (suggestion count from the settings, GPU on) filled in.
+    """
+    if not isinstance(model_type, str) or model_type not in MODEL_TYPES:
+        raise OptionError("model_type must be one of " + ", ".join(sorted(MODEL_TYPES)))
+    if not isinstance(options, dict):
+        raise OptionError("options must be an object")
+    required = _WD14_REQUIRED if model_type in WD14_MODELS else _LLM_REQUIRED
+    keys = set(options)
+    if not required <= keys or not keys <= required | _COMMON:
+        raise OptionError(
+            f"options for {model_type} take exactly "
+            + ", ".join(sorted(required))
+            + " and optionally "
+            + ", ".join(sorted(_COMMON))
+        )
+    limit = options.get("suggestion_limit", default_limit)
+    if type(limit) is not int or not 1 <= limit <= max_limit:
+        raise OptionError(f"suggestion_limit must be an integer from 1 to {max_limit}")
+    use_gpu = options.get("use_gpu", True)
+    if not isinstance(use_gpu, bool):
+        raise OptionError("use_gpu must be true or false")
+    normalised = {"use_gpu": use_gpu, "suggestion_limit": limit}
+    if model_type in WD14_MODELS:
+        for name in sorted(_WD14_REQUIRED):
+            normalised[name] = _threshold(name, options[name])
+    else:
+        normalised["instructions"] = _instructions(options["instructions"])
+    return model_type, normalised
+
+
+def unavailable_message(model_name: str) -> str:
+    return f"{model_name} is not downloaded. {UNAVAILABLE_HINT}"
+
+
+def run_suggestions(service, capture_path: str, model_type: str, options: dict):
+    """Blocking: load (or swap to) *model_type*, then suggest for the capture.
+
+    The caller holds the model gate for the whole call; nothing is ever
+    downloaded (``load_model`` refuses a model that is not on disk).
+    """
+    if not service.is_model_loaded() or service.get_loaded_model_type() != model_type:
+        service.load_model(model_type, options["use_gpu"])
+    return service.generate_tag_suggestions(
+        capture_path,
+        limit=options["suggestion_limit"],
+        prompt=options.get("instructions"),
+        general_threshold=options.get("general_threshold"),
+        character_threshold=options.get("character_threshold"),
+    )
+
+
+def suggestion_rows(pairs) -> list:
+    """``[{name, confidence}]``; a non-finite score is reported as unknown."""
+    rows = []
+    for name, score in pairs:
+        confidence = None
+        if score is not None and math.isfinite(score):
+            confidence = float(score)
+        rows.append({"name": str(name), "confidence": confidence})
+    return rows

py/api/annotations.py

@@ -16,8 +16,18 @@ from pathlib import Path
 
 from aiohttp import web
 
+from .annotation_suggest import (
+    SUGGEST_FIELDS,
+    OptionError,
+    run_suggestions,
+    suggestion_rows,
+    unavailable_message,
+    validate_options,
+)
+from .autotag_routes import model_busy_response
 from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title
 from .prompts import bad_request
+from .. import model_gate
 
 try:
     from ...database.annotations import (
@@ -87,6 +97,7 @@ _STATUS = {
     "storage-changed": 409,
     "tag-conflict": 409,
     "capture-expired": 409,
+    "model-unavailable": 409,
 }
 _MESSAGES = {
     "source-changed": (
@@ -190,6 +201,9 @@ class AnnotationRoutesMixin:
         routes.get(CAPTURE_URL + "{token}")(self.get_capture_preview)
         routes.delete(CAPTURE_URL + "{token}")(self.release_capture)
         routes.put("/prompt_manager/image-annotations/tags")(self.save_image_tags)
+        routes.post("/prompt_manager/image-annotations/suggest")(
+            self.suggest_image_tags
+        )
 
     # -- capture store lifetime ------------------------------------------------
 
@@ -474,6 +488,121 @@ class AnnotationRoutesMixin:
             {"success": True, **receipt}, headers={"Cache-Control": "no-store"}
         )
 
+    # -- suggest ---------------------------------------------------------------
+
+    async def suggest_image_tags(self, request: web.Request) -> web.Response:
+        """POST: the top N tags of the captured bytes from one legacy engine.
+
+        Runs only on the private capture (pinned for the whole inference),
+        under the shared model gate (``model-busy`` while anything else holds
+        it, including a legacy batch), never downloads (``model-unavailable``)
+        and never writes.
+        """
+        from ..autotag import MODELS, get_autotag_service
+        from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig
+
+        data, error = await self._annotation_body(request, SUGGEST_FIELDS)
+        if error is not None:
+            return error
+        target, error = _target(data["target"])
+        if error is not None:
+            return error
+        try:
+            model_type, options = validate_options(
+                data["model_type"],
+                data["options"],
+                default_limit=AutotagConfig.suggestion_limit(),
+                max_limit=MAX_SUGGESTION_LIMIT,
+            )
+        except OptionError as exc:
+            return bad_request(str(exc))
+        token, epoch = data["capture_token"], data["expected_epoch"]
+        if not isinstance(token, str) or _TOKEN.fullmatch(token) is None:
+            return bad_request("Invalid capture token")
+        if not isinstance(epoch, str) or _UUID.fullmatch(epoch) is None:
+            return bad_request("expected_epoch must be a canonical UUID")
+        index, scope, relative = target
+        root, error = self._scoped_root(index, scope)
+        if error is not None:
+            return error
+        store = self._captures
+        capture = store.acquire(token) if store is not None else None
+        if capture is None:
+            return _coded("capture-expired")
+        started = False
+        try:
+            error = await self._suggest_guards(capture, target, root, epoch)
+            if error is not None:
+                return error
+            service = get_autotag_service()
+            status = await self._run_in_executor(service.get_models_status)
+            if not status.get(model_type, {}).get("downloaded"):
+                name = MODELS.get(model_type, {}).get("name", model_type)
+                return _coded("model-unavailable", unavailable_message(name))
+            owner = model_gate.new_owner("image-annotations")
+            if not model_gate.try_acquire(owner):
+                return model_busy_response()
+            try:
+                started = True
+                pairs = await model_gate.run(
+                    owner,
+                    _pinned_suggestions,
+                    store,
+                    token,
+                    service,
+                    capture,
+                    model_type,
+                    options,
+                )
+            except Exception:
+                self.logger.exception("Tag suggestion failed")
+                return _error(
+                    500,
+                    "suggest-failed",
+                    "Auto-tagging failed. Check server logs for details.",
+                )
+            finally:
+                model_gate.release(owner)
+        finally:
+            if not started:
+                store.release_use(token)
+        return web.json_response(
+            {
+                "success": True,
+                "model_type": model_type,
+                "options": options,
+                "suggestions": suggestion_rows(pairs),
+            },
+            headers={"Cache-Control": "no-store"},
+        )
+
+    async def _suggest_guards(self, capture, target, root, epoch):
+        """The capture must be this target's, of this epoch, and inferable."""
+        index, scope, relative = target
+        expected = os.path.normcase(
+            os.path.join(os.path.realpath(root), *relative.split("/"))
+        )
+        if (
+            capture.root_index != index
+            or capture.scope != scope
+            or capture.locator != expected
+        ):
+            return _coded("capture-expired", _OTHER_FILE)
+        if epoch != capture.epoch:
+            current = await self._run_in_executor(self.db.get_execution_epoch)
+            return _coded("storage-changed", library_epoch=current)
+        if not capture.inferable:
+            return _coded("unsupported-media", capture.inferable_reason)
+        return None
+
+
+def _pinned_suggestions(store, token, service, capture, model_type, options):
+    """Worker thread: infer on the capture, then end its pin (even on error)."""
+    try:
+        return run_suggestions(service, capture.path, model_type, options)
+    finally:
+        store.release_use(token)
+
 
 def _verify_captured_file(store: CaptureStore, capture, root, relative: str):
     """Re-hash the current file, or None when the token names another file.

py/api/autotag_routes.py

@@ -9,11 +9,24 @@ from typing import Optional
 from urllib.parse import quote
 
 from .prompts import thumbnail_url_for
+from .. import model_gate
 
 from aiohttp import web
 
 _INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
 _BUSY_ERROR = "Auto-tagging is already running"
+MODEL_BUSY_ERROR = (
+    "A tagging model is busy with another operation. Try again when it finishes."
+)
+
+
+def model_busy_response() -> web.Response:
+    """``409 model-busy``: another operation holds the shared model gate."""
+    return web.json_response(
+        {"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR},
+        status=409,
+    )
+
 
 # Only these may be handed to the tagging engine or listed by the scanner.
 _MEDIA_EXTENSIONS = frozenset(
@@ -86,6 +99,7 @@ class AutotagRoutesMixin:
         """Get status of available AutoTag models."""
         try:
             from ..autotag import get_autotag_service
+            from ..config import AutotagConfig
 
             service = get_autotag_service()
             models_status = service.get_models_status()
@@ -99,6 +113,7 @@ class AutotagRoutesMixin:
                     "loaded_model_type": service.get_loaded_model_type(),
                     "wd14_general_threshold": service.wd14_general_threshold,
                     "wd14_character_threshold": service.wd14_character_threshold,
+                    "suggestion_limit": AutotagConfig.suggestion_limit(),
                 }
             )
 
@@ -138,6 +153,7 @@ class AutotagRoutesMixin:
     async def download_autotag_model(self, request):
         """Download an AutoTag model with streaming progress (POST only)."""
         model_type = request.match_info.get("model_type")
+        owner = model_gate.new_owner("legacy-download")
 
         async def events():
             try:
@@ -160,8 +176,8 @@ class AutotagRoutesMixin:
                 def progress_callback(status: str, progress: float):
                     pass
 
-                success = await self._run_in_executor(
-                    service.download_model, model_type, progress_callback
+                success = await model_gate.run(
+                    owner, service.download_model, model_type, progress_callback
                 )
                 if success:
                     yield {
@@ -178,7 +194,12 @@ class AutotagRoutesMixin:
                     "message": _INTERNAL_ERROR,
                 }
 
-        return await self._stream_sse(request, events())
+        if not model_gate.try_acquire(owner):
+            return model_busy_response()
+        try:
+            return await self._stream_sse(request, events())
+        finally:
+            model_gate.release(owner)
 
     async def _read_start_params(self, request) -> Optional[dict]:
         """Merge query-string and optional JSON-body parameters for autotag/start.
@@ -239,8 +260,9 @@ class AutotagRoutesMixin:
             return
 
         yield {"type": "progress", "progress": 0, "status": "Loading model..."}
+        run = self._gated_runner(opts.get("gate_owner"))
         try:
-            await self._run_in_executor(service.load_model, model_type, True)
+            await run(service.load_model, model_type, True)
         except Exception:
             self.logger.exception("Failed to load model")
             yield {
@@ -338,7 +360,7 @@ class AutotagRoutesMixin:
                     continue
 
                 try:
-                    tags = await self._run_in_executor(
+                    tags = await run(
                         service.generate_tags,
                         allowed_path,
                         general_threshold=opts["general_threshold"],
@@ -423,10 +445,26 @@ class AutotagRoutesMixin:
                 {"success": False, "error": _BUSY_ERROR}, status=409
             )
         try:
-            return await self._stream_sse(request, events())
+            opts["gate_owner"] = model_gate.new_owner("legacy-batch")
+            if not model_gate.try_acquire(opts["gate_owner"]):
+                return model_busy_response()
+            try:
+                return await self._stream_sse(request, events())
+            finally:
+                model_gate.release(opts["gate_owner"])
         finally:
             _autotag_batch_lock.release()
 
+    def _gated_runner(self, owner):
+        """Executor calls under the model gate *owner* holds (plain when None)."""
+        if owner is None:
+            return self._run_in_executor
+
+        async def run(func, *args, **kwargs):
+            return await model_gate.run(owner, func, *args, **kwargs)
+
+        return run
+
     def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]:
         """Return the real path when it is a media file inside an output directory."""
         if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS:
@@ -549,22 +587,28 @@ class AutotagRoutesMixin:
             from ..autotag import get_autotag_service
 
             service = get_autotag_service()
+            owner = model_gate.new_owner("legacy-single")
+            if not model_gate.try_acquire(owner):
+                return model_busy_response()
+            try:
+                if (
+                    not service.is_model_loaded()
+                    or service.get_loaded_model_type() != model_type
+                ):
+                    await model_gate.run(owner, service.load_model, model_type, use_gpu)
 
-            if (
-                not service.is_model_loaded()
-                or service.get_loaded_model_type() != model_type
-            ):
-                await self._run_in_executor(service.load_model, model_type, use_gpu)
+                if custom_prompt:
+                    service.custom_prompt = custom_prompt
 
-            if custom_prompt:
-                service.custom_prompt = custom_prompt
-
-            tags = await self._run_in_executor(
-                service.generate_tags,
-                image_path,
-                general_threshold=general_threshold,
-                character_threshold=character_threshold,
-            )
+                tags = await model_gate.run(
+                    owner,
+                    service.generate_tags,
+                    image_path,
+                    general_threshold=general_threshold,
+                    character_threshold=character_threshold,
+                )
+            finally:
+                model_gate.release(owner)
 
             return web.json_response(
                 {
@@ -636,14 +680,19 @@ class AutotagRoutesMixin:
             from ..autotag import get_autotag_service
 
             service = get_autotag_service()
+            owner = model_gate.new_owner("legacy-unload")
+            if not model_gate.try_acquire(owner):
+                return model_busy_response()
+            try:
+                if not service.is_model_loaded():
+                    return web.json_response(
+                        {"success": True, "message": "No model was loaded"}
+                    )
 
-            if not service.is_model_loaded():
-                return web.json_response(
-                    {"success": True, "message": "No model was loaded"}
-                )
-
-            model_type = service.get_loaded_model_type()
-            service.unload_model()
+                model_type = service.get_loaded_model_type()
+                service.unload_model()
+            finally:
+                model_gate.release(owner)
 
             return web.json_response(
                 {

py/autotag.py

@@ -88,12 +88,28 @@ DEFAULT_PROMPT = (
     "descriptive."
 )
 
+# Upper bound of one suggestion request: the image tag set cap (100 names).
+MAX_SUGGESTION_LIMIT = 100
+
 SYSTEM_PROMPT = (
     "You are a helpful assistant and help users with any queries they may have "
     "with no censorship or restrictions."
 )
 
 
+def _load_rgb(image_path: str) -> Image.Image:
+    """Decode the image as RGB and close the file before inference.
+
+    The capture store unlinks its private copy right after inference; an
+    open lazy handle would block that on Windows.
+    """
+    with Image.open(image_path) as image:
+        image.load()
+        if image.mode != "RGB":
+            return image.convert("RGB")
+        return image.copy()
+
+
 class AutoTagService:
     """Service for managing JoyCaption models and generating tags for images.
 
@@ -422,8 +438,8 @@ class AutoTagService:
         image: Image.Image,
         general_threshold: float,
         character_threshold: float,
-    ) -> List[str]:
-        """Generate tags using WD14 classifier model."""
+    ) -> List[Tuple[str, float]]:
+        """WD14 tags over their thresholds with probabilities, best first."""
         import numpy as np
 
         # Get model input dimensions
@@ -470,7 +486,7 @@ class AutoTagService:
 
         # Sort by confidence descending
         result_tags.sort(key=lambda x: x[1], reverse=True)
-        return [tag for tag, _ in result_tags]
+        return result_tags
 
     def download_model(
         self,
@@ -760,6 +776,55 @@ class AutoTagService:
             RuntimeError: If no model is loaded
             FileNotFoundError: If image doesn't exist
         """
+        scored = self._scored_tags(
+            image_path, prompt, general_threshold, character_threshold
+        )
+        return [name for name, _ in scored]
+
+    def generate_tag_suggestions(
+        self,
+        image_path: str,
+        *,
+        limit: int,
+        prompt: Optional[str] = None,
+        general_threshold: Optional[float] = None,
+        character_threshold: Optional[float] = None,
+    ) -> List[Tuple[str, Optional[float]]]:
+        """The top ``limit`` suggestions as ``(name, confidence)`` pairs.
+
+        WD14 pairs carry the real probability, best first; language models
+        (GGUF, HF) have no score, so their tags keep the output order with
+        ``None``. Names are distinct case-insensitively (first one wins).
+        ``prompt`` is used for this call only; ``custom_prompt`` is never
+        changed.
+        """
+        if type(limit) is not int or not 1 <= limit <= MAX_SUGGESTION_LIMIT:
+            raise ValueError(
+                f"limit must be an integer from 1 to {MAX_SUGGESTION_LIMIT}"
+            )
+        scored = self._scored_tags(
+            image_path, prompt, general_threshold, character_threshold
+        )
+        suggestions: List[Tuple[str, Optional[float]]] = []
+        seen = set()
+        for name, score in scored:
+            key = name.casefold()
+            if key in seen:
+                continue
+            seen.add(key)
+            suggestions.append((name, score))
+            if len(suggestions) == limit:
+                break
+        return suggestions
+
+    def _scored_tags(
+        self,
+        image_path: str,
+        prompt: Optional[str],
+        general_threshold: Optional[float],
+        character_threshold: Optional[float],
+    ) -> List[Tuple[str, Optional[float]]]:
+        """Run the loaded model: WD14 ``(name, probability)``, LLM ``(name, None)``."""
         if self._tagger is None:
             raise RuntimeError("No model loaded. Call load_model() first.")
 
@@ -767,18 +832,14 @@ class AutoTagService:
             raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path))
 
         use_prompt = prompt or self._custom_prompt
-
-        # Load image
-        image = Image.open(image_path)
-        if image.mode != "RGB":
-            image = image.convert("RGB")
+        image = _load_rgb(image_path)
 
         # Generate based on model type
         model_type, tagger_obj = self._tagger
 
         if model_type == "gguf":
             raw_tags = self._generate_gguf(tagger_obj, image, use_prompt)
-            return self._parse_tags(raw_tags)
+            return [(tag, None) for tag in self._parse_tags(raw_tags)]
         elif model_type == "wd14":
             session, tags_list = tagger_obj
             gen_thresh = (
@@ -799,7 +860,7 @@ class AutoTagService:
             raw_tags = self._generate_hf(
                 model, processor, device, compute_dtype, image, use_prompt
             )
-            return self._parse_tags(raw_tags)
+            return [(tag, None) for tag in self._parse_tags(raw_tags)]
 
     def _generate_gguf(self, model, image: Image.Image, prompt: str) -> str:
         """Generate tags using GGUF model."""

py/config.py

@@ -465,6 +465,59 @@ class IntegrationConfig:
             cls.CIVITAI_API_KEY = lora["civitai_api_key"]
 
 
+# Number of tag suggestions the Tag dialogs request by default (D68). The cap
+# is the image tag set limit: more suggestions could never all be kept.
+DEFAULT_SUGGESTION_LIMIT = 10
+MAX_SUGGESTION_LIMIT = 100
+
+
+class AutotagConfig:
+    """Persisted tagging defaults under ``autotag.defaults`` in config.json.
+
+    ``suggestion_limit`` is how many suggestions Generate returns (the top N
+    by score), an integer in ``1..MAX_SUGGESTION_LIMIT``. A hand-edited value
+    outside that range, or of another type, falls back to the default with a
+    warning rather than failing the load.
+    """
+
+    SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
+
+    @classmethod
+    def suggestion_limit(cls) -> int:
+        return cls.SUGGESTION_LIMIT
+
+    @staticmethod
+    def valid_suggestion_limit(value: Any) -> bool:
+        return (
+            type(value) is int and 1 <= value <= MAX_SUGGESTION_LIMIT
+        )  # bool, float and str are rejected
+
+    @classmethod
+    def get_config(cls) -> Dict[str, Any]:
+        return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}}
+
+    @classmethod
+    def update_config(cls, new_config: Any) -> None:
+        if not isinstance(new_config, dict):
+            config_logger.warning("Ignoring autotag settings that are not an object")
+            return
+        defaults = new_config.get("defaults", {})
+        if not isinstance(defaults, dict):
+            config_logger.warning("Ignoring autotag.defaults that is not an object")
+            return
+        if "suggestion_limit" not in defaults:
+            return
+        value = defaults["suggestion_limit"]
+        if cls.valid_suggestion_limit(value):
+            cls.SUGGESTION_LIMIT = value
+            return
+        config_logger.warning(
+            "autotag.defaults.suggestion_limit must be an integer from 1 to "
+            f"{MAX_SUGGESTION_LIMIT}; using {DEFAULT_SUGGESTION_LIMIT}"
+        )
+        cls.SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
+
+
 class PromptManagerConfig:
     """Main configuration class for PromptManager core functionality.
 
@@ -566,6 +619,7 @@ class PromptManagerConfig:
             },
             "gallery": GalleryConfig.get_config(),
             "integrations": IntegrationConfig.get_config(),
+            "autotag": AutotagConfig.get_config(),
         }
 
     @classmethod
@@ -706,6 +760,9 @@ class PromptManagerConfig:
         if "integrations" in new_config:
             IntegrationConfig.update_config(new_config["integrations"])
 
+        if "autotag" in new_config:
+            AutotagConfig.update_config(new_config["autotag"])
+
 
 # Load configuration on import
 try:

py/model_gate.py

@@ -0,0 +1,148 @@
+"""One process-wide, non-blocking gate around the tagging models (#249 B2).
+
+Loading, swapping, running, downloading and unloading a tagging model all
+happen under this gate, whether the caller is a legacy Auto-tag route, the
+Tag dialog's Generate or (later) a bulk job. A caller that cannot take the
+gate is refused at once (``409 model-busy``), never queued.
+
+Work that runs in a worker thread is admitted with :meth:`ModelGate.run`. A
+release requested while such work still runs (a client disconnect cancels
+the waiting coroutine, not the thread) takes effect only when the last
+admitted call really exits, so nobody can swap the model under a running
+inference.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+import uuid
+from typing import Any, Callable, Optional
+
+
+class GateNotHeld(RuntimeError):
+    """Raised when work is admitted for an owner that does not hold the gate."""
+
+
+class ModelGate:
+    """A single-owner gate whose release waits for admitted worker calls."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._holder: Optional[str] = None
+        self._inflight = 0
+        self._releasing = False
+
+    def try_acquire(self, owner: str) -> bool:
+        """Take the gate for *owner*; False at once when anyone holds it."""
+        with self._lock:
+            if self._holder is not None:
+                return False
+            self._holder = owner
+            self._inflight = 0
+            self._releasing = False
+            return True
+
+    def release(self, owner: str) -> None:
+        """Let go; deferred until admitted work of *owner* has exited."""
+        with self._lock:
+            if self._holder != owner:
+                return
+            if self._inflight:
+                self._releasing = True
+            else:
+                self._clear_locked()
+
+    def holder(self) -> Optional[str]:
+        with self._lock:
+            return self._holder
+
+    def begin(self, owner: str) -> None:
+        """Admit one worker call; the holder must still hold the gate."""
+        with self._lock:
+            if self._holder != owner or self._releasing:
+                raise GateNotHeld(f"{owner} does not hold the model gate")
+            self._inflight += 1
+
+    def end(self, owner: str) -> None:
+        """One admitted call exited; completes a deferred release."""
+        with self._lock:
+            self._end_locked(owner)
+
+    async def run(self, owner: str, func: Callable[..., Any], *args, **kwargs):
+        """Run ``func`` in the default executor while *owner* holds the gate.
+
+        The executor future is shielded: cancelling the caller never cancels
+        a call that has not started yet (which would leave it admitted
+        forever) and never ends the admission of a call that is running.
+        """
+        self.begin(owner)
+        state = {"started": False, "finished": False}
+
+        def finish() -> None:
+            with self._lock:
+                if not state["finished"]:
+                    state["finished"] = True
+                    self._end_locked(owner)
+
+        def work():
+            with self._lock:
+                state["started"] = True
+            try:
+                return func(*args, **kwargs)
+            finally:
+                finish()
+
+        def never_ran(future) -> None:
+            # Only a loop shutdown cancels the shielded future; a call that
+            # never started must still give its admission back.
+            if future.cancelled():
+                with self._lock:
+                    skipped = not state["started"] and not state["finished"]
+                    if skipped:
+                        state["finished"] = True
+                        self._end_locked(owner)
+
+        try:
+            future = asyncio.get_running_loop().run_in_executor(None, work)
+        except BaseException:
+            finish()
+            raise
+        future.add_done_callback(never_ran)
+        return await asyncio.shield(future)
+
+    def _end_locked(self, owner: str) -> None:
+        if self._holder != owner or self._inflight == 0:
+            return
+        self._inflight -= 1
+        if self._releasing and self._inflight == 0:
+            self._clear_locked()
+
+    def _clear_locked(self) -> None:
+        self._holder = None
+        self._inflight = 0
+        self._releasing = False
+
+
+_GATE = ModelGate()
+
+
+def new_owner(kind: str) -> str:
+    """A unique owner name such as ``legacy-single:<uuid>``."""
+    return f"{kind}:{uuid.uuid4()}"
+
+
+def try_acquire(owner: str) -> bool:
+    return _GATE.try_acquire(owner)
+
+
+def release(owner: str) -> None:
+    _GATE.release(owner)
+
+
+def holder() -> Optional[str]:
+    return _GATE.holder()
+
+
+async def run(owner: str, func: Callable[..., Any], *args, **kwargs):
+    return await _GATE.run(owner, func, *args, **kwargs)

tests/browser/server.py

@@ -80,6 +80,21 @@ class TestModel:
     def generate_tags(self, image_path: str, **kwargs: object) -> list[str]:
         return ["generated", Path(image_path).parent.parent.name]
 
+    def generate_tag_suggestions(
+        self, image_path: str, *, limit: int, **kwargs: object
+    ) -> list[tuple[str, float | None]]:
+        """Fixed scored suggestions (WD14) or unscored ones (language models)."""
+        scored = self.loaded is not None and self.loaded.startswith("wd14")
+        names = ["sandstone", "desert", "arch", "sky", "rock", "dune", "canyon"]
+        names += [f"tag_{index:02d}" for index in range(100)]
+        return [
+            (name, round(0.98 - index * 0.008, 3) if scored else None)
+            for index, name in enumerate(names[:limit])
+        ]
+
+    def download_model(self, model_type: str, progress_callback=None) -> bool:
+        raise AssertionError("browser tests never download a model")
+
 
 REVEAL_OUTCOMES = ("opened", "unavailable", "fail")
 

tests/test_api_annotation_suggest.py

@@ -0,0 +1,537 @@
+"""Scored top-N tag suggestions for a captured image (#249 B2).
+
+``POST /prompt_manager/image-annotations/suggest`` runs the selected legacy
+engine on the private capture (never the current file), behind the one
+process-wide model gate shared with the legacy Auto-tag routes. It never
+downloads a model and never writes. No real model is loaded: the engine is a
+fake service.
+"""
+
+import asyncio
+import json
+import math
+import os
+import threading
+import unittest
+import uuid
+from unittest.mock import MagicMock, patch
+
+try:
+    from test_api_annotations import AnnotationAPITestCase
+    from test_api_images import make_png
+except ImportError:
+    from tests.test_api_annotations import AnnotationAPITestCase
+    from tests.test_api_images import make_png
+
+try:  # guarded so a missing module fails assertions, not test loading
+    from py import model_gate
+except ImportError:
+    model_gate = None
+
+from py import config as config_module
+
+AutotagConfig = getattr(config_module, "AutotagConfig", None)
+
+SUGGEST = "/prompt_manager/image-annotations/suggest"
+WAIT = 30.0
+WD14 = {"general_threshold": 0.35, "character_threshold": 0.85}
+LLM = {"instructions": "Generate tags."}
+UNAVAILABLE = "Download it from the legacy Auto-tag page."
+
+SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [
+    (f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150)
+]
+
+
+class FakeTagger:
+    """Stand-in for ``AutoTagService`` with the suggestion entry point."""
+
+    default_prompt = "server default prompt"
+    custom_prompt = "legacy prompt"
+    wd14_general_threshold = 0.35
+    wd14_character_threshold = 0.85
+
+    def __init__(self):
+        self.models_config = {
+            "wd14-swinv2": {},
+            "wd14-vit": {},
+            "gguf": {},
+            "hf": {},
+        }
+        self.downloaded = {name: True for name in self.models_config}
+        self.loaded = None
+        self.load_calls = []
+        self.unload_calls = 0
+        self.calls = []
+        self.read_paths = []
+        self.read_bytes = []
+        self.download_model = MagicMock(side_effect=AssertionError("downloaded"))
+        self.load_error = None
+        self.block = None  # (entered, proceed) events
+
+    def get_models_status(self):
+        return {
+            name: {"downloaded": self.downloaded[name], "model_path": None}
+            for name in self.models_config
+        }
+
+    def is_model_loaded(self):
+        return self.loaded is not None
+
+    def get_loaded_model_type(self):
+        return self.loaded
+
+    def load_model(self, model_type, use_gpu=True):
+        self.load_calls.append((model_type, use_gpu))
+        if self.load_error is not None:
+            raise self.load_error
+        self.loaded = model_type
+        return True
+
+    def unload_model(self):
+        self.unload_calls += 1
+        self.loaded = None
+
+    def generate_tags(self, image_path, **kwargs):
+        return ["legacy"]
+
+    def generate_tag_suggestions(
+        self,
+        image_path,
+        *,
+        limit,
+        prompt=None,
+        general_threshold=None,
+        character_threshold=None,
+    ):
+        self.read_paths.append(image_path)
+        with open(image_path, "rb") as handle:
+            self.read_bytes.append(handle.read())
+        self.calls.append(
+            {
+                "limit": limit,
+                "prompt": prompt,
+                "general_threshold": general_threshold,
+                "character_threshold": character_threshold,
+                "model": self.loaded,
+            }
+        )
+        if self.block is not None:
+            entered, proceed = self.block
+            entered.set()
+            proceed.wait(WAIT)
+        if self.loaded.startswith("wd14"):
+            return list(SCORED[:limit])
+        return [(name, None) for name, _ in SCORED[:limit]]
+
+
+class SuggestTestCase(AnnotationAPITestCase):
+    async def get_application(self):
+        app = await super().get_application()
+        self.assertIsNotNone(model_gate, "py.model_gate is not implemented")
+        self.assertIsNotNone(AutotagConfig, "AutotagConfig is not implemented")
+        self.tagger = FakeTagger()
+        patcher = patch("py.autotag.get_autotag_service", return_value=self.tagger)
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        AutotagConfig.SUGGESTION_LIMIT = 10
+        self.addCleanup(setattr, AutotagConfig, "SUGGESTION_LIMIT", 10)
+        return app
+
+    async def tearDownAsync(self):
+        leaked = model_gate.holder() if model_gate is not None else None
+        if leaked is not None:
+            model_gate.release(leaked)
+        await super().tearDownAsync()
+        self.assertIsNone(leaked, "the suggest route left the model gate held")
+
+    async def observed(self, path="sub/a.png", size=(8, 6)):
+        make_png(self.output_dir / path, size=size)
+        return await self.observe_ok(path)
+
+    def suggest_body(self, observed, model_type="wd14-vit", options=None, **changes):
+        if options is None:
+            options = dict(WD14 if model_type.startswith("wd14") else LLM)
+        return {
+            "target": dict(observed["target"]),
+            "capture_token": observed["capture"]["token"],
+            "expected_epoch": observed["library_epoch"],
+            "model_type": model_type,
+            "options": options,
+            **changes,
+        }
+
+    async def suggest(self, body):
+        response = await self.client.post(SUGGEST, json=body)
+        return response.status, await self.json_of(response)
+
+    async def suggest_ok(self, body):
+        status, data = await self.suggest(body)
+        self.assertEqual(status, 200, data)
+        return data
+
+    def capture_path(self, observed):
+        return self.api._captures.get(observed["capture"]["token"]).path
+
+    async def wait_until_free(self):
+        for _ in range(3000):
+            if model_gate.holder() is None:
+                return
+            await asyncio.sleep(0.01)
+        self.fail(f"model gate still held by {model_gate.holder()}")
+
+
+class TestSuggestions(SuggestTestCase):
+    async def test_wd14_suggestions_are_scored_top_n_from_the_capture(self):
+        observed = await self.observed()
+
+        body = await self.suggest_ok(self.suggest_body(observed))
+
+        self.assertEqual(
+            body["suggestions"][:2],
+            [
+                {"name": "sandstone", "confidence": 0.98},
+                {"name": "desert", "confidence": 0.96},
+            ],
+        )
+        self.assertEqual(len(body["suggestions"]), 10)
+        self.assertEqual(body["model_type"], "wd14-vit")
+        self.assertEqual(
+            body["options"],
+            {
+                "general_threshold": 0.35,
+                "character_threshold": 0.85,
+                "use_gpu": True,
+                "suggestion_limit": 10,
+            },
+        )
+        self.assertEqual(self.tagger.read_paths, [self.capture_path(observed)])
+        self.assertEqual(
+            self.tagger.calls[0],
+            {
+                "limit": 10,
+                "prompt": None,
+                "general_threshold": 0.35,
+                "character_threshold": 0.85,
+                "model": "wd14-vit",
+            },
+        )
+        self.assertEqual(self.tagger.load_calls, [("wd14-vit", True)])
+        self.tagger.download_model.assert_not_called()
+        self.assertIsNone(model_gate.holder())
+        self.assert_no_paths(body)
+
+    async def test_suggestions_keep_the_requested_limit(self):
+        observed = await self.observed()
+        for limit in (1, 10, 100):
+            with self.subTest(limit=limit):
+                options = {**WD14, "suggestion_limit": limit}
+                body = await self.suggest_ok(
+                    self.suggest_body(observed, options=options)
+                )
+                self.assertEqual(len(body["suggestions"]), limit)
+                self.assertEqual(body["options"]["suggestion_limit"], limit)
+        # Omitted: the persisted autotag.defaults.suggestion_limit.
+        AutotagConfig.SUGGESTION_LIMIT = 3
+        body = await self.suggest_ok(self.suggest_body(observed))
+        self.assertEqual(len(body["suggestions"]), 3)
+        self.assertEqual(body["options"]["suggestion_limit"], 3)
+
+    async def test_llm_suggestions_send_instructions_and_null_confidence(self):
+        observed = await self.observed()
+        self.tagger.custom_prompt = "legacy prompt"
+        for model_type in ("gguf", "hf"):
+            with self.subTest(model_type=model_type):
+                options = {"instructions": f"tags for {model_type}", "use_gpu": False}
+                body = await self.suggest_ok(
+                    self.suggest_body(observed, model_type, options)
+                )
+                self.assertEqual(body["model_type"], model_type)
+                self.assertEqual(body["options"], {**options, "suggestion_limit": 10})
+                self.assertEqual(
+                    body["suggestions"][:2],
+                    [
+                        {"name": "sandstone", "confidence": None},
+                        {"name": "desert", "confidence": None},
+                    ],
+                )
+                call = self.tagger.calls[-1]
+                self.assertEqual(call["prompt"], f"tags for {model_type}")
+                self.assertIsNone(call["general_threshold"])
+                self.assertEqual(self.tagger.load_calls[-1], (model_type, False))
+        self.assertEqual(self.tagger.custom_prompt, "legacy prompt")
+
+    async def test_loaded_model_is_reused_and_another_is_swapped_in(self):
+        observed = await self.observed()
+        self.tagger.loaded = "wd14-vit"
+        await self.suggest_ok(self.suggest_body(observed))
+        self.assertEqual(self.tagger.load_calls, [])
+        await self.suggest_ok(self.suggest_body(observed, "wd14-swinv2"))
+        self.assertEqual(self.tagger.load_calls, [("wd14-swinv2", True)])
+        self.assertEqual(self.tagger.calls[-1]["model"], "wd14-swinv2")
+
+    async def test_suggest_uses_capture_bytes_not_current_file(self):
+        observed = await self.observed()
+        original = (self.output_dir / "sub" / "a.png").read_bytes()
+        make_png(self.output_dir / "sub" / "a.png", size=(9, 9), color=(0, 0, 255))
+        self.assertNotEqual((self.output_dir / "sub" / "a.png").read_bytes(), original)
+
+        await self.suggest_ok(self.suggest_body(observed))
+
+        capture_path = self.capture_path(observed)
+        self.assertEqual(self.tagger.read_paths, [capture_path])
+        self.assertEqual(self.tagger.read_bytes, [original])
+        self.assertNotEqual(
+            os.path.normcase(capture_path),
+            os.path.normcase(str(self.output_dir / "sub" / "a.png")),
+        )
+
+    async def test_suggest_writes_nothing(self):
+        observed = await self.observed()
+        before = self.db_snapshot()
+        rows = self.bytes_rows()
+
+        await self.suggest_ok(self.suggest_body(observed))
+        await self.suggest_ok(self.suggest_body(observed, "gguf"))
+
+        self.assertEqual(self.db_snapshot(), before)
+        self.assertEqual(self.bytes_rows(), rows)
+        # The capture stays usable for Save after Generate.
+        self.assertIsNotNone(self.api._captures.get(observed["capture"]["token"]))
+
+
+class TestModelAvailability(SuggestTestCase):
+    async def test_suggest_never_downloads_and_reports_model_unavailable(self):
+        observed = await self.observed()
+        self.tagger.downloaded["wd14-vit"] = False
+
+        status, body = await self.suggest(self.suggest_body(observed))
+
+        self.assertEqual((status, body["code"]), (409, "model-unavailable"))
+        self.assertIn(UNAVAILABLE, body["error"])
+        self.assertIn("WD14 ViT", body["error"])
+        self.tagger.download_model.assert_not_called()
+        self.assertEqual(self.tagger.load_calls, [])
+        self.assertEqual(self.tagger.calls, [])
+        self.assertIsNone(model_gate.holder())
+
+    async def test_load_failure_is_a_generic_error_and_frees_everything(self):
+        observed = await self.observed()
+        self.tagger.load_error = RuntimeError(f"cannot load {self.output_dir}")
+
+        status, body = await self.suggest(self.suggest_body(observed))
+
+        self.assertEqual((status, body["code"]), (500, "suggest-failed"))
+        self.assert_no_paths(body)
+        self.assertIsNone(model_gate.holder())
+        entry = self.api._captures._entries[observed["capture"]["token"]]
+        self.assertEqual(entry.uses, 0)
+
+
+class TestModelGate(SuggestTestCase):
+    async def _blocked_suggest(self, observed, model_type="wd14-vit"):
+        entered, proceed = threading.Event(), threading.Event()
+        self.tagger.block = (entered, proceed)
+        task = asyncio.ensure_future(
+            self.client.post(SUGGEST, json=self.suggest_body(observed, model_type))
+        )
+        await asyncio.get_running_loop().run_in_executor(None, entered.wait, WAIT)
+        self.assertTrue(entered.is_set())
+        return task, proceed
+
+    async def test_suggest_reports_busy_while_legacy_batch_runs(self):
+        observed = await self.observed()
+        prompt_id = self._save_prompt("batch prompt")
+        self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png"))
+        started, proceed = threading.Event(), threading.Event()
+
+        def slow_generate(image_path, **kwargs):
+            started.set()
+            proceed.wait(WAIT)
+            return ["batch"]
+
+        self.tagger.generate_tags = slow_generate
+        batch = asyncio.ensure_future(
+            self.client.post("/prompt_manager/autotag/start?model_type=wd14-vit")
+        )
+        try:
+            await asyncio.get_running_loop().run_in_executor(None, started.wait, WAIT)
+            status, body = await self.suggest(self.suggest_body(observed, "gguf"))
+            self.assertEqual((status, body["code"]), (409, "model-busy"))
+            self.assertEqual(self.tagger.calls, [])
+            self.assertEqual(self.tagger.loaded, "wd14-vit")  # not swapped
+        finally:
+            proceed.set()
+            response = await batch
+            await response.read()
+        await self.wait_until_free()
+        await self.suggest_ok(self.suggest_body(observed))
+
+    async def test_legacy_unload_returns_model_busy_while_native_generate_runs(self):
+        observed = await self.observed()
+        task, proceed = await self._blocked_suggest(observed)
+        try:
+            response = await self.client.post("/prompt_manager/autotag/unload")
+            self.assertEqual(response.status, 409)
+            self.assertEqual((await response.json())["code"], "model-busy")
+            response = await self.client.post(
+                "/prompt_manager/autotag/single",
+                json={"root_index": 0, "path": "sub/a.png", "model_type": "gguf"},
+            )
+            self.assertEqual(response.status, 409)
+            status, body = await self.suggest(self.suggest_body(observed, "gguf"))
+            self.assertEqual((status, body["code"]), (409, "model-busy"))
+            self.assertEqual(self.tagger.unload_calls, 0)
+            self.assertEqual(self.tagger.loaded, "wd14-vit")
+        finally:
+            proceed.set()
+            response = await task
+        self.assertEqual(response.status, 200)
+        await self.wait_until_free()
+        self.assertEqual([call["model"] for call in self.tagger.calls], ["wd14-vit"])
+
+    async def test_gate_held_after_client_disconnect_until_executor_finishes(self):
+        observed = await self.observed()
+        token = observed["capture"]["token"]
+        task, proceed = await self._blocked_suggest(observed)
+        task.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await task
+        try:
+            # The client is gone; the inference still runs, so the gate and
+            # the capture pin stay until the worker thread exits.
+            self.assertIsNotNone(model_gate.holder())
+            response = await self.client.post("/prompt_manager/autotag/unload")
+            self.assertEqual(response.status, 409)
+            self.assertEqual(self.api._captures._entries[token].uses, 1)
+        finally:
+            proceed.set()
+        await self.wait_until_free()
+        for _ in range(3000):
+            if self.api._captures._entries[token].uses == 0:
+                break
+            await asyncio.sleep(0.01)
+        self.assertEqual(self.api._captures._entries[token].uses, 0)
+        response = await self.client.post("/prompt_manager/autotag/unload")
+        self.assertEqual(response.status, 200)
+
+
+class TestRequestValidation(SuggestTestCase):
+    async def test_strict_option_validation(self):
+        observed = await self.observed()
+        bad_options = {
+            "wd14-vit": [
+                {**WD14, "general_threshold": True},
+                {**WD14, "general_threshold": math.nan},
+                {**WD14, "general_threshold": math.inf},
+                {**WD14, "general_threshold": "0.3"},
+                {**WD14, "general_threshold": 1.5},
+                {**WD14, "character_threshold": -0.1},
+                {"general_threshold": 0.35},
+                {**WD14, "instructions": "x"},
+                {**WD14, "unknown": 1},
+                {**WD14, "use_gpu": "yes"},
+                {**WD14, "suggestion_limit": 0},
+                {**WD14, "suggestion_limit": 101},
+                {**WD14, "suggestion_limit": 2.5},
+                {**WD14, "suggestion_limit": "10"},
+                {**WD14, "suggestion_limit": True},
+            ],
+            "gguf": [
+                {**LLM, "general_threshold": 0.35},
+                {**LLM, "character_threshold": 0.85},
+                {"instructions": "x" * 4001},
+                {"instructions": "\ud800"},  # not encodable as UTF-8
+                {"instructions": "   "},
+                {"instructions": 5},
+                {},
+                {**LLM, "suggestion_limit": 0},
+            ],
+        }
+        for model_type, cases in bad_options.items():
+            for options in cases:
+                with self.subTest(model_type=model_type, options=repr(options)[:80]):
+                    status, body = await self.suggest(
+                        self.suggest_body(observed, model_type, options)
+                    )
+                    self.assertEqual(status, 400, body)
+        for body in (
+            {**self.suggest_body(observed), "extra": 1},
+            {k: v for k, v in self.suggest_body(observed).items() if k != "options"},
+            self.suggest_body(observed, options=[]),
+            self.suggest_body(observed, model_type="ollama"),
+            self.suggest_body(observed, model_type=7, options=dict(WD14)),
+            self.suggest_body(observed, capture_token="../x"),
+            self.suggest_body(observed, expected_epoch="nope"),
+            self.suggest_body(observed, target={"root_index": 0}),
+        ):
+            with self.subTest(body=json.dumps(body)[:120]):
+                status, _ = await self.suggest(body)
+                self.assertEqual(status, 400)
+        # The edges that are allowed.
+        await self.suggest_ok(
+            self.suggest_body(
+                observed,
+                options={
+                    "general_threshold": 0,
+                    "character_threshold": 1,
+                    "suggestion_limit": 100,
+                },
+            )
+        )
+        await self.suggest_ok(
+            self.suggest_body(observed, "hf", options={"instructions": "x" * 4000})
+        )
+        self.assertEqual(len(self.tagger.calls), 2)
+        self.tagger.download_model.assert_not_called()
+
+    async def test_capture_token_must_match_target_and_epoch(self):
+        observed = await self.observed()
+        make_png(self.output_dir / "sub" / "b.png")
+        other = await self.observe_ok("sub/b.png")
+
+        status, body = await self.suggest(
+            self.suggest_body(observed, capture_token="x" * 43)
+        )
+        self.assertEqual((status, body["code"]), (409, "capture-expired"))
+        status, body = await self.suggest(
+            self.suggest_body(observed, capture_token=other["capture"]["token"])
+        )
+        self.assertEqual((status, body["code"]), (409, "capture-expired"))
+        status, body = await self.suggest(
+            self.suggest_body(observed, expected_epoch=str(uuid.uuid4()))
+        )
+        self.assertEqual((status, body["code"]), (409, "storage-changed"))
+        self.assertEqual(body["library_epoch"], self.epoch)
+        status, body = await self.suggest(
+            self.suggest_body(
+                observed, target={**observed["target"], "root_scope": "0" * 64}
+            )
+        )
+        self.assertEqual((status, body["code"]), (409, "source-changed"))
+        self.api._captures.release(observed["capture"]["token"])
+        status, body = await self.suggest(self.suggest_body(observed))
+        self.assertEqual((status, body["code"]), (409, "capture-expired"))
+        self.assertEqual(self.tagger.calls, [])
+        for token in (observed, other):
+            entry = self.api._captures._entries.get(token["capture"]["token"])
+            self.assertTrue(entry is None or entry.uses == 0)
+
+    async def test_video_and_oversized_dimensions_unsupported_media(self):
+        (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 4)
+        video = await self.observe_ok("clip.mp4")
+        self.assertFalse(video["capture"]["generate"]["supported"])
+        big = await self.observed("big.png", size=(4097, 1))
+        self.assertFalse(big["capture"]["generate"]["supported"])
+
+        for observed in (video, big):
+            status, body = await self.suggest(self.suggest_body(observed))
+            self.assertEqual((status, body["code"]), (415, "unsupported-media"))
+            self.assertEqual(body["error"], observed["capture"]["generate"]["reason"])
+        self.assertEqual(self.tagger.calls, [])
+        self.assertEqual(self.tagger.load_calls, [])
+        self.assertIsNone(model_gate.holder())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_api_annotations.py

@@ -986,6 +986,7 @@ class TestEnvelopes(AnnotationAPITestCase):
                 "model_loaded",
                 "models",
                 "success",
+                "suggestion_limit",  # B2 adds it (D68); nothing else changes
                 "wd14_character_threshold",
                 "wd14_general_threshold",
             ],

tests/test_api_autotag.py

@@ -37,6 +37,11 @@ from database.operations import PromptDatabase
 from py.api import PromptManagerAPI
 from utils.hashing import generate_prompt_hash
 
+try:  # guarded so a missing module fails assertions, not test loading
+    from py import model_gate
+except ImportError:
+    model_gate = None
+
 
 class FakeAutotagService:
     """Stand-in for AutoTagService that records calls and returns fixed tags."""
@@ -128,6 +133,11 @@ class AutotagAPITestCase(AioHTTPTestCase):
     async def tearDownAsync(self):
         await super().tearDownAsync()  # closes the aiohttp test client
         self._service_patch.stop()
+        if model_gate is not None:
+            leaked = model_gate.holder()
+            if leaked is not None:  # never leak the process-wide gate
+                model_gate.release(leaked)
+            self.assertIsNone(leaked, "a route left the model gate held")
         self.api.db.close_all()
         for path in (
             self._temp_db.name,
@@ -947,6 +957,159 @@ class TestUnloadAutotag(AutotagAPITestCase):
         self.assertEqual(resp.status, 500)
 
 
+class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase):
+    """While another model operation holds the gate, every legacy model route
+    answers ``409 model-busy`` and touches no model (the one allowed legacy
+    change)."""
+
+    BUSY = {"success": False, "code": "model-busy"}
+
+    def setUp(self):
+        super().setUp()
+        self.assertIsNotNone(model_gate, "py.model_gate is not implemented")
+
+    async def _busy(self, resp):
+        self.assertEqual(resp.status, 409)
+        data = await resp.json()
+        self.assertEqual({k: data[k] for k in self.BUSY}, self.BUSY)
+        self.assertIn("busy", data["error"].lower())
+        return data
+
+    async def test_legacy_single_download_start_return_model_busy_while_gate_held(
+        self,
+    ):
+        image = self._make_image(self.output_dir)
+        self.service._loaded_type = "wd14-vit"
+        self.assertTrue(model_gate.try_acquire("test-holder"))
+        try:
+            await self._busy(
+                await self._post_json(
+                    "/prompt_manager/autotag/single",
+                    {"path": str(image), "model_type": "gguf"},
+                )
+            )
+            await self._busy(
+                await self.client.request(
+                    "POST", "/prompt_manager/autotag/download/gguf"
+                )
+            )
+            await self._busy(
+                await self.client.request(
+                    "POST", "/prompt_manager/autotag/start?model_type=gguf"
+                )
+            )
+            await self._busy(
+                await self.client.request("POST", "/prompt_manager/autotag/unload")
+            )
+            self.assertEqual(model_gate.holder(), "test-holder")
+        finally:
+            model_gate.release("test-holder")
+
+        self.assertEqual(self.service.load_calls, [])
+        self.assertEqual(self.service.generate_calls, [])
+        self.assertEqual(self.service.download_calls, [])
+        self.assertEqual(self.service.unload_calls, 0)
+        self.assertEqual(self.service.get_loaded_model_type(), "wd14-vit")
+
+        # A refused start released the batch lock: the next start runs.
+        resp = await self.client.request(
+            "POST", "/prompt_manager/autotag/start?model_type=gguf"
+        )
+        self.assertEqual(resp.status, 200)
+        self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete")
+
+    async def test_legacy_routes_hold_the_gate_only_while_they_run(self):
+        image = self._make_image(self.output_dir)
+        seen = []
+
+        def generate(image_path, **kwargs):
+            seen.append(model_gate.holder())
+            return ["x"]
+
+        def download(model_type, progress_callback=None):
+            seen.append(model_gate.holder())
+            return True
+
+        def load(model_type, use_gpu=True):
+            seen.append(model_gate.holder())
+            self.service._loaded_type = model_type
+            return True
+
+        self.service.generate_tags = generate
+        self.service.download_model = download
+        self.service.load_model = load
+
+        resp = await self._post_json(
+            "/prompt_manager/autotag/single", {"path": str(image)}
+        )
+        self.assertEqual(resp.status, 200)
+        self.assertIsNone(model_gate.holder())
+        resp = await self.client.request(
+            "POST", "/prompt_manager/autotag/download/gguf"
+        )
+        self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete")
+        self.assertIsNone(model_gate.holder())
+        resp = await self.client.request("POST", "/prompt_manager/autotag/unload")
+        self.assertEqual(resp.status, 200)
+        self.assertIsNone(model_gate.holder())
+
+        self.assertEqual(len(seen), 3)
+        self.assertTrue(all(owner and owner.startswith("legacy-") for owner in seen))
+
+    async def test_validation_errors_win_over_the_gate(self):
+        self.assertTrue(model_gate.try_acquire("test-holder"))
+        try:
+            resp = await self._post_json("/prompt_manager/autotag/single", {})
+            self.assertEqual(resp.status, 400)
+            resp = await self.client.request(
+                "POST", "/prompt_manager/autotag/start?general_threshold=x"
+            )
+            self.assertEqual(resp.status, 400)
+        finally:
+            model_gate.release("test-holder")
+
+    async def test_batch_holds_the_gate_until_it_completes(self):
+        prompt = self._save_prompt("p")
+        self._link_image(prompt, self._make_image(self.output_dir, "a.png"))
+        started = threading.Event()
+        proceed = threading.Event()
+
+        def slow_generate(image_path, **kwargs):
+            started.set()
+            proceed.wait(30)
+            return ["slow"]
+
+        self.service.generate_tags = slow_generate
+        run = asyncio.ensure_future(
+            self.client.request("POST", "/prompt_manager/autotag/start")
+        )
+        try:
+            await asyncio.get_running_loop().run_in_executor(None, started.wait, 30)
+            self.assertTrue(model_gate.holder().startswith("legacy-batch:"))
+            await self._busy(
+                await self.client.request("POST", "/prompt_manager/autotag/unload")
+            )
+        finally:
+            proceed.set()
+            resp = await run
+            events = self._sse_events(await resp.text())
+        self.assertEqual(events[-1]["type"], "complete")
+        self.assertIsNone(model_gate.holder())
+
+    async def test_models_response_carries_the_suggestion_limit(self):
+        from py.config import AutotagConfig
+
+        resp = await self.client.get("/prompt_manager/autotag/models")
+        data = await resp.json()
+        self.assertEqual(data["suggestion_limit"], 10)
+        AutotagConfig.SUGGESTION_LIMIT = 25
+        try:
+            resp = await self.client.get("/prompt_manager/autotag/models")
+            self.assertEqual((await resp.json())["suggestion_limit"], 25)
+        finally:
+            AutotagConfig.SUGGESTION_LIMIT = 10
+
+
 class TestScanOutputDir(AutotagAPITestCase):
 
     async def test_lists_images_with_thumbnails(self):

tests/test_autotag_engine.py

@@ -301,6 +301,143 @@ class TestGgufNormalisation(EngineTestCase):
         )
 
 
+SCORED_TAGS = [
+    {"name": "desert", "category": "0"},
+    {"name": "sandstone", "category": "0"},
+    {"name": "Sandstone", "category": "0"},
+    {"name": "hatsune_miku", "category": "4"},
+    {"name": "rating:safe", "category": "9"},
+] + [{"name": f"tag_{i:03d}", "category": "0"} for i in range(120)]
+SCORED_PROBABILITIES = [0.96, 0.98, 0.5, 0.9, 0.99] + [
+    0.9 - i * 0.004 for i in range(120)
+]
+
+
+class TestTagSuggestions(EngineTestCase):
+    """``generate_tag_suggestions``: scored, ordered, bounded, distinct."""
+
+    def _wd14(self, tags=SCORED_TAGS, probabilities=SCORED_PROBABILITIES):
+        _touch_model_files(self.models_dir, "wd14-vit")
+        self.session = FakeOnnxSession(probabilities)
+        service = self._service(("wd14", (self.session, tags)))
+        service.load_model("wd14-vit")
+        return service
+
+    def _gguf(self, text):
+        _touch_model_files(self.models_dir, "gguf")
+        self.model = FakeGgufModel(text)
+        service = self._service(("gguf", self.model))
+        service.load_model("gguf")
+        return service
+
+    def test_wd14_suggestions_carry_real_probabilities_sorted(self):
+        service = self._wd14()
+
+        suggestions = service.generate_tag_suggestions(
+            str(self.image_path),
+            limit=10,
+            general_threshold=0.35,
+            character_threshold=0.85,
+        )
+
+        self.assertEqual(suggestions[:2], [("sandstone", 0.98), ("desert", 0.96)])
+        # The character tag passes its own threshold; rating tags never show.
+        self.assertIn(("hatsune_miku", 0.9), suggestions)
+        self.assertNotIn("rating:safe", [name for name, _ in suggestions])
+        scores = [score for _, score in suggestions]
+        self.assertEqual(scores, sorted(scores, reverse=True))
+        self.assertTrue(all(isinstance(score, float) for score in scores))
+        # "Sandstone" (0.5) repeats "sandstone" case-insensitively: dropped.
+        self.assertEqual(
+            [name.casefold() for name, _ in suggestions].count("sandstone"), 1
+        )
+
+    def test_suggestions_keep_the_requested_limit(self):
+        service = self._wd14()
+        for limit in (1, 10, 100):
+            with self.subTest(limit=limit):
+                suggestions = service.generate_tag_suggestions(
+                    str(self.image_path),
+                    limit=limit,
+                    general_threshold=0.0,
+                    character_threshold=0.0,
+                )
+                self.assertEqual(len(suggestions), limit)
+        for bad in (0, -1, 101, 2.5, "10", True, None):
+            with self.subTest(bad=bad), self.assertRaises(ValueError):
+                service.generate_tag_suggestions(str(self.image_path), limit=bad)
+
+    def test_thresholds_default_to_the_service_values(self):
+        service = self._wd14()
+        service.wd14_general_threshold = 0.95
+        service.wd14_character_threshold = 0.95
+
+        suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10)
+
+        self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)])
+
+    def test_llm_suggestions_have_null_confidence_in_output_order(self):
+        service = self._gguf("Zebra, apple, ZEBRA, mango tree, kiwi, plum")
+
+        suggestions = service.generate_tag_suggestions(
+            str(self.image_path), limit=3, prompt="tags please"
+        )
+
+        self.assertEqual(
+            suggestions, [("zebra", None), ("apple", None), ("mango_tree", None)]
+        )
+
+    def test_generate_tags_return_shape_unchanged(self):
+        service = self._wd14()
+        self.assertEqual(
+            service.generate_tags(
+                str(self.image_path), general_threshold=0.95, character_threshold=0.95
+            ),
+            ["sandstone", "desert"],
+        )
+        gguf = self._gguf("Smile, smile, long hair")
+        self.assertEqual(
+            gguf.generate_tags(str(self.image_path)), ["smile", "long_hair"]
+        )
+
+    def test_explicit_prompt_does_not_mutate_custom_prompt(self):
+        service = self._gguf("tag")
+        service.custom_prompt = "saved legacy prompt"
+
+        service.generate_tag_suggestions(
+            str(self.image_path), limit=10, prompt="dialog instructions"
+        )
+
+        self.assertEqual(
+            self.model.calls[0]["messages"][1]["content"][0]["text"],
+            "dialog instructions",
+        )
+        self.assertEqual(service.custom_prompt, "saved legacy prompt")
+
+    def test_suggestions_need_a_loaded_model(self):
+        service = self._service(("gguf", FakeGgufModel("x")))
+        with self.assertRaises(RuntimeError):
+            service.generate_tag_suggestions(str(self.image_path), limit=10)
+
+    def test_image_file_is_closed_after_inference(self):
+        service = self._gguf("tag")
+        opened = []
+        real_open = autotag.Image.open
+
+        def tracking_open(*args, **kwargs):
+            image = real_open(*args, **kwargs)
+            opened.append(image)
+            return image
+
+        with patch.object(autotag.Image, "open", tracking_open):
+            service.generate_tag_suggestions(str(self.image_path), limit=10)
+            service.generate_tags(str(self.image_path))
+
+        self.assertEqual(len(opened), 2)
+        for image in opened:
+            self.assertIsNone(getattr(image, "fp", None))
+
+
 class TestGenerateTagsErrors(EngineTestCase):
 
     def test_no_model_loaded(self):

tests/test_config.py

@@ -19,6 +19,7 @@ sys.modules.setdefault("server", _mock_server)
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
+from py import config as config_module
 from py.config import GalleryConfig, IntegrationConfig, PromptManagerConfig
 
 
@@ -404,5 +405,73 @@ class TestConfigPathOverride(unittest.TestCase):
             self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 77)
 
 
+class TestAutotagConfig(unittest.TestCase):
+    """``autotag.defaults.suggestion_limit``: integer 1..100, default 10 (D68)."""
+
+    def setUp(self):
+        self.config_class = getattr(config_module, "AutotagConfig", None)
+        self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented")
+        self._orig = PromptManagerConfig.get_config()
+
+    def tearDown(self):
+        self.config_class.SUGGESTION_LIMIT = 10
+        PromptManagerConfig.update_config(self._orig)
+
+    def _load(self, value):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "config.json")
+            with open(path, "w") as f:
+                # json writes NaN for float("nan"), which json.load reads back.
+                json.dump({"autotag": {"defaults": {"suggestion_limit": value}}}, f)
+            with self.assertLogs("prompt_manager.config", level="WARNING") as logs:
+                PromptManagerConfig.load_from_file(path)
+        return logs
+
+    def test_default_is_ten_and_persisted_under_autotag_defaults(self):
+        self.assertEqual(config_module.DEFAULT_SUGGESTION_LIMIT, 10)
+        self.assertEqual(config_module.MAX_SUGGESTION_LIMIT, 100)
+        self.assertEqual(self.config_class.suggestion_limit(), 10)
+        self.assertEqual(
+            PromptManagerConfig.get_config()["autotag"],
+            {"defaults": {"suggestion_limit": 10}},
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, "config.json")
+            PromptManagerConfig.update_config(
+                {"autotag": {"defaults": {"suggestion_limit": 25}}}
+            )
+            PromptManagerConfig.save_to_file(path)
+            self.config_class.SUGGESTION_LIMIT = 10
+            PromptManagerConfig.load_from_file(path)
+        self.assertEqual(self.config_class.suggestion_limit(), 25)
+
+    def test_valid_limits_are_kept(self):
+        for value in (1, 10, 100):
+            with self.subTest(value=value):
+                PromptManagerConfig.update_config(
+                    {"autotag": {"defaults": {"suggestion_limit": value}}}
+                )
+                self.assertEqual(self.config_class.suggestion_limit(), value)
+
+    def test_suggestion_limit_setting_validated_on_load(self):
+        for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]):
+            with self.subTest(value=value):
+                self.config_class.SUGGESTION_LIMIT = 42
+                logs = self._load(value)
+                self.assertEqual(self.config_class.suggestion_limit(), 10)
+                self.assertIn("suggestion_limit", "\n".join(logs.output))
+
+    def test_missing_or_malformed_sections_keep_the_current_value(self):
+        self.config_class.SUGGESTION_LIMIT = 7
+        for config in ({}, {"autotag": {}}, {"autotag": {"defaults": {}}}):
+            PromptManagerConfig.update_config(config)
+            self.assertEqual(self.config_class.suggestion_limit(), 7)
+        with self.assertLogs("prompt_manager.config", level="WARNING"):
+            PromptManagerConfig.update_config({"autotag": "nope"})
+        with self.assertLogs("prompt_manager.config", level="WARNING"):
+            PromptManagerConfig.update_config({"autotag": {"defaults": [1]}})
+        self.assertEqual(self.config_class.suggestion_limit(), 7)
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_media_capture.py

@@ -596,6 +596,42 @@ class CaptureStoreTestCase(unittest.TestCase):
             cap.sha256, hashlib.sha256(b"\xff\xd8not really a jpeg").hexdigest()
         )
 
+    def test_metadata_over_inspection_budgets_does_not_block_generate(self):
+        # B2 carry-over (A M1): inferable is decided from the image header
+        # alone, so prompt metadata over the inspection budgets (here a 3 MiB
+        # text chunk) leaves a decodable image inferable.
+        base = png_bytes(64, 32)
+        huge = _chunk(b"tEXt", b"prompt\x00" + b"x" * (3 * 1024 * 1024))
+        self.write("big-metadata.png", base[:-12] + huge + base[-12:])
+        store = self.make_store()
+
+        cap = self.capture(store, "big-metadata.png")
+
+        self.assertEqual((cap.width, cap.height), (64, 32))
+        self.assertTrue(cap.inferable)
+        self.assertIsNone(cap.inferable_reason)
+        assert_closed(self, self.parent)
+
+    def test_decompression_bomb_header_is_not_inferable(self):
+        self.write("bomb.png", png_bytes(20_000, 20_000))
+        store = self.make_store()
+
+        cap = self.capture(store, "bomb.png")
+
+        self.assertFalse(cap.inferable)
+        self.assertIn("4096", cap.inferable_reason)
+        assert_closed(self, self.parent)
+
+    def test_header_dimensions_reads_only_the_header(self):
+        path = self.write("dims.png", png_bytes(7, 5))
+        self.assertEqual(capture_module.header_dimensions(str(path)), (7, 5))
+        assert_closed(self, self.parent)
+        with self.assertRaises(capture_module.HeaderUnreadable):
+            capture_module.header_dimensions(
+                str(self.write("bad.png", b"\x89PNG\r\n\x1a\nnope"))
+            )
+        assert_closed(self, self.parent)
+
     def test_video_capture_is_not_inferable(self):
         data = b"\x00\x00\x00\x18ftypmp42" + os.urandom(32)
         self.write("clip.mp4", data)
@@ -604,7 +640,7 @@ class CaptureStoreTestCase(unittest.TestCase):
         def never(path):  # pragma: no cover - guard only
             raise AssertionError("video captures are never probed")
 
-        with patch.object(capture_module, "read_inspection_file", never):
+        with patch.object(capture_module, "header_dimensions", never):
             cap = self.capture(store, "clip.mp4")
         self.assertEqual((cap.media_type, cap.width, cap.height), ("video", None, None))
         self.assertFalse(cap.inferable)

tests/test_model_gate.py

@@ -0,0 +1,176 @@
+"""The shared, non-blocking model gate (#249 B2).
+
+One owner at a time may load, swap, run, download or unload a tagging model.
+A release requested while gated work still runs in a worker thread takes
+effect only when that work really exits, so a client disconnect never lets a
+second owner swap the model under a running inference.
+"""
+
+import asyncio
+import os
+import sys
+import threading
+import unittest
+import unittest.mock
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+try:  # guarded so a missing module fails assertions, not test loading
+    from py import model_gate
+except ImportError:
+    model_gate = None
+
+WAIT = 30.0  # ceiling for cross-thread handoffs, never a short clock
+
+
+class GateTestCase(unittest.TestCase):
+    def setUp(self):
+        self.assertIsNotNone(model_gate, "py.model_gate is not implemented")
+        self.gate = model_gate.ModelGate()
+
+
+class TestGate(GateTestCase):
+    def test_gate_is_nonblocking_and_single_owner(self):
+        self.assertIsNone(self.gate.holder())
+        self.assertTrue(self.gate.try_acquire("legacy-batch"))
+        self.assertEqual(self.gate.holder(), "legacy-batch")
+
+        # A second owner is refused at once, never queued.
+        self.assertFalse(self.gate.try_acquire("image-annotations"))
+        self.assertFalse(self.gate.try_acquire("legacy-batch"))
+        self.assertEqual(self.gate.holder(), "legacy-batch")
+
+        # Only the holder can release.
+        self.gate.release("image-annotations")
+        self.assertEqual(self.gate.holder(), "legacy-batch")
+        self.gate.release("legacy-batch")
+        self.assertIsNone(self.gate.holder())
+        self.gate.release("legacy-batch")  # idempotent
+
+        self.assertTrue(self.gate.try_acquire("image-annotations"))
+        self.assertEqual(self.gate.holder(), "image-annotations")
+
+    def test_release_waits_for_running_work(self):
+        self.assertTrue(self.gate.try_acquire("owner"))
+        self.gate.begin("owner")
+        self.gate.release("owner")
+        self.assertEqual(self.gate.holder(), "owner")
+        self.assertFalse(self.gate.try_acquire("other"))
+        # No new work is admitted once the owner let go.
+        with self.assertRaises(model_gate.GateNotHeld):
+            self.gate.begin("owner")
+        self.gate.end("owner")
+        self.assertIsNone(self.gate.holder())
+
+    def test_begin_requires_the_holder(self):
+        with self.assertRaises(model_gate.GateNotHeld):
+            self.gate.begin("nobody")
+        self.assertTrue(self.gate.try_acquire("owner"))
+        with self.assertRaises(model_gate.GateNotHeld):
+            self.gate.begin("other")
+        self.gate.end("other")  # a stray end never disturbs the holder
+        self.assertEqual(self.gate.holder(), "owner")
+
+    def test_module_gate_is_one_process_wide_instance(self):
+        self.assertIsNone(model_gate.holder())
+        self.assertTrue(model_gate.try_acquire("a"))
+        try:
+            self.assertFalse(model_gate.try_acquire("b"))
+            self.assertEqual(model_gate.holder(), "a")
+        finally:
+            model_gate.release("a")
+        self.assertIsNone(model_gate.holder())
+
+    def test_new_owner_names_are_unique(self):
+        first = model_gate.new_owner("legacy-single")
+        second = model_gate.new_owner("legacy-single")
+        self.assertNotEqual(first, second)
+        self.assertTrue(first.startswith("legacy-single:"))
+
+
+class TestGatedExecutor(GateTestCase):
+    def test_cancelled_waiter_keeps_the_gate_until_the_thread_exits(self):
+        started = threading.Event()
+        proceed = threading.Event()
+        finished = threading.Event()
+
+        def blocking_work():
+            started.set()
+            proceed.wait(WAIT)
+            finished.set()
+            return "done"
+
+        async def scenario():
+            self.assertTrue(self.gate.try_acquire("native"))
+            task = asyncio.ensure_future(self.gate.run("native", blocking_work))
+            try:
+                await asyncio.get_running_loop().run_in_executor(
+                    None, started.wait, WAIT
+                )
+                task.cancel()
+                with self.assertRaises(asyncio.CancelledError):
+                    await task
+            finally:
+                self.gate.release("native")
+            # The waiter is gone but the inference still runs: still held.
+            self.assertEqual(self.gate.holder(), "native")
+            self.assertFalse(self.gate.try_acquire("legacy-unload"))
+            proceed.set()
+            await asyncio.get_running_loop().run_in_executor(None, finished.wait, WAIT)
+            for _ in range(1000):
+                if self.gate.holder() is None:
+                    break
+                await asyncio.sleep(0.01)
+            self.assertIsNone(self.gate.holder())
+
+        asyncio.run(scenario())
+
+    def test_run_returns_the_result_and_propagates_errors(self):
+        async def scenario():
+            self.assertTrue(self.gate.try_acquire("owner"))
+            try:
+                self.assertEqual(await self.gate.run("owner", lambda x: x * 2, 21), 42)
+
+                def boom():
+                    raise ValueError("bad")
+
+                with self.assertRaises(ValueError):
+                    await self.gate.run("owner", boom)
+                self.assertEqual(self.gate.holder(), "owner")
+            finally:
+                self.gate.release("owner")
+            self.assertIsNone(self.gate.holder())
+            with self.assertRaises(model_gate.GateNotHeld):
+                await self.gate.run("owner", lambda: None)
+
+        asyncio.run(scenario())
+
+    def test_work_that_never_starts_gives_its_admission_back(self):
+        async def scenario():
+            loop = asyncio.get_running_loop()
+            self.assertTrue(self.gate.try_acquire("owner"))
+
+            def refused(executor, func):
+                raise RuntimeError("executor shut down")
+
+            with unittest.mock.patch.object(loop, "run_in_executor", refused):
+                with self.assertRaises(RuntimeError):
+                    await self.gate.run("owner", lambda: None)
+
+            def cancelled(executor, func):
+                future = loop.create_future()
+                future.cancel()  # a loop shutdown drops the queued call
+                return future
+
+            with unittest.mock.patch.object(loop, "run_in_executor", cancelled):
+                with self.assertRaises(asyncio.CancelledError):
+                    await self.gate.run("owner", lambda: None)
+            await asyncio.sleep(0)  # let the done callback run
+            self.gate.release("owner")
+            self.assertIsNone(self.gate.holder())
+
+        asyncio.run(scenario())
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/media_capture.py

@@ -28,16 +28,13 @@ import threading
 import time
 import warnings
 from dataclasses import dataclass
-from pathlib import Path
 from typing import Callable
 
 from PIL import Image
 
 try:
-    from .prompt_inspector import MetadataReadError, read_inspection_file
     from .video_metadata import VIDEO_EXTENSIONS, is_video_path
 except ImportError:  # imported with the repo root on sys.path
-    from utils.prompt_inspector import MetadataReadError, read_inspection_file
     from utils.video_metadata import VIDEO_EXTENSIONS, is_video_path
 
 logger = logging.getLogger(__name__)
@@ -578,20 +575,17 @@ class CaptureStore:
     def _probe(self, path: str, media_type: str):
         """Header-only dimensions of the capture; never decodes pixels.
 
-        Uses the inspection reader, so a PNG whose prompt metadata is over the
-        inspection budgets is reported as unreadable (not inferable) although
-        its pixels decode; B2 may switch to a dimensions-only header read.
+        Only the image header decides whether Generate is offered: prompt
+        metadata over the inspection budgets does not make a decodable image
+        unreadable (the metadata itself is read elsewhere, under its budgets).
         """
         if media_type == "video":
             return None, None, False, VIDEO_ONLY_REASON
         try:
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore", Image.DecompressionBombWarning)
-                read = read_inspection_file(Path(path))
-        except (MetadataReadError, Image.DecompressionBombError):
-            return None, None, False, HEADER_REASON
-        width, height = read.width, read.height
-        if width is None or height is None:
+            width, height = header_dimensions(path)
+        except HeaderTooLarge:
+            return None, None, False, DIMENSION_REASON
+        except HeaderUnreadable:
             return None, None, False, HEADER_REASON
         if (
             width > MAX_INFERENCE_SIDE
@@ -602,6 +596,34 @@ class CaptureStore:
         return width, height, True, None
 
 
+class HeaderUnreadable(Exception):
+    """The file has no image header Pillow can read."""
+
+
+class HeaderTooLarge(HeaderUnreadable):
+    """The header declares a decompression bomb (no pixel is decoded)."""
+
+
+def header_dimensions(path: str) -> tuple[int, int]:
+    """``(width, height)`` from the image header alone; the file is closed.
+
+    ``Image.open`` reads only the header; a decompression-bomb warning is an
+    error here, so an absurd size is refused before anything is decoded.
+    """
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", Image.DecompressionBombWarning)
+            with Image.open(path) as image:
+                width, height = image.size
+    except (Image.DecompressionBombError, Image.DecompressionBombWarning) as exc:
+        raise HeaderTooLarge(str(exc)) from None
+    except Exception as exc:  # any decoder failure means "no usable header"
+        raise HeaderUnreadable(str(exc)) from None
+    if not width or not height:
+        raise HeaderUnreadable("The header declares no size")
+    return int(width), int(height)
+
+
 def _relative_parts(relative) -> tuple[list[str], str]:
     """Split a safe relative POSIX media path; anything else is denied."""
     if not isinstance(relative, str) or not relative: