1use super::rate_limit_5h::PRIORITY;
7use std::collections::BTreeMap;
8
9use super::rate_limit_format::{
10 apply_common_extras, format_jsonl_tokens, format_percent, parse_bool, parse_percent_format,
11 render_error, CommonRateLimitConfig, PercentFormat,
12};
13use super::{RenderContext, RenderResult, RenderedSegment, Segment, SegmentDefaults};
14use crate::data_context::{DataContext, DataDep, UsageData};
15use crate::theme::Role;
16
17#[non_exhaustive]
18pub struct RateLimit7dSegment {
19 pub format: PercentFormat,
20 pub invert: bool,
21 pub config: CommonRateLimitConfig,
22}
23
24impl Default for RateLimit7dSegment {
25 fn default() -> Self {
26 Self {
27 format: PercentFormat::Percent,
28 invert: false,
29 config: CommonRateLimitConfig::new("7d"),
30 }
31 }
32}
33
34impl RateLimit7dSegment {
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", warn);
42 if let Some(f) = parse_percent_format(extras, "rate_limit_7d", warn) {
43 seg.format = f;
44 }
45 if let Some(b) = parse_bool(extras, "invert", "rate_limit_7d", warn) {
46 seg.invert = b;
47 }
48 if seg.config.invalid_progress_width {
49 seg.format = PercentFormat::Percent;
50 }
51 seg
52 }
53}
54
55impl Segment for RateLimit7dSegment {
56 fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
57 let usage = ctx.usage();
58 let text = match &*usage {
59 Ok(UsageData::Endpoint(e)) => match &e.seven_day {
60 Some(bucket) => format_percent(bucket, self.format, self.invert, &self.config),
61 None => {
62 crate::lsm_debug!("rate_limit_7d: endpoint usage.seven_day absent; hiding");
63 return Ok(None);
64 }
65 },
66 Ok(UsageData::Jsonl(j)) => {
69 format_jsonl_tokens(j.seven_day.tokens.total(), &self.config)
70 }
71 Err(err) => render_error(err, &self.config),
72 };
73 Ok(Some(RenderedSegment::new(text).with_role(Role::Info)))
74 }
75
76 fn data_deps(&self) -> &'static [DataDep] {
77 &[DataDep::Usage]
78 }
79
80 fn defaults(&self) -> SegmentDefaults {
81 SegmentDefaults::with_priority(PRIORITY)
82 }
83}
84
85#[cfg(test)]
86mod tests {
87 use super::*;
88 use crate::data_context::{
89 EndpointUsage, JsonlUsage, SevenDayWindow, TokenCounts, UsageBucket, UsageData, UsageError,
90 };
91 use crate::input::{ModelInfo, Percent, StatusContext, Tool, WorkspaceInfo};
92 use std::path::PathBuf;
93 use std::sync::Arc;
94
95 fn rc() -> RenderContext {
96 RenderContext::new(80)
97 }
98
99 fn ctx_with_usage(usage: Result<UsageData, UsageError>) -> DataContext {
100 let dc = DataContext::new(StatusContext {
101 tool: Tool::ClaudeCode,
102 model: Some(ModelInfo {
103 display_name: "X".into(),
104 }),
105 workspace: Some(WorkspaceInfo {
106 project_dir: PathBuf::from("/repo"),
107 git_worktree: None,
108 }),
109 context_window: None,
110 cost: None,
111 effort: None,
112 vim: None,
113 output_style: None,
114 agent_name: None,
115 version: None,
116 raw: Arc::new(serde_json::Value::Null),
117 });
118 dc.preseed_usage(usage).expect("seed");
119 dc
120 }
121
122 fn endpoint_data_with_seven_day(pct: f32) -> UsageData {
123 UsageData::Endpoint(EndpointUsage {
124 five_hour: None,
125 seven_day: Some(UsageBucket {
126 utilization: Percent::new(pct).unwrap(),
127 resets_at: None,
128 }),
129 seven_day_opus: None,
130 seven_day_sonnet: None,
131 seven_day_oauth_apps: None,
132 extra_usage: None,
133 unknown_buckets: std::collections::HashMap::new(),
134 })
135 }
136
137 fn jsonl_data_with_seven_day_tokens(total: u64) -> UsageData {
138 let tokens = TokenCounts::from_parts(total, 0, 0, 0);
139 UsageData::Jsonl(JsonlUsage::new(None, SevenDayWindow::new(tokens)))
140 }
141
142 #[test]
143 fn hidden_when_seven_day_bucket_absent() {
144 let data = UsageData::Endpoint(EndpointUsage {
145 five_hour: None,
146 seven_day: None,
147 seven_day_opus: None,
148 seven_day_sonnet: None,
149 seven_day_oauth_apps: None,
150 extra_usage: None,
151 unknown_buckets: std::collections::HashMap::new(),
152 });
153 assert_eq!(
154 RateLimit7dSegment::default()
155 .render(&ctx_with_usage(Ok(data)), &rc())
156 .unwrap(),
157 None,
158 );
159 }
160
161 #[test]
162 fn renders_percent_happy_path() {
163 let dc = ctx_with_usage(Ok(endpoint_data_with_seven_day(33.0)));
164 let rendered = RateLimit7dSegment::default()
165 .render(&dc, &rc())
166 .unwrap()
167 .expect("visible");
168 assert_eq!(rendered.text(), "7d: 33.0%");
169 }
170
171 #[test]
172 fn renders_inverted_percent_when_configured() {
173 let dc = ctx_with_usage(Ok(endpoint_data_with_seven_day(33.0)));
174 let seg = RateLimit7dSegment {
175 invert: true,
176 ..Default::default()
177 };
178 let rendered = seg.render(&dc, &rc()).unwrap().expect("visible");
179 assert_eq!(rendered.text(), "7d: 67.0%");
180 }
181
182 #[test]
183 fn jsonl_mode_renders_compact_tokens_with_stale_marker() {
184 let dc = ctx_with_usage(Ok(jsonl_data_with_seven_day_tokens(1_200_000)));
185 let rendered = RateLimit7dSegment::default()
186 .render(&dc, &rc())
187 .unwrap()
188 .expect("visible");
189 assert_eq!(rendered.text(), "~7d: 1.2M");
190 }
191
192 #[test]
193 fn jsonl_mode_still_renders_on_zero_tokens() {
194 let dc = ctx_with_usage(Ok(jsonl_data_with_seven_day_tokens(0)));
198 let rendered = RateLimit7dSegment::default()
199 .render(&dc, &rc())
200 .unwrap()
201 .expect("visible");
202 assert_eq!(rendered.text(), "~7d: 0");
203 }
204
205 #[test]
206 fn renders_error_when_usage_fails() {
207 let dc = ctx_with_usage(Err(UsageError::Unauthorized));
208 let rendered = RateLimit7dSegment::default()
209 .render(&dc, &rc())
210 .unwrap()
211 .expect("visible");
212 assert_eq!(rendered.text(), "7d: [Unauthorized]");
213 }
214
215 #[test]
216 fn declares_usage_as_its_only_data_dep() {
217 assert_eq!(RateLimit7dSegment::default().data_deps(), &[DataDep::Usage],);
218 }
219
220 #[test]
221 fn from_extras_applies_percent_format_knobs() {
222 let mut extras = std::collections::BTreeMap::new();
223 extras.insert("format".into(), toml::Value::String("progress".into()));
224 extras.insert("invert".into(), toml::Value::Boolean(true));
225 extras.insert("label".into(), toml::Value::String("week".into()));
226 let mut warnings = Vec::new();
227 let seg = RateLimit7dSegment::from_extras(&extras, &mut |m| warnings.push(m.to_string()));
228 assert!(warnings.is_empty(), "{warnings:?}");
229 assert_eq!(seg.format, PercentFormat::Progress);
230 assert!(seg.invert);
231 assert_eq!(seg.config.label, "week");
232 }
233}