From 0ee7b26b11b2625cdd4f37d462d5deb772acfae2 Mon Sep 17 00:00:00 2001
From: Tomas Slusny <slusnucky@gmail.com>
Date: Mon, 15 Sep 2025 17:30:02 +0200
Subject: [PATCH] fix(client): correct history handling for headless ask

Refactors ask request construction to properly handle history in headless
mode. Removes prompt parameter from internal request generation and ensures
the prompt is included as a user message only when appropriate. Also
updates token counting and history trimming logic to avoid removing the
current prompt in headless scenarios.

Resolves: #1415
---
 lua/CopilotChat/client.lua | 28 ++++++++--------------------
 lua/CopilotChat/init.lua   | 15 +++++++++++++--
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/lua/CopilotChat/client.lua b/lua/CopilotChat/client.lua
index cf240a55..93e1c91d 100644
--- a/lua/CopilotChat/client.lua
+++ b/lua/CopilotChat/client.lua
@@ -142,11 +142,10 @@ local function generate_resource_messages(resources)
 end
 
 --- Generate ask request
---- @param prompt string
 --- @param system_prompt string
 --- @param history table<CopilotChat.client.Message>
 --- @param generated_messages table<CopilotChat.client.Message>
-local function generate_ask_request(prompt, system_prompt, history, generated_messages)
+local function generate_ask_request(system_prompt, history, generated_messages)
   local messages = {}
 
   system_prompt = vim.trim(system_prompt)
@@ -162,15 +161,6 @@ local function generate_ask_request(prompt, system_prompt, history, generated_me
   -- Include generated messages and history
   vim.list_extend(messages, generated_messages)
   vim.list_extend(messages, history)
-
-  -- Include user prompt if we have no history
-  if not utils.empty(prompt) and utils.empty(history) then
-    table.insert(messages, {
-      content = prompt,
-      role = constants.ROLE.USER,
-    })
-  end
-
   return messages
 end
 
@@ -303,10 +293,9 @@ function Client:info()
 end
 
 --- Ask a question to Copilot
----@param prompt string: The prompt to send to Copilot
 ---@param opts CopilotChat.client.AskOptions: Options for the request
 ---@return CopilotChat.client.AskResponse?
-function Client:ask(prompt, opts)
+function Client:ask(opts)
   opts = opts or {}
   local job_id = utils.uuid()
 
@@ -350,20 +339,20 @@ function Client:ask(prompt, opts)
     notify.publish(notify.STATUS, 'Generating request')
   end
 
-  local history = not opts.headless and vim.deepcopy(opts.history) or {}
+  local history = vim.deepcopy(opts.history)
   local tool_calls = orderedmap()
   local generated_messages = {}
   local resource_messages = generate_resource_messages(opts.resources)
 
   if max_tokens then
     -- Count required tokens that we cannot reduce
-    local prompt_tokens = tiktoken:count(prompt)
     local system_tokens = tiktoken:count(opts.system_prompt)
+    local prompt_tokens = #history > 0 and tiktoken:count(history[#history].content) or 0
     local resource_tokens = #resource_messages > 0 and tiktoken:count(resource_messages[1].content) or 0
     local required_tokens = prompt_tokens + system_tokens + resource_tokens
 
-    log.debug('Prompt tokens:', prompt_tokens)
     log.debug('System tokens:', system_tokens)
+    log.debug('Prompt tokens:', prompt_tokens)
     log.debug('Resource tokens:', resource_tokens)
 
     -- Calculate how many tokens we can use for history
@@ -373,8 +362,8 @@ function Client:ask(prompt, opts)
       history_tokens = history_tokens + tiktoken:count(msg.content)
     end
 
-    -- Remove history messages until we are under the limit
-    while history_tokens > history_limit and #history > 0 do
+    -- Remove history messages except prompt until we are under the limit
+    while history_tokens > history_limit and #history > 1 do
       local entry = table.remove(history, 1)
       history_tokens = history_tokens - tiktoken:count(entry.content)
     end
@@ -522,8 +511,7 @@ function Client:ask(prompt, opts)
   end
 
   local headers = self:authenticate(provider_name)
-  local request =
-    provider.prepare_input(generate_ask_request(prompt, opts.system_prompt, history, generated_messages), options)
+  local request = provider.prepare_input(generate_ask_request(opts.system_prompt, history, generated_messages), options)
   local is_stream = request.stream
 
   local args = {
diff --git a/lua/CopilotChat/init.lua b/lua/CopilotChat/init.lua
index f473750a..d50251f1 100644
--- a/lua/CopilotChat/init.lua
+++ b/lua/CopilotChat/init.lua
@@ -888,9 +888,20 @@ function M.ask(prompt, config)
         return
       end
 
-      local ask_response = client.ask(client, prompt, {
+      -- Build history, when in headless mode its just current prompt
+      local history
+      if not config.headless then
+        history = M.chat:get_messages()
+      else
+        history = {
+          content = prompt,
+          role = constants.ROLE.USER,
+        }
+      end
+
+      local ask_response = client:ask({
         headless = config.headless,
-        history = M.chat:get_messages(),
+        history = history,
         resources = resolved_resources,
         tools = selected_tools,
         system_prompt = system_prompt,