AstroResearch/src/agent/runtime/error_recovery/mod.rs

// src/agent/runtime/error_recovery/mod.rs
//
// LLM 调用错误恢复引擎。
// 子模块：
//   classification — FailoverReason / ErrorKind / classify_* 分类函数族
//   overflow      — 上下文溢出解析与安全 Token 计算

pub(crate) mod classification;
pub(crate) mod overflow;

use tracing::{info, warn};

use crate::agent::runtime::token_budget::TokenBudget;

// Re-export 公共类型，保持外部引用路径不变
pub use classification::{
    classify_error, classify_failover, recovery_suggestion, ErrorKind, FailoverReason,
};
pub use overflow::{
    calculate_safe_max_tokens, parse_context_overflow, parse_retry_after, ContextOverflowInfo,
};

#[derive(Debug, Clone, PartialEq)]
pub enum RecoveryStep {
    /// 调整 max_tokens 以适配上下文容量（ContextOverflow 首选策略）。
    /// 参考 Claude Code — 解析 "X + Y > Z" 错误后计算安全值。
    AdjustMaxTokens { new_max_tokens: usize },
    /// 尝试更激进的 micro_compact
    AggressiveCompact,
    /// 使用 LLM 摘要压缩
    ReactiveCompact,
    /// 提升 token 上限
    EscalateTokens { new_hard_limit: usize },
    /// 分轮恢复（注入 meta 消息）
    MultiTurn,
    /// 放弃，暴露错误
    Surface,
    /// 指数退避重试（用于 429/529 瞬态错误）
    RetryWithBackoff { attempt: u32, delay_ms: u64 },
}

/// 恢复尝试追踪
#[derive(Debug, Clone)]
pub struct RecoveryAttempts {
    pub aggressive_compact: bool,
    pub reactive_compact: bool,
    pub escalate_tokens: bool,
    pub multi_turn: bool,
}

impl RecoveryAttempts {
    pub fn new() -> Self {
        RecoveryAttempts {
            aggressive_compact: false,
            reactive_compact: false,
            escalate_tokens: false,
            multi_turn: false,
        }
    }

    /// 是否有未尝试的恢复步骤
    pub fn has_remaining(&self) -> bool {
        !self.aggressive_compact
            || !self.reactive_compact
            || !self.escalate_tokens
            || !self.multi_turn
    }

    /// 获取下一个应尝试的恢复步骤。
    ///
    /// 对于 ContextOverflow 类错误，恢复优先级为：
    ///   1. AdjustMaxTokens — 从错误消息解析容量，精确计算安全的 max_tokens（参考 Claude Code）
    ///   2. AggressiveCompact — 激进微压缩
    ///   3. ReactiveCompact  — LLM 摘要压缩
    ///   4. EscalateTokens   — 提升 hard_limit
    ///   5. MultiTurn        — 分轮
    pub fn next_step(
        &mut self,
        error_kind: &ErrorKind,
        overflow_info: Option<&ContextOverflowInfo>,
    ) -> Option<RecoveryStep> {
        if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) {
            return Some(RecoveryStep::RetryWithBackoff {
                attempt: 0,
                delay_ms: 500,
            });
        }

        match error_kind {
            ErrorKind::PromptTooLong | ErrorKind::TokenExhausted => {
                // Step 0: 尝试精确调整 max_tokens（仅在有 Overflow info 时）
                if let Some(info) = overflow_info {
                    if let Some(safe_tokens) = calculate_safe_max_tokens(info) {
                        // 仅在使用默认路径（尚未尝试过其他恢复方式）时触发
                        if !self.aggressive_compact
                            && !self.reactive_compact
                            && !self.escalate_tokens
                        {
                            return Some(RecoveryStep::AdjustMaxTokens {
                                new_max_tokens: safe_tokens,
                            });
                        }
                    }
                }
                // Step 1: 激进微压缩
                if !self.aggressive_compact {
                    self.aggressive_compact = true;
                    return Some(RecoveryStep::AggressiveCompact);
                }
                // Step 2: LLM 摘要压缩
                if !self.reactive_compact {
                    self.reactive_compact = true;
                    return Some(RecoveryStep::ReactiveCompact);
                }
                // Step 3: 提升硬限制
                if !self.escalate_tokens {
                    self.escalate_tokens = true;
                    return Some(RecoveryStep::EscalateTokens {
                        new_hard_limit: 64_000,
                    });
                }
                // Step 4: 分轮
                if !self.multi_turn {
                    self.multi_turn = true;
                    return Some(RecoveryStep::MultiTurn);
                }
            }
            _ => {
                if !self.multi_turn {
                    self.multi_turn = true;
                    return Some(RecoveryStep::Surface);
                }
            }
        }
        None
    }

    /// 向后兼容的无 overflow_info 版本。
    pub fn next_step_compat(&mut self, error_kind: &ErrorKind) -> Option<RecoveryStep> {
        self.next_step(error_kind, None)
    }
}

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

