Skip to main content

weft_core/defaults/
error_handler.rs

1use crate::layers::error_handler::{ErrorAction, ErrorHandlerLayer, RequestError};
2use async_trait::async_trait;
3
4pub struct DefaultErrorHandler {
5    pub max_retries: u32,
6}
7
8#[async_trait]
9impl ErrorHandlerLayer for DefaultErrorHandler {
10    async fn handle(&self, error: &RequestError) -> ErrorAction {
11        match error.status {
12            // Rate limited -> switch key
13            Some(429) => ErrorAction::SwitchKey,
14
15            // Auth error -> switch key (key might be revoked)
16            Some(401) | Some(403) => ErrorAction::SwitchKey,
17
18            // Server error -> retry with backoff
19            Some(s) if s >= 500 => {
20                if error.retry_count < self.max_retries {
21                    ErrorAction::Retry {
22                        delay_ms: 1000 * (error.retry_count as u64 + 1),
23                    }
24                } else {
25                    ErrorAction::SwitchProvider
26                }
27            }
28
29            // Connection error (no status) -> retry
30            None => {
31                if error.retry_count < self.max_retries {
32                    ErrorAction::Retry { delay_ms: 500 }
33                } else {
34                    ErrorAction::SwitchProvider
35                }
36            }
37
38            // Client errors -> fail immediately
39            _ => ErrorAction::Fail {
40                message: error.message.clone(),
41            },
42        }
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    fn make_error(status: Option<u16>, retry_count: u32) -> RequestError {
51        RequestError {
52            status,
53            message: "test error".into(),
54            provider: "test".into(),
55            key_index: 0,
56            retry_count,
57        }
58    }
59
60    #[tokio::test]
61    async fn test_429_switches_key() {
62        let h = DefaultErrorHandler { max_retries: 2 };
63        let action = h.handle(&make_error(Some(429), 0)).await;
64        assert!(matches!(action, ErrorAction::SwitchKey));
65    }
66
67    #[tokio::test]
68    async fn test_500_retries_then_switches_provider() {
69        let h = DefaultErrorHandler { max_retries: 2 };
70        let a1 = h.handle(&make_error(Some(500), 0)).await;
71        assert!(matches!(a1, ErrorAction::Retry { .. }));
72        let a2 = h.handle(&make_error(Some(500), 2)).await;
73        assert!(matches!(a2, ErrorAction::SwitchProvider));
74    }
75
76    #[tokio::test]
77    async fn test_400_fails() {
78        let h = DefaultErrorHandler { max_retries: 2 };
79        let action = h.handle(&make_error(Some(400), 0)).await;
80        assert!(matches!(action, ErrorAction::Fail { .. }));
81    }
82}