diff --git a/src/openhuman/memory_queue/store.rs b/src/openhuman/memory_queue/store.rs index 8bfef4dc1..f4deb2d33 100644 --- a/src/openhuman/memory_queue/store.rs +++ b/src/openhuman/memory_queue/store.rs @@ -325,12 +325,41 @@ pub fn recover_stale_locks(config: &Config) -> Result { params![now_ms], )?; if n > 0 { - log::warn!("[memory::jobs] recovered {n} stale-locked job(s) at startup"); + // Info, not warn: with graceful shutdown releasing in-flight + // locks (see `release_running_locks`), a non-empty recovery now + // means the previous process was hard-killed — expected and + // self-healing, not actionable (bug-report-2026-05-26 I2). + log::info!("[memory::jobs] recovered {n} stale-locked job(s) at startup"); } Ok(n) }) } +/// Release this process's in-flight job locks on a *graceful* shutdown: +/// flip every `running` row back to `ready` so the work is immediately +/// re-claimable on next launch instead of waiting out the lease and +/// surfacing as a stale-lock recovery. The core runs a single worker pool, +/// so any `running` row at clean-shutdown time was claimed by us. +/// Returns the number of rows released (bug-report-2026-05-26 I2). +pub fn release_running_locks(config: &Config) -> Result { + with_connection(config, |conn| { + // TODO(multi-process): this releases *every* `running` row, with no + // process-scoped guard (e.g. a `worker_pid` / session-token column). + // Correct today because a single worker pool is the invariant, but if + // multi-process workers are ever introduced, a rolling restart with + // two overlapping processes would let one shutdown release the other's + // in-flight locks. Add a per-owner predicate here when that happens. + let n = conn.execute( + "UPDATE mem_tree_jobs + SET status = 'ready', + locked_until_ms = NULL + WHERE status = 'running'", + [], + )?; + Ok(n) + }) +} + /// Quick count helper for tests / diagnostics. pub fn count_by_status(config: &Config, status: JobStatus) -> Result { with_connection(config, |conn| { @@ -610,6 +639,27 @@ mod tests { assert_eq!(row.status, JobStatus::Ready); } + #[test] + fn release_running_locks_resets_running_rows_regardless_of_lease() { + let (_tmp, cfg) = test_config(); + let payload = ExtractChunkPayload { + chunk_id: "c-release".into(), + }; + let nj = NewJob::extract_chunk(&payload).unwrap(); + let id = enqueue(&cfg, &nj).unwrap().unwrap(); + + // Claim with a long, still-valid lease. `recover_stale_locks` would + // NOT touch this (not expired), but a graceful shutdown release must + // so a clean restart re-claims it immediately (bug-report-2026-05-26 I2). + let _ = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + + let released = release_running_locks(&cfg).unwrap(); + assert_eq!(released, 1); + let row = get_job(&cfg, &id).unwrap().unwrap(); + assert_eq!(row.status, JobStatus::Ready); + assert!(row.locked_until_ms.is_none()); + } + /// Happy path: a non-stale settlement still succeeds after the claim-token /// check is applied. Regression guard so the common case isn't broken. #[test] diff --git a/src/openhuman/memory_queue/worker.rs b/src/openhuman/memory_queue/worker.rs index fb6574243..4a41d4794 100644 --- a/src/openhuman/memory_queue/worker.rs +++ b/src/openhuman/memory_queue/worker.rs @@ -18,7 +18,7 @@ use crate::openhuman::config::Config; use crate::openhuman::memory_queue::handlers; use crate::openhuman::memory_queue::redact::scrub_for_log; use crate::openhuman::memory_queue::store::{ - claim_next, mark_deferred, mark_done, mark_failed, recover_stale_locks, + claim_next, mark_deferred, mark_done, mark_failed, recover_stale_locks, release_running_locks, DEFAULT_LOCK_DURATION_MS, }; use crate::openhuman::memory_queue::types::JobOutcome; @@ -70,6 +70,36 @@ pub fn start(config: Config) { log::warn!("[memory::jobs] recover_stale_locks failed at startup: {err:#}"); } + // Release in-flight locks on graceful shutdown so a clean restart + // re-claims the work immediately instead of waiting out the lease + // (which surfaced as a stale-lock recovery warn on every launch). + // Hard kills still fall back to lease-expiry recovery at startup + // (bug-report-2026-05-26 I2). + let shutdown_cfg = config.clone(); + crate::core::shutdown::register(move || { + // NOTE: `shutdown::register` is bound `F: Fn() -> Fut`, so this + // closure may be invoked more than once; each call must hand the + // returned future its own owned `Config`. Moving `shutdown_cfg` + // in directly is `E0507` (cannot move out of an `Fn` closure), so + // the per-call clone is required, not redundant. + let cfg = shutdown_cfg.clone(); + async move { + match release_running_locks(&cfg) { + Ok(n) if n > 0 => { + log::info!( + "[memory::jobs] released {n} in-flight job lock(s) on graceful shutdown" + ); + } + Ok(_) => {} + Err(err) => { + log::warn!( + "[memory::jobs] failed to release job locks on shutdown: {err:#}" + ); + } + } + } + }); + for idx in 0..WORKER_COUNT { let notify = notify.clone(); let cfg = config.clone(); diff --git a/src/openhuman/memory_tree/score/extract/llm.rs b/src/openhuman/memory_tree/score/extract/llm.rs index 51d075053..78d948485 100644 --- a/src/openhuman/memory_tree/score/extract/llm.rs +++ b/src/openhuman/memory_tree/score/extract/llm.rs @@ -213,6 +213,20 @@ impl LlmEntityExtractor { let parsed: LlmExtractionOutput = match serde_json::from_str(&raw) { Ok(v) => v, + Err(e) if e.is_eof() => { + // Truncated mid-JSON: the response stream closed before the + // closing brace (token budget or a server-side timeout cut + // it short). Unlike a wrong-shape body, this is transient — + // signal a retryable failure (`None`) so `extract`'s backoff + // loop tries again rather than silently dropping the whole + // chunk (bug-report-2026-05-26 I1). + log::warn!( + "[memory_tree::extract::llm] LLM response truncated mid-JSON ({e}); \ + response_bytes={} — retrying", + raw.len() + ); + return None; + } Err(e) => { log::warn!( "[memory_tree::extract::llm] LLM returned non-JSON or wrong-shape \ diff --git a/src/openhuman/memory_tree/score/extract/llm_tests.rs b/src/openhuman/memory_tree/score/extract/llm_tests.rs index b72ac3da1..4e6d6a2e2 100644 --- a/src/openhuman/memory_tree/score/extract/llm_tests.rs +++ b/src/openhuman/memory_tree/score/extract/llm_tests.rs @@ -428,3 +428,93 @@ fn truncate_for_log_long_input_appends_ellipsis() { assert_eq!(out.chars().count(), 11); // 10 + "…" assert!(out.ends_with('…')); } + +#[tokio::test] +async fn extract_retries_on_truncated_response() { + // First response is truncated mid-JSON (serde EOF) — a stream cutoff, + // not a wrong-shape body. It must be treated as retryable rather than + // silently dropped; the second (complete) response then recovers the + // entities (bug-report-2026-05-26 I1). + use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct TruncatedThenCompleteProvider { + calls: AtomicUsize, + } + #[async_trait] + impl ChatProvider for TruncatedThenCompleteProvider { + fn name(&self) -> &str { + "test:truncated" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + let n = self.calls.fetch_add(1, Ordering::SeqCst); + if n == 0 { + // Array never closes → serde reports EOF (is_eof()). + Ok(r#"{"entities":[{"kind":"person","text":"Alice"}"#.to_string()) + } else { + Ok(r#"{"entities":[{"kind":"person","text":"Alice"}],"importance":0.5,"importance_reason":"r"}"#.to_string()) + } + } + } + + let mock = Arc::new(TruncatedThenCompleteProvider { + calls: AtomicUsize::new(0), + }); + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); + let out = ex.extract("Alice met Bob.").await.unwrap(); + assert_eq!( + mock.calls.load(Ordering::SeqCst), + 2, + "truncation should trigger a retry, not a silent drop" + ); + assert_eq!(out.entities.len(), 1); + assert_eq!(out.entities[0].text, "Alice"); +} + +#[tokio::test] +async fn extract_does_not_retry_on_wrong_shape_response() { + // Companion to the truncation test: a *complete* but wrong-shape body is + // a serde error that is NOT EOF (`is_eof() == false`). Unlike a mid-JSON + // cutoff it's deterministic and won't fix itself on retry, so it must + // return immediately (`calls == 1`) with an empty extraction rather than + // entering the retry loop. Guards the `Err(e) if e.is_eof()` split so a + // future change can't accidentally retry every parse error + // (bug-report-2026-05-26 I1). + use crate::openhuman::memory::chat::{ChatPrompt, ChatProvider}; + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct WrongShapeProvider { + calls: AtomicUsize, + } + #[async_trait] + impl ChatProvider for WrongShapeProvider { + fn name(&self) -> &str { + "test:wrong-shape" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + // `entities` must be an array; a scalar is a complete-input type + // error (not a truncation), so serde reports a non-EOF error. + Ok(r#"{"entities":123}"#.to_string()) + } + } + + let mock = Arc::new(WrongShapeProvider { + calls: AtomicUsize::new(0), + }); + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); + let out = ex.extract("Alice met Bob.").await.unwrap(); + assert_eq!( + mock.calls.load(Ordering::SeqCst), + 1, + "a non-EOF wrong-shape response must not trigger a retry" + ); + assert!( + out.entities.is_empty(), + "wrong-shape response should yield an empty extraction" + ); +}