/// 错误恢复器
pub struct ErrorRecovery {
    pub attempts: RecoveryAttempts,
    pub token_budget: TokenBudget,
}

impl ErrorRecovery {
    pub fn new(token_budget: TokenBudget) -> Self {
        ErrorRecovery {
            attempts: RecoveryAttempts::new(),
            token_budget,
        }
    }

    pub fn try_recover(
        &mut self,
        error_kind: &ErrorKind,
        overflow_info: Option<&ContextOverflowInfo>,
    ) -> Option<RecoveryStep> {
        let step = self.attempts.next_step(error_kind, overflow_info);
        match &step {
            Some(RecoveryStep::AdjustMaxTokens { new_max_tokens }) => {
                info!(
                    "[ErrorRecovery] 尝试步骤 0/5: AdjustMaxTokens → {} (从错误消息计算)",
                    new_max_tokens
                );
                // 将硬限制下调至安全值，使下一次 LLM 调用能在容量内完成
                self.token_budget.hard_limit = *new_max_tokens;
            }
            Some(RecoveryStep::AggressiveCompact) => {
                info!("[ErrorRecovery] 尝试步骤 1/5: AggressiveCompact");
            }
            Some(RecoveryStep::ReactiveCompact) => {
                info!("[ErrorRecovery] 尝试步骤 2/5: ReactiveCompact");
            }
            Some(RecoveryStep::EscalateTokens { new_hard_limit }) => {
                info!(
                    "[ErrorRecovery] 尝试步骤 3/5: EscalateTokens → {}",
                    new_hard_limit
                );
                self.token_budget.escalate_hard_limit(*new_hard_limit);
            }
            Some(RecoveryStep::RetryWithBackoff { attempt, delay_ms }) => {
                info!(
                    "[ErrorRecovery] 退避重试: attempt={}, delay={}ms",
                    attempt, delay_ms
                );
            }
            Some(RecoveryStep::MultiTurn) => {
                info!("[ErrorRecovery] 尝试步骤 4/5: MultiTurn");
            }
            Some(RecoveryStep::Surface) => {
                warn!("[ErrorRecovery] 无法恢复，暴露错误");
            }
            None => {
                warn!("[ErrorRecovery] 所有恢复步骤已尝试完毕");
            }
        }
        step
    }

    /// 向后兼容的无 overflow_info 版本。
    pub fn try_recover_compat(&mut self, error_kind: &ErrorKind) -> Option<RecoveryStep> {
        self.try_recover(error_kind, None)
    }

    pub fn multi_turn_message() -> String {
        "由于 token 限制，当前回答被截断。请基于已收集的信息继续分析，\
         重点关注尚未完成的部分。你可以：\n\
         1. 总结已有发现\n\
         2. 使用 compress_context 手动压缩上下文\n\
         3. 分步完成剩余工作"
            .to_string()
    }

    pub fn is_recoverable(error_kind: &ErrorKind) -> bool {
        matches!(
            error_kind,
            ErrorKind::PromptTooLong
                | ErrorKind::TokenExhausted
                | ErrorKind::RateLimited
                | ErrorKind::Overloaded
        )
    }
}

