1use std::time::Duration;
4
5use serde::{Deserialize, Serialize};
6
7#[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#[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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
44pub struct ProviderClassification {
45 pub error_type: UpstreamErrorKind,
46 pub code: UpstreamErrorKind,
47 pub http_status: Option<u16>,
49}
50
51#[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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
73pub struct RetryAccounting {
74 pub attempts: u64,
77 pub completed_backoff: Duration,
80 pub elapsed: Duration,
82}
83
84#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
86pub struct RetryProgress {
87 pub route: ProviderRoute,
88 pub reason: ProviderFailureReason,
89 pub upstream: ProviderClassification,
90 pub accounting: RetryAccounting,
93 pub next_delay: Duration,
94}
95
96#[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 pub last_attempt_reason: Option<ProviderFailureReason>,
105 pub upstream: ProviderClassification,
106 pub accounting: RetryAccounting,
107}
108
109#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
115#[non_exhaustive]
116pub enum ProviderRetryEvent {
117 Scheduled(RetryProgress),
118 Stopped(ProviderFailure),
120 Succeeded {
122 route: ProviderRoute,
123 accounting: RetryAccounting,
124 },
125}
126
127pub 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(¤t).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}