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