Skip to main content

gemini_cli/starship/
mod.rs

1use std::path::Path;
2
3use nils_common::env as shared_env;
4
5use crate::auth;
6use crate::paths;
7use crate::rate_limits;
8use crate::rate_limits::client::{UsageRequest, fetch_usage};
9use crate::rate_limits::render as rate_render;
10
11mod render;
12
13pub use render::CacheEntry;
14
15#[derive(Clone, Debug, Default)]
16pub struct StarshipOptions {
17    pub no_5h: bool,
18    pub ttl: Option<String>,
19    pub time_format: Option<String>,
20    pub show_timezone: bool,
21    pub refresh: bool,
22    pub is_enabled: bool,
23}
24
25const DEFAULT_TTL_SECONDS: u64 = 300;
26const DEFAULT_TIME_FORMAT: &str = "%m-%d %H:%M";
27const DEFAULT_TIME_FORMAT_WITH_TIMEZONE: &str = "%m-%d %H:%M %:z";
28const DEFAULT_CODE_ASSIST_ENDPOINT: &str = "https://cloudcode-pa.googleapis.com";
29const DEFAULT_CODE_ASSIST_API_VERSION: &str = "v1internal";
30const DEFAULT_CODE_ASSIST_PROJECT: &str = "projects/default";
31
32pub fn run(options: &StarshipOptions) -> i32 {
33    if options.is_enabled {
34        return if starship_enabled() { 0 } else { 1 };
35    }
36
37    let ttl_seconds = match resolve_ttl_seconds(options.ttl.as_deref()) {
38        Ok(value) => value,
39        Err(_) => {
40            print_ttl_usage();
41            return 2;
42        }
43    };
44
45    if !starship_enabled() {
46        return 0;
47    }
48
49    let target_file = match paths::resolve_auth_file() {
50        Some(path) => path,
51        None => return 0,
52    };
53    if !target_file.is_file() {
54        return 0;
55    }
56
57    let show_5h =
58        shared_env::env_truthy_or("GEMINI_STARSHIP_SHOW_5H_ENABLED", true) && !options.no_5h;
59    let time_format = match options.time_format.as_deref() {
60        Some(value) => value,
61        None if options.show_timezone => DEFAULT_TIME_FORMAT_WITH_TIMEZONE,
62        None => DEFAULT_TIME_FORMAT,
63    };
64    let stale_suffix =
65        std::env::var("GEMINI_STARSHIP_STALE_SUFFIX").unwrap_or_else(|_| " (stale)".to_string());
66
67    let prefix = resolve_name_prefix(&target_file);
68
69    if options.refresh {
70        if let Some(entry) = refresh_blocking(&target_file)
71            && let Some(line) = render::render_line(&entry, &prefix, show_5h, time_format)
72            && !line.trim().is_empty()
73        {
74            println!("{line}");
75        }
76        return 0;
77    }
78
79    let (cached, is_stale) = read_cached_entry(&target_file, ttl_seconds);
80    if let Some(entry) = cached.clone()
81        && let Some(mut line) = render::render_line(&entry, &prefix, show_5h, time_format)
82    {
83        if is_stale {
84            line.push_str(&stale_suffix);
85        }
86        if !line.trim().is_empty() {
87            println!("{line}");
88        }
89    }
90
91    if cached.is_none() || is_stale {
92        let _ = refresh_blocking(&target_file);
93    }
94
95    0
96}
97
98fn refresh_blocking(target_file: &Path) -> Option<render::CacheEntry> {
99    let connect_timeout = env_u64("GEMINI_STARSHIP_CURL_CONNECT_TIMEOUT_SECONDS", 2);
100    let max_time = env_u64("GEMINI_STARSHIP_CURL_MAX_TIME_SECONDS", 8);
101
102    let usage_request = UsageRequest {
103        target_file: target_file.to_path_buf(),
104        refresh_on_401: false,
105        endpoint: code_assist_endpoint(),
106        api_version: code_assist_api_version(),
107        project: code_assist_project(),
108        connect_timeout_seconds: connect_timeout,
109        max_time_seconds: max_time,
110    };
111
112    let usage = fetch_usage(&usage_request).ok()?;
113    let usage_data = rate_render::parse_usage(&usage.body)?;
114    let values = rate_render::render_values(&usage_data);
115    let weekly = rate_render::weekly_values(&values);
116
117    let fetched_at_epoch = now_epoch();
118    if fetched_at_epoch > 0 {
119        let _ = rate_limits::write_starship_cache(
120            target_file,
121            fetched_at_epoch,
122            &weekly.non_weekly_label,
123            weekly.non_weekly_remaining,
124            weekly.weekly_remaining,
125            weekly.weekly_reset_epoch,
126            weekly.non_weekly_reset_epoch,
127        );
128    }
129
130    Some(render::CacheEntry {
131        fetched_at_epoch,
132        non_weekly_label: weekly.non_weekly_label,
133        non_weekly_remaining: weekly.non_weekly_remaining,
134        non_weekly_reset_epoch: weekly.non_weekly_reset_epoch,
135        weekly_remaining: weekly.weekly_remaining,
136        weekly_reset_epoch: weekly.weekly_reset_epoch,
137    })
138}
139
140fn read_cached_entry(target_file: &Path, ttl_seconds: u64) -> (Option<CacheEntry>, bool) {
141    let cache_file = match rate_limits::cache_file_for_target(target_file) {
142        Ok(value) => value,
143        Err(_) => return (None, false),
144    };
145    if !cache_file.is_file() {
146        return (None, false);
147    }
148
149    let entry = render::read_cache_file(&cache_file);
150    let Some(entry) = entry else {
151        return (None, false);
152    };
153
154    let now = now_epoch();
155    if now <= 0 || entry.fetched_at_epoch <= 0 {
156        return (Some(entry), true);
157    }
158
159    let ttl = i64::try_from(ttl_seconds).unwrap_or(i64::MAX);
160    let stale = now.saturating_sub(entry.fetched_at_epoch) > ttl;
161    (Some(entry), stale)
162}
163
164fn resolve_ttl_seconds(cli_ttl: Option<&str>) -> Result<u64, ()> {
165    if let Some(raw) = cli_ttl {
166        return parse_duration_seconds(raw).ok_or(());
167    }
168
169    if let Ok(raw) = std::env::var("GEMINI_STARSHIP_TTL")
170        && let Some(value) = parse_duration_seconds(&raw)
171    {
172        return Ok(value);
173    }
174
175    Ok(DEFAULT_TTL_SECONDS)
176}
177
178fn parse_duration_seconds(raw: &str) -> Option<u64> {
179    let raw = raw.trim();
180    if raw.is_empty() {
181        return None;
182    }
183
184    let normalized = raw.to_ascii_lowercase();
185    let (num_part, multiplier) = match normalized.chars().last()? {
186        's' => (&normalized[..normalized.len().saturating_sub(1)], 1u64),
187        'm' => (&normalized[..normalized.len().saturating_sub(1)], 60u64),
188        'h' => (
189            &normalized[..normalized.len().saturating_sub(1)],
190            60u64 * 60u64,
191        ),
192        'd' => (
193            &normalized[..normalized.len().saturating_sub(1)],
194            60u64 * 60u64 * 24u64,
195        ),
196        'w' => (
197            &normalized[..normalized.len().saturating_sub(1)],
198            60u64 * 60u64 * 24u64 * 7u64,
199        ),
200        ch if ch.is_ascii_digit() => (normalized.as_str(), 1u64),
201        _ => return None,
202    };
203
204    let num = num_part.trim().parse::<u64>().ok()?;
205    if num == 0 {
206        return None;
207    }
208    num.checked_mul(multiplier)
209}
210
211fn starship_enabled() -> bool {
212    shared_env::env_truthy("GEMINI_STARSHIP_ENABLED")
213}
214
215fn print_ttl_usage() {
216    eprintln!("gemini-cli starship: invalid --ttl");
217    eprintln!(
218        "usage: gemini-cli starship [--no-5h] [--ttl <duration>] [--time-format <strftime>] [--show-timezone] [--refresh] [--is-enabled]"
219    );
220}
221
222fn resolve_name_prefix(target_file: &Path) -> String {
223    let name = resolve_name(target_file);
224    match name {
225        Some(value) if !value.trim().is_empty() => format!("{} ", value.trim()),
226        _ => String::new(),
227    }
228}
229
230fn resolve_name(target_file: &Path) -> Option<String> {
231    let source = std::env::var("GEMINI_STARSHIP_NAME_SOURCE")
232        .ok()
233        .map(|value| value.to_ascii_lowercase())
234        .unwrap_or_else(|| "secret".to_string());
235    let show_fallback = shared_env::env_truthy("GEMINI_STARSHIP_SHOW_FALLBACK_NAME_ENABLED");
236    let show_full_email = shared_env::env_truthy("GEMINI_STARSHIP_SHOW_FULL_EMAIL_ENABLED");
237
238    if source == "email" {
239        if let Ok(Some(email)) = auth::email_from_auth_file(target_file) {
240            return Some(format_email_name(&email, show_full_email));
241        }
242        if show_fallback && let Ok(Some(identity)) = auth::identity_from_auth_file(target_file) {
243            return Some(format_email_name(&identity, show_full_email));
244        }
245        return None;
246    }
247
248    if let Some(secret_name) = rate_limits::secret_name_for_target(target_file) {
249        return Some(secret_name);
250    }
251
252    if show_fallback && let Ok(Some(identity)) = auth::identity_from_auth_file(target_file) {
253        return Some(format_email_name(&identity, show_full_email));
254    }
255
256    None
257}
258
259fn format_email_name(raw: &str, show_full_email: bool) -> String {
260    let trimmed = raw.trim();
261    if show_full_email {
262        return trimmed.to_string();
263    }
264    trimmed.split('@').next().unwrap_or(trimmed).to_string()
265}
266
267fn env_u64(key: &str, default: u64) -> u64 {
268    std::env::var(key)
269        .ok()
270        .and_then(|raw| raw.trim().parse::<u64>().ok())
271        .unwrap_or(default)
272}
273
274fn code_assist_endpoint() -> String {
275    env_non_empty("CODE_ASSIST_ENDPOINT")
276        .or_else(|| env_non_empty("GEMINI_CODE_ASSIST_ENDPOINT"))
277        .unwrap_or_else(|| DEFAULT_CODE_ASSIST_ENDPOINT.to_string())
278}
279
280fn code_assist_api_version() -> String {
281    env_non_empty("CODE_ASSIST_API_VERSION")
282        .or_else(|| env_non_empty("GEMINI_CODE_ASSIST_API_VERSION"))
283        .unwrap_or_else(|| DEFAULT_CODE_ASSIST_API_VERSION.to_string())
284}
285
286fn code_assist_project() -> String {
287    let raw = env_non_empty("GEMINI_CODE_ASSIST_PROJECT")
288        .or_else(|| env_non_empty("GOOGLE_CLOUD_PROJECT"))
289        .or_else(|| env_non_empty("GOOGLE_CLOUD_PROJECT_ID"))
290        .unwrap_or_else(|| DEFAULT_CODE_ASSIST_PROJECT.to_string());
291
292    if raw.starts_with("projects/") {
293        raw
294    } else {
295        format!("projects/{raw}")
296    }
297}
298
299fn env_non_empty(key: &str) -> Option<String> {
300    std::env::var(key)
301        .ok()
302        .map(|raw| raw.trim().to_string())
303        .filter(|raw| !raw.is_empty())
304}
305
306fn now_epoch() -> i64 {
307    std::time::SystemTime::now()
308        .duration_since(std::time::UNIX_EPOCH)
309        .ok()
310        .and_then(|duration| i64::try_from(duration.as_secs()).ok())
311        .unwrap_or(0)
312}