diff --git a/plugins/vi-mode/vi-mode.plugin.zsh b/plugins/vi-mode/vi-mode.plugin.zsh
index c1cdd2851..a82e41439 100644
--- a/plugins/vi-mode/vi-mode.plugin.zsh
+++ b/plugins/vi-mode/vi-mode.plugin.zsh
@@ -1,14 +1,12 @@
 # Updates editor information when the keymap changes.
 function zle-keymap-select() {
+  # update keymap variable for the prompt
+  VI_KEYMAP=$KEYMAP
+
   zle reset-prompt
   zle -R
 }
 
-# Ensure that the prompt is redrawn when the terminal size changes.
-TRAPWINCH() {
-  zle && { zle -R; zle reset-prompt }
-}
-
 zle -N zle-keymap-select
 zle -N edit-command-line
 
@@ -42,7 +40,7 @@ if [[ "$MODE_INDICATOR" == "" ]]; then
 fi
 
 function vi_mode_prompt_info() {
-  echo "${${KEYMAP/vicmd/$MODE_INDICATOR}/(main|viins)/}"
+  echo "${${VI_KEYMAP/vicmd/$MODE_INDICATOR}/(main|viins)/}"
 }
 
 # define right prompt, if it wasn't defined by a theme