mod.rs - DeepSeek-TUI-2026

返回 DeepSeek-TUI-2026

mod.rs

1	//! LLM Client Trait and Retry Logic
2	//!
3	//! This module provides a unified interface for LLM providers with robust retry logic,
4	//! exponential backoff, and proper error classification.
5	//!
6	//! # Architecture
7	//!
8	//! - `LlmClient` trait: Async interface for LLM providers (DeepSeek, `OpenAI`, etc.)
9	//! - `RetryConfig`: Configurable retry behavior with exponential backoff and jitter
10	//! - `LlmError`: Classified errors with retryability information
11
12	//! - `with_retry`: Generic retry wrapper for any async operation
13	//!
14	//! # Example
15	//!
16	//! ```ignore
17	//! use crate::llm_client::{LlmClient, RetryConfig, with_retry};
18	//!
19	//! let config = RetryConfig::default();
20	//! let result = with_retry(&config, \|\| async {
21	//! client.create_message(request).await
22	//! }, None).await;
23	//! ```
24
25	use crate::config::RetryPolicy;
26	use crate::models::{MessageRequest, MessageResponse, StreamEvent};
27	use anyhow::Result;
28	use std::future::Future;
29	use std::pin::Pin;
30	use std::time::{Duration, Instant};
31	use uuid::Uuid;
32
33	#[cfg(test)]
34	pub mod mock;
35
36	// === LlmClient Trait ===
37
38	/// Type alias for boxed stream of SSE events
39	pub type StreamEventBox =
40	Pin<Box<dyn futures_util::Stream<Item = Result<StreamEvent>> + Send + 'static>>;
41
42	/// Unified interface for LLM providers.
43	///
44	/// This trait abstracts over different LLM APIs (DeepSeek, `OpenAI`, etc.)
45	/// allowing the agent to work with any provider that implements this interface.
46	///
47	/// # Implementation Notes
48	///
49	/// - All methods are async and require `Send + Sync` for thread safety
50	/// - The `create_message_stream` method returns a pinned boxed stream for SSE
51	/// - Implementations should handle their own authentication and base URL configuration
52	#[allow(async_fn_in_trait, dead_code)] // Trait methods are part of the LLM provider interface
53	pub trait LlmClient: Send + Sync {
54	/// Returns the provider name (e.g., "openai", "deepseek")
55	fn provider_name(&self) -> &'static str;
56
57	/// Returns the model identifier being used
58	fn model(&self) -> &str;
59
60	/// Creates a non-streaming message completion
61	fn create_message(
62	&self,
63	request: MessageRequest,
64	) -> impl Future<Output = Result<MessageResponse>> + Send;
65
66	/// Creates a streaming message completion
67	///
68	/// Returns a stream of SSE events that should be consumed until completion.
69	async fn create_message_stream(&self, request: MessageRequest) -> Result<StreamEventBox>;
70
71	/// Optional health check to verify API connectivity
72	async fn health_check(&self) -> Result<bool> {
73	Ok(true)
74	}
75	}
76
77	/// Trait for clients that support configurable retry behavior
78	#[allow(dead_code)] // Part of LLM provider interface, will be used by additional providers
79	pub trait RetryConfigurable {
80	fn retry_config(&self) -> &RetryConfig;
81	fn set_retry_config(&mut self, config: RetryConfig);
82	}
83
84	// === LlmError - Classified Error Types ===
85
86	/// Classified LLM errors with retryability information.
87	///
88	/// This enum categorizes API errors to enable smart retry decisions.
89	/// Some errors (rate limits, transient server errors) are retryable,
90	/// while others (auth failures, invalid requests) should fail immediately.
91	#[derive(Debug)]
92	pub enum LlmError {
93	/// Rate limit exceeded (HTTP 429)
94	/// Contains optional Retry-After duration from server
95	RateLimited {
96	message: String,
97	retry_after: Option<Duration>,
98	},
99
100	/// Server error (HTTP 5xx)
101	ServerError { status: u16, message: String },
102
103	/// Network connectivity error
104	NetworkError(String),
105
106	/// Request timed out
107	Timeout(Duration),
108
109	/// Authentication failed (HTTP 401, 403)
110	AuthenticationError(String),
111
112	/// Invalid request parameters (HTTP 400)
113	InvalidRequest { status: u16, message: String },
114
115	/// Model-specific error (model not found, etc.)
116	ModelError(String),
117
118	/// Content policy violation (safety filters)
119	ContentPolicyError(String),
120
121	/// Failed to parse API response
122	ParseError(String),
123
124	/// Context length exceeded
125	ContextLengthError(String),
126
127	/// Catch-all for other errors
128	Other(String),
129	}
130
131	impl std::fmt::Display for LlmError {
132	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133	match self {
134	LlmError::RateLimited { message, .. } => write!(f, "Rate limit exceeded: {message}"),
135	LlmError::ServerError { status, message } => {
136	write!(f, "Server error ({status}): {message}")
137	}
138	LlmError::NetworkError(msg) => write!(f, "Network error: {msg}"),
139	LlmError::Timeout(d) => write!(f, "Request timed out after {d:?}"),
140	LlmError::AuthenticationError(msg) => write!(f, "Authentication failed: {msg}"),
141	LlmError::InvalidRequest { status, message } => {
142	write!(f, "Invalid request ({status}): {message}")
143	}
144	LlmError::ModelError(msg) => write!(f, "Model error: {msg}"),
145	LlmError::ContentPolicyError(msg) => write!(f, "Content policy violation: {msg}"),
146	LlmError::ParseError(msg) => write!(f, "Response parsing error: {msg}"),
147	LlmError::ContextLengthError(msg) => write!(f, "Context length exceeded: {msg}"),
148	LlmError::Other(msg) => write!(f, "LLM error: {msg}"),
149	}
150	}
151	}
152
153	impl std::error::Error for LlmError {}
154
155	impl LlmError {
156	/// Determines if this error is potentially transient and worth retrying.
157	///
158	/// Retryable errors:
159	/// - Rate limits (with backoff)
160	/// - Server errors (5xx)
161	/// - Network errors (connection issues)
162	/// - Timeouts
163	///
164	/// Non-retryable errors:
165	/// - Authentication failures
166	/// - Invalid requests
167	/// - Content policy violations
168	/// - Context length errors
169	pub fn is_retryable(&self) -> bool {
170	matches!(
171	self,
172	LlmError::RateLimited { .. }
173	\| LlmError::ServerError { .. }
174	\| LlmError::NetworkError(_)
175	\| LlmError::Timeout(_)
176	)
177	}
178
179	/// Returns the server-suggested retry delay if available.
180	///
181	/// This is typically present for rate limit errors when the server
182	/// provides a Retry-After header.
183	pub fn suggested_retry_delay(&self) -> Option<Duration> {
184	match self {
185	LlmError::RateLimited { retry_after, .. } => *retry_after,
186	_ => None,
187	}
188	}
189
190	/// Constructs an `LlmError` from HTTP status code and response body.
191	///
192	/// Performs heuristic classification based on:
193	/// - Status code (429 = rate limit, 401/403 = auth, 5xx = server error)
194	/// - Response body keywords (`context_length`, `content_policy`, safety, etc.)
195	pub fn from_http_response(status: u16, body: &str) -> Self {
196	match status {
197	429 => LlmError::RateLimited {
198	message: body.to_string(),
199	retry_after: None,
200	},
201	401 \| 403 => LlmError::AuthenticationError(body.to_string()),
202	400 => {
203	// Classify 400 errors by examining the response body
204	let body_lower = body.to_lowercase();
205	if body_lower.contains("context_length")
206	\|\| body_lower.contains("token")
207	\|\| body_lower.contains("too long")
208	\|\| body_lower.contains("maximum")
209	{
210	LlmError::ContextLengthError(body.to_string())
211	} else if body_lower.contains("content_policy")
212	\|\| body_lower.contains("safety")
213	\|\| body_lower.contains("harmful")
214	\|\| body_lower.contains("inappropriate")
215	{
216	LlmError::ContentPolicyError(body.to_string())
217	} else if body_lower.contains("model") && body_lower.contains("not found") {
218	LlmError::ModelError(body.to_string())
219	} else {
220	LlmError::InvalidRequest {
221	status,
222	message: body.to_string(),
223	}
224	}
225	}
226	404 => {
227	if body.to_lowercase().contains("model") {
228	LlmError::ModelError(body.to_string())
229	} else {
230	LlmError::InvalidRequest {
231	status,
232	message: body.to_string(),
233	}
234	}
235	}
236	500..=599 => LlmError::ServerError {
237	status,
238	message: body.to_string(),
239	},
240	_ => LlmError::Other(format!("HTTP {status}: {body}")),
241	}
242	}
243
244	/// Constructs an `LlmError` from HTTP status code, body, and optional Retry-After header.
245	pub fn from_http_response_with_retry_after(
246	status: u16,
247	body: &str,
248	retry_after: Option<Duration>,
249	) -> Self {
250	let mut error = Self::from_http_response(status, body);
251	if let LlmError::RateLimited {
252	retry_after: ref mut ra,
253	..
254	} = error
255	{
256	*ra = retry_after;
257	}
258	error
259	}
260
261	/// Constructs an `LlmError` from a reqwest error.
262	pub fn from_reqwest(err: &reqwest::Error) -> Self {
263	if err.is_timeout() {
264	LlmError::Timeout(Duration::from_secs(0))
265	} else if err.is_connect() {
266	LlmError::NetworkError(format!("Connection failed: {err}"))
267	} else if err.is_request() {
268	LlmError::NetworkError(format!("Request failed: {err}"))
269	} else {
270	LlmError::Other(err.to_string())
271	}
272	}
273	}
274
275	impl From<reqwest::Error> for LlmError {
276	fn from(err: reqwest::Error) -> Self {
277	LlmError::from_reqwest(&err)
278	}
279	}
280
281	impl From<serde_json::Error> for LlmError {
282	fn from(err: serde_json::Error) -> Self {
283	LlmError::ParseError(err.to_string())
284	}
285	}
286
287	// === RetryConfig - Exponential Backoff Configuration ===
288
289	/// Configuration for retry behavior with exponential backoff.
290	///
291	/// This struct controls how retries are performed:
292	/// - Number of retry attempts
293	/// - Delay calculation (exponential backoff with optional jitter)
294	/// - Which HTTP status codes are retryable
295	/// - Timeout handling
296	///
297	/// # Default Values
298	///
299	/// - `enabled`: true
300	/// - `max_retries`: 3
301	/// - `initial_delay`: 1.0 seconds
302	/// - `max_delay`: 60.0 seconds
303	/// - `exponential_base`: 2.0
304	/// - `jitter`: true (adds randomness to prevent thundering herd)
305	/// - `jitter_factor`: 0.1 (10% variation)
306	/// - `retryable_status_codes`: [429, 500, 502, 503, 504]
307	#[derive(Debug, Clone)]
308	pub struct RetryConfig {
309	/// Whether retry logic is enabled
310	pub enabled: bool,
311
312	/// Maximum number of retry attempts (0 = no retries, 3 = up to 4 total attempts)
313	pub max_retries: u32,
314
315	/// Initial delay before first retry (seconds)
316	pub initial_delay: f64,
317
318	/// Maximum delay between retries (seconds)
319	pub max_delay: f64,
320
321	/// Base for exponential backoff (delay = initial * base^attempt)
322	pub exponential_base: f64,
323
324	/// Whether to add random jitter to delays
325	pub jitter: bool,
326
327	/// Jitter factor (0.1 = +/- 10% variation)
328	pub jitter_factor: f64,
329
330	/// Whether to respect server's Retry-After header
331	pub respect_retry_after: bool,
332
333	/// HTTP status codes that should trigger a retry
334	#[allow(dead_code)] // Used in tests via is_retryable_status()
335	pub retryable_status_codes: Vec<u16>,
336
337	/// Timeout for individual requests (seconds, 0 = no timeout)
338	#[allow(dead_code)] // Configuration field for retry consumers
339	pub request_timeout: f64,
340
341	/// Total timeout for all retry attempts (seconds, 0 = no total timeout)
342	pub total_timeout: f64,
343	}
344
345	impl Default for RetryConfig {
346	fn default() -> Self {
347	Self {
348	enabled: true,
349	max_retries: 3,
350	initial_delay: 1.0,
351	max_delay: 60.0,
352	exponential_base: 2.0,
353	jitter: true,
354	jitter_factor: 0.1,
355	respect_retry_after: true,
356	retryable_status_codes: vec![429, 500, 502, 503, 504],
357	request_timeout: 120.0,
358	total_timeout: 0.0, // No total timeout by default
359	}
360	}
361	}
362
363	#[allow(dead_code)] // Public builder API, used in tests
364	impl RetryConfig {
365	/// Creates a new `RetryConfig` with default values
366	pub fn new() -> Self {
367	Self::default()
368	}
369
370	/// Creates a config with retry disabled
371	pub fn disabled() -> Self {
372	Self {
373	enabled: false,
374	..Default::default()
375	}
376	}
377
378	/// Builder method to set max retries
379	pub fn with_max_retries(mut self, max_retries: u32) -> Self {
380	self.max_retries = max_retries;
381	self
382	}
383
384	/// Builder method to set initial delay
385	pub fn with_initial_delay(mut self, delay: f64) -> Self {
386	self.initial_delay = delay;
387	self
388	}
389
390	/// Builder method to set max delay
391	pub fn with_max_delay(mut self, delay: f64) -> Self {
392	self.max_delay = delay;
393	self
394	}
395
396	/// Builder method to enable/disable jitter
397	pub fn with_jitter(mut self, enabled: bool) -> Self {
398	self.jitter = enabled;
399	self
400	}
401
402	/// Builder method to set request timeout
403	pub fn with_request_timeout(mut self, timeout: f64) -> Self {
404	self.request_timeout = timeout;
405	self
406	}
407
408	/// Builder method to set total timeout
409	pub fn with_total_timeout(mut self, timeout: f64) -> Self {
410	self.total_timeout = timeout;
411	self
412	}
413
414	/// Calculates the delay for a given retry attempt.
415	///
416	/// Uses exponential backoff: delay = `initial_delay` * `exponential_base^attempt`
417	/// The result is capped at `max_delay` and optionally has jitter applied.
418	///
419	/// # Arguments
420	///
421	/// * `attempt` - Zero-based attempt number (0 = first retry)
422	///
423	/// # Returns
424	///
425	/// Duration to wait before the next retry attempt
426	pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
427	let exponent = i32::try_from(attempt).unwrap_or(i32::MAX);
428	let base_delay = self.initial_delay * self.exponential_base.powi(exponent);
429	let capped_delay = base_delay.min(self.max_delay);
430
431	let final_delay = if self.jitter {
432	// Add random jitter to prevent thundering herd problem
433	let jitter_range = capped_delay * self.jitter_factor;
434	// Use UUID v4 entropy for jitter randomness.
435	let bytes = *Uuid::new_v4().as_bytes();
436	let sample = u16::from_le_bytes([bytes[0], bytes[1]]);
437	let random_factor = f64::from(sample) / f64::from(u16::MAX); // 0.0 to 1.0
438	let jitter = jitter_range * (2.0 * random_factor - 1.0); // -range to +range
439
440	(capped_delay + jitter).max(0.0)
441	} else {
442	capped_delay
443	};
444
445	Duration::from_secs_f64(final_delay)
446	}
447
448	/// Checks if a given HTTP status code should trigger a retry
449	pub fn is_retryable_status(&self, status: u16) -> bool {
450	self.retryable_status_codes.contains(&status)
451	}
452	}
453
454	/// Converts from the existing `RetryPolicy` in config
455	impl From<RetryPolicy> for RetryConfig {
456	fn from(policy: RetryPolicy) -> Self {
457	Self {
458	enabled: policy.enabled,
459	max_retries: policy.max_retries,
460	initial_delay: policy.initial_delay,
461	max_delay: policy.max_delay,
462	exponential_base: policy.exponential_base,
463	..Default::default()
464	}
465	}
466	}
467
468	/// Converts back to `RetryPolicy` for compatibility
469	impl From<RetryConfig> for RetryPolicy {
470	fn from(config: RetryConfig) -> Self {
471	Self {
472	enabled: config.enabled,
473	max_retries: config.max_retries,
474	initial_delay: config.initial_delay,
475	max_delay: config.max_delay,
476	exponential_base: config.exponential_base,
477	}
478	}
479	}
480
481	// === Retry Error and Result Types ===
482
483	/// Error returned when all retry attempts have been exhausted.
484	#[derive(Debug)]
485	pub struct RetryError {
486	/// The last error encountered
487	pub last_error: LlmError,
488
489	/// Total number of attempts made
490	pub attempts: u32,
491
492	/// Total time spent across all attempts
493	pub total_time: Duration,
494	}
495
496	impl std::fmt::Display for RetryError {
497	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
498	write!(
499	f,
500	"Retry exhausted after {} attempts ({:?}): {}",
501	self.attempts, self.total_time, self.last_error
502	)
503	}
504	}
505
506	impl std::error::Error for RetryError {
507	fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
508	Some(&self.last_error)
509	}
510	}
511
512	/// Result type for retry operations
513	pub type RetryResult<T> = Result<T, RetryError>;
514
515	/// Callback type for retry notifications
516	///
517	/// Called before each retry with:
518	/// - The error that triggered the retry
519	/// - The attempt number (0-based)
520	/// - The delay before the next attempt
521	pub type RetryCallback = Box<dyn Fn(&LlmError, u32, Duration) + Send + Sync>;
522
523	// === with_retry - Generic Retry Wrapper ===
524
525	/// Executes an async operation with configurable retry logic.
526	///
527	/// This function wraps any async operation that returns `Result<T, LlmError>`
528	/// and automatically retries on transient failures using exponential backoff.
529	///
530	/// # Arguments
531	///
532	/// * `config` - Retry configuration (delays, max attempts, etc.)
533	/// * `operation` - Async closure to execute (will be called multiple times on retry)
534	/// * `callback` - Optional callback for retry notifications (logging, metrics, etc.)
535	///
536	/// # Returns
537	///
538	/// * `Ok(T)` - The successful result from the operation
539	/// * `Err(RetryError)` - All retries exhausted or non-retryable error encountered
540	///
541	/// # Example
542	///
543	/// ```ignore
544	/// let result = with_retry(
545	/// &config,
546	/// \|\| async { client.send_request(&req).await },
547	/// Some(Box::new(\|err, attempt, delay\| {
548	/// eprintln!("Retry {} after {:?}: {}", attempt, delay, err);
549	/// })),
550	/// ).await;
551	/// ```
552	pub async fn with_retry<F, Fut, T>(
553	config: &RetryConfig,
554	mut operation: F,
555	callback: Option<RetryCallback>,
556	) -> RetryResult<T>
557	where
558	F: FnMut() -> Fut,
559	Fut: Future<Output = Result<T, LlmError>>,
560	{
561	// If retries are disabled, just run once
562	if !config.enabled {
563	return operation().await.map_err(\|e\| RetryError {
564	last_error: e,
565	attempts: 1,
566	total_time: Duration::ZERO,
567	});
568	}
569
570	let start_time = Instant::now();
571	let total_timeout = if config.total_timeout > 0.0 {
572	Some(Duration::from_secs_f64(config.total_timeout))
573	} else {
574	None
575	};
576
577	let mut last_error: Option<LlmError> = None;
578
579	// Attempt 0 is the first try, then up to max_retries additional attempts
580	for attempt in 0..=config.max_retries {
581	// Check total timeout
582	if let Some(timeout) = total_timeout
583	&& start_time.elapsed() >= timeout
584	{
585	return Err(RetryError {
586	last_error: last_error.unwrap_or(LlmError::Timeout(timeout)),
587	attempts: attempt,
588	total_time: start_time.elapsed(),
589	});
590	}
591
592	match operation().await {
593	Ok(result) => return Ok(result),
594	Err(err) => {
595	// Non-retryable errors fail immediately
596	if !err.is_retryable() {
597	return Err(RetryError {
598	last_error: err,
599	attempts: attempt + 1,
600	total_time: start_time.elapsed(),
601	});
602	}
603
604	// Last attempt - no more retries
605	if attempt >= config.max_retries {
606	return Err(RetryError {
607	last_error: err,
608	attempts: attempt + 1,
609	total_time: start_time.elapsed(),
610	});
611	}
612
613	// Calculate delay
614	// Use server's Retry-After if available and configured
615	let base_delay = config.delay_for_attempt(attempt);
616	let delay = if config.respect_retry_after {
617	err.suggested_retry_delay().unwrap_or(base_delay)
618	} else {
619	base_delay
620	};
621
622	// Notify callback if provided
623	if let Some(ref cb) = callback {
624	cb(&err, attempt, delay);
625	}
626
627	last_error = Some(err);
628
629	// Wait before retrying
630	tokio::time::sleep(delay).await;
631	}
632	}
633	}
634
635	// Should not reach here, but handle gracefully
636	Err(RetryError {
637	last_error: last_error.unwrap_or(LlmError::Other("Unknown retry error".to_string())),
638	attempts: config.max_retries + 1,
639	total_time: start_time.elapsed(),
640	})
641	}
642
643	/// Simplified version of `with_retry` without callback
644	#[allow(dead_code)] // Convenience wrapper for with_retry
645	pub async fn with_retry_simple<F, Fut, T>(config: &RetryConfig, operation: F) -> RetryResult<T>
646	where
647	F: FnMut() -> Fut,
648	Fut: Future<Output = Result<T, LlmError>>,
649	{
650	with_retry(config, operation, None).await
651	}
652
653	// === Utility Functions ===
654
655	/// Parses the Retry-After header value into a Duration.
656	///
657	/// Supports both:
658	/// - Seconds as integer: "120" -> 120 seconds
659	/// - HTTP-date format: "Wed, 21 Oct 2015 07:28:00 GMT" (not implemented, returns None)
660	pub fn parse_retry_after(value: &str) -> Option<Duration> {
661	// Try parsing as seconds
662	if let Ok(seconds) = value.parse::<u64>() {
663	return Some(Duration::from_secs(seconds));
664	}
665
666	// Try parsing as float seconds
667	if let Ok(seconds) = value.parse::<f64>() {
668	return Some(Duration::from_secs_f64(seconds));
669	}
670
671	// HTTP-date format not supported yet
672	// Could use chrono or httpdate crate if needed
673	None
674	}
675
676	/// Extracts Retry-After duration from response headers
677	pub fn extract_retry_after(headers: &reqwest::header::HeaderMap) -> Option<Duration> {
678	headers
679	.get(reqwest::header::RETRY_AFTER)
680	.and_then(\|v\| v.to_str().ok())
681	.and_then(parse_retry_after)
682	}
683
684	// === Tests ===
685
686	#[cfg(test)]
687	mod tests {
688	use super::*;
689
690	fn assert_f64_eq(actual: f64, expected: f64) {
691	assert!(
692	(actual - expected).abs() < f64::EPSILON,
693	"expected {expected}, got {actual}"
694	);
695	}
696
697	#[test]
698	fn test_retry_config_defaults() {
699	let config = RetryConfig::default();
700	assert!(config.enabled);
701	assert_eq!(config.max_retries, 3);
702	assert_f64_eq(config.initial_delay, 1.0);
703	assert_f64_eq(config.max_delay, 60.0);
704	assert_f64_eq(config.exponential_base, 2.0);
705	assert!(config.jitter);
706	}
707
708	#[test]
709	fn test_retry_config_disabled() {
710	let config = RetryConfig::disabled();
711	assert!(!config.enabled);
712	}
713
714	#[test]
715	fn test_retry_config_builder() {
716	let config = RetryConfig::new()
717	.with_max_retries(5)
718	.with_initial_delay(2.0)
719	.with_max_delay(120.0)
720	.with_jitter(false);
721
722	assert_eq!(config.max_retries, 5);
723	assert_f64_eq(config.initial_delay, 2.0);
724	assert_f64_eq(config.max_delay, 120.0);
725	assert!(!config.jitter);
726	}
727
728	#[test]
729	fn test_delay_for_attempt_exponential() {
730	let config = RetryConfig::new().with_jitter(false);
731
732	// delay = initial * base^attempt
733	// 1.0 * 2^0 = 1.0
734	let d0 = config.delay_for_attempt(0);
735	assert_eq!(d0, Duration::from_secs_f64(1.0));
736
737	// 1.0 * 2^1 = 2.0
738	let d1 = config.delay_for_attempt(1);
739	assert_eq!(d1, Duration::from_secs_f64(2.0));
740
741	// 1.0 * 2^2 = 4.0
742	let d2 = config.delay_for_attempt(2);
743	assert_eq!(d2, Duration::from_secs_f64(4.0));
744
745	// 1.0 * 2^3 = 8.0
746	let d3 = config.delay_for_attempt(3);
747	assert_eq!(d3, Duration::from_secs_f64(8.0));
748	}
749
750	#[test]
751	fn test_delay_for_attempt_capped() {
752	let config = RetryConfig::new().with_jitter(false).with_max_delay(5.0);
753
754	// 1.0 * 2^3 = 8.0, but capped at 5.0
755	let d3 = config.delay_for_attempt(3);
756	assert_eq!(d3, Duration::from_secs_f64(5.0));
757	}
758
759	#[test]
760	fn test_delay_for_attempt_with_jitter() {
761	let config = RetryConfig::new().with_jitter(true);
762
763	// With jitter, delays should vary slightly
764	let d1 = config.delay_for_attempt(1);
765	let d2 = config.delay_for_attempt(1);
766
767	// Both should be close to 2.0 seconds (within 10% jitter)
768	let base = 2.0;
769	let range = base * 0.1;
770	assert!(d1.as_secs_f64() >= base - range);
771	assert!(d1.as_secs_f64() <= base + range);
772	assert!(d2.as_secs_f64() >= base - range);
773	assert!(d2.as_secs_f64() <= base + range);
774	}
775
776	#[test]
777	fn test_is_retryable_status() {
778	let config = RetryConfig::default();
779
780	assert!(config.is_retryable_status(429)); // Rate limit
781	assert!(config.is_retryable_status(500)); // Internal server error
782	assert!(config.is_retryable_status(502)); // Bad gateway
783	assert!(config.is_retryable_status(503)); // Service unavailable
784	assert!(config.is_retryable_status(504)); // Gateway timeout
785
786	assert!(!config.is_retryable_status(400)); // Bad request
787	assert!(!config.is_retryable_status(401)); // Unauthorized
788	assert!(!config.is_retryable_status(403)); // Forbidden
789	assert!(!config.is_retryable_status(404)); // Not found
790	}
791
792	#[test]
793	fn test_llm_error_retryable() {
794	// Retryable errors
795	assert!(
796	LlmError::RateLimited {
797	message: "too many requests".to_string(),
798	retry_after: None
799	}
800	.is_retryable()
801	);
802	assert!(
803	LlmError::ServerError {
804	status: 500,
805	message: "internal error".to_string()
806	}
807	.is_retryable()
808	);
809	assert!(LlmError::NetworkError("connection refused".to_string()).is_retryable());
810	assert!(LlmError::Timeout(Duration::from_secs(30)).is_retryable());
811
812	// Non-retryable errors
813	assert!(!LlmError::AuthenticationError("invalid key".to_string()).is_retryable());
814	assert!(
815	!LlmError::InvalidRequest {
816	status: 400,
817	message: "bad json".to_string()
818	}
819	.is_retryable()
820	);
821	assert!(!LlmError::ContentPolicyError("unsafe content".to_string()).is_retryable());
822	assert!(!LlmError::ContextLengthError("too long".to_string()).is_retryable());
823	}
824
825	#[test]
826	fn test_llm_error_from_http_response() {
827	// Rate limit
828	let err = LlmError::from_http_response(429, "rate limit exceeded");
829	assert!(matches!(err, LlmError::RateLimited { .. }));
830
831	// Auth errors
832	let err = LlmError::from_http_response(401, "invalid api key");
833	assert!(matches!(err, LlmError::AuthenticationError(_)));
834
835	let err = LlmError::from_http_response(403, "forbidden");
836	assert!(matches!(err, LlmError::AuthenticationError(_)));
837
838	// Server errors
839	let err = LlmError::from_http_response(500, "internal server error");
840	assert!(matches!(err, LlmError::ServerError { status: 500, .. }));
841
842	let err = LlmError::from_http_response(503, "service unavailable");
843	assert!(matches!(err, LlmError::ServerError { status: 503, .. }));
844
845	// Context length
846	let err = LlmError::from_http_response(400, "context_length_exceeded");
847	assert!(matches!(err, LlmError::ContextLengthError(_)));
848
849	// Content policy
850	let err = LlmError::from_http_response(400, "content_policy_violation");
851	assert!(matches!(err, LlmError::ContentPolicyError(_)));
852
853	// Generic 400
854	let err = LlmError::from_http_response(400, "invalid json");
855	assert!(matches!(err, LlmError::InvalidRequest { status: 400, .. }));
856	}
857
858	#[test]
859	fn test_llm_error_suggested_retry_delay() {
860	let err = LlmError::RateLimited {
861	message: "slow down".to_string(),
862	retry_after: Some(Duration::from_secs(60)),
863	};
864	assert_eq!(err.suggested_retry_delay(), Some(Duration::from_secs(60)));
865
866	let err = LlmError::ServerError {
867	status: 500,
868	message: "error".to_string(),
869	};
870	assert_eq!(err.suggested_retry_delay(), None);
871	}
872
873	#[test]
874	fn test_parse_retry_after() {
875	// Integer seconds
876	assert_eq!(parse_retry_after("120"), Some(Duration::from_secs(120)));
877	assert_eq!(parse_retry_after("0"), Some(Duration::from_secs(0)));
878
879	// Float seconds
880	assert_eq!(parse_retry_after("1.5"), Some(Duration::from_secs_f64(1.5)));
881
882	// Invalid
883	assert_eq!(parse_retry_after("invalid"), None);
884	assert_eq!(parse_retry_after(""), None);
885	}
886
887	#[test]
888	fn test_retry_policy_conversion() {
889	let policy = RetryPolicy {
890	enabled: true,
891	max_retries: 5,
892	initial_delay: 2.0,
893	max_delay: 30.0,
894	exponential_base: 3.0,
895	};
896
897	let config: RetryConfig = policy.clone().into();
898	assert_eq!(config.enabled, policy.enabled);
899	assert_eq!(config.max_retries, policy.max_retries);
900	assert_f64_eq(config.initial_delay, policy.initial_delay);
901	assert_f64_eq(config.max_delay, policy.max_delay);
902	assert_f64_eq(config.exponential_base, policy.exponential_base);
903
904	// Convert back
905	let policy2: RetryPolicy = config.into();
906	assert_eq!(policy2.enabled, policy.enabled);
907	assert_eq!(policy2.max_retries, policy.max_retries);
908	}
909
910	#[tokio::test]
911	async fn test_with_retry_success_first_attempt() {
912	let config = RetryConfig::default();
913	let mut call_count = 0;
914
915	let result = with_retry(
916	&config,
917	\|\| {
918	call_count += 1;
919	async { Ok::<_, LlmError>(42) }
920	},
921	None,
922	)
923	.await;
924
925	assert!(result.is_ok());
926	assert_eq!(result.unwrap(), 42);
927	assert_eq!(call_count, 1);
928	}
929
930	#[tokio::test]
931	async fn test_with_retry_disabled() {
932	let config = RetryConfig::disabled();
933	let mut call_count = 0;
934
935	let result: RetryResult<i32> = with_retry(
936	&config,
937	\|\| {
938	call_count += 1;
939	async {
940	Err(LlmError::ServerError {
941	status: 500,
942	message: "error".to_string(),
943	})
944	}
945	},
946	None,
947	)
948	.await;
949
950	assert!(result.is_err());
951	assert_eq!(call_count, 1); // No retries when disabled
952	}
953
954	#[tokio::test]
955	async fn test_with_retry_non_retryable_error() {
956	let config = RetryConfig::default();
957	let mut call_count = 0;
958
959	let result: RetryResult<i32> = with_retry(
960	&config,
961	\|\| {
962	call_count += 1;
963	async { Err(LlmError::AuthenticationError("bad key".to_string())) }
964	},
965	None,
966	)
967	.await;
968
969	assert!(result.is_err());
970	assert_eq!(call_count, 1); // Auth errors are not retried
971	}
972
973	#[tokio::test]
974	async fn test_with_retry_eventual_success() {
975	let config = RetryConfig::new()
976	.with_max_retries(3)
977	.with_initial_delay(0.01); // Fast for testing
978
979	let call_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
980	let cc = call_count.clone();
981
982	let result = with_retry(
983	&config,
984	\|\| {
985	let count = cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
986	async move {
987	if count < 2 {
988	Err(LlmError::ServerError {
989	status: 500,
990	message: "temporary error".to_string(),
991	})
992	} else {
993	Ok::<_, LlmError>(42)
994	}
995	}
996	},
997	None,
998	)
999	.await;
1000
1001	assert!(result.is_ok());
1002	assert_eq!(result.unwrap(), 42);
1003	assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 3); // 2 failures + 1 success
1004	}
1005
1006	#[tokio::test]
1007	async fn test_with_retry_exhausted() {
1008	let config = RetryConfig::new()
1009	.with_max_retries(2)
1010	.with_initial_delay(0.01);
1011
1012	let call_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1013	let cc = call_count.clone();
1014
1015	let result: RetryResult<i32> = with_retry(
1016	&config,
1017	\|\| {
1018	cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1019	async {
1020	Err(LlmError::ServerError {
1021	status: 500,
1022	message: "persistent error".to_string(),
1023	})
1024	}
1025	},
1026	None,
1027	)
1028	.await;
1029
1030	assert!(result.is_err());
1031	let err = result.unwrap_err();
1032	assert_eq!(err.attempts, 3); // 1 initial + 2 retries
1033	assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 3);
1034	}
1035
1036	#[tokio::test]
1037	async fn test_with_retry_callback() {
1038	let config = RetryConfig::new()
1039	.with_max_retries(2)
1040	.with_initial_delay(0.01);
1041
1042	let callback_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
1043	let cc = callback_count.clone();
1044
1045	let _: RetryResult<i32> = with_retry(
1046	&config,
1047	\|\| async {
1048	Err(LlmError::ServerError {
1049	status: 500,
1050	message: "error".to_string(),
1051	})
1052	},
1053	Some(Box::new(move \|_err, _attempt, _delay\| {
1054	cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1055	})),
1056	)
1057	.await;
1058
1059	// Callback called once per retry (not for the final failure)
1060	assert_eq!(callback_count.load(std::sync::atomic::Ordering::SeqCst), 2);
1061	}
1062
1063	#[test]
1064	fn test_retry_error_display() {
1065	let err = RetryError {
1066	last_error: LlmError::ServerError {
1067	status: 500,
1068	message: "internal error".to_string(),
1069	},
1070	attempts: 4,
1071	total_time: Duration::from_secs(10),
1072	};
1073
1074	let display = format!("{err}");
1075	assert!(display.contains("4 attempts"));
1076	assert!(display.contains("10"));
1077	assert!(display.contains("Server error"));
1078	}
1079	}
1080

1080 lines RUST