refactor: implement retry logic with exponential backoff and jitter for subagent and provider calls
This commit is contained in:
@@ -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}");
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -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<String> {
|
||||
// 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::<anyhow::Result<String>>();
|
||||
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);
|
||||
|
||||
Reference in New Issue
Block a user