From b02754acd269b82fc98d5d188eae8ea195f765ee Mon Sep 17 00:00:00 2001 From: asepharyana Date: Fri, 17 Jul 2026 15:03:08 +0700 Subject: [PATCH] refactor: implement retry logic with exponential backoff and jitter for subagent and provider calls --- .../zesdex-backend/src/app/subagent/engine.rs | 188 +++++++---- .../src/app/workflow/engine/mod.rs | 58 +++- crates/zesdex-backend/src/service/provider.rs | 313 ++++++++++++++++-- 3 files changed, 459 insertions(+), 100 deletions(-) diff --git a/crates/zesdex-backend/src/app/subagent/engine.rs b/crates/zesdex-backend/src/app/subagent/engine.rs index 821ea69..60e2ac7 100644 --- a/crates/zesdex-backend/src/app/subagent/engine.rs +++ b/crates/zesdex-backend/src/app/subagent/engine.rs @@ -15,9 +15,49 @@ use crate::dto::chat::message::ChatMessage; use crate::dto::provider::request::ToolDef; use crate::tool::tool_is_risky; use sha2::Digest; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::mpsc; use zesdex_cms::domain::repository::EditLogRepository; +/// Tiny jitter helper so retry backoffs don't arrive in lockstep. +fn retry_jitter_ns(range_ns: u64) -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos() as u64 + % range_ns +} + +/// Exponential backoff with ±25% jitter for subagent step retries. +fn step_retry_delay(attempt: u32) -> Duration { + let base_secs = (2u64).pow(attempt).min(16); // 2s, 4s, 8s, 16s cap + let quarter = (base_secs * 250_000_000).max(100_000_000); + let offset = retry_jitter_ns(quarter); + let ns = base_secs * 1_000_000_000 + offset - quarter / 2; + Duration::from_nanos(ns) +} + +/// Heuristic to decide whether the error is worth retrying. +fn should_retry_subagent_step(err_str: &str) -> bool { + let lower = err_str.to_lowercase(); + // Never retry auth/billing failures + if err_str.contains("API error 401") + || err_str.contains("API error 402") + || err_str.contains("API error 403") + || lower.contains("unauthorized") + || lower.contains("forbidden") + || lower.contains("authentication failed") + { + return false; + } + // Never retry abort or user cancellation + if lower.contains("aborted") { + return false; + } + // Everything else (timeout, 5xx, rate-limit, network blip) is retryable + true +} + fn format_subagent_progress(prefix: &str, text: &str) -> String { let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).collect(); if lines.is_empty() { @@ -118,71 +158,93 @@ pub fn run_subagent( // Use streaming API so the abort flag is checked per SSE event, // making the subagent responsive to cancellation even during an // LLM call (non-streaming would block for 10-30s unchecked). - let stream_result = client.chat_with_tools_streaming( - &messages, - tdefs_opt.clone(), - Some(0.7), - Some(4096), - |event| -> bool { - // Check abort on every SSE event for responsive cancellation. - if ctx - .abort_flag - .as_ref() - .is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) - { - return false; // signals provider to abort - } - match event { - crate::app::runtime::stream::StreamEvent::Reasoning(text) => { - current_thinking.push_str(text); - let prog = format_subagent_progress("thinking", ¤t_thinking); - let _ = tx_clone.blocking_send(SubagentEvent::Progress(prog)); - } - crate::app::runtime::stream::StreamEvent::Token(text) => { - current_token.push_str(text); - let prog = format_subagent_progress("replying", ¤t_token); - let _ = tx_clone.blocking_send(SubagentEvent::Progress(prog)); - } - crate::app::runtime::stream::StreamEvent::Usage { - prompt_tokens, - completion_tokens, - .. - } => { - // Capture usage so the drain thread can route it - // to the parent's `UsageStats::review_tokens`. - // Last writer wins — providers send exactly one - // Usage event per streaming call. - step_usage = Some((*prompt_tokens, *completion_tokens)); - } - _ => {} - } - true - }, - ); + // Retry the LLM call at the step level (up to 3 attempts) so a + // transient network blip doesn't kill the subagent. The underlying + // `chat_with_tools_streaming` already has its own retry loop (5 + + // non-streaming fallback), so this loop is a second safety net for + // rare cases where the combined 5+10 retries are all exhausted. + let max_step_retries = 3; + let mut step_attempt = 0u32; - let (response, returned_usage) = match stream_result { - Ok(result) => result, - Err(e) => { - let is_abort = ctx - .abort_flag - .as_ref() - .is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) - || e.to_string().contains("aborted"); - let _ = tx.blocking_send(SubagentEvent::StepFailed { - step, - error: if is_abort { - "subagent aborted by user".to_string() - } else { - e.to_string() - }, - }); - if is_abort { - anyhow::bail!("subagent aborted by parent at step {step}"); + let (response, returned_usage) = loop { + step_attempt += 1; + let stream_result = client.chat_with_tools_streaming( + &messages, + tdefs_opt.clone(), + Some(0.7), + Some(4096), + |event| -> bool { + // Check abort on every SSE event for responsive cancellation. + if ctx + .abort_flag + .as_ref() + .is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) + { + return false; // signals provider to abort + } + match event { + crate::app::runtime::stream::StreamEvent::Reasoning(text) => { + current_thinking.push_str(text); + let prog = format_subagent_progress("thinking", ¤t_thinking); + let _ = tx_clone.blocking_send(SubagentEvent::Progress(prog)); + } + crate::app::runtime::stream::StreamEvent::Token(text) => { + current_token.push_str(text); + let prog = format_subagent_progress("replying", ¤t_token); + let _ = tx_clone.blocking_send(SubagentEvent::Progress(prog)); + } + crate::app::runtime::stream::StreamEvent::Usage { + prompt_tokens, + completion_tokens, + .. + } => { + // Capture usage so the drain thread can route it + // to the parent's `UsageStats::review_tokens`. + // Last writer wins — providers send exactly one + // Usage event per streaming call. + step_usage = Some((*prompt_tokens, *completion_tokens)); + } + _ => {} + } + true + }, + ); + + match stream_result { + Ok(result) => break result, + Err(e) => { + let err_str = e.to_string(); + let is_abort = ctx + .abort_flag + .as_ref() + .is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) + || err_str.contains("aborted"); + + if is_abort || !should_retry_subagent_step(&err_str) || step_attempt >= max_step_retries { + let _ = tx.blocking_send(SubagentEvent::StepFailed { + step, + error: if is_abort { + "subagent aborted by user".to_string() + } else { + err_str.clone() + }, + }); + if is_abort { + anyhow::bail!("subagent aborted by parent at step {step}"); + } + anyhow::bail!("subagent call failed at step {step} after {step_attempt} attempt(s): {err_str}"); + } + + let delay = step_retry_delay(step_attempt); + tracing::warn!( + "[subagent] step {step} attempt {step_attempt}/{max_step_retries} failed: {err_str}. \ + retrying in {delay:?}...", + ); + let _ = tx.blocking_send(SubagentEvent::Progress(format!( + "retrying step {step} ({step_attempt}/{max_step_retries}) after error…", + ))); + std::thread::sleep(delay); } - // No non-streaming fallback — API must support streaming. - // Non-streaming calls block for up to 1 min without checking - // abort_flag, making cancellation unresponsive. - anyhow::bail!("subagent call failed at step {step}: {e}"); } }; diff --git a/crates/zesdex-backend/src/app/workflow/engine/mod.rs b/crates/zesdex-backend/src/app/workflow/engine/mod.rs index cb73b8f..9eaef8f 100644 --- a/crates/zesdex-backend/src/app/workflow/engine/mod.rs +++ b/crates/zesdex-backend/src/app/workflow/engine/mod.rs @@ -28,7 +28,7 @@ use std::sync::{ atomic::{AtomicBool, Ordering}, Arc, Mutex, }; -use std::time::Duration; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; /// The lifecycle state of an agent within a workflow run. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -387,13 +387,67 @@ fn spawn_single_agent(sp: SpawnCtx<'_>) -> anyhow::Result { // Run subagent on a separate thread so the abort flag can be polled. // If abort is requested while the subagent is running, we abandon the // thread (Rust threads cannot be forcibly killed) and return early. + // + // Retry: wrap `run_subagent` with up to 2 attempts so a transient + // network blip doesn't kill the whole pipeline. Auth and abort errors + // are not retried. let (done_tx, done_rx) = std::sync::mpsc::channel::>(); let bg_ctx = ctx; let bg_tx = tx; let bg_name = sp.agent_name.to_string(); let bg_abort = sp.abort_flag.clone(); + let bg_abort_thread = bg_abort.clone(); + let bg_name_thread = bg_name.clone(); std::thread::spawn(move || { - let _ = done_tx.send(run_subagent(&bg_ctx, &bg_tx)); + // Retry wrapper: jittered backoff 1s → 2s. + let retry_backoff = |attempt: u32| { + let base_secs = (2u64).pow(attempt).min(8); + let quarter = (base_secs * 250_000_000).max(100_000_000); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos() as u64; + let offset = nanos % quarter; + let ns = base_secs * 1_000_000_000 + offset - quarter / 2; + Duration::from_nanos(ns) + }; + + for attempt in 1..=2 { + // Don't retry if aborted. + if bg_abort_thread + .as_ref() + .is_some_and(|f| f.load(Ordering::SeqCst)) + { + let _ = done_tx.send(Err(anyhow::anyhow!( + "subagent '{bg_name_thread}' aborted by user" + ))); + return; + } + match run_subagent(&bg_ctx, &bg_tx) { + Ok(output) => { + let _ = done_tx.send(Ok(output)); + return; + } + Err(e) => { + let err_str = e.to_string(); + let is_auth = err_str.contains("API error 401") + || err_str.contains("API error 402") + || err_str.contains("API error 403"); + // Auth errors are permanent — don't retry. + if is_auth || attempt >= 2 { + let _ = done_tx.send(Err(e)); + return; + } + tracing::warn!( + "[workflow] agent '{bg_name_thread}' attempt {attempt}/2 failed: {err_str}. retrying...", + ); + std::thread::sleep(retry_backoff(attempt)); + } + } + } + // Should be unreachable because the loop returns on success or final + // error, but keep the compiler happy. + unreachable!() }); let poll_interval = Duration::from_millis(200); diff --git a/crates/zesdex-backend/src/service/provider.rs b/crates/zesdex-backend/src/service/provider.rs index dff4565..1f502be 100644 --- a/crates/zesdex-backend/src/service/provider.rs +++ b/crates/zesdex-backend/src/service/provider.rs @@ -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); } } }