fix(memory): retry truncated extractions + release job locks on graceful shutdown (#2684)

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

src/openhuman/memory_queue/store.rs

@@ -325,12 +325,41 @@ pub fn recover_stale_locks(config: &Config) -> Result<usize> {
             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<usize> {
+    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<u64> {
     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]

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();

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 \

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<String> {
+            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<String> {
+            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"
+    );
+}