1use std::collections::BTreeMap;
13
14use super::config::{
15 apply_common_extras, parse_percent_format, parse_reset_format, CommonRateLimitConfig,
16 PercentFormat, ResetFormat, PRIORITY,
17};
18use super::format::{format_jsonl_tokens, format_percent, format_reset, render_error, ResetWindow};
19use super::window::{resolve_seven_day_reset, UsageWindow, WindowResolution};
20use crate::data_context::{DataContext, DataDep};
21use crate::segments::extras::parse_bool;
22use crate::segments::{RenderContext, RenderResult, RenderedSegment, Segment, SegmentDefaults};
23use crate::theme::Role;
24
25#[non_exhaustive]
26pub struct RateLimit7dSegment {
27 pub format: PercentFormat,
28 pub invert: bool,
29 pub config: CommonRateLimitConfig,
30}
31
32impl Default for RateLimit7dSegment {
33 fn default() -> Self {
34 Self {
35 format: PercentFormat::Percent,
36 invert: false,
37 config: CommonRateLimitConfig::new("7d"),
38 }
39 }
40}
41
42impl RateLimit7dSegment {
43 #[must_use]
44 pub fn from_extras(
45 extras: &BTreeMap<String, toml::Value>,
46 warn: &mut impl FnMut(&str),
47 ) -> Self {
48 let mut seg = Self::default();
49 apply_common_extras(&mut seg.config, extras, "rate_limit_7d", warn);
50 if let Some(f) = parse_percent_format(extras, "rate_limit_7d", warn) {
51 seg.format = f;
52 }
53 if let Some(b) = parse_bool(extras, "invert", "rate_limit_7d", warn) {
54 seg.invert = b;
55 }
56 if seg.config.invalid_progress_width {
57 seg.format = PercentFormat::Percent;
58 }
59 seg
60 }
61}
62
63impl Segment for RateLimit7dSegment {
64 fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
65 let usage = ctx.usage();
66 let text = match &*usage {
67 Ok(data) => match UsageWindow::SevenDay.resolve_percent(data) {
68 Ok(WindowResolution::Endpoint(bucket)) => {
69 format_percent(bucket, self.format, self.invert, &self.config)
70 }
71 Ok(WindowResolution::JsonlTokens(total)) => {
72 format_jsonl_tokens(total, &self.config)
73 }
74 Err(reason) => {
75 crate::lsm_debug!("rate_limit_7d: {reason}; hiding");
76 return Ok(None);
77 }
78 },
79 Err(err) => render_error(err, &self.config),
80 };
81 Ok(Some(RenderedSegment::new(text).with_role(Role::Info)))
82 }
83
84 fn data_deps(&self) -> &'static [DataDep] {
85 &[DataDep::Usage]
86 }
87
88 fn defaults(&self) -> SegmentDefaults {
89 SegmentDefaults::with_priority(PRIORITY)
90 }
91}
92
93#[non_exhaustive]
94pub struct RateLimit7dResetSegment {
95 pub format: ResetFormat,
96 pub compact: bool,
97 pub use_days: bool,
98 pub config: CommonRateLimitConfig,
99}
100
101impl Default for RateLimit7dResetSegment {
102 fn default() -> Self {
103 Self {
104 format: ResetFormat::Duration,
105 compact: false,
106 use_days: true,
107 config: CommonRateLimitConfig::new("7d reset"),
108 }
109 }
110}
111
112impl RateLimit7dResetSegment {
113 #[must_use]
114 pub fn from_extras(
115 extras: &BTreeMap<String, toml::Value>,
116 warn: &mut impl FnMut(&str),
117 ) -> Self {
118 let mut seg = Self::default();
119 apply_common_extras(&mut seg.config, extras, "rate_limit_7d_reset", warn);
120 if let Some(f) = parse_reset_format(extras, "rate_limit_7d_reset", warn) {
121 seg.format = f;
122 }
123 if let Some(b) = parse_bool(extras, "compact", "rate_limit_7d_reset", warn) {
124 seg.compact = b;
125 }
126 if let Some(b) = parse_bool(extras, "use_days", "rate_limit_7d_reset", warn) {
127 seg.use_days = b;
128 }
129 if seg.config.invalid_progress_width && matches!(seg.format, ResetFormat::Progress) {
134 seg.format = ResetFormat::Duration;
135 }
136 seg
137 }
138}
139
140impl Segment for RateLimit7dResetSegment {
141 fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
142 let usage = ctx.usage();
143 let text = match &*usage {
144 Ok(data) => {
145 let resets_at = match resolve_seven_day_reset(data) {
146 Ok(at) => at,
147 Err(reason) => {
148 crate::lsm_debug!("rate_limit_7d_reset: {reason}; hiding");
149 return Ok(None);
150 }
151 };
152 let remaining = resets_at.duration_since(jiff::Timestamp::now());
153 if remaining <= jiff::SignedDuration::ZERO {
154 crate::lsm_debug!(
155 "rate_limit_7d_reset: seven_day.resets_at in the past ({resets_at}); hiding"
156 );
157 return Ok(None);
158 }
159 format_reset(
160 resets_at,
161 remaining,
162 &self.format,
163 self.compact,
164 self.use_days,
165 ResetWindow::SevenDay,
166 false,
167 &self.config,
168 )
169 }
170 Err(err) => render_error(err, &self.config),
171 };
172 Ok(Some(RenderedSegment::new(text).with_role(Role::Info)))
173 }
174
175 fn data_deps(&self) -> &'static [DataDep] {
176 &[DataDep::Usage]
177 }
178
179 fn defaults(&self) -> SegmentDefaults {
180 SegmentDefaults::with_priority(PRIORITY)
181 }
182}
183
184#[cfg(test)]
185mod tests {
186 use super::*;
187 use crate::data_context::{
188 EndpointUsage, JsonlUsage, SevenDayWindow, TokenCounts, UsageBucket, UsageData, UsageError,
189 };
190 use crate::input::{ModelInfo, Percent, StatusContext, Tool, WorkspaceInfo};
191 use jiff::SignedDuration;
192 use std::path::PathBuf;
193 use std::sync::Arc;
194
195 fn rc() -> RenderContext {
196 RenderContext::new(80)
197 }
198
199 fn ctx_with_usage(usage: Result<UsageData, UsageError>) -> DataContext {
200 let dc = DataContext::new(StatusContext {
201 tool: Tool::ClaudeCode,
202 model: Some(ModelInfo {
203 display_name: "X".into(),
204 }),
205 workspace: Some(WorkspaceInfo {
206 project_dir: PathBuf::from("/repo"),
207 git_worktree: None,
208 }),
209 context_window: None,
210 cost: None,
211 effort: None,
212 vim: None,
213 output_style: None,
214 agent_name: None,
215 version: None,
216 raw: Arc::new(serde_json::Value::Null),
217 });
218 dc.preseed_usage(usage).expect("seed");
219 dc
220 }
221
222 fn endpoint_data_with_seven_day(pct: f32) -> UsageData {
223 UsageData::Endpoint(EndpointUsage {
224 five_hour: None,
225 seven_day: Some(UsageBucket {
226 utilization: Percent::new(pct).unwrap(),
227 resets_at: None,
228 }),
229 seven_day_opus: None,
230 seven_day_sonnet: None,
231 seven_day_oauth_apps: None,
232 extra_usage: None,
233 unknown_buckets: std::collections::HashMap::new(),
234 })
235 }
236
237 fn jsonl_data_with_seven_day_tokens(total: u64) -> UsageData {
238 let tokens = TokenCounts::from_parts(total, 0, 0, 0);
239 UsageData::Jsonl(JsonlUsage::new(None, SevenDayWindow::new(tokens)))
240 }
241
242 fn data_with_reset_in(duration: SignedDuration) -> UsageData {
243 let slack = if duration > SignedDuration::ZERO {
246 SignedDuration::from_secs(30)
247 } else {
248 SignedDuration::ZERO
249 };
250 UsageData::Endpoint(EndpointUsage {
251 five_hour: None,
252 seven_day: Some(UsageBucket {
253 utilization: Percent::new(33.0).unwrap(),
254 resets_at: Some(jiff::Timestamp::now() + duration + slack),
255 }),
256 seven_day_opus: None,
257 seven_day_sonnet: None,
258 seven_day_oauth_apps: None,
259 extra_usage: None,
260 unknown_buckets: std::collections::HashMap::new(),
261 })
262 }
263
264 #[test]
265 fn hidden_when_seven_day_bucket_absent() {
266 let data = UsageData::Endpoint(EndpointUsage {
267 five_hour: None,
268 seven_day: None,
269 seven_day_opus: None,
270 seven_day_sonnet: None,
271 seven_day_oauth_apps: None,
272 extra_usage: None,
273 unknown_buckets: std::collections::HashMap::new(),
274 });
275 assert_eq!(
276 RateLimit7dSegment::default()
277 .render(&ctx_with_usage(Ok(data)), &rc())
278 .unwrap(),
279 None,
280 );
281 }
282
283 #[test]
284 fn renders_percent_happy_path() {
285 let dc = ctx_with_usage(Ok(endpoint_data_with_seven_day(33.0)));
286 let rendered = RateLimit7dSegment::default()
287 .render(&dc, &rc())
288 .unwrap()
289 .expect("visible");
290 assert_eq!(rendered.text(), "7d: 33.0%");
291 }
292
293 #[test]
294 fn renders_inverted_percent_when_configured() {
295 let dc = ctx_with_usage(Ok(endpoint_data_with_seven_day(33.0)));
296 let seg = RateLimit7dSegment {
297 invert: true,
298 ..Default::default()
299 };
300 let rendered = seg.render(&dc, &rc()).unwrap().expect("visible");
301 assert_eq!(rendered.text(), "7d: 67.0%");
302 }
303
304 #[test]
305 fn jsonl_mode_renders_compact_tokens_with_stale_marker() {
306 let dc = ctx_with_usage(Ok(jsonl_data_with_seven_day_tokens(1_200_000)));
307 let rendered = RateLimit7dSegment::default()
308 .render(&dc, &rc())
309 .unwrap()
310 .expect("visible");
311 assert_eq!(rendered.text(), "~7d: 1.2M");
312 }
313
314 #[test]
315 fn jsonl_mode_still_renders_on_zero_tokens() {
316 let dc = ctx_with_usage(Ok(jsonl_data_with_seven_day_tokens(0)));
320 let rendered = RateLimit7dSegment::default()
321 .render(&dc, &rc())
322 .unwrap()
323 .expect("visible");
324 assert_eq!(rendered.text(), "~7d: 0");
325 }
326
327 #[test]
328 fn renders_error_when_usage_fails() {
329 let dc = ctx_with_usage(Err(UsageError::Unauthorized));
330 let rendered = RateLimit7dSegment::default()
331 .render(&dc, &rc())
332 .unwrap()
333 .expect("visible");
334 assert_eq!(rendered.text(), "7d: [Unauthorized]");
335 }
336
337 #[test]
338 fn declares_usage_as_its_only_data_dep() {
339 assert_eq!(RateLimit7dSegment::default().data_deps(), &[DataDep::Usage],);
340 }
341
342 #[test]
343 fn from_extras_applies_percent_format_knobs() {
344 let mut extras = std::collections::BTreeMap::new();
345 extras.insert("format".into(), toml::Value::String("progress".into()));
346 extras.insert("invert".into(), toml::Value::Boolean(true));
347 extras.insert("label".into(), toml::Value::String("week".into()));
348 let mut warnings = Vec::new();
349 let seg = RateLimit7dSegment::from_extras(&extras, &mut |m| warnings.push(m.to_string()));
350 assert!(warnings.is_empty(), "{warnings:?}");
351 assert_eq!(seg.format, PercentFormat::Progress);
352 assert!(seg.invert);
353 assert_eq!(seg.config.label, "week");
354 }
355
356 #[test]
359 fn reset_renders_countdown_with_days_by_default() {
360 let dc = ctx_with_usage(Ok(data_with_reset_in(
361 SignedDuration::from_hours(4 * 24) + SignedDuration::from_hours(8),
362 )));
363 let rendered = RateLimit7dResetSegment::default()
364 .render(&dc, &rc())
365 .unwrap()
366 .expect("visible");
367 assert_eq!(rendered.text(), "7d reset: 4d 8hr");
368 }
369
370 #[test]
371 fn reset_use_days_false_emits_hours_only() {
372 let seg = RateLimit7dResetSegment {
373 use_days: false,
374 ..Default::default()
375 };
376 let dc = ctx_with_usage(Ok(data_with_reset_in(
377 SignedDuration::from_hours(24) + SignedDuration::from_hours(3),
378 )));
379 let rendered = seg.render(&dc, &rc()).unwrap().expect("visible");
380 assert_eq!(rendered.text(), "7d reset: 27hr");
381 }
382
383 #[test]
384 fn reset_hidden_when_resets_at_in_past() {
385 let dc = ctx_with_usage(Ok(data_with_reset_in(SignedDuration::from_mins(-10))));
386 assert_eq!(
387 RateLimit7dResetSegment::default()
388 .render(&dc, &rc())
389 .unwrap(),
390 None,
391 );
392 }
393
394 #[test]
395 fn reset_hidden_when_seven_day_bucket_absent() {
396 let data = UsageData::Endpoint(EndpointUsage {
397 five_hour: None,
398 seven_day: None,
399 seven_day_opus: None,
400 seven_day_sonnet: None,
401 seven_day_oauth_apps: None,
402 extra_usage: None,
403 unknown_buckets: std::collections::HashMap::new(),
404 });
405 assert_eq!(
406 RateLimit7dResetSegment::default()
407 .render(&ctx_with_usage(Ok(data)), &rc())
408 .unwrap(),
409 None,
410 );
411 }
412
413 #[test]
414 fn reset_renders_error_when_usage_fails() {
415 let dc = ctx_with_usage(Err(UsageError::RateLimited { retry_after: None }));
416 let rendered = RateLimit7dResetSegment::default()
417 .render(&dc, &rc())
418 .unwrap()
419 .expect("visible");
420 assert_eq!(rendered.text(), "7d reset: [Rate limited]");
421 }
422
423 #[test]
424 fn reset_hidden_under_jsonl_fallback() {
425 let data = UsageData::Jsonl(JsonlUsage::new(
430 None,
431 SevenDayWindow::new(TokenCounts::default()),
432 ));
433 let dc = ctx_with_usage(Ok(data));
434 assert_eq!(
435 RateLimit7dResetSegment::default()
436 .render(&dc, &rc())
437 .unwrap(),
438 None,
439 );
440 }
441
442 #[test]
443 fn reset_progress_format_divides_by_seven_day_window_not_five_hour() {
444 let dc = ctx_with_usage(Ok(data_with_reset_in(SignedDuration::from_hours(4))));
449 let seg = RateLimit7dResetSegment {
450 format: ResetFormat::Progress,
451 ..Default::default()
452 };
453 let rendered = seg.render(&dc, &rc()).unwrap().expect("visible");
454 let pct_str = rendered
455 .text()
456 .rsplit(' ')
457 .next()
458 .expect("percent suffix")
459 .trim_end_matches('%');
460 let pct: f64 = pct_str.parse().expect("numeric percent");
461 assert!(
462 (96.0..=98.5).contains(&pct),
463 "expected ~97% elapsed, got {pct}% from {:?}",
464 rendered.text(),
465 );
466 }
467
468 #[test]
469 fn reset_hidden_when_resets_at_missing() {
470 let data = UsageData::Endpoint(EndpointUsage {
471 five_hour: None,
472 seven_day: Some(UsageBucket {
473 utilization: Percent::new(33.0).unwrap(),
474 resets_at: None,
475 }),
476 seven_day_opus: None,
477 seven_day_sonnet: None,
478 seven_day_oauth_apps: None,
479 extra_usage: None,
480 unknown_buckets: std::collections::HashMap::new(),
481 });
482 assert_eq!(
483 RateLimit7dResetSegment::default()
484 .render(&ctx_with_usage(Ok(data)), &rc())
485 .unwrap(),
486 None,
487 );
488 }
489
490 #[test]
491 fn reset_declares_usage_as_its_only_data_dep() {
492 assert_eq!(
493 RateLimit7dResetSegment::default().data_deps(),
494 &[DataDep::Usage],
495 );
496 }
497
498 #[test]
499 fn reset_from_extras_applies_duration_format_knobs() {
500 let mut extras = std::collections::BTreeMap::new();
501 extras.insert("format".into(), toml::Value::String("progress".into()));
502 extras.insert("compact".into(), toml::Value::Boolean(true));
503 let mut warnings = Vec::new();
504 let seg =
505 RateLimit7dResetSegment::from_extras(&extras, &mut |m| warnings.push(m.to_string()));
506 assert!(warnings.is_empty(), "{warnings:?}");
507 assert_eq!(seg.format, ResetFormat::Progress);
508 assert!(seg.compact);
509 }
510
511 #[test]
512 fn reset_from_extras_warns_on_percent_format_string() {
513 let mut extras = std::collections::BTreeMap::new();
517 extras.insert("format".into(), toml::Value::String("percent".into()));
518 let mut warnings = Vec::new();
519 let _ =
520 RateLimit7dResetSegment::from_extras(&extras, &mut |m| warnings.push(m.to_string()));
521 assert_eq!(warnings.len(), 1);
522 assert!(warnings[0].contains("format"), "{:?}", warnings[0]);
523 }
524
525 #[test]
526 fn reset_invalid_progress_width_does_not_clobber_absolute_format() {
527 let mut extras = std::collections::BTreeMap::new();
530 extras.insert("format".into(), toml::Value::String("absolute".into()));
531 extras.insert("progress_width".into(), toml::Value::Integer(0));
532 let mut warnings = Vec::new();
533 let seg =
534 RateLimit7dResetSegment::from_extras(&extras, &mut |m| warnings.push(m.to_string()));
535 assert!(
536 matches!(seg.format, ResetFormat::Absolute(_)),
537 "absolute survived: {:?}",
538 seg.format
539 );
540 assert!(
541 warnings.iter().any(|w| w.contains("progress_width")),
542 "expected progress_width warning: {warnings:?}"
543 );
544 }
545}