Skip to main content

codex_cli/starship/
render.rs

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