vtcode-core 0.136.3

//! Orchestrator retry and error handling module
//!
//! This module provides robust error handling for orchestrator response failures,
//! including retry mechanisms with exponential backoff and fallback strategies.

use crate::config::models::ModelId;
use crate::error::{ErrorCode, Result as VtCodeResult, VtCodeError};
use crate::retry::{RetryEvent, RetryPolicy, RetryPolicyCoreExt, RetryStep};
use std::future::Future;
use std::result::Result as StdResult;
use std::time::{Duration, Instant};
use tokio::time::sleep;
use tracing::{info, warn};

use vtcode_config::constants::execution::{
    DEFAULT_ORCHESTRATOR_MAX_RETRIES, DEFAULT_ORCHESTRATOR_RETRY_INITIAL_DELAY_MS,
    DEFAULT_ORCHESTRATOR_RETRY_MAX_DELAY_SECS, DEFAULT_ORCHESTRATOR_RETRY_MULTIPLIER,
};

/// Statistics about retry attempts
#[derive(Debug, Clone, Default)]
pub struct RetryStats {
    /// Total number of attempts made (including first attempts).
    pub total_attempts: u32,
    /// Number of operations that succeeded after at least one retry.
    pub successful_retries: u32,
    /// Number of retry sequences that exhausted all attempts.
    pub failed_retries: u32,
    /// Number of times the fallback model was activated.
    pub fallback_activations: u32,
    /// Cumulative time spent waiting during backoff delays.
    pub total_backoff_time: Duration,
}

/// Retry manager for orchestrator operations.
///
/// Uses [`RetryPolicy`] directly for retry decisions, eliminating the
/// redundant `RetryConfig` → `RetryPolicy` conversion overhead.
#[derive(Debug)]
pub struct RetryManager {
    policy: RetryPolicy,
    stats: RetryStats,
}

impl Default for RetryManager {
    fn default() -> Self {
        Self::new()
    }
}

impl RetryManager {
    /// Create a new retry manager with default configuration
    pub fn new() -> Self {
        Self {
            policy: RetryPolicy::from_retries(
                DEFAULT_ORCHESTRATOR_MAX_RETRIES,
                Duration::from_millis(DEFAULT_ORCHESTRATOR_RETRY_INITIAL_DELAY_MS),
                Duration::from_secs(DEFAULT_ORCHESTRATOR_RETRY_MAX_DELAY_SECS),
                DEFAULT_ORCHESTRATOR_RETRY_MULTIPLIER,
            ),
            stats: RetryStats::default(),
        }
    }

    /// Create a new retry manager with a custom retry policy
    pub fn with_policy(policy: RetryPolicy) -> Self {
        Self { policy, stats: RetryStats::default() }
    }

    /// Get the current retry statistics
    pub fn stats(&self) -> &RetryStats {
        &self.stats
    }

    /// Reset retry statistics
    pub fn reset_stats(&mut self) {
        self.stats = RetryStats::default();
    }

