feat: highlight the selection in the source buffer

README.md

@@ -210,6 +210,7 @@ Also see [here](/lua/CopilotChat/config.lua):
   auto_follow_cursor = true, -- Auto-follow cursor in chat
   auto_insert_mode = false, -- Automatically enter insert mode when opening window and if auto follow cursor is enabled on new prompt
   clear_chat_on_new_prompt = false, -- Clears chat on every new prompt
+  highlight_selection = true, -- Highlight selection in the source buffer when in the chat window
 
   context = nil, -- Default context to use, 'buffers', 'buffer' or none (can be specified manually in prompt via @).
   history_path = vim.fn.stdpath('data') .. '/copilotchat_history', -- Default path to stored history

lua/CopilotChat/config.lua

@@ -69,6 +69,7 @@ local select = require('CopilotChat.select')
 ---@field auto_follow_cursor boolean?
 ---@field auto_insert_mode boolean?
 ---@field clear_chat_on_new_prompt boolean?
+---@field highlight_selection boolean?
 ---@field context string?
 ---@field history_path string?
 ---@field callback fun(response: string, source: CopilotChat.config.source)?
@@ -95,6 +96,7 @@ return {
   auto_follow_cursor = true, -- Auto-follow cursor in chat
   auto_insert_mode = false, -- Automatically enter insert mode when opening window and if auto follow cursor is enabled on new prompt
   clear_chat_on_new_prompt = false, -- Clears chat on every new prompt
+  highlight_selection = true, -- Highlight selection
 
   context = nil, -- Default context to use, 'buffers', 'buffer' or none (can be specified manually in prompt via @).
   history_path = vim.fn.stdpath('data') .. '/copilotchat_history', -- Default path to stored history

lua/CopilotChat/init.lua

@@ -275,6 +275,34 @@ function M.prompts(skip_system)
   return prompts_to_use
 end
 
+local selection_ns = nil ---@type number
+
+--- Highlights the selection in the source buffer.
+---@param clear? boolean
+function M.highlight_selection(clear)
+  for _, buf in ipairs(vim.api.nvim_list_bufs()) do
+    vim.api.nvim_buf_clear_namespace(buf, selection_ns, 0, -1)
+  end
+  if clear then
+    return
+  end
+  local selection = get_selection()
+  if selection then
+    vim.api.nvim_buf_set_extmark(
+      state.source.bufnr,
+      selection_ns,
+      selection.start_row - 1,
+      selection.start_col - 1,
+      {
+        hl_group = 'CopilotChatSelection',
+        end_row = selection.end_row - 1,
+        end_col = selection.end_col,
+        strict = false,
+      }
+    )
+  end
+end
+
 --- Open the chat window.
 ---@param config CopilotChat.config|CopilotChat.config.prompt|nil
 ---@param source CopilotChat.config.source?
@@ -578,11 +606,13 @@ function M.setup(config)
   state.copilot = Copilot(M.config.proxy, M.config.allow_insecure)
   M.debug(M.config.debug)
 
+  selection_ns = vim.api.nvim_create_namespace('copilot-chat-selection')
   local mark_ns = vim.api.nvim_create_namespace('copilot-chat-marks')
   local hl_ns = vim.api.nvim_create_namespace('copilot-chat-highlights')
   vim.api.nvim_set_hl(hl_ns, '@diff.plus', { bg = blend_color_with_neovim_bg('DiffAdd', 20) })
   vim.api.nvim_set_hl(hl_ns, '@diff.minus', { bg = blend_color_with_neovim_bg('DiffDelete', 20) })
   vim.api.nvim_set_hl(hl_ns, '@diff.delta', { bg = blend_color_with_neovim_bg('DiffChange', 20) })
+  vim.api.nvim_set_hl(0, 'CopilotChatSelection', { link = 'Visual', default = true })
 
   local overlay_help = ''
   if M.config.mappings.close then
@@ -790,6 +820,15 @@ function M.setup(config)
       end
     end)
 
+    vim.api.nvim_create_autocmd({ 'BufEnter', 'BufLeave' }, {
+      buffer = state.chat.bufnr,
+      callback = function(ev)
+        if state.config.highlight_selection then
+          M.highlight_selection(ev.event == 'BufLeave')
+        end
+      end,
+    })
+
     append(M.config.question_header .. M.config.separator .. '\n\n')
     state.chat:finish()
   end)