diff --git a/src/widgets.zsh b/src/widgets.zsh
index 56b0cca..190ddce 100644
--- a/src/widgets.zsh
+++ b/src/widgets.zsh
@@ -166,14 +166,13 @@ _zsh_autosuggest_partial_accept() {
 
 	# In vicmd keymap, the cursor is placed on the current character
 	# instead of after it. We increment it by 1 to keep the following logic
-	# working even in this case.
-	CURSOR_POS=$CURSOR
-	if [[ "$KEYMAP" = "vicmd" ]]; then
-		(( CURSOR_POS++ ))
+	# working even in thi case.
+	if [ $KEYMAP = 'vicmd' ]; then
+		(( CURSOR++ ))
 	fi
 
 	# If we've moved past the end of the original buffer
-	if (( $CURSOR_POS > $#original_buffer )); then
+	if (( $CURSOR > $#original_buffer )); then
 		# Set POSTDISPLAY to text right of the cursor
 		POSTDISPLAY="$RBUFFER"
 
@@ -184,6 +183,11 @@ _zsh_autosuggest_partial_accept() {
 		BUFFER="$original_buffer"
 	fi
 
+	# Restore CURSOR
+	if [ $KEYMAP = 'vicmd' ]; then
+		(( CURSOR-- ))
+	fi
+
 	return $retval
 }
 
diff --git a/zsh-autosuggestions.zsh b/zsh-autosuggestions.zsh
index 4aa85df..931598f 100644
--- a/zsh-autosuggestions.zsh
+++ b/zsh-autosuggestions.zsh
@@ -444,14 +444,13 @@ _zsh_autosuggest_partial_accept() {
 
 	# In vicmd keymap, the cursor is placed on the current character
 	# instead of after it. We increment it by 1 to keep the following logic
-	# working even in this case.
-	CURSOR_POS=$CURSOR
-	if [[ "$KEYMAP" = "vicmd" ]]; then
-		(( CURSOR_POS++ ))
+	# working even in thi case.
+	if [ $KEYMAP = 'vicmd' ]; then
+		(( CURSOR++ ))
 	fi
 
 	# If we've moved past the end of the original buffer
-	if (( $CURSOR_POS > $#original_buffer )); then
+	if (( $CURSOR > $#original_buffer )); then
 		# Set POSTDISPLAY to text right of the cursor
 		POSTDISPLAY="$RBUFFER"
 
@@ -462,6 +461,11 @@ _zsh_autosuggest_partial_accept() {
 		BUFFER="$original_buffer"
 	fi
 
+	# Restore CURSOR
+	if [ $KEYMAP = 'vicmd' ]; then
+		(( CURSOR-- ))
+	fi
+
 	return $retval
 }