From 3d640b9efa3561d709e67ff818ada800a45b38f3 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 27 Aug 2026 12:49:07 -0600 Subject: [PATCH] fix: checkpoint sessions that crash for easy resuming --- src/config/request_context.rs | 73 +++++++++++++++++++++++++++++++++- src/config/session.rs | 74 ++++++++++++++++++++++++++++++++++- src/repl/mod.rs | 7 +--- 3 files changed, 144 insertions(+), 10 deletions(-) diff --git a/src/config/request_context.rs b/src/config/request_context.rs index a4a2836..b8a7368 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -1,6 +1,6 @@ use super::bundles::BundleStore; use super::rag_cache::{RagCache, RagKey}; -use super::session::Session; +use super::session::{INTERRUPTED_RESPONSE_TEXT, Session}; use super::skill::{SKILL_SCAFFOLD, Skill}; use super::skill_policy::SkillPolicy; use super::skill_registry::SkillRegistry; @@ -958,13 +958,25 @@ impl RequestContext { let mut i = input.clone(); i.clear_patch(); if let Some(session) = i.session_mut(&mut self.session) { - let _ = session.add_message(&i, "[Response interrupted due to error]"); + let _ = session.add_message(&i, INTERRUPTED_RESPONSE_TEXT); if !app.dry_run && session.save_session() == Some(true) { let _ = session.flush(); } } } + pub fn has_recoverable_interruption(&self) -> bool { + let live = self + .last_message + .as_ref() + .map(|v| v.continuous && v.input.with_session()) + .unwrap_or(false); + live || self + .session + .as_ref() + .is_some_and(Session::has_interrupted_error_checkpoint) + } + pub fn discontinuous_last_message(&mut self) { if let Some(last_message) = self.last_message.as_mut() { last_message.continuous = false; @@ -6739,6 +6751,63 @@ mod tests { ); } + #[test] + fn has_recoverable_interruption_false_by_default() { + let ctx = create_test_ctx(); + assert!(!ctx.has_recoverable_interruption()); + } + + #[test] + fn has_recoverable_interruption_true_for_live_interrupted_turn() { + let mut ctx = create_test_ctx(); + ctx.session = Some(Session::default()); + let app = Arc::clone(&ctx.app.config); + let input = Input::from_str(&ctx, "hello", None).unwrap(); + + ctx.on_chat_completion_error(app.as_ref(), &input); + + assert!(ctx.has_recoverable_interruption()); + } + + #[test] + fn has_recoverable_interruption_false_after_normal_exchange() { + let mut ctx = create_test_ctx(); + ctx.session = Some(Session::default()); + let input = Input::from_str(&ctx, "hello", None).unwrap(); + ctx.session + .as_mut() + .unwrap() + .add_message(&input, "all done") + .unwrap(); + + assert!(!ctx.has_recoverable_interruption()); + } + + #[test] + fn has_recoverable_interruption_survives_session_save_and_reload() { + // Simulates the full user flow: a turn crashes mid-tool-loop (the + // checkpoint lands in the session), the user runs `.save session`, + // exits, and a fresh process resumes the session. The in-memory + // last_message is gone; only the persisted checkpoint can mark the + // session recoverable for `.recover`. + let mut ctx = create_test_ctx(); + ctx.session = Some(Session::default()); + let app = Arc::clone(&ctx.app.config); + let input = Input::from_str(&ctx, "hello", None).unwrap(); + ctx.on_chat_completion_error(app.as_ref(), &input); + + let yaml = serde_yaml::to_string(ctx.session.as_ref().unwrap()).unwrap(); + let reloaded: Session = serde_yaml::from_str(&yaml).unwrap(); + + let mut resumed_ctx = create_test_ctx(); + resumed_ctx.session = Some(reloaded); + assert!(resumed_ctx.last_message.is_none()); + assert!( + resumed_ctx.has_recoverable_interruption(), + ".recover must work on a session resumed after an interrupted turn" + ); + } + #[test] fn after_chat_completion_sweeps_auto_unload_skills_at_turn_end() { let mut ctx = create_test_ctx(); diff --git a/src/config/session.rs b/src/config/session.rs index 1e3b042..3413294 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -16,6 +16,8 @@ use std::sync::LazyLock; static RE_AUTONAME_PREFIX: LazyLock = LazyLock::new(|| Regex::new(r"\d{8}T\d{6}-").unwrap()); +pub const INTERRUPTED_RESPONSE_TEXT: &str = "[Response interrupted due to error]"; + #[derive(Debug, Clone, Default, Deserialize, Serialize)] pub struct Session { #[serde(rename(serialize = "model", deserialize = "model"))] @@ -175,6 +177,13 @@ impl Session { self.messages.is_empty() && self.compressed_messages.is_empty() } + pub fn has_interrupted_error_checkpoint(&self) -> bool { + self.messages.last().is_some_and(|message| { + message.role.is_assistant() + && matches!(&message.content, MessageContent::Text(text) if text == INTERRUPTED_RESPONSE_TEXT) + }) + } + pub fn messages(&self) -> &[Message] { &self.messages } @@ -934,9 +943,9 @@ impl AutoName { #[cfg(test)] mod tests { use super::*; - use crate::client::{Message, MessageContent, MessageRole, Model}; + use crate::client::{Message, MessageContent, MessageContentToolCalls, MessageRole, Model}; use crate::config::{AppConfig, AppState, RequestContext, WorkingMode}; - use crate::function::Functions; + use crate::function::{Functions, ToolCall, ToolResult}; use std::sync::Arc; #[test] @@ -949,6 +958,67 @@ mod tests { assert!(!session.dirty()); } + fn push_interrupted_turn(session: &mut Session) { + session.messages.push(Message::new( + MessageRole::User, + MessageContent::Text("do things".to_string()), + )); + session.messages.push(Message::new( + MessageRole::Tool, + MessageContent::ToolCalls(MessageContentToolCalls::new( + vec![ToolResult::new(ToolCall::default(), json!("ok"))], + String::new(), + )), + )); + session.messages.push(Message::new( + MessageRole::Assistant, + MessageContent::Text(INTERRUPTED_RESPONSE_TEXT.to_string()), + )); + } + + #[test] + fn session_has_interrupted_error_checkpoint_detects_sentinel() { + let mut session = Session::default(); + assert!(!session.has_interrupted_error_checkpoint()); + + session.messages.push(Message::new( + MessageRole::User, + MessageContent::Text("hi".to_string()), + )); + session.messages.push(Message::new( + MessageRole::Assistant, + MessageContent::Text("hello".to_string()), + )); + assert!( + !session.has_interrupted_error_checkpoint(), + "a normal completed exchange is not an interruption" + ); + + push_interrupted_turn(&mut session); + assert!(session.has_interrupted_error_checkpoint()); + } + + #[test] + fn session_interrupted_checkpoint_with_tool_calls_survives_yaml_round_trip() { + let mut session = Session::default(); + push_interrupted_turn(&mut session); + + let yaml = serde_yaml::to_string(&session).unwrap(); + let reloaded: Session = serde_yaml::from_str(&yaml).unwrap(); + + assert!( + reloaded.has_interrupted_error_checkpoint(), + "interruption checkpoint must survive save/reload" + ); + assert!( + reloaded.messages.iter().any(|m| matches!( + &m.content, + MessageContent::ToolCalls(tc) if tc.tool_results.len() == 1 + )), + "tool calls made before the crash must survive save/reload" + ); + } + #[test] fn session_enabled_macros_absent_is_none() { let session: Session = serde_yaml::from_str("model: provider:test\nmessages: []").unwrap(); diff --git a/src/repl/mod.rs b/src/repl/mod.rs index c53202c..e0a573e 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -1191,12 +1191,7 @@ pub async fn run_repl_command( ask(ctx, abort_signal.clone(), input, true).await?; } ".recover" => { - let has_recoverable = ctx - .last_message - .as_ref() - .map(|v| v.continuous && v.input.with_session()) - .unwrap_or(false); - if !has_recoverable { + if !ctx.has_recoverable_interruption() { bail!("Unable to recover: no interrupted session response to recover from"); } let recovery_text = args