feat(chat): add streaming callback support

Add stream callback function to process tokens as they are received from Copilot
while callback function is used to process the complete final response. Both
functions can modify the text that gets displayed and stored in history.

This separates streaming token processing from final response handling, giving
more control over how chat outputs are processed and displayed.

Signed-off-by: Tomas Slusny <slusnucky@gmail.com>

Author: deathbeam

README.md

@@ -460,8 +460,9 @@ Below are all available configuration options with their default values:
   sticky = nil, -- Default sticky prompt or array of sticky prompts to use at start of every new chat.
 
   temperature = 0.1, -- GPT result temperature
-  headless = false, -- Do not write to chat buffer and use history(useful for using callback for custom processing)
-  callback = nil, -- Callback to use when ask response is received
+  headless = false, -- Do not write to chat buffer and use history (useful for using custom processing)
+  stream = nil, -- Function called when receiving stream updates (returned string is appended to the chat buffer)
+  callback = nil, -- Function called when full response is received (retuned string is stored to history)
   remember_as_sticky = true, -- Remember model/agent/context as sticky prompts when asking questions
 
   -- default selection
@@ -748,6 +749,7 @@ require("CopilotChat").open()
 require("CopilotChat").ask("Explain this code", {
   callback = function(response)
     vim.notify("Got response: " .. response:sub(1, 50) .. "...")
+    return response
   end,
   context = "#buffer"
 })

lua/CopilotChat/client.lua

@@ -1,6 +1,6 @@
 ---@class CopilotChat.Client.ask
 ---@field load_history boolean
----@field store_history boolean
+---@field summarize_history boolean
 ---@field contexts table<string, string>?
 ---@field selection CopilotChat.select.selection?
 ---@field embeddings table<CopilotChat.context.embed>?
@@ -530,7 +530,7 @@ function Client:ask(prompt, opts)
 
     -- If we're over history limit, trigger summarization
     if history_tokens > history_limit then
-      if opts.store_history and #history >= 4 then
+      if opts.summarize_history and #history >= 4 then
         self:summarize_history(opts.model)
 
         -- Recalculate history and tokens
@@ -755,18 +755,6 @@ function Client:ask(prompt, opts)
     return
   end
 
-  if opts.store_history then
-    table.insert(self.history, {
-      content = prompt,
-      role = 'user',
-    })
-
-    table.insert(self.history, {
-      content = response_text,
-      role = 'assistant',
-    })
-  end
-
   return response_text, references:values(), last_message and last_message.total_tokens or 0, max_tokens
 end
 
@@ -943,7 +931,7 @@ Ensure all critical code samples, commands, and configuration snippets are prese
 
   local response = self:ask('Create a technical summary of our conversation for future context', {
     load_history = true,
-    store_history = false,
+    summarize_history = false,
     model = model,
     temperature = 0,
     system_prompt = system_prompt,

lua/CopilotChat/config.lua

@@ -20,7 +20,8 @@ local select = require('CopilotChat.select')
 ---@field sticky string|table<string>|nil
 ---@field temperature number?
 ---@field headless boolean?
----@field callback fun(response: string, source: CopilotChat.source)?
+---@field stream nil|fun(chunk: string, source: CopilotChat.source):string
+---@field callback nil|fun(response: string, source: CopilotChat.source):string
 ---@field remember_as_sticky boolean?
 ---@field selection false|nil|fun(source: CopilotChat.source):CopilotChat.select.selection?
 ---@field window CopilotChat.config.window?
@@ -63,8 +64,9 @@ return {
   sticky = nil, -- Default sticky prompt or array of sticky prompts to use at start of every new chat.
 
   temperature = 0.1, -- GPT result temperature
-  headless = false, -- Do not write to chat buffer and use history(useful for using callback for custom processing)
-  callback = nil, -- Callback to use when ask response is received
+  headless = false, -- Do not write to chat buffer and use history (useful for using custom processing)
+  stream = nil, -- Function called when receiving stream updates (returned string is appended to the chat buffer)
+  callback = nil, -- Function called when full response is received (retuned string is stored to history)
   remember_as_sticky = true, -- Remember model/agent/context as sticky prompts when asking questions
 
   -- default selection

lua/CopilotChat/config/prompts.lua

@@ -141,6 +141,7 @@ return {
         end
       end
       vim.diagnostic.set(vim.api.nvim_create_namespace('copilot-chat-diagnostics'), source.bufnr, diagnostics)
+      return response
     end,
   },
 

lua/CopilotChat/init.lua

@@ -197,8 +197,7 @@ end
 
 --- Show an error in the chat window.
 ---@param err string|table|nil
----@param append_newline boolean?
-local function show_error(err, append_newline)
+local function show_error(err)
   err = err or 'Unknown error'
 
   if type(err) == 'string' then
@@ -213,11 +212,7 @@ local function show_error(err, append_newline)
     err = utils.make_string(err)
   end
 
-  if append_newline then
-    M.chat:append('\n')
-  end
-
-  M.chat:append(M.config.error_header .. '\n```error\n' .. err .. '\n```')
+  M.chat:append('\n' .. M.config.error_header .. '\n```error\n' .. err .. '\n```')
   finish()
 end
 
@@ -876,22 +871,21 @@ function M.ask(prompt, config)
     local selected_model, prompt = M.resolve_model(prompt, config)
     local embeddings, prompt = M.resolve_context(prompt, config)
 
-    local has_output = false
     local query_ok, filtered_embeddings =
       pcall(context.filter_embeddings, prompt, selected_model, config.headless, embeddings)
 
     if not query_ok then
       utils.schedule_main()
       log.error(filtered_embeddings)
       if not config.headless then
-        show_error(filtered_embeddings, has_output)
+        show_error(filtered_embeddings)
       end
       return
     end
 
     local ask_ok, response, references, token_count, token_max_count = pcall(client.ask, client, prompt, {
       load_history = not config.headless,
-      store_history = not config.headless,
+      summarize_history = not config.headless,
       contexts = contexts,
       selection = selection,
       embeddings = filtered_embeddings,
@@ -900,10 +894,16 @@ function M.ask(prompt, config)
       agent = selected_agent,
       temperature = config.temperature,
       on_progress = vim.schedule_wrap(function(token)
-        if not config.headless then
+        local to_print = not config.headless and token
+        if to_print and config.stream then
+          local out = config.stream(token, state.source)
+          if out ~= nil then
+            to_print = out
+          end
+        end
+        if to_print and to_print ~= '' then
           M.chat:append(token)
         end
-        has_output = true
       end),
     })
 
@@ -912,11 +912,30 @@ function M.ask(prompt, config)
     if not ask_ok then
       log.error(response)
       if not config.headless then
-        show_error(response, has_output)
+        show_error(response)
       end
       return
     end
 
+    -- Call the callback function and store to history
+    local to_store = not config.headless and response
+    if to_store and config.callback then
+      local out = config.callback(response, state.source)
+      if out ~= nil then
+        to_store = out
+      end
+    end
+    if to_store and to_store ~= '' then
+      table.insert(client.history, {
+        content = prompt,
+        role = 'user',
+      })
+      table.insert(client.history, {
+        content = to_store,
+        role = 'assistant',
+      })
+    end
+
     if not config.headless then
       state.last_response = response
       M.chat.references = references
@@ -932,10 +951,6 @@ function M.ask(prompt, config)
 
       finish()
     end
-
-    if config.callback then
-      config.callback(response, state.source)
-    end
   end)
 
   if not ok then