fix(infinity-daemon): settle tool calls orphaned by a daemon restart

crates/infinity-agent-core/src/system/mod.rs

@@ -41,7 +41,7 @@ pub use defer::{DeferQueue, InMemoryDeferQueue, NoDeferral};
 pub use events::{AgentEvent, ReplaySnapshot, UserChoice};
 pub use model::{ModelSource, ResolvedModel, StaticModel};
 pub use observer::{EventCollector, ThreadObserver};
-pub use thread::StepOutcome;
+pub use thread::{StepOutcome, tool_call_expects_result};
 
 #[cfg(test)]
 pub(crate) mod test_support;

crates/infinity-agent-core/src/system/tests.rs

@@ -382,6 +382,27 @@ async fn driver_idles_after_close_thread_tool_call() {
         .await;
 }
 
+/// Likewise for `sleep_until_event_or_input`: the next event or input
+/// respawns the driver.
+#[tokio::test(flavor = "current_thread")]
+async fn driver_idles_after_sleep_until_event_or_input_call() {
+    let local = tokio::task::LocalSet::new();
+    local
+        .run_until(async {
+            let (mut running, mut rx, mut ctrl, _conv) = start_system(vec![], None);
+            running
+                .send_user_text(ThreadId::from_ref("t1"), "wait")
+                .await;
+            let _req = ctrl.next_request().await;
+            ctrl.send_tool_call("tc-1", "sleep_until_event_or_input", serde_json::json!({}));
+            ctrl.finish();
+            collect_until_finished(&mut rx).await;
+            wait_idle(&mut running).await;
+            assert!(running.is_idle());
+        })
+        .await;
+}
+
 /// Thread reports arriving while waiting for a non-passive async tool result
 /// are deferred until the tool result is processed.
 #[tokio::test(flavor = "current_thread")]

crates/infinity-agent-core/src/system/thread.rs

@@ -29,6 +29,16 @@ type BoxError = Box<dyn std::error::Error + Send + Sync>;
 /// "waiting for a tool result".
 const CLOSE_THREAD_TOOL: &str = "close_thread";
 
+/// The name of the built-in indefinite-wait tool. Its no-op call is never
+/// answered either: the next event or input interrupts it, respawning the
+/// driver through the router.
+const SLEEP_UNTIL_EVENT_OR_INPUT_TOOL: &str = "sleep_until_event_or_input";
+
+/// Whether an unanswered call to `tool_name` will ever receive a result.
+pub fn tool_call_expects_result(tool_name: &str) -> bool {
+    tool_name != CLOSE_THREAD_TOOL && tool_name != SLEEP_UNTIL_EVENT_OR_INPUT_TOOL
+}
+
 /// The result of one step (see [`AgentSystem::step`](super::AgentSystem::step)).
 #[derive(Debug)]
 pub enum StepOutcome {
@@ -226,15 +236,15 @@ where
     }
 
     /// Whether the thread must stay resident: the result of an unanswered
-    /// tool call (other than `close_thread`, whose result never arrives) is
-    /// still on its way. Active subscriptions do not keep a thread resident;
-    /// their events wake it through the router when they arrive.
+    /// tool call (see [`tool_call_expects_result`]) is still on its way.
+    /// Active subscriptions do not keep a thread resident; their events wake
+    /// it through the router when they arrive.
     ///
     /// Synchronous so a driver can make its exit decision atomically with
     /// respect to the router on the shared `LocalSet`.
     pub fn awaiting_tool_result(&self) -> bool {
         self.pending_tool_call()
-            .is_some_and(|(_, name)| name != CLOSE_THREAD_TOOL)
+            .is_some_and(|(_, name)| tool_call_expects_result(&name))
     }
 
     /// Apply the deferral policy to a batch of inputs.

crates/infinity-daemon/Cargo.toml

@@ -26,7 +26,7 @@ rap-client = { path = "../rap-client", version = "^0.1.0" }
 rap-protocol = { path = "../rap-protocol", version = "^0.1.0" }
 
 serde = { workspace = true }
-serde_json = { workspace = true }
+serde_json = { workspace = true, features = ["raw_value"] }
 tokio = { workspace = true, features = ["rt-multi-thread", "io-util", "sync", "process", "net", "signal", "macros", "time"] }
 uuid = { workspace = true }
 tracing = { workspace = true }

crates/infinity-daemon/src/memory_store.rs

@@ -18,10 +18,10 @@ use infinity_agent_core::message::InfinityMessage;
 use infinity_agent_core::stores::{
     self as core_stores, CompactionSummary, ThreadInfo, ThreadState,
 };
