diff --git a/lua/CopilotChat/config/mappings.lua b/lua/CopilotChat/config/mappings.lua index 46f47e01..23867f61 100644 --- a/lua/CopilotChat/config/mappings.lua +++ b/lua/CopilotChat/config/mappings.lua @@ -279,9 +279,10 @@ return { normal = 'gqa', callback = function() local items = {} - for i, message in ipairs(copilot.chat.messages) do + local messages = copilot.chat:get_messages() + for i, message in ipairs(messages) do if message.section and message.role == constants.ROLE.ASSISTANT then - local prev_message = copilot.chat.messages[i - 1] + local prev_message = messages[i - 1] local text = '' if prev_message then text = prev_message.content @@ -305,8 +306,8 @@ return { normal = 'gqd', callback = function(source) local items = {} - - for _, message in ipairs(copilot.chat.messages) do + local messages = copilot.chat:get_messages() + for _, message in ipairs(messages) do if message.section then for _, block in ipairs(message.section.blocks) do local diff = get_diff(source.bufnr, block) diff --git a/lua/CopilotChat/init.lua b/lua/CopilotChat/init.lua index d87b90e7..69d6ac77 100644 --- a/lua/CopilotChat/init.lua +++ b/lua/CopilotChat/init.lua @@ -356,7 +356,7 @@ function M.resolve_functions(prompt, config) local resolved_resources = {} local resolved_tools = {} local tool_calls = {} - for _, message in ipairs(M.chat.messages) do + for _, message in ipairs(M.chat:get_messages()) do if message.tool_calls then for _, tool_call in ipairs(message.tool_calls) do table.insert(tool_calls, tool_call) @@ -868,7 +868,7 @@ function M.ask(prompt, config) local ask_response = client.ask(client, prompt, { headless = config.headless, - history = M.chat.messages, + history = M.chat:get_messages(), resources = resolved_resources, tools = selected_tools, system_prompt = system_prompt, @@ -948,7 +948,7 @@ function M.save(name, history_path) return end - local history = vim.deepcopy(M.chat.messages) + local history = vim.deepcopy(M.chat:get_messages()) for _, message in ipairs(history) do message.section = nil end diff --git a/lua/CopilotChat/ui/chat.lua b/lua/CopilotChat/ui/chat.lua index 4681779b..95d7fc1e 100644 --- a/lua/CopilotChat/ui/chat.lua +++ b/lua/CopilotChat/ui/chat.lua @@ -4,6 +4,7 @@ local constants = require('CopilotChat.constants') local notify = require('CopilotChat.notify') local utils = require('CopilotChat.utils') local class = require('CopilotChat.utils.class') +local orderedmap = require('CopilotChat.utils.orderedmap') function CopilotChatFoldExpr(lnum, separator) local to_match = separator .. '$' @@ -93,7 +94,7 @@ end ---@field config CopilotChat.config.Shared ---@field token_count number? ---@field token_max_count number? ----@field messages table +---@field private messages OrderedMap ---@field private layout CopilotChat.config.Layout? ---@field private headers table ---@field private separator string @@ -106,7 +107,7 @@ local Chat = class(function(self, config, on_buf_create) self.config = config self.token_count = nil self.token_max_count = nil - self.messages = {} + self.messages = orderedmap() self.layout = nil self.headers = {} @@ -168,6 +169,9 @@ end ---@param cursor boolean? If true, returns the block closest to the cursor position ---@return CopilotChat.ui.chat.Block? function Chat:get_block(role, cursor) + self:parse() + local messages = self:get_messages() + if cursor then if not self:visible() then return nil @@ -178,7 +182,7 @@ function Chat:get_block(role, cursor) local closest_block = nil local max_line_below_cursor = -1 - for _, message in ipairs(self.messages) do + for _, message in ipairs(messages) do local section = message.section local matches_role = not role or message.role == role if matches_role and section and section.blocks then @@ -194,8 +198,8 @@ function Chat:get_block(role, cursor) return closest_block end - for i = #self.messages, 1, -1 do - local message = self.messages[i] + for i = #messages, 1, -1 do + local message = messages[i] local matches_role = not role or message.role == role if matches_role and message.section and message.section.blocks and #message.section.blocks > 0 then return message.section.blocks[#message.section.blocks] @@ -203,12 +207,19 @@ function Chat:get_block(role, cursor) end end +--- Get list of all chat messages +---@return table +function Chat:get_messages() + return self.messages:values() +end + --- Get last message by role in the chat window. ---@param role string? If specified, only considers sections of the given role ---@param cursor boolean? If true, returns the message closest to the cursor position ---@return CopilotChat.ui.chat.Message? function Chat:get_message(role, cursor) self:parse() + local messages = self:get_messages() if cursor then if not self:visible() then @@ -220,7 +231,7 @@ function Chat:get_message(role, cursor) local closest_message = nil local max_line_below_cursor = -1 - for _, message in ipairs(self.messages) do + for _, message in ipairs(messages) do local section = message.section local matches_role = not role or message.role == role if matches_role and section.start_line <= cursor_line and section.start_line > max_line_below_cursor then @@ -232,8 +243,8 @@ function Chat:get_message(role, cursor) return closest_message end - for i = #self.messages, 1, -1 do - local message = self.messages[i] + for i = #messages, 1, -1 do + local message = messages[i] local matches_role = not role or message.role == role if matches_role then return message @@ -479,7 +490,8 @@ end function Chat:add_message(message, replace) self:parse() - local current_message = self.messages[#self.messages] + local messages = self:get_messages() + local current_message = messages[#messages] local is_new = not current_message or current_message.role ~= message.role or (message.id and current_message.id ~= message.id) @@ -488,7 +500,7 @@ function Chat:add_message(message, replace) -- Add appropriate header based on role and generate a new ID if not provided message.id = message.id or utils.uuid() local header = self.headers[message.role] - table.insert(self.messages, message) + self.messages:set(message.id, message) if current_message then self:append('\n') @@ -546,12 +558,7 @@ function Chat:remove_message(role, cursor) vim.bo[self.bufnr].modifiable = modifiable -- Remove the message from the messages list - for i, msg in ipairs(self.messages) do - if msg.id == message.id then - table.remove(self.messages, i) - break - end - end + self.messages:remove(message.id) end --- Append text to the chat window. @@ -585,7 +592,7 @@ function Chat:clear() self:validate() self.token_count = nil self.token_max_count = nil - self.messages = {} + self.messages = orderedmap() local modifiable = vim.bo[self.bufnr].modifiable vim.bo[self.bufnr].modifiable = true @@ -718,15 +725,8 @@ function Chat:parse() -- Finish last message current_message.section.end_line = vim.api.nvim_buf_line_count(self.bufnr) - -- Build lookup table for previous messages by id - local old_messages_by_id = {} - for _, msg in ipairs(self.messages or {}) do - if msg.id then - old_messages_by_id[msg.id] = msg - end - end - -- Format new messages and preserve extra fields from old messages + local messages = orderedmap() for _, message in ipairs(new_messages) do message.content = vim.trim(table.concat(message.content, '\n')) if message.section then @@ -735,7 +735,7 @@ function Chat:parse() end end - local old = old_messages_by_id[message.id] + local old = self.messages:get(message.id) if old then for k, v in pairs(old) do if message[k] == nil then @@ -743,9 +743,12 @@ function Chat:parse() end end end + + messages:set(message.id, message) end - self.messages = new_messages + -- Update messages + self.messages = messages end --- Render the chat window. @@ -757,7 +760,9 @@ function Chat:render() vim.api.nvim_buf_clear_namespace(self.bufnr, highlight_ns, 0, -1) -- Clear previous highlights self:show_help() -- Clear previous help - for i, message in ipairs(self.messages) do + local messages = self:get_messages() + + for i, message in ipairs(messages) do if self.config.highlight_headers then -- Overlay section header with nice display local header_value = self.headers[message.role] @@ -847,7 +852,7 @@ function Chat:render() end end - if i == #self.messages and message.role == constants.ROLE.USER then + if i == #messages and message.role == constants.ROLE.USER then -- Highlight tools in the last user message local assistant_msg = self:get_message(constants.ROLE.ASSISTANT) if assistant_msg and assistant_msg.tool_calls and #assistant_msg.tool_calls > 0 then @@ -883,7 +888,7 @@ function Chat:render() -- Auto fold non-assistant messages if enabled if self.config.auto_fold and self:visible() then - if message.role ~= constants.ROLE.ASSISTANT and message.section and i < #self.messages then + if message.role ~= constants.ROLE.ASSISTANT and message.section and i < #messages then vim.api.nvim_win_call(self.winnr, function() local fold_level = vim.fn.foldlevel(message.section.start_line) if fold_level > 0 and vim.fn.foldclosed(message.section.start_line) == -1 then diff --git a/lua/CopilotChat/utils/orderedmap.lua b/lua/CopilotChat/utils/orderedmap.lua index 778c686d..1907c161 100644 --- a/lua/CopilotChat/utils/orderedmap.lua +++ b/lua/CopilotChat/utils/orderedmap.lua @@ -1,6 +1,7 @@ ---@class OrderedMap ---@field set fun(self:OrderedMap, key:any, value:any) ---@field get fun(self:OrderedMap, key:any):any +---@field remove fun(self:OrderedMap, key:any) ---@field keys fun(self:OrderedMap):table ---@field values fun(self:OrderedMap):table @@ -22,6 +23,18 @@ local function orderedmap() return self._data[key] end, + remove = function(self, key) + if self._data[key] then + self._data[key] = nil + for i, k in ipairs(self._keys) do + if k == key then + table.remove(self._keys, i) + break + end + end + end + end, + keys = function(self) return self._keys end, diff --git a/tests/orderedmap_spec.lua b/tests/orderedmap_spec.lua index 9000915c..b5fa5a37 100644 --- a/tests/orderedmap_spec.lua +++ b/tests/orderedmap_spec.lua @@ -25,4 +25,13 @@ describe('CopilotChat.utils.orderedmap', function() assert.are.same({ 'a' }, map:keys()) assert.are.same({ 2 }, map:values()) end) + + it('removes values and updates order', function() + local map = orderedmap() + map:set('a', 1) + map:set('b', 2) + map:remove('a') + assert.are.same({ 'b' }, map:keys()) + assert.are.same({ 2 }, map:values()) + end) end)