refactor(history): reduce computational complexity when handling history

When updating chat history to be used in LLM request there are several
instances where we do O(n^2) operations: scanning all messages to locate
a tool "use" and for each use scan messages again to locate
corresponding "result".

Refactor the code to scan messages once collecting tool uses and results
together, and then do 2nd scan to drop incomplete tool invocations and
refresh "view" and "edit" results with the latest content.

Also reduce number of pre-scan loops (where we discard partially
generated messages or messages that are not interesting or too-old) by
combining them when possible.

This reduces time to scan initial 417 messages on my system (which
result in 576 final messages) from 0.32 to 0.12 seconds.

lua/avante/history/helpers.lua

@@ -67,20 +67,6 @@ function M.get_tool_use_message(message, messages)
   end
 end
 
----@param tool_use_message avante.HistoryMessage | nil
-function M.is_edit_func_call_message(tool_use_message)
-  local is_replace_func_call = false
-  local is_str_replace_editor_func_call = false
-  local is_str_replace_based_edit_tool_func_call = false
-  local path = nil
-  if tool_use_message and M.is_tool_use_message(tool_use_message) then
-    local tool_use = tool_use_message.message.content[1]
-    ---@cast tool_use AvanteLLMToolUse
-    return Utils.is_edit_func_call_tool_use(tool_use)
-  end
-  return is_replace_func_call, is_str_replace_editor_func_call, is_str_replace_based_edit_tool_func_call, path
-end
-
 ---Given a tool use message locate corresponding tool result message
 ---@param message avante.HistoryMessage
 ---@param messages avante.HistoryMessage[]

lua/avante/history/init.lua

@@ -41,62 +41,144 @@ function M.get_history_messages(history)
   return messages
 end
 
+---Represents information about tool use: invocation, result, affected file (for "view" or "edit" tools).
+---@class HistoryToolInfo
+---@field kind "edit" | "view" | "other"
+---@field use AvanteLLMToolUse
+---@field result? AvanteLLMToolResult
+---@field result_message? avante.HistoryMessage Complete result message
+---@field path? string Uniform (normalized) path of the affected file
+
+---@class HistoryFileInfo
+---@field last_tool_id? string ID of the tool with most up-to-date state of the file
+---@field edit_tool_id? string ID of the last tool done edit on the file
+
+---Collects information about all uses of tools in the history: their invocations, results, and affected files.
 ---@param messages avante.HistoryMessage[]
----@param using_ReAct_prompt boolean
+---@return table<string, HistoryToolInfo>
----@param add_diagnostic boolean Mix in LSP diagnostic info for affected files
+---@return table<string, HistoryFileInfo>
+local function collect_tool_info(messages)
+  ---@type table<string, HistoryToolInfo> Maps tool ID to tool information
+  local tools = {}
+  ---@type table<string, HistoryFileInfo> Maps file path to file information
+  local files = {}
+
+  -- Collect invocations of all tools, and also build a list of viewed or edited files.
+  for _, message in ipairs(messages) do
+    local use = Helpers.get_tool_use_data(message)
+    if use then
+      if use.name == "view" or Utils.is_edit_tool_use(use) then
+        if use.input.path then
+          local path = Utils.uniform_path(use.input.path)
+          tools[use.id] = { kind = use.name == "view" and "view" or "edit", use = use, path = path }
+        end
+      else
+        tools[use.id] = { kind = "other", use = use }
+      end
+      goto continue
+    end
+
+    local result = Helpers.get_tool_result_data(message)
+    if result then
+      -- We assume that "result" entries always come after corresponding "use" entries.
+      local info = tools[result.tool_use_id]
+      if info then
+        info.result = result
+        info.result_message = message
+        if info.path then
+          local f = files[info.path]
+          if not f then
+            f = {}
+            files[info.path] = f
+          end
+          f.last_tool_id = result.tool_use_id
+          if info.kind == "edit" and not (result.is_error or result.is_user_declined) then
+            f.edit_tool_id = result.tool_use_id
+          end
+        end
+      end
+    end
+
+    ::continue::
+  end
+
+  return tools, files
+end
+
+---Converts a tool invocation (use + result) into a simple request/response pair of text messages
+---@param tool_info HistoryToolInfo
 ---@return avante.HistoryMessage[]
