refactor: implement retry logic with exponential backoff and jitter for subagent and provider calls
This commit is contained in:
asepharyana
@@ -1,7 +1,31 @@
|
||||
//! Blocking HTTP client for OpenAI/Anthropic-compatible chat completion APIs,
|
||||
//! supporting both non-streaming and SSE-streaming requests with automatic retry.
|
||||
//!
|
||||
//! # Retry policy
|
||||
//!
|
||||
//! Both paths use exponential backoff with ±25% jitter so retries spread out
|
||||
//! naturally instead of hammering the server in lockstep. Auth errors
|
||||
//! (401/402/403) are never retried — they indicate a bad key or billing issue
|
||||
//! that retrying won't fix. Rate-limit (429) responses get a longer backoff
|
||||
//! (base 5s instead of the usual 1s) so the server has time to drain its queue.
|
||||
//!
|
||||
//! ## Non-streaming (`chat_with_tools_non_streaming`)
|
||||
//! - Up to **10** attempts
|
||||
//! - Backoff: `1s, 2s, 4s, 8s, 16s, 30s(capped), 30s, …` + jitter
|
||||
//! - Auth errors → abort immediately on the **status code** embedded in the
|
||||
//! error message (avoids false positives from port numbers, model names etc.)
|
||||
//!
|
||||
//! ## Streaming (`chat_with_tools_streaming`)
|
||||
//! - Up to **5** attempts *before* any meaningful content (tokens / reasoning)
|
||||
//! - After meaningful content arrives, falls back to a **non-streaming retry**
|
||||
//! (the non-streaming call carries 10 retries of its own), so a mid-stream
|
||||
//! network blip is recovered instead of killing the whole turn.
|
||||
//! - The `started` flag still prevents retries on the raw SSE call once the
|
||||
//! stream has begun (partial content cannot be safely replayed), but the
|
||||
//! caller-level fallback handles that case.
|
||||
|
||||
use anyhow::Result;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::app::runtime::stream::turn::StreamedTurn;
|
||||
use crate::app::runtime::stream::{SseParser, StreamEvent};
|
||||
@@ -14,6 +38,72 @@ pub const DEFAULT_API_KEY: &str = "";
|
||||
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const REQUEST_TIMEOUT: Duration = Duration::from_mins(1);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Retry helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Return a pseudo-random jitter offset in the range [0, range_ns).
|
||||
fn jitter_ns(range_ns: u64) -> u64 {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos() as u64;
|
||||
nanos % range_ns
|
||||
}
|
||||
|
||||
/// Exponential backoff with ±25% jitter, capped at 30 seconds.
|
||||
///
|
||||
/// `attempt` is 1-based (first retry → attempt=1).
|
||||
fn backoff_duration(attempt: u32) -> Duration {
|
||||
let base_secs = (2u64).pow(attempt).min(30);
|
||||
let quarter = (base_secs * 250_000_000).max(100_000_000); // ±25%, min 100ms
|
||||
let offset = jitter_ns(quarter);
|
||||
// ±25% jitter: sometimes slightly less, sometimes slightly more
|
||||
let ns = base_secs * 1_000_000_000 + offset - quarter / 2;
|
||||
Duration::from_nanos(ns)
|
||||
}
|
||||
|
||||
/// Is the error an auth / billing failure that retrying won't fix?
|
||||
///
|
||||
/// Matches the structured "API error {status} from …" format used by the
|
||||
/// request builders below, plus well-known auth keywords in case the body
|
||||
/// contains them. This is intentionally tighter than `contains("401")`,
|
||||
/// which could false-positive on a URL port, model name, or body text.
|
||||
fn is_auth_error(err_str: &str) -> bool {
|
||||
let err_lower = err_str.to_lowercase();
|
||||
// Structured HTTP status patterns
|
||||
(err_str.contains("API error 401")
|
||||
|| err_str.contains("API error 402")
|
||||
|| err_str.contains("API error 403"))
|
||||
// Keyword fallback for non-standard error formats
|
||||
|| err_lower.contains("unauthorized")
|
||||
|| err_lower.contains("forbidden")
|
||||
|| err_lower.contains("authentication failed")
|
||||
}
|
||||
|
||||
/// Is the error a rate-limit response?
|
||||
fn is_rate_limit(err_str: &str) -> bool {
|
||||
err_str.contains("API error 429") || err_str.to_lowercase().contains("rate limit")
|
||||
}
|
||||
|
||||
/// Return a rate-appropriate backoff (longer for 429).
|
||||
fn backoff_for_error(attempt: u32, err_str: &str) -> Duration {
|
||||
if is_rate_limit(err_str) {
|
||||
// Rate limits need more time to drain — start at 5s instead of 2s.
|
||||
let base_secs = (5u64 * (2u64).pow(attempt.saturating_sub(1))).min(60);
|
||||
let quarter = (base_secs * 250_000_000).max(100_000_000);
|
||||
let offset = jitter_ns(quarter);
|
||||
let ns = base_secs * 1_000_000_000 + offset - quarter / 2;
|
||||
Duration::from_nanos(ns)
|
||||
} else {
|
||||
backoff_duration(attempt)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Client
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Blocking HTTP client for a single LLM provider endpoint.
|
||||
///
|
||||
/// Holds the reqwest client, credentials, and model/base URL selection used
|
||||
@@ -84,11 +174,13 @@ impl LlmClient {
|
||||
|
||||
/// Send a non-streaming chat completion request and return the assistant's reply.
|
||||
///
|
||||
/// Flow: build request → POST with retry loop (up to 10 attempts, 2s backoff)
|
||||
/// → parse JSON response → extract first choice's message and token usage.
|
||||
/// Flow: build request → POST with retry loop (up to 10 attempts, exponential
|
||||
/// backoff with jitter) → parse JSON response → extract first choice's message
|
||||
/// and token usage.
|
||||
///
|
||||
/// Why: retries transient failures but aborts immediately on 401/403, since
|
||||
/// those indicate a bad API key that retrying won't fix.
|
||||
/// Why: retries transient failures but aborts immediately on 401/402/403 (bad
|
||||
/// API key / billing issue — retrying won't fix). 429 (rate-limit) responses
|
||||
/// get a longer backoff so the server has time to recover.
|
||||
///
|
||||
/// Return: `Err` if all retries are exhausted, an auth error occurs, or the
|
||||
/// response has no choices.
|
||||
@@ -112,7 +204,7 @@ impl LlmClient {
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
let max_retries = 10;
|
||||
let mut attempt = 0;
|
||||
let mut attempt = 0u32;
|
||||
|
||||
loop {
|
||||
attempt += 1;
|
||||
@@ -160,28 +252,45 @@ impl LlmClient {
|
||||
Ok((msg, usage)) => return Ok((msg, usage)),
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
let err_lower = err_str.to_lowercase();
|
||||
let is_auth_error = err_str.contains("401")
|
||||
|| err_str.contains("403")
|
||||
|| err_lower.contains("unauthorized")
|
||||
|| err_lower.contains("forbidden")
|
||||
|| err_lower.contains("authentication failed");
|
||||
if attempt >= max_retries || is_auth_error {
|
||||
if attempt >= max_retries || is_auth_error(&err_str) {
|
||||
return Err(e);
|
||||
}
|
||||
tracing::warn!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
let delay = backoff_for_error(attempt, &err_str);
|
||||
tracing::warn!(
|
||||
"Warning: {}. Retrying {}/{}, sleeping {delay:?}...",
|
||||
e,
|
||||
attempt,
|
||||
max_retries,
|
||||
);
|
||||
std::thread::sleep(delay);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Streaming variant of `chat_with_tools`. Feeds SSE chunks into an `SseParser` /
|
||||
/// `StreamedTurn` and invokes `on_event` for every parsed `StreamEvent` as it arrives,
|
||||
/// so the caller can push incremental UI updates in real time. Returns the fully
|
||||
/// assembled assistant message plus token usage (prompt, completion) if the server
|
||||
/// reported it. Retries the whole request only if no event has been observed yet
|
||||
/// (once tokens start arriving, a partial turn cannot be safely replayed).
|
||||
/// Streaming variant of `chat_with_tools`. Feeds SSE chunks into an
|
||||
/// `SseParser` / `StreamedTurn` and invokes `on_event` for every parsed
|
||||
/// `StreamEvent` as it arrives, so the caller can push incremental UI
|
||||
/// updates in real time.
|
||||
///
|
||||
/// Returns the fully assembled assistant message plus token usage (prompt,
|
||||
/// completion) if the server reported it.
|
||||
///
|
||||
/// # Retry semantics
|
||||
///
|
||||
/// Retries the raw SSE request only *before* any meaningful content (text
|
||||
/// tokens or reasoning tokens) has been received — once the LLM has started
|
||||
/// generating, a partial stream cannot be safely replayed without duplicating
|
||||
/// or garbling output.
|
||||
///
|
||||
/// Once meaningful content has arrived and the stream fails, **this method
|
||||
/// falls back to a non-streaming call** (which carries its own 10-retry
|
||||
/// loop). The non-streaming call uses the same `messages` independently
|
||||
/// (no SSE state to replay), so the caller always gets a complete result if
|
||||
/// the provider is reachable.
|
||||
///
|
||||
/// Auth errors (401/402/403) are never retried on either path. Rate-limit
|
||||
/// (429) responses get a longer backoff.
|
||||
pub fn chat_with_tools_streaming(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
@@ -206,34 +315,168 @@ impl LlmClient {
|
||||
};
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
// Fewer retries on streaming because `run_agent_turn` has a
|
||||
// non-streaming fallback that also retries. Combined total is
|
||||
// capped implicitly by the per-turn timeout and step limits.
|
||||
let max_retries = 3;
|
||||
let mut attempt = 0;
|
||||
|
||||
// Phase 1: Retry the raw SSE call up to 5 times, but only before
|
||||
// meaningful content arrives. After that, fall back to non-streaming.
|
||||
let max_retries_stream = 5;
|
||||
let mut attempt = 0u32;
|
||||
let mut started = false;
|
||||
// Track whether we've emitted text/reasoning tokens (meaningful
|
||||
// content). Non-meaningful events (role/usage/done) are safe to
|
||||
// ignore for the retry decision.
|
||||
let mut meaningful_content = false;
|
||||
|
||||
loop {
|
||||
attempt += 1;
|
||||
let mut captured_content = false;
|
||||
let mut wrapped = |event: &StreamEvent| -> bool {
|
||||
started = true;
|
||||
match event {
|
||||
StreamEvent::Token(_) | StreamEvent::Reasoning(_) => {
|
||||
captured_content = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
on_event(event)
|
||||
};
|
||||
match self.try_stream_once(&req, &url, &mut wrapped) {
|
||||
Ok(result) => return Ok(result),
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
let err_lower = err_str.to_lowercase();
|
||||
let is_auth_error = err_str.contains("401")
|
||||
|| err_str.contains("403")
|
||||
|| err_lower.contains("unauthorized")
|
||||
|| err_lower.contains("forbidden")
|
||||
|| err_lower.contains("authentication failed");
|
||||
if started || attempt >= max_retries || is_auth_error {
|
||||
if is_auth_error(&err_str) {
|
||||
return Err(e);
|
||||
}
|
||||
tracing::warn!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
// Once meaningful content has been streamed, a raw SSE
|
||||
// retry would produce a different sequence — fall back
|
||||
// to non-streaming so the caller gets a clean,
|
||||
// reproducible answer.
|
||||
if captured_content || started && (attempt >= max_retries_stream) {
|
||||
meaningful_content = captured_content || meaningful_content;
|
||||
break;
|
||||
}
|
||||
if attempt >= max_retries_stream {
|
||||
return Err(e);
|
||||
}
|
||||
let delay = backoff_for_error(attempt, &err_str);
|
||||
tracing::warn!(
|
||||
"Warning: {}. Retrying stream {}/{}, sleeping {delay:?}...",
|
||||
e,
|
||||
attempt,
|
||||
max_retries_stream,
|
||||
);
|
||||
std::thread::sleep(delay);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2: If we got meaningful content via SSE but the stream
|
||||
// failed before completion, fall back to a non-streaming retry.
|
||||
// This preserves the conversation state because the messages
|
||||
// passed in are the same — we don't need the partial SSE output.
|
||||
if meaningful_content {
|
||||
tracing::warn!(
|
||||
"streaming failed after meaningful content — falling back to non-streaming retry",
|
||||
);
|
||||
// Use the same messages; pass None for tools (streaming already
|
||||
// included them) and let the non-streaming path handle retries.
|
||||
// The on_event callback is irrelevant for non-streaming, but we
|
||||
// signal a special synthetic Done event so callers aren't left
|
||||
// hanging waiting for stream completion.
|
||||
// Rebuild ChatRequest without streaming options.
|
||||
return self.chat_with_tools_non_streaming_retry(
|
||||
messages,
|
||||
max_tokens.unwrap_or(4096),
|
||||
temperature.unwrap_or(0.7),
|
||||
);
|
||||
}
|
||||
|
||||
Err(anyhow::anyhow!(
|
||||
"streaming request failed after {max_retries_stream} attempts"
|
||||
))
|
||||
}
|
||||
|
||||
/// Non-streaming fallback used by the streaming method after a partial
|
||||
/// stream failure. Same retry policy as `chat_with_tools_non_streaming`.
|
||||
fn chat_with_tools_non_streaming_retry(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
max_tokens: u32,
|
||||
temperature: f32,
|
||||
) -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
let req = ChatRequest {
|
||||
model: self.model.clone(),
|
||||
messages: messages.to_vec(),
|
||||
max_tokens: Some(max_tokens),
|
||||
temperature: Some(temperature),
|
||||
tools: None,
|
||||
stream: Some(false),
|
||||
stop: None,
|
||||
stream_options: None,
|
||||
tool_choice: None,
|
||||
top_p: None,
|
||||
};
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
let max_retries = 10;
|
||||
let mut attempt = 0u32;
|
||||
|
||||
loop {
|
||||
attempt += 1;
|
||||
|
||||
let mut http_req = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Content-Type", "application/json");
|
||||
|
||||
if !self.api_key.is_empty() {
|
||||
http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
|
||||
}
|
||||
|
||||
let result = (|| -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
let resp = http_req.json(&req).send().map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
anyhow::anyhow!("API request timed out after {REQUEST_TIMEOUT:?}. Check your network or try again.")
|
||||
} else if e.is_connect() {
|
||||
anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
|
||||
} else {
|
||||
anyhow::anyhow!("API request failed: {e}")
|
||||
}
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().unwrap_or_default();
|
||||
anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
|
||||
}
|
||||
|
||||
let data: crate::dto::provider::response::ChatResponse = resp.json()?;
|
||||
let usage = data
|
||||
.usage
|
||||
.map(|u| (u64::from(u.prompt_tokens), u64::from(u.completion_tokens)));
|
||||
let message = data
|
||||
.choices
|
||||
.into_iter()
|
||||
.next()
|
||||
.and_then(|c| c.message)
|
||||
.ok_or_else(|| anyhow::anyhow!("API response had no choices"))?;
|
||||
Ok((message, usage))
|
||||
})();
|
||||
|
||||
match result {
|
||||
Ok((msg, usage)) => return Ok((msg, usage)),
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
if attempt >= max_retries || is_auth_error(&err_str) {
|
||||
return Err(e);
|
||||
}
|
||||
let delay = backoff_for_error(attempt, &err_str);
|
||||
tracing::warn!(
|
||||
"Warning [non-streaming fallback]: {}. Retrying {}/{}, sleeping {delay:?}...",
|
||||
e,
|
||||
attempt,
|
||||
max_retries,
|
||||
);
|
||||
std::thread::sleep(delay);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user