Skip to main content

agentkit_loop/
retry.rs

1//! Sanitized, payload-free model retry observations. These are not effects provenance.
2
3use std::time::Duration;
4
5use serde::{Deserialize, Serialize};
6
7/// Static provider route; never an endpoint URL or account identifier.
8#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
9#[serde(rename_all = "snake_case")]
10#[non_exhaustive]
11pub enum ProviderRoute {
12    #[default]
13    Unknown,
14    OpenAiResponses,
15    OpenAiChatGptResponses,
16}
17
18/// Allowlisted provider type/code values. Unknown strings are never retained.
19#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
20#[serde(rename_all = "snake_case")]
21#[non_exhaustive]
22pub enum UpstreamErrorKind {
23    ServiceUnavailableError,
24    ServerIsOverloaded,
25    ServerError,
26    RateLimitError,
27    RateLimitExceeded,
28    TemporarilyUnavailable,
29    AuthenticationError,
30    InvalidApiKey,
31    InvalidAuthentication,
32    Unauthorized,
33    InvalidRequestError,
34    PermissionDenied,
35    InsufficientQuota,
36    ContentPolicyViolation,
37    #[default]
38    #[serde(other)]
39    Unknown,
40}
41
42/// Sanitized source classification, kept separate from the local stopping reason.
43#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
44pub struct ProviderClassification {
45    pub error_type: UpstreamErrorKind,
46    pub code: UpstreamErrorKind,
47    /// Source HTTP status, if present. No headers or response body are retained.
48    pub http_status: Option<u16>,
49}
50
51/// Local reason for a failed attempt or logical request.
52#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
53#[serde(rename_all = "snake_case")]
54#[non_exhaustive]
55pub enum ProviderFailureReason {
56    HttpStatus,
57    Transport,
58    ResponseFailed,
59    Protocol,
60    InvalidRequest,
61    Authentication,
62    AttemptTimeout,
63    IdleTimeout,
64    RetryExhausted,
65    RetryBudget,
66    RetryDisabled,
67    ReplayUnsafe,
68    Cancelled,
69}
70
71/// Per-logical-request accounting, independent of policy retry count.
72#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
73pub struct RetryAccounting {
74    /// Actual HTTP sends started, including a resend after authentication refresh.
75    /// Preflight/authentication failures can have zero attempts.
76    pub attempts: u64,
77    /// Sum of requested durations of fully completed backoff waits. Interrupted
78    /// waits contribute zero, even when they consumed wall-clock time.
79    pub completed_backoff: Duration,
80    /// Monotonic elapsed time since before initial authentication/preflight.
81    pub elapsed: Duration,
82}
83
84/// A nonterminal snapshot emitted before a retry wait or reactive refresh.
85#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
86pub struct RetryProgress {
87    pub route: ProviderRoute,
88    pub reason: ProviderFailureReason,
89    pub upstream: ProviderClassification,
90    /// Attempts already started. The planned next send is `attempts + 1`, but
91    /// cancellation/preflight failure can prevent it from ever starting.
92    pub accounting: RetryAccounting,
93    pub next_delay: Duration,
94}
95
96/// Payload-free terminal model failure. Display and Debug contain only typed data.
97#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, thiserror::Error)]
98#[error("provider request failed ({reason:?}; attempts: {attempts})", attempts = .accounting.attempts)]
99pub struct ProviderFailure {
100    pub route: ProviderRoute,
101    pub reason: ProviderFailureReason,
102    /// Last failed request-attempt category, retained across local budget/limit stops.
103    /// None when no request attempt failed (for example initial authentication).
104    pub last_attempt_reason: Option<ProviderFailureReason>,
105    pub upstream: ProviderClassification,
106    pub accounting: RetryAccounting,
107}
108
109/// Observational lifecycle; never a second model result or a tool-effects record.
110///
111/// Correlate through the enclosing `ObservedEvent.session_id` and current
112/// `AgentEvent::TurnStarted`. Direct session consumers own that association.
113/// Stable fatal event IDs belong to the host, not to this payload.
114#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
115#[non_exhaustive]
116pub enum ProviderRetryEvent {
117    Scheduled(RetryProgress),
118    /// Emitted once on explicit failure/cancellation, including zero-send failures.
119    Stopped(ProviderFailure),
120    /// Clears retry activity without introducing another successful model result.
121    Succeeded {
122        route: ProviderRoute,
123        accounting: RetryAccounting,
124    },
125}
126
127/// Synchronous, queue-free observer installed before `begin_turn`.
128///
129/// Implementations must not block or re-enter the session. Panics propagate, just
130/// like loop observers; delivery cannot be guaranteed if observers panic/block.
131/// Dropping a future or turn is not explicit cancellation and does not promise a
132/// terminal observation. Implementations should rate-limit Scheduled snapshots
133/// per logical turn, retaining exact accounting and unsuppressed terminal events.
134pub trait RetryObserver: Send + Sync {
135    fn on_retry_event(&self, event: ProviderRetryEvent);
136}
137
138impl<F: Fn(ProviderRetryEvent) + Send + Sync> RetryObserver for F {
139    fn on_retry_event(&self, event: ProviderRetryEvent) {
140        self(event);
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use crate::AgentEvent;
148
149    #[test]
150    fn event_roundtrip_and_legacy_shape_remain_compatible() {
151        let old = r#"{"RunStarted":{"session_id":"session"}}"#;
152        let event: AgentEvent = serde_json::from_str(old).unwrap();
153        assert_eq!(serde_json::to_string(&event).unwrap(), old);
154        let current = AgentEvent::ProviderRetry(ProviderRetryEvent::Stopped(ProviderFailure {
155            route: ProviderRoute::OpenAiChatGptResponses,
156            reason: ProviderFailureReason::RetryExhausted,
157            last_attempt_reason: Some(ProviderFailureReason::ResponseFailed),
158            upstream: ProviderClassification {
159                error_type: UpstreamErrorKind::ServiceUnavailableError,
160                code: UpstreamErrorKind::ServerIsOverloaded,
161                http_status: None,
162            },
163            accounting: RetryAccounting {
164                attempts: 3,
165                completed_backoff: Duration::from_millis(125),
166                elapsed: Duration::from_millis(250),
167            },
168        }));
169        let encoded = serde_json::to_string(&current).unwrap();
170        assert_eq!(
171            serde_json::from_str::<AgentEvent>(&encoded).unwrap(),
172            current
173        );
174        assert!(encoded.contains("service_unavailable_error"));
175        assert!(encoded.contains("server_is_overloaded"));
176        assert!(serde_json::from_str::<AgentEvent>(r#"{"ProviderRetry":{"Stopped":{}}}"#).is_err());
177        assert_eq!(
178            serde_json::from_str::<UpstreamErrorKind>(r#""future-private-value""#).unwrap(),
179            UpstreamErrorKind::Unknown
180        );
181    }
182}
183
184#[cfg(test)]
185mod cancellation_hook_tests {
186    use crate::{
187        Agent, LoopError, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, SessionConfig,
188        TurnRequest,
189    };
190    use agentkit_core::{CancellationController, Item, ItemKind, TurnCancellation, Usage};
191    use std::sync::{
192        Arc,
193        atomic::{AtomicUsize, Ordering},
194    };
195
196    #[derive(Clone)]
197    struct CancellingModel {
198        controller: Arc<CancellationController>,
199        notified: Arc<AtomicUsize>,
200    }
201
202    #[async_trait::async_trait]
203    impl ModelAdapter for CancellingModel {
204        type Session = Self;
205        async fn start_session(&self, _: SessionConfig) -> Result<Self, LoopError> {
206            Ok(self.clone())
207        }
208    }
209
210    #[async_trait::async_trait]
211    impl ModelSession for CancellingModel {
212        type Turn = Self;
213        async fn begin_turn(
214            &mut self,
215            _: TurnRequest,
216            _: Option<TurnCancellation>,
217        ) -> Result<Self, LoopError> {
218            Ok(self.clone())
219        }
220    }
221
222    #[async_trait::async_trait]
223    impl ModelTurn for CancellingModel {
224        fn on_cancelled(&mut self) {
225            self.notified.fetch_add(1, Ordering::SeqCst);
226        }
227        async fn next_event(
228            &mut self,
229            _: Option<TurnCancellation>,
230        ) -> Result<Option<ModelTurnEvent>, LoopError> {
231            self.controller.interrupt();
232            Ok(Some(ModelTurnEvent::Usage(Usage::default())))
233        }
234    }
235
236    #[tokio::test]
237    async fn driver_notifies_before_dropping_a_turn_cancelled_between_events() {
238        let controller = Arc::new(CancellationController::new());
239        let notified = Arc::new(AtomicUsize::new(0));
240        let agent = Agent::builder()
241            .model(CancellingModel {
242                controller: controller.clone(),
243                notified: notified.clone(),
244            })
245            .cancellation(controller.handle())
246            .input(vec![Item::text(ItemKind::User, "hello")])
247            .build()
248            .unwrap();
249        let mut driver = agent.start(SessionConfig::new("session")).await.unwrap();
250        driver.next().await.unwrap();
251        assert_eq!(notified.load(Ordering::SeqCst), 1);
252    }
253}