Skip to main content

rate_limits/
reset_time.rs

1use crate::convert;
2use crate::error::{Error, Result};
3use headers::HeaderValue;
4use time::format_description::well_known::{Iso8601, Rfc2822};
5use time::{Duration, OffsetDateTime, PrimitiveDateTime};
6
7/// The kind of rate limit reset time
8///
9/// There are different ways to denote rate limits reset times.
10/// Some vendors use seconds, others use a timestamp format for example.
11///
12/// This enum lists all known variants.
13#[derive(Copy, Clone, Debug, PartialEq)]
14#[non_exhaustive]
15pub enum ResetTimeKind {
16    /// Number of seconds until rate limit is lifted
17    Seconds,
18    /// Unix timestamp (UTC epoch seconds)
19    /// when rate limit will be lifted
20    Timestamp,
21    /// Unix timestamp in millisecond resolution (UTC epoch milliseconds)
22    /// when rate limit will be lifted
23    TimestampMillis,
24    /// RFC 2822 date when rate limit will be lifted
25    ImfFixdate,
26    /// ISO 8601 date when rate limit will be lifted
27    Iso8601,
28    /// OpenAI-style duration string (e.g. "1s", "6m0s") until rate limit is lifted
29    OpenAIDuration,
30}
31
32/// Reset time of rate limiting
33///
34/// There are different variants on how to specify reset times
35/// in rate limit headers. The most common ones are seconds and datetime.
36#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd)]
37pub enum ResetTime {
38    /// Number of seconds until rate limit is lifted
39    Seconds(usize),
40    /// Date when rate limit will be lifted
41    DateTime(OffsetDateTime),
42}
43
44impl ResetTime {
45    /// Create a new reset time from a header value and a reset time kind
46    ///
47    /// # Errors
48    ///
49    /// This function returns an error if the header value cannot be parsed
50    /// or if the reset time kind is unknown.
51    pub fn new(value: &HeaderValue, kind: ResetTimeKind) -> Result<Self> {
52        let value = value.to_str()?;
53        match kind {
54            ResetTimeKind::Seconds => Ok(ResetTime::Seconds(convert::to_usize(value)?)),
55            ResetTimeKind::Timestamp => Ok(Self::DateTime(
56                OffsetDateTime::from_unix_timestamp(convert::to_i64(value)?)
57                    .map_err(Error::Time)?,
58            )),
59            ResetTimeKind::TimestampMillis => Ok(Self::DateTime(
60                OffsetDateTime::from_unix_timestamp_nanos(convert::to_i128(value)? * 1_000_000)
61                    .map_err(Error::Time)?,
62            )),
63            ResetTimeKind::Iso8601 => {
64                // https://github.com/time-rs/time/issues/378
65                let d = PrimitiveDateTime::parse(value, &Iso8601::PARSING).map_err(Error::Parse)?;
66                Ok(ResetTime::DateTime(d.assume_utc()))
67            }
68            ResetTimeKind::ImfFixdate => {
69                let d = PrimitiveDateTime::parse(value, &Rfc2822).map_err(Error::Parse)?;
70                Ok(ResetTime::DateTime(d.assume_utc()))
71            }
72            ResetTimeKind::OpenAIDuration => {
73                let seconds = parse_openai_duration_to_seconds(value)?;
74                Ok(ResetTime::Seconds(seconds))
75            }
76        }
77    }
78
79    /// Get the number of seconds until the rate limit gets lifted.
80    #[must_use]
81    pub fn seconds(&self) -> usize {
82        match self {
83            ResetTime::Seconds(s) => *s,
84            // OffsetDateTime is not timezone aware, so we need to convert it to UTC
85            // and then convert it to seconds.
86            // If the reset time is in the past, we return 0.
87            #[allow(clippy::cast_possible_truncation)]
88            ResetTime::DateTime(d) => {
89                let diff = *d - OffsetDateTime::now_utc();
90                let seconds = diff.whole_seconds();
91                if seconds < 0 { 0 } else { seconds as usize }
92            }
93        }
94    }
95
96    /// Convert reset time to duration
97    #[must_use]
98    pub fn duration(&self) -> Duration {
99        match self {
100            ResetTime::Seconds(s) => Duration::seconds(*s as i64),
101            ResetTime::DateTime(d) => {
102                Duration::seconds((*d - OffsetDateTime::now_utc()).whole_seconds())
103            }
104        }
105    }
106}
107
108/// Parse OpenAI duration string into seconds
109///
110/// Examples: "1s", "6m0s", "1h30m", "10ms"
111fn parse_openai_duration_to_seconds(value: &str) -> Result<usize> {
112    let value = value.trim();
113    if value.is_empty() {
114        return Err(Error::InvalidDuration(value.to_string()));
115    }
116
117    let mut total_seconds = 0;
118    let mut current_number = String::new();
119
120    let mut chars = value.chars().peekable();
121
122    while let Some(c) = chars.next() {
123        if c.is_ascii_digit() {
124            current_number.push(c);
125        } else {
126            if current_number.is_empty() {
127                return Err(Error::InvalidDuration(value.to_string()));
128            }
129            let n: usize = current_number
130                .parse()
131                .map_err(|_| Error::InvalidDuration(value.to_string()))?;
132            current_number.clear();
133
134            match c {
135                'd' => total_seconds += n * 86400,
136                'h' => total_seconds += n * 3600,
137                'm' => {
138                    // Check if next is 's' for 'ms'
139                    if let Some('s') = chars.peek() {
140                        chars.next(); // consume 's'
141                        // If it's > 0 ms, we round up to 1s to be safe for rate limits.
142                        if n > 0 {
143                            total_seconds += 1;
144                        }
145                    } else {
146                        total_seconds += n * 60;
147                    }
148                }
149                's' => total_seconds += n,
150                _ => return Err(Error::InvalidDuration(value.to_string())),
151            }
152        }
153    }
154
155    if !current_number.is_empty() {
156        return Err(Error::InvalidDuration(value.to_string()));
157    }
158
159    Ok(total_seconds)
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use headers::HeaderValue;
166
167    #[test]
168    fn test_parse_openai_duration() {
169        // Invalid
170        assert!(parse_openai_duration_to_seconds("").is_err());
171        assert!(parse_openai_duration_to_seconds("🤖").is_err());
172        assert!(parse_openai_duration_to_seconds("1").is_err());
173        assert!(parse_openai_duration_to_seconds("s").is_err());
174        assert!(parse_openai_duration_to_seconds("1x").is_err());
175        assert!(parse_openai_duration_to_seconds("1m30").is_err());
176        assert!(parse_openai_duration_to_seconds("around 30s").is_err());
177        assert!(parse_openai_duration_to_seconds("1m30s hello").is_err());
178
179        assert_eq!(parse_openai_duration_to_seconds("1s").unwrap(), 1);
180        assert_eq!(parse_openai_duration_to_seconds("1s ").unwrap(), 1);
181        assert_eq!(parse_openai_duration_to_seconds("1m").unwrap(), 60);
182        assert_eq!(parse_openai_duration_to_seconds("1h").unwrap(), 3600);
183        assert_eq!(parse_openai_duration_to_seconds("1d").unwrap(), 86400);
184
185        // Combined
186        assert_eq!(parse_openai_duration_to_seconds("1m30s").unwrap(), 90);
187        assert_eq!(parse_openai_duration_to_seconds("1h1m1s").unwrap(), 3661);
188        assert_eq!(parse_openai_duration_to_seconds("6m0s").unwrap(), 360);
189
190        // Milliseconds
191        assert_eq!(parse_openai_duration_to_seconds("10ms").unwrap(), 1);
192        assert_eq!(parse_openai_duration_to_seconds("0ms").unwrap(), 0);
193        assert_eq!(parse_openai_duration_to_seconds("1000ms").unwrap(), 1);
194    }
195
196    #[test]
197    fn test_reset_time_new_openai_duration() {
198        let v = HeaderValue::from_str("1h30m").unwrap();
199        let rt = ResetTime::new(&v, ResetTimeKind::OpenAIDuration).unwrap();
200        assert_eq!(rt, ResetTime::Seconds(5400));
201    }
202}