/// 计算指数退避延迟（毫秒）
pub fn backoff_delay(attempt: u32, retry_after_secs: Option<u64>) -> u64 {
    if let Some(ra) = retry_after_secs {
        return ra * 1000;
    }
    let base = 500u64 * 2u64.pow(attempt.min(6));
    let base = base.min(32_000);
    let jitter = (base / 4) * (attempt as u64 % 5) / 5;
    base + jitter
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── FailoverReason tests ──

    #[test]
    fn test_is_retryable() {
        assert!(FailoverReason::Auth.is_retryable());
        assert!(FailoverReason::RateLimited.is_retryable());
        assert!(FailoverReason::Overloaded.is_retryable());
        assert!(!FailoverReason::AuthPermanent.is_retryable());
        assert!(!FailoverReason::ContentPolicyBlocked.is_retryable());
    }

    #[test]
    fn test_should_compress() {
        assert!(FailoverReason::ContextOverflow.should_compress());
        assert!(FailoverReason::PayloadTooLarge.should_compress());
        assert!(!FailoverReason::RateLimited.should_compress());
    }

    #[test]
    fn test_should_failover() {
        assert!(FailoverReason::Auth.should_failover());
        assert!(FailoverReason::Billing.should_failover());
        assert!(FailoverReason::ModelNotFound.should_failover());
        assert!(!FailoverReason::Timeout.should_failover());
    }

    #[test]
    fn test_is_permanent() {
        assert!(FailoverReason::AuthPermanent.is_permanent());
        assert!(FailoverReason::ContentPolicyBlocked.is_permanent());
        assert!(!FailoverReason::Auth.is_permanent());
    }

    #[test]
    fn test_is_user_error() {
        assert!(FailoverReason::ContextOverflow.is_user_error());
        assert!(FailoverReason::ImageTooLarge.is_user_error());
        assert!(!FailoverReason::Timeout.is_user_error());
    }

    #[test]
    fn test_failover_reason_display() {
        assert_eq!(FailoverReason::Auth.to_string(), "auth");
        assert_eq!(FailoverReason::RateLimited.to_string(), "rate_limited");
        assert_eq!(
            FailoverReason::Unknown("test".into()).to_string(),
            "unknown(test)"
        );
    }

    // ── Classification pipeline tests ──

    #[test]
    fn test_classify_rate_limited_429() {
        let reason = classify_failover("HTTP 429: Too Many Requests", None, None);
        assert_eq!(reason, FailoverReason::RateLimited);
    }

    #[test]
    fn test_classify_overloaded_529() {
        let reason = classify_failover("HTTP 529: Service Overloaded", None, None);
        assert_eq!(reason, FailoverReason::Overloaded);
    }

    #[test]
    fn test_classify_context_overflow() {
        let reason = classify_failover(
            "Error: prompt_too_long: context length exceeded",
            None,
            None,
        );
        assert_eq!(reason, FailoverReason::ContextOverflow);
    }

    #[test]
    fn test_classify_auth_401() {
        let reason = classify_failover("HTTP 401: Unauthorized - invalid api key", None, None);
        assert_eq!(reason, FailoverReason::Auth);
    }

    #[test]
    fn test_classify_billing_402() {
        let reason = classify_failover("HTTP 402: insufficient_quota", None, None);
        assert_eq!(reason, FailoverReason::Billing);
    }

    #[test]
    fn test_classify_content_policy_blocked() {
        let reason = classify_failover(
            "content_policy_violation: responsible_ai filter triggered",
            None,
            None,
        );
        assert_eq!(reason, FailoverReason::ContentPolicyBlocked);
    }

    #[test]
    fn test_classify_model_not_found() {
        let reason = classify_failover(
            "HTTP 404: model_not_found - deployment not found",
            None,
            None,
        );
        assert_eq!(reason, FailoverReason::ModelNotFound);
    }

    #[test]
    fn test_classify_transport_timeout() {
        let reason = classify_failover("Connection reset by peer", None, None);
        // Connection reset with no session info → timeout
        assert_eq!(reason, FailoverReason::Timeout);
    }

    #[test]
    fn test_classify_disconnect_with_large_session() {
        let reason = classify_failover(
            "connection reset",
            Some(8000),  // estimated tokens
            Some(10000), // context length (80% → >60%)
        );
        assert_eq!(reason, FailoverReason::ContextOverflow);
    }

    #[test]
    fn test_classify_disconnect_with_small_session() {
        let reason = classify_failover(
            "connection reset",
            Some(4000),  // estimated tokens
            Some(10000), // context length (40% → <60%)
        );
        assert_eq!(reason, FailoverReason::Timeout);
    }

    #[test]
    fn test_classify_priorities_content_policy_before_status() {
        // content_policy_violation should be detected before 400 status
        let reason = classify_failover(
            "HTTP 400: content_policy_violation - content management policy",
            None,
            None,
        );
        assert_eq!(reason, FailoverReason::ContentPolicyBlocked);
    }

    #[test]
    fn test_classify_unknown_fallback() {
        let reason = classify_failover("Some weird error that doesn't match anything", None, None);
        assert!(matches!(reason, FailoverReason::Unknown(_)));
    }

    #[test]
    fn test_classify_image_too_large() {
        let reason = classify_failover("HTTP 413: image too large, max 20MB", None, None);
        assert_eq!(reason, FailoverReason::ImageTooLarge);
    }

    #[test]
    fn test_classify_provider_policy_blocked() {
        let reason = classify_failover("HTTP 403: provider policy blocked this region", None, None);
        assert_eq!(reason, FailoverReason::ProviderPolicyBlocked);
    }

    #[test]
    fn test_classify_ssl_transient() {
        let reason = classify_failover("SSL handshake timeout: connection reset", None, None);
        assert_eq!(reason, FailoverReason::SslTransient);
    }

    // ── Backward compatibility tests ──

    #[test]
    fn test_classify_error_backward_compat() {
        let kind = classify_error("HTTP 429: Too Many Requests");
        assert_eq!(kind, ErrorKind::RateLimited);
    }

    #[test]
    fn test_classify_error_529_backward_compat() {
        let kind = classify_error("HTTP 529: Service Overloaded");
        assert_eq!(kind, ErrorKind::Overloaded);
    }

    // ── Recovery step tests ──

    #[test]
    fn test_all_steps_sequence() {
        let mut attempts = RecoveryAttempts::new();
        assert_eq!(
            attempts.next_step_compat(&ErrorKind::PromptTooLong),
            Some(RecoveryStep::AggressiveCompact)
        );
        assert_eq!(
            attempts.next_step_compat(&ErrorKind::PromptTooLong),
            Some(RecoveryStep::ReactiveCompact)
        );
        assert_eq!(
            attempts.next_step_compat(&ErrorKind::PromptTooLong),
            Some(RecoveryStep::EscalateTokens {
                new_hard_limit: 64_000
            })
        );
        assert_eq!(
            attempts.next_step_compat(&ErrorKind::PromptTooLong),
            Some(RecoveryStep::MultiTurn)
        );
        assert_eq!(attempts.next_step_compat(&ErrorKind::PromptTooLong), None);
    }

    #[test]
    fn test_model_error_goes_straight_to_surface() {
        let mut attempts = RecoveryAttempts::new();
        assert_eq!(
            attempts.next_step_compat(&ErrorKind::ModelError("test".into())),
            Some(RecoveryStep::Surface)
        );
        assert_eq!(
            attempts.next_step_compat(&ErrorKind::ModelError("test".into())),
            None
        );
    }

    #[test]
    fn test_has_remaining() {
        let mut attempts = RecoveryAttempts::new();
        assert!(attempts.has_remaining());
        for _ in 0..4 {
            attempts.next_step_compat(&ErrorKind::PromptTooLong);
        }
        assert!(!attempts.has_remaining());
    }

    #[test]
    fn test_is_recoverable() {
        assert!(ErrorRecovery::is_recoverable(&ErrorKind::PromptTooLong));
        assert!(ErrorRecovery::is_recoverable(&ErrorKind::TokenExhausted));
        assert!(ErrorRecovery::is_recoverable(&ErrorKind::RateLimited));
        assert!(ErrorRecovery::is_recoverable(&ErrorKind::Overloaded));
        assert!(!ErrorRecovery::is_recoverable(&ErrorKind::ModelError(
            "test".into()
        )));
        assert!(!ErrorRecovery::is_recoverable(&ErrorKind::Timeout));
    }

    #[test]
    fn test_parse_retry_after() {
        let err = "HTTP 429: retry_after=Some(30)";
        assert_eq!(parse_retry_after(err), Some(30));
    }

    #[test]
    fn test_parse_retry_after_header_format() {
        let err = "HTTP 429\nRetry-After: 60\n";
        assert_eq!(parse_retry_after(err), Some(60));
    }

    #[test]
    fn test_parse_retry_after_none() {
        let err = "HTTP 500: Internal Server Error";
        assert_eq!(parse_retry_after(err), None);
    }

    #[test]
    fn test_backoff_delay() {
        let d0 = backoff_delay(0, None);
        assert!(d0 >= 500 && d0 <= 700);

        let d3 = backoff_delay(3, None);
        assert!(d3 >= 4000 && d3 <= 5000);

        let d10 = backoff_delay(10, None);
        assert!(d10 <= 40_000);

        let d_ra = backoff_delay(0, Some(15));
        assert_eq!(d_ra, 15000);
    }

    #[test]
    fn test_rate_limited_goes_to_retry() {
        let mut attempts = RecoveryAttempts::new();
        let step = attempts.next_step_compat(&ErrorKind::RateLimited);
        assert!(matches!(step, Some(RecoveryStep::RetryWithBackoff { .. })));
    }

    #[test]
    fn test_token_budget_escalation() {
        let budget = TokenBudget::new(32_000, 40_000);
        let mut recovery = ErrorRecovery::new(budget);
        assert_eq!(recovery.token_budget.hard_limit, 40_000);

        recovery
            .attempts
            .next_step_compat(&ErrorKind::PromptTooLong); // aggressive
        recovery
            .attempts
            .next_step_compat(&ErrorKind::PromptTooLong); // reactive
        let step = recovery.try_recover_compat(&ErrorKind::PromptTooLong); // escalate

        assert_eq!(
            step,
            Some(RecoveryStep::EscalateTokens {
                new_hard_limit: 64_000
            })
        );
        assert_eq!(recovery.token_budget.hard_limit, 64_000);
    }

    #[test]
    fn test_recovery_suggestion_all_variants() {
        // Every variant should return a non-empty suggestion
        let variants = [
            FailoverReason::Auth,
            FailoverReason::AuthPermanent,
            FailoverReason::Billing,
            FailoverReason::RateLimited,
            FailoverReason::Overloaded,
            FailoverReason::ServerError,
            FailoverReason::Timeout,
            FailoverReason::ContextOverflow,
            FailoverReason::PayloadTooLarge,
            FailoverReason::ImageTooLarge,
            FailoverReason::ModelNotFound,
            FailoverReason::ProviderPolicyBlocked,
            FailoverReason::ContentPolicyBlocked,
            FailoverReason::FormatError,
            FailoverReason::InvalidEncryptedContent,
            FailoverReason::MultimodalToolContentUnsupported,
            FailoverReason::ThinkingSignature,
            FailoverReason::LongContextTier,
            FailoverReason::OauthLongContextBetaForbidden,
            FailoverReason::SslTransient,
            FailoverReason::Unknown("test".into()),
        ];
        for v in &variants {
            let suggestion = recovery_suggestion(v);
            assert!(!suggestion.is_empty(), "Empty suggestion for {:?}", v);
        }
    }

    // ── Context Overflow 解析测试 ──

    #[test]
    fn test_parse_overflow_anthropic_format() {
        let err = "input length and max_tokens exceed context limit: 180000 + 32000 > 200000. \
                   please reduce the length of the messages or the max_tokens.";
        let info = parse_context_overflow(err).unwrap();
        assert_eq!(info.input_length, 180000);
        assert_eq!(info.requested_max_tokens, 32000);
        assert_eq!(info.context_limit, 200000);
    }

    #[test]
    fn test_parse_overflow_anthropic_format_different_numbers() {
        let err = "Error: input length and max_tokens exceed context limit: 50000 + 16000 > 64000";
        let info = parse_context_overflow(err).unwrap();
        assert_eq!(info.input_length, 50000);
        assert_eq!(info.requested_max_tokens, 16000);
        assert_eq!(info.context_limit, 64000);
    }

    #[test]
    fn test_parse_overflow_generic_format() {
        let err = "context length exceeded: 130000 + 32000 > 128000";
        let info = parse_context_overflow(err).unwrap();
        assert_eq!(info.input_length, 130000);
        assert_eq!(info.requested_max_tokens, 32000);
        assert_eq!(info.context_limit, 128000);
    }

    #[test]
    fn test_parse_overflow_no_match() {
        let err = "Something went wrong but no overflow pattern here";
        assert!(parse_context_overflow(err).is_none());
    }

    #[test]
    fn test_calculate_safe_max_tokens_normal() {
        let info = ContextOverflowInfo {
            input_length: 180000,
            requested_max_tokens: 32000,
            context_limit: 200000,
        };
        // available = 200000 - 180000 = 20000
        // adjusted = 20000 - 1000 = 19000
        assert_eq!(calculate_safe_max_tokens(&info), Some(19000));
    }

    #[test]
    fn test_calculate_safe_max_tokens_no_room() {
        let info = ContextOverflowInfo {
            input_length: 199000,
            requested_max_tokens: 32000,
            context_limit: 200000,
        };
        // available = 200000 - 199000 = 1000
        // adjusted = 1000 - 1000 = 0 < MIN_ACCEPTABLE_TOKENS(500)
        assert_eq!(calculate_safe_max_tokens(&info), None);
    }

    #[test]
    fn test_calculate_safe_max_tokens_input_exceeds_context() {
        let info = ContextOverflowInfo {
            input_length: 210000,
            requested_max_tokens: 1000,
            context_limit: 200000,
        };
        // input > context → None (must compress)
        assert_eq!(calculate_safe_max_tokens(&info), None);
    }

    #[test]
    fn test_recovery_uses_adjust_max_tokens_first() {
        let mut attempts = RecoveryAttempts::new();
        let info = ContextOverflowInfo {
            input_length: 180000,
            requested_max_tokens: 32000,
            context_limit: 200000,
        };
        let step = attempts.next_step(&ErrorKind::PromptTooLong, Some(&info));
        assert_eq!(
            step,
            Some(RecoveryStep::AdjustMaxTokens {
                new_max_tokens: 19000
            })
        );
    }

    #[test]
    fn test_recovery_falls_back_to_compact_when_no_room() {
        let mut attempts = RecoveryAttempts::new();
        // Input almost fills context — no room to adjust max_tokens
        let info = ContextOverflowInfo {
            input_length: 199500,
            requested_max_tokens: 32000,
            context_limit: 200000,
        };
        let step = attempts.next_step(&ErrorKind::PromptTooLong, Some(&info));
        // Should skip AdjustMaxTokens and go to AggressiveCompact
        assert_eq!(step, Some(RecoveryStep::AggressiveCompact));
    }

    #[test]
    fn test_recovery_adjust_after_compact_exhausted() {
        let mut attempts = RecoveryAttempts::new();
        // Mark aggressive_compact as tried
        attempts.aggressive_compact = true;

        let info = ContextOverflowInfo {
            input_length: 180000,
            requested_max_tokens: 32000,
            context_limit: 200000,
        };
        // AdjustMaxTokens check is skipped when aggressive_compact is already true
        let step = attempts.next_step(&ErrorKind::PromptTooLong, Some(&info));
        assert_eq!(step, Some(RecoveryStep::ReactiveCompact));
    }

    #[test]
    fn test_try_recover_adjusts_token_budget() {
        let budget = TokenBudget::new(32_000, 40_000);
        let mut recovery = ErrorRecovery::new(budget);
        let info = ContextOverflowInfo {
            input_length: 180000,
            requested_max_tokens: 32000,
            context_limit: 200000,
        };
        let step = recovery.try_recover(&ErrorKind::PromptTooLong, Some(&info));
        assert_eq!(
            step,
            Some(RecoveryStep::AdjustMaxTokens {
                new_max_tokens: 19000
            })
        );
        // Token budget hard_limit should be adjusted down
        assert_eq!(recovery.token_budget.hard_limit, 19000);
    }
}