fix(mm): close TOCTOU race in _restore_incomplete_installs

[lstein]

Summary

This PR fixes an intermittent unit test failure in test_model_install.py and, following review, two further races in the same code.

• _restore_incomplete_installs snapshotted active sources under the lock, then iterated tmpdirs against the stale snapshot. A foreground import_model call landing during iteration could be missed, the same source enqueued twice, and the second download would fail with FileNotFoundError when trying to rename the .downloading file the first download had already moved.
• Fix: move the membership check inside the lock immediately before _install_jobs.append(job) so check-and-append is atomic.
• Review round 2 (JPPhoto):
  • The duplicate-tmpdir dedup ran before the active-source check, so restore could delete an active job's tmpdir when two markers existed for one source. The dedup now runs inside the locked section after the active check, and no tmpdir is deleted while its source has an active job.
  • The lock was one-sided: import_model checked/registered without holding _lock. _lock is now an RLock and import_model holds it across its duplicate check, job creation, and registration, making it atomic with restore.

Fixes #9141.

Test plan

• pytest tests/app/services/model_install/test_model_install.py::test_simple_download — ran 10× consecutively, all pass (was flaky ~30–50% on main)
• Regression test test_restore_skips_source_queued_during_restore: pauses restore in the race window, queues a foreground job for the same source, asserts no duplicate is enqueued — fails on main, passes with this fix
• Regression test test_restore_preserves_active_jobs_tmpdir: two markers for one source, active job owns the later-visited dir — asserts the active dir survives restore
• Regression test test_restore_preserves_tmpdir_of_job_in_download_cache: same, with the active job tracked only in _download_cache with a different tmpdir
• Regression test test_concurrent_import_and_restore_register_single_job: a real import_model() racing restore registers exactly one job and completes end-to-end; 15/15 consecutive runs
• All three review-round tests verified to fail against the pre-fix code
• pytest tests/app/services/model_install/ — 37/37 pass

🤖 Generated with Claude Code

[JPPhoto]

This looks good when inspecting visually, but I'd also like to see a test added for this in tests/app/services/model_install/test_model_install.py, especially since it's a race condition. It may be a little challenging to get the test right, though:

1. Create a tmp install dir with an active marker for a remote source.
2. Monkeypatch _resume_remote_download() so restore pauses at a controlled point.
3. Insert an active job for the same source into _install_jobs while restore is between marker parsing and the new locked duplicate check.
4. Let restore continue.
5. Assert restore did not append a duplicate job and did not enqueue/resume a second download.

That maps directly to the PR’s fix around invokeai/app/services/model_install/model_install_default.py:249.

The only slightly fiddly part is choosing the pause point. Because the fixed code does the duplicate check under _lock immediately before append, the test probably needs to monkeypatch either _guess_source() or _read_install_marker() to signal “restore has observed the marker”, then add the foreground job before restore reaches the locked block. That is manageable with two threading.Events, similar to the existing test_import_waits_for_startup_restore.

The test should not need real network downloads if _resume_remote_download() is monkeypatched to fail the test if called for the skipped duplicate.

[lstein]

Thanks for the detailed test sketch @JPPhoto — added in b953bba, following it closely:

1. The test writes a tmpinstall_* dir with a version-1 marker (status downloading) for a URL source before starting the installer.
2. _guess_source is monkeypatched to signal "restore has observed the marker" and then block on a threading.Event, which pauses restore exactly between marker parsing and the locked duplicate check (two events, same pattern as test_import_waits_for_startup_restore).
3. While restore is paused, the test appends an active job for the same source to _install_jobs under _lock, as a concurrent import_model would.
4. Restore is released and the test waits for it to complete.
5. It asserts the foreground job is the only job for that source and that _resume_remote_download (monkeypatched to record calls) was never invoked — no network needed.

To confirm it actually covers the race, I ran the test against the pre-fix model_install_default.py from main: it fails there (duplicate job appended and resume attempted) and passes with the fix. Full tests/app/services/model_install/ suite passes (34/34).

[JPPhoto]

• invokeai/app/services/model_install/model_install_default.py:222: _restore_incomplete_installs() now marks a source as seen before it checks whether that source is already active. If there are two tmpinstall_* dirs with markers for the same source and an active job is already using the second dir, restore can see the stale first dir, add the source to seen_sources, skip it as active at lines 249-255, then treat the active second dir as a duplicate and delete it at lines 222-224. That can remove the active job's partial download directory while the job is still tracked in _install_jobs or _download_cache, causing later resume/download/install work to fail or restart incorrectly. To expose this issue, add a test that creates two install marker dirs for the same source, registers an active job whose _install_tmpdir is the second dir, forces restore to visit the first dir before the active dir, then asserts the active dir still exists and no duplicate restore job was queued.

