From 5abe57901908a9b53c98a8a02f60bee29fb578f2 Mon Sep 17 00:00:00 2001
From: yetone <yetoneful@gmail.com>
Date: Thu, 27 Feb 2025 12:52:58 +0800
Subject: [PATCH] fix: retry when rate limited (#1417)

---
 lua/avante/llm.lua              | 22 ++++++++++++++++++++++
 lua/avante/providers/claude.lua |  1 +
 lua/avante/types.lua            |  3 ++-
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua
index 3ca6ff4..c9a79ab 100644
--- a/lua/avante/llm.lua
+++ b/lua/avante/llm.lua
@@ -198,6 +198,13 @@ function M._stream(opts)
         end
         return handle_next_tool_use(sorted_tool_use_list, 1, old_tool_histories)
       end
+      if stop_opts.reason == "rate_limit" then
+        local msg = "Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds ..."
+        opts.on_chunk("\n*[" .. msg .. "]*\n")
+        Utils.info("Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds", { title = "Avante" })
+        vim.defer_fn(function() M._stream(opts) end, stop_opts.retry_after * 1000)
+        return
+      end
       return opts.on_stop(stop_opts)
     end,
   }
@@ -301,6 +308,21 @@ function M._stream(opts)
         else
           Utils.error("API request failed with status " .. result.status, { once = true, title = "Avante" })
         end
+        if result.status == 429 then
+          local headers_map = vim.iter(result.headers):fold({}, function(acc, value)
+            local pieces = vim.split(value, ":")
+            local key = pieces[1]
+            local remain = vim.list_slice(pieces, 2)
+            if not remain then return acc end
+            local val = Utils.trim_spaces(table.concat(remain, ":"))
+            acc[key] = val
+            return acc
+          end)
+          local retry_after = 10
+          if headers_map["retry-after"] then retry_after = tonumber(headers_map["retry-after"]) or 10 end
+          handler_opts.on_stop({ reason = "rate_limit", retry_after = retry_after })
+          return
+        end
         vim.schedule(function()
           if not completed then
             completed = true
diff --git a/lua/avante/providers/claude.lua b/lua/avante/providers/claude.lua
index 7eb132f..a862fc3 100644
--- a/lua/avante/providers/claude.lua
+++ b/lua/avante/providers/claude.lua
@@ -257,6 +257,7 @@ function M.parse_curl_args(provider, prompt_opts)
 end
 
 function M.on_error(result)
+  if result.status == 429 then return end
   if not result.body then
     return Utils.error("API request failed with status " .. result.status, { once = true, title = "Avante" })
   end
diff --git a/lua/avante/types.lua b/lua/avante/types.lua
index e8dea6b..2ed5d26 100644
--- a/lua/avante/types.lua
+++ b/lua/avante/types.lua
@@ -231,10 +231,11 @@ vim.g.avante_login = vim.g.avante_login
 ---@field usage? AvanteLLMUsage
 ---
 ---@class AvanteLLMStopCallbackOptions
----@field reason "complete" | "tool_use" | "error"
+---@field reason "complete" | "tool_use" | "error" | "rate_limit"
 ---@field error? string | table
 ---@field usage? AvanteLLMUsage
 ---@field tool_use_list? AvanteLLMToolUse[]
+---@field retry_after? integer
 ---
 ---@alias AvanteStreamParser fun(line: string, handler_opts: AvanteHandlerOptions): nil
 ---@alias AvanteLLMStartCallback fun(opts: AvanteLLMStartCallbackOptions): nil