fix(mm): close TOCTOU race in _restore_incomplete_installs

invokeai/app/services/model_install/model_install_default.py

@@ -203,11 +203,6 @@ def _find_reusable_tmpdir(self, source: ModelSource) -> Optional[Path]:
     def _restore_incomplete_installs(self) -> None:
         path = self._app_config.models_path
         seen_sources: set[str] = set()
-        # Collect sources already tracked by active jobs (including those being downloaded right now).
-        # We must not re-queue these or delete their tmpdirs.
-        with self._lock:
-            active_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state}
-            active_sources.update(str(j.source) for j in self._download_cache.values() if not j.in_terminal_state)
         for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"):
             marker = self._read_install_marker(tmpdir)
             if not marker:
@@ -224,10 +219,6 @@ def _restore_incomplete_installs(self) -> None:
                 access_token = marker.get("access_token")
                 if isinstance(source, (HFModelSource, URLModelSource)) and isinstance(access_token, str):
                     source.access_token = access_token
-                if source_str in active_sources:
-                    # This tmpdir belongs to an install already in progress; leave it alone.
-                    self._logger.debug(f"Skipping restore for {source_str} - already being tracked")
-                    continue
                 if source_str in seen_sources:
                     self._logger.info(f"Removing duplicate temporary directory {tmpdir}")
                     self._safe_rmtree(tmpdir, self._logger)
@@ -249,7 +240,20 @@ def _restore_incomplete_installs(self) -> None:
             if files_meta:
                 job._resume_metadata = {f.get("url"): f for f in files_meta if f.get("url")}
             job.status = InstallStatus(status) if status else InstallStatus.WAITING
-            self._install_jobs.append(job)
+
+            # Atomically check that no other thread (e.g. import_model) has already
+            # queued this source, then append. Without this, a TOCTOU race against
+            # foreground import_model calls can enqueue the same source twice and
+            # cause a FileNotFoundError when the second download tries to rename
+            # the .downloading file the first one already moved.
+            with self._lock:
+                already_active = any(
+                    str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state
+                ) or any(str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state)
+                if already_active:
+                    self._logger.debug(f"Skipping restore for {source_str} - already being tracked")
+                    continue
+                self._install_jobs.append(job)
 
             if job.paused:
                 continue

tests/app/services/model_install/test_model_install.py

@@ -3,6 +3,7 @@
 """
 
 import gc
+import json
 import platform
 import shutil
 import threading
@@ -36,6 +37,11 @@
     ModelInstallJob,
     URLModelSource,
 )
+from invokeai.app.services.model_install.model_install_default import (
+    INSTALL_MARKER_FILENAME,
+    INSTALL_MARKER_VERSION,
+    TMPDIR_PREFIX,
+)
 from invokeai.app.services.model_records import ModelRecordChanges, UnknownModelException
 from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig
 from invokeai.backend.model_manager.taxonomy import (
@@ -372,6 +378,81 @@ def _blocked_restore() -> None:
         installer.stop()
 
 
+@pytest.mark.timeout(timeout=20, method="thread")
+def test_restore_skips_source_queued_during_restore(
+    mm2_app_config: InvokeAIAppConfig,
+    mm2_record_store,
+    mm2_download_queue,
+    mm2_session,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Regression test for #9141: if a foreground thread queues a job for a source after
+    restore has parsed that source's install marker but before restore appends its own job,
+    restore must notice the active job and skip the source instead of enqueuing it twice."""
+    installer = ModelInstallService(
+        app_config=mm2_app_config,
+        record_store=mm2_record_store,
+        download_queue=mm2_download_queue,
+        event_bus=TestEventService(),
+        session=mm2_session,
+    )
+    source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors"))
+
+    tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}race"
+    tmpdir.mkdir(parents=True)
+    marker = {
+        "version": INSTALL_MARKER_VERSION,
+        "source": str(source),
+        "access_token": None,
+        "config_in": {},
+        "status": InstallStatus.DOWNLOADING.value,
+        "updated_at": "",
+        "files": [],
+    }
+    with open(tmpdir / INSTALL_MARKER_FILENAME, "wt", encoding="utf-8") as f:
+        json.dump(marker, f)
+
+    marker_observed = threading.Event()
+    release_restore = threading.Event()
+    real_guess_source = installer._guess_source
+
+    def _pausing_guess_source(source_str: str):
+        result = real_guess_source(source_str)
+        marker_observed.set()
+        assert release_restore.wait(timeout=10)
+        return result
+
+    resumed: list[ModelInstallJob] = []
+    monkeypatch.setattr(installer, "_guess_source", _pausing_guess_source)
+    monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job))
+
+    try:
+        installer.start()
+        assert marker_observed.wait(timeout=10)
+
+        # Restore is now paused between parsing the marker and its locked duplicate
+        # check. Queue an active job for the same source, as a concurrent
+        # import_model call would.
+        foreground_job = ModelInstallJob(
+            id=installer._next_id(),
+            source=source,
+            config_in=ModelRecordChanges(),
+            local_path=tmpdir,
+        )
+        with installer._lock:
+            installer._install_jobs.append(foreground_job)
+
+        release_restore.set()
+        installer._wait_for_restore_complete()
+
+        jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)]
+        assert jobs_for_source == [foreground_job]
+        assert resumed == []
+    finally:
+        release_restore.set()
+        installer.stop()
+
+
 def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None:
     source = URLModelSource(
         url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors")