返回 CodeWhale
resource_telemetry.rs
根目录 / crates / tui / src / resource_telemetry.rs
1 #![allow(dead_code)]
2
3 //! Resource-usage telemetry for long-running CodeWhale tasks.
4 //!
5 //! This module is a pure, side-effect-free foundation for surfacing how many
6 //! tokens and how much wall-clock time a task has consumed, optionally relative
7 //! to a budget. It performs no I/O and no rendering; consumers (status lines,
8 //! the cost panel, the goal/budget tooling) are wired up separately so the
9 //! formatting and pressure logic can be unit-tested in isolation.
10 //!
11 //! The shape intentionally mirrors the budget vocabulary already used by the
12 //! goal tooling (`token_budget: Option<_>`) so a consumer can adapt between the
13 //! two without inventing new concepts. We keep a local type rather than reusing
14 //! `tools::goal` here to avoid coupling a presentation-layer helper to the tool
15 //! domain model (whose budgets are `u32` and carry unrelated bookkeeping).
16
17 use std::{
18 fmt::{self, Write as _},
19 time::Duration,
20 };
21
22 /// A coarse, three-level read on how close a task is to exhausting its budget.
23 ///
24 /// The level is derived from the *highest* pressure across all bounded
25 /// dimensions (tokens and time), so a task that is comfortable on tokens but
26 /// nearly out of time still reports [`PressureLevel::High`]. When nothing is
27 /// bounded, pressure is [`PressureLevel::Low`] by definition.
28 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
29 pub enum PressureLevel {
30 /// Plenty of headroom (under ~75% of every bounded budget).
31 Low,
32 /// Getting close (at/over ~75% but under 100% of some budget).
33 Medium,
34 /// At or over budget on some bounded dimension.
35 High,
36 }
37
38 impl PressureLevel {
39 /// Fraction at/above which a dimension is considered medium pressure.
40 const MEDIUM_THRESHOLD: f64 = 0.75;
41 /// Fraction at/above which a dimension is considered high pressure.
42 const HIGH_THRESHOLD: f64 = 1.0;
43
44 /// Classify a single budget fraction (e.g. `0.41` for 41% used).
45 ///
46 /// Negative or non-finite input is treated as [`PressureLevel::Low`]; the
47 /// telemetry helpers never produce such values, but classifying defensively
48 /// keeps this usable for arbitrary callers.
49 fn from_fraction(fraction: f64) -> Self {
50 if !fraction.is_finite() || fraction < Self::MEDIUM_THRESHOLD {
51 PressureLevel::Low
52 } else if fraction < Self::HIGH_THRESHOLD {
53 PressureLevel::Medium
54 } else {
55 PressureLevel::High
56 }
57 }
58
59 /// A short lowercase label suitable for compact status output.
60 pub fn label(self) -> &'static str {
61 match self {
62 PressureLevel::Low => "low",
63 PressureLevel::Medium => "medium",
64 PressureLevel::High => "high",
65 }
66 }
67 }
68
69 /// A snapshot of token and time usage for a single task, with optional budgets.
70 ///
71 /// All fields are plain counters; this type owns no clock and reads no
72 /// environment. Construct it from whatever the caller is already tracking and
73 /// use the helpers below to render or classify it.
74 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
75 pub struct ResourceTelemetry {
76 /// Total tokens consumed so far.
77 pub tokens_used: u64,
78 /// Total wall-clock seconds elapsed so far.
79 pub time_used_seconds: u64,
80 /// Optional token ceiling for the task; `None` means unbounded.
81 pub token_budget: Option<u64>,
82 /// Optional time ceiling in seconds; `None` means unbounded.
83 pub time_budget_seconds: Option<u64>,
84 }
85
86 impl ResourceTelemetry {
87 /// Create a telemetry snapshot with no budgets (fully unbounded).
88 pub fn new(tokens_used: u64, time_used_seconds: u64) -> Self {
89 Self {
90 tokens_used,
91 time_used_seconds,
92 token_budget: None,
93 time_budget_seconds: None,
94 }
95 }
96
97 /// Set the token budget, returning the updated snapshot (builder style).
98 pub fn with_token_budget(mut self, budget: u64) -> Self {
99 self.token_budget = Some(budget);
100 self
101 }
102
103 /// Set the time budget in seconds, returning the updated snapshot.
104 pub fn with_time_budget_seconds(mut self, seconds: u64) -> Self {
105 self.time_budget_seconds = Some(seconds);
106 self
107 }
108
109 /// Fraction of the token budget consumed, or `None` when unbounded.
110 ///
111 /// A zero budget yields `None` (a percentage of nothing is meaningless)
112 /// rather than infinity, keeping every downstream consumer safe.
113 pub fn token_fraction(&self) -> Option<f64> {
114 fraction(self.tokens_used, self.token_budget)
115 }
116
117 /// Fraction of the time budget consumed, or `None` when unbounded.
118 pub fn time_fraction(&self) -> Option<f64> {
119 fraction(self.time_used_seconds, self.time_budget_seconds)
120 }
121
122 /// The largest bounded budget fraction across tokens and time.
123 ///
124 /// Returns `None` only when *neither* dimension is bounded. When at least
125 /// one budget is present, the most-pressured bounded dimension wins.
126 pub fn budget_fraction(&self) -> Option<f64> {
127 match (self.token_fraction(), self.time_fraction()) {
128 (Some(t), Some(s)) => Some(t.max(s)),
129 (Some(t), None) => Some(t),
130 (None, Some(s)) => Some(s),
131 (None, None) => None,
132 }
133 }
134
135 /// Budget fraction expressed as a whole-number percent (rounded), or `None`
136 /// when unbounded. This is the value surfaced in the human summary.
137 pub fn budget_percent(&self) -> Option<u64> {
138 self.budget_fraction().map(|f| (f * 100.0).round() as u64)
139 }
140
141 /// Coarse pressure level derived from [`Self::budget_fraction`].
142 ///
143 /// Unbounded tasks are always [`PressureLevel::Low`].
144 pub fn pressure(&self) -> PressureLevel {
145 match self.budget_fraction() {
146 Some(fraction) => PressureLevel::from_fraction(fraction),
147 None => PressureLevel::Low,
148 }
149 }
150
151 /// A compact, human-readable one-liner, e.g. `12.3k tok · 4m12s · 41% budget`.
152 ///
153 /// Tokens are abbreviated with `k`/`M` suffixes, time is rendered as
154 /// `Hh Mm Ss` (dropping leading zero units), and the budget segment is
155 /// omitted entirely when the task is unbounded.
156 pub fn human_summary(&self) -> String {
157 let mut out = String::new();
158 // `write!` into a String is infallible; ignore the Result.
159 let _ = write!(
160 out,
161 "{} tok · {}",
162 format_tokens(self.tokens_used),
163 crate::elapsed::format_elapsed_secs(self.time_used_seconds),
164 );
165 if let Some(percent) = self.budget_percent() {
166 let _ = write!(out, " · {percent}% budget");
167 }
168 out
169 }
170 }
171
172 impl fmt::Display for ResourceTelemetry {
173 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174 f.write_str(&self.human_summary())
175 }
176 }
177
178 /// Output-token throughput for a live or completed turn.
179 #[derive(Debug, Clone, Copy, PartialEq)]
180 pub struct TokenThroughput {
181 pub output_tokens: u64,
182 pub elapsed_seconds: f64,
183 }
184
185 impl TokenThroughput {
186 pub fn new(output_tokens: u64, elapsed: Duration) -> Option<Self> {
187 let elapsed_seconds = elapsed.as_secs_f64();
188 if output_tokens == 0 || !elapsed_seconds.is_finite() || elapsed_seconds <= 0.0 {
189 return None;
190 }
191 Some(Self {
192 output_tokens,
193 elapsed_seconds,
194 })
195 }
196
197 pub fn from_estimated_text(text: &str, elapsed: Duration) -> Option<Self> {
198 Self::new(estimate_output_tokens_from_text(text), elapsed)
199 }
200
201 pub fn tokens_per_second(self) -> f64 {
202 self.output_tokens as f64 / self.elapsed_seconds
203 }
204
205 pub fn compact_rate(self) -> String {
206 let rate = self.tokens_per_second();
207 if rate < 10.0 {
208 format!("{rate:.1}")
209 } else {
210 format!("{rate:.0}")
211 }
212 }
213 }
214
215 /// Estimate output tokens from streamed text before provider usage arrives.
216 ///
217 /// Provider-reported usage remains canonical at turn completion. During a live
218 /// stream, this gives the footer a stable approximation without inspecting
219 /// provider-specific tokenizer internals.
220 pub fn estimate_output_tokens_from_text(text: &str) -> u64 {
221 let chars = text.chars().count() as u64;
222 if chars == 0 {
223 0
224 } else {
225 chars.saturating_add(3) / 4
226 }
227 }
228
229 /// Divide `used` by an optional budget, guarding against an absent or zero
230 /// budget. Returns `None` when the budget is `None` or `0`.
231 fn fraction(used: u64, budget: Option<u64>) -> Option<f64> {
232 match budget {
233 Some(budget) if budget > 0 => Some(used as f64 / budget as f64),
234 _ => None,
235 }
236 }
237
238 /// Format a token count with a `k`/`M` suffix once it crosses each threshold.
239 ///
240 /// Values under 1_000 are printed verbatim. Thousands use one decimal place
241 /// (`12.3k`), trimming a trailing `.0` so round values read cleanly (`5k`).
242 /// Millions follow the same rule (`1.5M`, `2M`).
243 fn format_tokens(tokens: u64) -> String {
244 const K: u64 = 1_000;
245 const M: u64 = 1_000_000;
246 if tokens >= M {
247 format_scaled(tokens, M, 'M')
248 } else if tokens >= K {
249 format_scaled(tokens, K, 'k')
250 } else {
251 tokens.to_string()
252 }
253 }
254
255 /// Render `value / divisor` to one decimal place with `suffix`, dropping a
256 /// trailing `.0`. The divisor is always one of the constants above (non-zero).
257 fn format_scaled(value: u64, divisor: u64, suffix: char) -> String {
258 let scaled = value as f64 / divisor as f64;
259 // Round to one decimal before deciding whether the fraction is ".0", so a
260 // value like 1_999_999 reads as "2M" rather than "1.9...M".
261 let rounded = (scaled * 10.0).round() / 10.0;
262 if (rounded.fract()).abs() < f64::EPSILON {
263 format!("{}{}", rounded as u64, suffix)
264 } else {
265 format!("{rounded:.1}{suffix}")
266 }
267 }
268
269 #[cfg(test)]
270 mod tests {
271 use super::*;
272
273 // ---- token formatting -------------------------------------------------
274
275 #[test]
276 fn format_tokens_under_a_thousand_is_verbatim() {
277 assert_eq!(format_tokens(0), "0");
278 assert_eq!(format_tokens(1), "1");
279 assert_eq!(format_tokens(999), "999");
280 }
281
282 #[test]
283 fn format_tokens_uses_k_suffix_with_trimmed_decimal() {
284 assert_eq!(format_tokens(1_000), "1k");
285 assert_eq!(format_tokens(1_500), "1.5k");
286 assert_eq!(format_tokens(12_345), "12.3k");
287 // Exactly on a round thousand trims the ".0".
288 assert_eq!(format_tokens(5_000), "5k");
289 // Just under the millions boundary stays in k.
290 assert_eq!(format_tokens(999_400), "999.4k");
291 }
292
293 #[test]
294 fn format_tokens_uses_m_suffix_for_millions() {
295 assert_eq!(format_tokens(1_000_000), "1M");
296 assert_eq!(format_tokens(1_500_000), "1.5M");
297 assert_eq!(format_tokens(2_340_000), "2.3M");
298 }
299
300 #[test]
301 fn format_tokens_rounds_up_across_a_unit_boundary() {
302 // 1_999_999 rounds to 2.0M -> "2M", not "1.9M" or "2.0M".
303 assert_eq!(format_tokens(1_999_999), "2M");
304 // 999_950 rounds to 1000.0k; still within the k branch and trims ".0".
305 assert_eq!(format_tokens(999_950), "1000k");
306 }
307
308 #[test]
309 fn format_tokens_handles_very_large_values() {
310 assert_eq!(format_tokens(u64::MAX), "18446744073709.6M");
311 }
312
313 // ---- throughput -------------------------------------------------------
314
315 #[test]
316 fn token_throughput_formats_compact_rates() {
317 let throughput = TokenThroughput::new(120, Duration::from_secs(6)).expect("throughput");
318 assert_eq!(throughput.tokens_per_second(), 20.0);
319 assert_eq!(throughput.compact_rate(), "20");
320
321 let slow = TokenThroughput::new(15, Duration::from_secs(4)).expect("throughput");
322 assert_eq!(slow.compact_rate(), "3.8");
323 }
324
325 #[test]
326 fn token_throughput_rejects_empty_or_zero_elapsed_samples() {
327 assert!(TokenThroughput::new(0, Duration::from_secs(5)).is_none());
328 assert!(TokenThroughput::new(5, Duration::ZERO).is_none());
329 }
330
331 #[test]
332 fn estimated_streaming_tokens_round_up_from_text_chars() {
333 assert_eq!(estimate_output_tokens_from_text(""), 0);
334 assert_eq!(estimate_output_tokens_from_text("abc"), 1);
335 assert_eq!(estimate_output_tokens_from_text("abcd"), 1);
336 assert_eq!(estimate_output_tokens_from_text("abcde"), 2);
337
338 let throughput =
339 TokenThroughput::from_estimated_text(&"x".repeat(400), Duration::from_secs(10))
340 .expect("estimated throughput");
341 assert_eq!(throughput.output_tokens, 100);
342 assert_eq!(throughput.compact_rate(), "10");
343 }
344
345 // ---- fraction / percent ----------------------------------------------
346
347 #[test]
348 fn fractions_are_none_when_unbounded() {
349 let t = ResourceTelemetry::new(5_000, 120);
350 assert_eq!(t.token_fraction(), None);
351 assert_eq!(t.time_fraction(), None);
352 assert_eq!(t.budget_fraction(), None);
353 assert_eq!(t.budget_percent(), None);
354 }
355
356 #[test]
357 fn zero_budget_yields_none_not_infinity() {
358 let t = ResourceTelemetry {
359 tokens_used: 100,
360 time_used_seconds: 0,
361 token_budget: Some(0),
362 time_budget_seconds: Some(0),
363 };
364 assert_eq!(t.token_fraction(), None);
365 assert_eq!(t.time_fraction(), None);
366 assert_eq!(t.budget_fraction(), None);
367 assert_eq!(t.pressure(), PressureLevel::Low);
368 }
369
370 #[test]
371 fn token_fraction_is_computed_when_bounded() {
372 let t = ResourceTelemetry::new(4_100, 0).with_token_budget(10_000);
373 let frac = t.token_fraction().expect("bounded");
374 assert!((frac - 0.41).abs() < 1e-9, "got {frac}");
375 assert_eq!(t.budget_percent(), Some(41));
376 }
377
378 #[test]
379 fn budget_fraction_takes_the_max_across_dimensions() {
380 // Tokens at 10%, time at 80% -> the time pressure dominates.
381 let t = ResourceTelemetry {
382 tokens_used: 1_000,
383 time_used_seconds: 80,
384 token_budget: Some(10_000),
385 time_budget_seconds: Some(100),
386 };
387 let frac = t.budget_fraction().expect("bounded");
388 assert!((frac - 0.80).abs() < 1e-9, "got {frac}");
389 assert_eq!(t.budget_percent(), Some(80));
390 }
391
392 #[test]
393 fn budget_fraction_present_when_only_one_dimension_bounded() {
394 let only_time = ResourceTelemetry::new(9_999, 50).with_time_budget_seconds(200);
395 assert_eq!(only_time.budget_percent(), Some(25));
396
397 let only_tokens = ResourceTelemetry::new(2_500, 9_999).with_token_budget(10_000);
398 assert_eq!(only_tokens.budget_percent(), Some(25));
399 }
400
401 #[test]
402 fn budget_percent_rounds_to_nearest_whole() {
403 // 333 / 1000 = 33.3% -> 33
404 let down = ResourceTelemetry::new(333, 0).with_token_budget(1_000);
405 assert_eq!(down.budget_percent(), Some(33));
406 // 336 / 1000 = 33.6% -> 34
407 let up = ResourceTelemetry::new(336, 0).with_token_budget(1_000);
408 assert_eq!(up.budget_percent(), Some(34));
409 }
410
411 // ---- pressure levels --------------------------------------------------
412
413 #[test]
414 fn pressure_low_when_unbounded_regardless_of_usage() {
415 let t = ResourceTelemetry::new(u64::MAX, u64::MAX);
416 assert_eq!(t.pressure(), PressureLevel::Low);
417 }
418
419 #[test]
420 fn pressure_thresholds_just_under_and_over() {
421 // 74% -> Low (just under the medium threshold).
422 let low = ResourceTelemetry::new(7_400, 0).with_token_budget(10_000);
423 assert_eq!(low.pressure(), PressureLevel::Low);
424
425 // Exactly 75% -> Medium (inclusive lower bound).
426 let medium_edge = ResourceTelemetry::new(7_500, 0).with_token_budget(10_000);
427 assert_eq!(medium_edge.pressure(), PressureLevel::Medium);
428
429 // 99% -> Medium (just under the high threshold).
430 let medium = ResourceTelemetry::new(9_900, 0).with_token_budget(10_000);
431 assert_eq!(medium.pressure(), PressureLevel::Medium);
432
433 // Exactly 100% -> High (at budget).
434 let high_edge = ResourceTelemetry::new(10_000, 0).with_token_budget(10_000);
435 assert_eq!(high_edge.pressure(), PressureLevel::High);
436
437 // Over budget -> High.
438 let over = ResourceTelemetry::new(12_500, 0).with_token_budget(10_000);
439 assert_eq!(over.pressure(), PressureLevel::High);
440 }
441
442 #[test]
443 fn pressure_level_labels_and_ordering() {
444 assert_eq!(PressureLevel::Low.label(), "low");
445 assert_eq!(PressureLevel::Medium.label(), "medium");
446 assert_eq!(PressureLevel::High.label(), "high");
447 // Ord derive: Low < Medium < High.
448 assert!(PressureLevel::Low < PressureLevel::Medium);
449 assert!(PressureLevel::Medium < PressureLevel::High);
450 }
451
452 #[test]
453 fn pressure_from_fraction_ignores_non_finite() {
454 assert_eq!(PressureLevel::from_fraction(f64::NAN), PressureLevel::Low);
455 assert_eq!(
456 PressureLevel::from_fraction(f64::INFINITY),
457 PressureLevel::Low
458 );
459 assert_eq!(PressureLevel::from_fraction(-0.5), PressureLevel::Low);
460 }
461
462 // ---- human summary ----------------------------------------------------
463
464 #[test]
465 fn human_summary_bounded_matches_example_shape() {
466 let t = ResourceTelemetry::new(12_345, 252).with_token_budget(30_000);
467 // 12_345 -> "12.3k", 252s -> "4m 12s", 12345/30000 = 41.15% -> 41%.
468 assert_eq!(t.human_summary(), "12.3k tok · 4m 12s · 41% budget");
469 }
470
471 #[test]
472 fn human_summary_unbounded_omits_budget_segment() {
473 let t = ResourceTelemetry::new(500, 5);
474 assert_eq!(t.human_summary(), "500 tok · 5s");
475 // Display delegates to human_summary.
476 assert_eq!(t.to_string(), "500 tok · 5s");
477 }
478
479 #[test]
480 fn human_summary_zero_everything() {
481 let t = ResourceTelemetry::default();
482 assert_eq!(t.human_summary(), "0 tok · 0s");
483 }
484
485 #[test]
486 fn human_summary_over_budget_can_exceed_one_hundred_percent() {
487 let t = ResourceTelemetry::new(15_000, 7_320).with_token_budget(10_000);
488 // 15000/10000 = 150%, 7320s -> 122m 00s.
489 assert_eq!(t.human_summary(), "15k tok · 122m 00s · 150% budget");
490 assert_eq!(t.pressure(), PressureLevel::High);
491 }
492
493 #[test]
494 fn human_summary_with_only_time_budget() {
495 let t = ResourceTelemetry::new(2_000_000, 300).with_time_budget_seconds(600);
496 // 2M tokens, 5m 00s, 300/600 = 50% budget.
497 assert_eq!(t.human_summary(), "2M tok · 5m 00s · 50% budget");
498 assert_eq!(t.pressure(), PressureLevel::Low);
499 }
500 }
501
501 lines RUST