Skip to main content

gemini_cli/starship/
render.rs

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