1use std::time::Duration;
4
5use fast_rands::Rand;
6use pyo3::prelude::*;
7
8use crate::streaming::StreamError;
9
10#[non_exhaustive]
12#[derive(Debug, Clone, thiserror::Error)]
13pub enum Error {
14 #[error("Agent is not started or has been shut down")]
16 AgentNotStarted,
17 #[error("Backend error: {message}")]
19 BackendError {
20 message: String,
22 },
23
24 #[error("Connection error: {message}")]
26 ConnectionError {
27 message: String,
29 },
30
31 #[error("Quota exceeded, retry after {retry_after:?}")]
33 QuotaExceeded {
34 retry_after: Duration,
36 },
37
38 #[error("Channel closed: {message}")]
40 ChannelClosed {
41 message: String,
43 },
44
45 #[error("Connection permanently closed: {message}")]
47 ConnectionClosed {
48 message: String,
50 },
51
52 #[error("Timeout after {duration:?}: {operation}")]
54 Timeout {
55 duration: Duration,
57 operation: String,
59 },
60
61 #[error(transparent)]
63 Stream(StreamError),
64
65 #[error("Invalid configuration: {message}")]
67 InvalidConfig {
68 message: String,
70 },
71
72 #[error("I/O error: {message}")]
74 Io {
75 message: String,
77 kind: std::io::ErrorKind,
79 },
80}
81
82impl Error {
83 #[must_use]
95 pub fn is_retryable(&self) -> bool {
96 match self {
97 Self::ConnectionError { .. } | Self::QuotaExceeded { .. } => true,
98 Self::BackendError { message } => message.contains("503"),
99 Self::Stream(se) => se.message.contains("503") || se.message.contains("429"),
100 _ => false,
101 }
102 }
103
104 #[must_use]
110 pub fn is_quota_error(&self) -> bool {
111 match self {
112 Self::QuotaExceeded { .. } => true,
113 Self::BackendError { message } => {
114 message.contains("429")
115 || message.contains("503")
116 || message.contains("RESOURCE_EXHAUSTED")
117 }
118 Self::Stream(se) => {
119 se.message.contains("429")
120 || se.message.contains("503")
121 || se.message.contains("quota")
122 || se.message.contains("RESOURCE_EXHAUSTED")
123 }
124 _ => false,
125 }
126 }
127}
128
129impl From<std::io::Error> for Error {
139 fn from(err: std::io::Error) -> Self {
140 Self::Io {
141 message: err.to_string(),
142 kind: err.kind(),
143 }
144 }
145}
146
147impl From<StreamError> for Error {
148 fn from(err: StreamError) -> Self {
149 Self::Stream(err)
150 }
151}
152
153#[doc(hidden)]
154impl From<PyErr> for Error {
155 fn from(err: PyErr) -> Self {
156 Python::attach(|py| classify_py_error(py, &err))
157 }
158}
159
160#[doc(hidden)]
161impl From<Error> for PyErr {
162 fn from(err: Error) -> Self {
163 pyo3::exceptions::PyRuntimeError::new_err(err.to_string())
164 }
165}
166
167pub(crate) fn classify_py_error(py: Python<'_>, err: &PyErr) -> Error {
173 if let Some(classified) = check_antigravity_error(py, err) {
174 return classified;
175 }
176 if let Some(classified) = check_pydantic_error(py, err) {
177 return classified;
178 }
179 if let Some(classified) = check_builtin_error(py, err) {
180 return classified;
181 }
182
183 let message = format_backend_error(py, err);
184 Error::BackendError { message }
185}
186
187fn check_antigravity_error(py: Python<'_>, err: &PyErr) -> Option<Error> {
188 match py.import("google.antigravity.types") {
189 Ok(types_mod) => {
190 if let Ok(conn_err_cls) = types_mod.getattr("AntigravityConnectionError")
192 && err.is_instance(py, &conn_err_cls)
193 {
194 return Some(Error::ConnectionError {
195 message: err.to_string(),
196 });
197 }
198 if let Ok(val_err_cls) = types_mod.getattr("AntigravityValidationError")
200 && err.is_instance(py, &val_err_cls)
201 {
202 return Some(Error::BackendError {
203 message: err.to_string(),
204 });
205 }
206 }
207 Err(import_err) => {
208 tracing::debug!(
209 error = %import_err,
210 "antigravity.types not available, skipping AntigravityError classification"
211 );
212 }
213 }
214 None
215}
216
217fn check_pydantic_error(py: Python<'_>, err: &PyErr) -> Option<Error> {
218 match py.import("pydantic") {
219 Ok(pydantic) => {
220 if let Ok(validation_err_cls) = pydantic.getattr("ValidationError")
222 && err.is_instance(py, &validation_err_cls)
223 {
224 return Some(Error::BackendError {
225 message: err.to_string(),
226 });
227 }
228 }
229 Err(import_err) => {
230 tracing::debug!(
231 error = %import_err,
232 "pydantic not available, skipping ValidationError classification"
233 );
234 }
235 }
236 None
237}
238
239fn check_builtin_error(py: Python<'_>, err: &PyErr) -> Option<Error> {
240 if let Ok(builtins) = py.import("builtins") {
242 if let Ok(import_err_cls) = builtins.getattr("ImportError")
244 && err.is_instance(py, &import_err_cls)
245 {
246 return Some(Error::BackendError {
247 message: err.to_string(),
248 });
249 }
250 } else {
251 tracing::warn!("Failed to import Python builtins module, skipping ImportError check");
252 }
253 None
254}
255
256fn format_backend_error(py: Python<'_>, err: &PyErr) -> String {
258 let formatted = py
260 .import("traceback")
261 .and_then(|tb_mod| {
262 tb_mod.call_method1(
263 "format_exception",
264 (err.get_type(py), err.value(py), err.traceback(py)),
265 )
266 })
267 .and_then(|lines| lines.extract::<Vec<String>>());
268
269 match formatted {
270 Ok(lines) => lines.join(""),
271 Err(fmt_err) => {
272 tracing::warn!(error = %fmt_err, "Failed to format backend traceback, using fallback");
273 let traceback = err.traceback(py);
275 traceback.as_ref().map_or_else(
276 || err.to_string(),
277 |tb| {
278 tb.format().map_or_else(
279 |tb_fmt_err| {
280 tracing::warn!(error = %tb_fmt_err, "Failed to format Python traceback");
281 err.to_string()
282 },
283 |tb_str| format!("{}\nTraceback:\n{}", err.value(py), tb_str),
284 )
285 },
286 )
287 }
288 }
289}
290
291pub async fn with_timeout<F, T>(timeout: Duration, operation: &str, f: F) -> Result<T, Error>
299where
300 F: std::future::Future<Output = Result<T, Error>>,
301{
302 match tokio::time::timeout(timeout, f).await {
303 Ok(result) => result,
304 Err(_elapsed) => Err(Error::Timeout {
305 duration: timeout,
306 operation: operation.to_string(),
307 }),
308 }
309}
310
311pub async fn with_retry<F, Fut, T>(max_retries: u32, operation: &str, mut f: F) -> Result<T, Error>
328where
329 F: FnMut() -> Fut,
330 Fut: std::future::Future<Output = Result<T, Error>>,
331{
332 let mut attempt = 0u32;
333 loop {
334 match f().await {
335 Ok(val) => return Ok(val),
336 Err(Error::ConnectionError { ref message }) => {
337 attempt += 1;
338 if attempt > max_retries {
339 tracing::error!(
340 attempts = attempt,
341 operation,
342 "All retries exhausted for connection error: {message}"
343 );
344 return Err(Error::ConnectionError {
345 message: message.clone(),
346 });
347 }
348 let backoff = backoff_duration(attempt);
349 tracing::warn!(
350 attempt,
351 max_retries,
352 backoff_ms = u64::try_from(backoff.as_millis()).unwrap_or_else(|e| {
353 tracing::warn!("Int conversion failed: {}", e);
354 u64::MAX
355 }),
356 operation,
357 "Connection error, retrying: {message}"
358 );
359 tokio::time::sleep(backoff).await;
360 }
361 Err(other) => return Err(other),
362 }
363 }
364}
365
366pub(crate) const MAX_BACKOFF_SECS: u64 = 120;
367
368const BACKOFF_EXPONENT_BASE: u64 = 2;
370const MILLISECONDS_PER_SECOND: u64 = 1000;
372const JITTER_TOTAL_SPREAD_DIVISOR: u64 = 2;
374const JITTER_MIN_SUBTRACT_DIVISOR: u64 = 4;
376
377pub(crate) fn backoff_duration(attempt: u32) -> Duration {
386 let attempt = attempt.max(1);
387 let base_secs = BACKOFF_EXPONENT_BASE
388 .checked_shl(attempt.saturating_sub(1))
389 .unwrap_or(MAX_BACKOFF_SECS)
390 .min(MAX_BACKOFF_SECS);
391 let base_ms = base_secs.saturating_mul(MILLISECONDS_PER_SECOND);
392 let jitter_range = base_ms / JITTER_TOTAL_SPREAD_DIVISOR; let jitter_min = base_ms.saturating_sub(base_ms / JITTER_MIN_SUBTRACT_DIVISOR);
395 let jittered_ms = if jitter_range == 0 {
396 base_ms
397 } else {
398 let limit = u32::try_from(jitter_range).unwrap_or_else(|e| {
399 tracing::warn!("Int conversion failed: {}", e);
400 u32::MAX
401 });
402 jitter_min
403 + (fast_rands::StdRand::new().between(0, limit.saturating_sub(1) as usize) as u64)
404 };
405 Duration::from_millis(jittered_ms)
406}
407
408#[cfg(test)]
409mod tests {
410 use std::sync::atomic::{AtomicU32, Ordering};
411
412 use super::*;
413
414 #[test]
415 fn test_stream_error_conversion() {
416 let safety_err = StreamError {
419 message: "Step error (status=ERROR): Candidate blocked by safety".to_string(),
420 };
421 let mapped_safety = Error::from(safety_err);
422 assert!(
423 matches!(mapped_safety, Error::Stream(_)),
424 "StreamError with 'safety' should pass through as Error::Stream"
425 );
426
427 let max_tokens_err = StreamError {
428 message: "Step error (status=ERROR): Max tokens reached".to_string(),
429 };
430 let mapped_max_tokens = Error::from(max_tokens_err);
431 assert!(
432 matches!(mapped_max_tokens, Error::Stream(_)),
433 "StreamError with 'max tokens' should pass through as Error::Stream"
434 );
435
436 let other_err = StreamError {
437 message: "Some other connection issue".to_string(),
438 };
439 let mapped_other = Error::from(other_err);
440 match mapped_other {
441 Error::Stream(e) => {
442 assert_eq!(e.message, "Some other connection issue");
443 }
444 other => panic!("Expected Error::Stream, got: {other:?}"),
445 }
446 }
447
448 #[test]
449 fn test_backend_error_from_pyerr() {
450 Python::initialize();
451 let err = Python::attach(|py| {
452 let result: PyResult<()> = py.run(c"raise ValueError('test error 42')", None, None);
453 result.unwrap_err()
454 });
455
456 let bridge_err: Error = err.into();
457 match &bridge_err {
458 Error::BackendError { message } => {
459 assert!(
460 message.contains("ValueError"),
461 "Expected 'ValueError' in message, got: {message}"
462 );
463 assert!(
464 message.contains("test error 42"),
465 "Expected 'test error 42' in message, got: {message}"
466 );
467 }
468 other => panic!("Expected BackendError, got: {other:?}"),
469 }
470 }
471
472 #[tokio::test]
473 async fn test_timeout_triggers() {
474 let short_timeout = Duration::from_millis(50);
475 let result: Result<(), Error> = with_timeout(short_timeout, "test_op", async {
476 tokio::time::sleep(Duration::from_secs(10)).await;
477 Ok(())
478 })
479 .await;
480
481 match result {
482 Err(Error::Timeout {
483 duration,
484 operation,
485 }) => {
486 assert_eq!(duration, short_timeout);
487 assert_eq!(operation, "test_op");
488 }
489 other => panic!("Expected Timeout, got: {other:?}"),
490 }
491 }
492
493 #[tokio::test]
494 async fn test_timeout_succeeds_when_fast() {
495 let result = with_timeout(Duration::from_secs(5), "fast_op", async { Ok(42) }).await;
496 assert_eq!(result.unwrap(), 42);
497 }
498
499 #[tokio::test]
500 async fn test_retry_succeeds_after_failures() {
501 let counter = AtomicU32::new(0);
502 let result = with_retry(3, "test_retry", || {
503 let attempt = counter.fetch_add(1, Ordering::SeqCst);
504 async move {
505 if attempt < 2 {
506 Err(Error::ConnectionError {
507 message: "transient".to_string(),
508 })
509 } else {
510 Ok(42)
511 }
512 }
513 })
514 .await;
515
516 assert_eq!(result.unwrap(), 42);
517 assert_eq!(counter.load(Ordering::SeqCst), 3);
518 }
519
520 #[tokio::test]
521 async fn test_retry_exhausted() {
522 let counter = AtomicU32::new(0);
523 let result: Result<i32, Error> = with_retry(2, "doomed", || {
524 counter.fetch_add(1, Ordering::SeqCst);
525 async {
526 Err(Error::ConnectionError {
527 message: "always fails".to_string(),
528 })
529 }
530 })
531 .await;
532
533 assert!(matches!(result, Err(Error::ConnectionError { .. })));
534 assert_eq!(counter.load(Ordering::SeqCst), 3);
536 }
537
538 #[tokio::test]
539 async fn test_retry_does_not_retry_non_connection_errors() {
540 let counter = AtomicU32::new(0);
541 let result: Result<i32, Error> = with_retry(5, "python_err", || {
542 counter.fetch_add(1, Ordering::SeqCst);
543 async {
544 Err(Error::BackendError {
545 message: "kaboom".to_string(),
546 })
547 }
548 })
549 .await;
550
551 assert!(matches!(result, Err(Error::BackendError { .. })));
552 assert_eq!(counter.load(Ordering::SeqCst), 1);
553 }
554
555 #[test]
556 fn test_backoff_duration_progression() {
557 let bases_ms: [(u32, u64); 6] = [
559 (1, 2_000),
560 (2, 4_000),
561 (3, 8_000),
562 (4, 16_000),
563 (7, 120_000), (100, 120_000), ];
566 for (attempt, base_ms) in bases_ms {
567 let d = backoff_duration(attempt);
568 let lo = base_ms * 3 / 4;
569 let hi = base_ms * 5 / 4;
570 assert!(
571 d.as_millis() >= u128::from(lo) && d.as_millis() <= u128::from(hi),
572 "backoff_duration({attempt}) = {d:?} outside [{lo}ms, {hi}ms]"
573 );
574 }
575 }
576
577 #[test]
578 fn test_error_display_messages() {
579 let err = Error::BackendError {
580 message: "test".to_string(),
581 };
582 assert_eq!(format!("{err}"), "Backend error: test");
583
584 let err = Error::ConnectionError {
585 message: "lost".to_string(),
586 };
587 assert_eq!(format!("{err}"), "Connection error: lost");
588
589 let err = Error::QuotaExceeded {
590 retry_after: Duration::from_secs(5),
591 };
592 assert!(format!("{err}").contains("5s"));
593
594 let err = Error::ChannelClosed {
595 message: "cmd".to_string(),
596 };
597 assert_eq!(format!("{err}"), "Channel closed: cmd");
598
599 let err = Error::Timeout {
600 duration: Duration::from_secs(30),
601 operation: "chat".to_string(),
602 };
603 assert!(format!("{err}").contains("chat"));
604 }
605
606 #[test]
607 fn test_backoff_duration_zero_attempt() {
608 let d = backoff_duration(0);
610 assert!(
611 d.as_millis() >= 1500 && d.as_millis() <= 2500,
612 "backoff_duration(0) = {d:?} outside [1500ms, 2500ms]"
613 );
614 }
615
616 #[test]
617 fn test_backoff_duration_large_attempt_capped() {
618 let d = backoff_duration(u32::MAX);
620 assert!(
621 d.as_millis() >= 90_000 && d.as_millis() <= 150_000,
622 "backoff_duration(u32::MAX) = {d:?} outside [90s, 150s]"
623 );
624 }
625
626 #[tokio::test]
627 async fn test_timeout_propagates_inner_error() {
628 let result: Result<(), Error> = with_timeout(Duration::from_secs(10), "inner_err", async {
629 Err(Error::BackendError {
630 message: "inner failure".to_string(),
631 })
632 })
633 .await;
634
635 match result {
636 Err(Error::BackendError { message }) => {
637 assert_eq!(message, "inner failure");
638 }
639 other => panic!("Expected BackendError, got: {other:?}"),
640 }
641 }
642
643 #[tokio::test]
644 async fn test_retry_zero_max_retries_still_runs_once() {
645 let counter = AtomicU32::new(0);
646 let result: Result<i32, Error> = with_retry(0, "no_retries", || {
647 counter.fetch_add(1, Ordering::SeqCst);
648 async {
649 Err(Error::ConnectionError {
650 message: "fail".to_string(),
651 })
652 }
653 })
654 .await;
655
656 assert!(matches!(result, Err(Error::ConnectionError { .. })));
657 assert_eq!(counter.load(Ordering::SeqCst), 1);
659 }
660
661 #[tokio::test]
662 async fn test_retry_succeeds_on_first_attempt() {
663 let counter = AtomicU32::new(0);
664 let result = with_retry(5, "instant_success", || {
665 counter.fetch_add(1, Ordering::SeqCst);
666 async { Ok(99) }
667 })
668 .await;
669
670 assert_eq!(result.unwrap(), 99);
671 assert_eq!(counter.load(Ordering::SeqCst), 1);
672 }
673
674 #[tokio::test]
675 async fn test_retry_quota_exceeded_not_retried() {
676 let counter = AtomicU32::new(0);
677 let result: Result<i32, Error> = with_retry(5, "quota", || {
678 counter.fetch_add(1, Ordering::SeqCst);
679 async {
680 Err(Error::QuotaExceeded {
681 retry_after: Duration::from_secs(1),
682 })
683 }
684 })
685 .await;
686
687 assert!(matches!(result, Err(Error::QuotaExceeded { .. })));
688 assert_eq!(counter.load(Ordering::SeqCst), 1);
690 }
691
692 #[tokio::test]
693 async fn test_retry_timeout_not_retried() {
694 let counter = AtomicU32::new(0);
695 let result: Result<i32, Error> = with_retry(5, "timeout", || {
696 counter.fetch_add(1, Ordering::SeqCst);
697 async {
698 Err(Error::Timeout {
699 duration: Duration::from_secs(10),
700 operation: "test".to_string(),
701 })
702 }
703 })
704 .await;
705
706 assert!(matches!(result, Err(Error::Timeout { .. })));
707 assert_eq!(counter.load(Ordering::SeqCst), 1);
708 }
709
710 #[tokio::test]
711 async fn test_retry_channel_closed_not_retried() {
712 let counter = AtomicU32::new(0);
713 let result: Result<i32, Error> = with_retry(5, "channel", || {
714 counter.fetch_add(1, Ordering::SeqCst);
715 async {
716 Err(Error::ChannelClosed {
717 message: "gone".to_string(),
718 })
719 }
720 })
721 .await;
722
723 assert!(matches!(result, Err(Error::ChannelClosed { .. })));
724 assert_eq!(counter.load(Ordering::SeqCst), 1);
725 }
726
727 #[test]
728 fn test_error_debug_format() {
729 let err = Error::BackendError {
730 message: "debug test".to_string(),
731 };
732 let debug = format!("{err:?}");
733 assert!(debug.contains("BackendError"));
734 assert!(debug.contains("debug test"));
735 }
736
737 #[test]
738 fn test_backoff_duration_full_progression() {
739 let base_secs: [u64; 8] = [2, 4, 8, 16, 32, 64, 120, 120];
741 for (i, base) in base_secs.iter().enumerate() {
742 let attempt = u32::try_from(i + 1).unwrap();
743 let d = backoff_duration(attempt);
744 let base_ms = base * 1000;
745 let lo = base_ms * 3 / 4;
746 let hi = base_ms * 5 / 4;
747 assert!(
748 d.as_millis() >= u128::from(lo) && d.as_millis() <= u128::from(hi),
749 "backoff_duration({attempt}) = {d:?} outside [{lo}ms, {hi}ms]"
750 );
751 }
752 }
753
754 #[test]
755 fn test_stream_error_from_conversion() {
756 let stream_err = StreamError {
757 message: "connection reset".to_string(),
758 };
759 let bridge_err = Error::from(stream_err);
760 match &bridge_err {
761 Error::Stream(inner) => {
762 assert_eq!(inner.message, "connection reset");
763 }
764 other => panic!("Expected Stream variant, got: {other:?}"),
765 }
766 }
767
768 #[test]
769 fn test_stream_error_display_through_bridge() {
770 let stream_err = StreamError {
771 message: "quota exceeded".to_string(),
772 };
773 let bridge_err = Error::from(stream_err);
774 let display = format!("{bridge_err}");
775 assert!(
776 display.contains("quota exceeded"),
777 "Expected 'quota exceeded' in display, got: {display}"
778 );
779 }
780
781 #[test]
782 fn test_is_retryable_connection_error() {
783 let err = Error::ConnectionError {
784 message: "timeout".to_string(),
785 };
786 assert!(err.is_retryable());
787 }
788
789 #[test]
790 fn test_quota_exceeded_is_retryable() {
791 let err = Error::QuotaExceeded {
792 retry_after: Duration::from_secs(5),
793 };
794 assert!(err.is_retryable());
795 }
796
797 #[test]
798 fn test_is_not_retryable_backend_error() {
799 let err = Error::BackendError {
800 message: "kaboom".to_string(),
801 };
802 assert!(!err.is_retryable());
803 }
804
805 #[test]
806 fn test_is_not_retryable_channel_closed() {
807 let err = Error::ChannelClosed {
808 message: "gone".to_string(),
809 };
810 assert!(!err.is_retryable());
811 }
812
813 #[test]
814 fn test_is_not_retryable_timeout() {
815 let err = Error::Timeout {
816 duration: Duration::from_secs(30),
817 operation: "chat".to_string(),
818 };
819 assert!(!err.is_retryable());
820 }
821
822 #[test]
823 fn test_is_not_retryable_stream() {
824 let err = Error::Stream(StreamError {
825 message: "stream failed".to_string(),
826 });
827 assert!(!err.is_retryable());
828 }
829
830 #[test]
831 fn test_is_retryable_503_backend_error() {
832 let err = Error::BackendError {
833 message: "request failed (code 503): high demand".to_string(),
834 };
835 assert!(err.is_retryable());
836 }
837
838 #[test]
839 fn test_is_quota_error_quota_exceeded() {
840 let err = Error::QuotaExceeded {
841 retry_after: Duration::from_secs(5),
842 };
843 assert!(err.is_quota_error());
844 }
845
846 #[test]
847 fn test_is_quota_error_backend_429() {
848 let err = Error::BackendError {
849 message: "HTTP 429 Too Many Requests".to_string(),
850 };
851 assert!(err.is_quota_error());
852 }
853
854 #[test]
855 fn test_is_quota_error_resource_exhausted() {
856 let err = Error::BackendError {
857 message: "RESOURCE_EXHAUSTED: quota exceeded".to_string(),
858 };
859 assert!(err.is_quota_error());
860 }
861
862 #[test]
863 fn test_is_not_quota_error_connection() {
864 let err = Error::ConnectionError {
865 message: "timeout".to_string(),
866 };
867 assert!(!err.is_quota_error());
868 }
869
870 #[test]
871 fn test_is_not_quota_error_normal_backend() {
872 let err = Error::BackendError {
873 message: "something else".to_string(),
874 };
875 assert!(!err.is_quota_error());
876 }
877
878 #[test]
879 fn test_is_quota_error_503_high_demand() {
880 let err = Error::BackendError {
881 message: "request failed (code 503): This model is currently experiencing high demand"
882 .to_string(),
883 };
884 assert!(err.is_quota_error());
885 }
886}