-M.update_history_messages = function(messages, using_ReAct_prompt, add_diagnostic)
+local function convert_tool_to_text(tool_info)
-  local tool_id_to_tool_name = {}
+  return {
-  local tool_id_to_path = {}
+    Message:new_assistant_synthetic(
-  local tool_id_to_start_line = {}
+      string.format("Tool use %s(%s)", tool_info.use.name, vim.json.encode(tool_info.use.input))
-  local tool_id_to_end_line = {}
+    ),
-  local viewed_files = {}
+    Message:new_user_synthetic({
-  local last_modified_files = {}
+      type = "text",
-  local history_messages = {}
+      text = string.format(
+        "Tool use [%s] is successful: %s",
+        tool_info.use.name,
+        tostring(not tool_info.result.is_error)
+      ),
+    }),
+  }
+end
 
-  for idx, message in ipairs(messages) do
+---Generates a fake file "content" telling LLM to look further for up-to-date data
-    if Helpers.is_tool_result_message(message) then
+---@param path string
-      local tool_use_message = Helpers.get_tool_use_message(message, messages)
+---@return string
+local function stale_view_content(path)
+  return string.format("The file %s has been updated. Please use the latest `view` tool result!", path)
+end
 
-      local is_edit_func_call, _, _, path = Helpers.is_edit_func_call_message(tool_use_message)
+---Updates the result of "view" tool invocation with latest contents of a buffer or file,
+---or a stub message if this result will be superseded by another one.
+---@param tool_info HistoryToolInfo
+---@param stale_view boolean
+local function update_view_result(tool_info, stale_view)
+  local use = tool_info.use
+  local result = tool_info.result
 
-      -- Only track as successful modification if not an error AND not user-declined
+  if stale_view then
-      if
+    result.content = stale_view_content(tool_info.path)
-        is_edit_func_call
+  else
-        and path
+    local view_result, view_error = require("avante.llm_tools.view").func(
-        and not message.message.content[1].is_error
+      { path = tool_info.path, start_line = use.input.start_line, end_line = use.input.end_line },
-        and not message.message.content[1].is_user_declined
+      {}
-      then
+    )
-        local uniformed_path = Utils.uniform_path(path)
+    result.content = view_error and ("Error: " .. view_error) or view_result
-        last_modified_files[uniformed_path] = idx
+    result.is_error = view_error ~= nil
-      end
   end
+end
+
+---Generates synthetic "view" tool invocation to tell LLM to refresh its view of a file after editing
+---@param tool_use AvanteLLMToolUse
+---@param path any
+---@param stale_view any
+---@return avante.HistoryMessage[]
+local function generate_view_messages(tool_use, path, stale_view)
+  local view_result, view_error
+  if stale_view then
+    view_result = stale_view_content(path)
+  else
+    view_result, view_error = require("avante.llm_tools.view").func({ path = path }, {})
   end
 
-  for idx, message in ipairs(messages) do
-    table.insert(history_messages, message)
-    if Helpers.is_tool_result_message(message) then
-      local tool_use_message = Helpers.get_tool_use_message(message, messages)
-      local is_edit_func_call, is_str_replace_editor_func_call, is_str_replace_based_edit_tool_func_call, path =
-        Helpers.is_edit_func_call_message(tool_use_message)
-      --- For models like gpt-4o, the input parameter of replace_in_file is treated as the latest file content, so here we need to insert a fake view tool call to ensure it uses the latest file content
-      if is_edit_func_call and path and not message.message.content[1].is_error then
-        local uniformed_path = Utils.uniform_path(path)
-        local view_result, view_error = require("avante.llm_tools.view").func({ path = path }, {})
   if view_error then view_result = "Error: " .. view_error end
