diff --git a/src/llm/anthropic_oauth.rs b/src/llm/anthropic_oauth.rs index ae6674dc..8c701101 100644 --- a/src/llm/anthropic_oauth.rs +++ b/src/llm/anthropic_oauth.rs @@ -22,6 +22,7 @@ use crate::llm::provider::{ ToolCompletionRequest, ToolCompletionResponse, strip_unsupported_completion_params, strip_unsupported_tool_params, }; +use crate::llm::retry::cap_retry_after; const ANTHROPIC_API_URL: &str = "https://api.anthropic.com/v1/messages"; /// OAuth beta requires 2023-06-01; the 2024-10-22 version is not valid with the beta flag. @@ -150,6 +151,7 @@ impl AnthropicOAuthProvider { .and_then(|v| v.to_str().ok()) .and_then(|v| v.parse::().ok()) .map(std::time::Duration::from_secs) + .map(cap_retry_after) .or(Some(std::time::Duration::from_secs(60))); let response_text = response @@ -766,9 +768,14 @@ mod tests { #[test] fn test_retry_after_large_number() { - // Verify large numbers are accepted + // Verify large numbers are capped to the safe maximum let duration = parse_retry_after_anthropic_for_test("7200"); // 2 hours - assert_eq!(duration, Some(std::time::Duration::from_secs(7200))); + assert_eq!( + duration, + Some(std::time::Duration::from_secs( + crate::llm::retry::MAX_RETRY_AFTER_SECS + )) + ); } /// Helper function to test Retry-After header parsing logic for Anthropic @@ -779,6 +786,7 @@ mod tests { .parse::() .ok() .map(std::time::Duration::from_secs) + .map(cap_retry_after) .or(Some(std::time::Duration::from_secs(60))) } } diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index 0a9e1fdc..f0d711a9 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -22,7 +22,7 @@ use crate::llm::provider::{ ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, }; -use crate::llm::{costs, session::SessionManager}; +use crate::llm::{costs, retry::cap_retry_after, session::SessionManager}; /// Information about an available model from NEAR AI API. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -252,7 +252,7 @@ impl NearAiChatProvider { .and_then(|v| { // Try delay-seconds first (most common from API providers) if let Ok(secs) = v.trim().parse::() { - return Some(std::time::Duration::from_secs(secs)); + return Some(cap_retry_after(std::time::Duration::from_secs(secs))); } // Try HTTP-date (e.g. "Mon, 02 Mar 2026 18:00:00 GMT") if let Ok(dt) = chrono::DateTime::parse_from_rfc2822(v.trim()) { @@ -260,9 +260,9 @@ impl NearAiChatProvider { let delta = dt.signed_duration_since(now); // Use max(0) so past/present dates yield Duration::ZERO // rather than None (which would cause an immediate retry). - return Some(std::time::Duration::from_secs( - delta.num_seconds().max(0) as u64 - )); + return Some(cap_retry_after(std::time::Duration::from_secs( + delta.num_seconds().max(0) as u64, + ))); } None }) @@ -2306,9 +2306,17 @@ mod tests { #[test] fn test_retry_after_large_number() { - // Verify large numbers are accepted + // Verify large numbers are capped to the safe maximum let duration = parse_retry_after_for_test("3600"); // 1 hour assert_eq!(duration, Some(std::time::Duration::from_secs(3600))); + + let huge = parse_retry_after_for_test("18446744073709551615"); + assert_eq!( + huge, + Some(std::time::Duration::from_secs( + crate::llm::retry::MAX_RETRY_AFTER_SECS + )) + ); } /// Helper function to test Retry-After header parsing logic @@ -2316,13 +2324,13 @@ mod tests { fn parse_retry_after_for_test(header_value: &str) -> Option { let trimmed = header_value.trim(); let parsed = if let Ok(secs) = trimmed.parse::() { - Some(std::time::Duration::from_secs(secs)) + Some(cap_retry_after(std::time::Duration::from_secs(secs))) } else if let Ok(dt) = chrono::DateTime::parse_from_rfc2822(trimmed) { let now = chrono::Utc::now(); let delta = dt.signed_duration_since(now); - Some(std::time::Duration::from_secs( - delta.num_seconds().max(0) as u64 - )) + Some(cap_retry_after(std::time::Duration::from_secs( + delta.num_seconds().max(0) as u64, + ))) } else { None }; diff --git a/src/llm/retry.rs b/src/llm/retry.rs index 2875fbd3..6250de33 100644 --- a/src/llm/retry.rs +++ b/src/llm/retry.rs @@ -19,6 +19,12 @@ use crate::llm::provider::{ ToolCompletionResponse, }; +/// Upper bound for provider-suggested `Retry-After` delays. +/// +/// This prevents malicious or malformed headers from turning a retryable +/// response into an effectively unbounded sleep. +pub(crate) const MAX_RETRY_AFTER_SECS: u64 = 3600; + /// Returns `true` if the `LlmError` is transient and the request should be retried. /// /// Used by `RetryProvider` (retry the same provider) and `FailoverProvider` @@ -67,6 +73,11 @@ pub(crate) fn retry_backoff_delay(attempt: u32) -> Duration { Duration::from_millis(delay_ms) } +/// Clamp a provider-suggested retry delay to a safe maximum. +pub(crate) fn cap_retry_after(duration: Duration) -> Duration { + duration.min(Duration::from_secs(MAX_RETRY_AFTER_SECS)) +} + /// Configuration for the retry decorator. #[derive(Debug, Clone)] pub struct RetryConfig { @@ -421,4 +432,16 @@ mod tests { panic!("Expected RateLimited error"); } } + + #[test] + fn cap_retry_after_clamps_huge_delays() { + assert_eq!( + cap_retry_after(Duration::from_secs(u64::MAX)), + Duration::from_secs(MAX_RETRY_AFTER_SECS) + ); + assert_eq!( + cap_retry_after(Duration::from_secs(0)), + Duration::from_secs(0) + ); + } } diff --git a/src/workspace/embeddings.rs b/src/workspace/embeddings.rs index 96fe144b..a8ed0a3e 100644 --- a/src/workspace/embeddings.rs +++ b/src/workspace/embeddings.rs @@ -6,6 +6,8 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; +use crate::llm::retry::cap_retry_after; + /// Error type for embedding operations. #[derive(Debug, thiserror::Error)] pub enum EmbeddingError { @@ -232,6 +234,7 @@ impl EmbeddingProvider for OpenAiEmbeddings { .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()) .map(std::time::Duration::from_secs) + .map(cap_retry_after) .or(Some(std::time::Duration::from_secs(60))); return Err(EmbeddingError::RateLimited { retry_after }); } @@ -374,6 +377,7 @@ impl EmbeddingProvider for NearAiEmbeddings { .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()) .map(std::time::Duration::from_secs) + .map(cap_retry_after) .or(Some(std::time::Duration::from_secs(60))); return Err(EmbeddingError::RateLimited { retry_after }); } @@ -690,6 +694,7 @@ mod tests { .parse::() .ok() .map(std::time::Duration::from_secs) + .map(cap_retry_after) .or(Some(std::time::Duration::from_secs(60))) } }