feat: dismissal and acceptWord commands

README.md

@@ -168,6 +168,28 @@ inoremap <c-y> <cmd>call augment#Accept()<cr>
 inoremap <cr> <cmd>call augment#Accept("\n")<cr>
 ```
 
+You can also accept only the next word of a suggestion using `augment#AcceptWord()`.
+This function also takes an optional fallback argument:
+
+```vim
+" Accept the next word of a suggestion with Ctrl-Right
+inoremap <c-right> <cmd>call augment#AcceptWord()<cr>
+```
+
+Or, with a fallback:
+
+```vim
+" Accept the next word, falling back to moving the cursor right
+inoremap <c-right> <cmd>call augment#AcceptWord("\<c-right>")<cr>
+```
+
+To dismiss a suggestion manually, use `augment#Dismiss()`:
+
+```vim
+" Dismiss the current suggestion with Escape
+inoremap <esc> <cmd>call augment#Dismiss()<cr><esc>
+```
+
 The default tab mapping can be disabled by setting
 `g:augment_disable_tab_mapping = v:true` before the plugin is loaded.
 

autoload/augment.vim

@@ -296,12 +296,20 @@ function! augment#OnTextChanged() abort
 endfunction
 
 function! augment#OnTextChangedI() abort
-    " Since CursorMovedI is always called before TextChangedI, the suggestion will already be cleared
     call s:UpdateBuffer()
+    " During AcceptWord, skip requesting new completions so they don't
+    " overwrite the partially-accepted suggestion
+    if get(b:, '_augment_suggestion_skip_clear', v:false)
+        return
+    endif
     call s:RequestCompletion()
 endfunction
 
 function! augment#OnCursorMovedI() abort
+    " If skip_clear flag is set, don't clear the suggestion
+    if get(b:, '_augment_suggestion_skip_clear', v:false)
+        return
+    endif
     call augment#suggestion#Clear()
 endfunction
 
@@ -311,9 +319,24 @@ function! augment#OnInsertEnter() abort
 endfunction
 
 function! augment#OnInsertLeavePre() abort
+    let b:_augment_suggestion_skip_clear = v:false
+    call augment#suggestion#Clear()
+endfunction
+
+" Dismiss the currently active suggestion
+function! augment#Dismiss() abort
+    let b:_augment_suggestion_skip_clear = v:false
     call augment#suggestion#Clear()
 endfunction
 
+" Accept the next word of the suggestion, with optional fallback
+function! augment#AcceptWord(...) abort
+    let fallback = a:0 >= 1 ? a:1 : ''
+    if !augment#suggestion#AcceptWord()
+        call feedkeys(fallback, 'nt')
+    endif
+endfunction
+
 " Accept the currently active suggestion if one is available, otherwise insert
 " the fallback text provided as the first argument
 function! augment#Accept(...) abort

autoload/augment/suggestion.vim

@@ -3,14 +3,19 @@
 
 " Functions for interacting with augment suggestions
 
-" Clear the suggestion
-function! augment#suggestion#Clear(...) abort
+" Remove ghost text visual elements without touching suggestion state
+function! augment#suggestion#ClearGhostText() abort
     if has('nvim')
         let ns_id = nvim_create_namespace('AugmentSuggestion')
         call nvim_buf_clear_namespace(0, ns_id, 0, -1)
     else
         call prop_remove({'type': 'AugmentSuggestion', 'all': v:true})
     endif
+endfunction
+
+" Clear the suggestion
+function! augment#suggestion#Clear(...) abort
+    call augment#suggestion#ClearGhostText()
 
     let current = exists('b:_augment_suggestion') ? b:_augment_suggestion : {}
     let b:_augment_suggestion = {}
@@ -28,26 +33,15 @@ function! augment#suggestion#Clear(...) abort
     return current
 endfunction
 
-" Show a suggestion
-function! augment#suggestion#Show(text, request_id, req_line, req_col, req_changedtick) abort
-    if len(a:text) == 0
+" Render ghost text for the given lines at the current cursor position
+function! augment#suggestion#Render(lines) abort
+    if empty(a:lines)
         return
     endif
 
-    call augment#suggestion#Clear()
-
-    " Save the suggestion information in a buffer-local variable
-    let b:_augment_suggestion = {
-                \ 'lines': split(a:text, "\n", 1),
-                \ 'request_id': a:request_id,
-                \ 'req_line': a:req_line,
-                \ 'req_col': a:req_col,
-                \ 'req_changedtick': a:req_changedtick,
-                \ }
-
     " Text properties don't render tabs, so manually add the correct spacing
     let tab_spaces = repeat(' ', &tabstop)
