Fix undo bug