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 ContextLengthExceeded,
39 #[default]
40 #[serde(other)]
41 Unknown,
42}
43
44#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
46pub struct ProviderClassification {
47 pub error_type: UpstreamErrorKind,
48 pub code: UpstreamErrorKind,
49 pub http_status: Option<u16>,
51}
52
53#[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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
75pub struct RetryAccounting {
76 pub attempts: u64,
79 pub completed_backoff: Duration,
82 pub elapsed: Duration,
84}
85
86#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
88pub struct RetryProgress {
89 pub route: ProviderRoute,
90 pub reason: ProviderFailureReason,
91 pub upstream: ProviderClassification,
92 pub accounting: RetryAccounting,
95 pub next_delay: Duration,
96}
97
98#[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 pub last_attempt_reason: Option<ProviderFailureReason>,
107 pub upstream: ProviderClassification,
108 pub accounting: RetryAccounting,
109}
110
111#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
117#[non_exhaustive]
118pub enum ProviderRetryEvent {
119 Scheduled(RetryProgress),
120 Stopped(ProviderFailure),
122 Succeeded {
124 route: ProviderRoute,
125 accounting: RetryAccounting,
126 },
127}
128
129pub 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(¤t).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}