1use super::rate_limit_5h::PRIORITY;
5use std::collections::BTreeMap;
6
7use super::rate_limit_format::{
8 apply_common_extras, format_duration, parse_bool, parse_duration_format, render_error,
9 CommonRateLimitConfig, DurationFormat, ResetWindow,
10};
11use super::{RenderContext, RenderResult, RenderedSegment, Segment, SegmentDefaults};
12use crate::data_context::{DataContext, DataDep, UsageData};
13use crate::theme::Role;
14
15#[non_exhaustive]
16pub struct RateLimit7dResetSegment {
17 pub format: DurationFormat,
18 pub compact: bool,
19 pub use_days: bool,
20 pub config: CommonRateLimitConfig,
21}
22
23impl Default for RateLimit7dResetSegment {
24 fn default() -> Self {
25 Self {
26 format: DurationFormat::Duration,
27 compact: false,
28 use_days: true,
29 config: CommonRateLimitConfig::new("7d reset"),
30 }
31 }
32}
33
34impl RateLimit7dResetSegment {
35 #[must_use]
36 pub fn from_extras(
37 extras: &BTreeMap<String, toml::Value>,
38 warn: &mut impl FnMut(&str),
39 ) -> Self {
40 let mut seg = Self::default();
41 apply_common_extras(&mut seg.config, extras, "rate_limit_7d_reset", warn);
42 if let Some(f) = parse_duration_format(extras, "rate_limit_7d_reset", warn) {
43 seg.format = f;
44 }
45 if let Some(b) = parse_bool(extras, "compact", "rate_limit_7d_reset", warn) {
46 seg.compact = b;
47 }
48 if let Some(b) = parse_bool(extras, "use_days", "rate_limit_7d_reset", warn) {
49 seg.use_days = b;
50 }
51 if seg.config.invalid_progress_width {
52 seg.format = DurationFormat::Duration;
53 }
54 seg
55 }
56}
57
58impl Segment for RateLimit7dResetSegment {
59 fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
60 let usage = ctx.usage();
61 let text = match &*usage {
62 Ok(UsageData::Endpoint(e)) => {
63 let Some(bucket) = e.seven_day.as_ref() else {
64 crate::lsm_debug!(
65 "rate_limit_7d_reset: endpoint usage.seven_day absent; hiding"
66 );
67 return Ok(None);
68 };
69 let Some(resets_at) = bucket.resets_at else {
70 crate::lsm_debug!("rate_limit_7d_reset: seven_day.resets_at absent; hiding");
71 return Ok(None);
72 };
73 let remaining = resets_at.signed_duration_since(chrono::Utc::now());
74 if remaining <= chrono::Duration::zero() {
75 crate::lsm_debug!(
76 "rate_limit_7d_reset: seven_day.resets_at in the past ({resets_at}); hiding"
77 );
78 return Ok(None);
79 }
80 format_duration(
81 remaining,
82 self.format,
83 self.compact,
84 self.use_days,
85 ResetWindow::SevenDay,
86 false,
87 &self.config,
88 )
89 }
90 Ok(UsageData::Jsonl(_)) => {
94 crate::lsm_debug!("rate_limit_7d_reset: jsonl fallback has no hard reset; hiding");
95 return Ok(None);
96 }
97 Err(err) => render_error(err, &self.config),
98 };
99 Ok(Some(RenderedSegment::new(text).with_role(Role::Info)))
100 }
101
102 fn data_deps(&self) -> &'static [DataDep] {
103 &[DataDep::Usage]
104 }
105
106 fn defaults(&self) -> SegmentDefaults {
107 SegmentDefaults::with_priority(PRIORITY)
108 }
109}
110
111#[cfg(test)]
112mod tests {
113 use super::*;
114 use crate::data_context::{
115 EndpointUsage, JsonlUsage, SevenDayWindow, TokenCounts, UsageBucket, UsageData, UsageError,
116 };
117 use crate::input::{ModelInfo, Percent, StatusContext, Tool, WorkspaceInfo};
118 use chrono::Duration;
119 use std::path::PathBuf;
120 use std::sync::Arc;
121
122 fn rc() -> RenderContext {
123 RenderContext::new(80)
124 }
125
126 fn ctx_with_usage(usage: Result<UsageData, UsageError>) -> DataContext {
127 let dc = DataContext::new(StatusContext {
128 tool: Tool::ClaudeCode,
129 model: Some(ModelInfo {
130 display_name: "X".into(),
131 }),
132 workspace: Some(WorkspaceInfo {
133 project_dir: PathBuf::from("/repo"),
134 git_worktree: None,
135 }),
136 context_window: None,
137 cost: None,
138 effort: None,
139 vim: None,
140 output_style: None,
141 agent_name: None,
142 version: None,
143 raw: Arc::new(serde_json::Value::Null),
144 });
145 dc.preseed_usage(usage).expect("seed");
146 dc
147 }
148
149 fn data_with_reset_in(duration: Duration) -> UsageData {
150 let slack = if duration > Duration::zero() {
153 Duration::seconds(30)
154 } else {
155 Duration::zero()
156 };
157 UsageData::Endpoint(EndpointUsage {
158 five_hour: None,
159 seven_day: Some(UsageBucket {
160 utilization: Percent::new(33.0).unwrap(),
161 resets_at: Some(chrono::Utc::now() + duration + slack),
162 }),
163 seven_day_opus: None,
164 seven_day_sonnet: None,
165 seven_day_oauth_apps: None,
166 extra_usage: None,
167 unknown_buckets: std::collections::HashMap::new(),
168 })
169 }
170
171 #[test]
172 fn renders_countdown_with_days_by_default() {
173 let dc = ctx_with_usage(Ok(data_with_reset_in(
174 Duration::days(4) + Duration::hours(8),
175 )));
176 let rendered = RateLimit7dResetSegment::default()
177 .render(&dc, &rc())
178 .unwrap()
179 .expect("visible");
180 assert_eq!(rendered.text(), "7d reset: 4d 8hr");
181 }
182
183 #[test]
184 fn use_days_false_emits_hours_only() {
185 let seg = RateLimit7dResetSegment {
186 use_days: false,
187 ..Default::default()
188 };
189 let dc = ctx_with_usage(Ok(data_with_reset_in(
190 Duration::days(1) + Duration::hours(3),
191 )));
192 let rendered = seg.render(&dc, &rc()).unwrap().expect("visible");
193 assert_eq!(rendered.text(), "7d reset: 27hr");
194 }
195
196 #[test]
197 fn hidden_when_resets_at_in_past() {
198 let dc = ctx_with_usage(Ok(data_with_reset_in(Duration::minutes(-10))));
199 assert_eq!(
200 RateLimit7dResetSegment::default()
201 .render(&dc, &rc())
202 .unwrap(),
203 None,
204 );
205 }
206
207 #[test]
208 fn hidden_when_seven_day_bucket_absent() {
209 let data = UsageData::Endpoint(EndpointUsage {
210 five_hour: None,
211 seven_day: None,
212 seven_day_opus: None,
213 seven_day_sonnet: None,
214 seven_day_oauth_apps: None,
215 extra_usage: None,
216 unknown_buckets: std::collections::HashMap::new(),
217 });
218 assert_eq!(
219 RateLimit7dResetSegment::default()
220 .render(&ctx_with_usage(Ok(data)), &rc())
221 .unwrap(),
222 None,
223 );
224 }
225
226 #[test]
227 fn renders_error_when_usage_fails() {
228 let dc = ctx_with_usage(Err(UsageError::RateLimited { retry_after: None }));
229 let rendered = RateLimit7dResetSegment::default()
230 .render(&dc, &rc())
231 .unwrap()
232 .expect("visible");
233 assert_eq!(rendered.text(), "7d reset: [Rate limited]");
234 }
235
236 #[test]
237 fn hidden_under_jsonl_fallback() {
238 let data = UsageData::Jsonl(JsonlUsage::new(
243 None,
244 SevenDayWindow::new(TokenCounts::default()),
245 ));
246 let dc = ctx_with_usage(Ok(data));
247 assert_eq!(
248 RateLimit7dResetSegment::default()
249 .render(&dc, &rc())
250 .unwrap(),
251 None,
252 );
253 }
254
255 #[test]
256 fn progress_format_divides_by_seven_day_window_not_five_hour() {
257 let dc = ctx_with_usage(Ok(data_with_reset_in(Duration::hours(4))));
262 let seg = RateLimit7dResetSegment {
263 format: DurationFormat::Progress,
264 ..Default::default()
265 };
266 let rendered = seg.render(&dc, &rc()).unwrap().expect("visible");
267 let pct_str = rendered
268 .text()
269 .rsplit(' ')
270 .next()
271 .expect("percent suffix")
272 .trim_end_matches('%');
273 let pct: f64 = pct_str.parse().expect("numeric percent");
274 assert!(
275 (96.0..=98.5).contains(&pct),
276 "expected ~97% elapsed, got {pct}% from {:?}",
277 rendered.text(),
278 );
279 }
280
281 #[test]
282 fn hidden_when_resets_at_missing() {
283 let data = UsageData::Endpoint(EndpointUsage {
284 five_hour: None,
285 seven_day: Some(UsageBucket {
286 utilization: Percent::new(33.0).unwrap(),
287 resets_at: None,
288 }),
289 seven_day_opus: None,
290 seven_day_sonnet: None,
291 seven_day_oauth_apps: None,
292 extra_usage: None,
293 unknown_buckets: std::collections::HashMap::new(),
294 });
295 assert_eq!(
296 RateLimit7dResetSegment::default()
297 .render(&ctx_with_usage(Ok(data)), &rc())
298 .unwrap(),
299 None,
300 );
301 }
302
303 #[test]
304 fn declares_usage_as_its_only_data_dep() {
305 assert_eq!(
306 RateLimit7dResetSegment::default().data_deps(),
307 &[DataDep::Usage],
308 );
309 }
310
311 #[test]
312 fn from_extras_applies_duration_format_knobs() {
313 let mut extras = std::collections::BTreeMap::new();
314 extras.insert("format".into(), toml::Value::String("progress".into()));
315 extras.insert("compact".into(), toml::Value::Boolean(true));
316 let mut warnings = Vec::new();
317 let seg =
318 RateLimit7dResetSegment::from_extras(&extras, &mut |m| warnings.push(m.to_string()));
319 assert!(warnings.is_empty(), "{warnings:?}");
320 assert_eq!(seg.format, DurationFormat::Progress);
321 assert!(seg.compact);
322 }
323}