Skip to main content

codex_cli/starship/
render.rs

1use chrono::{Local, TimeZone};
2use nils_common::env as shared_env;
3use std::path::Path;
4
5use crate::rate_limits::ansi;
6
7#[derive(Clone, Debug)]
8pub struct CacheEntry {
9    pub fetched_at_epoch: i64,
10    pub non_weekly_label: String,
11    pub non_weekly_remaining: i64,
12    pub non_weekly_reset_epoch: Option<i64>,
13    pub weekly_remaining: i64,
14    pub weekly_reset_epoch: i64,
15}
16
17pub fn read_cache_file(path: &Path) -> Option<CacheEntry> {
18    let content = std::fs::read_to_string(path).ok()?;
19    parse_cache_kv(&content)
20}
21
22fn parse_cache_kv(content: &str) -> Option<CacheEntry> {
23    let mut fetched_at_epoch: Option<i64> = None;
24    let mut non_weekly_label: Option<String> = None;
25    let mut non_weekly_remaining: Option<i64> = None;
26    let mut non_weekly_reset_epoch: Option<i64> = None;
27    let mut weekly_remaining: Option<i64> = None;
28    let mut weekly_reset_epoch: Option<i64> = None;
29
30    for line in content.lines() {
31        if let Some(value) = line.strip_prefix("fetched_at=") {
32            fetched_at_epoch = value.parse::<i64>().ok();
33        } else if let Some(value) = line.strip_prefix("non_weekly_label=") {
34            non_weekly_label = Some(value.to_string());
35        } else if let Some(value) = line.strip_prefix("non_weekly_remaining=") {
36            non_weekly_remaining = value.parse::<i64>().ok();
37        } else if let Some(value) = line.strip_prefix("non_weekly_reset_epoch=") {
38            non_weekly_reset_epoch = value.parse::<i64>().ok();
39        } else if let Some(value) = line.strip_prefix("weekly_remaining=") {
40            weekly_remaining = value.parse::<i64>().ok();
41        } else if let Some(value) = line.strip_prefix("weekly_reset_epoch=") {
42            weekly_reset_epoch = value.parse::<i64>().ok();
43        }
44    }
45
46    let fetched_at_epoch = fetched_at_epoch?;
47    let non_weekly_label = non_weekly_label?;
48    if non_weekly_label.trim().is_empty() {
49        return None;
50    }
51    let non_weekly_remaining = non_weekly_remaining?;
52    let weekly_remaining = weekly_remaining?;
53    let weekly_reset_epoch = weekly_reset_epoch?;
54
55    Some(CacheEntry {
56        fetched_at_epoch,
57        non_weekly_label,
58        non_weekly_remaining,
59        non_weekly_reset_epoch,
60        weekly_remaining,
61        weekly_reset_epoch,
62    })
63}
64
65pub fn render_line(
66    entry: &CacheEntry,
67    prefix: &str,
68    show_5h: bool,
69    weekly_reset_time_format: &str,
70) -> Option<String> {
71    let weekly_reset_time = format_epoch_local(entry.weekly_reset_epoch, weekly_reset_time_format)
72        .unwrap_or_else(|| "?".to_string());
73
74    let color_enabled = should_color();
75    let weekly_token = ansi::format_percent_token(
76        &format!("W:{}%", entry.weekly_remaining),
77        Some(color_enabled),
78    );
79
80    if show_5h {
81        let non_weekly_token = ansi::format_percent_token(
82            &format!("{}:{}%", entry.non_weekly_label, entry.non_weekly_remaining),
83            Some(color_enabled),
84        );
85        return Some(format!(
86            "{prefix}{non_weekly_token} {weekly_token} {weekly_reset_time}"
87        ));
88    }
89
90    Some(format!("{prefix}{weekly_token} {weekly_reset_time}"))
91}
92
93fn format_epoch_local(epoch: i64, fmt: &str) -> Option<String> {
94    let dt = Local.timestamp_opt(epoch, 0).single()?;
95    Some(dt.format(fmt).to_string())
96}
97
98fn should_color() -> bool {
99    shared_env::starship_color_enabled("CODEX_STARSHIP_COLOR_ENABLED")
100}
101
102#[cfg(test)]
103mod tests {
104    use super::should_color;
105    use nils_test_support::{EnvGuard, GlobalStateLock};
106
107    #[test]
108    fn should_color_no_color_has_highest_priority() {
109        let lock = GlobalStateLock::new();
110        let _no_color = EnvGuard::set(&lock, "NO_COLOR", "1");
111        let _explicit = EnvGuard::set(&lock, "CODEX_STARSHIP_COLOR_ENABLED", "true");
112        let _session = EnvGuard::set(&lock, "STARSHIP_SESSION_KEY", "session");
113        assert!(!should_color());
114    }
115
116    #[test]
117    fn should_color_explicit_truthy_and_falsey_values_are_stable() {
118        let lock = GlobalStateLock::new();
119        let _no_color = EnvGuard::remove(&lock, "NO_COLOR");
120        let _session = EnvGuard::remove(&lock, "STARSHIP_SESSION_KEY");
121        let _shell = EnvGuard::remove(&lock, "STARSHIP_SHELL");
122
123        for value in ["1", " true ", "YES", "on"] {
124            let _explicit = EnvGuard::set(&lock, "CODEX_STARSHIP_COLOR_ENABLED", value);
125            assert!(should_color(), "expected truthy value: {value}");
126        }
127
128        for value in ["", " ", "0", "false", "no", "off", "y", "enabled"] {
129            let _explicit = EnvGuard::set(&lock, "CODEX_STARSHIP_COLOR_ENABLED", value);
130            assert!(!should_color(), "expected falsey value: {value}");
131        }
132    }
133
134    #[test]
135    fn should_color_falls_back_to_starship_markers_when_not_overridden() {
136        let lock = GlobalStateLock::new();
137        let _no_color = EnvGuard::remove(&lock, "NO_COLOR");
138        let _explicit = EnvGuard::remove(&lock, "CODEX_STARSHIP_COLOR_ENABLED");
139        let _session = EnvGuard::set(&lock, "STARSHIP_SESSION_KEY", "session");
140        assert!(should_color());
141    }
142}