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(
200        &self,
201        input: I,
202        config: Option<RunnableConfig>,
203    ) -> Result<O, Self::Error> {
204        // Check cancellation before starting
205        if config.as_ref().is_some_and(|c| c.is_cancelled()) {
206            return Err(LcelError::Other("Operation cancelled".to_string()));
207        }
208
209        let mut last_error = None;
210
211        for attempt in 0..=self.retry_config.max_retries {
212            // Check cancellation before each attempt
213            if attempt > 0 && config.as_ref().is_some_and(|c| c.is_cancelled()) {
214                return Err(LcelError::Other("Operation cancelled".to_string()));
215            }
216
217            // Delay before retry (not on first attempt)
218            if attempt > 0 {
219                let delay = self.retry_config.delay_for_attempt(attempt - 1);
220                tokio::time::sleep(delay).await;
221            }
222
223            match self.runnable.invoke_any(Box::new(input.clone()), config.clone()).await {
224                Ok(result) => {
225                    return result
226                        .downcast::<O>()
227                        .map(|boxed| *boxed)
228                        .map_err(|_| LcelError::Other("Type mismatch in retry result".to_string()));
229                }
230                Err(e) => {
231                    let error_str = e.to_string();
232                    if attempt < self.retry_config.max_retries && self.retry_config.should_retry(&error_str) {
233                        last_error = Some(e);
234                        continue;
235                    }
236                    return Err(e);
237                }
238            }
239        }
240
241        Err(last_error.unwrap_or_else(|| LcelError::Other("Retry exhausted with no error recorded".to_string())))
242    }
243
244    async fn stream(
245        &self,
246        input: I,
247        config: Option<RunnableConfig>,
248    ) -> Result<Pin<Box<dyn Stream<Item = Result<O, Self::Error>> + Send>>, Self::Error> {
249        // For stream, we retry the stream setup (construction), not individual tokens.
250        // Once the stream is established, token-level errors propagate normally.
251        let mut last_error = None;
252
253        for attempt in 0..=self.retry_config.max_retries {
254            if attempt > 0 && config.as_ref().is_some_and(|c| c.is_cancelled()) {
255                return Err(LcelError::Other("Operation cancelled".to_string()));
256            }
257
258            if attempt > 0 {
259                let delay = self.retry_config.delay_for_attempt(attempt - 1);
260                tokio::time::sleep(delay).await;
261            }
262
263            match self.runnable.stream_any(Box::new(input.clone()), config.clone()).await {
264                Ok(stream) => {
265                    // Convert the type-erased stream to a typed stream
266                    let typed_stream = stream.map(|result| {
267                        result
268                            .and_then(|boxed| {
269                                boxed
270                                    .downcast::<O>()
271                                    .map(|b| *b)
272                                    .map_err(|_| LcelError::Other("Type mismatch in retry stream".to_string()))
273                            })
274                    });
275                    return Ok(Box::pin(typed_stream));
276                }
277                Err(e) => {
278                    let error_str = e.to_string();
279                    if attempt < self.retry_config.max_retries && self.retry_config.should_retry(&error_str) {
280                        last_error = Some(e);
281                        continue;
282                    }
283                    return Err(e);
284                }
285            }
286        }
287
288        Err(last_error.unwrap_or_else(|| LcelError::Other("Retry exhausted with no error recorded".to_string())))
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295    use crate::runnables::{CancellationToken, RunnableConfig, RunnableLambda, RunnableExt};
296    use std::sync::atomic::{AtomicUsize, Ordering};
297    use std::sync::Arc;
298
299    #[test]
300    fn test_retry_config_default() {
301        let config = RetryConfig::default();
302        assert_eq!(config.max_retries, 3);
303        assert_eq!(config.initial_delay, Duration::from_millis(500));
304        assert_eq!(config.max_delay, Duration::from_secs(10));
305        assert!(config.backoff_multiplier - 2.0 < f64::EPSILON);
306    }
307
308    #[test]
309    fn test_delay_for_attempt() {
310        let config = RetryConfig::default();
311        assert_eq!(config.delay_for_attempt(0), Duration::from_millis(500));
312        assert_eq!(config.delay_for_attempt(1), Duration::from_secs(1));
313        assert_eq!(config.delay_for_attempt(2), Duration::from_secs(2));
314        // Should be capped at max_delay
315        assert_eq!(config.delay_for_attempt(10), Duration::from_secs(10));
316    }
317
318    #[test]
319    fn test_is_transient_error() {
320        assert!(is_transient_error("HTTP 429: Too Many Requests"));
321        assert!(is_transient_error("HTTP 503: Service Unavailable"));
322        assert!(is_transient_error("rate limit exceeded"));
323        assert!(is_transient_error("Connection timeout"));
324        assert!(is_transient_error("internal server error"));
325
326        assert!(!is_transient_error("HTTP 401: Unauthorized"));
327        assert!(!is_transient_error("HTTP 403: Forbidden"));
328        assert!(!is_transient_error("invalid API key"));
329        assert!(!is_transient_error("model not found"));
330    }
331
332    #[tokio::test]
333    async fn test_retry_succeeds_on_second_attempt() {
334        let call_count = Arc::new(AtomicUsize::new(0));
335        let count_clone = call_count.clone();
336
337        let runnable = RunnableLambda::new_async(move |_: i32| {
338            let count = count_clone.clone();
339            async move {
340                let n = count.fetch_add(1, Ordering::SeqCst);
341                if n == 0 {
342                    Err(LcelError::Other("HTTP 503: Service Unavailable".to_string()))
343                } else {
344                    Ok(42)
345                }
346            }
347        });
348
349        let retry = runnable.with_retry(RetryConfig::new(2));
350        let result: Result<i32, _> = retry.invoke(1, None).await;
351        assert_eq!(result.unwrap(), 42);
352        assert_eq!(call_count.load(Ordering::SeqCst), 2);
353    }
354
355    #[tokio::test]
356    async fn test_retry_exhausts_all_attempts() {
357        let call_count = Arc::new(AtomicUsize::new(0));
358        let count_clone = call_count.clone();
359
360        let runnable = RunnableLambda::new_async(move |_: i32| {
361            let count = count_clone.clone();
362            async move {
363                count.fetch_add(1, Ordering::SeqCst);
364                Err(LcelError::Other("HTTP 503: Service Unavailable".to_string()))
365            }
366        });
367
368        let retry = runnable.with_retry(RetryConfig::new(2));
369        let result: Result<i32, _> = retry.invoke(1, None).await;
370        assert!(result.is_err());
371        assert_eq!(call_count.load(Ordering::SeqCst), 3); // 1 initial + 2 retries
372    }
373
374    #[tokio::test]
375    async fn test_retry_non_retriable_error_fails_immediately() {
376        let call_count = Arc::new(AtomicUsize::new(0));
377        let count_clone = call_count.clone();
378
379        let runnable = RunnableLambda::new_async(move |_: i32| {
380            let count = count_clone.clone();
381            async move {
382                count.fetch_add(1, Ordering::SeqCst);
383                Err(LcelError::Other("HTTP 401: Unauthorized".to_string()))
384            }
385        });
386
387        let retry = runnable.with_retry(RetryConfig::new(3));
388        let result: Result<i32, _> = retry.invoke(1, None).await;
389        assert!(result.is_err());
390        assert_eq!(call_count.load(Ordering::SeqCst), 1); // No retry for 401
391    }
392
393    #[tokio::test]
394    async fn test_retry_succeeds_on_first_attempt() {
395        let runnable = RunnableLambda::new_sync(|x: i32| x * 2);
396        let retry = runnable.with_retry(RetryConfig::new(3));
397        let result: Result<i32, _> = retry.invoke(5, None).await;
398        assert_eq!(result.unwrap(), 10);
399    }
400
401    #[tokio::test]
402    async fn test_retry_respects_cancellation() {
403        let token = CancellationToken::new();
404        token.cancel();
405
406        let runnable = RunnableLambda::new_sync(|x: i32| x * 2);
407        let retry = runnable.with_retry(RetryConfig::new(3));
408
409        let config = RunnableConfig::new().with_cancellation_token(token);
410        let result: Result<i32, _> = retry.invoke(5, Some(config)).await;
411        assert!(result.is_err());
412        assert!(result.unwrap_err().to_string().contains("cancelled"));
413    }
414}
lc_core/runnables/retry.rs

lc_core/runnables/
retry.rs