    /// Execute an operation with retry and fallback logic
    pub async fn execute_with_retry<F, Fut, T, E>(
        &mut self,
        operation_name: &str,
        primary_model: &ModelId,
        fallback_model: Option<&ModelId>,
        operation: F,
    ) -> VtCodeResult<T>
    where
        F: Fn(ModelId) -> Fut,
        Fut: Future<Output = StdResult<T, E>>,
        E: Into<VtCodeError>,
        T: Clone,
    {
        let start_time = Instant::now();
        let policy = &self.policy;
        let mut last_error: Option<VtCodeError> = None;
        let mut observer = RetryObserver {
            stats: &mut self.stats,
            operation_name,
            primary_model,
            max_attempts: policy.max_attempts,
        };

        // Primary model retry loop — structure kept inline because the
        // orchestrator's `operation: Fn` cannot be passed to
        // `run_with_retry` (which requires `FnMut` + `Send` futures).
        // The `RetryObserver` centralises the bookkeeping so the loop
        // stays DRY at the observability layer.
        for attempt in 0..policy.max_attempts {
            observer.observe(RetryEvent::AttemptStart { attempt, max_attempts: policy.max_attempts });

            let err = match operation(primary_model.clone()).await {
                Ok(result) => {
                    observer.observe(RetryEvent::Success { attempt });
                    return Ok(result);
                }
                Err(err) => {
                    let err: VtCodeError = err.into();
                    err
                }
            };

            let category_was_retryable = err.category.is_retryable();
            let step = policy.step_for_vtcode_error(err, attempt, None);

            match step {
                RetryStep::GiveUp { decision, error } => {
                    observer.observe(RetryEvent::GiveUp {
                        attempt,
                        error: &error,
                        decision: &decision,
                        category_was_retryable,
                    });
                    return Err(error);
                }
                RetryStep::Backoff { delay, decision, error } => {
                    observer.observe(RetryEvent::Backoff {
                        attempt,
                        error: &error,
                        decision: &decision,
                        delay,
                        category_was_retryable,
                    });
                    last_error = Some(error);
                    sleep(delay).await;
                }
            }
        }

        // If we have a fallback model and primary failed, try fallback
        if let Some(fallback) = fallback_model {
            warn!(
                operation = operation_name,
                attempts = policy.max_attempts,
                primary_model = ?primary_model,
                fallback_model = ?fallback,
                "primary model failed; attempting fallback"
            );
            self.stats.fallback_activations += 1;

            match operation(fallback.clone()).await {
                Ok(result) => {
                    info!(operation = operation_name, model = ?fallback, "fallback model succeeded");
                    return Ok(result);
                }
                Err(err) => {
                    let err: VtCodeError = err.into();
                    let fallback_err = err
                        .with_context(format!("fallback model '{fallback}' failed for operation '{operation_name}'"));
                    warn!(
                        operation = operation_name,
                        model = ?fallback,
                        error = %fallback_err,
                        "fallback model failed"
                    );
                    last_error = Some(fallback_err);
                }
            }
        }

        let total_time = start_time.elapsed();
        warn!(
            operation = operation_name,
            attempts = policy.max_attempts,
            total_time = %humantime::format_duration(total_time),
            primary_model = ?primary_model,
            fallback_model = ?fallback_model,
            "operation failed after retries"
        );

        Err(last_error.unwrap_or_else(|| {
            VtCodeError::execution(
                ErrorCode::ToolExecutionFailed,
                format!("operation '{operation_name}' failed after {} attempts", policy.max_attempts),
            )
            .with_context(format!("primary model: {primary_model}, fallback model: {fallback_model:?}"))
        }))
    }
}

/// Lightweight observer that updates [`RetryStats`] and emits tracing
/// events for each retry lifecycle event. Used by
/// [`RetryManager::execute_with_retry`] so the bookkeeping stays
/// co-located with the retry decision rather than scattered across
/// the loop body.
struct RetryObserver<'a> {
    stats: &'a mut RetryStats,
    operation_name: &'a str,
    primary_model: &'a ModelId,
    max_attempts: u32,
}

impl RetryObserver<'_> {
    fn observe(&mut self, event: RetryEvent<'_>) {
        match event {
            RetryEvent::AttemptStart { attempt, .. } => {
                self.stats.total_attempts += 1;
                info!(
                    attempt = attempt + 1,
                    max_attempts = self.max_attempts,
                    operation = self.operation_name,
                    model = ?self.primary_model,
                    "retry attempt starting"
                );
            }
            RetryEvent::Success { attempt } if attempt > 0 => {
                self.stats.successful_retries += 1;
                info!(
                    attempt = attempt + 1,
                    operation = self.operation_name,
                    model = ?self.primary_model,
                    "operation succeeded after retry"
                );
            }
            RetryEvent::Success { .. } => {}
            RetryEvent::GiveUp { attempt, error, decision, category_was_retryable } => {
                warn!(
                    attempt = attempt + 1,
                    max_attempts = self.max_attempts,
                    operation = self.operation_name,
                    model = ?self.primary_model,
                    error = %error,
                    category = ?decision.category,
                    retryable = category_was_retryable,
                    "non-retryable error, giving up"
                );
                if attempt + 1 == self.max_attempts && category_was_retryable {
                    self.stats.failed_retries += 1;
                }
            }
            RetryEvent::Backoff { attempt, error, decision, delay, .. } => {
                warn!(
                    attempt = attempt + 1,
                    max_attempts = self.max_attempts,
                    operation = self.operation_name,
                    model = ?self.primary_model,
                    error = %error,
                    category = ?decision.category,
                    "operation attempt failed"
                );
                self.stats.total_backoff_time += delay;
                if attempt + 2 == self.max_attempts {
                    self.stats.failed_retries += 1;
                }
                info!(
                    delay_ms = delay.as_millis() as u64,
                    next_attempt = attempt + 2,
                    operation = self.operation_name,
                    category = ?decision.category,
                    "backing off before retry"
                );
            }
            RetryEvent::Exhausted { .. } => {}
        }
    }
}

