Add status column icons for inline prompt

[lucobellic]

With context added to the InlineStarted event data payload in #1293, it is now easy to add status columns icons or any type of extmarks when starting an inline prompt

You can also draw inspiration from #640 to add virtual text, such as a loading indicator.

Here is a code snippet to produce the result shown in the screenshot by adding status columns icons during the InlineStarted event and clearing them during the InlineFinished event:

--- @class CodeCompanion.InlineExtmark
--- @field unique_line_sign_text string Text used for sign when there's only a single line
--- @field first_line_sign_text string Text used for sign on the first line of multi-line section
--- @field last_line_sign_text string Text used for sign on the last line of multi-line section
--- @field extmark vim.api.keyset.set_extmark Extmark options passed to nvim_buf_set_extmark

local M = {}

--- @type CodeCompanion.InlineExtmark
local default_opts = {
  unique_line_sign_text = '',
  first_line_sign_text = '┌',
  last_line_sign_text = '└',
  extmark = {
    sign_hl_group = 'DiagnosticWarn',
    sign_text = '│',
    priority = 1000,
  },
}

--- Helper function to set a line extmark with specified sign text
--- @param bufnr number
--- @param ns_id number
--- @param line_num number Line number
--- @param opts vim.api.keyset.set_extmark Extmark options
--- @param sign_type string Key in opts for the sign text to use
local function set_line_extmark(bufnr, ns_id, line_num, opts, sign_type)
  vim.api.nvim_buf_set_extmark(
    bufnr,
    ns_id,
    line_num - 1, -- Convert to 0-based index
    0,
    vim.tbl_deep_extend('force', opts.extmark or {}, {
      sign_text = opts[sign_type] or opts.extmark.sign_text,
    })
  )
end

--- Creates extmarks for inline code annotations
--- @param opts CodeCompanion.InlineExtmark Configuration options for the extmarks
--- @param data CodeCompanion.InlineArgs Data containing context information about the code block
--- @param ns_id number unique namespace id for the extmarks
local function create_extmarks(opts, data, ns_id)
  --- @type {bufnr: number, start_line: number, end_line: number}
  local context = data.context

  -- Handle the case where start and end lines are the same (unique line)
  if context.start_line == context.end_line then
    set_line_extmark(context.bufnr, ns_id, context.start_line, opts, 'unique_line_sign_text')
    return
  end

  -- Set extmark for the first line with special options
  set_line_extmark(context.bufnr, ns_id, context.start_line, opts, 'first_line_sign_text')

  -- Set extmarks for the middle lines with standard options
  for i = context.start_line + 1, context.end_line - 1 do
    vim.api.nvim_buf_set_extmark(context.bufnr, ns_id, i - 1, 0, opts.extmark)
  end

  -- Set extmark for the last line with special options
  if context.end_line > context.start_line then
    set_line_extmark(context.bufnr, ns_id, context.end_line, opts, 'last_line_sign_text')
  end
end

--- Creates autocmds for CodeCompanionRequest events
--- @param opts CodeCompanion.InlineExtmark Configuration options passed from setup
local function create_autocmds(opts)
  vim.api.nvim_create_autocmd({ 'User' }, {
    pattern = { 'CodeCompanionRequest*' },
    callback =
      --- @param args {buf: number, data : CodeCompanion.InlineArgs, match: string}
      function(args)
        local data = args.data or {}
        local context = data and data.context or {}
        if data and data.context then
          local ns_id = vim.api.nvim_create_namespace('CodeCompanionInline_' .. data.id)
          if args.match:find('StartedInline') then
            create_extmarks(opts, data, ns_id)
          elseif args.match:find('FinishedInline') then
            vim.api.nvim_buf_clear_namespace(context.bufnr, ns_id, 0, -1)
          end
        end
      end,
  })
end

--- @param opts? CodeCompanion.InlineExtmark Optional configuration to override defaults
function M.setup(opts) create_autocmds(vim.tbl_deep_extend('force', default_opts, opts or {})) end

return M

Usage:

return {
  'olimorris/codecompanion.nvim',
  -- ...
  config = function(_, opts)
      require('codecompanion').setup(opts)
      require('plugins.codecompanion.utils.extmarks').setup()
  end
}

[lucobellic]

Here is another fancy way to use extmarks with inline assistant (not necessarily practical) just to get some ideas with code:

codecompanion-snippet.mp4

extmarks.lua

