feat(auto-retry): preliminary support for auto retry on 429 (#721)

* feat(auto-retry): preliminary support for auto retry on 429

* style: simplification of async closure

* fix: bring back status check, default options

Author: georgeh0

src/llm/anthropic.rs

@@ -1,14 +1,12 @@
+use crate::prelude::*;
+use base64::prelude::*;
+
 use crate::llm::{
     LlmGenerateRequest, LlmGenerateResponse, LlmGenerationClient, OutputFormat,
     ToJsonSchemaOptions, detect_image_mime_type,
 };
-use anyhow::{Context, Result, bail};
-use async_trait::async_trait;
-use base64::prelude::*;
+use anyhow::Context;
 use json5;
-use serde_json::Value;
-
-use crate::api_bail;
 use urlencoding::encode;
 
 pub struct Client {
@@ -91,23 +89,26 @@ impl LlmGenerationClient for Client {
 
         let encoded_api_key = encode(&self.api_key);
 
-        let resp = self
-            .client
-            .post(url)
-            .header("x-api-key", encoded_api_key.as_ref())
-            .header("anthropic-version", "2023-06-01")
-            .json(&payload)
-            .send()
-            .await
-            .context("HTTP error")?;
+        let resp = retryable::run(
+            || {
+                self.client
+                    .post(url)
+                    .header("x-api-key", encoded_api_key.as_ref())
+                    .header("anthropic-version", "2023-06-01")
+                    .json(&payload)
+                    .send()
+            },
+            &retryable::HEAVY_LOADED_OPTIONS,
+        )
+        .await?;
         if !resp.status().is_success() {
             bail!(
                 "Anthropic API error: {:?}\n{}\n",
                 resp.status(),
                 resp.text().await?
             );
         }
-        let mut resp_json: Value = resp.json().await.context("Invalid JSON")?;
+        let mut resp_json: serde_json::Value = resp.json().await.context("Invalid JSON")?;
         if let Some(error) = resp_json.get("error") {
             bail!("Anthropic API error: {:?}", error);
         }
@@ -117,11 +118,11 @@ impl LlmGenerationClient for Client {
 
         let resp_content = &resp_json["content"];
         let tool_name = "report_result";
-        let mut extracted_json: Option<Value> = None;
+        let mut extracted_json: Option<serde_json::Value> = None;
         if let Some(array) = resp_content.as_array() {
             for item in array {
-                if item.get("type") == Some(&Value::String("tool_use".to_string()))
-                    && item.get("name") == Some(&Value::String(tool_name.to_string()))
+                if item.get("type") == Some(&serde_json::Value::String("tool_use".to_string()))
+                    && item.get("name") == Some(&serde_json::Value::String(tool_name.to_string()))
                 {
                     if let Some(input) = item.get("input") {
                         extracted_json = Some(input.clone());
@@ -136,7 +137,7 @@ impl LlmGenerationClient for Client {
         } else {
             // Fallback: try text if no tool output found
             match &mut resp_json["content"][0]["text"] {
-                Value::String(s) => {
+                serde_json::Value::String(s) => {
                     // Try strict JSON parsing first
                     match serde_json::from_str::<serde_json::Value>(s) {
                         Ok(_) => std::mem::take(s),

src/llm/gemini.rs

@@ -113,13 +113,11 @@ impl LlmGenerationClient for Client {
         }
 
         let url = self.get_api_url(request.model, "generateContent");
-        let resp = self
-            .client
-            .post(&url)
-            .json(&payload)
-            .send()
-            .await
-            .context("HTTP error")?;
+        let resp = retryable::run(
+            || self.client.post(&url).json(&payload).send(),
+            &retryable::HEAVY_LOADED_OPTIONS,
+        )
+        .await?;
         if !resp.status().is_success() {
             bail!(
                 "Gemini API error: {:?}\n{}\n",
@@ -174,13 +172,11 @@ impl LlmEmbeddingClient for Client {
         if let Some(task_type) = request.task_type {
             payload["taskType"] = serde_json::Value::String(task_type.into());
         }
-        let resp = self
-            .client
-            .post(&url)
-            .json(&payload)
-            .send()
-            .await
-            .context("HTTP error")?;
+        let resp = retryable::run(
+            || self.client.post(&url).json(&payload).send(),
+            &retryable::HEAVY_LOADED_OPTIONS,
+        )
+        .await?;
         if !resp.status().is_success() {
             bail!(
                 "Gemini API error: {:?}\n{}\n",

src/llm/ollama.rs

@@ -65,12 +65,16 @@ impl LlmGenerationClient for Client {
             system: request.system_prompt.as_ref().map(|s| s.as_ref()),
             stream: Some(false),
         };
-        let res = self
-            .reqwest_client
-            .post(self.generate_url.as_str())
-            .json(&req)
-            .send()
-            .await?;
+        let res = retryable::run(
+            || {
+                self.reqwest_client
+                    .post(self.generate_url.as_str())
+                    .json(&req)
+                    .send()
+            },
+            &retryable::HEAVY_LOADED_OPTIONS,
+        )
+        .await?;
         if !res.status().is_success() {
             bail!(
                 "Ollama API error: {:?}\n{}\n",

src/llm/voyage.rs

@@ -75,14 +75,17 @@ impl LlmEmbeddingClient for Client {
             payload["input_type"] = serde_json::Value::String(task_type.into());
         }
 
-        let resp = self
-            .client
-            .post(url)
-            .header("Authorization", format!("Bearer {}", self.api_key))
-            .json(&payload)
-            .send()
-            .await
-            .context("HTTP error")?;
+        let resp = retryable::run(
+            || {
+                self.client
+                    .post(url)
+                    .header("Authorization", format!("Bearer {}", self.api_key))
+                    .json(&payload)
+                    .send()
+            },
+            &retryable::HEAVY_LOADED_OPTIONS,
+        )
+        .await?;
 
         if !resp.status().is_success() {
             bail!(

src/utils/retryable.rs

@@ -28,6 +28,12 @@ impl IsRetryable for Error {
     }
 }
 
+impl IsRetryable for reqwest::Error {
+    fn is_retryable(&self) -> bool {
+        self.status() == Some(reqwest::StatusCode::TOO_MANY_REQUESTS)
+    }
+}
+
 impl Error {
     pub fn always_retryable(error: anyhow::Error) -> Self {
         Self {
@@ -77,13 +83,19 @@ pub struct RetryOptions {
 impl Default for RetryOptions {
     fn default() -> Self {
         Self {
-            max_retries: Some(5),
+            max_retries: Some(10),
             initial_backoff: Duration::from_millis(100),
             max_backoff: Duration::from_secs(10),
         }
     }
 }
 
+pub static HEAVY_LOADED_OPTIONS: RetryOptions = RetryOptions {
+    max_retries: Some(10),
+    initial_backoff: Duration::from_secs(1),
+    max_backoff: Duration::from_secs(60),
+};
+
 pub async fn run<
     Ok,
     Err: std::fmt::Display + IsRetryable,