From e499795b8c859efcfa95fe9aeabd6b946c5351a5 Mon Sep 17 00:00:00 2001 From: AI-Reviewer-QS Date: Wed, 18 Feb 2026 00:35:19 +0800 Subject: [PATCH] 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 Co-authored-by: firat.sertgoz --- src/agent/agent_loop.rs | 26 ++++--- src/agent/undo.rs | 152 +++++++++++++++++++++++++++++++++++----- 2 files changed, 152 insertions(+), 26 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 968f0ba9..b1015fe8 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -1579,6 +1579,9 @@ impl Agent { session: Arc>, thread_id: Uuid, ) -> Result { + // 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>, thread_id: Uuid, ) -> Result { + // 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 {}.", diff --git a/src/agent/undo.rs b/src/agent/undo.rs index 5307a77a..892f7c88 100644 --- a/src/agent/undo.rs +++ b/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, @@ -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, - ) -> Option<&Checkpoint> { + ) -> Option { 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 { + /// + /// 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, + ) -> Option { + 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 = 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); + } }