Skip to main content

gemini_cli/starship/
render.rs

1use nils_common::env as shared_env;
2use std::path::Path;
3use std::{io, io::IsTerminal};
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 = crate::rate_limits::render::format_epoch_local(
72        entry.weekly_reset_epoch,
73        weekly_reset_time_format,
74    )
75    .unwrap_or_else(|| "?".to_string());
76
77    let color_enabled = should_color();
78    let weekly_token = ansi::format_percent_token(
79        &format!("W:{}%", entry.weekly_remaining),
80        Some(color_enabled),
81    );
82
83    if show_5h {
84        let non_weekly_token = ansi::format_percent_token(
85            &format!("{}:{}%", entry.non_weekly_label, entry.non_weekly_remaining),
86            Some(color_enabled),
87        );
88        return Some(format!(
89            "{prefix}{non_weekly_token} {weekly_token} {weekly_reset_time}"
90        ));
91    }
92
93    Some(format!("{prefix}{weekly_token} {weekly_reset_time}"))
94}
95
96fn should_color() -> bool {
97    if shared_env::no_color_enabled() {
98        return false;
99    }
100
101    if std::env::var("GEMINI_STARSHIP_COLOR_ENABLED").is_ok() {
102        return shared_env::env_truthy("GEMINI_STARSHIP_COLOR_ENABLED");
103    }
104
105    if std::env::var_os("STARSHIP_SESSION_KEY").is_some()
106        || std::env::var_os("STARSHIP_SHELL").is_some()
107    {
108        return true;
109    }
110
111    io::stdout().is_terminal()
112}
113
114#[cfg(test)]
115mod tests {
116    use super::should_color;
117    use nils_test_support::{EnvGuard, GlobalStateLock};
118
119    #[test]
120    fn should_color_no_color_has_highest_priority() {
121        let lock = GlobalStateLock::new();
122        let _no_color = EnvGuard::set(&lock, "NO_COLOR", "1");
123        let _explicit = EnvGuard::set(&lock, "GEMINI_STARSHIP_COLOR_ENABLED", "true");
124        let _session = EnvGuard::set(&lock, "STARSHIP_SESSION_KEY", "session");
125        assert!(!should_color());
126    }
127
128    #[test]
129    fn should_color_explicit_truthy_and_falsey_values_are_stable() {
130        let lock = GlobalStateLock::new();
131        let _no_color = EnvGuard::remove(&lock, "NO_COLOR");
132        let _session = EnvGuard::remove(&lock, "STARSHIP_SESSION_KEY");
133        let _shell = EnvGuard::remove(&lock, "STARSHIP_SHELL");
134
135        for value in ["1", " true ", "YES", "on"] {
136            let _explicit = EnvGuard::set(&lock, "GEMINI_STARSHIP_COLOR_ENABLED", value);
137            assert!(should_color(), "expected truthy value: {value}");
138        }
139
140        for value in ["", " ", "0", "false", "no", "off", "y", "enabled"] {
141            let _explicit = EnvGuard::set(&lock, "GEMINI_STARSHIP_COLOR_ENABLED", value);
142            assert!(!should_color(), "expected falsey value: {value}");
143        }
144    }
145
146    #[test]
147    fn should_color_falls_back_to_starship_markers_when_not_overridden() {
148        let lock = GlobalStateLock::new();
149        let _no_color = EnvGuard::remove(&lock, "NO_COLOR");
150        let _explicit = EnvGuard::remove(&lock, "GEMINI_STARSHIP_COLOR_ENABLED");
151        let _session = EnvGuard::set(&lock, "STARSHIP_SESSION_KEY", "session");
152        assert!(should_color());
153    }
154}