-use infinity_agent_core::system::UserChoice;
+use infinity_agent_core::system::{UserChoice, tool_call_expects_result};
 use infinity_agent_core::traits::{ConversationStore, SpawnContext, StateStore};
 use infinity_protocol::ModelRef;
-use infinity_provider_protocol::message::Message;
+use infinity_provider_protocol::message::{Message, ToolCall};
 use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
 use std::path::{Path, PathBuf};
@@ -129,6 +129,30 @@ fn unset_model_ref() -> ModelRef {
     }
 }
 
+/// The trailing tool call of the thread snapshot at `path`, if it expects a
+/// result. Only the tail message is deserialized, so probing every thread at
+/// startup stays cheap. Other files and legacy snapshots have none.
+fn tail_call_expecting_result(path: &Path) -> Option<ToolCall> {
+    #[derive(Deserialize)]
+    struct Probe<'a> {
+        #[serde(borrow)]
+        messages: Vec<(&'a serde_json::value::RawValue, serde::de::IgnoredAny)>,
+    }
+    let json = std::fs::read(path)
+        .inspect_err(|e| tracing::error!("cannot read {}: {e}", path.display()))
+        .ok()?;
+    let (tail, _) = serde_json::from_slice::<Probe>(&json)
+        .ok()?
+        .messages
+        .pop()?;
+    match serde_json::from_str(tail.get()).ok()? {
+        InfinityMessage::ToolCall { call, .. } if tool_call_expects_result(&call.function.name) => {
+            Some(call)
+        }
+        _ => None,
+    }
+}
+
 /// Per-thread snapshot written to `{dir}/{thread_id}.json`.
 #[derive(Serialize)]
 pub(crate) struct ThreadSnapshot {
@@ -456,6 +480,54 @@ impl PersistentConversationStore {
         self.core.thread_info(thread_id).is_some()
     }
 
+    /// Cold start: settle each thread's trailing tool call whose result can
+    /// no longer arrive (its RAP server, input queue or timer died with the
+    /// previous daemon) with an [`InfinityMessage::Interrupt`], as an
+    /// interrupting message would. Otherwise any driver woken for the thread, even just to
+    /// replay it to a viewer, parks on the call and holds its session
+    /// running. Must run before the agent system starts, so that every
+    /// unanswered call is the previous process's. Returns how many calls were
+    /// settled.
+    pub async fn settle_orphaned_tool_calls(&self) -> usize {
+        let Some(dir) = &self.dir else { return 0 };
+        let entries = match std::fs::read_dir(dir) {
+            Ok(entries) => entries,
+            Err(e) => {
+                tracing::error!("cannot scan {} for orphaned tool calls: {e}", dir.display());
+                return 0;
+            }
+        };
+        let mut settled = 0;
+        for entry in entries {
+            let path = match entry {
+                Ok(entry) => entry.path(),
+                Err(e) => {
+                    tracing::error!("cannot scan {} for orphaned tool calls: {e}", dir.display());
+                    continue;
+                }
+            };
+            let Some(thread_id) = path.file_stem().and_then(|s| s.to_str()) else {
+                continue;
+            };
+            let Some(call) = tail_call_expecting_result(&path) else {
+                continue;
+            };
+            tracing::info!(%thread_id, tool = %call.function.name, "settling tool call orphaned by daemon restart");
+            let interrupt = (
+                InfinityMessage::Interrupt,
+                format!("{}-orphaned-by-restart", call.id),
+            );
+            match self
+                .append_messages(ThreadId::from_ref(thread_id), vec![interrupt])
+                .await
+            {
+                Ok(()) => settled += 1,
+                Err(e) => tracing::error!(%thread_id, "failed to settle orphaned tool call: {e}"),
+            }
+        }
+        settled
+    }
+
     /// Resolve a thread ID to its root thread ID (i.e. the session ID).
     pub fn get_root_thread_id(&self, thread_id: &ThreadId<str>) -> ThreadId {
         self.ensure_thread_metadata_loaded(thread_id);
@@ -1628,4 +1700,61 @@ mod tests {
             assert_eq!(t.text, "c3");
         }
     }
+
+    /// Cold-start settlement answers only calls whose results can no longer
+    /// arrive, persists the interrupt, and is idempotent.
+    #[tokio::test]
+    async fn settle_orphaned_tool_calls_answers_only_orphaned_calls() {
+        let dir = tempfile::tempdir().expect("create temp dir");
+        let new_store = || {
+            PersistentConversationStore::new_with_dir(
+                dir.path(),
+                test_model(),
+                Arc::new(crate::ids::UuidIdSource),
+            )
+        };
+        let store = new_store();
+        for tool in [
+            "execute_command",
+            "sleep_until_event_or_input",
+            "close_thread",
+        ] {
+            let id = ThreadId::from_ref(tool);
+            store
+                .ensure_root_thread(id)
+                .await
+                .expect("ensure root thread");
+            let call = InfinityMessage::ToolCall {
+                call: ToolCall {
+                    id: format!("tc-{tool}"),
+                    call_id: None,
+                    function: infinity_provider_protocol::message::ToolFunction {
+                        name: tool.to_owned(),
+                        arguments: serde_json::json!({}),
+                    },
+                },
+                display_as: None,
+            };
+            store
+                .append_messages(id, vec![(call, "c".to_owned())])
+                .await
+                .expect("append tail tool call");
+        }
+
+        // Restart: a fresh store over the same directory.
+        let store = new_store();
+        assert_eq!(store.settle_orphaned_tool_calls().await, 1);
+        let tail = store
+            .core
+            .thread_messages(ThreadId::from_ref("execute_command"))
+            .and_then(|m| m.last().cloned())
+            .expect("thread has history");
+        assert!(
+            matches!(tail.0, InfinityMessage::Interrupt),
+            "expected an interrupt settling the orphaned call, got {tail:?}"
+        );
+
+        // Persisted, so a second restart finds nothing left to settle.
+        assert_eq!(new_store().settle_orphaned_tool_calls().await, 0);
+    }
 }