-        local get_diagnostics_tool_use_id = Utils.uuid()
+
   local view_tool_use_id = Utils.uuid()
   local view_tool_name = "view"
   local view_tool_input = { path = path }
-        if is_str_replace_editor_func_call then
+
+  if tool_use.name == "str_replace_editor" and tool_use.input.command == "str_replace" then
     view_tool_name = "str_replace_editor"
-          view_tool_input = { command = "view", path = path }
+    view_tool_input.command = "view"
-        end
+  elseif tool_use.name == "str_replace_based_edit_tool" and tool_use.input.command == "str_replace" then
-        if is_str_replace_based_edit_tool_func_call then
     view_tool_name = "str_replace_based_edit_tool"
-          view_tool_input = { command = "view", path = path }
+    view_tool_input.command = "view"
   end
-        history_messages = vim.list_extend(history_messages, {
+
+  return {
     Message:new_assistant_synthetic(string.format("Viewing file %s to get the latest content", path)),
     Message:new_assistant_synthetic({
       type = "tool_use",
@@ -111,15 +193,18 @@ M.update_history_messages = function(messages, using_ReAct_prompt, add_diagnosti
       is_error = view_error ~= nil,
       is_user_declined = false,
     }),
-        })
+  }
-        if last_modified_files[uniformed_path] == idx and add_diagnostic then
+end
+
+---Generates "diagnostic" for a file after it has been edited to help catching errors
+---@param path string
+---@return avante.HistoryMessage[]
+local function generate_diagnostic_messages(path)
+  local get_diagnostics_tool_use_id = Utils.uuid()
   local diagnostics = Utils.lsp.get_diagnostics_from_filepath(path)
-          history_messages = vim.list_extend(history_messages, {
+  return {
     Message:new_assistant_synthetic(
-              string.format(
+      string.format("The file %s has been modified, let me check if there are any errors in the changes.", path)
-                "The file %s has been modified, let me check if there are any errors in the changes.",
-                path
-              )
     ),
     Message:new_assistant_synthetic({
       type = "tool_use",
@@ -134,126 +219,108 @@ M.update_history_messages = function(messages, using_ReAct_prompt, add_diagnosti
       is_error = false,
       is_user_declined = false,
     }),
-          })
+  }
-        end
+end
-      end
+
-    end
+---Iterate through history messages and generate a new list containing updated history
-  end
+---that has up-to-date file contents and potentially updated diagnostic for modified
-  for _, message in ipairs(history_messages) do
+---files.
-    local content = message.message.content
+---@param messages avante.HistoryMessage[]
-    if type(content) ~= "table" then goto continue end
+---@param tools HistoryToolInfo[]
-    for _, item in ipairs(content) do
+---@param files HistoryFileInfo[]
-      if type(item) ~= "table" then goto continue1 end
+---@param add_diagnostic boolean Whether to generate and add diagnostic info to "edit" invocations
-      if item.type ~= "tool_use" then goto continue1 end
+---@param tools_to_text integer Number of tool invocations to be converted to simple text
-      local tool_name = item.name
+---@return avante.HistoryMessage[]
-      if tool_name ~= "view" then goto continue1 end
+local function refresh_history(messages, tools, files, add_diagnostic, tools_to_text)
-      local path = item.input.path
+  ---@type avante.HistoryMessage[]
-      tool_id_to_tool_name[item.id] = tool_name
+  local updated_messages = {}
-      if path then
+  local tool_count = 0
-        local uniform_path = Utils.uniform_path(path)
+
-        tool_id_to_path[item.id] = uniform_path
+  for _, message in ipairs(messages) do
-        tool_id_to_start_line[item.id] = item.input.start_line
+    local use = Helpers.get_tool_use_data(message)
-        tool_id_to_end_line[item.id] = item.input.end_line
+    if use then
-        viewed_files[uniform_path] = item.id
+      -- This is a tool invocation message. We will be handling both use and result together.
-      end
+      local tool_info = tools[use.id]
-      ::continue1::
+      if not tool_info.result then goto continue end
-    end
+
-    ::continue::
+      if tool_count < tools_to_text then
-  end
+        local text_msgs = convert_tool_to_text(tool_info)
-  for _, message in ipairs(history_messages) do
+        Utils.debug("Converted", use.name, "invocation to", #text_msgs, "messages")
-    local content = message.message.content
+        updated_messages = vim.list_extend(updated_messages, text_msgs)
-    if type(content) == "table" then
-      for _, item in ipairs(content) do
-        if type(item) ~= "table" then goto continue end
-        if item.type ~= "tool_result" then goto continue end
-        local tool_name = tool_id_to_tool_name[item.tool_use_id]
-        if tool_name ~= "view" then goto continue end
-        if item.is_error then goto continue end
-        local path = tool_id_to_path[item.tool_use_id]
-        local latest_tool_id = viewed_files[path]
-        if not latest_tool_id then goto continue end
-        if latest_tool_id ~= item.tool_use_id then
-          item.content = string.format("The file %s has been updated. Please use the latest `view` tool result!", path)
       else
-          local start_line = tool_id_to_start_line[item.tool_use_id]
+        table.insert(updated_messages, message)
-          local end_line = tool_id_to_end_line[item.tool_use_id]
+        table.insert(updated_messages, tool_info.result_message)
-          local view_result, view_error = require("avante.llm_tools.view").func(
+        tool_count = tool_count + 1
-            { path = path, start_line = start_line, end_line = end_line },
+
-            {}
+        if tool_info.kind == "view" then
-          )
+          local path = tool_info.path
-          if view_error then view_result = "Error: " .. view_error end
+          assert(path, "encountered 'view' tool invocation without path")
-          item.content = view_result
+          update_view_result(tool_info, use.id ~= files[tool_info.path].last_tool_id)
-          item.is_error = view_error ~= nil
-        end
-        ::continue::
-      end
         end
       end
 
-  if not using_ReAct_prompt then
+      if tool_info.kind == "edit" then
-    local picked_messages = {}
+        local path = tool_info.path
-    local max_tool_use_count = 25
+        assert(path, "encountered 'edit' tool invocation without path")
-    local tool_use_count = 0
+        local file_info = files[path]
-    for idx = #history_messages, 1, -1 do
+
-      local msg = history_messages[idx]
+        -- If this is the last operation for this file, generate synthetic "view"
-      if tool_use_count > max_tool_use_count then
+        -- invocation to provide the up-to-date file contents.
-        if Helpers.is_tool_result_message(msg) then
+        if not tool_info.result.is_error then
-          local tool_use_message = Helpers.get_tool_use_message(msg, history_messages)
+          local view_msgs = generate_view_messages(use, path, use.id == file_info.last_tool_id)
-          if tool_use_message then
+          Utils.debug("Added", #view_msgs, "'view' tool messages for", path)
-            table.insert(
+          updated_messages = vim.list_extend(updated_messages, view_msgs)
-              picked_messages,
+          tool_count = tool_count + 1
-              1,
-              Message:new_user_synthetic({
-                type = "text",
-                text = string.format(
-                  "Tool use [%s] is successful: %s",
-                  tool_use_message.message.content[1].name,
-                  tostring(not msg.message.content[1].is_error)
-                ),
-              })
-            )
-            table.insert(
-              picked_messages,
-              1,
-              Message:new_assistant_synthetic(
-                string.format(
-                  "Tool use %s(%s)",
-                  tool_use_message.message.content[1].name,
-                  vim.json.encode(tool_use_message.message.content[1].input)
-                )
-              )
-            )
         end
-        elseif Helpers.is_tool_use_message(msg) then
+
-          tool_use_count = tool_use_count + 1
+        if add_diagnostic and use.id == file_info.edit_tool_id then
-          goto continue
+          local diag_msgs = generate_diagnostic_messages(path)
-        else
+          Utils.debug("Added", #diag_msgs, "'diagnostics' tool messages for", path)
-          table.insert(picked_messages, 1, msg)
+          updated_messages = vim.list_extend(updated_messages, diag_msgs)
+          tool_count = tool_count + 1
         end
-      else
-        if Helpers.is_tool_use_message(msg) then tool_use_count = tool_use_count + 1 end
-        table.insert(picked_messages, 1, msg)
       end
+    elseif not Helpers.get_tool_result_data(message) then
+      -- Skip the tool result messages (since we process them together with their "use"s.
+      -- All other (non-tool-related) messages we simply keep.
+      table.insert(updated_messages, message)
+    end
+
     ::continue::
   end
 
-    history_messages = picked_messages
+  return updated_messages
+end
+
+---Analyzes the history looking for tool invocations, drops incomplete invocations,
+---and updates complete ones with the latest data available.
+---@param messages avante.HistoryMessage[]
+---@param max_tool_use integer | nil Maximum number of tool invocations to keep
+---@param add_diagnostic boolean Mix in LSP diagnostic info for affected files
+---@return avante.HistoryMessage[]
+M.update_tool_invocation_history = function(messages, max_tool_use, add_diagnostic)
+  local tools, files = collect_tool_info(messages)
+
+  -- Figure number of tool invocations that should be converted to simple "text"
+  -- messages to reduce prompt costs.
+  local tools_to_text = 0
+  if max_tool_use then
+    local n_edits = vim.iter(files):fold(
+      0,
+      ---@param count integer
+      ---@param file_info HistoryFileInfo
+      function(count, file_info)
+        if file_info.edit_tool_id then count = count + 1 end
+        return count
+      end
+    )
+    -- Each valid "edit" invocation will result in synthetic "view" and also
+    -- in "diagnostic" if it is requested by the caller.
+    local expected = #tools + n_edits + (add_diagnostic and n_edits or 0)
+    tools_to_text = expected - max_tool_use
   end
 
-  local final_history_messages = {}
+  return refresh_history(messages, tools, files, add_diagnostic, tools_to_text)
-  for _, msg in ipairs(history_messages) do
-    local tool_result_message
-    if Helpers.is_tool_use_message(msg) then
-      tool_result_message = Helpers.get_tool_result_message(msg, history_messages)
-      if not tool_result_message then goto continue end
-    end
-    if Helpers.is_tool_result_message(msg) then goto continue end
-    table.insert(final_history_messages, msg)
-    if tool_result_message then table.insert(final_history_messages, tool_result_message) end
-    ::continue::
-  end
-
-  return final_history_messages
 end
 
 ---Scans message history backwards, looking for tool invocations that have not been executed yet

lua/avante/llm.lua

@@ -824,7 +824,7 @@ function M._stream(opts)
           table.insert(tool_results, tool_result)
           return handle_next_tool_use(tool_uses, tool_use_messages, tool_use_index + 1, tool_results)
         end
-        local is_edit_tool_use = Utils.is_edit_func_call_tool_use(partial_tool_use)
+        local is_edit_tool_use = Utils.is_edit_tool_use(partial_tool_use)
         local support_streaming = false
         local llm_tool = vim.iter(prompt_opts.tools):find(function(tool) return tool.name == partial_tool_use.name end)
         if llm_tool then support_streaming = llm_tool.support_streaming == true end

lua/avante/sidebar.lua

@@ -2255,35 +2255,50 @@ end
 ---@return avante.HistoryMessage[]
 function Sidebar:get_history_messages_for_api(opts)
   opts = opts or {}
-  local history_messages0 = History.get_history_messages(self.chat_history)
+  local messages = History.get_history_messages(self.chat_history)
 
-  history_messages0 = vim
+  -- Scan the initial set of messages, filtering out "uninteresting" ones, but also
-    .iter(history_messages0)
+  -- check if the last message mentioned in the chat memory is actually present.
-    :filter(function(message) return not message.just_for_display and not message.is_compacted end)
+  local last_message = self.chat_history.memory and self.chat_history.memory.last_message_uuid
+  local last_message_present = false
+  messages = vim
+    .iter(messages)
+    :filter(function(message)
+      if message.just_for_display or message.is_compacted then return false end
+      if not opts.all then
+        if message.state == "generating" then return false end
+        if last_message and message.uuid == last_message then last_message_present = true end
+      end
+      return true
+    end)
     :totable()
 
-  if opts.all then return history_messages0 end
+  if not opts.all then
-
+    if last_message and last_message_present then
-  history_messages0 = vim
+      -- Drop all old messages preceding the "last" one from the memory
-    .iter(history_messages0)
+      local last_message_seen = false
-    :filter(function(message) return message.state ~= "generating" end)
+      messages = vim
+        .iter(messages)
+        :filter(function(message)
+          if not last_message_seen then
+            if message.uuid == last_message then last_message_seen = true end
+            return false
+          end
+          return true
+        end)
         :totable()
-
-  if self.chat_history and self.chat_history.memory then
-    local picked_messages = {}
-    for idx = #history_messages0, 1, -1 do
-      local message = history_messages0[idx]
-      if message.uuid == self.chat_history.memory.last_message_uuid then break end
-      table.insert(picked_messages, 1, message)
-    end
-    history_messages0 = picked_messages
     end
 
-  return History.update_history_messages(
+    local tool_limit
-    history_messages0,
+    if Providers[Config.provider].use_ReAct_prompt then
-    Providers[Config.provider].use_ReAct_prompt ~= nil,
+      tool_limit = nil
-    Config.behaviour.auto_check_diagnostics
+    else
-  )
+      tool_limit = 25
+    end
+    messages = History.update_tool_invocation_history(messages, tool_limit, Config.behaviour.auto_check_diagnostics)
+  end
+
+  return messages
 end
 
 ---@param request string

lua/avante/utils/init.lua

@@ -1496,30 +1496,10 @@ function M.tool_use_to_xml(tool_use)
 end
 
 ---@param tool_use AvanteLLMToolUse
-function M.is_edit_func_call_tool_use(tool_use)
+function M.is_edit_tool_use(tool_use)
-  local is_replace_func_call = false
+  return tool_use.name == "replace_in_file"
-  local is_str_replace_editor_func_call = false
+    or (tool_use.name == "str_replace_editor" and tool_use.input.command == "str_replace")
-  local is_str_replace_based_edit_tool_func_call = false
+    or (tool_use.name == "str_replace_based_edit_tool" and tool_use.input.command == "str_replace")
-  local path = nil
-  if tool_use.name == "replace_in_file" then
-    is_replace_func_call = true
-    path = tool_use.input.path
-  end
-  if tool_use.name == "str_replace_editor" then
-    if tool_use.input.command == "str_replace" then
-      is_replace_func_call = true
-      is_str_replace_editor_func_call = true
-      path = tool_use.input.path
-    end
-  end
-  if tool_use.name == "str_replace_based_edit_tool" then
-    if tool_use.input.command == "str_replace" then
-      is_replace_func_call = true
-      is_str_replace_based_edit_tool_func_call = true
-      path = tool_use.input.path
-    end
-  end
-  return is_replace_func_call, is_str_replace_editor_func_call, is_str_replace_based_edit_tool_func_call, path
 end
 
 ---Counts number of strings in text, accounting for possibility of a trailing newline