claude_code_stats/source/
cli_probe.rs

1use std::io::{Read, Write};
2use std::os::fd::{AsFd, AsRawFd, FromRawFd};
3use std::process::Command;
4use std::time::{Duration, Instant};
5
6use anyhow::{Result, anyhow};
7use regex::Regex;
8
9use crate::source::SourceError;
10use crate::source::UsageSource;
11use crate::types::{ApiResponse, ApiWindow};
12
13const PROBE_TIMEOUT: Duration = Duration::from_secs(15);
14
15pub struct CliProbeSource;
16
17impl UsageSource for CliProbeSource {
18    fn name(&self) -> &'static str {
19        "cli_probe"
20    }
21
22    fn try_fetch(&self) -> Result<ApiResponse, SourceError> {
23        let which = Command::new("which")
24            .arg("claude")
25            .output()
26            .map_err(|e| SourceError::NotAvailable(format!("cannot check for claude: {e}")))?;
27
28        if !which.status.success() {
29            return Err(SourceError::NotAvailable(
30                "claude CLI not found".to_string(),
31            ));
32        }
33
34        run_cli_probe().map_err(SourceError::Failed)
35    }
36}
37
38fn run_cli_probe() -> Result<ApiResponse> {
39    let pty = nix::pty::openpty(None, None).map_err(|e| anyhow!("openpty failed: {e}"))?;
40
41    let slave_fd = pty.slave.as_raw_fd();
42    let master_fd = pty.master.as_raw_fd();
43
44    match unsafe { nix::unistd::fork() } {
45        Ok(nix::unistd::ForkResult::Child) => {
46            drop(pty.master);
47            let _ = nix::unistd::setsid();
48            let _ = nix::unistd::dup2(slave_fd, 0);
49            let _ = nix::unistd::dup2(slave_fd, 1);
50            let _ = nix::unistd::dup2(slave_fd, 2);
51            drop(pty.slave);
52
53            let err = nix::unistd::execvp(c"claude", &[c"claude", c"--allowed-tools", c""]);
54            eprintln!("execvp failed: {err:?}");
55            std::process::exit(1);
56        }
57        Ok(nix::unistd::ForkResult::Parent { child }) => {
58            drop(pty.slave);
59
60            let result = interact_with_pty(master_fd);
61
62            let _ = nix::sys::signal::kill(child, nix::sys::signal::Signal::SIGTERM);
63            let _ = nix::sys::wait::waitpid(child, None);
64
65            result
66        }
67        Err(e) => Err(anyhow!("fork failed: {e}")),
68    }
69}
70
71fn interact_with_pty(master_fd: i32) -> Result<ApiResponse> {
72    let mut master = unsafe { std::fs::File::from_raw_fd(master_fd) };
73    let start = Instant::now();
74    let mut buffer = Vec::new();
75    let mut read_buf = [0u8; 4096];
76    let mut sent_usage = false;
77    let mut usage_output_start = 0;
78
79    // Set non-blocking via the BorrowedFd
80    let borrowed = master.as_fd();
81    let flags = nix::fcntl::fcntl(borrowed.as_raw_fd(), nix::fcntl::FcntlArg::F_GETFL)?;
82    let mut oflags = nix::fcntl::OFlag::from_bits_truncate(flags);
83    oflags.insert(nix::fcntl::OFlag::O_NONBLOCK);
84    nix::fcntl::fcntl(borrowed.as_raw_fd(), nix::fcntl::FcntlArg::F_SETFL(oflags))?;
85
86    loop {
87        if start.elapsed() > PROBE_TIMEOUT {
88            return Err(anyhow!("CLI probe timed out after {PROBE_TIMEOUT:?}"));
89        }
90
91        match master.read(&mut read_buf) {
92            Ok(0) => break,
93            Ok(n) => {
94                buffer.extend_from_slice(&read_buf[..n]);
95            }
96            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
97                std::thread::sleep(Duration::from_millis(100));
98            }
99            Err(e) => return Err(anyhow!("read error: {e}")),
100        }
101
102        let text = String::from_utf8_lossy(&buffer);
103
104        if text.contains("Do you trust") {
105            let _ = master.write_all(b"y\n");
106            let _ = master.flush();
107            buffer.clear();
108            continue;
109        }
110
111        if text.contains("Quick safety check") || text.contains("press Enter") {
112            let _ = master.write_all(b"\n");
113            let _ = master.flush();
114            buffer.clear();
115            continue;
116        }
117
118        if !sent_usage
119            && (text.contains("> ") || text.contains("\u{276F}") || text.contains("claude>"))
120        {
121            let _ = master.write_all(b"/usage\n");
122            let _ = master.flush();
123            sent_usage = true;
124            usage_output_start = buffer.len();
125            continue;
126        }
127
128        if sent_usage {
129            let output_text = String::from_utf8_lossy(&buffer[usage_output_start..]);
130            if output_text.contains("> ")
131                || output_text.contains("\u{276F}")
132                || output_text.contains("claude>")
133            {
134                let full_output = String::from_utf8_lossy(&buffer);
135                let cleaned = strip_ansi_codes(&full_output);
136                return parse_usage_output(&cleaned);
137            }
138        }
139    }
140
141    let full_output = String::from_utf8_lossy(&buffer);
142    let cleaned = strip_ansi_codes(&full_output);
143    parse_usage_output(&cleaned)
144}
145
146fn strip_ansi_codes(s: &str) -> String {
147    let re = Regex::new(r"\x1B\[[0-?]*[ -/]*[@-~]").expect("invalid regex");
148    re.replace_all(s, "").to_string()
149}
150
151fn parse_usage_output(output: &str) -> Result<ApiResponse> {
152    let pct_re = Regex::new(r"(\d{1,3}(?:\.\d+)?)\s*%").expect("invalid regex");
153    let resets_re = Regex::new(r"(?i)resets?\s+(.*?)(?:\n|$)").expect("invalid regex");
154
155    let mut five_hour: Option<ApiWindow> = None;
156    let mut seven_day: Option<ApiWindow> = None;
157    let mut seven_day_opus: Option<ApiWindow> = None;
158
159    let sections = [
160        ("Current session", &mut five_hour),
161        ("Current week (all models)", &mut seven_day),
162        ("Current week (Opus", &mut seven_day_opus),
163    ];
164
165    for (label, target) in sections {
166        if let Some(pos) = output.find(label) {
167            let section = &output[pos..];
168            let section_end = section[label.len()..]
169                .find("Current")
170                .map(|p| p + label.len())
171                .unwrap_or(section.len());
172            let section_text = &section[..section_end];
173
174            if let Some(caps) = pct_re.captures(section_text) {
175                let raw_pct: f64 = caps[1].parse().unwrap_or(0.0);
176
177                let lower = section_text.to_lowercase();
178                let utilization = if lower.contains("left") || lower.contains("remaining") {
179                    100.0 - raw_pct
180                } else {
181                    raw_pct
182                };
183
184                let resets_at = resets_re
185                    .captures(section_text)
186                    .map(|c| c[1].trim().to_string());
187
188                *target = Some(ApiWindow {
189                    utilization: Some(utilization),
190                    resets_at,
191                });
192            }
193        }
194    }
195
196    if five_hour.is_none() && seven_day.is_none() && seven_day_opus.is_none() {
197        return Err(anyhow!("could not parse usage from CLI output:\n{output}"));
198    }
199
200    Ok(ApiResponse {
201        five_hour,
202        seven_day,
203        seven_day_sonnet: None,
204        seven_day_opus,
205        extra_usage: None,
206    })
207}
claude_code_stats/source/cli_probe.rs

claude_code_stats/source/
cli_probe.rs
