gemini_cli/prompt_segment/
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;
8
9mod refresh;
10mod render;
11
12pub use render::CacheEntry;
13
14#[derive(Clone, Debug, Default)]
15pub struct PromptSegmentOptions {
16    pub no_5h: bool,
17    pub ttl: Option<String>,
18    pub time_format: Option<String>,
19    pub show_timezone: bool,
20    pub refresh: bool,
21    pub is_enabled: bool,
22}
23
24const DEFAULT_TTL_SECONDS: u64 = 300;
25const DEFAULT_TIME_FORMAT: &str = "%m-%d %H:%M";
26const DEFAULT_TIME_FORMAT_WITH_TIMEZONE: &str = "%m-%d %H:%M %:z";
27const DEFAULT_CODE_ASSIST_ENDPOINT: &str = "https://cloudcode-pa.googleapis.com";
28const DEFAULT_CODE_ASSIST_API_VERSION: &str = "v1internal";
29const DEFAULT_CODE_ASSIST_PROJECT: &str = "projects/default";
30
31pub fn run(options: &PromptSegmentOptions) -> i32 {
32    if options.is_enabled {
33        return if prompt_segment_enabled() { 0 } else { 1 };
34    }
35
36    let ttl_seconds = match resolve_ttl_seconds(options.ttl.as_deref()) {
37        Ok(value) => value,
38        Err(_) => {
39            print_ttl_usage();
40            return 2;
41        }
42    };
43
44    if !prompt_segment_enabled() {
45        return 0;
46    }
47
48    let target_file = match paths::resolve_auth_file() {
49        Some(path) => path,
50        None => return 0,
51    };
52    if !target_file.is_file() {
53        return 0;
54    }
55
56    let show_5h =
57        shared_env::env_truthy_or("GEMINI_PROMPT_SEGMENT_SHOW_5H_ENABLED", true) && !options.no_5h;
58    let time_format = match options.time_format.as_deref() {
59        Some(value) => value,
60        None if options.show_timezone => DEFAULT_TIME_FORMAT_WITH_TIMEZONE,
61        None => DEFAULT_TIME_FORMAT,
62    };
63    let stale_suffix = std::env::var("GEMINI_PROMPT_SEGMENT_STALE_SUFFIX")
64        .unwrap_or_else(|_| " (stale)".to_string());
65
66    let prefix = resolve_name_prefix(&target_file);
67
68    if options.refresh {
69        if let Some(entry) = refresh::refresh_blocking(&target_file)
70            && let Some(line) = render::render_line(&entry, &prefix, show_5h, time_format)
71            && !line.trim().is_empty()
72        {
73            println!("{line}");
74        }
75        return 0;
76    }
77
78    let (cached, is_stale) = read_cached_entry(&target_file, ttl_seconds);
79    if let Some(entry) = cached.clone()
80        && let Some(mut line) = render::render_line(&entry, &prefix, show_5h, time_format)
81    {
82        if is_stale {
83            line.push_str(&stale_suffix);
84        }
85        if !line.trim().is_empty() {
86            println!("{line}");
87        }
88    }
89
90    if cached.is_none() || is_stale {
91        refresh::enqueue_background_refresh(&target_file);
92    }
93
94    0
95}
96
97fn read_cached_entry(target_file: &Path, ttl_seconds: u64) -> (Option<CacheEntry>, bool) {
98    let cache_file = match rate_limits::cache_file_for_target(target_file) {
99        Ok(value) => value,
100        Err(_) => return (None, false),
101    };
102    if !cache_file.is_file() {
103        return (None, false);
104    }
105
106    let entry = render::read_cache_file(&cache_file);
107    let Some(entry) = entry else {
108        return (None, false);
109    };
110
111    let now = now_epoch();
112    if now <= 0 || entry.fetched_at_epoch <= 0 {
113        return (Some(entry), true);
114    }
115
116    let ttl = i64::try_from(ttl_seconds).unwrap_or(i64::MAX);
117    let stale = now.saturating_sub(entry.fetched_at_epoch) > ttl;
118    (Some(entry), stale)
119}
120
121fn resolve_ttl_seconds(cli_ttl: Option<&str>) -> Result<u64, ()> {
122    if let Some(raw) = cli_ttl {
123        return shared_env::parse_duration_seconds(raw).ok_or(());
124    }
125
126    if let Ok(raw) = std::env::var("GEMINI_PROMPT_SEGMENT_TTL")
127        && let Some(value) = shared_env::parse_duration_seconds(&raw)
128    {
129        return Ok(value);
130    }
131
132    Ok(DEFAULT_TTL_SECONDS)
133}
134
135fn prompt_segment_enabled() -> bool {
136    shared_env::env_truthy("GEMINI_PROMPT_SEGMENT_ENABLED")
137}
138
139fn print_ttl_usage() {
140    eprintln!("gemini-cli prompt-segment: invalid --ttl");
141    eprintln!(
142        "usage: gemini-cli prompt-segment [--no-5h] [--ttl <duration>] [--time-format <strftime>] [--show-timezone] [--refresh] [--is-enabled]"
143    );
144}
145
146fn resolve_name_prefix(target_file: &Path) -> String {
147    let name = resolve_name(target_file);
148    match name {
149        Some(value) if !value.trim().is_empty() => format!("{} ", value.trim()),
150        _ => String::new(),
151    }
152}
153
154fn resolve_name(target_file: &Path) -> Option<String> {
155    let source = std::env::var("GEMINI_PROMPT_SEGMENT_NAME_SOURCE")
156        .ok()
157        .map(|value| value.to_ascii_lowercase())
158        .unwrap_or_else(|| "secret".to_string());
159    let show_fallback = shared_env::env_truthy("GEMINI_PROMPT_SEGMENT_SHOW_FALLBACK_NAME_ENABLED");
160    let show_full_email = shared_env::env_truthy("GEMINI_PROMPT_SEGMENT_SHOW_FULL_EMAIL_ENABLED");
161
162    if source == "email" {
163        if let Ok(Some(email)) = auth::email_from_auth_file(target_file) {
164            return Some(format_email_name(&email, show_full_email));
165        }
166        if show_fallback && let Ok(Some(identity)) = auth::identity_from_auth_file(target_file) {
167            return Some(format_email_name(&identity, show_full_email));
168        }
169        return None;
170    }
171
172    if let Some(secret_name) = rate_limits::secret_name_for_target(target_file) {
173        return Some(secret_name);
174    }
175
176    if show_fallback && let Ok(Some(identity)) = auth::identity_from_auth_file(target_file) {
177        return Some(format_email_name(&identity, show_full_email));
178    }
179
180    None
181}
182
183fn format_email_name(raw: &str, show_full_email: bool) -> String {
184    let trimmed = raw.trim();
185    if show_full_email {
186        return trimmed.to_string();
187    }
188    trimmed.split('@').next().unwrap_or(trimmed).to_string()
189}
190
191fn env_u64(key: &str, default: u64) -> u64 {
192    std::env::var(key)
193        .ok()
194        .and_then(|raw| raw.trim().parse::<u64>().ok())
195        .unwrap_or(default)
196}
197
198fn code_assist_endpoint() -> String {
199    env_non_empty("CODE_ASSIST_ENDPOINT")
200        .or_else(|| env_non_empty("GEMINI_CODE_ASSIST_ENDPOINT"))
201        .unwrap_or_else(|| DEFAULT_CODE_ASSIST_ENDPOINT.to_string())
202}
203
204fn code_assist_api_version() -> String {
205    env_non_empty("CODE_ASSIST_API_VERSION")
206        .or_else(|| env_non_empty("GEMINI_CODE_ASSIST_API_VERSION"))
207        .unwrap_or_else(|| DEFAULT_CODE_ASSIST_API_VERSION.to_string())
208}
209
210fn code_assist_project() -> String {
211    let raw = env_non_empty("GEMINI_CODE_ASSIST_PROJECT")
212        .or_else(|| env_non_empty("GOOGLE_CLOUD_PROJECT"))
213        .or_else(|| env_non_empty("GOOGLE_CLOUD_PROJECT_ID"))
214        .unwrap_or_else(|| DEFAULT_CODE_ASSIST_PROJECT.to_string());
215
216    if raw.starts_with("projects/") {
217        raw
218    } else {
219        format!("projects/{raw}")
220    }
221}
222
223fn env_non_empty(key: &str) -> Option<String> {
224    std::env::var(key)
225        .ok()
226        .map(|raw| raw.trim().to_string())
227        .filter(|raw| !raw.is_empty())
228}
229
230fn now_epoch() -> i64 {
231    std::time::SystemTime::now()
232        .duration_since(std::time::UNIX_EPOCH)
233        .ok()
234        .and_then(|duration| i64::try_from(duration.as_secs()).ok())
235        .unwrap_or(0)
236}
gemini_cli/prompt_segment/mod.rs

gemini_cli/prompt_segment/
mod.rs