crates/infinity-daemon/src/session/mod.rs

@@ -171,6 +171,8 @@ impl SessionManager {
             id_source.clone(),
         );
         conversation_store.set_change_tx(change_tx_for_conv);
+        // Before the agent system starts, so every unanswered call is orphaned.
+        conversation_store.settle_orphaned_tool_calls().await;
         let state_store = PersistentStateStore::new(
             state_dir.join("state"),
             conversation_store.clone(),

crates/infinity-daemon/src/session/tests.rs

@@ -778,3 +778,71 @@ async fn shut_down_session_events_do_not_wake_threads() {
         })
         .await;
 }
+
+/// A thread whose history ends in a tool call from the previous daemon
+/// process is settled with an interrupt when the session manager starts.
+#[tokio::test(flavor = "current_thread")]
+async fn startup_settles_orphaned_tool_calls() {
+    let local = tokio::task::LocalSet::new();
+    local
+        .run_until(async {
+            let (model, _ctrl) = mock_model();
+            let entry = ModelEntry {
+                model_id: "model1".to_owned(),
+                display_name: "model1".to_owned(),
+                context_window: 100_000,
+                max_output_tokens: None,
+                supports_image_input: false,
+            };
+            let state_dir = tempfile::tempdir().expect("create daemon state dir");
+            let threads_dir = state_dir.path().join("threads");
+            let conv = PersistentConversationStore::new_with_dir(
+                &threads_dir,
+                model1_ref(),
+                Arc::new(SequentialIdSource::new()),
+            );
+            let killed = ThreadId::from_ref("killed");
+            conv.ensure_root_thread(killed).await.expect("ensure root");
+            let call = infinity_agent_core::message::InfinityMessage::ToolCall {
+                call: infinity_provider_protocol::message::ToolCall {
+                    id: "tc-exec".to_owned(),
+                    call_id: None,
+                    function: infinity_provider_protocol::message::ToolFunction {
+                        name: "execute_command".to_owned(),
+                        arguments: serde_json::json!({}),
+                    },
+                },
+                display_as: None,
+            };
+            conv.append_messages(killed, vec![(call, "c1".to_owned())])
+                .await
+                .expect("append orphaned call");
+
+            SessionManager::with_providers(
+                SessionManagerConfig {
+                    state_dir: state_dir.path().to_path_buf(),
+                    callback_url: "http://127.0.0.1:0".to_owned(),
+                    user_rap_config: None,
+                    id_source: Arc::new(SequentialIdSource::new()),
+                },
+                vec![(
+                    "provider1".to_owned(),
+                    Arc::new(SingleModelProvider::new(entry, model)) as Arc<dyn ModelProvider>,
+                )],
+                vec![],
+            )
+            .await
+            .expect("build session manager");
+
+            let snapshot: serde_json::Value = serde_json::from_str(
+                &std::fs::read_to_string(threads_dir.join("killed.json"))
+                    .expect("read thread snapshot"),
+            )
+            .expect("parse thread snapshot");
+            assert_eq!(
+                snapshot["messages"][1][0]["kind"], "interrupt",
+                "the orphaned call must be settled at startup: {snapshot}"
+            );
+        })
+        .await;
+}