Add support for vi mode