Keep auto-suggestions available in Vi navigation mode

src/prompt_toolkit/key_binding/bindings/auto_suggest.py

@@ -7,7 +7,7 @@
 import re
 
 from prompt_toolkit.application.current import get_app
-from prompt_toolkit.filters import Condition, emacs_mode
+from prompt_toolkit.filters import Condition, emacs_mode, vi_navigation_mode
 from prompt_toolkit.key_binding.key_bindings import KeyBindings
 from prompt_toolkit.key_binding.key_processor import KeyPressEvent
 
@@ -18,6 +18,26 @@
 E = KeyPressEvent
 
 
+@Condition
+def suggestion_available() -> bool:
+    app = get_app()
+    buffer = app.current_buffer
+    document = buffer.document
+
+    return (
+        buffer.suggestion is not None
+        and len(buffer.suggestion.text) > 0
+        and (
+            document.is_cursor_at_the_end
+            or (
+                vi_navigation_mode()
+                and document.cursor_position_row == document.line_count - 1
+                and document.cursor_position_col == len(document.current_line) - 1
+            )
+        )
+    )
+
+
 def load_auto_suggest_bindings() -> KeyBindings:
     """
     Key bindings for accepting auto suggestion text.
@@ -29,15 +49,7 @@ def load_auto_suggest_bindings() -> KeyBindings:
     key_bindings = KeyBindings()
     handle = key_bindings.add
 
-    @Condition
-    def suggestion_available() -> bool:
-        app = get_app()
-        return (
-            app.current_buffer.suggestion is not None
-            and len(app.current_buffer.suggestion.text) > 0
-            and app.current_buffer.document.is_cursor_at_the_end
-        )
-
+    @handle("l", filter=suggestion_available & vi_navigation_mode)
     @handle("c-f", filter=suggestion_available)
     @handle("c-e", filter=suggestion_available)
     @handle("right", filter=suggestion_available)
@@ -49,6 +61,8 @@ def _accept(event: E) -> None:
         suggestion = b.suggestion
 
         if suggestion:
+            if vi_navigation_mode():
+                b.cursor_position = len(b.text)
             b.insert_text(suggestion.text)
 
     @handle("escape", "f", filter=suggestion_available & emacs_mode)

src/prompt_toolkit/layout/processors.py

@@ -16,7 +16,12 @@
 from prompt_toolkit.application.current import get_app
 from prompt_toolkit.cache import SimpleCache
 from prompt_toolkit.document import Document
-from prompt_toolkit.filters import FilterOrBool, to_filter, vi_insert_multiple_mode
+from prompt_toolkit.filters import (
+    FilterOrBool,
+    to_filter,
+    vi_insert_multiple_mode,
+    vi_navigation_mode,
+)
 from prompt_toolkit.formatted_text import (
     AnyFormattedText,
     StyleAndTextTuples,
@@ -591,7 +596,13 @@ def apply_transformation(self, ti: TransformationInput) -> Transformation:
         if ti.lineno == ti.document.line_count - 1:
             buffer = ti.buffer_control.buffer
 
-            if buffer.suggestion and ti.document.is_cursor_at_the_end:
+            at_end = ti.document.is_cursor_at_the_end or (
+                vi_navigation_mode()
+                and ti.document.cursor_position_row == ti.document.line_count - 1
+                and ti.document.cursor_position_col == len(ti.document.current_line) - 1
+            )
+
+            if buffer.suggestion and at_end:
                 suggestion = buffer.suggestion.text
             else:
                 suggestion = ""

tests/test_cli.py

@@ -9,6 +9,7 @@
 
 import pytest
 
+from prompt_toolkit.auto_suggest import Suggestion
 from prompt_toolkit.clipboard import ClipboardData, InMemoryClipboard
 from prompt_toolkit.enums import EditingMode
 from prompt_toolkit.filters import ViInsertMode
@@ -17,6 +18,8 @@
 from prompt_toolkit.input.vt100_parser import ANSI_SEQUENCES
 from prompt_toolkit.key_binding.bindings.named_commands import prefix_meta
 from prompt_toolkit.key_binding.key_bindings import KeyBindings
+from prompt_toolkit.key_binding.vi_state import InputMode
+from prompt_toolkit.layout.processors import AppendAutoSuggestion, TransformationInput
 from prompt_toolkit.output import DummyOutput
 from prompt_toolkit.shortcuts import PromptSession
 
@@ -893,6 +896,41 @@ def test_vi_temp_navigation_mode():
     assert result.cursor_position == 5
 
 
+@pytest.mark.parametrize("accept_key", ["l", "\x1b[C"])
+def test_vi_navigation_keeps_and_accepts_suggestion(accept_key):
+    with create_pipe_input() as inp:
+        inp.send_text(accept_key + "\r")
+        session = PromptSession(
+            input=inp,
+            output=DummyOutput(),
+            editing_mode=EditingMode.VI,
+        )
+
+        displayed = []
+
+        def pre_run():
+            buffer = session.default_buffer
+            buffer.suggestion = Suggestion(" hello")
+            buffer.cursor_position = len(buffer.text) - 1
+            session.app.vi_state.input_mode = InputMode.NAVIGATION
+
+            transformation = AppendAutoSuggestion().apply_transformation(
+                TransformationInput(
+                    session.app.layout.current_control,
+                    buffer.document,
+                    0,
+                    lambda i: i,
+                    [("", buffer.text)],
+                    80,
+                    24,
+                )
+            )
+            displayed.extend(text for _, text in transformation.fragments)
+
+        assert session.prompt(default="echo", pre_run=pre_run) == "echo hello"
+        assert "".join(displayed) == "echo hello"
+
+
 def test_vi_macros():
     feed = partial(_feed_cli_with_input, editing_mode=EditingMode.VI)