Add Neovim floating chat input (:Augment chat-input)

CHANGELOG.md

@@ -6,6 +6,9 @@ include any changes that may impact the user experience.
 
 ## Unreleased
 
+- Add the `:Augment chat-input` command, which opens a floating window for
+  composing a chat message before sending it (Neovim only). It is range-aware
+  like `:Augment chat`, and falls back to the standard `input()` prompt on Vim.
 - Add the `:Augment help [command]` command, which lists the available commands
   or shows more detailed help for a specific command.
 

README.md

@@ -79,6 +79,7 @@ The following commands are provided:
 :Augment signout       " Sign out of Augment
 :Augment log           " View the plugin log
 :Augment chat          " Send a chat message to Augment AI
+:Augment chat-input    " Compose a chat message in a floating window (Neovim only)
 :Augment chat-new      " Start a new chat conversation
 :Augment chat-toggle   " Toggle the chat panel visibility
 :Augment help          " List the available commands, or `:Augment help <command>` for details
@@ -146,6 +147,33 @@ You can interact with chat in two ways:
 
 The response will appear in a separate chat buffer with markdown formatting.
 
+### Floating chat input (Neovim only)
+
+The `:Augment chat-input` command opens a centered floating window with a
+markdown scratch buffer where you can compose a chat message before sending it.
+This is handy for writing longer, multi-line prompts. The window opens in insert
+mode, and its title shows the available keys:
+
+- `<C-s>` (insert or normal mode) or `<CR>` (normal mode) submits the message
+- `<Esc>` (normal mode) or `<C-c>` (insert or normal mode) cancels
+
+Like `:Augment chat`, it is range-aware: invoking it from visual mode (or with a
+range) includes the selected text in the chat request once you submit.
+
+If an input window is already open, running the command again refocuses it
+rather than opening a new one, so you won't lose what you've typed if focus
+moves away.
+
+This command requires Neovim's floating window support. In Vim it falls back to
+the standard `input()` prompt used by `:Augment chat`, with no change to
+existing behavior. The plugin does not define a default mapping for it, so map
+it yourself if you'd like a shortcut, for example:
+
+```vim
+nnoremap <leader>ai :Augment chat-input<CR>
+vnoremap <leader>ai :Augment chat-input<CR>
+```
+
 To start a new conversation, use the `:Augment chat-new` command. This will
 clear the chat history from your context.
 

autoload/augment.vim

@@ -186,8 +186,9 @@ function! s:CommandChat(range, args) abort
     " prompt the user for a message
     let message = empty(a:args) ? input('Message: ') : a:args
 
-    " Handle cancellation or empty input
-    if message ==# '' || message =~# '^\s*$'
+    " Handle cancellation or empty input. \_s matches whitespace including
+    " newlines, so a message that is only blank lines is treated as cancel.
+    if message ==# '' || message =~# '^\_s*$'
         redraw
         echo 'Chat cancelled'
         return
@@ -225,6 +226,65 @@ function! s:CommandChat(range, args) abort
     call augment#client#Client().Request('augment/chat', params)
 endfunction
 
