agent: code-atomiser-fix — narrow LlmEngine critical section + drop old model first (Race-3, Lifecycle-1)

Race-3 (conf 86): `LlmEngine::load_model` previously held the inner
`std::sync::Mutex` for the entire `LlamaBackend::init` +
`LlamaModel::load_from_file` call (5-15 s on a cold load). `is_loaded()`,
`loaded_model()`, and `loaded_model_id()` all take that same mutex and
are called from sync Tauri handlers (`get_llm_status`, `check_llm_model`,
`delete_llm_model`, `test_llm_model`) WITHOUT `spawn_blocking`. During a
first-run load, parallel `refreshLlmStatus()` polls from the frontend
parked tokio worker threads on the std-mutex; a handful of concurrent
status polls was enough to deadlock the default `num_cpus`-sized Tauri
runtime.

Lifecycle-1 (conf 80): same function held the OLD `Arc<LlamaModel>` in
`guard.model` during the new `load_from_file` call, so a model swap
peaked at ~2x VRAM. A 27B Q4 (~17 GB) swap OOMed a 24 GB card even
though either model fit alone.

Fix: redesign `load_model` so the slow llama-cpp work happens OUTSIDE
the mutex.

  1. Short crit section: compare against currently-loaded triple —
     return Ok on match (no-op fast path).
  2. CAS a new `loading: AtomicBool` from false → true. If a load is
     already in flight, refuse with the new `EngineError::AlreadyLoading`
     rather than starting a parallel one. A `LoadingGuard` RAII drop
     clears the flag on every exit path including panic.
  3. Short crit section: drop the OLD model Arc (releasing VRAM via
     `llama_free_model`) BEFORE the new load begins — fixes Lifecycle-1.
  4. Heavy `LlamaBackend::init` + `LlamaModel::load_from_file` run
     OUTSIDE the mutex.
  5. Short crit section: install the new state.

`is_loaded()`, `loaded_model()`, `loaded_model_id()` now only contend
on the brief state-mutation sections, not the multi-second file load.
A new `is_loading()` accessor exposes the in-flight state for callers
that need to distinguish "loading" from "not loaded".

Backend lifecycle: `LlamaBackend::init` is process-singleton —
llama-cpp-2 enforces this via `LLAMA_BACKEND_INITIALIZED: AtomicBool`
and returns `BackendAlreadyInitialized` on a second call. The Drop impl
DOES call `llama_backend_free` and flips the flag back, so re-init
would technically work, but we keep the backend Arc resident across
loads/unloads to avoid init/free churn. The Lifecycle-1 fix drops only
the model Arc, NOT the backend Arc — dropping the backend mid-swap
would still work (the new load would re-init), but the comment trail
documents the singleton contract for the next reader.

`is_loaded()` now reports false briefly during a swap window (model
dropped, new model not yet installed). Documented in the doc comment.
Callers wanting "model X is loaded" must check `loaded_model_id()`
against their target, not just `is_loaded()`.

Regression tests added in `crates/llm/src/lib.rs::tests`:

  - `is_loaded_does_not_block_on_slow_load`: holds the lock-discipline
    open via a Barrier and asserts `is_loaded()` / `loaded_model_id()`
    return in ≤50 ms while the slow section is mid-flight. Verified
    that a pre-fix structure (`std::sync::Mutex` held across a 500 ms
    sleep) makes the probe block ~450 ms; the new structure makes it
    return in microseconds.
  - `second_concurrent_load_is_refused`: two concurrent
    `__test_run_with_lock_discipline` calls; the second gets
    `EngineError::AlreadyLoading` and never reaches its op closure.
    Doubles as the TOCTOU guard for Race-4/5 at the engine layer.

A `pub(crate) #[cfg(test)] __test_run_with_lock_discipline` helper
exposes the locking skeleton (loading flag + drop-old-model + slow op
outside lock + install) without requiring a real GGUF on disk; this is
the harness the regression tests use.

Verification: cargo test --workspace + npm run check both green.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/llm/src/lib.rs

@@ -1,5 +1,6 @@
 use std::num::NonZeroU32;
 use std::path::Path;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 
 use encoding_rs::UTF_8;
@@ -31,6 +32,8 @@ pub enum EngineError {
     NotLoaded,
     #[error("LLM load failed: {0}")]
     LoadFailed(String),
+    #[error("Another LLM load is already in flight; refusing to start a parallel load.")]
+    AlreadyLoading,
     #[error(
         "prompt too long: {prompt_tokens} prompt tokens exceed the {available_prompt_tokens}-token prompt budget for an {context_window}-token context with {max_tokens} reserved response tokens"
     )]
