From c2e4ae5ef6b607080c98df41a744102563abc1f5 Mon Sep 17 00:00:00 2001
From: Dmitry Torokhov <dmitry.torokhov@gmail.com>
Date: Thu, 17 Jul 2025 22:11:20 -0700
Subject: [PATCH] fix(llm): fix rate limit handling

Rate limit handling seems to be broken: it starts an one-shot timer
with callback that reschedules the timer each second and prints/updates
messages. Simultaneously it schedules a deferred by 1 second function
that cancels the timer and resumes the stream. This results in the timer
executing at most once, and stream resume happening way too early.

Fix this by switching to use repeating timer with first instance
executing immediately. All message handling is moved into the timer
callback. Once countdown is complete the same callback will stop and
destroy the timer and resume the stream.
---
 lua/avante/llm.lua | 59 +++++++++++++++++++++++-----------------------
 1 file changed, 30 insertions(+), 29 deletions(-)

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)