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.
This commit is contained in:
Dmitry Torokhov
2025-07-17 22:11:20 -07:00
parent b4c9246461
commit c2e4ae5ef6

View File

@@ -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)