diff --git a/lua/parrot/chat_handler.lua b/lua/parrot/chat_handler.lua index 1d3d8b0..2d1d689 100644 --- a/lua/parrot/chat_handler.lua +++ b/lua/parrot/chat_handler.lua @@ -844,13 +844,20 @@ function ChatHandler:_chat_respond(params) -- determine chat params or fallback to {} local chat_cfg = self.providers[query_prov.name] or {} local chat_params = (chat_cfg.params and chat_cfg.params.chat) or {} + local response_handler = ResponseHandler:new( + self.queries, + buf, + win, + utils.last_content_line(buf), + true, + "", + not self.options.chat_free_cursor + ) self:query( buf, query_prov, utils.prepare_payload(messages, model_obj.name, chat_params), - ResponseHandler - :new(self.queries, buf, win, utils.last_content_line(buf), true, "", not self.options.chat_free_cursor) - :create_handler(), + response_handler:create_handler(), vim.schedule_wrap(function(qid) if self.options.enable_spinner and spinner then spinner:stop() @@ -918,6 +925,8 @@ function ChatHandler:_chat_respond(params) if self.options.enable_spinner and topic_spinner then topic_spinner:stop() end + -- Cleanup topic response handler + topic_resp_handler:cleanup() -- get topic from invisible buffer local topic = vim.api.nvim_buf_get_lines(topic_buf, 0, -1, false)[1] -- close invisible buffer @@ -944,6 +953,11 @@ function ChatHandler:_chat_respond(params) local line = vim.api.nvim_buf_line_count(buf) utils.cursor_to_line(line, buf, win) end + + -- Ensure final update is flushed and cleanup response handler + response_handler:flush_updates(qid) + response_handler:cleanup() + vim.cmd("doautocmd User PrtDone") end) ) @@ -1397,7 +1411,9 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h logger.debug("LAST COMMAND in use " .. self.history.last_command) command = self.history.last_command end - -- dummy handler + -- response handler for managing streaming updates + local response_handler = nil + -- dummy handler (will be replaced by response_handler:create_handler()) local handler = function() end -- default on_exit strips trailing backticks if response was markdown snippet local on_exit = function(qid) @@ -1534,21 +1550,24 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h -- delete selection vim.api.nvim_buf_set_lines(buf, start_line - 1, end_line - 1, false, {}) -- prepare handler - handler = ResponseHandler:new(self.queries, buf, win, start_line - 1, true, prefix, cursor):create_handler() + response_handler = ResponseHandler:new(self.queries, buf, win, start_line - 1, true, prefix, cursor) + handler = response_handler:create_handler() elseif target == ui.Target.append then -- move cursor to the end of the selection vim.api.nvim_win_set_cursor(0, { end_line, 0 }) -- put newline after selection vim.api.nvim_put({ "" }, "l", true, true) -- prepare handler - handler = ResponseHandler:new(self.queries, buf, win, end_line, true, prefix, cursor):create_handler() + response_handler = ResponseHandler:new(self.queries, buf, win, end_line, true, prefix, cursor) + handler = response_handler:create_handler() elseif target == ui.Target.prepend then -- move cursor to the start of the selection vim.api.nvim_win_set_cursor(0, { start_line, 0 }) -- put newline before selection vim.api.nvim_put({ "" }, "l", false, true) -- prepare handler - handler = ResponseHandler:new(self.queries, buf, win, start_line - 1, true, prefix, cursor):create_handler() + response_handler = ResponseHandler:new(self.queries, buf, win, start_line - 1, true, prefix, cursor) + handler = response_handler:create_handler() elseif target == ui.Target.popup then self:toggle_close(self._toggle_kind.popup) -- create a new buffer @@ -1576,7 +1595,8 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h -- better text wrapping vim.api.nvim_command("setlocal wrap linebreak") -- prepare handler - handler = ResponseHandler:new(self.queries, buf, win, 0, false, "", false):create_handler() + response_handler = ResponseHandler:new(self.queries, buf, win, 0, false, "", false) + handler = response_handler:create_handler() self:toggle_add(self._toggle_kind.popup, { win = win, buf = buf, close = popup_close }) elseif type(target) == "table" then if target.type == ui.Target.new().type then @@ -1607,7 +1627,8 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h vim.api.nvim_set_option_value("filetype", "markdown", { buf = buf }) - handler = ResponseHandler:new(self.queries, buf, win, 0, false, "", cursor):create_handler() + response_handler = ResponseHandler:new(self.queries, buf, win, 0, false, "", cursor) + handler = response_handler:create_handler() end -- call the model and write the response @@ -1637,6 +1658,13 @@ function ChatHandler:prompt(params, target, model_obj, prompt, template, reset_h spinner:stop() end on_exit(qid) + + -- Ensure final update is flushed and cleanup response handler + if response_handler then + response_handler:flush_updates(qid) + response_handler:cleanup() + end + vim.cmd("doautocmd User PrtDone") end) ) diff --git a/lua/parrot/response_handler.lua b/lua/parrot/response_handler.lua index b29de2a..f919eb2 100644 --- a/lua/parrot/response_handler.lua +++ b/lua/parrot/response_handler.lua @@ -12,9 +12,18 @@ local logger = require("parrot.logger") ---@field prefix string ---@field cursor boolean ---@field hl_handler_group string +---@field chunk_buffer string +---@field update_timer any +---@field pending_chunks boolean +---@field last_update_time number local ResponseHandler = {} ResponseHandler.__index = ResponseHandler +-- Configuration for buffering and debouncing +local BUFFER_TIMEOUT_MS = 16 -- ~60fps +local MIN_CHUNK_SIZE = 10 -- minimum characters to trigger immediate update +local MAX_BUFFER_SIZE = 1000 -- maximum characters to buffer before forced update + ---Creates a new ResponseHandler ---@param queries table ---@param buffer number|nil @@ -36,6 +45,12 @@ function ResponseHandler:new(queries, buffer, window, line, first_undojoin, pref self.queries = queries self.skip_first_undojoin = not first_undojoin + -- Buffering and debouncing fields + self.chunk_buffer = "" + self.update_timer = nil + self.pending_chunks = false + self.last_update_time = 0 + self.hl_handler_group = "PrtHandlerStandout" vim.api.nvim_set_hl(0, self.hl_handler_group, { link = "CursorLine" }) @@ -48,7 +63,7 @@ function ResponseHandler:new(queries, buffer, window, line, first_undojoin, pref return self end ----Handles a chunk of response +---Handles a chunk of response with buffering and debouncing ---@param qid any ---@param chunk string function ResponseHandler:handle_chunk(qid, chunk) @@ -57,16 +72,111 @@ function ResponseHandler:handle_chunk(qid, chunk) return end + -- Add chunk to buffer + if chunk and chunk ~= "" then + self.chunk_buffer = self.chunk_buffer .. chunk + self.response = self.response .. chunk + self.pending_chunks = true + + qt.ns_id = qt.ns_id or self.ns_id + qt.ex_id = qt.ex_id or self.ex_id + qt.response = self.response + end + + -- Determine if we should update immediately or wait + local should_update_immediately = false + local current_time = vim.loop.hrtime() / 1000000 -- Convert to milliseconds + + -- Force update if chunk buffer is too large + if #self.chunk_buffer >= MAX_BUFFER_SIZE then + should_update_immediately = true + end + + -- Force update if chunk contains newlines (likely end of sentence/paragraph) + if chunk and chunk:find("\n") then + should_update_immediately = true + end + + -- Force update if chunk is large enough + if chunk and #chunk >= MIN_CHUNK_SIZE then + should_update_immediately = true + end + + if should_update_immediately then + self:flush_updates(qid) + else + -- Schedule a debounced update + self:schedule_update(qid) + end +end + +---Schedule a debounced update +---@param qid any +function ResponseHandler:schedule_update(qid) + -- Cancel existing timer if it exists + if self.update_timer then + self.update_timer:stop() + self.update_timer:close() + end + + -- Schedule new update + self.update_timer = vim.loop.new_timer() + self.update_timer:start( + BUFFER_TIMEOUT_MS, + 0, + vim.schedule_wrap(function() + if self.pending_chunks then + self:flush_updates(qid) + end + if self.update_timer then + self.update_timer:close() + self.update_timer = nil + end + end) + ) +end + +---Flush all pending updates to the buffer +---@param qid any +function ResponseHandler:flush_updates(qid) + if not self.pending_chunks or not vim.api.nvim_buf_is_valid(self.buffer) then + return + end + + local qt = self.queries:get(qid) + if not qt then + return + end + + -- Cancel any pending timer + if self.update_timer then + self.update_timer:stop() + self.update_timer:close() + self.update_timer = nil + end + + -- Perform batch update if not self.skip_first_undojoin then utils.undojoin(self.buffer) end self.skip_first_undojoin = false - qt.ns_id = qt.ns_id or self.ns_id - qt.ex_id = qt.ex_id or self.ex_id - - self.first_line = vim.api.nvim_buf_get_extmark_by_id(self.buffer, self.ns_id, self.ex_id, {})[1] + -- Safely get extmark position with fallback + local extmark_pos = vim.api.nvim_buf_get_extmark_by_id(self.buffer, self.ns_id, self.ex_id, {}) + if extmark_pos and #extmark_pos > 0 then + self.first_line = extmark_pos[1] + elseif not self.first_line then + -- Fallback to current cursor position if extmark is lost + if self.window and vim.api.nvim_win_is_valid(self.window) then + local cursor_pos = vim.api.nvim_win_get_cursor(self.window) + self.first_line = math.max(0, cursor_pos[1] - 1) + else + -- Ultimate fallback - use end of buffer + self.first_line = vim.api.nvim_buf_line_count(self.buffer) + end + end + -- Clear previous response lines to avoid duplication local line_count = #vim.split(self.response, "\n") vim.api.nvim_buf_set_lines( self.buffer, @@ -76,28 +186,24 @@ function ResponseHandler:handle_chunk(qid, chunk) {} ) - self:update_response(chunk) self:update_buffer() self:update_highlighting(qt) self:update_query_object(qt) self:move_cursor() -end ----Updates the response with a new chunk ----@param chunk string -function ResponseHandler:update_response(chunk) - if chunk ~= nil then - self.response = self.response .. chunk - logger.debug("ResponseHandler:update_response", { - response = self.response, - chunk = chunk, - }) - utils.undojoin(self.buffer) - end + -- Reset buffer state + self.chunk_buffer = "" + self.pending_chunks = false + self.last_update_time = vim.loop.hrtime() / 1000000 end ---Updates the buffer with the current response function ResponseHandler:update_buffer() + -- Safety check for first_line + if not self.first_line then + return + end + local lines = vim.split(self.response, "\n") local prefixed_lines = vim.tbl_map(function(l) return self.prefix .. l @@ -115,20 +221,44 @@ function ResponseHandler:update_buffer() ) end ----Updates the highlighting for new lines +---Updates the highlighting for new lines (batch operation) ---@param qt table function ResponseHandler:update_highlighting(qt) + -- Safety check for first_line + if not self.first_line then + return + end + local lines = vim.split(self.response, "\n") local new_finished_lines = math.max(0, #lines - 1) - for i = self.finished_lines, new_finished_lines do - vim.api.nvim_buf_add_highlight(self.buffer, qt.ns_id, self.hl_handler_group, self.first_line + i, 0, -1) + + -- Batch highlight updates to reduce flicker + if new_finished_lines > self.finished_lines then + -- Clear existing highlights in the range to avoid duplicates + vim.api.nvim_buf_clear_namespace( + self.buffer, + qt.ns_id, + self.first_line + self.finished_lines, + self.first_line + new_finished_lines + 1 + ) + + -- Add highlights for the new range + for i = self.finished_lines, new_finished_lines do + vim.api.nvim_buf_add_highlight(self.buffer, qt.ns_id, self.hl_handler_group, self.first_line + i, 0, -1) + end end + self.finished_lines = new_finished_lines end ---Updates the query object with new line information ---@param qt table function ResponseHandler:update_query_object(qt) + -- Safety check for first_line + if not self.first_line then + return + end + local end_line = self.first_line + #vim.split(self.response, "\n") qt.first_line = self.first_line qt.last_line = end_line - 1 @@ -136,12 +266,21 @@ end ---Moves the cursor to the end of the response if needed function ResponseHandler:move_cursor() - if self.cursor then + if self.cursor and self.first_line then local end_line = self.first_line + #vim.split(self.response, "\n") utils.cursor_to_line(end_line, self.buffer, self.window) end end +---Cleanup method to ensure timers are properly closed +function ResponseHandler:cleanup() + if self.update_timer then + self.update_timer:stop() + self.update_timer:close() + self.update_timer = nil + end +end + ---Creates a handler function ---@return function function ResponseHandler:create_handler() diff --git a/tests/parrot/response_handler_spec.lua b/tests/parrot/response_handler_spec.lua index a4d195e..60f5483 100644 --- a/tests/parrot/response_handler_spec.lua +++ b/tests/parrot/response_handler_spec.lua @@ -16,8 +16,25 @@ describe("ResponseHandler", function() nvim_buf_set_lines = stub.new(), nvim_buf_add_highlight = stub.new(), nvim_win_get_cursor = stub.new().returns({ 1, 0 }), + nvim_set_hl = stub.new(), + nvim_buf_clear_namespace = stub.new(), + nvim_buf_line_count = stub.new().returns(10), + nvim_win_is_valid = stub.new().returns(true), }, split = stub.new().returns({ "test" }), + list_slice = stub.new().returns({ "test" }), + tbl_map = stub.new().returns({ "test" }), + loop = { + new_timer = stub.new().returns({ + start = stub.new(), + stop = stub.new(), + close = stub.new(), + }), + hrtime = stub.new().returns(1000000000), -- 1 second in nanoseconds + }, + schedule_wrap = function(fn) + return fn + end, -- Remove vim.cmd to avoid potential issues with Vim options } @@ -31,14 +48,21 @@ describe("ResponseHandler", function() get = stub.new().returns({ ns_id = 1, ex_id = 1 }), } + -- Mock logger module + local mock_logger = { + debug = stub.new(), + } + -- Use package.loaded instead of _G.vim package.loaded.vim = mock_vim package.loaded["parrot.utils"] = mock_utils + package.loaded["parrot.logger"] = mock_logger end) after_each(function() package.loaded.vim = nil package.loaded["parrot.utils"] = nil + package.loaded["parrot.logger"] = nil end) it("should create a new ResponseHandler with default values", function() @@ -80,13 +104,6 @@ describe("ResponseHandler", function() assert.are.same("", handler.response) end) - it("should update the response with a new chunk", function() - local handler = ResponseHandler:new(mock_queries) - handler:update_response("test chunk") - handler:update_response(" test chunk") - assert.are.same("test chunk test chunk", handler.response) - end) - it("should not move the cursor when cursor is false", function() local handler = ResponseHandler:new(mock_queries) handler.response = "line1\nline2" @@ -99,4 +116,159 @@ describe("ResponseHandler", function() local handler_func = handler:create_handler() assert.is_function(handler_func) end) + + it("should schedule updates with timer", function() + local handler = ResponseHandler:new(mock_queries) + local mock_timer = { + start = stub.new(), + stop = stub.new(), + close = stub.new(), + } + + -- Mock vim.loop.new_timer + local original_new_timer = vim.loop.new_timer + vim.loop.new_timer = stub.new().returns(mock_timer) + + handler:schedule_update(1) + + assert.stub(vim.loop.new_timer).was_called() + assert.stub(mock_timer.start).was_called() + + -- Restore original function + vim.loop.new_timer = original_new_timer + end) + + it("should flush updates to buffer", function() + local handler = ResponseHandler:new(mock_queries) + handler.response = "test response" + handler.pending_chunks = true + handler.first_line = 1 + handler.finished_lines = 0 + + -- Setup fresh mocks for this test + mock_queries.get.returns({ ns_id = 1, ex_id = 1, response = "" }) + mock_vim.split.returns({ "test response" }) + mock_vim.api.nvim_buf_get_extmark_by_id.returns({ 1, 0 }) + mock_vim.api.nvim_buf_is_valid.returns(true) + + handler:flush_updates(1) + + -- Test behavior rather than implementation + assert.are.same(false, handler.pending_chunks) + assert.are.same("", handler.chunk_buffer) + end) + + it("should update buffer with response lines", function() + local handler = ResponseHandler:new(mock_queries) + handler.response = "line1\nline2\nline3" + handler.first_line = 1 + handler.finished_lines = 0 + + -- Setup mocks for this test + mock_vim.split.returns({ "line1", "line2", "line3" }) + mock_vim.tbl_map.returns({ "line1", "line2", "line3" }) + mock_vim.list_slice.returns({ "line1", "line2", "line3" }) + + -- Call the method - we can't easily test the vim API calls due to mocking complexity + -- but we can test that the method doesn't crash + handler:update_buffer() + + -- Test that state is maintained + assert.are.same("line1\nline2\nline3", handler.response) + assert.are.same(1, handler.first_line) + end) + + it("should update highlighting for new lines", function() + local handler = ResponseHandler:new(mock_queries) + handler.response = "line1\nline2\nline3" + handler.first_line = 1 + handler.finished_lines = 0 + + -- Setup mocks for this test + mock_vim.split.returns({ "line1", "line2", "line3" }) + + local qt = { ns_id = 1 } + handler:update_highlighting(qt) + + -- Test that finished_lines is updated correctly + assert.are.same(2, handler.finished_lines) + end) + + it("should update query object with line information", function() + local handler = ResponseHandler:new(mock_queries) + handler.response = "line1\nline2\nline3" + handler.first_line = 1 + + -- Mock vim.split to return the lines + mock_vim.split.returns({ "line1", "line2", "line3" }) + + local qt = {} + handler:update_query_object(qt) + + assert.are.same(1, qt.first_line) + assert.are.same(3, qt.last_line) + end) + + it("should move cursor when cursor is true", function() + local handler = ResponseHandler:new(mock_queries, nil, 1, 1, true, "", true) + handler.response = "line1\nline2" + handler.first_line = 1 + + -- Setup mocks for this test + mock_vim.split.returns({ "line1", "line2" }) + + handler:move_cursor() + + -- Test that cursor property is maintained + assert.are.same(true, handler.cursor) + assert.are.same(1, handler.first_line) + end) + + it("should cleanup timers properly", function() + local handler = ResponseHandler:new(mock_queries) + local mock_timer = { + stop = stub.new(), + close = stub.new(), + } + handler.update_timer = mock_timer + + handler:cleanup() + + assert.stub(mock_timer.stop).was_called() + assert.stub(mock_timer.close).was_called() + assert.is_nil(handler.update_timer) + end) + + it("should not update buffer when first_line is nil", function() + local handler = ResponseHandler:new(mock_queries) + handler.first_line = nil + handler.response = "test" + + handler:update_buffer() + + assert.stub(mock_vim.api.nvim_buf_set_lines).was_not_called() + end) + + it("should not update highlighting when first_line is nil", function() + local handler = ResponseHandler:new(mock_queries) + handler.first_line = nil + handler.response = "test" + + local qt = { ns_id = 1 } + handler:update_highlighting(qt) + + assert.stub(mock_vim.api.nvim_buf_clear_namespace).was_not_called() + end) + + it("should not update query object when first_line is nil", function() + local handler = ResponseHandler:new(mock_queries) + handler.first_line = nil + handler.response = "test" + + local qt = {} + handler:update_query_object(qt) + + assert.is_nil(qt.first_line) + assert.is_nil(qt.last_line) + end) end)