1use chrono::{DateTime, Utc};
12
13use super::client::{RunCodexBarUsageOptions, run_codexbar_usage};
14use super::errors::CodexBarError;
15use super::types::{CodexBarUsageResponse, CodexBarWindow};
16
17#[derive(Debug, Clone)]
19pub enum AgentLimit {
20 NotLimited,
22 Limited { reset_time: Option<DateTime<Utc>> },
25}
26
27const 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 .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
74pub async fn check_limit(provider: &str) -> Result<AgentLimit, CodexBarError> {
80 check_limit_with(provider, &RunCodexBarUsageOptions::default()).await
81}
82
83pub 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 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 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}