Skip to main content

ai_agent/services/
rate_limit_messages.rs

1//! Rate limit messages - centralized rate limit message generation.
2//!
3////! Translates rateLimitMessages.ts from claude code.
4
5use crate::services::claude_ai_limits::{ClaudeAILimits, QuotaStatus, RateLimitType};
6
7pub const RATE_LIMIT_ERROR_PREFIXES: &[&str] = &[
8    "You've hit your",
9    "You've used",
10    "You're now using extra usage",
11    "You're close to",
12    "You're out of extra usage",
13];
14
15pub fn is_rate_limit_error_message(text: &str) -> bool {
16    RATE_LIMIT_ERROR_PREFIXES
17        .iter()
18        .any(|prefix| text.starts_with(prefix))
19}
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub enum MessageSeverity {
23    Error,
24    Warning,
25}
26
27#[derive(Debug, Clone)]
28pub struct RateLimitMessage {
29    pub message: String,
30    pub severity: MessageSeverity,
31}
32
33pub fn get_rate_limit_message(limits: &ClaudeAILimits, _model: &str) -> Option<RateLimitMessage> {
34    if limits.is_using_overage == Some(true) {
35        if limits.overage_status == Some(QuotaStatus::AllowedWarning) {
36            return Some(RateLimitMessage {
37                message: "You're close to your extra usage spending limit".to_string(),
38                severity: MessageSeverity::Warning,
39            });
40        }
41        return None;
42    }
43
44    if limits.status == QuotaStatus::Rejected {
45        return Some(RateLimitMessage {
46            message: get_limit_reached_text(limits, _model),
47            severity: MessageSeverity::Error,
48        });
49    }
50
51    if limits.status == QuotaStatus::AllowedWarning {
52        const WARNING_THRESHOLD: f64 = 0.7;
53
54        if let Some(util) = limits.utilization {
55            if util < WARNING_THRESHOLD {
56                return None;
57            }
58        }
59
60        let text = get_early_warning_text(limits);
61        if let Some(text) = text {
62            return Some(RateLimitMessage {
63                message: text,
64                severity: MessageSeverity::Warning,
65            });
66        }
67    }
68
69    None
70}
71
72pub fn get_rate_limit_error_message(limits: &ClaudeAILimits, model: &str) -> Option<String> {
73    get_rate_limit_message(limits, model)
74        .filter(|m| m.severity == MessageSeverity::Error)
75        .map(|m| m.message)
76}
77
78pub fn get_rate_limit_warning(limits: &ClaudeAILimits, model: &str) -> Option<String> {
79    get_rate_limit_message(limits, model)
80        .filter(|m| m.severity == MessageSeverity::Warning)
81        .map(|m| m.message)
82}
83
84fn get_limit_reached_text(limits: &ClaudeAILimits, model: &str) -> String {
85    let reset_message = limits
86        .resets_at
87        .map(|r| format!(" · resets {}", format_reset_time(r, true)))
88        .unwrap_or_default();
89
90    if limits.overage_status == Some(QuotaStatus::Rejected) {
91        if limits.overage_disabled_reason
92            == Some(crate::services::claude_ai_limits::OverageDisabledReason::OutOfCredits)
93        {
94            return format!("You're out of extra usage{}", reset_message);
95        }
96        return format_limit_reached_text("limit", &reset_message, model);
97    }
98
99    match &limits.rate_limit_type {
100        Some(RateLimitType::SevenDaySonnet) => {
101            format_limit_reached_text("Sonnet limit", &reset_message, model)
102        }
103        Some(RateLimitType::SevenDayOpus) => {
104            format_limit_reached_text("Opus limit", &reset_message, model)
105        }
106        Some(RateLimitType::SevenDay) => {
107            format_limit_reached_text("weekly limit", &reset_message, model)
108        }
109        Some(RateLimitType::FiveHour) => {
110            format_limit_reached_text("session limit", &reset_message, model)
111        }
112        _ => format_limit_reached_text("usage limit", &reset_message, model),
113    }
114}
115
116fn get_early_warning_text(limits: &ClaudeAILimits) -> Option<String> {
117    let limit_name = match &limits.rate_limit_type {
118        Some(RateLimitType::SevenDay) => "weekly limit",
119        Some(RateLimitType::FiveHour) => "session limit",
120        Some(RateLimitType::SevenDayOpus) => "Opus limit",
121        Some(RateLimitType::SevenDaySonnet) => "Sonnet limit",
122        Some(RateLimitType::Overage) => "extra usage",
123        None => return None,
124    };
125
126    let used = limits.utilization.map(|u| (u * 100.0) as u32);
127
128    let reset_time = limits.resets_at.map(|r| format_reset_time(r, true));
129
130    let base = match (used, reset_time) {
131        (Some(u), Some(t)) => format!("You've used {}% of your {} · resets {}", u, limit_name, t),
132        (Some(u), None) => format!("You've used {}% of your {}", u, limit_name),
133        (None, Some(t)) => format!("Approaching {} · resets {}", limit_name, t),
134        (None, None) => format!("Approaching {}", limit_name),
135    };
136
137    Some(base)
138}
139
140fn format_limit_reached_text(limit: &str, reset_message: &str, _model: &str) -> String {
141    format!("You've hit your {}{}", limit, reset_message)
142}
143
144fn format_reset_time(resets_at: u64, _short: bool) -> String {
145    let now = std::time::SystemTime::now()
146        .duration_since(std::time::UNIX_EPOCH)
147        .unwrap_or_default()
148        .as_secs();
149
150    let diff = resets_at.saturating_sub(now);
151
152    if diff < 60 {
153        return "less than a minute".to_string();
154    }
155
156    let minutes = diff / 60;
157    if minutes < 60 {
158        return format!("{} minute{}", minutes, if minutes == 1 { "" } else { "s" });
159    }
160
161    let hours = minutes / 60;
162    if hours < 24 {
163        return format!("{} hour{}", hours, if hours == 1 { "" } else { "s" });
164    }
165
166    let days = hours / 24;
167    format!("{} day{}", days, if days == 1 { "" } else { "s" })
168}
169
170pub fn get_using_overage_text(limits: &ClaudeAILimits) -> String {
171    let reset_time = limits
172        .resets_at
173        .map(|r| format_reset_time(r, true))
174        .unwrap_or_default();
175
176    let limit_name = match &limits.rate_limit_type {
177        Some(RateLimitType::FiveHour) => "session limit",
178        Some(RateLimitType::SevenDay) => "weekly limit",
179        Some(RateLimitType::SevenDayOpus) => "Opus limit",
180        Some(RateLimitType::SevenDaySonnet) => "Sonnet limit",
181        _ => "",
182    };
183
184    if limit_name.is_empty() {
185        return "Now using extra usage".to_string();
186    }
187
188    let reset_msg = if !reset_time.is_empty() {
189        format!(" · Your {} resets {}", limit_name, reset_time)
190    } else {
191        String::new()
192    };
193
194    format!("You're now using extra usage{}", reset_msg)
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn test_is_rate_limit_error_message() {
203        assert!(is_rate_limit_error_message("You've hit your session limit"));
204        assert!(is_rate_limit_error_message(
205            "You've used 75% of your weekly limit"
206        ));
207        assert!(is_rate_limit_error_message("You're now using extra usage"));
208        assert!(!is_rate_limit_error_message("Something else"));
209    }
210
211    #[test]
212    fn test_format_reset_time() {
213        let now = std::time::SystemTime::now()
214            .duration_since(std::time::UNIX_EPOCH)
215            .unwrap_or_default()
216            .as_secs();
217
218        assert_eq!(format_reset_time(now + 30, false), "less than a minute");
219        assert_eq!(format_reset_time(now + 120, false), "2 minutes");
220        assert_eq!(format_reset_time(now + 3600, false), "1 hour");
221        assert_eq!(format_reset_time(now + 86400, false), "1 day");
222    }
223}