--- @class CodeCompanion.InlineExtmark
--- @field unique_line_sign_text string Text used for sign when there's only a single line
--- @field first_line_sign_text string Text used for sign on the first line of multi-line section
--- @field last_line_sign_text string Text used for sign on the last line of multi-line section
--- @field extmark vim.api.keyset.set_extmark Extmark options passed to nvim_buf_set_extmark

local M = {}

local hl_group = 'DiagnosticVirtualTextWarn'
local priority = 2048
local repeat_interval = 100

--- @type CodeCompanion.InlineExtmark
local default_opts = {
  unique_line_sign_text = '',
  first_line_sign_text = '┌',
  last_line_sign_text = '└',
  extmark = {
    sign_hl_group = hl_group,
    sign_text = '│',
    priority = priority,
  },
}

--- @type VirtualTextBlockSpinner.SpinnerOpts
--- @diagnostic disable-next-line: missing-fields
local virtual_text_spinners_opts = {
  hl_group = 'Comment',
  repeat_interval = repeat_interval,
  extmark = {
    virt_text_pos = 'overlay',
    priority = priority,
  },
}

--- @type {number: [VirtualTextSpinner, VirtualTextBlockSpinner]}
local virtual_text_spinners = {}

--- Helper function to set a line extmark with specified sign text
--- @param bufnr number
--- @param ns_id number
--- @param line_num number Line number
--- @param opts CodeCompanion.InlineExtmark Extmark options
--- @param sign_type string Key in opts for the sign text to use
local function set_line_extmark(bufnr, ns_id, line_num, opts, sign_type)
  vim.api.nvim_buf_set_extmark(
    bufnr,
    ns_id,
    line_num - 1, -- Convert to 0-based index
    0,
    vim.tbl_deep_extend('force', opts.extmark or {}, {
      sign_text = opts[sign_type] or opts.extmark.sign_text,
    })
  )
end

--- Start animated spinners for a buffer region
--- This function creates and starts two spinner objects:
---  - A block spinner that spans multiple lines
---  - A centered spinner that appears in the middle of the block
---
--- @param bufnr number Buffer number where the spinners will be displayed
--- @param ns_id number Namespace ID for the extmarks
--- @param start_line number Starting line of the region (0-indexed)
--- @param end_line number Ending line of the region (0-indexed)
local function start_spinners(bufnr, ns_id, start_line, end_line)
  local block_spinner = require('plugins.codecompanion.utils.block_spinner').new({
    bufnr = bufnr,
    ns_id = ns_id,
    start_line = start_line,
    end_line = end_line,
    opts = virtual_text_spinners_opts,
  })

  local spinner = require('plugins.codecompanion.utils.spinner').new({
    bufnr = bufnr,
    ns_id = ns_id,
    line_num = start_line + math.floor((end_line - start_line) / 2),
    width = block_spinner.width,
    opts = {
      repeat_interval = repeat_interval,
      extmark = { virt_text_pos = 'overlay', priority = priority + 1 },
    },
  })

  spinner:start()
  block_spinner:start()
  virtual_text_spinners[ns_id] = { spinner, block_spinner }
end

--- Stop active spinners associated with a namespace ID
--- @param ns_id number The namespace ID to stop spinners
local function stop_spinner(ns_id)
  local block_spinner, spinner = unpack(virtual_text_spinners[ns_id])
  if spinner then
    spinner:stop()
  end
  if block_spinner then
    block_spinner:stop()
  end
  virtual_text_spinners[ns_id] = nil
end

--- Creates extmarks for inline code annotations
--- @param opts CodeCompanion.InlineExtmark Configuration options for the extmarks
--- @param data CodeCompanion.InlineArgs Data containing context information about the code block
--- @param ns_id number unique namespace id for the extmarks
local function create_extmarks(opts, data, ns_id)
  --- @type {bufnr: number, start_line: number, end_line: number}
  local context = data.context

  -- Start a spinner on first line at the end of line
  start_spinners(context.bufnr, ns_id, context.start_line, context.end_line)

  -- Handle the case where start and end lines are the same (unique line)
  if context.start_line == context.end_line then
    set_line_extmark(context.bufnr, ns_id, context.start_line, opts, 'unique_line_sign_text')
    return
  end

  -- Set extmark for the first line with special options
  set_line_extmark(context.bufnr, ns_id, context.start_line, opts, 'first_line_sign_text')

  -- Set extmarks for the middle lines with standard options
  for i = context.start_line + 1, context.end_line - 1 do
    vim.api.nvim_buf_set_extmark(context.bufnr, ns_id, i - 1, 0, opts.extmark)
  end

  -- Set extmark for the last line with special options
  if context.end_line > context.start_line then
    set_line_extmark(context.bufnr, ns_id, context.end_line, opts, 'last_line_sign_text')
  end
