From e6dc24beb5d516c9780532f221822e6f69c9f90a Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Tue, 4 Aug 2026 12:02:31 -0600 Subject: [PATCH] feat: add a new .recover command for sessions to recover from errors --- src/config/input.rs | 4 +++ src/config/request_context.rs | 68 +++++++++++++++++++++++++++++++++++ src/config/session.rs | 13 +++++++ src/repl/mod.rs | 60 +++++++++++++++++++++++-------- 4 files changed, 130 insertions(+), 15 deletions(-) diff --git a/src/config/input.rs b/src/config/input.rs index 27ecbd5..71f82d5 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -321,6 +321,10 @@ impl Input { } } + pub fn with_session(&self) -> bool { + self.with_session + } + pub fn with_agent(&self) -> bool { self.with_agent } diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 0b1b20e..2cc61cd 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -728,6 +728,26 @@ impl RequestContext { Ok(()) } + pub fn on_chat_completion_error(&mut self, app: &AppConfig, input: &Input) { + self.last_message = Some(LastMessage::new(input.clone(), String::new())); + if input.session(&self.session).is_none() { + if let Some(lm) = self.last_message.as_mut() { + lm.continuous = false; + } + + return; + } + + 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]"); + if !app.dry_run && session.save_session() == Some(true) { + let _ = session.flush(); + } + } + } + pub fn discontinuous_last_message(&mut self) { if let Some(last_message) = self.last_message.as_mut() { last_message.continuous = false; @@ -5424,6 +5444,54 @@ mod tests { assert!(lm.continuous); } + #[test] + fn on_chat_completion_error_without_session_sets_last_message_discontinuous() { + let mut ctx = create_test_ctx(); + 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 lm = ctx.last_message.as_ref().unwrap(); + assert_eq!(lm.output, ""); + assert!(!lm.continuous, "no session means recovery is not possible"); + } + + #[test] + fn on_chat_completion_error_with_session_sets_last_message_continuous() { + let mut ctx = create_test_ctx(); + ctx.app = Arc::new(AppState { + config: Arc::new(AppConfig { + dry_run: true, + ..(*ctx.app.config).clone() + }), + ..(*ctx.app).clone() + }); + 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 lm = ctx.last_message.as_ref().unwrap(); + assert_eq!(lm.output, ""); + assert!(lm.continuous, "session present means .recover is available"); + } + + #[test] + fn on_chat_completion_error_with_session_checkpoints_session_messages() { + let mut ctx = create_test_ctx(); + ctx.session = Some(Session::default()); + assert!(ctx.session.as_ref().unwrap().is_empty()); + 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.session.as_ref().unwrap().is_empty(), + "session should have the interrupted turn checkpointed" + ); + } + #[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 98cc304..0b593b7 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -699,6 +699,19 @@ impl Session { Ok(()) } + pub fn flush(&mut self) -> Result<()> { + if !self.dirty { + return Ok(()); + } + + if let Some(path) = self.path.clone() { + let name = self.name.clone(); + self.save(&name, Path::new(&path), false)?; + } + + Ok(()) + } + pub fn guard_empty(&self) -> Result<()> { if !self.is_empty() { bail!( diff --git a/src/repl/mod.rs b/src/repl/mod.rs index e636947..3e84f10 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -53,7 +53,7 @@ pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {" 4. Continue with the next pending item now. Call tools immediately." }; -static REPL_COMMANDS: LazyLock<[ReplCommand; 58]> = LazyLock::new(|| { +static REPL_COMMANDS: LazyLock<[ReplCommand; 59]> = LazyLock::new(|| { [ ReplCommand::new(".help", "Show this help guide", AssertState::pass()), ReplCommand::new(".info", "Show system info", AssertState::pass()), @@ -278,6 +278,11 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 58]> = LazyLock::new(|| { "Continue previous response", AssertState::pass(), ), + ReplCommand::new( + ".recover", + "Recover interrupted session after API error or Ctrl+C", + AssertState::pass(), + ), ReplCommand::new( ".regenerate", "Regenerate last response", @@ -1101,6 +1106,21 @@ pub async fn run_repl_command( input.set_continue_output(&output); 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 { + bail!("Unable to recover: no interrupted session response to recover from"); + } + let recovery_text = args + .unwrap_or("Please continue from where you left off.") + .to_string(); + let recovery_input = Input::from_str(ctx, &recovery_text, None)?; + ask(ctx, abort_signal.clone(), recovery_input, false).await?; + } ".regenerate" => { let LastMessage { mut input, .. } = match ctx.last_message.as_ref().filter(|v| v.continuous).cloned() { @@ -1307,18 +1327,28 @@ async fn ask( let client = input.create_client()?; ctx.before_chat_completion(&input)?; - let (output, tool_results) = if input.stream() { - call_chat_completions_streaming(&input, client.as_ref(), ctx, abort_signal.clone()).await? - } else { - call_chat_completions( - &input, - true, - false, - client.as_ref(), - ctx, - abort_signal.clone(), - ) - .await? + let (output, tool_results) = { + let result = if input.stream() { + call_chat_completions_streaming(&input, client.as_ref(), ctx, abort_signal.clone()) + .await + } else { + call_chat_completions( + &input, + true, + false, + client.as_ref(), + ctx, + abort_signal.clone(), + ) + .await + }; + match result { + Ok(v) => v, + Err(err) => { + ctx.on_chat_completion_error(app.as_ref(), &input); + return Err(err); + } + } }; ctx.after_chat_completion(app.as_ref(), &input, &output, &tool_results)?; if !tool_results.is_empty() { @@ -1681,8 +1711,8 @@ mod tests { } #[test] - fn repl_commands_has_58_entries() { - assert_eq!(REPL_COMMANDS.len(), 58); + fn repl_commands_has_59_entries() { + assert_eq!(REPL_COMMANDS.len(), 59); } #[test]