From 4c7de650c00bff792c2a737b84f5eaaa3717e1bb Mon Sep 17 00:00:00 2001
From: Alex Clarke <alex.j.tusa@gmail.com>
Date: Wed, 15 Jul 2026 16:23:59 -0600
Subject: [PATCH] feat: Added an .undo command to the REPL to let users have
 more control over the conversation

---
 src/config/request_context.rs | 20 ++++++++++++++++++++
 src/config/session.rs         |  9 +++++++++
 src/repl/mod.rs               | 14 +++++++++++++-
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/config/request_context.rs b/src/config/request_context.rs
index 0aabfdd..0781c24 100644
--- a/src/config/request_context.rs
+++ b/src/config/request_context.rs
@@ -124,6 +124,7 @@ pub struct RequestContext {
     pub todo_list: TodoList,
     pub skill_registry: SkillRegistry,
     pub last_continuation_response: Option<String>,
+    pub pending_prefill: Option<String>,
 
     pub render_mode: RenderMode,
 }
@@ -154,6 +155,7 @@ impl RequestContext {
             todo_list: TodoList::default(),
             skill_registry: SkillRegistry::default(),
             last_continuation_response: None,
+            pending_prefill: None,
             render_mode: RenderMode::default(),
         }
     }
@@ -210,6 +212,7 @@ impl RequestContext {
             todo_list: TodoList::default(),
             skill_registry: SkillRegistry::default(),
             last_continuation_response: None,
+            pending_prefill: None,
             render_mode: RenderMode::default(),
         })
     }
@@ -253,6 +256,7 @@ impl RequestContext {
             todo_list: self.todo_list.clone(),
             skill_registry: self.skill_registry.clone(),
             last_continuation_response: None,
+            pending_prefill: None,
             render_mode: self.render_mode,
         }
     }
@@ -294,6 +298,7 @@ impl RequestContext {
             todo_list: TodoList::default(),
             skill_registry: SkillRegistry::default(),
             last_continuation_response: None,
+            pending_prefill: None,
             render_mode: parent.render_mode,
         }
     }
@@ -607,6 +612,21 @@ impl RequestContext {
         Ok(())
     }
 
+    pub fn undo_last_exchange(&mut self) -> Result<()> {
+        let text = match self.session.as_mut() {
+            Some(session) => session.pop_last_exchange(),
+            None => bail!("No session"),
+        };
+        match text {
+            Some(text) => {
+                self.pending_prefill = Some(text);
+                self.discontinuous_last_message();
+                Ok(())
+            }
+            None => bail!("Nothing to undo"),
+        }
+    }
+
     pub fn set_save_session_this_time(&mut self) -> Result<()> {
         if let Some(session) = self.session.as_mut() {
             session.set_save_session_this_time();
diff --git a/src/config/session.rs b/src/config/session.rs
index 7e9c24f..37d6694 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -732,6 +732,15 @@ impl Session {
         self.update_tokens();
     }
 
+    pub fn pop_last_exchange(&mut self) -> Option<String> {
+        let user_idx = self.messages.iter().rposition(|m| m.role.is_user())?;
+        let user_text = self.messages[user_idx].content.as_text()?.to_string();
+        self.messages.truncate(user_idx);
+        self.dirty = true;
+        self.update_tokens();
+        Some(user_text)
+    }
+
     pub fn echo_messages(&self, input: &Input) -> String {
         let messages = self.build_messages(input);
         serde_yaml::to_string(&messages).unwrap_or_else(|_| "Unable to echo message".into())
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index ecd1b56..bf8b49b 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -52,7 +52,7 @@ pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {"
     4. Continue with the next pending item now. Call tools immediately."
 };
 
-static REPL_COMMANDS: LazyLock<[ReplCommand; 50]> = LazyLock::new(|| {
+static REPL_COMMANDS: LazyLock<[ReplCommand; 51]> = LazyLock::new(|| {
     [
         ReplCommand::new(".help", "Show this help guide", AssertState::pass()),
         ReplCommand::new(".info", "Show system info", AssertState::pass()),
@@ -125,6 +125,11 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 50]> = LazyLock::new(|| {
             "Clear session messages",
             AssertState::True(StateFlags::SESSION),
         ),
+        ReplCommand::new(
+            ".undo",
+            "Undo the last exchange and restore the prompt",
+            AssertState::True(StateFlags::SESSION),
+        ),
         ReplCommand::new(
             ".compress session",
             "Compress session messages",
@@ -390,6 +395,10 @@ Type ".help" for additional help.
                             if exit {
                                 break;
                             }
+                            if let Some(text) = self.ctx.write().pending_prefill.take() {
+                                self.editor
+                                    .run_edit_commands(&[EditCommand::InsertString(text)]);
+                            }
                         }
                         Err(err) => {
                             render_error(err);
@@ -966,6 +975,9 @@ pub async fn run_repl_command(
                     println!(r#"Usage: .empty session"#)
                 }
             },
+            ".undo" => {
+                ctx.undo_last_exchange()?;
+            }
             ".rebuild" => match args {
                 Some("rag") => {
                     ctx.rebuild_rag(abort_signal.clone()).await?;