-    let lines = a:text->substitute("\t", tab_spaces, 'g')->split("\n", 1)
+    let lines = mapnew(a:lines, {_, v -> substitute(v, "\t", tab_spaces, 'g')})
 
     " Show the suggestion in ghost text
     if has('nvim')
@@ -80,6 +74,166 @@ function! augment#suggestion#Show(text, request_id, req_line, req_col, req_chang
     endif
 endfunction
 
+" Show a suggestion
+function! augment#suggestion#Show(text, request_id, req_line, req_col, req_changedtick) abort
+    if len(a:text) == 0
+        return
+    endif
+
+    call augment#suggestion#Clear()
+
+    " Save the suggestion information in a buffer-local variable
+    let b:_augment_suggestion = {
+                \ 'lines': split(a:text, "\n", 1),
+                \ 'request_id': a:request_id,
+                \ 'req_line': a:req_line,
+                \ 'req_col': a:req_col,
+                \ 'req_changedtick': a:req_changedtick,
+                \ }
+
+    call augment#suggestion#Render(b:_augment_suggestion.lines)
+endfunction
+
+" Compute remaining suggestion lines after accepting a word.
+function! s:ComputeRemainingLines(first_line, word, lines) abort
+    let remaining_first_line = strpart(a:first_line, len(a:word))
+    if !empty(remaining_first_line)
+        return [remaining_first_line] + a:lines[1:]
+    endif
+    if len(a:lines) > 1
+        return a:lines[1:]
+    endif
+    return []
+endfunction
+
+" Extract the next word from the suggestion lines.
+" Returns [word, first_line] on success, or empty list if nothing to extract.
+" word='' means consume a newline boundary.
+function! s:ExtractNextWord(lines) abort
+    if empty(a:lines)
+        return []
+    endif
+
+    " Newline boundary: first line empty but more lines follow
+    if empty(a:lines[0]) && len(a:lines) > 1
+        return ['', '']
+    endif
+
+    if empty(a:lines[0])
+        return []
+    endif
+
+    let first_line = a:lines[0]
+    let first_char = first_line[0]
+
+    if first_char =~ '\s'
+        let word = matchstr(first_line, '^\s\+')
+    elseif first_char =~ '\k'
+        let word = matchstr(first_line, '^\k\+')
+    else
+        let word = matchstr(first_line, '^\%(\k\@!\S\)\+')
+    endif
+
+    if empty(word)
+        return []
+    endif
+
+    return [word, first_line]
+endfunction
+
+" Accept the next word of the currently active suggestion, returning true
+" if there was a word to accept and false otherwise
+function! augment#suggestion#AcceptWord() abort
+    " Get current suggestion without clearing it
+    if !exists('b:_augment_suggestion') || empty(b:_augment_suggestion)
+        return v:false
+    endif
+    let info = b:_augment_suggestion
+
+    " Check buffer state is as expected
+    if line('.') != info.req_line || col('.') != info.req_col || b:changedtick != info.req_changedtick
+        let buf_state = '{line=' . line('.') . ', col=' . col('.') . ', changedtick=' . b:changedtick . '}'
+        let buf_expected = '{line=' . info.req_line . ', col=' . info.req_col . ', changedtick=' . info.req_changedtick . '}'
+        call augment#log#Warn(
+                    \ 'Attempted to accept word from completion "' . string(info.lines)
+                    \ . '" with buffer state ' . buf_state
+                    \ . ' and expected ' . buf_expected
+                    \ )
+        return v:false
+    endif
+
+    if empty(info.lines)
+        return v:false
+    endif
+
+    let extracted = s:ExtractNextWord(info.lines)
+    if empty(extracted)
+        return v:false
+    endif
+    let [word, first_line] = extracted
+
+    " Set the skip_clear flag to prevent autocommands from interfering
+    let b:_augment_suggestion_skip_clear = v:true
+    let l:bufnr = bufnr('%')
+
+    try
+        if empty(word)
+            " Empty word means we're consuming a newline boundary: split the
+            " current line at the cursor and move to the new line
+            let before = strpart(getline(line('.')), 0, col('.') - 1)
+            let after = strpart(getline(line('.')), col('.') - 1)
+            call setline(line('.'), before)
+            call append(line('.'), after)
+            call cursor(line('.') + 1, 1)
+            let remaining_lines = info.lines[1:]
+        else
+            " Insert the word into the buffer
+            let before = strpart(getline(line('.')), 0, col('.') - 1)
+            let after = strpart(getline(line('.')), col('.') - 1)
+            call setline(line('.'), before . word . after)
+
+            " Move cursor to the end of the inserted word
+            call cursor(line('.'), col('.') + len(word))
+
+            let remaining_lines = s:ComputeRemainingLines(first_line, word, info.lines)
+        endif
+
+        " Clear the old ghost text and update with remaining suggestion
+        call augment#suggestion#ClearGhostText()
+
+        " Update the suggestion state BEFORE unsetting skip flag
+        if !empty(remaining_lines)
+            let b:_augment_suggestion = {
+                        \ 'lines': remaining_lines,
+                        \ 'request_id': info.request_id,
+                        \ 'req_line': line('.'),
+                        \ 'req_col': col('.'),
+                        \ 'req_changedtick': b:changedtick,
+                        \ }
+            " Render the remaining ghost text
+            call augment#suggestion#Render(remaining_lines)
+
+            call augment#log#Debug('AcceptWord: Updated suggestion state - remaining=' . string(remaining_lines) . ' req_line=' . line('.') . ' req_col=' . col('.') . ' req_changedtick=' . b:changedtick)
+        else
+            " No remaining suggestion, clear state and send accept resolution
+            let b:_augment_suggestion = {}
+            call augment#client#Client().Notify('augment/resolveCompletion', {
+                        \ 'requestId': info.request_id,
+                        \ 'accept': v:true,
+                        \ })
+            call augment#log#Debug('Accepted completion (via AcceptWord) with request_id=' . info.request_id . ' text=' . string(info.lines))
+        endif
+    finally
+        " Unset the skip_clear flag after a short delay to allow autocommands to settle
+        if exists('b:_augment_suggestion_skip_clear_timer')
+            call timer_stop(b:_augment_suggestion_skip_clear_timer)
+        endif
+        let b:_augment_suggestion_skip_clear_timer = timer_start(10, {-> setbufvar(l:bufnr, '_augment_suggestion_skip_clear', v:false)})
+    endtry
+
+    return v:true
+endfunction
+
 " Accept the currently active suggestion if one is available, returning true
 " if there was a suggestion to accept and false otherwise
 function! augment#suggestion#Accept() abort

