Skip to main content

seher/codexbar/
limit.rs

1//! Maps a codexbar usage payload to an [`AgentLimit`].
2//!
3//! Mirrors `seher-ts/packages/sdk/src/codexbar/limit.ts`: every rate window
4//! (primary/secondary/tertiary + extra windows) at `usedPercent >= 100` counts
5//! as limited, and the earliest reset is returned so the agent waits the minimum
6//! amount of time. Extension over seher-ts: a window whose `resetsAt` has
7//! already passed is treated as a stale snapshot (the window has presumably
8//! already reset server-side) rather than as evidence of an active limit, so
9//! it is excluded from consideration.
10
11use chrono::{DateTime, Utc};
12
13use super::client::{RunCodexBarUsageOptions, run_codexbar_usage};
14use super::errors::CodexBarError;
15use super::types::{CodexBarUsageResponse, CodexBarWindow};
16
17/// Outcome of a rate-limit check for a single provider.
18#[derive(Debug, Clone)]
19pub enum AgentLimit {
20    /// The provider has quota available.
21    NotLimited,
22    /// The provider is at-limit; `reset_time` is the earliest moment it frees up
23    /// (when known).
24    Limited { reset_time: Option<DateTime<Utc>> },
25}
26
27/// Reset fallback when codexbar reports a limited window without a parseable
28/// `resetsAt` (matches seher-ts's 5-minute fallback).
29const FALLBACK_RESET_SECS: i64 = 5 * 60;
30
31fn parse_resets_at(resets_at: Option<&str>, now: DateTime<Utc>) -> DateTime<Utc> {
32    if let Some(s) = resets_at
33        && let Ok(parsed) = DateTime::parse_from_rfc3339(s)
34    {
35        return parsed.with_timezone(&Utc);
36    }
37    now + chrono::Duration::seconds(FALLBACK_RESET_SECS)
38}
39
40fn is_limited(window: &CodexBarWindow) -> bool {
41    window.used_percent >= 100.0
42}
43
44fn classify(response: &CodexBarUsageResponse, now: DateTime<Utc>) -> AgentLimit {
45    let usage = &response.usage;
46    let mut windows: Vec<&CodexBarWindow> = Vec::new();
47    windows.extend(usage.primary.as_ref());
48    windows.extend(usage.secondary.as_ref());
49    windows.extend(usage.tertiary.as_ref());
50    if let Some(extra) = &usage.extra_rate_windows {
51        windows.extend(extra.iter().map(|named| &named.window));
52    }
53
54    let earliest = windows
55        .into_iter()
56        .filter(|w| is_limited(w))
57        .map(|w| parse_resets_at(w.resets_at.as_deref(), now))
58        // A window whose resetsAt has already passed is a stale snapshot (it
59        // has presumably reset server-side already), not an active limit.
60        // Windows with no parseable resetsAt fall back to `now + 5m` (see
61        // `parse_resets_at`), which is always in the future, so this filter
62        // never drops the no-resetsAt fallback case.
63        .filter(|reset| *reset > now)
64        .min();
65
66    match earliest {
67        Some(reset_time) => AgentLimit::Limited {
68            reset_time: Some(reset_time),
69        },
70        None => AgentLimit::NotLimited,
71    }
72}
73
74/// Determine whether `provider` is rate-limited by invoking codexbar with default options.
75///
76/// # Errors
77///
78/// Propagates [`CodexBarError`] from [`run_codexbar_usage`].
79pub async fn check_limit(provider: &str) -> Result<AgentLimit, CodexBarError> {
80    check_limit_with(provider, &RunCodexBarUsageOptions::default()).await
81}
82
83/// Like [`check_limit`] but with explicit [`RunCodexBarUsageOptions`].
84///
85/// # Errors
86///
87/// Propagates [`CodexBarError`] from [`run_codexbar_usage`].
88pub async fn check_limit_with(
89    provider: &str,
90    opts: &RunCodexBarUsageOptions,
91) -> Result<AgentLimit, CodexBarError> {
92    let response = run_codexbar_usage(provider, opts).await?;
93    Ok(classify(&response, Utc::now()))
94}
95
96#[cfg(test)]
97#[expect(clippy::expect_used, reason = "tests may panic on unexpected fixtures")]
98mod tests {
99    use super::*;
100
101    fn window(used_percent: f64, resets_at: Option<&str>) -> CodexBarWindow {
102        CodexBarWindow {
103            used_percent,
104            window_minutes: None,
105            resets_at: resets_at.map(ToString::to_string),
106            reset_description: None,
107            next_regen_percent: None,
108        }
109    }
110
111    fn response(primary: CodexBarWindow, secondary: CodexBarWindow) -> CodexBarUsageResponse {
112        CodexBarUsageResponse {
113            provider: "codex".to_string(),
114            usage: super::super::types::CodexBarUsage {
115                primary: Some(primary),
116                secondary: Some(secondary),
117                tertiary: None,
118                extra_rate_windows: None,
119            },
120        }
121    }
122
123    #[test]
124    fn not_limited_when_all_windows_under_100() {
125        let resp = response(window(50.0, None), window(30.0, None));
126        assert!(matches!(
127            classify(&resp, Utc::now()),
128            AgentLimit::NotLimited
129        ));
130    }
131
132    #[test]
133    fn limited_when_any_window_at_100() {
134        let resp = response(
135            window(100.0, Some("2099-01-01T00:00:00Z")),
136            window(30.0, None),
137        );
138        match classify(&resp, Utc::now()) {
139            AgentLimit::Limited { reset_time } => {
140                let reset = reset_time.expect("reset present");
141                assert_eq!(
142                    reset,
143                    DateTime::parse_from_rfc3339("2099-01-01T00:00:00Z")
144                        .expect("parse")
145                        .with_timezone(&Utc)
146                );
147            }
148            AgentLimit::NotLimited => panic!("expected limited"),
149        }
150    }
151
152    #[test]
153    fn picks_earliest_reset_across_limited_windows() {
154        let resp = response(
155            window(100.0, Some("2099-06-01T00:00:00Z")),
156            window(100.0, Some("2099-01-01T00:00:00Z")),
157        );
158        match classify(&resp, Utc::now()) {
159            AgentLimit::Limited { reset_time } => {
160                let reset = reset_time.expect("reset present");
161                assert_eq!(
162                    reset,
163                    DateTime::parse_from_rfc3339("2099-01-01T00:00:00Z")
164                        .expect("parse")
165                        .with_timezone(&Utc)
166                );
167            }
168            AgentLimit::NotLimited => panic!("expected limited"),
169        }
170    }
171
172    #[test]
173    fn limited_window_without_resets_at_uses_fallback() {
174        let now = DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
175            .expect("parse")
176            .with_timezone(&Utc);
177        let resp = response(window(100.0, None), window(10.0, None));
178        match classify(&resp, now) {
179            AgentLimit::Limited { reset_time } => {
180                let reset = reset_time.expect("reset present");
181                assert_eq!(reset, now + chrono::Duration::seconds(FALLBACK_RESET_SECS));
182            }
183            AgentLimit::NotLimited => panic!("expected limited"),
184        }
185    }
186
187    #[test]
188    fn not_limited_when_resets_at_already_passed() {
189        // A window at 100% whose resetsAt is in the past is a stale snapshot
190        // (it has presumably already reset server-side), not an active limit.
191        let now = DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
192            .expect("parse")
193            .with_timezone(&Utc);
194        let resp = response(
195            window(100.0, Some("2025-01-01T00:00:00Z")),
196            window(10.0, None),
197        );
198        assert!(matches!(classify(&resp, now), AgentLimit::NotLimited));
199    }
200
201    #[test]
202    fn ignores_stale_reset_but_limits_on_future_reset() {
203        // One window is 100% with a past resetsAt (stale, ignored) and the
204        // other is 100% with a future resetsAt (still active) -- the result
205        // should be Limited, using the future window's reset time.
206        let now = DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z")
207            .expect("parse")
208            .with_timezone(&Utc);
209        let resp = response(
210            window(100.0, Some("2025-01-01T00:00:00Z")),
211            window(100.0, Some("2099-01-01T00:00:00Z")),
212        );
213        match classify(&resp, now) {
214            AgentLimit::Limited { reset_time } => {
215                let reset = reset_time.expect("reset present");
216                assert_eq!(
217                    reset,
218                    DateTime::parse_from_rfc3339("2099-01-01T00:00:00Z")
219                        .expect("parse")
220                        .with_timezone(&Utc)
221                );
222            }
223            AgentLimit::NotLimited => panic!("expected limited"),
224        }
225    }
226
227    #[test]
228    fn counts_extra_rate_windows() {
229        let mut resp = response(window(10.0, None), window(20.0, None));
230        resp.usage.extra_rate_windows = Some(vec![super::super::types::NamedCodexBarWindow {
231            id: "daily".to_string(),
232            title: "Daily".to_string(),
233            window: window(100.0, Some("2099-03-03T00:00:00Z")),
234        }]);
235        assert!(matches!(
236            classify(&resp, Utc::now()),
237            AgentLimit::Limited { .. }
238        ));
239    }
240}