From d0f0580d64c391de7b87160a9e66191677e02f54 Mon Sep 17 00:00:00 2001 From: Dmitry Torokhov Date: Wed, 6 Aug 2025 22:06:30 -0700 Subject: [PATCH] fix: sanitize loaded chat history Ensure that chat history loaded from a file has resemblance of correct data. Namely title and timestamp are present and are strings, and entires, messages, and todos are lists. In case of inconsistencies replace with empty/default data. This should help with #2584. More changes are needed to sanitize individual entries. --- lua/avante/llm_tools/update_todo_status.lua | 2 +- lua/avante/path.lua | 43 ++++++++++++++------- lua/avante/sidebar.lua | 10 ++--- lua/avante/types.lua | 6 +-- 4 files changed, 36 insertions(+), 25 deletions(-) diff --git a/lua/avante/llm_tools/update_todo_status.lua b/lua/avante/llm_tools/update_todo_status.lua index 6f3e39b..8f266d6 100644 --- a/lua/avante/llm_tools/update_todo_status.lua +++ b/lua/avante/llm_tools/update_todo_status.lua @@ -48,7 +48,7 @@ function M.func(input, opts) local sidebar = require("avante").get() if not sidebar then return false, "Avante sidebar not found" end local todos = sidebar.chat_history.todos - if not todos or #todos == 0 then return false, "No todos found" end + if #todos == 0 then return false, "No todos found" end for _, todo in ipairs(todos) do if tostring(todo.id) == tostring(input.id) then todo.status = input.status diff --git a/lua/avante/path.lua b/lua/avante/path.lua index 794d02e..ac2fcea 100644 --- a/lua/avante/path.lua +++ b/lua/avante/path.lua @@ -54,10 +54,8 @@ function History.list(bufnr) for _, filename in ipairs(files) do if not filename:match("metadata.json") then local filepath = Path:new(filename) - local content = filepath:read() - local history = vim.json.decode(content) - history.filename = filepath_to_filename(filepath) - table.insert(res, history) + local history = History.from_file(filepath) + if history then table.insert(res, history) end end end --- sort by timestamp @@ -134,12 +132,37 @@ function History.new(bufnr) local history = { title = "untitled", timestamp = Utils.get_timestamp(), + entries = {}, messages = {}, + todos = {}, filename = filepath_to_filename(filepath), } return history end +---Attempts to load chat history from a given file +---@param filepath Path +---@return avante.ChatHistory|nil +function History.from_file(filepath) + if filepath:exists() then + local content = filepath:read() + if content ~= nil then + local decode_ok, history = pcall(vim.json.decode, content) + if decode_ok and type(history) == "table" then + if not history.title or history.title ~= "string" then history.title = "untitled" end + if not history.timestamp or history.timestamp ~= "string" then history.timestamp = Utils.get_timestamp() end + -- TODO: sanitize individual entries of the lists below as well. + if not vim.islist(history.entries) then history.entries = {} end + if not vim.islist(history.messages) then history.messages = {} end + if not vim.islist(history.todos) then history.todos = {} end + ---@cast history avante.ChatHistory + history.filename = filepath_to_filename(filepath) + return history + end + end + end +end + -- Loads the chat history for the given buffer. ---@param bufnr integer ---@param filename string? @@ -147,21 +170,13 @@ end function History.load(bufnr, filename) local history_filepath = filename and History.get_filepath(bufnr, filename) or History.get_latest_filepath(bufnr, false) - if history_filepath:exists() then - local content = history_filepath:read() - if content ~= nil then - local history = vim.json.decode(content) - history.filename = filepath_to_filename(history_filepath) - return history - end - end - return History.new(bufnr) + return History.from_file(history_filepath) or History.new(bufnr) end -- Saves the chat history for the given buffer. ---@param bufnr integer ---@param history avante.ChatHistory -History.save = function(bufnr, history) +function History.save(bufnr, history) local history_filepath = History.get_filepath(bufnr, history.filename) history_filepath:write(vim.json.encode(history), "w") History.save_latest_filename(bufnr, history.filename) diff --git a/lua/avante/sidebar.lua b/lua/avante/sidebar.lua index 25abf86..1d0cd4e 100644 --- a/lua/avante/sidebar.lua +++ b/lua/avante/sidebar.lua @@ -2838,7 +2838,7 @@ function Sidebar:create_input_container() get_history_messages = function(opts) return self:get_history_messages_for_api(opts) end, get_todos = function() local history = Path.history.load(self.code.bufnr) - return history and history.todos or {} + return history.todos end, update_todos = function(todos) self:update_todos(todos) end, session_ctx = {}, @@ -3096,7 +3096,7 @@ end function Sidebar:get_todos_container_height() local history = Path.history.load(self.code.bufnr) - if not history or not history.todos or #history.todos == 0 then return 0 end + if #history.todos == 0 then return 0 end return 3 end @@ -3389,7 +3389,7 @@ end function Sidebar:create_todos_container() local history = Path.history.load(self.code.bufnr) - if not history or not history.todos or #history.todos == 0 then + if #history.todos == 0 then if self.containers.todos and Utils.is_valid_container(self.containers.todos) then self.containers.todos:unmount() end @@ -3438,10 +3438,6 @@ function Sidebar:create_todos_container() local total_count = #history.todos local focused_idx = 1 local todos_content_lines = {} - if type(history.todos) ~= "table" then - Utils.debug("Invalid todos type", history.todos) - history.todos = {} - end for idx, todo in ipairs(history.todos) do local status_content = "[ ]" if todo.status == "done" then diff --git a/lua/avante/types.lua b/lua/avante/types.lua index 1997844..c03a5af 100644 --- a/lua/avante/types.lua +++ b/lua/avante/types.lua @@ -503,9 +503,9 @@ vim.g.avante_login = vim.g.avante_login ---@class avante.ChatHistory ---@field title string ---@field timestamp string ----@field messages avante.HistoryMessage[] | nil ----@field entries avante.ChatHistoryEntry[] | nil ----@field todos avante.TODO[] | nil +---@field messages avante.HistoryMessage[] +---@field entries avante.ChatHistoryEntry[] +---@field todos avante.TODO[] ---@field memory avante.ChatMemory | nil ---@field filename string ---@field system_prompt string | nil