diff --git a/Cargo.toml b/Cargo.toml index 74ef8836f..9366a37a3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,7 +40,7 @@ tempfile = "3.24.0" termcolor = "1.4.1" thiserror = "2.0.17" tiktoken-rs = "0.6.0" -tokio = { version = "1.48.0", features = ["full"] } +tokio = { version = "1.48.0", features = ["full", "test-util"] } tokio-util = { version = "0.7.17", features = ["codec", "io"] } toml = "0.8" tower-http = { version = "0.6.8", features = ["cors", "fs", "trace"] } diff --git a/Settings.toml b/Settings.toml index ba61fbc62..813f8e797 100644 --- a/Settings.toml +++ b/Settings.toml @@ -29,6 +29,10 @@ model = "gemini-3.1-pro-preview" max_input_tokens = 900000 max_interactions = 50 temperature = 1.0 +# global_backoff = true +# quota_backoff_secs = 60 +# transient_backoff_type = "exponential" +# transient_flat_backoff_secs = 30 # To use Claude API directly: # provider = "claude" diff --git a/designs/ai_error_handling.md b/designs/ai_error_handling.md index 653e8c000..141fc4276 100644 --- a/designs/ai_error_handling.md +++ b/designs/ai_error_handling.md @@ -1,118 +1,90 @@ # AI Provider Error Handling and Rate Limiting Design ## Objective -Currently, Sashiko's handling of AI provider errors (like 503 Service Unavailable, 529 Overloaded, and 429 Too Many Requests) is fragmented and localized. Some errors cause individual worker threads to sleep and retry locally (which can lead to thread exhaustion and timeout issues), while only specific quota errors trigger the global `QuotaManager`. - -The goal of this redesign is to: -1. **Unify error handling:** All transient and rate-limiting errors from any LLM provider (Gemini, Claude) should trigger a global backoff/delay mechanism. -2. **Global Pause:** When an AI provider is overloaded or returns a transient error, *all* AI interactions across *all* workers should be delayed to allow the provider to recover and to prevent exacerbating the issue. -3. **Infinite Retries:** There should be no hard limit on the number of retries for transient/overloaded errors. The system should patiently wait for the provider to become available again. -4. **Timeout Exemption:** The time spent waiting for global delays (due to quota or transient errors) must *not* be counted towards any per-worker or per-task timeouts. Worker timeouts should only apply to actual processing and execution time. - -## Proposed Architecture - -### 1. Unified Global AI Rate Limiter / Quota Manager -The existing `QuotaManager` in `src/ai/proxy.rs` (or potentially moved to a better location like `src/ai/mod.rs` or `src/ai/quota.rs`) will be expanded and renamed to `GlobalAiBackoffManager` (or similar). -* It will manage a global `blocked_until` timestamp. -* It will keep track of consecutive transient failures to implement global exponential backoff. -* It will provide a way for any worker to report *any* type of retryable AI error (Quota Exceeded, Overloaded, Transient 5xx). -* It will provide a `wait_for_access()` method that workers must call *before* attempting any AI request. -* **Crucially**, `wait_for_access()` will return the `Duration` that the task actually spent sleeping. -* **Global Pause:** When *any* worker receives a 503 (or other retryable error), it will report it to the manager. The manager will update the global `blocked_until` timestamp. This ensures that *all other workers* will pause and wait when they call `wait_for_access()`, preventing the system from bombarding an overloaded provider. - -### 2. Error Categorization & Backoff Strategy -Errors from providers will be categorized into: -* **Fatal Errors (400, 401, 403 - except Gemini cache expiration, 404):** These should fail immediately. No retries. -* **Rate Limit / Quota (429):** Providers often include a `Retry-After` header. If present, the global delay is set to this value. If not, a default (e.g., 60s) is used. -* **Overloaded / Transient (500, 502, 503, 504, 529):** These will trigger a global exponential backoff. - * The backoff will start small (e.g., 5s) and increase (e.g., multiply by 1.5x or 2x) on consecutive failures, up to a maximum cap (e.g., 5 minutes). - * A successful request will reset the consecutive failure counter. - -### 3. Provider Client Changes (`src/ai/gemini.rs`, `src/ai/claude.rs`) -* Remove the internal `loop` and retry logic (`max_retries`) from `post_request_with_retry` (Claude) and `generate_content` / `generate_content_with_cache` (Gemini). -* The clients should simply execute the request once and return the result or the categorized error. Both clients should map HTTP status codes to specific Rust enum variants (e.g., `GeminiError::QuotaExceeded`, `GeminiError::TransientError`). - -### 4. Worker Timeout Exemption (Dynamic Deadline) -In `src/reviewer.rs`, the AI interaction is currently wrapped in a fixed `tokio::time::timeout`: -```rust -let interaction_result = tokio::time::timeout( - std::time::Duration::from_secs(settings.review.timeout_seconds), - async { ... } -).await; +Sashiko interacts with external LLM providers (like Gemini and Claude) which can occasionally return transient errors (e.g., 503 Service Unavailable, 529 Overloaded, or connection timeouts) or rate-limiting errors (429 Too Many Requests). + +To ensure robustness and high availability, Sashiko implements a unified, configurable error handling and rate limiting strategy with the following goals: +1. **Unified Error Handling:** All transient and rate-limiting errors from any provider trigger a consistent backoff/delay mechanism. +2. **Configurable Scope (Global vs. Local):** Support both global rate limiting (blocking all workers to protect API quota) and local rate limiting (per-worker backoff to maximize concurrency). +3. **Flexible Backoff Strategies:** Support both exponential backoff (for gradual recovery) and flat backoff (for aggressive, predictable retries). +4. **Infinite Retries for Transients:** Transient errors are retried indefinitely to ensure reviews eventually complete. +5. **Timeout Exemption:** Time spent waiting for rate limits or transient backoffs does *not* count towards review timeouts. The active deadline is dynamically extended by the duration of the sleep. + +## Architecture + +### 1. AI Rate Limiter / Quota Manager (`QuotaManager`) +The `QuotaManager` (in `src/ai/quota.rs`) manages the backoff state: +* It tracks a `blocked_until` timestamp. +* It tracks consecutive transient failures to implement exponential backoff. +* It provides `wait_for_access()` which blocks the caller if the rate limit is active, returning the `Duration` spent sleeping. +* It supports two backoff types for transient errors: + * **Exponential:** Backoff starts at 1s and doubles on consecutive failures (`1s, 2s, 4s, 8s...`) capped at 60 seconds. + * **Flat:** Backoff always waits a configured flat duration (e.g., 30 seconds). + +### 2. Global vs. Local Backoff +The scope of the backoff is determined by how the `QuotaManager` is shared: +* **Global Backoff (`global_backoff = true`):** A single `QuotaManager` is shared as a singleton (`Arc`) across all workers. If one worker hits a rate limit or transient error, it blocks *all* workers. This is useful to protect API keys from getting banned under severe quota limits. +* **Local Backoff (`global_backoff = false`):** Each review task instantiates its own local `QuotaManager`. If a worker handling a review hits an error, it backs off individually without affecting other active workers. This maximizes throughput and retries "as soon as possible" across different reviews. + +### 3. Configuration (`Settings.toml`) +The behavior is fully configurable under the `[ai]` section: +```toml +[ai] +# Scope of the backoff +global_backoff = false # If false, workers back off independently (local backoff) + +# Quota (429) error handling +quota_backoff_secs = 60 # Flat delay when hitting 429 + +# Transient (503, timeout) error handling +transient_backoff_type = "flat" # "exponential" or "flat" +transient_flat_backoff_secs = 30 # Flat delay when using "flat" transient backoff ``` -If `wait_for_access()` sleeps for 5 minutes inside that `async` block, the timeout will fire incorrectly. -* **Solution:** We will replace `tokio::time::timeout` with a custom implementation or loop that tracks the *active* deadline. -* Because `tokio::time::timeout` cannot have its deadline extended once started, we should track the deadline manually using `tokio::time::sleep` combined with `tokio::select!`. -* Alternatively, we can track the accumulated "sleep time" inside the async block. -* **Best approach for Sashiko:** The AI loop is already yielding events (reading lines from the child process). -* We can create a struct `ActiveTimeout`: +### 4. Standardized Error Categorization (`AiError`) +To avoid fragile string-matching in the orchestration layer (`reviewer.rs`), Sashiko abstracts provider-specific errors at the API boundary. + +We introduce a unified `AiError` enum in `src/ai/mod.rs`: ```rust -struct ActiveTimeout { - base_duration: Duration, - total_sleep_added: Duration, - start: Instant, -} -impl ActiveTimeout { - fn new(base: Duration) -> Self { ... } - fn add_sleep(&mut self, d: Duration) { self.total_sleep_added += d; } - fn is_expired(&self) -> bool { - self.start.elapsed() > (self.base_duration + self.total_sleep_added) - } - fn remaining(&self) -> Duration { ... } +#[derive(Debug, thiserror::Error, Clone)] +pub enum AiError { + #[error("Quota exceeded: retry after {0:?}")] + QuotaExceeded(std::time::Duration), + #[error("Transient error: {1}, retry after {0:?}")] + Transient(std::time::Duration, String), + #[error("Fatal AI error: {0}")] + Fatal(String), } ``` -* And wrap the reading of lines in a `timeout` that is re-calculated for each line or each AI request. Or better yet, run the entire child process reading loop inside a custom future that we poll alongside a sleep future that gets its deadline updated. -Let's refine the timeout approach. Instead of wrapping the whole block in `tokio::time::timeout`, we can do: +Each provider client (Gemini, Claude, OpenAI) is responsible for mapping its raw HTTP status codes or client-specific errors into this standardized enum before returning them: +* **Gemini:** Maps `GeminiError::QuotaExceeded` and `GeminiError::TransientError`. +* **Claude:** Maps `ClaudeError::RateLimitExceeded` and `ClaudeError::OverloadedError`. +* **OpenAI:** Maps `OpenAiCompatError::RateLimitExceeded` and `OpenAiCompatError::TransientError`. +In `reviewer.rs`, the orchestration layer downcasts the returned `anyhow::Error` to `AiError`. This allows robust, type-safe classification to trigger the appropriate backoff: +* `AiError::QuotaExceeded(delay)`: Triggers `report_quota_error(delay)`. +* `AiError::Transient(delay, msg)`: Triggers `report_transient_error()`. +* `AiError::Fatal(msg)`: Fails the review immediately (no retries). + +A fallback string-matching mechanism is preserved only for backward compatibility with custom or legacy external tools. + +### 5. Worker Timeout Exemption (Dynamic Deadline) +To prevent reviews from timing out while waiting for API recovery, the review deadline is dynamically extended. + +In `src/reviewer.rs` (`run_review_tool`), the active timeout is managed using a mutable `deadline` (`TokioInstant`): ```rust -let deadline = Instant::now() + Duration::from_secs(settings.review.timeout_seconds); -// ... inside the loop handling messages ... - let resp_payload = loop { - let slept = quota_manager.wait_for_access().await; - deadline += slept; // Extend deadline by the time we slept! - - // Check timeout manually before heavy work, or rely on read timeouts - if Instant::now() > deadline { - break Err(anyhow!("Review tool timed out (active time exceeded)")); - } - - match provider.generate_content(req.clone()).await { - Ok(resp) => { - quota_manager.report_success().await; - break Ok(resp); - } - Err(e) => { - if is_retryable(&e) { - quota_manager.report_error(e).await; - continue; // Loop indefinitely until success or fatal error - } - break Err(e); - } - } - } +let mut deadline = TokioInstant::now() + Duration::from_secs(settings.review.timeout_seconds); ``` -However, the `tokio::time::timeout` also guards against the *child process* hanging while producing stdout. We still need a timeout around the `lines.next_line().await`. - -**Updated Timeout Strategy:** -1. We will use a mutable `deadline` variable: `let mut deadline = Instant::now() + Duration::from_secs(timeout_seconds);` -2. We wrap the `lines.next_line().await` (and other I/O) in `tokio::time::timeout_at(deadline.into(), ...)` -3. Whenever `quota_manager.wait_for_access().await` returns a `Duration > 0`, we do `deadline += duration`. - -## Execution Plan - -1. **Refactor `QuotaManager`:** - * Rename to `QuotaManager` (keep name to minimize churn) but add tracking for consecutive transient failures. - * Update `wait_for_access` to return the `Duration` slept. - * Add `report_transient_error()` to trigger global exponential backoff. - * Add `report_success()` to reset transient error counter. -2. **Update `gemini.rs` and `claude.rs`:** - * Remove local `retry_count` loops and `tokio::time::sleep`. - * Return errors directly (ensure 503/529 map to `TransientError` and 429 maps to `QuotaExceeded`). -3. **Update Callers (`reviewer.rs`, `proxy.rs`):** - * Implement infinite loop around `provider.generate_content`. - * Call `wait_for_access()` and capture the sleep duration. - * Extend timeouts (`deadline += slept`). - * Report errors (`report_quota_error` or `report_transient_error`) or success (`report_success`). -4. **Implement `timeout_at` logic in `reviewer.rs`:** Replace the single block `tokio::time::timeout` with loop-level `timeout_at` to allow deadline extension. \ No newline at end of file + +During the AI interaction loop, before each request, the worker calls `wait_for_access()` and extends the deadline by the slept duration: +```rust +let slept = quota_manager.wait_for_access().await; +deadline += slept; // Extend deadline by the time we slept! + +if TokioInstant::now() > deadline { + return Err(anyhow!("Review tool timed out (active time exceeded)")); +} +``` + +This ensures that the `timeout_seconds` config applies strictly to **active processing time** (e.g., waiting for child process I/O, local processing), while all time spent waiting for external AI provider availability is exempt. \ No newline at end of file diff --git a/sashiko.dev/base/app/sashiko-k8s.yaml b/sashiko.dev/base/app/sashiko-k8s.yaml index 86ec357d5..c22d55e64 100644 --- a/sashiko.dev/base/app/sashiko-k8s.yaml +++ b/sashiko.dev/base/app/sashiko-k8s.yaml @@ -102,6 +102,17 @@ spec: value: "false" - name: SASHIKO__REVIEW__EMAIL_POLICY_PATH value: "/app/email_policy.toml" + # Enable per-worker (local) backoff isolation + - name: SASHIKO__AI__GLOBAL_BACKOFF + value: "false" + # Configure flat backoffs instead of exponential + - name: SASHIKO__AI__TRANSIENT_BACKOFF_TYPE + value: "flat" + - name: SASHIKO__AI__TRANSIENT_FLAT_BACKOFF_SECS + value: "30" + # Set flat quota retry to 65s too (optional, defaults to 60s if omitted) + - name: SASHIKO__AI__QUOTA_BACKOFF_SECS + value: "65" - name: LLM_API_KEY valueFrom: secretKeyRef: diff --git a/src/ai/claude.rs b/src/ai/claude.rs index 42520f572..bc8289fe3 100644 --- a/src/ai/claude.rs +++ b/src/ai/claude.rs @@ -574,7 +574,27 @@ impl AiProvider for ClaudeClient { claude_req.model = self.model.clone(); // 3. Make API call - let response = self.post_request(&claude_req).await?; + let response = match self.post_request(&claude_req).await { + Ok(resp) => resp, + Err(e) => { + if let Some(claude_err) = e.downcast_ref::() { + match claude_err { + ClaudeError::RateLimitExceeded(d) => { + return Err(crate::ai::AiError::QuotaExceeded(*d).into()); + } + ClaudeError::OverloadedError(d) => { + return Err(crate::ai::AiError::Transient( + *d, + "Claude API overloaded".to_string(), + ) + .into()); + } + _ => {} + } + } + return Err(e); + } + }; // 4. Translate response back to generic format translate_ai_response(&response) diff --git a/src/ai/gemini.rs b/src/ai/gemini.rs index c0f412788..d82e5708b 100644 --- a/src/ai/gemini.rs +++ b/src/ai/gemini.rs @@ -756,7 +756,39 @@ fn estimate_tokens_generic(request: &AiRequest) -> usize { impl AiProvider for GeminiClient { async fn generate_content(&self, request: AiRequest) -> Result { let gen_req = translate_ai_request(request)?; - let resp = GenAiClient::generate_content(self, gen_req).await?; + let resp = match GenAiClient::generate_content(self, gen_req).await { + Ok(resp) => resp, + Err(e) => { + if let Some(gemini_err) = e.downcast_ref::() { + match gemini_err { + GeminiError::QuotaExceeded(d) => { + return Err(crate::ai::AiError::QuotaExceeded(*d).into()); + } + GeminiError::TransientError(d, msg) => { + return Err(crate::ai::AiError::Transient(*d, msg.clone()).into()); + } + GeminiError::PermissionDenied(msg) => { + return Err(crate::ai::AiError::Fatal(format!( + "Permission denied: {}", + msg + )) + .into()); + } + GeminiError::ApiError(status, msg) => { + return Err(crate::ai::AiError::Fatal(format!( + "API Error ({}): {}", + status, msg + )) + .into()); + } + GeminiError::Other(msg) => { + return Err(crate::ai::AiError::Fatal(msg.clone()).into()); + } + } + } + return Err(e); + } + }; translate_ai_response(resp) } diff --git a/src/ai/mod.rs b/src/ai/mod.rs index 367e5bce1..7d5c6cbd0 100644 --- a/src/ai/mod.rs +++ b/src/ai/mod.rs @@ -176,6 +176,17 @@ pub struct CacheStats { pub tokens_saved_prev_session: u64, } +/// Generic errors to all AI providers +#[derive(Debug, thiserror::Error, Clone)] +pub enum AiError { + #[error("Quota exceeded: retry after {0:?}")] + QuotaExceeded(std::time::Duration), + #[error("Transient error: {1}, retry after {0:?}")] + Transient(std::time::Duration, String), + #[error("Fatal AI error: {0}")] + Fatal(String), +} + /// Trait defining the standard interface for all AI providers in Sashiko. #[async_trait] pub trait AiProvider: Send + Sync { diff --git a/src/ai/openai.rs b/src/ai/openai.rs index c6173970a..a79337b14 100644 --- a/src/ai/openai.rs +++ b/src/ai/openai.rs @@ -482,7 +482,18 @@ impl AiProvider for OpenAiCompatClient { openai_req.model = self.model.clone(); let resp_body = serde_json::to_value(&openai_req)?; - let resp = self.post_request(&resp_body).await?; + let resp = match self.post_request(&resp_body).await { + Ok(resp) => resp, + Err(e) => match e { + OpenAiCompatError::RateLimitExceeded(d) => { + return Err(crate::ai::AiError::QuotaExceeded(d).into()); + } + OpenAiCompatError::TransientError(d, msg) => { + return Err(crate::ai::AiError::Transient(d, msg).into()); + } + _ => return Err(e.into()), + }, + }; translate_ai_response(resp) } diff --git a/src/ai/quota.rs b/src/ai/quota.rs index 4ffa4f620..8a5d9aefb 100644 --- a/src/ai/quota.rs +++ b/src/ai/quota.rs @@ -12,16 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::time::{Duration, Instant}; +use std::time::Duration; use tokio::sync::Mutex; +use tokio::time::Instant; use tracing::{info, warn}; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BackoffType { + Exponential, + Flat, +} + pub struct QuotaManager { // Stores the time when we can resume making requests. // If None or in the past, we are free to go. blocked_until: Mutex>, // Track consecutive transient errors for exponential backoff. consecutive_transient_errors: Mutex, + transient_backoff_type: BackoffType, + transient_flat_delay: Duration, } impl Default for QuotaManager { @@ -32,9 +41,18 @@ impl Default for QuotaManager { impl QuotaManager { pub fn new() -> Self { + Self::new_with_settings(BackoffType::Exponential, Duration::from_secs(30)) + } + + pub fn new_with_settings( + transient_backoff_type: BackoffType, + transient_flat_delay: Duration, + ) -> Self { Self { blocked_until: Mutex::new(None), consecutive_transient_errors: Mutex::new(0), + transient_backoff_type, + transient_flat_delay: transient_flat_delay.max(Duration::from_secs(1)), } } @@ -76,8 +94,9 @@ impl QuotaManager { } pub async fn report_quota_error(&self, retry_after: Duration) { + let delay = retry_after.max(Duration::from_secs(1)); let mut guard = self.blocked_until.lock().await; - let resume_time = Instant::now() + retry_after; + let resume_time = Instant::now() + delay; if let Some(current) = *guard { if resume_time > current { @@ -89,7 +108,7 @@ impl QuotaManager { warn!( "Quota exhausted! Blocking all LLM requests for {:.2}s", - retry_after.as_secs_f64() + delay.as_secs_f64() ); } @@ -98,9 +117,13 @@ impl QuotaManager { *count_guard += 1; let count = *count_guard; - // Exponential backoff: 1s, 2s, 4s, 8s... capped at 60s - let backoff_secs = (1.0 * (2.0_f64.powi((count - 1) as i32))).min(60.0); - let backoff = Duration::from_secs_f64(backoff_secs); + let backoff = match self.transient_backoff_type { + BackoffType::Exponential => { + let backoff_secs = (1.0 * (2.0_f64.powi((count - 1) as i32))).min(60.0); + Duration::from_secs_f64(backoff_secs).max(Duration::from_secs(1)) + } + BackoffType::Flat => self.transient_flat_delay, + }; let mut block_guard = self.blocked_until.lock().await; let resume_time = Instant::now() + backoff; @@ -115,7 +138,122 @@ impl QuotaManager { warn!( "AI provider transient error (streak: {}). Globally backing off for {:.2}s", - count, backoff_secs + count, + backoff.as_secs_f64() ); } } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + #[tokio::test] + async fn test_quota_manager_flat_backoff() { + tokio::time::pause(); // Pause the real clock + let manager = Arc::new(QuotaManager::new_with_settings( + BackoffType::Flat, + Duration::from_secs(5), + )); + + // First transient error -> triggers 5s flat delay + manager.report_transient_error().await; + + let manager_clone = manager.clone(); + let handle = tokio::spawn(async move { manager_clone.wait_for_access().await }); + + // Yield execution to let the spawned task run and hit the sleep + tokio::task::yield_now().await; + // Sleep a tiny virtual duration to ensure the task has transitioned to sleeping + tokio::time::sleep(Duration::from_millis(1)).await; + + // Fast-forward 5s instantly + tokio::time::advance(Duration::from_secs(5)).await; + let slept = handle.await.unwrap(); + + // It should have slept for at least 5s (capped at the total advanced time) + assert!(slept >= Duration::from_secs(5)); + } + + #[tokio::test] + async fn test_quota_manager_exponential_backoff() { + tokio::time::pause(); + let manager = Arc::new(QuotaManager::new_with_settings( + BackoffType::Exponential, + Duration::from_secs(30), + )); + + // Attempt 1 -> 1s delay + manager.report_transient_error().await; + let manager_clone = manager.clone(); + let handle1 = tokio::spawn(async move { manager_clone.wait_for_access().await }); + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(1)).await; + tokio::time::advance(Duration::from_secs(1)).await; + assert!(handle1.await.unwrap() >= Duration::from_secs(1)); + + // Attempt 2 -> 2s delay + manager.report_transient_error().await; + let manager_clone = manager.clone(); + let handle2 = tokio::spawn(async move { manager_clone.wait_for_access().await }); + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(1)).await; + tokio::time::advance(Duration::from_secs(2)).await; + assert!(handle2.await.unwrap() >= Duration::from_secs(2)); + + // Attempt 3 -> 4s delay + manager.report_transient_error().await; + let manager_clone = manager.clone(); + let handle3 = tokio::spawn(async move { manager_clone.wait_for_access().await }); + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(1)).await; + tokio::time::advance(Duration::from_secs(4)).await; + assert!(handle3.await.unwrap() >= Duration::from_secs(4)); + } + + #[tokio::test] + async fn test_quota_manager_min_delay_cap() { + tokio::time::pause(); + // Configure with 0s flat delay (invalid, should be capped at 1s) + let manager = Arc::new(QuotaManager::new_with_settings( + BackoffType::Flat, + Duration::from_secs(0), + )); + + manager.report_transient_error().await; + + let manager_clone = manager.clone(); + let handle = tokio::spawn(async move { manager_clone.wait_for_access().await }); + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(1)).await; + tokio::time::advance(Duration::from_secs(1)).await; + let slept = handle.await.unwrap(); + + assert!(slept >= Duration::from_secs(1)); + } + + #[tokio::test] + async fn test_quota_manager_report_success_resets() { + tokio::time::pause(); + let manager = Arc::new(QuotaManager::new_with_settings( + BackoffType::Exponential, + Duration::from_secs(30), + )); + + // Attempt 1 -> 1s delay + manager.report_transient_error().await; + + // Report success -> resets + manager.report_success().await; + + // Attempt 1 (after reset) -> should be 1s again (not 2s) + manager.report_transient_error().await; + let manager_clone = manager.clone(); + let handle = tokio::spawn(async move { manager_clone.wait_for_access().await }); + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(1)).await; + tokio::time::advance(Duration::from_secs(1)).await; + assert!(handle.await.unwrap() >= Duration::from_secs(1)); + } +} diff --git a/src/reviewer.rs b/src/reviewer.rs index aebf3859a..961400af3 100644 --- a/src/reviewer.rs +++ b/src/reviewer.rs @@ -13,8 +13,8 @@ // limitations under the License. use crate::ReviewStatus; -use crate::ai::quota::QuotaManager; -use crate::ai::{AiProvider, AiRequest, create_provider_cached}; +use crate::ai::quota::{BackoffType, QuotaManager}; +use crate::ai::{AiError, AiProvider, AiRequest, create_provider_cached}; use crate::baseline::{BaselineRegistry, BaselineResolution, extract_files_from_diff}; use crate::db::{AiInteractionParams, Database, Finding, PatchsetRow, Severity, ToolUsage}; use crate::email_policy::EmailPolicyConfig; @@ -119,12 +119,18 @@ impl Reviewer { .await .expect("Failed to create AI provider"); + let transient_type = match settings.ai.transient_backoff_type.as_str() { + "flat" => BackoffType::Flat, + _ => BackoffType::Exponential, + }; + let flat_delay = std::time::Duration::from_secs(settings.ai.transient_flat_backoff_secs); + let quota_manager = Arc::new(QuotaManager::new_with_settings(transient_type, flat_delay)); Self { db, settings, semaphore: Arc::new(Semaphore::new(concurrency)), baseline_registry, - quota_manager: Arc::new(QuotaManager::new()), + quota_manager, provider, } } @@ -1480,6 +1486,17 @@ async fn run_review_tool( worktree_path: Option<&Path>, provider: Arc, ) -> Result { + let quota_manager = if settings.ai.global_backoff { + quota_manager + } else { + let transient_type = match settings.ai.transient_backoff_type.as_str() { + "flat" => BackoffType::Flat, + _ => BackoffType::Exponential, + }; + let flat_delay = std::time::Duration::from_secs(settings.ai.transient_flat_backoff_secs); + Arc::new(QuotaManager::new_with_settings(transient_type, flat_delay)) + }; + let mut cmd = if let Some(ref override_bin) = settings.review.review_tool_override { Command::new(override_bin) } else { @@ -1678,22 +1695,26 @@ async fn run_review_tool( break Ok(resp); } Err(e) => { - let err_str = e.to_string(); - - // Categorize and report errors to QuotaManager - if err_str.contains("Quota exceeded") - || err_str.contains("429") - { - quota_manager - .report_quota_error(Duration::from_secs(60)) - .await; + let (is_quota, is_transient, server_delay) = if let Some(ai_err) = e.downcast_ref::() { + match ai_err { + AiError::QuotaExceeded(d) => (true, false, Some(*d)), + AiError::Transient(d, _) => (false, true, Some(*d)), + AiError::Fatal(_) => (false, false, None), + } + } else { + // Fallback string matching + let err_str = e.to_string(); + let q = err_str.contains("Quota exceeded") || err_str.contains("429"); + let t = err_str.contains("Transient error") || err_str.contains("503") || err_str.contains("529") || err_str.contains("overloaded"); + (q, t, None) + }; + + if is_quota { + let delay = server_delay.unwrap_or(Duration::from_secs(settings.ai.quota_backoff_secs)); + quota_manager.report_quota_error(delay).await; continue; } - if err_str.contains("Transient error") - || err_str.contains("503") - || err_str.contains("529") - || err_str.contains("overloaded") - { + if is_transient { quota_manager.report_transient_error().await; continue; } @@ -2510,7 +2531,12 @@ echo '{"patchset_id": 1, "patches": [{"index": 1, "status": "applied"}]}' let db = Arc::new(Database::new(&settings.database).await?); db.migrate().await?; - let quota_manager = Arc::new(QuotaManager::new()); + let transient_type = match settings.ai.transient_backoff_type.as_str() { + "flat" => BackoffType::Flat, + _ => BackoffType::Exponential, + }; + let flat_delay = std::time::Duration::from_secs(settings.ai.transient_flat_backoff_secs); + let quota_manager = Arc::new(QuotaManager::new_with_settings(transient_type, flat_delay)); let thread_id = db.create_thread("msg_id_1", "Subject", 1000).await?; db.create_message( @@ -2606,4 +2632,233 @@ echo '{"patchset_id": 1, "patches": [{"index": 1, "status": "applied"}]}' ); Ok(()) } + + struct MockTransientThenSuccessProvider { + attempts: Arc>, + } + #[async_trait] + impl AiProvider for MockTransientThenSuccessProvider { + async fn generate_content(&self, _request: AiRequest) -> Result { + let mut att = self.attempts.lock().await; + *att += 1; + if *att == 1 { + return Err(AiError::Transient( + std::time::Duration::from_secs(5), + "Simulated transient".to_string(), + ) + .into()); + } + Ok(AiResponse { + content: Some("Mocked AI response after transient".to_string()), + thought: None, + thought_signature: None, + tool_calls: None, + usage: None, + }) + } + fn estimate_tokens(&self, _request: &AiRequest) -> usize { + 0 + } + fn get_capabilities(&self) -> ProviderCapabilities { + ProviderCapabilities { + model_name: "mock".to_string(), + context_window_size: 1000, + } + } + } + + struct MockQuotaProvider { + attempts: Arc>, + } + #[async_trait] + impl AiProvider for MockQuotaProvider { + async fn generate_content(&self, _request: AiRequest) -> Result { + let mut att = self.attempts.lock().await; + *att += 1; + if *att == 1 { + return Err(AiError::QuotaExceeded(std::time::Duration::from_secs(5)).into()); + } + Ok(AiResponse { + content: Some("Mocked AI response after quota".to_string()), + thought: None, + thought_signature: None, + tool_calls: None, + usage: None, + }) + } + fn estimate_tokens(&self, _request: &AiRequest) -> usize { + 0 + } + fn get_capabilities(&self) -> ProviderCapabilities { + ProviderCapabilities { + model_name: "mock".to_string(), + context_window_size: 1000, + } + } + } + + async fn run_one_request_mock( + mut settings: Settings, + provider: Arc, + ) -> Result<()> { + let temp_dir = tempdir()?; + let bin_path = temp_dir.path().join("mock_review"); + + // Mock binary: sends one AI request then a final result. + let mock_script = r#"#!/bin/bash +read -r input +echo '{"type": "ai_request", "payload": {"messages": [{"role": "user", "content": "hello"}], "temperature": 0.5}}' +read -r ai_response +echo '{"patchset_id": 1, "patches": [{"index": 1, "status": "applied"}]}' +"#; + std::fs::write(&bin_path, mock_script)?; + std::fs::set_permissions(&bin_path, Permissions::from_mode(0o755))?; + settings.review.review_tool_override = Some(bin_path.clone()); + + let db = Arc::new(Database::new(&settings.database).await?); + db.migrate().await?; + let transient_type = match settings.ai.transient_backoff_type.as_str() { + "flat" => BackoffType::Flat, + _ => BackoffType::Exponential, + }; + let flat_delay = std::time::Duration::from_secs(settings.ai.transient_flat_backoff_secs); + let quota_manager = Arc::new(QuotaManager::new_with_settings(transient_type, flat_delay)); + + let thread_id = db.create_thread("msg_id_1", "Subject", 1000).await?; + db.create_message( + "msg_id_p1", + thread_id, + None, + "Author", + "Subject", + 1000, + "Body", + "", + "", + None, + None, + ) + .await?; + let ps_id = db + .create_patchset( + thread_id, None, "msg_id_1", "Subject", "Author", 1000, 1, 1, "", "", None, 1, + None, false, None, None, + ) + .await? + .unwrap(); + let p_id = db + .create_patch(ps_id, "msg_id_p1", 1, "diff --git a/foo.c b/foo.c\n+int x;") + .await?; + let review_id = db + .create_review(ps_id, Some(p_id), "mock", "mock", None, None) + .await?; + + run_review_tool( + ps_id, + &json!({}), + &settings, + db, + "HEAD", + Some(1), + None, + quota_manager, + review_id, + None, + provider, + ) + .await + .map(|_| ()) + } + + #[tokio::test] + async fn test_review_timeout_exemption() -> Result<()> { + // Run in real-time because child process spawning is asynchronous to tokio virtual clock. + let mut settings = Settings::new()?; + settings.database.url = ":memory:".to_string(); + settings.review.timeout_seconds = 1; // 1s active timeout + settings.ai.transient_backoff_type = "flat".to_string(); + settings.ai.transient_flat_backoff_secs = 2; // 2s transient backoff (longer than timeout) + + let provider = Arc::new(MockTransientThenSuccessProvider { + attempts: Arc::new(tokio::sync::Mutex::new(0)), + }); + + let start = std::time::Instant::now(); + let res = run_one_request_mock(settings, provider).await; + let elapsed = start.elapsed(); + + assert!( + res.is_ok(), + "Expected review to succeed despite sleeping longer than timeout. Error: {:?}", + res.err() + ); + assert!( + elapsed >= std::time::Duration::from_secs(2), + "Review finished too quickly, expected to sleep for backoff: {:?}", + elapsed + ); + Ok(()) + } + + #[tokio::test] + async fn test_review_local_backoff_isolation() -> Result<()> { + // Run in real-time because child process spawning is asynchronous to tokio virtual clock. + let mut settings = Settings::new()?; + settings.database.url = ":memory:".to_string(); + settings.ai.global_backoff = false; // LOCAL backoff! + settings.ai.quota_backoff_secs = 3; // 3s quota delay + + let provider_quota = Arc::new(MockQuotaProvider { + attempts: Arc::new(tokio::sync::Mutex::new(0)), + }); + let provider_success = Arc::new(MockProvider); // succeeds instantly + + let start = std::time::Instant::now(); + + let settings_quota = settings.clone(); + let handle_quota = + tokio::spawn(async move { run_one_request_mock(settings_quota, provider_quota).await }); + + let settings_success = settings.clone(); + let handle_success = + tokio::spawn( + async move { run_one_request_mock(settings_success, provider_success).await }, + ); + + // Wait for handle_success to finish (should finish very quickly, way before 3s quota backoff) + let res_success = handle_success.await.unwrap(); + assert!( + res_success.is_ok(), + "Success task failed: {:?}", + res_success.err() + ); + let success_elapsed = start.elapsed(); + assert!( + success_elapsed < std::time::Duration::from_secs(2), + "Success task took too long, was likely blocked by quota: {:?}", + success_elapsed + ); + + // Quota task should still be running (sleeping for 3s) + assert!( + !handle_quota.is_finished(), + "Quota task finished too early (should be sleeping)" + ); + + // Wait for handle_quota to finish + let res_quota = handle_quota.await.unwrap(); + assert!( + res_quota.is_ok(), + "Quota task failed: {:?}", + res_quota.err() + ); + let quota_elapsed = start.elapsed(); + assert!( + quota_elapsed >= std::time::Duration::from_secs(3), + "Quota task finished too quickly: {:?}", + quota_elapsed + ); + + Ok(()) + } } diff --git a/src/settings.rs b/src/settings.rs index 52cd9d501..b7a4f1409 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -212,6 +212,14 @@ pub struct AiSettings { pub response_cache: bool, #[serde(default = "default_response_cache_ttl_days")] pub response_cache_ttl_days: u64, + #[serde(default = "default_global_backoff")] + pub global_backoff: bool, + #[serde(default = "default_quota_backoff_secs")] + pub quota_backoff_secs: u64, + #[serde(default = "default_transient_backoff_type")] + pub transient_backoff_type: String, + #[serde(default = "default_transient_flat_backoff_secs")] + pub transient_flat_backoff_secs: u64, // Provider-specific settings pub claude: Option, pub gemini: Option, @@ -226,6 +234,22 @@ fn default_response_cache_ttl_days() -> u64 { 7 } +fn default_global_backoff() -> bool { + true +} + +fn default_quota_backoff_secs() -> u64 { + 60 +} + +fn default_transient_backoff_type() -> String { + "exponential".to_string() +} + +fn default_transient_flat_backoff_secs() -> u64 { + 30 +} + fn default_api_timeout_secs() -> u64 { 300 }