/// Check if a response is considered empty or invalid
pub fn is_empty_response(response: &serde_json::Value) -> bool {
    match response {
        serde_json::Value::Null => true,
        serde_json::Value::String(s) => s.trim().is_empty(),
        serde_json::Value::Object(obj) => {
            obj.is_empty() ||
            // Check for common empty response patterns
            (obj.get("candidates").is_some_and(|c| c.as_array().is_some_and(|arr| arr.is_empty()))) ||
            (obj.get("content").is_some_and(|c| match c {
                serde_json::Value::String(s) => s.trim().is_empty(),
                serde_json::Value::Array(arr) => arr.is_empty(),
                _ => false,
            }))
        }
        serde_json::Value::Array(arr) => arr.is_empty(),
        _ => false,
    }
}

/// Detect if an error indicates a temporary failure that should be retried.
/// Uses the shared VT Code retry policy for typed and fallback classification.
pub fn is_retryable_error(error: &anyhow::Error) -> bool {
    RetryPolicy::default().decision_for_anyhow(error, 0, None).retryable
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::{ErrorCode, VtCodeError};
    use anyhow::anyhow;
    use serde_json::json;
    use std::sync::{Arc, Mutex};

    #[test]
    fn test_empty_response_detection() {
        assert!(is_empty_response(&serde_json::Value::Null));
        assert!(is_empty_response(&json!("")));
        assert!(is_empty_response(&json!("  ")));
        assert!(is_empty_response(&json!({})));
        assert!(is_empty_response(&json!([])));
        assert!(is_empty_response(&json!({"candidates": []})));
        assert!(is_empty_response(&json!({"content": ""})));
        assert!(is_empty_response(&json!({"content": []})));

        assert!(!is_empty_response(&json!("hello")));
        assert!(!is_empty_response(&json!({"content": "hello"})));
        assert!(!is_empty_response(&json!({"candidates": [{"content": "hello"}]})));
    }

    #[test]
    fn test_retryable_error_detection() {
        assert!(is_retryable_error(&anyhow!("Connection timeout")));
        assert!(is_retryable_error(&anyhow!("Rate limit exceeded")));
        assert!(is_retryable_error(&anyhow!("HTTP 503 Service Unavailable")));
        assert!(is_retryable_error(&anyhow!("Network error")));

        assert!(is_retryable_error(&anyhow!("HTTP 429 Too Many Requests")));
        assert!(is_retryable_error(&anyhow!("Error 429: rate limited")));

        assert!(!is_retryable_error(&anyhow!("Invalid API key")));
        assert!(!is_retryable_error(&anyhow!("Permission denied")));
        assert!(!is_retryable_error(&anyhow!("Invalid model")));
        assert!(!is_retryable_error(&anyhow!("You exceeded your current quota")));
        assert!(!is_retryable_error(&anyhow!("insufficient_quota")));
        assert!(!is_retryable_error(&anyhow!("429 quota exceeded")));
    }

    #[tokio::test]
    async fn test_retry_manager_success_first_attempt() {
        let mut manager = RetryManager::new();
        let result = manager
            .execute_with_retry(
                "test_operation",
                &ModelId::Gemini35Flash,
                Some(&ModelId::Gemini31ProPreview),
                |_model| async { Ok::<String, anyhow::Error>("success".to_owned()) },
            )
            .await;

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "success");
        assert_eq!(manager.stats().total_attempts, 1);
        assert_eq!(manager.stats().successful_retries, 0);
        assert_eq!(manager.stats().fallback_activations, 0);
    }

    #[tokio::test]
    async fn test_retry_manager_success_after_retry() {
        let mut manager = RetryManager::with_policy(RetryPolicy::from_retries(
            2,
            Duration::from_secs(0), // No delay for test
            Duration::from_secs(1),
            2.0,
        ));

        let attempt_count = Arc::new(Mutex::new(0));
        let attempt_count_clone = attempt_count.clone();
        let result = manager
            .execute_with_retry(
                "test_operation",
                &ModelId::Gemini35Flash,
                Some(&ModelId::Gemini31ProPreview),
                move |_model| {
                    let attempt_count = attempt_count_clone.clone();
                    async move {
                        let mut count = attempt_count.lock().unwrap();
                        *count += 1;
                        if *count < 2 {
                            Err(VtCodeError::network(ErrorCode::ConnectionFailed, "temporary failure"))
                        } else {
                            Ok::<String, VtCodeError>("success".to_owned())
                        }
                    }
                },
            )
            .await;

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "success");
        assert_eq!(manager.stats().total_attempts, 2);
        assert_eq!(manager.stats().successful_retries, 1);
    }
}