feat: add async AI/AIEdit completion and :AIStop command

README.md

@@ -102,6 +102,7 @@ To use an AI command, type the command followed by an instruction prompt. You ca
 :AI          complete text
 :AIEdit      edit text
 :AIChat      continue or open new chat
+:AIStop      stop async generation for :AI and :AIEdit
 :AIStopChat  stop the generation of the AI response for the AIChat
 :AIImage     generate image
 
@@ -115,9 +116,9 @@ To use an AI command, type the command followed by an instruction prompt. You ca
 :help vim-ai
 ```
 
-**Tip:** Press `Ctrl-c` anytime to cancel `:AI` and `:AIEdit` completion
+**Tip:** Press `Ctrl-c` anytime to cancel synchronous `:AI` and `:AIEdit` completion. Use `:AIStop` for async runs.
 
-**Tip:** Use command shortcuts - `:AIE`, `:AIC`, `:AIS`,`:AIR`, `:AII` or setup your own [key bindings](#key-bindings)
+**Tip:** Use command shortcuts - `:AIE`, `:AIC`, `:AIR`, `:AII` or setup your own [key bindings](#key-bindings)
 
 **Tip:** Define and use [custom roles](#roles), e.g. `:AIEdit /grammar`.
 
@@ -543,6 +544,9 @@ let g:vim_ai_proxy = 'http://your-proxy-server:port'
 " enable/disable asynchronous AIChat (enabled by default)
 let g:vim_ai_async_chat = 1
 
+" enable/disable asynchronous AI and AIEdit (enabled by default)
+let g:vim_ai_async_complete = 1
+
 " enables/disables full markdown highlighting in aichat files
 " NOTE: code syntax highlighting works out of the box without this option enabled
 " NOTE: highlighting may be corrupted when using together with the `preservim/vim-markdown`
@@ -619,6 +623,9 @@ This plugin does not set any key binding. Create your own bindings in the `.vimr
 " stop async chat generation
 nnoremap <leader>s :AIStopChat<CR>
 
+" stop async AI/AIEdit generation
+nnoremap <leader>x :AIStop<CR>
+
 " complete text on the current line or in visual selection
 nnoremap <leader>a :AI<CR>
 xnoremap <leader>a :AI<CR>

autoload/vim_ai.vim

@@ -11,7 +11,6 @@ let s:last_command = ""
 let s:last_config = {}
 
 let s:scratch_buffer_name = ">>> AI chat"
-let s:chat_redraw_interval = 250 " milliseconds
 
 function! s:ImportPythonModules()
   for py_module in ['types', 'utils', 'context', 'chat', 'complete', 'roles', 'image']
@@ -21,6 +20,10 @@ function! s:ImportPythonModules()
   endfor
 endfunction
 
+function! vim_ai#ImportPythonModules() abort
+  call s:ImportPythonModules()
+endfunction
+
 function! s:StartsWith(longer, shorter) abort
   return a:longer[0:len(a:shorter)-1] ==# a:shorter
 endfunction
@@ -149,6 +152,13 @@ function! vim_ai#AIRun(uses_range, config, ...) range abort
   \}
   let l:context = py3eval("make_ai_context(unwrap('l:config_input'))")
   let l:config = l:context['config']
+  let l:bufnr = bufnr()
+  let l:is_async = g:vim_ai_async_complete == 1
+  if l:is_async && py3eval("ai_completion_job_pool.is_job_done(unwrap('l:bufnr'))") == 0
+    echoerr "Operation in progress, wait or stop it with :AIStop"
+    return
+  endif
+  let l:context['bufnr'] = l:bufnr
 
   let s:last_command = "complete"
   let s:last_config = a:config
@@ -158,17 +168,35 @@ function! vim_ai#AIRun(uses_range, config, ...) range abort
   let s:last_lastline = a:lastline
 
   let l:cursor_on_empty_line = empty(getline('.'))
+  let l:started = 0
   try
-    call s:set_paste(l:config)
+    if l:is_async
+      call vim_ai_async#EnablePasteMode(l:config)
+    else
+      call s:set_paste(l:config)
+    endif
     if l:cursor_on_empty_line
       execute "normal! " . a:lastline . "GA"
     else
       execute "normal! " . a:lastline . "Go"
     endif
