From 3f34ce4da401a18dc36c95f0c41280b544b02b6a Mon Sep 17 00:00:00 2001
From: CrispyConductor <crispy@cluenet.org>
Date: Mon, 15 Mar 2021 18:32:49 -0400
Subject: [PATCH] fix(vi-mode): fix keymap indicator on zle-line-finish (#9620)

* Update VI_KEYMAP on zle-line-finish.  Fixes next prompt draw.

* Also fix vi-mode prompt redraw in case previous prompt was canceled

* Do not redraw fix prompt if VI_MODE_RESET_PROMPT_ON_MODE_CHANGE is unset

* Use a local variable instead of global in function
---
 plugins/vi-mode/vi-mode.plugin.zsh | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/plugins/vi-mode/vi-mode.plugin.zsh b/plugins/vi-mode/vi-mode.plugin.zsh
index 7f14961ad..dd4afdd8b 100644
--- a/plugins/vi-mode/vi-mode.plugin.zsh
+++ b/plugins/vi-mode/vi-mode.plugin.zsh
@@ -50,13 +50,17 @@ zle -N zle-keymap-select
 # These "echoti" statements were originally set in lib/key-bindings.zsh
 # Not sure the best way to extend without overriding.
 function zle-line-init() {
+  local prev_vi_keymap
+  prev_vi_keymap="${VI_KEYMAP:-}"
   typeset -g VI_KEYMAP=main
+  [[ "$prev_vi_keymap" != 'main' ]] && [[ "${VI_MODE_RESET_PROMPT_ON_MODE_CHANGE:-}" = true ]] && zle reset-prompt
   (( ! ${+terminfo[smkx]} )) || echoti smkx
   _vi-mode-set-cursor-shape-for-keymap "${VI_KEYMAP}"
 }
 zle -N zle-line-init
 
 function zle-line-finish() {
+  typeset -g VI_KEYMAP=main
   (( ! ${+terminfo[rmkx]} )) || echoti rmkx
   _vi-mode-set-cursor-shape-for-keymap default
 }