1 files changed, 26 insertions, 20 deletions
diff --git a/src/config/session.rs b/src/config/session.rs
index 0b48cbd..2149ba5 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -353,30 +353,33 @@ impl Session {
     }
 
     pub fn add_message(&mut self, input: &Input, output: &str) -> Result<()> {
-        match input.continue_output() {
-            Some(_) => {
-                if let Some(message) = self.messages.last_mut() {
-                    if let MessageContent::Text(text) = &mut message.content {
-                        *text = format!("{text}{output}");
-                    }
+        if input.continue_output().is_some() {
+            if let Some(message) = self.messages.last_mut() {
+                if let MessageContent::Text(text) = &mut message.content {
+                    *text = format!("{text}{output}");
                 }
             }
-            None => {
-                let mut need_add_msg = true;
-                if self.messages.is_empty() {
-                    self.messages.extend(input.role().build_messages(input));
-                    need_add_msg = false;
+        } else if input.regenerate() {
+            if let Some(message) = self.messages.last_mut() {
+                if let MessageContent::Text(text) = &mut message.content {
+                    *text = output.to_string();
                 }
-                if need_add_msg {
-                    self.messages
-                        .push(Message::new(MessageRole::User, input.message_content()));
-                }
-                self.data_urls.extend(input.data_urls());
-                self.messages.push(Message::new(
-                    MessageRole::Assistant,
-                    MessageContent::Text(output.to_string()),
-                ));
             }
+        } else {
+            let mut need_add_msg = true;
+            if self.messages.is_empty() {
+                self.messages.extend(input.role().build_messages(input));
+                need_add_msg = false;
+            }
+            if need_add_msg {
+                self.messages
+                    .push(Message::new(MessageRole::User, input.message_content()));
+            }
+            self.data_urls.extend(input.data_urls());
+            self.messages.push(Message::new(
+                MessageRole::Assistant,
+                MessageContent::Text(output.to_string()),
+            ));
         }
         self.dirty = true;
         Ok(())
@@ -398,6 +401,9 @@ impl Session {
         let mut messages = self.messages.clone();
         if input.continue_output().is_some() {
             return messages;
+        } else if input.regenerate() {
+            messages.pop();
+            return messages;
         }
         let mut need_add_msg = true;
         let len = messages.len();