lc_core/runnables/
retry.rs

1// lc-core/src/runnables/retry.rs
2//! RunnableRetry - retry a Runnable with exponential backoff.
3//!
4//! # Example
5//!
6//! ```rust,ignore
7//! use lc_core::runnables::{RetryConfig, RetryOn, RunnableExt};
8//! use std::time::Duration;
9//!
10//! let chain = prompt.pipe(llm).pipe(parser)
11//!     .with_retry(RetryConfig {
12//!         max_retries: 3,
13//!         initial_delay: Duration::from_millis(500),
14//!         max_delay: Duration::from_secs(10),
15//!         backoff_multiplier: 2.0,
16//!         retry_on: RetryOn::TransientErrors,
17//!     });
18//! ```
19
20use std::pin::Pin;
21use std::sync::Arc;
22use std::time::Duration;
23
24use async_trait::async_trait;
25use futures_util::{Stream, StreamExt};
26
27use super::any::RunnableAny;
28use super::config::RunnableConfig;
29use super::error::LcelError;
30use super::runnable_trait::Runnable;
31
32/// Determines which errors should trigger a retry.
33#[derive(Clone)]
34pub enum RetryOn {
35    /// Retry on any error.
36    AllErrors,
37    /// Retry only on transient errors (rate limits, timeouts, server errors).
38    /// Specifically: HTTP 429, 500, 502, 503, 504 and timeout errors.
39    TransientErrors,
40    /// Retry only when a custom predicate returns true.
41    Custom(Arc<dyn Fn(&str) -> bool + Send + Sync>),
42}
43
44impl std::fmt::Debug for RetryOn {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        match self {
47            Self::AllErrors => write!(f, "AllErrors"),
48            Self::TransientErrors => write!(f, "TransientErrors"),
49            Self::Custom(_) => write!(f, "Custom(<closure>)"),
50        }
51    }
52}
53
54/// Configuration for retry behavior.
55#[derive(Debug, Clone)]
56pub struct RetryConfig {
57    /// Maximum number of retry attempts (not counting the initial call).
58    pub max_retries: usize,
59
60    /// Initial delay before the first retry.
61    pub initial_delay: Duration,
62
63    /// Maximum delay between retries (caps exponential growth).
64    pub max_delay: Duration,
65
66    /// Multiplier applied to the delay after each attempt.
67    /// A value of 2.0 means each retry waits twice as long as the previous.
68    pub backoff_multiplier: f64,
69
70    /// Which errors should trigger a retry.
71    pub retry_on: RetryOn,
72}
73
74impl Default for RetryConfig {
75    fn default() -> Self {
76        Self {
77            max_retries: 3,
78            initial_delay: Duration::from_millis(500),
79            max_delay: Duration::from_secs(10),
80            backoff_multiplier: 2.0,
81            retry_on: RetryOn::TransientErrors,
82        }
83    }
84}
85
86impl RetryConfig {
87    /// Creates a new RetryConfig with the specified max retries and defaults for other fields.
88    pub fn new(max_retries: usize) -> Self {
89        Self {
90            max_retries,
91            ..Default::default()
92        }
93    }
94
95    /// Sets the initial delay.
96    pub fn with_initial_delay(mut self, delay: Duration) -> Self {
97        self.initial_delay = delay;
98        self
99    }
100
101    /// Sets the maximum delay.
102    pub fn with_max_delay(mut self, delay: Duration) -> Self {
103        self.max_delay = delay;
104        self
105    }
106
107    /// Sets the backoff multiplier.
108    pub fn with_backoff_multiplier(mut self, multiplier: f64) -> Self {
109        self.backoff_multiplier = multiplier;
110        self
111    }
112
113    /// Sets which errors should trigger a retry.
114    pub fn with_retry_on(mut self, retry_on: RetryOn) -> Self {
115        self.retry_on = retry_on;
116        self
117    }
118
119    /// Checks if an error should trigger a retry.
120    fn should_retry(&self, error: &str) -> bool {
121        match &self.retry_on {
122            RetryOn::AllErrors => true,
123            RetryOn::TransientErrors => is_transient_error(error),
124            RetryOn::Custom(predicate) => predicate(error),
125        }
126    }
127
128    /// Calculates the delay for a given attempt number (0-based).
129    fn delay_for_attempt(&self, attempt: usize) -> Duration {
130        let multiplier = self.backoff_multiplier.powi(attempt as i32);
131        let delay = self.initial_delay.as_secs_f64() * multiplier;
132        let delay = delay.min(self.max_delay.as_secs_f64());
133        Duration::from_secs_f64(delay)
134    }
135}
136
137/// Check if an error looks like a transient (retriable) error.
138fn is_transient_error(error: &str) -> bool {
139    let error_lower = error.to_lowercase();
140
141    // HTTP status codes that are retriable
142    for code in &["429", "500", "502", "503", "504"] {
143        if error_lower.contains(code) {
144            return true;
145        }
146    }
147
148    // Common transient error patterns
149    let transient_patterns = [
150        "rate limit",
151        "rate_limit",
152        "ratelimit",
153        "too many requests",
154        "timeout",
155        "timed out",
156        "connection reset",
157        "connection refused",
158        "temporary failure",
159        "service unavailable",
160        "internal server error",
161        "overloaded",
162        "capacity",
163    ];
164
165    for pattern in &transient_patterns {
166        if error_lower.contains(pattern) {
167            return true;
168        }
169    }
170
171    false
172}
173
174/// A Runnable wrapper that retries the inner Runnable on failure.
175pub struct RunnableRetry<I: Send + Sync + 'static, O: Send + Sync + 'static> {
176    runnable: Arc<dyn RunnableAny>,
177    retry_config: RetryConfig,
178    _marker: std::marker::PhantomData<(I, O)>,
179}
180
181impl<I: Send + Sync + 'static, O: Send + Sync + 'static> RunnableRetry<I, O> {
182    /// Creates a new RunnableRetry from a boxed RunnableAny.
183    pub fn new(runnable: Box<dyn RunnableAny>, retry_config: RetryConfig) -> Self {
184        Self {
185            runnable: Arc::from(runnable),
186            retry_config,
187            _marker: std::marker::PhantomData,
188        }
189    }
190}
191
192#[async_trait]
193impl<I: Send + Sync + 'static, O: Send + Sync + 'static> Runnable<I, O> for RunnableRetry<I, O>
194where
195    I: Clone,
196{
197    type Error = LcelError;
198
199    async fn invoke(&self, input: I, config: Option<RunnableConfig>) -> Result<O, Self::Error> {
200        // Check cancellation before starting
201        if config.as_ref().is_some_and(|c| c.is_cancelled()) {
202            return Err(LcelError::Other("Operation cancelled".to_string()));
203        }
204
205        let mut last_error = None;
206
207        for attempt in 0..=self.retry_config.max_retries {
208            // Check cancellation before each attempt
209            if attempt > 0 && config.as_ref().is_some_and(|c| c.is_cancelled()) {
210                return Err(LcelError::Other("Operation cancelled".to_string()));
211            }
212
213            // Delay before retry (not on first attempt)
214            if attempt > 0 {
215                let delay = self.retry_config.delay_for_attempt(attempt - 1);
216                tokio::time::sleep(delay).await;
217            }
218
219            match self
220                .runnable
221                .invoke_any(Box::new(input.clone()), config.clone())
222                .await
223            {
224                Ok(result) => {
225                    return result.downcast::<O>().map(|boxed| *boxed).map_err(|_| {
226                        LcelError::Other("Type mismatch in retry result".to_string())
227                    });
228                }
229                Err(e) => {
230                    let error_str = e.to_string();
231                    if attempt < self.retry_config.max_retries
232                        && self.retry_config.should_retry(&error_str)
233                    {
234                        last_error = Some(e);
235                        continue;
236                    }
237                    return Err(e);
238                }
239            }
240        }
241
242        Err(last_error.unwrap_or_else(|| {
243            LcelError::Other("Retry exhausted with no error recorded".to_string())
244        }))
245    }
246
247    async fn stream(
248        &self,
249        input: I,
250        config: Option<RunnableConfig>,
251    ) -> Result<Pin<Box<dyn Stream<Item = Result<O, Self::Error>> + Send>>, Self::Error> {
252        // For stream, we retry the stream setup (construction), not individual tokens.
253        // Once the stream is established, token-level errors propagate normally.
254        let mut last_error = None;
255
256        for attempt in 0..=self.retry_config.max_retries {
257            if attempt > 0 && config.as_ref().is_some_and(|c| c.is_cancelled()) {
258                return Err(LcelError::Other("Operation cancelled".to_string()));
259            }
260
261            if attempt > 0 {
262                let delay = self.retry_config.delay_for_attempt(attempt - 1);
263                tokio::time::sleep(delay).await;
264            }
265
266            match self
267                .runnable
268                .stream_any(Box::new(input.clone()), config.clone())
269                .await
270            {
271                Ok(stream) => {
272                    // Convert the type-erased stream to a typed stream
273                    let typed_stream = stream.map(|result| {
274                        result.and_then(|boxed| {
275                            boxed.downcast::<O>().map(|b| *b).map_err(|_| {
276                                LcelError::Other("Type mismatch in retry stream".to_string())
277                            })
278                        })
279                    });
280                    return Ok(Box::pin(typed_stream));
281                }
282                Err(e) => {
283                    let error_str = e.to_string();
284                    if attempt < self.retry_config.max_retries
285                        && self.retry_config.should_retry(&error_str)
286                    {
287                        last_error = Some(e);
288                        continue;
289                    }
290                    return Err(e);
291                }
292            }
293        }
294
295        Err(last_error.unwrap_or_else(|| {
296            LcelError::Other("Retry exhausted with no error recorded".to_string())
297        }))
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    use crate::runnables::{CancellationToken, RunnableConfig, RunnableExt, RunnableLambda};
305    use std::sync::atomic::{AtomicUsize, Ordering};
306    use std::sync::Arc;
307
308    #[test]
309    fn test_retry_config_default() {
310        let config = RetryConfig::default();
311        assert_eq!(config.max_retries, 3);
312        assert_eq!(config.initial_delay, Duration::from_millis(500));
313        assert_eq!(config.max_delay, Duration::from_secs(10));
314        assert!((config.backoff_multiplier - 2.0).abs() < f64::EPSILON);
315    }
316
317    #[test]
318    fn test_delay_for_attempt() {
319        let config = RetryConfig::default();
320        assert_eq!(config.delay_for_attempt(0), Duration::from_millis(500));
321        assert_eq!(config.delay_for_attempt(1), Duration::from_secs(1));
322        assert_eq!(config.delay_for_attempt(2), Duration::from_secs(2));
323        // Should be capped at max_delay
324        assert_eq!(config.delay_for_attempt(10), Duration::from_secs(10));
325    }
326
327    #[test]
328    fn test_is_transient_error() {
329        assert!(is_transient_error("HTTP 429: Too Many Requests"));
330        assert!(is_transient_error("HTTP 503: Service Unavailable"));
331        assert!(is_transient_error("rate limit exceeded"));
332        assert!(is_transient_error("Connection timeout"));
333        assert!(is_transient_error("internal server error"));
334
335        assert!(!is_transient_error("HTTP 401: Unauthorized"));
336        assert!(!is_transient_error("HTTP 403: Forbidden"));
337        assert!(!is_transient_error("invalid API key"));
338        assert!(!is_transient_error("model not found"));
339    }
340
341    #[tokio::test]
342    async fn test_retry_succeeds_on_second_attempt() {
343        let call_count = Arc::new(AtomicUsize::new(0));
344        let count_clone = call_count.clone();
345
346        let runnable = RunnableLambda::new_async(move |_: i32| {
347            let count = count_clone.clone();
348            async move {
349                let n = count.fetch_add(1, Ordering::SeqCst);
350                if n == 0 {
351                    Err(LcelError::Other(
352                        "HTTP 503: Service Unavailable".to_string(),
353                    ))
354                } else {
355                    Ok(42)
356                }
357            }
358        });
359
360        let retry = runnable.with_retry(RetryConfig::new(2));
361        let result: Result<i32, _> = retry.invoke(1, None).await;
362        assert_eq!(result.unwrap(), 42);
363        assert_eq!(call_count.load(Ordering::SeqCst), 2);
364    }
365
366    #[tokio::test]
367    async fn test_retry_exhausts_all_attempts() {
368        let call_count = Arc::new(AtomicUsize::new(0));
369        let count_clone = call_count.clone();
370
371        let runnable = RunnableLambda::new_async(move |_: i32| {
372            let count = count_clone.clone();
373            async move {
374                count.fetch_add(1, Ordering::SeqCst);
375                Err(LcelError::Other(
376                    "HTTP 503: Service Unavailable".to_string(),
377                ))
378            }
379        });
380
381        let retry = runnable.with_retry(RetryConfig::new(2));
382        let result: Result<i32, _> = retry.invoke(1, None).await;
383        assert!(result.is_err());
384        assert_eq!(call_count.load(Ordering::SeqCst), 3); // 1 initial + 2 retries
385    }
386
387    #[tokio::test]
388    async fn test_retry_non_retriable_error_fails_immediately() {
389        let call_count = Arc::new(AtomicUsize::new(0));
390        let count_clone = call_count.clone();
391
392        let runnable = RunnableLambda::new_async(move |_: i32| {
393            let count = count_clone.clone();
394            async move {
395                count.fetch_add(1, Ordering::SeqCst);
396                Err(LcelError::Other("HTTP 401: Unauthorized".to_string()))
397            }
398        });
399
400        let retry = runnable.with_retry(RetryConfig::new(3));
401        let result: Result<i32, _> = retry.invoke(1, None).await;
402        assert!(result.is_err());
403        assert_eq!(call_count.load(Ordering::SeqCst), 1); // No retry for 401
404    }
405
406    #[tokio::test]
407    async fn test_retry_succeeds_on_first_attempt() {
408        let runnable = RunnableLambda::new_sync(|x: i32| x * 2);
409        let retry = runnable.with_retry(RetryConfig::new(3));
410        let result: Result<i32, _> = retry.invoke(5, None).await;
411        assert_eq!(result.unwrap(), 10);
412    }
413
414    #[tokio::test]
415    async fn test_retry_respects_cancellation() {
416        let token = CancellationToken::new();
417        token.cancel();
418
419        let runnable = RunnableLambda::new_sync(|x: i32| x * 2);
420        let retry = runnable.with_retry(RetryConfig::new(3));
421
422        let config = RunnableConfig::new().with_cancellation_token(token);
423        let result: Result<i32, _> = retry.invoke(5, Some(config)).await;
424        assert!(result.is_err());
425        assert!(result.unwrap_err().to_string().contains("cancelled"));
426    }
427}
lc_core/runnables/retry.rs

lc_core/runnables/
retry.rs