end

--- Creates autocmds for CodeCompanionRequest events
--- @param opts CodeCompanion.InlineExtmark Configuration options passed from setup
local function create_autocmds(opts)
  vim.api.nvim_create_autocmd('User', {
    pattern = 'CodeCompanionRequest*',
    callback = function(args)
      local data = args.data or {}
      local context = data.context or {}

      if vim.tbl_isempty(context) then
        return
      end

      local ns_id = vim.api.nvim_create_namespace('CodeCompanionInline_' .. data.id)

      if args.match:find('StartedInline') then
        create_extmarks(opts, data, ns_id)
      elseif args.match:find('FinishedInline') then
        stop_spinner(ns_id)
        vim.api.nvim_buf_clear_namespace(context.bufnr, ns_id, 0, -1)
      end
    end,
  })
end

--- @param opts? CodeCompanion.InlineExtmark Optional configuration to override defaults
function M.setup(opts) create_autocmds(vim.tbl_deep_extend('force', default_opts, opts or {})) end

return M

block_spinner.lua

--- @class VirtualTextBlockSpinner.SpinnerOpts
--- @field hl_group string Highlight group for the spinner
--- @field repeat_interval number Interval in milliseconds to update the spinner
--- @field extmark vim.api.keyset.set_extmark Extmark options passed to nvim_buf_set_extmark
--- @field patterns? table<string> Table of spinner patterns to cycle through
local spinner_opts = {
  hl_group = 'Comment',
  repeat_interval = 100,
  extmark = {
    virt_text_pos = 'inline',
    priority = 1000,
    virt_text_repeat_linebreak = true,
  },
}

--- @class VirtualTextBlockSpinner
--- @field bufnr number Buffer number where the text block spinner is shown
--- @field ns_id number The namespace ID for the extmark
--- @field start_line number Starting line (0-indexed) for the spinner
--- @field end_line number Ending line (0-indexed) for the spinner
--- @field ids table<number, number> Table of extmark IDs indexed by line numbers
--- @field patterns table<string> Table of spinner patterns to cycle through
--- @field current_index number Current index in the spinner patterns array
--- @field timer uv.uv_timer_t | nil Timer used to update spinner animation
--- @field opts VirtualTextBlockSpinner.SpinnerOpts Configuration options for the spinner
--- @field width number The width of the spinner content
--- @field height number The height of the spinner content
--- @field border_top string Top border of the spinner
--- @field border_bottom string Bottom border of the spinner
local VirtualTextBlockSpinner = {
  bufnr = 0,
  ns_id = 0,
  start_line = 0,
  end_line = 0,
  ids = {},
  patterns = {},
  current_index = 1,
  timer = nil,
  opts = spinner_opts,
  width = 0,
  height = 0,
  border_top = "",
  border_bottom = "",
}

--- @class VirtualTextBlockSpinner.Opts
--- @field bufnr number Buffer number where the spinner will be shown
--- @field ns_id number Namespace ID for the extmarks
--- @field start_line number Starting line (1-indexed) for the spinner
--- @field end_line number Ending line (1-indexed) for the spinner
--- @field opts? VirtualTextBlockSpinner.SpinnerOpts Optional configuration options

