diff --git a/crates/llm/src/lib.rs b/crates/llm/src/lib.rs index f215f58..aec5c9d 100644 --- a/crates/llm/src/lib.rs +++ b/crates/llm/src/lib.rs @@ -32,7 +32,10 @@ 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.")] + #[error( + "Another LLM load is already in flight; refusing to start a parallel load \ + or modify engine state mid-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" @@ -244,6 +247,18 @@ impl LlmEngine { } 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(); guard.model = None; // 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. 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"); + } }