+" Open a floating window to compose a chat message before sending it. The
+" floating input is Neovim-only; in Vim (and when a message is supplied
+" directly) this falls back to the standard chat command, which prompts for a
+" message via input() when none is given.
+function! s:CommandChatInput(range, args) abort
+    if !s:IsRunning()
+        echohl WarningMsg
+        echo s:NOT_RUNNING_MSG
+        echohl None
+        return
+    endif
+
+    " Determine whether a selection range is active. Leave visual mode so the
+    " '< and '> marks are set for the chat flow to pick up on submit.
+    let was_visual = index(['v', 'V', "\<C-v>"], mode()) >= 0
+    if was_visual
+        execute "normal! \<Esc>"
+    endif
+    let ranged = a:range == 2 || was_visual
+
+    " A message passed directly on the command line skips the floating input.
+    " Vim has no editable floating window, so it falls back to the input()
+    " prompt provided by the standard chat command.
+    if !empty(a:args) || !has('nvim')
+        call s:CommandChat(ranged ? 2 : 0, a:args)
+        return
+    endif
+
+    let source_win = win_getid()
+    let Callback = function('s:ChatInputSubmit', [source_win, ranged])
+    call augment#chat#OpenInputWindow(Callback)
+endfunction
+
+" Handle a message submitted from the floating chat input
+function! s:ChatInputSubmit(source_win, ranged, message) abort
+    " \_s matches whitespace including newlines, so a buffer of only blank
+    " lines is treated as cancel rather than sending an empty message.
+    if a:message ==# '' || a:message =~# '^\_s*$'
+        redraw
+        echo 'Chat cancelled'
+        return
+    endif
+
+    " Restore focus to the window the input was opened from
+    if win_id2win(a:source_win) != 0
+        call win_gotoid(a:source_win)
+    endif
+
+    " Re-select the original range so it is passed through to the chat request,
+    " mirroring the behavior of `:Augment chat` in visual mode. The '< and '>
+    " marks were set when the command left visual mode, so `gv` works whether
+    " invoked from visual mode or via an explicit `:'<,'>` range.
+    if a:ranged
+        normal! gv
+    endif
+
+    call s:CommandChat(a:ranged ? 2 : 0, a:message)
+endfunction
+
 function! s:CommandChatNew(range, args) abort
     call augment#chat#Reset()
 endfunction
@@ -257,6 +317,13 @@ let s:command_help = [
     \     'be included in the chat request. If no message is provided, you will',
     \     'be prompted to enter one.',
     \ ]},
