From 3661c2bb1264bf95ba86998672168f39c9e79fec Mon Sep 17 00:00:00 2001 From: Konfekt Date: Tue, 24 Feb 2026 12:59:05 +0100 Subject: [PATCH 1/2] feat: add async AI/AIEdit completion and :AIStop command --- README.md | 11 +- autoload/vim_ai.vim | 151 ++++++++++++---------------- autoload/vim_ai_async.vim | 119 ++++++++++++++++++++++ autoload/vim_ai_config.vim | 3 + doc/tags | 1 + doc/vim-ai.txt | 10 ++ plugin/vim-ai.vim | 3 +- py/complete.py | 199 +++++++++++++++++++++++++++++++++++-- 8 files changed, 397 insertions(+), 100 deletions(-) create mode 100644 autoload/vim_ai_async.vim diff --git a/README.md b/README.md index 5a29311..7fd10b8 100644 --- a/README.md +++ b/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 s :AIStopChat +" stop async AI/AIEdit generation +nnoremap x :AIStop + " complete text on the current line or in visual selection nnoremap a :AI xnoremap a :AI diff --git a/autoload/vim_ai.vim b/autoload/vim_ai.vim index 1b0bf5f..5906e5a 100644 --- a/autoload/vim_ai.vim +++ b/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,10 +380,10 @@ function! vim_ai#AIChatRun(uses_range, config, ...) range abort " will clean undo history after returning back augroup AichatUndo au! - autocmd BufEnter call s:AIChatUndoCleanup() + autocmd BufEnter 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 @@ -368,56 +391,6 @@ function! vim_ai#AIChatRun(uses_range, config, ...) range abort 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(...) diff --git a/autoload/vim_ai_async.vim b/autoload/vim_ai_async.vim new file mode 100644 index 0000000..2097556 --- /dev/null +++ b/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 diff --git a/autoload/vim_ai_config.vim b/autoload/vim_ai_config.vim index 92f9d97..f676419 100644 --- a/autoload/vim_ai_config.vim +++ b/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 diff --git a/doc/tags b/doc/tags index 4e4aa6e..9bb3e12 100644 --- a/doc/tags +++ b/doc/tags @@ -3,6 +3,7 @@ :AIEdit vim-ai.txt /*:AIEdit* :AIImage vim-ai.txt /*:AIImage* :AIRedo vim-ai.txt /*:AIRedo* +:AIStop vim-ai.txt /*:AIStop* :AIStopChat vim-ai.txt /*:AIStopChat* :AIUtilDebugOff vim-ai.txt /*:AIUtilDebugOff* :AIUtilDebugOn vim-ai.txt /*:AIUtilDebugOn* diff --git a/doc/vim-ai.txt b/doc/vim-ai.txt index b5c82d5..cdba559 100644 --- a/doc/vim-ai.txt +++ b/doc/vim-ai.txt @@ -115,6 +115,12 @@ Options: > Check OpenAI docs for more information: https://platform.openai.com/docs/api-reference/completions + *:AIStop* +AIStop Cancel the currently running async AI/AIEdit + generation for the active buffer. + If no task is running or if it has already + completed, this command has no effect. + *:AIChat* :AIChat continue or start a new conversation. @@ -269,6 +275,10 @@ You can also customize the options in the chat header: > generate a paragraph of lorem ipsum ... +Async options: > + let g:vim_ai_async_chat = 1 + let g:vim_ai_async_complete = 1 + ROLES Roles are defined in the `.ini` file: > diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim index 0b0bee4..19d106f 100644 --- a/plugin/vim-ai.vim +++ b/plugin/vim-ai.vim @@ -17,7 +17,8 @@ command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletionChat AIChat < command! -range -nargs=? -complete=customlist,vim_ai#RoleCompletionImage AIImage ,call vim_ai#AIImageRun(, {}, ) command! -nargs=? AINewChat call vim_ai#AINewChatDeprecatedRun() command! AIRedo call vim_ai#AIRedoRun() -command! AIStopChat call vim_ai#AIChatStopRun() +command! AIStopChat call vim_ai_async#AIChatStopRun() +command! AIStop call vim_ai_async#AIStopRun() command! AIUtilRolesOpen call vim_ai#AIUtilRolesOpen() command! AIUtilDebugOn call vim_ai#AIUtilSetDebug(1) command! AIUtilDebugOff call vim_ai#AIUtilSetDebug(0) diff --git a/py/complete.py b/py/complete.py index 6270fdd..756a2b8 100644 --- a/py/complete.py +++ b/py/complete.py @@ -1,14 +1,178 @@ import vim +import copy +import threading +import traceback complete_py_imported = True +class AI_completion_job(threading.Thread): + def __init__(self, bufnr, messages, provider, provider_name, append_to_eol, insert_before_cursor, row, col): + threading.Thread.__init__(self) + self.bufnr = bufnr + self.messages = messages + self.provider = provider + self.provider_name = provider_name + self.append_to_eol = append_to_eol + self.insert_before_cursor = insert_before_cursor + self.chunks = [] + self.cancelled = False + self.done = False + self.error = None + self.render_state = { + 'started': False, + 'full_text': '', + 'insert_before_cursor': insert_before_cursor, + 'append_to_eol': append_to_eol, + 'row': row, + 'col': col, + } + self.lock = threading.RLock() + + def run(self): + print_debug("AI_completion_job thread STARTED") + try: + for chunk in self.provider.request(self.messages): + with self.lock: + if self.cancelled: + break + if chunk.get('type') != 'assistant': + continue + content = chunk.get('content') + if content: + self.chunks.append(content) + except Exception as e: + with self.lock: + self.error = e + finally: + with self.lock: + self.done = True + print_debug("AI_completion_job thread DONE") + + def pickup_chunks(self): + with self.lock: + chunks = copy.deepcopy(self.chunks) + self.chunks = [] + return chunks + + def is_done(self): + with self.lock: + done = self.done + return done + + def cancel(self): + with self.lock: + self.cancelled = True + +class AI_completion_jobs_pool(object): + def __init__(self): + self.pool = {} + + def new_job(self, bufnr, messages, provider, provider_name, append_to_eol, insert_before_cursor, row, col): + bufnr = int(bufnr) + self.pool[bufnr] = AI_completion_job( + bufnr, + messages, + provider, + provider_name, + append_to_eol, + insert_before_cursor, + row, + col, + ) + self.pool[bufnr].start() + return self.pool[bufnr] + + def get_job(self, bufnr): + return self.pool.get(int(bufnr)) + + def pickup_chunks(self, bufnr): + job = self.pool.get(int(bufnr)) + return job.pickup_chunks() if job else [] + + def is_job_done(self, bufnr): + job = self.pool.get(int(bufnr)) + return job.is_done() if job else True + + def cancel_job(self, bufnr): + job = self.pool.get(int(bufnr)) + if not job: + return False + if not job.is_done(): + job.cancel() + return True + return False + +ai_completion_job_pool = AI_completion_jobs_pool() + +def _insert_text_into_buffer(buffer, row, col, text): + lines = text.split("\n") + current_line = buffer[row] + before = current_line[:col] + after = current_line[col:] + if len(lines) == 1: + buffer[row] = before + lines[0] + after + return row, col + len(lines[0]) + + new_lines = [before + lines[0]] + if len(lines) > 2: + new_lines.extend(lines[1:-1]) + new_lines.append(lines[-1] + after) + buffer[row:row + 1] = new_lines + return row + len(new_lines) - 1, len(lines[-1]) + +def _render_text_chunks_incremental(bufnr, chunks, state): + if not chunks: + return + buffer = vim.buffers[int(bufnr)] + for text in chunks: + if not state['started']: + text = text.lstrip() + if not text: + continue + state['started'] = True + row = state['row'] + col = state['col'] + if state['append_to_eol']: + col = len(buffer[row]) + elif state['insert_before_cursor']: + col = max(col - 1, 0) + state['insert_before_cursor'] = False + row, col = _insert_text_into_buffer(buffer, row, col, text) + state['row'] = row + state['col'] = col + state['full_text'] += text + +def apply_ai_completion_job(bufnr): + bufnr = int(bufnr) + job = ai_completion_job_pool.get_job(bufnr) + if not job: + return True + + chunks = job.pickup_chunks() + if chunks: + _render_text_chunks_incremental(bufnr, chunks, job.render_state) + + done = job.is_done() + if done: + if job.error: + handle_completion_error(job.provider_name, job.error) + elif job.cancelled: + print_info_message("Completion cancelled...") + elif not job.render_state['full_text'].strip(): + handle_completion_error( + job.provider_name, + KnownError('Empty response received. Tip: You can try modifying the prompt and retry.'), + ) + clear_echo_message() + return done + + def run_ai_completition(context): update_thread_shared_variables() command_type = context['command_type'] prompt = context['prompt'] config = make_config(context['config']) config_options = config['options'] - config_ui = config['ui'] roles = context['roles'] try: @@ -27,16 +191,35 @@ def run_ai_completition(context): provider_class = load_provider(config['provider']) provider = provider_class(command_type, config_options, ai_provider_utils) - response_chunks = provider.request(messages) - text_chunks = map( - lambda c: c.get("content"), - filter(lambda c: c['type'] == 'assistant', response_chunks), # omit `thinking` section - ) + if vim.eval('g:vim_ai_async_complete') == '1': + cursor_pos = vim.eval("getpos('.')") + row = int(cursor_pos[1]) - 1 + col = int(cursor_pos[2]) + ai_completion_job_pool.new_job( + int(context['bufnr']), + messages, + provider, + config['provider'], + command_type == 'complete', + need_insert_before_cursor(), + row, + col, + ) + else: + response_chunks = provider.request(messages) + + text_chunks = map( + lambda c: c.get('content'), + filter(lambda c: c['type'] == 'assistant', response_chunks), + ) - render_text_chunks(text_chunks, append_to_eol=command_type == 'complete') + render_text_chunks(text_chunks, append_to_eol=command_type == 'complete') - clear_echo_message() + clear_echo_message() + return True + return False except BaseException as error: handle_completion_error(config['provider'], error) print_debug("[{}] error: {}", command_type, traceback.format_exc()) + return False From 385618b622ff8e45f730050b60dbbdd66d2962ec Mon Sep 17 00:00:00 2001 From: Konfekt Date: Fri, 27 Feb 2026 11:48:56 +0100 Subject: [PATCH 2/2] fix: redraw after rendering AI completion chunks --- py/complete.py | 1 + tests/complete_test.py | 43 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+) create mode 100644 tests/complete_test.py diff --git a/py/complete.py b/py/complete.py index 756a2b8..037c8ea 100644 --- a/py/complete.py +++ b/py/complete.py @@ -151,6 +151,7 @@ def apply_ai_completion_job(bufnr): chunks = job.pickup_chunks() if chunks: _render_text_chunks_incremental(bufnr, chunks, job.render_state) + vim.command("redraw") done = job.is_done() if done: diff --git a/tests/complete_test.py b/tests/complete_test.py new file mode 100644 index 0000000..15cfd4c --- /dev/null +++ b/tests/complete_test.py @@ -0,0 +1,43 @@ +from types import SimpleNamespace + +import complete + + +class FakeVim: + def __init__(self, buffers): + self.buffers = buffers + self.commands = [] + + def command(self, cmd): + self.commands.append(cmd) + + +def make_job(chunks, done=False): + return SimpleNamespace( + render_state={ + 'started': False, + 'full_text': '', + 'insert_before_cursor': False, + 'append_to_eol': True, + 'row': 0, + 'col': 0, + }, + error=None, + cancelled=False, + provider_name='openai', + pickup_chunks=lambda: chunks, + is_done=lambda: done, + ) + + +def test_apply_ai_completion_job_redraws_after_chunk_render(monkeypatch): + fake_vim = FakeVim({1: ['hello ']}) + job = make_job(['world'], done=False) + monkeypatch.setattr(complete, 'vim', fake_vim) + monkeypatch.setattr(complete.ai_completion_job_pool, 'get_job', lambda _: job) + + done = complete.apply_ai_completion_job(1) + + assert done is False + assert fake_vim.buffers[1] == ['hello world'] + assert fake_vim.commands == ['redraw']