-    py3 run_ai_completition(unwrap('l:context'))
-    execute "normal! " . a:lastline . "G"
+    if l:is_async
+      stopinsert
+    endif
+    let l:started = py3eval("run_ai_completition(unwrap('l:context'))")
+    if l:is_async && l:started
+      call timer_start(0, function('vim_ai_async#AICompletionWatch', [l:bufnr]))
+    elseif !l:is_async
+      execute "normal! " . a:lastline . "G"
+    endif
   finally
-    call s:set_nopaste(l:config)
+    if l:is_async
+      if !l:started
+        call vim_ai_async#DisablePasteMode()
+      endif
+    else
+      call s:set_nopaste(l:config)
+    endif
   endtry
 endfunction
 
@@ -192,6 +220,13 @@ function! vim_ai#AIEditRun(uses_range, config, ...) range abort
   \}
   let l:context = py3eval("make_ai_context(unwrap('l:config_input'))")
   let l:config = l:context['config']
+  let l:bufnr = bufnr()
+  let l:is_async = g:vim_ai_async_complete == 1
+  if l:is_async && py3eval("ai_completion_job_pool.is_job_done(unwrap('l:bufnr'))") == 0
+    echoerr "Operation in progress, wait or stop it with :AIStop"
+    return
+  endif
+  let l:context['bufnr'] = l:bufnr
 
   let s:last_command = "edit"
   let s:last_config = a:config
@@ -200,13 +235,30 @@ function! vim_ai#AIEditRun(uses_range, config, ...) range abort
   let s:last_firstline = a:firstline
   let s:last_lastline = a:lastline
 
+  let l:started = 0
   try
-    call s:set_paste(l:config)
+    if l:is_async
+      call vim_ai_async#EnablePasteMode(l:config)
+    else
+      call s:set_paste(l:config)
+    endif
     call s:SelectSelectionOrRange(l:is_selection, a:firstline, a:lastline)
     execute "normal! c"
-    py3 run_ai_completition(unwrap('l:context'))
+    if l:is_async
+      stopinsert
+    endif
+    let l:started = py3eval("run_ai_completition(unwrap('l:context'))")
+    if l:is_async && l:started
+      call timer_start(0, function('vim_ai_async#AICompletionWatch', [l:bufnr]))
+    endif
   finally
-    call s:set_nopaste(l:config)
+    if l:is_async
+      if !l:started
+        call vim_ai_async#DisablePasteMode()
+      endif
+    else
+      call s:set_nopaste(l:config)
+    endif
   endtry
 endfunction
 
@@ -281,35 +333,6 @@ function! s:ReuseOrCreateChatWindow(config)
   endif
 endfunction
 
-" Undo history is cluttered when using async chat.
-" There doesn't seem to be a way to use standard undojoin feature,
-" therefore working around with undoing and pasting changes manually.
-function! s:AIChatUndoCleanup()
-  let l:bufnr = bufnr()
-  let l:done = py3eval("ai_job_pool.is_job_done(unwrap('l:bufnr'))")
-  let l:chat_initiation_line = getbufvar(l:bufnr, 'vim_ai_chat_start_last_line', -1)
-  let l:undo_cleaned = l:chat_initiation_line == -1
-  if !l:done || l:undo_cleaned
-    return
-  endif
-
-  let l:current_line_num = line('.')
-  " navigate to the line where it started generating answer
-  execute l:chat_initiation_line
-  execute 'normal! j'
-  " copy whole assistant message to the `d` register
-  execute 'normal! "dyG'
-  " undo until user message
-  while line('$') > l:chat_initiation_line
-    execute "normal! u"
-  endwhile
-  " paste assistat message as a whole
-  execute 'normal! "dp'
-  execute l:current_line_num
-
-  call setbufvar(l:bufnr, 'vim_ai_chat_start_last_line', -1)
-endfunction
-
 " Start and answer the chat
 " - uses_range   - truty if range passed
 " - config       - function scoped vim_ai_chat config
@@ -357,67 +380,17 @@ function! vim_ai#AIChatRun(uses_range, config, ...) range abort
         " will clean undo history after returning back
         augroup AichatUndo
           au!
-          autocmd BufEnter <buffer> call s:AIChatUndoCleanup()
+          autocmd BufEnter <buffer> call vim_ai_async#AIChatUndoCleanup()
         augroup END
         execute "normal! Go\n<<< answering"
