agent: lumotia — Phase B.7 close unload-during-load TOCTOU on LlmEngine

Phase B.7 audit of commit cde985d (LlmEngine critical-section narrowing
+ drop-old-model-first; Race-3 + Lifecycle-1). Existing coverage is
strong: is_loaded_does_not_block_on_slow_load proves probes return in
< 50 ms while the slow section runs (Race-3); second_concurrent_load_is_refused
proves a parallel load attempt is rejected with EngineError::AlreadyLoading
without reaching the heavy op (Race-3/4 TOCTOU at the engine layer);
the test harness __test_run_with_lock_discipline mirrors load_model_with's
discipline (claim loading flag, clear engine state, run op outside the
inner mutex, then install). The Lifecycle-1 visible side-effect
(is_loaded reports false mid-swap) is covered by the first test.

One real residual found.

unload() does not consult the `loading` flag. When load_model_with is
mid-flight (step 3 has already cleared model + loaded, step 5 has not
yet installed the new state), a concurrent unload() takes the inner
mutex, sees model + loaded already None, no-op-clears, and returns Ok.
The slow load then completes step 5 and installs the new state —
silently overwriting the unload the caller already saw success for.

Concrete attack shape: app startup auto-loads the default LLM in the
background via download_llm_model + load_model. User opens Settings,
clicks "Delete Model X". delete_llm_model checks loaded_model_id()
(returns None mid-load) and skips the unload branch, then calls
model_manager::delete_model(X) which removes the GGUF file from disk.
The load completes via mmap (which on Linux holds the inode alive
after unlink) and installs state pointing at a deleted file path. The
user sees "Model X loaded" in the UI even though they just deleted it.

Same `loading` AtomicBool that guards load-vs-load needs to guard
unload-vs-load.

Fix:
  * unload() now checks is_loading() at entry. Returns
    EngineError::AlreadyLoading when a load is mid-flight; caller can
    retry once is_loading() reports false.
  * EngineError::AlreadyLoading message generalised from "refusing to
    start a parallel load" to "refusing to start a parallel load or
    modify engine state mid-load", since the variant now fires from
    both directions. The variant name itself remains accurate (the
    state of being already loading).

Behavioural diff for unload during quiescent state: unchanged.

Behavioural diff for unload mid-load: Err(AlreadyLoading) instead of
Ok with silent overwrite.

Callers checked:
  * unload_llm_model (Tauri command) — converts EngineError → String
    via .map_err and surfaces to the frontend. New error string is
    self-explanatory; no frontend code matches on the old message
    substring.
  * delete_llm_model — calls unload only when loaded_model_id matches.
    If unload returns AlreadyLoading the delete also fails;
    .map_err(|e| e.to_string())? propagates. The user gets a clear
    "cannot unload while loading" toast and can retry; better than the
    silent contract-violation the old code allowed.
  * No other callers exist for LlmEngine::unload (whisper/parakeet
    engines have their own unload methods on a different type).

New regression test: unload_during_load_is_refused. Spins a loader
thread on the existing __test_run_with_lock_discipline harness, blocks
mid-slow-section via a Barrier, fires unload() from the main thread,
asserts AlreadyLoading. After releasing the load, unload() succeeds —
proving the flag-clear discipline on the happy path.

Verification:
  * cargo test -p lumotia-llm --lib
      → 26/26 pass including the new test.
  * cargo fmt --check → clean (applied fmt after the edit).
  * cargo clippy -p lumotia-llm --all-targets -- -D warnings → clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-14 19:41:57 +01:00
parent 7f0e1b0375
commit f252c1b50e

View File

@@ -32,7 +32,10 @@ pub enum EngineError {
NotLoaded, NotLoaded,
#[error("LLM load failed: {0}")] #[error("LLM load failed: {0}")]
LoadFailed(String), LoadFailed(String),
#[error("Another LLM load is already in flight; refusing to start a parallel load.")] #[error(
"Another LLM load is already in flight; refusing to start a parallel load \
or modify engine state mid-load."
)]
AlreadyLoading, AlreadyLoading,
#[error( #[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" "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"
@@ -244,6 +247,18 @@ impl LlmEngine {
} }
pub fn unload(&self) -> Result<(), EngineError> { pub fn unload(&self) -> Result<(), EngineError> {
// Refuse to unload mid-load. Without this check, `load_model_with`
// is mid-flight (it has cleared `model` / `loaded` in step 3 and
// is about to install new state in step 5); a concurrent unload
// would do nothing (the state is already None), return Ok, and
// then the load's step 5 silently overwrites — the caller saw
// unload success but the engine ends up loaded. Phase B.7 audit
// residual (2026-05-14): the load-vs-load TOCTOU was closed by
// `AlreadyLoading` in cde985d but the unload-vs-load race was
// left open. Same flag covers both directions.
if self.is_loading() {
return Err(EngineError::AlreadyLoading);
}
let mut guard = self.inner.lock().unwrap(); let mut guard = self.inner.lock().unwrap();
guard.model = None; guard.model = None;
// Backend is process-singleton (llama-cpp-2 enforces this via // Backend is process-singleton (llama-cpp-2 enforces this via
@@ -935,4 +950,59 @@ mod tests {
// After the first load completes, a fresh attempt is allowed. // After the first load completes, a fresh attempt is allowed.
assert!(engine.__test_run_with_lock_discipline(|| {}).is_ok()); assert!(engine.__test_run_with_lock_discipline(|| {}).is_ok());
} }
/// Phase B.7 audit regression (2026-05-14). The cde985d fix
/// introduced the `loading` AtomicBool to refuse a second concurrent
/// load, but left `unload()` blind to the flag. A concurrent unload
/// during a load's slow window observed `model == None` (the load's
/// step 3 had already cleared state), no-op-cleared the same nulls,
/// returned Ok — and then the load's step 5 silently installed the
/// new state. The caller saw unload-success but the engine ended up
/// loaded.
///
/// Post-fix: an unload mid-load is refused with
/// `EngineError::AlreadyLoading`. The caller can retry once the
/// load completes (signalled by `is_loading() == false`).
#[test]
fn unload_during_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();
});
// Wait until the loader is mid-slow-section (loading flag claimed,
// engine state cleared).
load_started.wait();
// Unload while the load is in flight MUST be refused, not silently
// no-op'd and then overwritten by the load's install step.
let result = engine.unload();
assert!(
matches!(result, Err(EngineError::AlreadyLoading)),
"unload during a load must surface AlreadyLoading, got {result:?}"
);
release_load.wait();
loader_handle.join().unwrap();
// After the load completes the flag clears and unload succeeds.
assert!(!engine.is_loading());
engine
.unload()
.expect("unload after load completes must succeed");
}
} }