linesmith_core/segments/
tokens.rs

1//! Per-turn token segments: `tokens_input`, `tokens_output`,
2//! `tokens_cached`, `tokens_total`. Each reads the corresponding field
3//! from `ctx.status.context_window.current_usage` (the per-turn
4//! breakdown from the most recent API call, see `TurnUsage`). All hide
5//! when `current_usage` is None — either before the first API call in
6//! a session, or when the schema didn't supply the key.
7//!
8//! Rather than one `tokens` segment with show/hide config toggles, we
9//! split into four segments so users compose via `[line.segments]` the
10//! same way they order other built-ins. The split matches ccstatusline
11//! (see `docs/research/jsonl-data-source.md` §7 "ccstatusline widget
12//! catalog").
13
14use super::rate_limit::format::format_tokens;
15use super::{RenderContext, RenderResult, RenderedSegment, Segment, SegmentDefaults};
16use crate::data_context::DataContext;
17use crate::input::TurnUsage;
18use crate::theme::Role;
19
20/// Same bucket as `cost` (192): per-turn token counts are useful
21/// context but not time-critical, so they yield first under width
22/// pressure.
23const PRIORITY: u8 = 192;
24
25pub struct TokensInputSegment;
26pub struct TokensOutputSegment;
27pub struct TokensCachedSegment;
28pub struct TokensTotalSegment;
29
30impl Segment for TokensInputSegment {
31    fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
32        let Some(usage) = current_usage(ctx) else {
33            crate::lsm_debug!("tokens_input: current_usage absent; hiding");
34            return Ok(None);
35        };
36        Ok(Some(render("in", usage.input_tokens)))
37    }
38    fn defaults(&self) -> SegmentDefaults {
39        SegmentDefaults::with_priority(PRIORITY)
40    }
41}
42
43impl Segment for TokensOutputSegment {
44    fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
45        let Some(usage) = current_usage(ctx) else {
46            crate::lsm_debug!("tokens_output: current_usage absent; hiding");
47            return Ok(None);
48        };
49        Ok(Some(render("out", usage.output_tokens)))
50    }
51    fn defaults(&self) -> SegmentDefaults {
52        SegmentDefaults::with_priority(PRIORITY)
53    }
54}
55
56impl Segment for TokensCachedSegment {
57    fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
58        let Some(usage) = current_usage(ctx) else {
59            crate::lsm_debug!("tokens_cached: current_usage absent; hiding");
60            return Ok(None);
61        };
62        // Cache-related tokens for the turn: creation + read. Both are
63        // input-side, so a single "cache" number tracks "how much did
64        // caching do for us this turn" without the user having to sum
65        // two fields themselves.
66        let cached = usage
67            .cache_creation_input_tokens
68            .saturating_add(usage.cache_read_input_tokens);
69        Ok(Some(render("cache", cached)))
70    }
71    fn defaults(&self) -> SegmentDefaults {
72        SegmentDefaults::with_priority(PRIORITY)
73    }
74}
75
76impl Segment for TokensTotalSegment {
77    fn render(&self, ctx: &DataContext, _rc: &RenderContext) -> RenderResult {
78        let Some(usage) = current_usage(ctx) else {
79            crate::lsm_debug!("tokens_total: current_usage absent; hiding");
80            return Ok(None);
81        };
82        let total = usage
83            .input_tokens
84            .saturating_add(usage.output_tokens)
85            .saturating_add(usage.cache_creation_input_tokens)
86            .saturating_add(usage.cache_read_input_tokens);
87        Ok(Some(render("total", total)))
88    }
89    fn defaults(&self) -> SegmentDefaults {
90        SegmentDefaults::with_priority(PRIORITY)
91    }
92}
93
94fn current_usage(ctx: &DataContext) -> Option<&TurnUsage> {
95    ctx.status.context_window.as_ref()?.current_usage.as_ref()
96}
97
98fn render(label: &'static str, count: u64) -> RenderedSegment {
99    // Reuse the project-wide compact formatter so `1_900` renders as
100    // `1.9k` (not `1k`) and a session's per-turn count stays consistent
101    // with the rate-limit / JSONL token formatting elsewhere.
102    RenderedSegment::new(format!("{label} {}", format_tokens(count))).with_role(Role::Muted)
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use crate::input::{
109        ContextWindow, ModelInfo, Percent, StatusContext, Tool, TurnUsage, WorkspaceInfo,
110    };
111    use std::path::PathBuf;
112    use std::sync::Arc;
113
114    fn rc() -> RenderContext {
115        RenderContext::new(80)
116    }
117
118    fn ctx(usage: Option<TurnUsage>) -> DataContext {
119        DataContext::new(StatusContext {
120            tool: Tool::ClaudeCode,
121            model: Some(ModelInfo {
122                display_name: "X".into(),
123            }),
124            workspace: Some(WorkspaceInfo {
125                project_dir: PathBuf::from("/repo"),
126                git_worktree: None,
127            }),
128            context_window: Some(ContextWindow {
129                used: Some(Percent::new(0.0).unwrap()),
130                size: Some(200_000),
131                total_input_tokens: Some(0),
132                total_output_tokens: Some(0),
133                current_usage: usage,
134            }),
135            cost: None,
136            effort: None,
137            vim: None,
138            output_style: None,
139            agent_name: None,
140            version: None,
141            raw: Arc::new(serde_json::Value::Null),
142        })
143    }
144
145    fn usage(input: u64, output: u64, cache_creation: u64, cache_read: u64) -> TurnUsage {
146        TurnUsage {
147            input_tokens: input,
148            output_tokens: output,
149            cache_creation_input_tokens: cache_creation,
150            cache_read_input_tokens: cache_read,
151        }
152    }
153
154    fn ctx_without_context_window() -> DataContext {
155        DataContext::new(StatusContext {
156            tool: Tool::ClaudeCode,
157            model: Some(ModelInfo {
158                display_name: "X".into(),
159            }),
160            workspace: Some(WorkspaceInfo {
161                project_dir: PathBuf::from("/repo"),
162                git_worktree: None,
163            }),
164            context_window: None,
165            cost: None,
166            effort: None,
167            vim: None,
168            output_style: None,
169            agent_name: None,
170            version: None,
171            raw: Arc::new(serde_json::Value::Null),
172        })
173    }
174
175    // format_tokens itself is covered in `rate_limit::format` tests;
176    // here we only assert the segment-level integration.
177
178    // --- tokens_input ---
179
180    #[test]
181    fn tokens_input_renders_current_usage_input_with_muted_role() {
182        assert_eq!(
183            TokensInputSegment
184                .render(&ctx(Some(usage(2_000, 500, 0, 500))), &rc())
185                .unwrap(),
186            Some(RenderedSegment::new("in 2.0k").with_role(Role::Muted))
187        );
188    }
189
190    #[test]
191    fn tokens_input_hidden_when_current_usage_absent() {
192        assert_eq!(TokensInputSegment.render(&ctx(None), &rc()).unwrap(), None);
193    }
194
195    #[test]
196    fn tokens_input_hidden_when_context_window_absent() {
197        // Double-None path: no context_window at all short-circuits
198        // before we even reach current_usage.
199        assert_eq!(
200            TokensInputSegment
201                .render(&ctx_without_context_window(), &rc())
202                .unwrap(),
203            None
204        );
205    }
206
207    // --- tokens_output ---
208
209    #[test]
210    fn tokens_output_renders_current_usage_output_with_muted_role() {
211        assert_eq!(
212            TokensOutputSegment
213                .render(&ctx(Some(usage(2_000, 500, 0, 500))), &rc())
214                .unwrap(),
215            Some(RenderedSegment::new("out 500").with_role(Role::Muted))
216        );
217    }
218
219    #[test]
220    fn tokens_output_hidden_when_current_usage_absent() {
221        assert_eq!(TokensOutputSegment.render(&ctx(None), &rc()).unwrap(), None);
222    }
223
224    #[test]
225    fn tokens_output_hidden_when_context_window_absent() {
226        assert_eq!(
227            TokensOutputSegment
228                .render(&ctx_without_context_window(), &rc())
229                .unwrap(),
230            None
231        );
232    }
233
234    // --- tokens_cached ---
235
236    #[test]
237    fn tokens_cached_sums_cache_creation_and_read_with_muted_role() {
238        assert_eq!(
239            TokensCachedSegment
240                .render(&ctx(Some(usage(2_000, 500, 300, 700))), &rc())
241                .unwrap(),
242            Some(RenderedSegment::new("cache 1.0k").with_role(Role::Muted))
243        );
244    }
245
246    #[test]
247    fn tokens_cached_renders_zero_when_both_fields_zero() {
248        assert_eq!(
249            TokensCachedSegment
250                .render(&ctx(Some(usage(2_000, 500, 0, 0))), &rc())
251                .unwrap(),
252            Some(RenderedSegment::new("cache 0").with_role(Role::Muted))
253        );
254    }
255
256    #[test]
257    fn tokens_cached_hidden_when_current_usage_absent() {
258        assert_eq!(TokensCachedSegment.render(&ctx(None), &rc()).unwrap(), None);
259    }
260
261    #[test]
262    fn tokens_cached_hidden_when_context_window_absent() {
263        assert_eq!(
264            TokensCachedSegment
265                .render(&ctx_without_context_window(), &rc())
266                .unwrap(),
267            None
268        );
269    }
270
271    // --- tokens_total ---
272
273    #[test]
274    fn tokens_total_sums_all_four_fields() {
275        assert_eq!(
276            TokensTotalSegment
277                .render(&ctx(Some(usage(2_000, 500, 300, 700))), &rc())
278                .unwrap(),
279            Some(RenderedSegment::new("total 3.5k").with_role(Role::Muted))
280        );
281    }
282
283    #[test]
284    fn tokens_total_saturates_on_overflow() {
285        // Four u64::MAX values can't co-exist in a real payload, but
286        // the saturating chain must not panic on the adversarial input.
287        assert_eq!(
288            TokensTotalSegment
289                .render(
290                    &ctx(Some(usage(u64::MAX, u64::MAX, u64::MAX, u64::MAX))),
291                    &rc()
292                )
293                .unwrap(),
294            Some(
295                RenderedSegment::new(format!("total {}", format_tokens(u64::MAX)))
296                    .with_role(Role::Muted)
297            )
298        );
299    }
300
301    #[test]
302    fn tokens_total_hidden_when_current_usage_absent() {
303        assert_eq!(TokensTotalSegment.render(&ctx(None), &rc()).unwrap(), None);
304    }
305
306    #[test]
307    fn tokens_total_hidden_when_context_window_absent() {
308        assert_eq!(
309            TokensTotalSegment
310                .render(&ctx_without_context_window(), &rc())
311                .unwrap(),
312            None
313        );
314    }
315
316    // --- defaults ---
317
318    #[test]
319    fn all_four_segments_share_priority() {
320        assert_eq!(TokensInputSegment.defaults().priority, PRIORITY);
321        assert_eq!(TokensOutputSegment.defaults().priority, PRIORITY);
322        assert_eq!(TokensCachedSegment.defaults().priority, PRIORITY);
323        assert_eq!(TokensTotalSegment.defaults().priority, PRIORITY);
324    }
325}
linesmith_core/segments/tokens.rs

linesmith_core/segments/
tokens.rs