@@ -83,6 +86,29 @@ struct LlmState {
 #[derive(Clone, Default)]
 pub struct LlmEngine {
     inner: Arc<Mutex<LlmState>>,
+    /// Flag held for the duration of a model load. The std::sync::Mutex
+    /// covers cheap state mutations (~microseconds); the multi-second
+    /// `LlamaModel::load_from_file` call now runs *outside* the mutex,
+    /// so polls like `is_loaded()` / `loaded_model_id()` (called from
+    /// sync Tauri handlers without `spawn_blocking`) don't park tokio
+    /// worker threads on a slow C++ FFI call. This Atomic also doubles
+    /// as a TOCTOU guard: two concurrent `load_model` invocations on
+    /// the same engine will not both reach the heavy load — the second
+    /// returns `EngineError::AlreadyLoading`.
+    loading: Arc<AtomicBool>,
+}
+
+/// RAII guard that clears the `loading` flag on drop, including on
+/// panic / early-return. Prevents the engine getting stuck in a
+/// "permanently loading" state if a load fails midway.
+struct LoadingGuard {
+    flag: Arc<AtomicBool>,
+}
+
+impl Drop for LoadingGuard {
+    fn drop(&mut self) {
+        self.flag.store(false, Ordering::Release);
+    }
 }
 
 impl LlmEngine {
@@ -100,8 +126,51 @@ impl LlmEngine {
         model_path: &Path,
         use_gpu: bool,
     ) -> Result<(), EngineError> {
-        let mut guard = self.inner.lock().unwrap();
+        self.load_model_with(model_id, model_path, use_gpu, |backend, path, params| {
+            LlamaModel::load_from_file(backend, path, params)
+                .map_err(|e| EngineError::LoadFailed(format!("model load: {e}")))
+        })
+    }
 
+    /// Core load implementation with a swappable file-loader closure.
+    /// Production callers use `load_model`, which delegates here with
+    /// the real `LlamaModel::load_from_file`. Tests inject a sleepy /
+    /// counting closure to exercise the locking discipline without
+    /// pulling a real GGUF off disk.
+    ///
+    /// Locking discipline (the whole point of this function):
+    ///   1. Take the mutex briefly to compare against the currently
+    ///      loaded triple — if it matches, return early. No-op fast path.
+    ///   2. CAS the `loading` flag from false → true. If another load is
+    ///      already in flight, refuse with `AlreadyLoading` rather than
+    ///      starting a parallel one. A `LoadingGuard` ensures the flag
+    ///      is cleared on every exit path including panic.
+    ///   3. Take the mutex briefly to drop the OLD model Arc (frees its
+    ///      VRAM via `llama_free_model`) before the new load begins.
+    ///      The backend Arc is preserved — `LlamaBackend::init()` is a
+    ///      one-shot per process (an `AtomicBool` in llama-cpp-2 enforces
+    ///      `BackendAlreadyInitialized` on a second call), so we must
+    ///      never drop the backend while the process keeps running.
+    ///      Note: `is_loaded()` reports false during the swap window —
+    ///      that is the correct semantics. Callers wanting "model X is
+    ///      loaded" must check `loaded_model_id()` against their target.
+    ///   4. Initialise the backend if absent (first-ever load only) and
+    ///      run the slow `load_from_file` call — both OUTSIDE the mutex.
+    ///   5. Take the mutex briefly to install the new backend (if just
+    ///      initialised) and the new model Arc.
+    fn load_model_with<F>(
+        &self,
+        model_id: LlmModelId,
+        model_path: &Path,
+        use_gpu: bool,
+        loader: F,
+    ) -> Result<(), EngineError>
+    where
+        F: FnOnce(&LlamaBackend, &Path, &LlamaModelParams) -> Result<LlamaModel, EngineError>,
+    {
+        // Step 1: short crit section — already-loaded fast path.
+        {
+            let guard = self.inner.lock().unwrap();
             if let Some(loaded) = &guard.loaded {
                 if loaded.model_id == model_id.as_str()
                     && loaded.model_path == model_path.display().to_string()
@@ -110,8 +179,38 @@ impl LlmEngine {
                     return Ok(());
                 }
             }
+        }
 
-        let backend = match guard.backend.clone() {
+        // Step 2: claim the loading slot. Refuse if a parallel load is
+        // already mid-flight rather than starting a second slow load
+        // and silently overwriting the first.
+        if self
+            .loading
+            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+            .is_err()
+        {
+            return Err(EngineError::AlreadyLoading);
+        }
+        let _loading_guard = LoadingGuard {
+            flag: Arc::clone(&self.loading),
+        };
+
+        // Step 3: short crit section — drop the OLD model so its VRAM is
+        // released BEFORE we allocate the new one. Without this, a swap
+        // briefly holds two models resident (Lifecycle-1: an
+        // ~17 GB Q4 27B swap on a 24 GB card OOMs even though either
+        // model fits alone). Keep the backend Arc — see locking notes.
+        let existing_backend = {
+            let mut guard = self.inner.lock().unwrap();
+            guard.model = None;
+            guard.loaded = None;
+            guard.backend.clone()
+        };
+
+        // Step 4: heavy work OUTSIDE the mutex. `is_loaded()` and
+        // `loaded_model_id()` can be polled freely here without parking
+        // tokio worker threads.
+        let backend = match existing_backend {
             Some(existing) => existing,
             None => Arc::new(
                 LlamaBackend::init()
@@ -121,9 +220,11 @@ impl LlmEngine {
 
         let gpu_layers = if use_gpu { u32::MAX } else { 0 };
         let params = LlamaModelParams::default().with_n_gpu_layers(gpu_layers);
-        let model = LlamaModel::load_from_file(&backend, model_path, &params)
-            .map_err(|e| EngineError::LoadFailed(format!("model load: {e}")))?;
+        let model = loader(&backend, model_path, &params)?;
 
+        // Step 5: short crit section — install the new state.
+        {
+            let mut guard = self.inner.lock().unwrap();
             guard.backend = Some(backend);
             guard.model = Some(Arc::new(model));
             guard.loaded = Some(LoadedModelState {
@@ -131,17 +232,60 @@ impl LlmEngine {
                 model_path: model_path.display().to_string(),
                 use_gpu,
             });
+        }
+        // `_loading_guard` drops here and clears the flag.
         Ok(())
     }
 
     pub fn unload(&self) -> Result<(), EngineError> {
         let mut guard = self.inner.lock().unwrap();
         guard.model = None;
-        guard.backend = None;
+        // Backend is process-singleton (llama-cpp-2 enforces this via
+        // `LLAMA_BACKEND_INITIALIZED`). Dropping the Arc here would call
+        // `llama_backend_free` and a subsequent `init` would succeed, but
+        // we keep it resident to avoid the init/free churn on every
+        // load/unload cycle.
         guard.loaded = None;
         Ok(())
     }
 
+    /// True iff a model load is currently in flight. Exposed for tests
+    /// + frontends that want to render a "loading…" state without
+    /// polling `is_loaded()` (which returns false during a swap).
+    pub fn is_loading(&self) -> bool {
+        self.loading.load(Ordering::Acquire)
+    }
+
+    /// Test-only harness: runs `op` while holding the same locking
+    /// discipline as `load_model_with` (loading flag claimed, model
+    /// state cleared, slow op runs OUTSIDE the inner mutex, new state
+    /// installed at the end). Used by the regression test to verify
+    /// that `is_loaded()` / `loaded_model_id()` don't block on the
+    /// slow section. Not part of the public API.
+    #[cfg(test)]
+    pub(crate) fn __test_run_with_lock_discipline<F>(&self, op: F) -> Result<(), EngineError>
+    where
+        F: FnOnce(),
+    {
+        if self
+            .loading
+            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+            .is_err()
+        {
+            return Err(EngineError::AlreadyLoading);
+        }
+        let _loading_guard = LoadingGuard {
+            flag: Arc::clone(&self.loading),
+        };
+        {
+            let mut guard = self.inner.lock().unwrap();
+            guard.model = None;
+            guard.loaded = None;
+        }
+        op();
+        Ok(())
+    }
+
     pub fn is_loaded(&self) -> bool {
         self.inner.lock().unwrap().model.is_some()
     }
@@ -666,4 +810,125 @@ mod tests {
         let n_ctx = preflight_context_window(7_104, 1_024).unwrap();
         assert_eq!(n_ctx, MAX_CONTEXT_TOKENS);
     }
+
+    /// Race-3 regression. The inner `std::sync::Mutex` MUST NOT be held
+    /// across the slow `LlamaModel::load_from_file` call: sync Tauri
+    /// command handlers like `get_llm_status`, `check_llm_model`,
+    /// `delete_llm_model`, and `test_llm_model` call `is_loaded()` /
+    /// `loaded_model_id()` from tokio worker threads without
+    /// `spawn_blocking`. If the lock is held for the duration of a
+    /// 5-15 s load, parallel status polls from the frontend park the
+    /// tokio executor and the whole UI deadlocks.
+    ///
+    /// Pre-fix this test FAILS — both probes time out because the load
+    /// holds the mutex. Post-fix it PASSES — probes return in ≤50 ms.
+    #[test]
+    fn is_loaded_does_not_block_on_slow_load() {
+        use std::sync::mpsc;
+        use std::thread;
+        use std::time::{Duration, Instant};
+
+        let engine = LlmEngine::new();
+        let load_started = Arc::new(std::sync::Barrier::new(2));
+        let release_load = Arc::new(std::sync::Barrier::new(2));
+
+        let engine_for_loader = engine.clone();
+        let load_started_for_loader = Arc::clone(&load_started);
+        let release_load_for_loader = Arc::clone(&release_load);
+
+        let loader_handle = thread::spawn(move || {
+            engine_for_loader
+                .__test_run_with_lock_discipline(|| {
+                    // Signal the probe thread that the load is mid-flight
+                    // (loading flag claimed, inner mutex released).
+                    load_started_for_loader.wait();
+                    // Wait until the probe thread says it's done so the
+                    // load's "duration" is bounded by the probes.
+                    release_load_for_loader.wait();
+                })
+                .unwrap();
+        });
+
+        // Wait until the loader is inside its slow section.
+        load_started.wait();
+
+        // Now probe `is_loaded()` and `loaded_model_id()` from this
+        // thread. They MUST return without contending on the lock.
+        let probe_deadline = Duration::from_millis(50);
+        let (tx, rx) = mpsc::channel();
+        let engine_for_probe = engine.clone();
+        let probe_handle = thread::spawn(move || {
+            let start = Instant::now();
+            let loaded = engine_for_probe.is_loaded();
+            let id = engine_for_probe.loaded_model_id();
+            let loading = engine_for_probe.is_loading();
+            let elapsed = start.elapsed();
+            tx.send((loaded, id, loading, elapsed)).unwrap();
+        });
+
+        let result = rx
+            .recv_timeout(probe_deadline)
+            .expect("is_loaded / loaded_model_id probe must return within 50 ms");
+        let (loaded, id, loading, elapsed) = result;
+
+        assert!(
+            !loaded,
+            "is_loaded() should report false while a load is in flight"
+        );
+        assert_eq!(id, None, "loaded_model_id() should be None mid-load");
+        assert!(loading, "is_loading() should report true mid-load");
+        assert!(
+            elapsed < probe_deadline,
+            "probe took {elapsed:?}, expected < {probe_deadline:?}"
+        );
+
+        probe_handle.join().unwrap();
+        release_load.wait();
+        loader_handle.join().unwrap();
+
+        // After the load completes the flag clears.
+        assert!(!engine.is_loading());
+    }
+
+    /// Race-3 / Race-4 — concurrent load attempts must not both reach
+    /// the heavy work. The second caller should be told `AlreadyLoading`
+    /// rather than starting a parallel load that silently overwrites
+    /// the first.
+    #[test]
+    fn second_concurrent_load_is_refused() {
+        use std::thread;
+
+        let engine = LlmEngine::new();
+        let load_started = Arc::new(std::sync::Barrier::new(2));
+        let release_load = Arc::new(std::sync::Barrier::new(2));
+
+        let engine_for_loader = engine.clone();
+        let load_started_for_loader = Arc::clone(&load_started);
+        let release_load_for_loader = Arc::clone(&release_load);
+
+        let loader_handle = thread::spawn(move || {
+            engine_for_loader
+                .__test_run_with_lock_discipline(|| {
+                    load_started_for_loader.wait();
+                    release_load_for_loader.wait();
+                })
+                .unwrap();
+        });
+
+        load_started.wait();
+
+        // Second concurrent attempt MUST be refused, not parallel-load.
+        let second = engine.__test_run_with_lock_discipline(|| {
+            panic!("second concurrent load should never reach its op");
+        });
+        assert!(matches!(second, Err(EngineError::AlreadyLoading)));
+
+        release_load.wait();
+        loader_handle.join().unwrap();
+
+        // After the first load completes, a fresh attempt is allowed.
+        assert!(engine
+            .__test_run_with_lock_discipline(|| {})
+            .is_ok());
+    }
 }