doc/augment.txt

@@ -8,6 +8,8 @@ Table of Contents                                  *augment-table-of-contents*
 2. Commands                                                 |augment-commands|
 3. Options                                                   |augment-options|
 4. Alternate Keybinds                             |augment-alternate-keybinds|
+   - augment#AcceptWord()                              |augment#AcceptWord()|
+   - augment#Dismiss()                                    |augment#Dismiss()|
 5. Highlighting                                         |augment-highlighting|
 
 ------------------------------------------------------------------------------
@@ -122,7 +124,7 @@ Alternate Keybinds                             *augment-alternate-keybinds*
 
 By default, tab is used to accept a suggestion. If you want to use a different
 key, create a mapping that calls `augment#Accept()`. The function takes an
-optional arugment used to specify the fallback text to insert if no suggestion
+optional argument used to specify the fallback text to insert if no suggestion
 is available.
 
 >vim
@@ -134,6 +136,30 @@ is available.
     inoremap <cr> <cmd>call augment#Accept("\n")<cr>
 <
 
+                                                            *augment#AcceptWord()*
+You can also accept only the next word of a suggestion using
+`augment#AcceptWord()`. This function also takes an optional fallback argument.
+
+>vim
+    " Accept the next word of a suggestion with Ctrl-Right
+    inoremap <c-right> <cmd>call augment#AcceptWord()<cr>
+<
+
+Or, with a fallback:
+
+>vim
+    " Accept the next word, falling back to moving the cursor right
+    inoremap <c-right> <cmd>call augment#AcceptWord("\<c-right>")<cr>
+<
+
+                                                            *augment#Dismiss()*
+To dismiss a suggestion manually, use `augment#Dismiss()`:
+
+>vim
+    " Dismiss the current suggestion with Escape
+    inoremap <esc> <cmd>call augment#Dismiss()<cr><esc>
+<
+
 ------------------------------------------------------------------------------
 Highlighting                                         *augment-highlighting*