Skip to main content

oino_core/session/
retry.rs

1//! Automatic retry of transient provider errors (`_prepareRetry`, `_finishCancelledRetry`).
2
3use oino_ai::AssistantMessage;
4use tokio_util::sync::CancellationToken;
5
6use super::{AgentSession, lock};
7use crate::compaction::retry_delay_ms;
8use crate::events::AgentSessionEvent;
9
10impl AgentSession {
11    /// F i n i s h   C a n c e l l e d   R e t r y.
12    pub(crate) fn finish_cancelled_retry(&self) {
13        let attempt = std::mem::take(&mut lock(&self.inner.book).retry_attempt);
14        if attempt == 0 {
15            return;
16        }
17        self.emit(AgentSessionEvent::AutoRetryEnd {
18            success: false,
19            attempt,
20            final_error: Some("Retry cancelled".to_owned()),
21        });
22    }
23
24    /// Schedules a retry with exponential backoff; `false` when the budget
25    /// is spent, retries are disabled or the wait was cancelled.
26    pub(crate) async fn prepare_retry(&self, message: &AssistantMessage) -> bool {
27        let settings = lock(&self.inner.settings).settings().retry_settings();
28        if !settings.enabled {
29            return false;
30        }
31        let attempt = {
32            let mut book = lock(&self.inner.book);
33            book.retry_attempt += 1;
34            if u64::from(book.retry_attempt) > settings.max_retries {
35                book.retry_attempt -= 1;
36                return false;
37            }
38            book.retry_attempt
39        };
40        let delay_ms = retry_delay_ms(
41            settings.base_delay_ms,
42            Some(settings.max_agent_delay_ms),
43            attempt,
44        );
45        self.emit(AgentSessionEvent::AutoRetryStart {
46            attempt,
47            max_attempts: u32::try_from(settings.max_retries).unwrap_or(u32::MAX),
48            delay_ms,
49            error_message: message
50                .error_message
51                .clone()
52                .filter(|m| !m.is_empty())
53                .unwrap_or_else(|| "Unknown error".to_owned()),
54        });
55        self.omit_recovery_attempt(message);
56        let cancel = CancellationToken::new();
57        lock(&self.inner.cancels).retry = Some(cancel.clone());
58        let slept = tokio::select! {
59            () = cancel.cancelled() => false,
60            () = tokio::time::sleep(std::time::Duration::from_millis(delay_ms)) => true,
61        };
62        lock(&self.inner.cancels).retry = None;
63        if !slept {
64            self.finish_cancelled_retry();
65        }
66        slept
67    }
68
69    /// Enables or disables automatic retry (settings).
70    pub fn set_auto_retry_enabled(&self, enabled: bool) {
71        if let Err(error) = self.with_settings(|s| s.set_retry_enabled(enabled)) {
72            tracing::warn!("saving retry.enabled failed: {error}");
73        }
74    }
75}