-        call timer_start(0, function('vim_ai#AIChatWatch', [l:bufnr, 0]))
+        call timer_start(0, function('vim_ai_async#AIChatWatch', [l:bufnr, 0]))
       endif
     endif
   finally
     call s:set_nopaste(l:config)
   endtry
 endfunction
 
-" Stop current chat job
-function! vim_ai#AIChatStopRun() abort
-  if &filetype !=# 'aichat'
-    echoerr "Not in an AI chat buffer."
-    return
-  endif
-  let l:bufnr = bufnr('%')
-  call s:ImportPythonModules() " Ensure chat.py is loaded
-  py3 ai_job_pool.cancel_job(unwrap('l:bufnr'))
-  call s:AIChatUndoCleanup()
-endfunction
-
-
-" Function called in a timer that check if there are new lines from AI and
-" appned them in a buffer. It ends when AI thread is finished (or when
-" stopped).
-function! vim_ai#AIChatWatch(bufnr, anim_index, timerid) abort
-  " inject new lines, first check if it is done to avoid data race, we do not
-  " mind if we run the timer one more time, but we want all the data
-  let l:done = py3eval("ai_job_pool.is_job_done(unwrap('a:bufnr'))")
-  let l:result = py3eval("ai_job_pool.pickup_lines(unwrap('a:bufnr'))")
-
-  " if user scroling over chat while answering, do not auto-scroll
-  let l:should_prevent_autoscroll = bufnr('%') == a:bufnr && line('.') != line('$')
-
-  call deletebufline(a:bufnr, '$')
-  call deletebufline(a:bufnr, '$')
-  call appendbufline(a:bufnr, '$', l:result)
-
-  " if not done, queue timer and animate
-  if l:done == 0
-    call timer_start(s:chat_redraw_interval, function('vim_ai#AIChatWatch', [a:bufnr, a:anim_index + 1]))
-    call appendbufline(a:bufnr, '$', "")
-    let l:animations = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
-    let l:current_animation = l:animations[a:anim_index % len(l:animations)]
-    call appendbufline(a:bufnr, '$', "<<< answering " . l:current_animation)
-  else
-    call s:AIChatUndoCleanup()
-    " Clear message
-    " https://neovim.discourse.group/t/how-to-clear-the-echo-message-in-the-command-line/268/3
-    call feedkeys(':','nx')
-  end
-
-  " if window is visible and user not scrolling, auto-scroll down
-  let winid = bufwinid(a:bufnr)
-  if winid != -1 && !l:should_prevent_autoscroll
-    call win_execute(winid, "normal! G")
-  endif
-endfunction
-
 " Start a new chat
 " a:1 - optional preset shorcut (below, right, tab)
 function! vim_ai#AINewChatDeprecatedRun(...)

autoload/vim_ai_async.vim

