Skip to content

Improve incremental session save reliability #62

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yiwang merged 1 commit into from
Feb 21, 2026
Merged

Improve incremental session save reliability #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 61 additions & 57 deletions crates/core/src/agent/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -644,97 +644,98 @@ 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<String> {
// 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?;
}
}
}
}

/// 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<String> {
// Add user message
// Add user message and start out saved session file
self.session.add_message(Message {
role: Role::User,
content: message.to_string(),
tool_calls: None,
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();

Expand All @@ -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)
}
Expand Down
8 changes: 5 additions & 3 deletions crates/core/src/heartbeat/runner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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("<Unknown>"));
Expand All @@ -422,6 +423,7 @@ impl HeartbeatRunner {
}

// Determine status based on response
let response = res?;
if is_heartbeat_ok(&response) {
return Ok((response, HeartbeatStatus::Ok));
}
Expand Down
Loading