diff --git a/spec/integrations/vi_mode_spec.rb b/spec/integrations/vi_mode_spec.rb
new file mode 100644
index 0000000..95f9293
--- /dev/null
+++ b/spec/integrations/vi_mode_spec.rb
@@ -0,0 +1,21 @@
+describe 'when using vi mode' do
+  let(:before_sourcing) do
+    -> do
+      session.run_command('bindkey -v')
+    end
+  end
+
+  describe 'moving the cursor after exiting insert mode' do
+    it 'should not clear the current suggestion' do
+      with_history('foobar foo') do
+        session.
+          send_string('foo').
+          send_keys('escape').
+          send_keys('h')
+
+        wait_for { session.content }.to eq('foobar foo')
+      end
+    end
+  end
+end
+
diff --git a/src/widgets.zsh b/src/widgets.zsh
index 8a31133..2004518 100644
--- a/src/widgets.zsh
+++ b/src/widgets.zsh
@@ -55,6 +55,7 @@ _zsh_autosuggest_modify() {
 
 	# Don't fetch a new suggestion if there's more input to be read immediately
 	if (( $PENDING > 0 )) || (( $KEYS_QUEUED_COUNT > 0 )); then
+		POSTDISPLAY="$orig_postdisplay"
 		return $retval
 	fi
 
diff --git a/zsh-autosuggestions.zsh b/zsh-autosuggestions.zsh
index a1037ba..bc84bfd 100644
--- a/zsh-autosuggestions.zsh
+++ b/zsh-autosuggestions.zsh
@@ -333,6 +333,7 @@ _zsh_autosuggest_modify() {
 
 	# Don't fetch a new suggestion if there's more input to be read immediately
 	if (( $PENDING > 0 )) || (( $KEYS_QUEUED_COUNT > 0 )); then
+		POSTDISPLAY="$orig_postdisplay"
 		return $retval
 	fi