fix: undo() peeks without popping, breaking repeated undo and leaking redo stack (#71)

* fix: undo() peeks without popping, breaking repeated undo and leaking redo stack

undo() used self.undo_stack.back() (peek) instead of pop_back(), so
repeated undo always returned the same checkpoint while pushing to
the redo stack unboundedly.

Additionally, redo() did not save the current state to the undo stack,
breaking the undo/redo cycle.

Changes:
- undo(): change back() to pop_back(), return owned Checkpoint
- redo(): accept current_turn/current_messages params, save current
  state to undo stack before popping from redo stack
- Update process_undo/process_redo callers in agent_loop.rs
- Add tests for repeated undo, undo/redo cycling, stack size invariant

* fix: standardize lock ordering and extract push_undo helper

Address review feedback:
- Standardize lock order (Session before UndoManager) in process_undo
  and process_redo to match process_user_input and prevent deadlocks
- Extract push_undo() helper to deduplicate push-and-trim logic shared
  by checkpoint() and redo()

* docs: add move-semantics notes and stack invariant to UndoManager

Address review feedback requesting documentation about the ownership
semantics of undo/redo parameters and the stack size invariant.

---------

Co-authored-by: Yi LIU <[email protected]>
Co-authored-by: firat.sertgoz <[email protected]>

src/agent/agent_loop.rs

@@ -1579,6 +1579,9 @@ impl Agent {
         session: Arc<Mutex<Session>>,
         thread_id: Uuid,
     ) -> Result<SubmissionResult, Error> {
+        // Lock session first, then undo manager -- consistent with process_user_input
+        // to avoid potential deadlocks.
+        let mut sess = session.lock().await;
         let undo_mgr = self.session_manager.get_undo_manager(thread_id).await;
         let mut mgr = undo_mgr.lock().await;
 
@@ -1586,7 +1589,6 @@ impl Agent {
             return Ok(SubmissionResult::ok_with_message("Nothing to undo."));
         }
 
-        let mut sess = session.lock().await;
         let thread = sess
             .threads
             .get_mut(&thread_id)
@@ -1597,12 +1599,10 @@ impl Agent {
         let current_turn = thread.turn_number();
 
         if let Some(checkpoint) = mgr.undo(current_turn, current_messages) {
-            // Extract values before consuming the reference
             let turn_number = checkpoint.turn_number;
-            let messages = checkpoint.messages.clone();
             let undo_count = mgr.undo_count();
             // Restore thread from checkpoint
-            thread.restore_from_messages(messages);
+            thread.restore_from_messages(checkpoint.messages);
             Ok(SubmissionResult::ok_with_message(format!(
                 "Undone to turn {}. {} undo(s) remaining.",
                 turn_number, undo_count
@@ -1617,6 +1617,9 @@ impl Agent {
         session: Arc<Mutex<Session>>,
         thread_id: Uuid,
     ) -> Result<SubmissionResult, Error> {
+        // Lock session first, then undo manager -- consistent with process_user_input
+        // to avoid potential deadlocks.
+        let mut sess = session.lock().await;
         let undo_mgr = self.session_manager.get_undo_manager(thread_id).await;
         let mut mgr = undo_mgr.lock().await;
 
@@ -1624,12 +1627,15 @@ impl Agent {
             return Ok(SubmissionResult::ok_with_message("Nothing to redo."));
         }
 
-        if let Some(checkpoint) = mgr.redo() {
-            let mut sess = session.lock().await;
-            let thread = sess
-                .threads
-                .get_mut(&thread_id)
-                .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
+        let thread = sess
+            .threads
+            .get_mut(&thread_id)
+            .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
+
+        let current_messages = thread.messages();
+        let current_turn = thread.turn_number();
+
+        if let Some(checkpoint) = mgr.redo(current_turn, current_messages) {
             thread.restore_from_messages(checkpoint.messages);
             Ok(SubmissionResult::ok_with_message(format!(
                 "Redone to turn {}.",

src/agent/undo.rs

@@ -43,6 +43,10 @@ impl Checkpoint {
 }
 
 /// Manager for undo/redo functionality.
+///
+/// Each undo/redo operation pops from one stack and pushes the current state
+/// onto the other, so `undo_count() + redo_count()` stays constant across
+/// undo/redo cycles (only `checkpoint()` and `clear()` change the total).
 pub struct UndoManager {
     /// Stack of past checkpoints (for undo).
     undo_stack: VecDeque<Checkpoint>,
@@ -68,6 +72,14 @@ impl UndoManager {
         self
     }
 
+    /// Push a checkpoint onto the undo stack, trimming oldest entries if over limit.
+    fn push_undo(&mut self, checkpoint: Checkpoint) {
+        self.undo_stack.push_back(checkpoint);
+        while self.undo_stack.len() > self.max_checkpoints {
+            self.undo_stack.pop_front();
+        }
+    }
+
     /// Create a checkpoint at the current state.
     ///
     /// This clears the redo stack since we're creating a new history branch.
@@ -80,24 +92,23 @@ impl UndoManager {
         // Clear redo stack (new branch of history)
         self.redo_stack.clear();
 
-        // Create and push checkpoint
         let checkpoint = Checkpoint::new(turn_number, messages, description);
-        self.undo_stack.push_back(checkpoint);
-
-        // Trim if over limit
-        while self.undo_stack.len() > self.max_checkpoints {
-            self.undo_stack.pop_front();
-        }
+        self.push_undo(checkpoint);
     }
 
     /// Undo: pop the last checkpoint and return it.
     ///
-    /// The current state should be saved to redo stack before calling this.
+    /// Saves the current state to the redo stack and pops the most recent
+    /// checkpoint from the undo stack so that repeated undos walk backwards
+    /// through history.
+    ///
+    /// Takes ownership of `current_messages`; callers must clone first if
+    /// they need to retain a copy.
     pub fn undo(
         &mut self,
         current_turn: usize,
         current_messages: Vec<ChatMessage>,
-    ) -> Option<&Checkpoint> {
+    ) -> Option<Checkpoint> {
         if self.undo_stack.is_empty() {
             return None;
         }
@@ -110,9 +121,8 @@ impl UndoManager {
         );
         self.redo_stack.push(current);
 
-        // Return the most recent checkpoint without removing it
-        // (we keep it so multiple undos can work)
-        self.undo_stack.back()
+        // Pop and return the most recent checkpoint
+        self.undo_stack.pop_back()
     }
 
     /// Pop the last checkpoint from the undo stack.
@@ -121,7 +131,29 @@ impl UndoManager {
     }
 
     /// Redo: restore a previously undone state.
-    pub fn redo(&mut self) -> Option<Checkpoint> {
+    ///
+    /// Saves the current state to the undo stack and pops the most recent
+    /// checkpoint from the redo stack.
+    ///
+    /// Takes ownership of `current_messages`; callers must clone first if
+    /// they need to retain a copy.
+    pub fn redo(
+        &mut self,
+        current_turn: usize,
+        current_messages: Vec<ChatMessage>,
+    ) -> Option<Checkpoint> {
+        if self.redo_stack.is_empty() {
+            return None;
+        }
+
+        // Save current state to undo stack
+        let current = Checkpoint::new(
+            current_turn,
+            current_messages,
+            format!("Turn {}", current_turn),
+        );
+        self.push_undo(current);
+
         self.redo_stack.pop()
     }
 
@@ -214,14 +246,16 @@ mod tests {
         assert!(manager.can_undo());
         assert!(!manager.can_redo());
 
-        // Undo
+        // Undo - returns owned Checkpoint now
         let current = vec![ChatMessage::user("Hello"), ChatMessage::assistant("Hi")];
         let checkpoint = manager.undo(2, current);
         assert!(checkpoint.is_some());
+        let checkpoint = checkpoint.unwrap();
+        assert_eq!(checkpoint.turn_number, 1);
         assert!(manager.can_redo());
 
-        // Redo
-        let restored = manager.redo();
+        // Redo - now requires current state parameters
+        let restored = manager.redo(checkpoint.turn_number, checkpoint.messages);
         assert!(restored.is_some());
     }
 
@@ -249,4 +283,90 @@ mod tests {
         assert!(restored.is_some());
         assert_eq!(manager.undo_count(), 0);
     }
+
+    #[test]
+    fn test_repeated_undo_advances_through_stack() {
+        let mut manager = UndoManager::new();
+
+        // Create 3 checkpoints at turns 0, 1, 2
+        manager.checkpoint(0, vec![], "Turn 0");
+        manager.checkpoint(1, vec![ChatMessage::user("msg1")], "Turn 1");
+        manager.checkpoint(2, vec![ChatMessage::user("msg2")], "Turn 2");
+        assert_eq!(manager.undo_count(), 3);
+
+        // First undo: should return turn 2 checkpoint, stack shrinks to 2
+        let cp1 = manager
+            .undo(3, vec![ChatMessage::user("msg3")])
+            .expect("first undo should succeed");
+        assert_eq!(cp1.turn_number, 2);
+        assert_eq!(manager.undo_count(), 2);
+
+        // Second undo: should return turn 1 checkpoint (different!), stack shrinks to 1
+        let cp2 = manager
+            .undo(cp1.turn_number, cp1.messages)
+            .expect("second undo should succeed");
+        assert_eq!(cp2.turn_number, 1);
+        assert_eq!(manager.undo_count(), 1);
+
+        // Verify we walked backwards through distinct checkpoints
+        assert_ne!(cp1.turn_number, cp2.turn_number);
+    }
+
+    #[test]
+    fn test_undo_redo_cycle_preserves_state() {
+        let mut manager = UndoManager::new();
+
+        let msgs_t0: Vec<ChatMessage> = vec![];
+        let msgs_t1 = vec![ChatMessage::user("hello")];
+        let msgs_t2 = vec![ChatMessage::user("hello"), ChatMessage::assistant("hi")];
+
+        manager.checkpoint(0, msgs_t0, "Turn 0");
+        manager.checkpoint(1, msgs_t1, "Turn 1");
+
+        // Undo from turn 2 -> get turn 1 checkpoint
+        let cp_undo1 = manager
+            .undo(2, msgs_t2.clone())
+            .expect("undo should succeed");
+        assert_eq!(cp_undo1.turn_number, 1);
+
+        // Redo from turn 1 -> get turn 2 state back
+        let cp_redo = manager
+            .redo(cp_undo1.turn_number, cp_undo1.messages)
+            .expect("redo should succeed");
+        assert_eq!(cp_redo.turn_number, 2);
+        assert_eq!(cp_redo.messages.len(), 2);
+
+        // Undo again from turn 2 -> should go back to turn 1 again
+        let cp_undo2 = manager
+            .undo(cp_redo.turn_number, cp_redo.messages)
+            .expect("second undo should succeed");
+        assert_eq!(cp_undo2.turn_number, 1);
+    }
+
+    #[test]
+    fn test_undo_redo_stack_sizes_consistent() {
+        let mut manager = UndoManager::new();
+
+        manager.checkpoint(0, vec![], "Turn 0");
+        manager.checkpoint(1, vec![ChatMessage::user("a")], "Turn 1");
+        manager.checkpoint(2, vec![ChatMessage::user("b")], "Turn 2");
+
+        // Start: undo=3, redo=0, total=3
+        let total = manager.undo_count() + manager.redo_count();
+        assert_eq!(total, 3);
+
+        // After undo: total should still be 3 (one moved from undo to redo,
+        // plus the current state pushed to redo)
+        // Actually: undo pops one (3->2), pushes current to redo (0->1), total=3
+        let cp = manager.undo(3, vec![]).unwrap();
+        assert_eq!(manager.undo_count() + manager.redo_count(), 3);
+
+        // After redo: redo pops one (1->0), pushes current to undo (2->3), total=3
+        let cp2 = manager.redo(cp.turn_number, cp.messages).unwrap();
+        assert_eq!(manager.undo_count() + manager.redo_count(), 3);
+
+        // After another undo: same invariant
+        let _cp3 = manager.undo(cp2.turn_number, cp2.messages).unwrap();
+        assert_eq!(manager.undo_count() + manager.redo_count(), 3);
+    }
 }