From d2478abee50820057e780cc2e5f030dd703fc01b Mon Sep 17 00:00:00 2001 From: Joshua T Corbin Date: Sat, 21 Feb 2026 11:18:17 -0500 Subject: [PATCH] Improve incremental session save reliability - loop instead of recursion so that there are no incidental limits on tool call depth - save session after *every* message append or compaction rewrite; between each tool call invocation itself, not just at the end of a batch of tool calls --- crates/core/src/agent/mod.rs | 118 ++++++++++++++-------------- crates/core/src/heartbeat/runner.rs | 8 +- 2 files changed, 66 insertions(+), 60 deletions(-) diff --git a/crates/core/src/agent/mod.rs b/crates/core/src/agent/mod.rs index 1a26090c..e6bc916a 100644 --- a/crates/core/src/agent/mod.rs +++ b/crates/core/src/agent/mod.rs @@ -644,70 +644,60 @@ impl Agent { /// Used by the heartbeat runner so the session log is visible while the heartbeat is running. async fn handle_response_saving_session( &mut self, - response: LLMResponse, + mut response: LLMResponse, agent_id: &str, ) -> Result { - // Track usage - self.add_usage(response.usage); - - match response.content { - LLMResponseContent::Text(text) => Ok(text), - LLMResponseContent::ToolCalls(calls) => { - // Execute tool calls - let mut results = Vec::new(); - - for call in &calls { - debug!( - "Executing tool: {} with args: {}", - call.name, call.arguments - ); - - let result = self.execute_tool(call).await; - let output = match result { - Ok((content, _warnings)) => content, - Err(e) => format!("Error: {}", e), - }; - results.push(ToolResult { - call_id: call.id.clone(), - output, - }); - } - - // Add tool call message - self.session.add_message(Message { - role: Role::Assistant, - content: String::new(), - tool_calls: Some(calls), - tool_call_id: None, - images: Vec::new(), - }); - - // Add tool results - for result in &results { + loop { + // Track usage + self.add_usage(response.usage); + + match response.content { + LLMResponseContent::Text(text) => return Ok(text), + LLMResponseContent::ToolCalls(calls) => { + // Add and save intent to call tools so it's visible during a long run self.session.add_message(Message { - role: Role::Tool, - content: result.output.clone(), - tool_calls: None, - tool_call_id: Some(result.call_id.clone()), + role: Role::Assistant, + content: String::new(), + tool_calls: Some(calls.clone()), + tool_call_id: None, images: Vec::new(), }); - } + if let Err(e) = self.session.save_for_agent(agent_id) { + debug!("Incremental session save failed: {}", e); + } - // Save session after each tool call round so it's visible during a long run - if let Err(e) = self.session.save_for_agent(agent_id) { - debug!("Incremental session save failed: {}", e); - } + // Execute each tool call, saving session after each so that partial it's visible + // during a long run, even partial progress if interrupted + for call in &calls { + debug!( + "Executing tool: {} with args: {}", + call.name, call.arguments + ); - // Continue conversation with tool results (with per-turn security block) - let messages = self.messages_for_api_call(); - let tool_schemas = self.tool_schemas_for_provider(); - let next_response = self - .provider - .chat(&messages, Some(tool_schemas.as_slice())) - .await?; + let result = self.execute_tool(call).await; + self.session.add_message(Message { + role: Role::Tool, + content: match result { + Ok((content, _warnings)) => content.clone(), + Err(e) => format!("Error: {}", e), + }, + tool_calls: None, + tool_call_id: Some(call.id.clone()), + images: Vec::new(), + }); + if let Err(e) = self.session.save_for_agent(agent_id) { + debug!("Incremental session save failed: {}", e); + } + } - // Recursively handle (in case of more tool calls) - Box::pin(self.handle_response_saving_session(next_response, agent_id)).await + // Continue conversation with tool results (with per-turn security block) + let messages = self.messages_for_api_call(); + let tool_schemas = self.tool_schemas_for_provider(); + response = self + .provider + .chat(&messages, Some(tool_schemas.as_slice())) + .await?; + } } } } @@ -715,7 +705,7 @@ impl Agent { /// Like `chat`, but saves the session log to `agent_id`'s sessions directory after each /// tool call round. Used by the heartbeat runner so in-progress sessions are visible. pub async fn chat_saving_session(&mut self, message: &str, agent_id: &str) -> Result { - // Add user message + // Add user message and start out saved session file self.session.add_message(Message { role: Role::User, content: message.to_string(), @@ -723,18 +713,29 @@ impl Agent { tool_call_id: None, images: Vec::new(), }); + if let Err(e) = self.session.save_for_agent(agent_id) { + debug!("Incremental session save failed: {}", e); + } // Check if we should run pre-compaction memory flush (soft threshold) if self.should_memory_flush() { info!("Running pre-compaction memory flush (soft threshold)"); self.memory_flush().await?; + if let Err(e) = self.session.save_for_agent(agent_id) { + debug!("Incremental session save failed: {}", e); + } } // Check if we need to compact (hard limit) if self.should_compact() { self.compact_session().await?; + if let Err(e) = self.session.save_for_agent(agent_id) { + debug!("Incremental session save failed: {}", e); + } } + // TODO y u no save + // Build messages for LLM (with per-turn security block) let messages = self.messages_for_api_call(); @@ -760,6 +761,9 @@ impl Agent { tool_call_id: None, images: Vec::new(), }); + if let Err(e) = self.session.save_for_agent(agent_id) { + debug!("Incremental session save failed: {}", e); + } Ok(final_response) } diff --git a/crates/core/src/heartbeat/runner.rs b/crates/core/src/heartbeat/runner.rs index ab1ad031..7a7c9f1b 100644 --- a/crates/core/src/heartbeat/runner.rs +++ b/crates/core/src/heartbeat/runner.rs @@ -407,11 +407,12 @@ impl HeartbeatRunner { // Send heartbeat prompt; save session after each tool call round so the log // is visible while the heartbeat is still running. let heartbeat_prompt = build_heartbeat_prompt(workspace_is_git); - let response = agent + let res = agent .chat_saving_session(&heartbeat_prompt, &self.agent_id) - .await?; + .await; - // Save final session log + // Save final session log, even if the chat failed, and even if this write if futile in the + // happy path, this ensures we at least save at the end match agent.save_session_for_agent(&self.agent_id).await { Ok(path) => { info!(name: "Heartbeat", "saved session: {:?}", path.to_str().unwrap_or("")); @@ -422,6 +423,7 @@ impl HeartbeatRunner { } // Determine status based on response + let response = res?; if is_heartbeat_ok(&response) { return Ok((response, HeartbeatStatus::Ok)); }