diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua index e24aced..ef6a637 100644 --- a/lua/avante/llm.lua +++ b/lua/avante/llm.lua @@ -936,44 +936,45 @@ function M._stream(opts) return handle_next_tool_use(pending_tools, pending_tool_use_messages, 1, {}, stop_opts.streaming_tool_use) end if stop_opts.reason == "rate_limit" then - local msg_content = "*[Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds ...]*" - if opts.on_chunk then opts.on_chunk("\n" .. msg_content .. "\n") end - local message - if opts.on_messages_add then - message = History.Message:new({ + local message = opts.on_messages_add + and History.Message:new({ role = "assistant", - content = "\n\n" .. msg_content, + content = "", -- Actual content will be set below }, { just_for_display = true, }) - opts.on_messages_add({ message }) - end + local timer = vim.loop.new_timer() if timer then - local retry_after = stop_opts.retry_after + local retry_count = stop_opts.retry_after + Utils.info("Rate limit reached. Retrying in " .. retry_count .. " seconds", { title = "Avante" }) + local function countdown() - timer:start( - 1000, - 0, - vim.schedule_wrap(function() - if retry_after > 0 then retry_after = retry_after - 1 end - local msg_content_ = "*[Rate limit reached. Retrying in " .. retry_after .. " seconds ...]*" - if opts.on_chunk then opts.on_chunk([[\033[1A\033[K]] .. "\n" .. msg_content_ .. "\n") end - if opts.on_messages_add and message then - message.message.content = "\n\n" .. msg_content_ - opts.on_messages_add({ message }) - end - countdown() - end) - ) + local msg_content = "*[Rate limit reached. Retrying in " .. retry_count .. " seconds ...]*" + if opts.on_chunk then + -- Use ANSI escape codes to clear line and move cursor up only for subsequent updates + local prefix = "" + if retry_count < stop_opts.retry_after then prefix = [[\033[1A\033[K]] end + opts.on_chunk(prefix .. "\n" .. msg_content .. "\n") + end + if opts.on_messages_add and message then + message.message.content = "\n\n" .. msg_content + opts.on_messages_add({ message }) + end + + if retry_count <= 0 then + timer:stop() + timer:close() + + Utils.info("Restarting stream after rate limi pause") + M._stream(opts) + else + retry_count = retry_count - 1 + end end - countdown() + + timer:start(0, 1000, vim.schedule_wrap(function() countdown() end)) end - Utils.info("Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds", { title = "Avante" }) - vim.defer_fn(function() - if timer then timer:stop() end - M._stream(opts) - end, stop_opts.retry_after * 1000) return end return opts.on_stop(stop_opts)