fixes vi_mode update

fixes Powerlevel9k#1201

generator/default.p9k

@@ -359,24 +359,35 @@ function __p9k_preexec() {
   _P9K_TIMER_START=${EPOCHREALTIME}
 }
 
-p9k::set_default P9K_PROMPT_ADD_NEWLINE false
-###############################################################
+##############################################################
 # @description
-#   This function is a hook that is run before the prompts are created.
-#   If sets all the required variables for the prompts and then
-#   calls the prompt segment building functions.
+#   This function is a hook that is run before __p9k_prepare_prompts.
+#   If sets the return value(s) that is/are used in prompt_status.
+#   Since we need to get the return value of the last command, this hook
+#   has to be run before __p9k_prepare_prompts does.
 ##
 # @noargs
 ##
-function __p9k_prepare_prompts() {
-  # Return values. These need to be global, because
+function __p9k_save_retvals() {
   # they are used in prompt_status. Also, we need
   # to get the return value of the last command at
   # very first in this function. Do not move the
   # lines down, otherwise the last command is not
   # what you expected it to be.
   RETVAL=$?
   RETVALS=( "$pipestatus[@]" )
+}
+
+p9k::set_default P9K_PROMPT_ADD_NEWLINE false
+###############################################################
+# @description
+#   This function is a hook that is run before the prompts are created.
+#   If sets all the required variables for the prompts and then
+#   calls the prompt segment building functions.
+##
+# @noargs
+##
+function __p9k_prepare_prompts() {
   local RPROMPT_PREFIX RPROMPT_SUFFIX
 
   _P9K_COMMAND_DURATION=$((EPOCHREALTIME - _P9K_TIMER_START))
@@ -520,6 +531,7 @@ function prompt_powerlevel9k_setup() {
   autoload -Uz add-zsh-hook
 
   # prepare prompts
+  add-zsh-hook precmd __p9k_save_retvals
   add-zsh-hook precmd __p9k_prepare_prompts
   add-zsh-hook preexec __p9k_preexec
 

segments/vi_mode/vi_mode.p9k

@@ -215,6 +215,7 @@ function zle-keymap-select {
       rebuild_vi_mode "${VIM_MODE_KEYMAP:-$KEYMAP}"
     ;;
     *)
+      __p9k_prepare_prompts
       # About .reset-prompt see:
       # https://github.com/sorin-ionescu/prezto/issues/1026
       # https://github.com/zsh-users/zsh-autosuggestions/issues/107#issuecomment-183824034