--- Creates a new VirtualTextBlockSpinner instance
--- @param opts VirtualTextBlockSpinner.Opts Options for the spinner
--- @return VirtualTextBlockSpinner self New spinner instance
function VirtualTextBlockSpinner.new(opts)
  local lines = vim.api.nvim_buf_get_lines(opts.bufnr, opts.start_line - 1, opts.end_line, false)
  local width = vim.fn.max(vim.iter(lines):map(function(line) return vim.fn.strdisplaywidth(line) end):totable())

  local merged_opts = vim.tbl_deep_extend('force', spinner_opts, opts.opts or {})
  local patterns = {}

  -- First determine the raw patterns
  local raw_patterns = {
    '╲  ',
    ' ╲ ',
    '  ╲',
  }

  if merged_opts.patterns and #merged_opts.patterns > 0 then
    raw_patterns = merged_opts.patterns --- @type string[]
  end

  -- Calculate required repetitions to match the content width
  --- @diagnostic disable-next-line: need-check-nil
  local pattern_width = vim.fn.strdisplaywidth(raw_patterns[1])
  local repetitions = pattern_width > 0 and math.ceil(width / pattern_width) or width
  local width = repetitions * pattern_width
  local horizontal_line = string.rep('─', width)

  -- Now create the final patterns with proper repetition
  for _, pattern in ipairs(raw_patterns) do
    table.insert(patterns, string.rep(pattern, repetitions))
  end

  return setmetatable({
    bufnr = opts.bufnr,
    ns_id = opts.ns_id,
    start_line = opts.start_line - 1,
    end_line = opts.end_line - 1,
    ids = {},
    patterns = patterns,
    current_index = 1,
    timer = vim.uv.new_timer(),
    opts = merged_opts,
    width = width,
    height = #lines,
    border_top = '╭' .. horizontal_line .. '╮',
    border_bottom = '╰' .. horizontal_line .. '╯',
  }, { __index = VirtualTextBlockSpinner })
end