• invokeai/app/services/model_install/model_install_default.py:249: The new lock is one-sided and does not make registration atomic with import_model(). The production import path checks _install_jobs at line 477, updates _download_cache at line 1264, and appends at line 495 without acquiring _lock. An import that passes _wait_for_restore_complete() immediately before start() clears the initially-set event can therefore race restoration: both paths can observe no active job and enqueue the same source. The new test masks this because it manually acquires _lock before appending at invokeai/tests/app/services/model_install/test_model_install.py:442, unlike the production path. To expose this issue, add a test that pauses a real import_model() call after its duplicate check, starts restoration for the same source, lets restore perform its locked check, and then allows both paths to register their jobs; assert that only one job and one download are created.

• invokeai/tests/app/services/model_install/test_model_install.py:381: The new regression test covers the single-marker duplicate-source case, but it does not cover the tmpdir ownership/deletion case above because it appends a foreground job using the same tmpdir as the marker at lines 436-443. It also bypasses the real remote import shape where _enqueue_remote_download() stores the job in _download_cache before import_model() appends it to _install_jobs (invokeai/app/services/model_install/model_install_default.py:1264 and :495). To expose this issue, add coverage where the active job is represented by _download_cache and has a different _install_tmpdir from an older duplicate marker for the same source.

[lstein]

Thanks @JPPhoto — all three findings were real, and they're addressed in ebaec12.

1. Restore deleting an active job's tmpdir. Confirmed exactly as you described: the seen_sources dedup ran before the active-source check, so a stale first dir marked the source seen and the active second dir was then rmtree'd as a "duplicate". The dedup now lives inside the locked section, after the active check, and when a source has an active job restore never deletes any of its tmpdirs (stale duplicates are collected on a later startup when the source is idle). The actual _safe_rmtree call happens after the lock is released. Covered by test_restore_preserves_active_jobs_tmpdir, which builds two marker dirs for one source, registers an active job owning the second, and forces restore to visit the stale dir first (sorted Path.glob).

2. One-sided lock. Also confirmed. _lock is now a threading.RLock (the import helpers call _next_id(), which acquires it), and import_model() holds it across its duplicate check, job creation, and registration — so import's check-and-register and restore's check-and-append are mutually atomic. One tradeoff worth flagging: the HF metadata fetch in _import_from_hf now runs under the lock, which briefly stalls download progress callbacks during an import's metadata lookup. Given imports are infrequent and human-triggered I think that's acceptable, but happy to hoist the fetch out with a re-check if you'd prefer. Covered by test_concurrent_import_and_restore_register_single_job, which drives a real import_model() (no manual locking this time): the import pauses inside _import_from_url after its duplicate check, start() runs in a second thread, and the test asserts exactly one job is registered, _resume_remote_download is never called, and the import then completes end-to-end. The test buffers _put_in_queue until after the assertions so the import job can't reach a terminal state (and delete its marker) before restore's active check — without that the test itself was flaky.

3. _download_cache coverage. Added test_restore_preserves_tmpdir_of_job_in_download_cache: the active job is tracked only in _download_cache with a different _install_tmpdir than an older duplicate marker for the same source.

Verification: all three new tests fail against the pre-fix code (checked by reverting model_install_default.py and re-running) and pass with the fixes; the concurrent test passed 15/15 consecutive runs; full tests/app/services/model_install/ suite passes 37/37; ruff clean.

[JPPhoto]

Two issues:

• invokeai/app/services/model_install/model_install_default.py:493: import_model() now holds the installer lock while remote import reaches _multifile_download(), which acquires the download queue lock through _next_id(). Download callbacks take these locks in the opposite order: invokeai/app/services/download/download_default.py:691 holds the download queue lock while invoking _download_progress_callback(), which acquires the installer lock at invokeai/app/services/model_install/model_install_default.py:1400. If a second remote import begins while an earlier download callback is running, the callback can wait for the installer lock while the import waits for the download queue lock, permanently deadlocking both threads. This can freeze active downloads and leave the new import request blocked indefinitely. To expose this issue, add a test that pauses an existing multifile callback while it holds the download queue lock, starts another remote import_model() call until it holds the installer lock and requests a download job ID, then releases the callback and asserts both threads finish within a timeout.

• invokeai/app/services/model_install/model_install_default.py:253: The new per-marker active check excludes terminal jobs, but job terminal transitions, marker deletion, and temporary-directory cleanup are not synchronized with this lock. Restore can read an active job's marker, the installer thread can then mark that job completed at line 1487, and restore can subsequently see no nonterminal owner and append a second job for the same marker at line 264. The original up-front active_sources snapshot removed by this PR would continue treating a job that was active when restoration began as owned. The original job then deletes the marker and temporary directory while the duplicate restore job resumes from that directory, causing a duplicate installation or filesystem failure. To expose this issue, add a test that starts restore with an active job owning the marker directory, pauses restore after reading the marker, transitions the owner to COMPLETED while leaving its marker and directory present, resumes restore, and asserts that no second job or download is queued for that source.