+    \ {'name': 'chat-input', 'usage': 'chat-input', 'summary': 'Compose a chat message in a floating window (Neovim only).', 'detail': [
+    \     'Open a centered floating window with a markdown scratch buffer for',
+    \     'composing a chat message before sending it. Submit with <C-s> or, in',
+    \     'normal mode, <CR>; cancel with <Esc> or <C-c>. Like ":Augment chat" it',
+    \     'is range-aware. Requires Neovim; in Vim it falls back to the input()',
+    \     'prompt used by ":Augment chat".',
+    \ ]},
     \ {'name': 'chat-new', 'usage': 'chat-new', 'summary': 'Start a new chat conversation.', 'detail': [
     \     'Start a new chat conversation with Augment AI, clearing the history',
     \     'from your context.',
@@ -326,6 +393,7 @@ let s:command_handlers = {
     \ 'disable': function('s:CommandDisable'),
     \ 'status': function('s:CommandStatus'),
     \ 'chat': function('s:CommandChat'),
+    \ 'chat-input': function('s:CommandChatInput'),
     \ 'chat-new': function('s:CommandChatNew'),
     \ 'chat-toggle': function('s:CommandChatToggle'),
     \ 'help': function('s:CommandHelp'),

autoload/augment/chat.vim

@@ -79,6 +79,97 @@ function! augment#chat#OpenChatPanel() abort
     call win_gotoid(current_win)
 endfunction
 
+" Open a centered floating window with a scratch markdown buffer for composing
+" a chat message. a:OnSubmit is a Funcref invoked with the composed message
+" when the user submits. This relies on Neovim's floating window API and should
+" only be called when running under Neovim.
+function! augment#chat#OpenInputWindow(OnSubmit) abort
+    " If an input window is already open, refocus it instead of opening a new
+    " one. This avoids orphaning the existing float (and losing any typed
+    " content) when the command is invoked again after focus moved away.
+    if exists('s:input_win') && s:input_win != -1 && nvim_win_is_valid(s:input_win)
+        call nvim_set_current_win(s:input_win)
+        startinsert
+        return
+    endif
+
+    let s:input_on_submit = a:OnSubmit
+
+    " Create an unlisted scratch buffer (buftype=nofile, noswapfile)
+    let buf = nvim_create_buf(v:false, v:true)
+
+    " Center the window, sizing it relative to the editor dimensions
+    let width = float2nr(&columns * 0.6)
+    let width = max([40, min([width, &columns - 4])])
+    let height = max([5, min([10, &lines - 4])])
+    let row = (&lines - height) / 2
+    let col = (&columns - width) / 2
+
+    let opts = {
+                \ 'relative': 'editor',
+                \ 'width': width,
+                \ 'height': height,
+                \ 'row': row,
+                \ 'col': col,
+                \ 'style': 'minimal',
+                \ 'border': 'rounded',
+                \ 'title': ' Augment Chat (<C-s>/<CR> submit, <Esc> cancel) ',
+                \ 'title_pos': 'center',
+                \ }
+
+    let s:input_win = nvim_open_win(buf, v:true, opts)
+    let s:input_buf = buf
+
+    setlocal filetype=markdown   " Use markdown syntax highlighting
+    setlocal bufhidden=wipe      " Discard the buffer when the window closes
+    setlocal wrap                " Wrap long lines
+    setlocal linebreak           " Wrap at word boundaries
+
+    " Submit with <C-s> (insert and normal) or <CR> (normal)
+    inoremap <buffer> <silent> <C-s> <Esc><Cmd>call <SID>InputSubmit()<CR>
+    nnoremap <buffer> <silent> <C-s> <Cmd>call <SID>InputSubmit()<CR>
+    nnoremap <buffer> <silent> <CR> <Cmd>call <SID>InputSubmit()<CR>
+    " Cancel with <Esc> (normal) or <C-c> (insert and normal)
+    nnoremap <buffer> <silent> <Esc> <Cmd>call <SID>InputCancel()<CR>
+    inoremap <buffer> <silent> <C-c> <Esc><Cmd>call <SID>InputCancel()<CR>
+    nnoremap <buffer> <silent> <C-c> <Cmd>call <SID>InputCancel()<CR>
+
+    " Start in insert mode so the user can type immediately
+    startinsert
+endfunction
+
+function! s:CloseInputWindow() abort
+    if exists('s:input_win') && s:input_win != -1 && nvim_win_is_valid(s:input_win)
+        call nvim_win_close(s:input_win, v:true)
+    endif
+    let s:input_win = -1
+endfunction
+
+" Join the buffer contents into a message, close the window, and invoke the
+" stored submit callback with the message.
+function! s:InputSubmit() abort
+    if !exists('s:input_buf') || !nvim_buf_is_valid(s:input_buf)
+        call s:CloseInputWindow()
+        return
+    endif
+
+    let lines = nvim_buf_get_lines(s:input_buf, 0, -1, v:false)
+    let message = join(lines, "\n")
+    let Callback = s:input_on_submit
+
+    call s:CloseInputWindow()
+
+    if type(Callback) == v:t_func
+        call Callback(message)
+    endif
+endfunction
+
+function! s:InputCancel() abort
+    call s:CloseInputWindow()
+    redraw
+    echo 'Chat cancelled'
+endfunction
+
 function! augment#chat#Reset() abort
     call s:ResetChatContents()
     call s:ResetHistory()

doc/augment.txt

@@ -46,6 +46,21 @@ The following commands are provided:
     Start a chat with Augment AI. In visual mode, the selected text will be
     included in the chat request.
 
+                                                            *:Augment_chat_input*
+`:Augment chat-input`
+    Open a centered floating window with a markdown scratch buffer to compose a
+    chat message before sending it. The window opens in insert mode. Submit the
+    message with `<C-s>` (insert or normal mode) or `<CR>` (normal mode); cancel
+    with `<Esc>` (normal mode) or `<C-c>` (insert or normal mode). Like
+    `:Augment chat`, it is range-aware: invoking it from visual mode (or with a
+    range) includes the selected text in the chat request. If an input window is
+    already open, running the command again refocuses it rather than opening a
+    new one, preserving any text you have typed.
+
+    The floating window requires Neovim. In Vim this command falls back to the
+    standard `input()` prompt used by `:Augment chat`. No default mapping is
+    defined for this command.
+
                                                             *:Augment_chat_new*
 `:Augment chat-new`
     Start a new chat conversation with Augment AI.