--- Gets the virtual text content for the spinner based on line and current animation frame
--- @param i number The line number to get virtual text for
--- @return table[] Virtual text for the extmark in the format required by nvim_buf_set_extmark
function VirtualTextBlockSpinner:get_virtual_text(i)
  local pattern_index = ((i + self.current_index - 1) % #self.patterns) + 1
  local pattern = self.patterns[pattern_index]

  if self.height <= 2 then
    return { { pattern, self.opts.hl_group } }
  end

  -- First line (top border)
  if i == self.start_line then
    return { { self.border_top, self.opts.hl_group } }
  end

  -- Last line (bottom border)
  if i == self.end_line then
    return { { self.border_bottom, self.opts.hl_group } }
  end

  -- Middle lines with spinner animation
  return { { '│' .. pattern .. '│', self.opts.hl_group } }
end

--- Sets up new extmarks for all lines in the spinner range
--- Creates extmarks for each line and stores their IDs
--- @private
function VirtualTextBlockSpinner:set_new_extmarks()
  self.ids = {}
  for i = self.start_line, self.end_line do
    self.ids[i] = vim.api.nvim_buf_set_extmark(
      self.bufnr,
      self.ns_id,
      i,
      0,
      vim.tbl_deep_extend('force', self.opts.extmark, { virt_text = self:get_virtual_text(i) })
    )
  end
end

--- Updates existing extmarks with new virtual text content based on the current animation frame
--- This method is called by the timer to update the spinner animation
--- @private
function VirtualTextBlockSpinner:set_extmarks()
  for i, id in pairs(self.ids) do
    local current_pos = vim.api.nvim_buf_get_extmark_by_id(self.bufnr, self.ns_id, id, {})
    pcall(
      function()
        vim.api.nvim_buf_set_extmark(
          self.bufnr,
          self.ns_id,
          current_pos[1],
          0,
          vim.tbl_deep_extend('force', self.opts.extmark, { virt_text = self:get_virtual_text(i), id = id })
        )
      end
    )
  end
end

--- Starts the spinner animation
--- Creates extmarks for each line in the range and starts the timer to update spinner animation
--- The spinner will continue to animate until the stop method is called
function VirtualTextBlockSpinner:start()
  self:set_new_extmarks()
  self.timer:start(
    0,
    self.opts.repeat_interval,
    vim.schedule_wrap(function()
      self.current_index = (self.current_index % #self.patterns) + 1
      self:set_extmarks()
    end)
  )
end

--- Stops the spinner animation
--- Cleans up the timer resources and removes all extmarks created for the spinner
--- This method should always be called when the spinner is no longer needed to prevent memory leaks
function VirtualTextBlockSpinner:stop()
  if self.timer then
    self.timer:stop()
    self.timer:close()
    self.timer = nil
  end
  if self.opts.extmark then
    for _, id in pairs(self.ids) do
      vim.schedule(function() vim.api.nvim_buf_del_extmark(self.bufnr, self.ns_id, id) end)
    end
  end
end

return VirtualTextBlockSpinner

spinner.lua

--- @class VirtualTextSpinner.SpinnerOpts
--- @field spinner_text string Text to display before the spinner
--- @field spinner_frames string[] Spinner frames to use for the spinner
--- @field hl_group string Highlight group for the spinner
--- @field repeat_interval number Interval in milliseconds to update the spinner
--- @field extmark vim.api.keyset.set_extmark Extmark options passed to nvim_buf_set_extmark
local spinner_opts = {
  spinner_text = '  Processing',
  spinner_frames = { '⣷', '⣯', '⣟', '⡿', '⢿', '⣻', '⣽', '⣾' },
  hl_group = 'DiagnosticVirtualTextWarn',
  repeat_interval = 100,
  extmark = {
    virt_text_pos = 'inline',
    priority = 1000,
  },
}

--- @class VirtualTextSpinner
--- @field bufnr number The buffer number where the spinner is displayed
--- @field ns_id number The namespace ID for the extmark
--- @field line_num number The line number where the spinner is displayed
--- @field current_index number Current index in the spinner frames array
--- @field timer uv.uv_timer_t | nil Timer used to update spinner animation
--- @field opts VirtualTextSpinner.SpinnerOpts Configuration options for the spinner
local VirtualTextSpinner = {
  bufnr = 0,
  ns_id = 0,
  line_num = 0,
  current_index = 1,
  timer = nil,
  opts = spinner_opts,
}

--- @class VirtualTextSpinner.Opts
--- @field bufnr number Buffer number to display the spinner in
--- @field ns_id number Namespace ID for the extmark
--- @field line_num number Line number to display the spinner on (1-indexed)
--- @field width? number Width of the spinner
--- @field opts? VirtualTextSpinner.SpinnerOpts Optional configuration options

--- Creates a new VirtualTextSpinner instance
--- @param opts VirtualTextSpinner.Opts Options for the spinner
--- @return VirtualTextSpinner self New spinner instance
function VirtualTextSpinner.new(opts)
  local width = opts.width or 0
  local spinner_opts = vim.tbl_deep_extend('force', spinner_opts, opts.opts or {})
  local width_center = width - spinner_opts.spinner_text:len()
  local col = width_center > 0 and math.floor(width_center / 2) or 0
  return setmetatable({
    bufnr = opts.bufnr,
    ns_id = opts.ns_id,
    line_num = opts.line_num - 1,
    current_index = 1,
    timer = vim.uv.new_timer(),
    opts = vim.tbl_deep_extend('force', spinner_opts, { extmark = { virt_text_win_col = col } }),
  }, { __index = VirtualTextSpinner })
end

--- Gets the virtual text content for the spinner
--- @return table[]
function VirtualTextSpinner:get_virtual_text()
  return { { self.opts.spinner_text .. ' ' .. self.opts.spinner_frames[self.current_index] .. ' ', self.opts.hl_group } }
end

--- @private
--- @return number id of the extmark
function VirtualTextSpinner:set_extmark()
  return vim.api.nvim_buf_set_extmark(self.bufnr, self.ns_id, self.line_num, 0, self.opts.extmark)
end

--- Starts the spinner animation
--- Creates the extmark and starts the timer to update the spinner frames
function VirtualTextSpinner:start()
  self.opts.extmark.virt_text = self:get_virtual_text()
  self.opts.extmark.id = self:set_extmark()
  self.timer:start(
    0,
    self.opts.repeat_interval,
    vim.schedule_wrap(function()
      self.current_index = self.current_index % #self.opts.spinner_frames + 1
      self.opts.extmark.virt_text = self:get_virtual_text()
      self:set_extmark()
    end)
  )
end

--- Stops the spinner animation
--- Cleans up the timer and removes the extmark
function VirtualTextSpinner:stop()
  if self.timer then
    self.timer:stop()
    self.timer:close()
    self.timer = nil
  end
  if self.opts.extmark then
    vim.schedule(function() vim.api.nvim_buf_del_extmark(self.bufnr, self.ns_id, self.opts.extmark.id) end)
  end
end

return VirtualTextSpinner

[ravitemer]

Sickkkkkk

[stuckinsnow]

With context added to the InlineStarted event data payload in #1293, it is now easy to add status columns icons or any type of extmarks when starting an inline prompt
...

This looks awesome! Unrelated but how did you configure your prompt to be near the cursor? It is nice!! haha

[lucobellic]

For the prompt position, I'm using snacks.nvim with input configured to use a style relative to cursor.
You can find my configuration here: https://github.com/lucobellic/nvim-config/blob/main/lua/plugins/snacks/snacks-input.lua