claude_cli/prompt_segment/
mod.rs

1use nils_common::cli_contract::exit;
2use nils_common::diag_output;
3use nils_common::env as shared_env;
4use serde::Serialize;
5use std::path::PathBuf;
6
7mod auth;
8mod cache;
9mod client;
10mod render;
11
12#[derive(Clone, Debug, Default)]
13pub struct PromptSegmentOptions {
14    pub ttl: Option<String>,
15    pub time_format: Option<String>,
16    pub refresh: bool,
17}
18
19const DEFAULT_TTL_SECONDS: u64 = 60;
20const DEFAULT_TIME_FORMAT: &str = "%m-%d %H:%M";
21const DEFAULT_STALE_SUFFIX: &str = " (stale)";
22const PROMPT_SEGMENT_SCHEMA_VERSION: &str = "claude-cli.prompt-segment.v1";
23
24pub fn run(options: &PromptSegmentOptions) -> i32 {
25    let ttl_seconds = match resolve_ttl_seconds(options.ttl.as_deref()) {
26        Ok(value) => value,
27        Err(_) => {
28            eprintln!("claude-cli prompt-segment: invalid --ttl");
29            return exit::USAGE;
30        }
31    };
32
33    let Some(cache_file) = cache::cache_file() else {
34        return exit::SUCCESS;
35    };
36
37    let force_refresh = options.refresh || ttl_seconds == 0;
38    let needs_refresh =
39        force_refresh || !cache_file.is_file() || cache::cache_stale(&cache_file, ttl_seconds);
40    let mut stale = false;
41
42    if needs_refresh {
43        match auth::resolve_access_token()
44            .map(|token| client::fetch_usage(&token.value))
45            .transpose()
46        {
47            Ok(Some(body)) => {
48                if cache::write_cache_file(&cache_file, &body).is_err() {
49                    stale = true;
50                }
51            }
52            _ => {
53                stale = true;
54            }
55        }
56    }
57
58    let Some(raw_cache) = cache::read_cache_file(&cache_file) else {
59        return exit::SUCCESS;
60    };
61
62    let time_format = options
63        .time_format
64        .as_deref()
65        .unwrap_or(DEFAULT_TIME_FORMAT);
66    let stale_suffix = resolve_stale_suffix();
67    if let Some(line) = render::render_usage_json(&raw_cache, time_format, stale, &stale_suffix)
68        && !line.trim().is_empty()
69    {
70        println!("{line}");
71    }
72
73    exit::SUCCESS
74}
75
76pub fn check() -> i32 {
77    if auth::resolve_access_token().is_some() {
78        exit::SUCCESS
79    } else {
80        exit::RUNTIME
81    }
82}
83
84pub fn status(output_json: bool) -> i32 {
85    let result = PromptSegmentStatusResult::inspect();
86
87    if output_json {
88        if diag_output::emit_success_result(
89            PROMPT_SEGMENT_SCHEMA_VERSION,
90            "prompt-segment status",
91            &result,
92        )
93        .is_err()
94        {
95            return exit::RUNTIME;
96        }
97    } else {
98        println!(
99            "claude: prompt-segment status authenticated={} would_render={} reason={}",
100            result.authenticated, result.would_render, result.reason
101        );
102    }
103
104    exit::SUCCESS
105}
106
107fn resolve_ttl_seconds(cli_ttl: Option<&str>) -> Result<u64, ()> {
108    if let Some(raw) = cli_ttl {
109        return parse_ttl_seconds(raw).ok_or(());
110    }
111
112    for key in ["CLAUDE_PROMPT_SEGMENT_TTL", "CLAUDE_PROMPT_TTL"] {
113        if let Ok(raw) = std::env::var(key)
114            && let Some(value) = parse_ttl_seconds(&raw)
115        {
116            return Ok(value);
117        }
118    }
119
120    Ok(DEFAULT_TTL_SECONDS)
121}
122
123fn parse_ttl_seconds(raw: &str) -> Option<u64> {
124    let raw = raw.trim();
125    if raw == "0" || raw.eq_ignore_ascii_case("0s") {
126        return Some(0);
127    }
128    shared_env::parse_duration_seconds(raw)
129}
130
131fn resolve_stale_suffix() -> String {
132    std::env::var("CLAUDE_PROMPT_SEGMENT_STALE_SUFFIX")
133        .ok()
134        .or_else(|| std::env::var("CLAUDE_PROMPT_STALE_SUFFIX").ok())
135        .unwrap_or_else(|| DEFAULT_STALE_SUFFIX.to_string())
136}
137
138#[derive(Debug, Clone, Serialize)]
139struct PromptSegmentStatusResult {
140    authenticated: bool,
141    auth_source: Option<String>,
142    cache_file: Option<String>,
143    cache_exists: bool,
144    cache_stale: bool,
145    would_render: bool,
146    reason: String,
147}
148
149impl PromptSegmentStatusResult {
150    fn inspect() -> Self {
151        let token = auth::resolve_access_token();
152        let cache_file = cache::cache_file();
153        let ttl_seconds = resolve_ttl_seconds(None).unwrap_or(DEFAULT_TTL_SECONDS);
154        let (cache_exists, cache_stale, would_render) =
155            inspect_cache(cache_file.as_ref(), ttl_seconds);
156
157        let authenticated = token.is_some();
158        let reason = if authenticated && would_render && !cache_stale {
159            "ready"
160        } else if authenticated && would_render {
161            "cache-stale"
162        } else if !authenticated {
163            "access-token-missing"
164        } else if !cache_exists {
165            "cache-missing"
166        } else {
167            "cache-empty-or-invalid"
168        };
169
170        Self {
171            authenticated,
172            auth_source: token.map(|token| match token.source {
173                auth::TokenSource::AccessTokenEnv => "access-token-env".to_string(),
174                auth::TokenSource::CredentialsJsonEnv => "credentials-json-env".to_string(),
175                auth::TokenSource::Keychain => "keychain".to_string(),
176            }),
177            cache_file: cache_file.map(display_path),
178            cache_exists,
179            cache_stale,
180            would_render,
181            reason: reason.to_string(),
182        }
183    }
184}
185
186fn inspect_cache(cache_file: Option<&PathBuf>, ttl_seconds: u64) -> (bool, bool, bool) {
187    let Some(cache_file) = cache_file else {
188        return (false, false, false);
189    };
190    if !cache_file.is_file() {
191        return (false, false, false);
192    }
193
194    let cache_stale = cache::cache_stale(cache_file, ttl_seconds);
195    let would_render = cache::read_cache_file(cache_file)
196        .as_deref()
197        .and_then(|raw| render::render_usage_json(raw, DEFAULT_TIME_FORMAT, false, ""))
198        .map(|line| !line.trim().is_empty())
199        .unwrap_or(false);
200
201    (true, cache_stale, would_render)
202}
203
204fn display_path(path: PathBuf) -> String {
205    path.to_string_lossy().to_string()
206}
207
208#[cfg(test)]
209mod tests {
210    use super::{parse_ttl_seconds, resolve_ttl_seconds};
211    use nils_test_support::{EnvGuard, GlobalStateLock};
212
213    #[test]
214    fn parse_ttl_seconds_accepts_zero_for_legacy_force_refresh() {
215        assert_eq!(parse_ttl_seconds("0"), Some(0));
216        assert_eq!(parse_ttl_seconds("0s"), Some(0));
217        assert_eq!(parse_ttl_seconds("60"), Some(60));
218        assert_eq!(parse_ttl_seconds("2m"), Some(120));
219    }
220
221    #[test]
222    fn resolve_ttl_prefers_segment_env_before_legacy_env() {
223        let lock = GlobalStateLock::new();
224        let _segment = EnvGuard::set(&lock, "CLAUDE_PROMPT_SEGMENT_TTL", "2m");
225        let _legacy = EnvGuard::set(&lock, "CLAUDE_PROMPT_TTL", "30");
226        assert_eq!(resolve_ttl_seconds(None), Ok(120));
227    }
228}
claude_cli/prompt_segment/mod.rs

claude_cli/prompt_segment/
mod.rs
