Skip to main content

linesmith_core/segments/
rate_limit_7d.rs

1//! 7-day rate-limit utilization segment. Mirrors `rate_limit_5h` but
2//! reads `data.seven_day`. Hidden when the bucket is absent (JSONL
3//! fallback always omits it per `rate-limit-segments.md`
4//! §JSONL-fallback display).
5
6use 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            // 7d window is always populated under JSONL (zero-valued
67            // on empty transcripts), so this branch never hides.
68            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        // The 7d window is always populated under JSONL (zero-valued
195        // on empty transcripts per `docs/specs/jsonl-aggregation.md`),
196        // so an empty-transcript user still sees `~7d: 0`, not a hide.
197        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}