@@ -0,0 +1,119 @@
+function! vim_ai_async#EnablePasteMode(config) abort
+  if !&l:paste && a:config['ui']['paste_mode'] == '1'
+    setlocal paste
+    let b:vim_ai_async_restore_paste = 1
+  endif
+endfunction
+
+function! vim_ai_async#DisablePasteMode() abort
+  if get(b:, 'vim_ai_async_restore_paste', 0)
+    setlocal nopaste
+    unlet b:vim_ai_async_restore_paste
+  endif
+endfunction
+
+function! vim_ai_async#DisablePasteModeForBuffer(bufnr) abort
+  if getbufvar(a:bufnr, 'vim_ai_async_restore_paste', 0)
+    call setbufvar(a:bufnr, '&paste', 0)
+    call setbufvar(a:bufnr, 'vim_ai_async_restore_paste', 0)
+  endif
+endfunction
+
+" Undo history is cluttered when using async chat.
+" There doesn't seem to be a way to use standard undojoin feature,
+" therefore working around with undoing and pasting changes manually.
+function! vim_ai_async#AIChatUndoCleanup() abort
+  let l:bufnr = bufnr()
+  let l:done = py3eval("ai_job_pool.is_job_done(unwrap('l:bufnr'))")
+  let l:chat_initiation_line = getbufvar(l:bufnr, 'vim_ai_chat_start_last_line', -1)
+  let l:undo_cleaned = l:chat_initiation_line == -1
+  if !l:done || l:undo_cleaned
+    return
+  endif
+
+  let l:current_line_num = line('.')
+  " navigate to the line where it started generating answer
+  execute l:chat_initiation_line
+  execute 'normal! j'
+  " copy whole assistant message to the `d` register
+  execute 'normal! "dyG'
+  " undo until user message
+  while line('$') > l:chat_initiation_line
+    execute 'normal! u'
+  endwhile
+  " paste assistat message as a whole
+  execute 'normal! "dp'
+  execute l:current_line_num
+
+  call setbufvar(l:bufnr, 'vim_ai_chat_start_last_line', -1)
+endfunction
+
+" Stop current chat job
+function! vim_ai_async#AIChatStopRun() abort
+  if &filetype !=# 'aichat'
+    echoerr 'Not in an AI chat buffer.'
+    return
+  endif
+  let l:bufnr = bufnr('%')
+  call vim_ai#ImportPythonModules()
+  py3 ai_job_pool.cancel_job(unwrap('l:bufnr'))
+  call vim_ai_async#AIChatUndoCleanup()
+endfunction
+
+" Function called in a timer that check if there are new lines from AI and
+" appned them in a buffer. It ends when AI thread is finished (or when
+" stopped).
+function! vim_ai_async#AIChatWatch(bufnr, anim_index, timerid) abort
+  " inject new lines, first check if it is done to avoid data race, we do not
+  " mind if we run the timer one more time, but we want all the data
+  let l:done = py3eval("ai_job_pool.is_job_done(unwrap('a:bufnr'))")
+  let l:result = py3eval("ai_job_pool.pickup_lines(unwrap('a:bufnr'))")
+
+  " if user scroling over chat while answering, do not auto-scroll
+  let l:should_prevent_autoscroll = bufnr('%') == a:bufnr && line('.') != line('$')
+
+  call deletebufline(a:bufnr, '$')
+  call deletebufline(a:bufnr, '$')
+  call appendbufline(a:bufnr, '$', l:result)
+
+  " if not done, queue timer and animate
+  if l:done == 0
+    call timer_start(250, function('vim_ai_async#AIChatWatch', [a:bufnr, a:anim_index + 1]))
+    call appendbufline(a:bufnr, '$', '')
+    let l:animations = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
+    let l:current_animation = l:animations[a:anim_index % len(l:animations)]
+    call appendbufline(a:bufnr, '$', '<<< answering ' . l:current_animation)
+  else
+    call vim_ai_async#AIChatUndoCleanup()
+    " Clear message
+    " https://neovim.discourse.group/t/how-to-clear-the-echo-message-in-the-command-line/268/3
+    call feedkeys(':','nx')
+  end
+
+  " if window is visible and user not scrolling, auto-scroll down
+  let winid = bufwinid(a:bufnr)
+  if winid != -1 && !l:should_prevent_autoscroll
+    call win_execute(winid, 'normal! G')
+  endif
+endfunction
+
+" Stop current completion/edit job
+function! vim_ai_async#AIStopRun() abort
+  call vim_ai#ImportPythonModules()
+  let l:bufnr = bufnr('%')
+  if py3eval("ai_completion_job_pool.is_job_done(unwrap('l:bufnr'))")
+    echoerr 'No async :AI or :AIEdit task is running.'
+    return
+  endif
+  py3 ai_completion_job_pool.cancel_job(unwrap('l:bufnr'))
+endfunction
+
+" Function called in a timer to insert async completion chunks.
+function! vim_ai_async#AICompletionWatch(bufnr, timerid) abort
+  let l:done = py3eval("apply_ai_completion_job(unwrap('a:bufnr'))")
+  if l:done == 0
+    call timer_start(150, function('vim_ai_async#AICompletionWatch', [a:bufnr]))
+  else
+    call vim_ai_async#DisablePasteModeForBuffer(a:bufnr)
+  endif
+endfunction

autoload/vim_ai_config.vim

@@ -157,6 +157,9 @@ endif
 if !exists("g:vim_ai_async_chat")
   let g:vim_ai_async_chat = 1
 endif
+if !exists("g:vim_ai_async_complete")
+  let g:vim_ai_async_complete = 1
+endif
 if !exists("g:vim_ai_proxy")
   let g:vim_ai_proxy = ""
 endif