refactor: make client module stateless

The client module has been refactored to be stateless by making it a singleton
and moving provider initialization to a separate method. This simplifies the
overall architecture and makes the code more maintainable.

Additional changes:
- Improved reference handling and generation
- Updated header naming convention for Copilot requests
- Fixed embedding context resolution to pass proper params
- Removed redundant client state from main module

Author: deathbeam

lua/CopilotChat/client.lua

@@ -187,6 +187,21 @@ local function generate_embedding_request(inputs, threshold)
   end, inputs)
 end
 
+local function generate_references(references)
+  local out = ''
+
+  for _, reference in ipairs(references) do
+    out = out .. '\n[' .. reference.name .. '](' .. reference.url .. ')'
+  end
+
+  if out == '' then
+    return out
+  end
+
+  out = '\n\n**`References:`**' .. out
+  return out
+end
+
 ---@class CopilotChat.Client : Class
 ---@field providers table<string, CopilotChat.Provider>
 ---@field provider_cache table<string, table>
@@ -198,17 +213,12 @@ end
 ---@field token table?
 ---@field sessionid string?
 ---@field machineid string
-local Client = class(function(self, providers)
-  self.providers = providers
+local Client = class(function(self)
+  self.providers = {}
   self.embedding_cache = {}
   self.models = nil
   self.agents = nil
-
   self.provider_cache = {}
-  for provider_name, _ in pairs(providers) do
-    self.provider_cache[provider_name] = {}
-  end
-
   self.current_job = nil
   self.expires_at = nil
   self.headers = nil
@@ -359,6 +369,14 @@ function Client:ask(prompt, opts)
   log.debug('Tokenizer: ', tokenizer)
   tiktoken.load(tokenizer)
 
+  local references = {}
+  for _, embed in ipairs(embeddings) do
+    table.insert(references, {
+      name = embed.filename,
+      url = embed.filename,
+    })
+  end
+
   local generated_messages = {}
   local selection_messages = generate_selection_messages(selection)
   local embeddings_messages = generate_embeddings_messages(embeddings)
@@ -408,7 +426,6 @@ function Client:ask(prompt, opts)
   local errored = false
   local finished = false
   local full_response = ''
-  local full_references = ''
 
   local function finish_stream(err, job)
     if err then
@@ -450,14 +467,10 @@ function Client:ask(prompt, opts)
       for _, reference in ipairs(content.copilot_references) do
         local metadata = reference.metadata
         if metadata and metadata.display_name and metadata.display_url then
-          full_references = full_references
-            .. '\n'
-            .. '['
-            .. metadata.display_name
-            .. ']'
-            .. '('
-            .. metadata.display_url
-            .. ')'
+          table.insert(references, {
+            name = metadata.display_name,
+            url = metadata.display_url,
+          })
         end
       end
     end
@@ -615,8 +628,9 @@ function Client:ask(prompt, opts)
     return
   end
 
+  local full_references = generate_references(references)
+  log.info('References: ', full_references)
   if full_references ~= '' then
-    full_references = '\n\n**`References:`**' .. full_references
     full_response = full_response .. full_references
     if on_progress then
       on_progress(full_references)
@@ -833,4 +847,13 @@ function Client:running()
   return self.current_job ~= nil
 end
 
-return Client
+--- Load providers to client
+function Client:load_providers(providers)
+  self.providers = providers
+  for provider_name, _ in pairs(providers) do
+    self.provider_cache[provider_name] = {}
+  end
+end
+
+--- @type CopilotChat.Client
+return Client()

lua/CopilotChat/config/contexts.lua

@@ -5,7 +5,7 @@ local utils = require('CopilotChat.utils')
 ---@class CopilotChat.config.context
 ---@field description string?
 ---@field input fun(callback: fun(input: string?), source: CopilotChat.source)?
----@field resolve fun(input: string?, source: CopilotChat.source):table<CopilotChat.context.embed>
+---@field resolve fun(input: string?, source: CopilotChat.source, prompt: string, model: string):table<CopilotChat.context.embed>
 
 ---@type table<string, CopilotChat.config.context>
 return {

lua/CopilotChat/config/providers.lua

@@ -31,9 +31,10 @@ local EDITOR_VERSION = 'Neovim/'
   .. '.'
   .. vim.version().patch
 
---- Get the github oauth cached token
----@return string|nil
 local cached_github_token = nil
+
+--- Get the github copilot oauth cached token (gu_ token)
+---@return string
 local function get_github_token()
   if cached_github_token then
     return cached_github_token
@@ -84,10 +85,11 @@ M.copilot = {
 
   get_headers = function(token)
     return {
-      ['Authorization'] = 'Bearer ' .. token,
-      ['Editor-Version'] = EDITOR_VERSION,
-      ['Copilot-Integration-Id'] = 'vscode-chat',
-      ['Content-Type'] = 'application/json',
+      ['authorization'] = 'Bearer ' .. token,
+      ['editor-version'] = EDITOR_VERSION,
+      ['editor-plugin-version'] = 'CopilotChat.nvim/*',
+      ['copilot-integration-id'] = 'vscode-chat',
+      ['content-type'] = 'application/json',
     }
   end,
 

lua/CopilotChat/context.lua

@@ -14,9 +14,11 @@
 ---@field outline string?
 ---@field symbols table<string, CopilotChat.context.symbol>?
 ---@field embedding table<number>?
+---@field score?
 
 local async = require('plenary.async')
 local log = require('plenary.log')
+local client = require('CopilotChat.client')
 local notify = require('CopilotChat.notify')
 local utils = require('CopilotChat.utils')
 local file_cache = {}
@@ -100,7 +102,7 @@ local function data_ranked_by_relatedness(query, data, top_n)
     return vim.tbl_extend(
       'force',
       item,
-      { score = spatial_distance_cosine(item.embedding, query.embedding) }
+      { score = item.score or spatial_distance_cosine(item.embedding, query.embedding) }
     )
   end, data)
 
@@ -583,12 +585,11 @@ function M.quickfix()
 end
 
 --- Filter embeddings based on the query
----@param client CopilotChat.Client
 ---@param prompt string
 ---@param model string
 ---@param embeddings table<CopilotChat.context.embed>
 ---@return table<CopilotChat.context.embed>
-function M.filter_embeddings(client, prompt, model, embeddings)
+function M.filter_embeddings(prompt, model, embeddings)
   -- If we dont need to embed anything, just return directly
   if #embeddings < MULTI_FILE_THRESHOLD then
     return embeddings

lua/CopilotChat/init.lua

@@ -1,6 +1,7 @@
 local async = require('plenary.async')
 local log = require('plenary.log')
 local context = require('CopilotChat.context')
+local client = require('CopilotChat.client')
 local utils = require('CopilotChat.utils')
 
 local M = {}
@@ -12,16 +13,13 @@ local WORD = '([^%s]+)'
 --- @field winnr number
 
 --- @class CopilotChat.state
---- @field client CopilotChat.Client?
 --- @field source CopilotChat.source?
 --- @field last_prompt string?
 --- @field last_response string?
 --- @field chat CopilotChat.ui.Chat?
 --- @field diff CopilotChat.ui.Diff?
 --- @field overlay CopilotChat.ui.Overlay?
 local state = {
-  client = nil,
-
   -- Current state tracking
   source = nil,
 
@@ -221,9 +219,10 @@ end
 
 --- Resolve the embeddings from the prompt.
 ---@param prompt string
+---@param model string
 ---@param config CopilotChat.config.shared
 ---@return table<CopilotChat.context.embed>, string
-function M.resolve_embeddings(prompt, config)
+function M.resolve_embeddings(prompt, model, config)
   local contexts = {}
   local function parse_context(prompt_context)
     local split = vim.split(prompt_context, ':')
@@ -262,7 +261,9 @@ function M.resolve_embeddings(prompt, config)
   local embeddings = utils.ordered_map()
   for _, context_data in ipairs(contexts) do
     local context_value = M.config.contexts[context_data.name]
-    for _, embedding in ipairs(context_value.resolve(context_data.input, state.source or {})) do
+    for _, embedding in
+      ipairs(context_value.resolve(context_data.input, state.source or {}, prompt, model))
+    do
       if embedding then
         embeddings:set(embedding.filename, embedding)
       end
@@ -276,7 +277,7 @@ end
 ---@param prompt string
 ---@param config CopilotChat.config.shared
 function M.resolve_agent(prompt, config)
-  local agents = vim.tbl_keys(state.client:list_agents())
+  local agents = vim.tbl_keys(client:list_agents())
   local selected_agent = config.agent
   prompt = prompt:gsub('@' .. WORD, function(match)
     if vim.tbl_contains(agents, match) then
@@ -295,7 +296,7 @@ end
 function M.resolve_model(prompt, config)
   local models = vim.tbl_map(function(model)
     return model.id
-  end, state.client:list_models())
+  end, client:list_models())
 
   local selected_model = config.model
   prompt = prompt:gsub('%$' .. WORD, function(match)
@@ -393,8 +394,8 @@ end
 ---@param callback function(table)
 function M.complete_items(callback)
   async.run(function()
-    local models = state.client:list_models()
-    local agents = state.client:list_agents()
+    local models = client:list_models()
+    local agents = client:list_agents()
     local prompts_to_use = M.prompts()
     local items = {}
 
@@ -536,7 +537,7 @@ end
 --- Select default Copilot GPT model.
 function M.select_model()
   async.run(function()
-    local models = state.client:list_models()
+    local models = client:list_models()
     local choices = vim.tbl_map(function(model)
       return {
         id = model.id,
@@ -567,7 +568,7 @@ end
 --- Select default Copilot agent.
 function M.select_agent()
   async.run(function()
-    local agents = state.client:list_agents()
+    local agents = client:list_agents()
     local choices = vim.tbl_map(function(agent)
       return {
         id = agent.id,
@@ -613,7 +614,7 @@ function M.ask(prompt, config)
   if not config.headless then
     if config.clear_chat_on_new_prompt then
       M.stop(true)
-    elseif state.client:stop() then
+    elseif client:stop() then
       finish()
     end
 
@@ -642,13 +643,13 @@ function M.ask(prompt, config)
   local selection = M.get_selection(config)
 
   local ok, err = pcall(async.run, function()
-    local embeddings, prompt = M.resolve_embeddings(prompt, config)
     local selected_agent, prompt = M.resolve_agent(prompt, config)
     local selected_model, prompt = M.resolve_model(prompt, config)
+    local embeddings, prompt = M.resolve_embeddings(prompt, selected_model, config)
 
     local has_output = false
     local query_ok, filtered_embeddings =
-      pcall(context.filter_embeddings, state.client, prompt, selected_model, embeddings)
+      pcall(context.filter_embeddings, prompt, selected_model, embeddings)
 
     if not query_ok then
       async.util.scheduler()
@@ -659,22 +660,21 @@ function M.ask(prompt, config)
       return
     end
 
-    local ask_ok, response, token_count, token_max_count =
-      pcall(state.client.ask, state.client, prompt, {
-        history = history,
-        selection = selection,
-        embeddings = filtered_embeddings,
-        system_prompt = system_prompt,
-        model = selected_model,
-        agent = selected_agent,
-        temperature = config.temperature,
-        on_progress = vim.schedule_wrap(function(token)
-          if not config.headless then
-            state.chat:append(token)
-          end
-          has_output = true
-        end),
-      })
+    local ask_ok, response, token_count, token_max_count = pcall(client.ask, client, prompt, {
+      history = history,
+      selection = selection,
+      embeddings = filtered_embeddings,
+      system_prompt = system_prompt,
+      model = selected_model,
+      agent = selected_agent,
+      temperature = config.temperature,
+      on_progress = vim.schedule_wrap(function(token)
+        if not config.headless then
+          state.chat:append(token)
+        end
+        has_output = true
+      end),
+    })
 
     async.util.scheduler()
 
@@ -716,12 +716,12 @@ end
 ---@param reset boolean?
 function M.stop(reset)
   if reset then
-    state.client:reset()
+    client:reset()
     state.chat:clear()
     state.last_prompt = nil
     state.last_response = nil
   else
-    state.client:stop()
+    client:stop()
   end
 
   finish(reset)
@@ -791,7 +791,7 @@ function M.load(name, history_path)
     },
   })
 
-  state.client:reset()
+  client:reset()
   state.chat:clear()
   state.chat:load_history(history)
   log.info('Loaded history from ' .. history_path)
@@ -873,10 +873,9 @@ function M.setup(config)
     proxy = M.config.proxy,
   })
 
-  if state.client then
-    state.client:stop()
-  end
-  state.client = require('CopilotChat.client')(M.config.providers)
+  -- Load the providers
+  client:stop()
+  client:load_providers(M.config.providers)
 
   if M.config.debug then
     M.log_level('debug')