context_bar_core/
usage_signal.rs

1//! Cross-project agent usage signals.
2//!
3//! Reads Claude Code (`~/.claude/projects/**/*.jsonl`) and Codex CLI
4//! (`~/.codex/sessions/**/*.jsonl`) transcript files to summarize token usage
5//! over a rolling 5-hour session and 7-day week, plus the most recent turn's
6//! context-window utilization. Output drives the HUD surface.
7//!
8//! Implementation note: the heavy lifting lives in `usage_signal.py` invoked
9//! through `process:exec`. The Rust side validates the JSON envelope and
10//! returns a typed snapshot. On systems without `python3` (or where the
11//! script aborts) the snapshot is empty and the HUD degrades gracefully.
12
13use serde::{Deserialize, Serialize};
14
15#[cfg(target_arch = "wasm32")]
16use zed_extension_api::{self as zed, process::Command, serde_json};
17
18/// One in-flight session for an agent — a JSONL file whose last turn is
19/// within ACTIVE_WINDOW (currently 30 minutes). Multiple of these can be
20/// live at the same time when the user runs 3-5 sessions in parallel.
21#[derive(Clone, Debug, Default, Serialize, Deserialize)]
22pub struct ActiveSession {
23    #[serde(default)]
24    pub id: String,
25    #[serde(default)]
26    pub tokens: u64,
27    #[serde(default)]
28    pub cost: f64,
29    #[serde(default)]
30    pub started_at: Option<String>,
31    #[serde(default)]
32    pub last_turn_at: Option<String>,
33    #[serde(default)]
34    pub model: Option<String>,
35    #[serde(default)]
36    pub cwd: Option<String>,
37    #[serde(default)]
38    pub project: Option<String>,
39    #[serde(default)]
40    pub context_pct: Option<f64>,
41    #[serde(default)]
42    pub context_window: Option<u64>,
43    #[serde(default)]
44    pub last_input_tokens: u64,
45}
46
47#[derive(Clone, Debug, Default, Serialize, Deserialize)]
48pub struct AgentUsage {
49    #[serde(default)]
50    pub session_5h_tokens: u64,
51    #[serde(default)]
52    pub session_5h_percent: Option<f64>,
53    #[serde(default)]
54    pub week_7d_tokens: u64,
55    #[serde(default)]
56    pub week_7d_percent: Option<f64>,
57    #[serde(default)]
58    pub cache_read_tokens_5h: u64,
59    #[serde(default)]
60    pub cache_read_tokens_7d: u64,
61    #[serde(default)]
62    pub cache_read_tokens_30d: u64,
63    #[serde(default)]
64    pub active_session_tokens: u64,
65    #[serde(default)]
66    pub active_session_cost: f64,
67    #[serde(default)]
68    pub active_session_file: Option<String>,
69    #[serde(default)]
70    pub last_turn_input_tokens: u64,
71    #[serde(default)]
72    pub last_turn_output_tokens: u64,
73    #[serde(default)]
74    pub last_model: Option<String>,
75    #[serde(default)]
76    pub last_context_window: Option<u64>,
77    #[serde(default)]
78    pub last_context_pct: Option<f64>,
79    #[serde(default)]
80    pub last_turn_at: Option<String>,
81    #[serde(default)]
82    pub last_cwd: Option<String>,
83    #[serde(default)]
84    pub active_session_started_at: Option<String>,
85
86    // Aggregates for the detail page. All optional/empty in the no-data case.
87    #[serde(default)]
88    pub total_tokens_30d: u64,
89    #[serde(default)]
90    pub total_sessions_30d: u64,
91    #[serde(default)]
92    pub max_session_minutes: f64,
93    // Estimated API-equivalent cost (USD). Subscription users aren't billed
94    // per token — these mirror what the metered API would charge.
95    #[serde(default)]
96    pub cost_5h: f64,
97    #[serde(default)]
98    pub cost_7d: f64,
99    #[serde(default)]
100    pub cost_today: f64,
101    #[serde(default)]
102    pub total_cost_30d: f64,
103    #[serde(default)]
104    pub total_input_30d: u64,
105    #[serde(default)]
106    pub total_output_30d: u64,
107    /// Net USD prompt caching saved over the 30-day window vs paying full input.
108    #[serde(default)]
109    pub cache_savings_30d: f64,
110    #[serde(default)]
111    pub by_day: Vec<TimeBucket>,
112    #[serde(default)]
113    pub by_week: Vec<TimeBucket>,
114    #[serde(default)]
115    pub by_month: Vec<TimeBucket>,
116    #[serde(default)]
117    pub by_model: Vec<NamedBucket>,
118    #[serde(default)]
119    pub by_project: Vec<NamedBucket>,
120    #[serde(default)]
121    pub by_day_project: Vec<DailyInstance>,
122    #[serde(default)]
123    pub recent_sessions: Vec<SessionRecord>,
124    #[serde(default)]
125    pub active_sessions: Vec<ActiveSession>,
126    #[serde(default)]
127    pub session_5h_resets_at: Option<String>,
128    #[serde(default)]
129    pub week_7d_resets_at: Option<String>,
130}
131
132#[derive(Clone, Debug, Default, Serialize, Deserialize)]
133pub struct TimeBucket {
134    #[serde(default, alias = "week", alias = "month")]
135    pub date: String,
136    #[serde(default)]
137    pub tokens: u64,
138    #[serde(default)]
139    pub sessions: u64,
140    // Token-category split + estimated USD cost (cost view).
141    #[serde(default)]
142    pub input: u64,
143    #[serde(default)]
144    pub output: u64,
145    #[serde(default)]
146    pub cache_creation: u64,
147    #[serde(default)]
148    pub cache_read: u64,
149    #[serde(default)]
150    pub cost: f64,
151}
152
153#[derive(Clone, Debug, Default, Serialize, Deserialize)]
154pub struct NamedBucket {
155    #[serde(default, alias = "project")]
156    pub model: String,
157    #[serde(default)]
158    pub tokens: u64,
159    #[serde(default)]
160    pub sessions: u64,
161    #[serde(default)]
162    pub input: u64,
163    #[serde(default)]
164    pub output: u64,
165    #[serde(default)]
166    pub cache_creation: u64,
167    #[serde(default)]
168    pub cache_read: u64,
169    #[serde(default)]
170    pub cost: f64,
171}
172
173#[derive(Clone, Debug, Default, Serialize, Deserialize)]
174pub struct SessionRecord {
175    #[serde(default)]
176    pub id: String,
177    #[serde(default)]
178    pub started_at: String,
179    #[serde(default)]
180    pub ended_at: String,
181    #[serde(default)]
182    pub duration_minutes: f64,
183    #[serde(default)]
184    pub tokens: u64,
185    #[serde(default)]
186    pub model: String,
187    #[serde(default)]
188    pub project: String,
189    #[serde(default)]
190    pub input: u64,
191    #[serde(default)]
192    pub output: u64,
193    #[serde(default)]
194    pub cache_creation: u64,
195    #[serde(default)]
196    pub cache_read: u64,
197    #[serde(default)]
198    pub cost: f64,
199}
200
201/// One (day × project) row — the `better-ccusage daily --instances` cross-tab.
202#[derive(Clone, Debug, Default, Serialize, Deserialize)]
203pub struct DailyInstance {
204    #[serde(default)]
205    pub date: String,
206    #[serde(default)]
207    pub project: String,
208    #[serde(default)]
209    pub models: Vec<String>,
210    #[serde(default)]
211    pub tokens: u64,
212    #[serde(default)]
213    pub sessions: u64,
214    #[serde(default)]
215    pub input: u64,
216    #[serde(default)]
217    pub output: u64,
218    #[serde(default)]
219    pub cache_creation: u64,
220    #[serde(default)]
221    pub cache_read: u64,
222    #[serde(default)]
223    pub cost: f64,
224}
225
226#[derive(Clone, Debug, Default, Serialize, Deserialize)]
227pub struct ToolSummary {
228    #[serde(default)]
229    pub name: String,
230    #[serde(default)]
231    pub sessions_7d: u64,
232    #[serde(default)]
233    pub sessions_today: u64,
234    #[serde(default)]
235    pub tokens_7d: u64,
236    #[serde(default)]
237    pub tokens_today: u64,
238    #[serde(default)]
239    pub last_used: Option<String>,
240    #[serde(default)]
241    pub last_model: Option<String>,
242}
243
244/// Subscription account read from `~/.claude/auth-*.json`.
245#[derive(Clone, Debug, Default, Serialize, Deserialize)]
246pub struct AccountInfo {
247    /// Filename stem, e.g. "hasan" from auth-hasan.json.
248    pub name: String,
249    /// "pro", "max", "free", etc.
250    pub subscription_type: String,
251    /// Raw tier string from the auth file.
252    pub rate_limit_tier: String,
253    /// Rolling 5-hour message limit (0 = unknown).
254    pub limit_5h_messages: u32,
255    /// Rolling 7-day message limit (0 = unknown).
256    pub limit_7d_messages: u32,
257    /// Whether this is the currently active account (matched via keychain).
258    pub is_active: bool,
259}
260
261#[cfg(not(target_arch = "wasm32"))]
262impl AccountInfo {
263    fn from_tier(name: String, subscription_type: String, rate_limit_tier: String) -> Self {
264        let (limit_5h_messages, limit_7d_messages) = match rate_limit_tier.as_str() {
265            t if t.contains("max_20x") => (900, 4500),
266            t if t.contains("max_5x") => (225, 1125),
267            t if t.contains("max") => (225, 1125),
268            _ => (45, 225),
269        };
270        Self { name, subscription_type, rate_limit_tier, limit_5h_messages, limit_7d_messages, is_active: false }
271    }
272}
273
274#[derive(Clone, Debug, Default, Serialize, Deserialize)]
275pub struct UsageSnapshot {
276    #[serde(default)]
277    pub claude: AgentUsage,
278    #[serde(default)]
279    pub codex: AgentUsage,
280    #[serde(default)]
281    pub others: Vec<ToolSummary>,
282    #[serde(default)]
283    pub accounts: Vec<AccountInfo>,
284    #[serde(default)]
285    pub collected_at: Option<String>,
286    #[serde(default)]
287    pub source: String,
288    /// Where the cost rate table came from: "live", "cache", or "fallback".
289    #[serde(default)]
290    pub pricing_source: Option<String>,
291    /// Always true: costs are API-equivalent estimates, not billed amounts.
292    #[serde(default)]
293    pub pricing_is_estimate: bool,
294}
295
296impl UsageSnapshot {
297    pub fn unavailable(reason: impl Into<String>) -> Self {
298        Self {
299            source: reason.into(),
300            ..Default::default()
301        }
302    }
303}
304
305#[cfg(target_arch = "wasm32")]
306const SCRIPT: &str = include_str!("usage_signal.py");
307
308#[cfg(target_arch = "wasm32")]
309pub fn collect(worktree: &zed::Worktree) -> UsageSnapshot {
310    let Some(python) = worktree
311        .which("python3")
312        .or_else(|| worktree.which("python"))
313    else {
314        return UsageSnapshot::unavailable("python3 not found on PATH");
315    };
316
317    let mut command = Command::new(python);
318    command = command.arg("-c").arg(SCRIPT);
319    command = command.envs(worktree.shell_env());
320
321    let output = match command.output() {
322        Ok(value) => value,
323        Err(error) => {
324            return UsageSnapshot::unavailable(format!("python spawn failed: {error}"));
325        }
326    };
327
328    if output.status != Some(0) {
329        let stderr = String::from_utf8_lossy(&output.stderr);
330        return UsageSnapshot::unavailable(format!(
331            "usage_signal.py exited with status {:?}: {}",
332            output.status,
333            stderr.trim()
334        ));
335    }
336
337    match serde_json::from_slice::<UsageSnapshot>(&output.stdout) {
338        Ok(snapshot) => snapshot,
339        Err(error) => UsageSnapshot::unavailable(format!("usage parse failed: {error}")),
340    }
341}
342
343/// Pure-Rust native snapshot — the engine, no `python3`. Reads `~/.claude` +
344/// `~/.codex` transcripts directly, prices with the live/cached LiteLLM table,
345/// applies the statusline + usage-API + Codex rate-limit overlays, and probes
346/// other AI tools. `accounts` is filled in by [`collect_native`].
347#[cfg(not(target_arch = "wasm32"))]
348fn collect_rust() -> UsageSnapshot {
349    use crate::aggregate::iso_utc;
350
351    let home = match std::env::var("HOME") {
352        Ok(h) => std::path::PathBuf::from(h),
353        Err(_) => return UsageSnapshot::unavailable("HOME not set"),
354    };
355    let now = std::time::SystemTime::now()
356        .duration_since(std::time::UNIX_EPOCH)
357        .map(|d| d.as_secs() as f64)
358        .unwrap_or(0.0);
359
360    let (table, pricing_source) = crate::pricing::load_pricing();
361    let claude = crate::collect::collect_claude_enriched(&home, now, &table);
362    let codex = crate::collect::collect_codex_enriched(&home, now, &table);
363    let others = crate::others::collect_others(&home, now);
364
365    UsageSnapshot {
366        claude,
367        codex,
368        others,
369        accounts: Vec::new(),
370        collected_at: Some(iso_utc(now)),
371        source: "rust".to_string(),
372        pricing_source: Some(pricing_source),
373        pricing_is_estimate: true,
374    }
375}
376
377/// Where we cache the full Python-emitted snapshot. Distinct from
378/// `usage_api_cache.json` (which Python uses for upstream API responses).
379#[cfg(not(target_arch = "wasm32"))]
380fn snapshot_cache_path() -> Option<std::path::PathBuf> {
381    let home = std::env::var("HOME").ok()?;
382    Some(std::path::PathBuf::from(home).join(".context-bar").join("usage.cache.json"))
383}
384
385/// TTL for the Rust-side snapshot cache. Matches the Python `CACHE_TTL_OK`
386/// upstream API window so we never spawn Python more often than the data
387/// changes anyway.
388#[cfg(not(target_arch = "wasm32"))]
389const SNAPSHOT_CACHE_TTL_SECS: u64 = 300;
390
391#[cfg(not(target_arch = "wasm32"))]
392fn load_snapshot_cache() -> Option<UsageSnapshot> {
393    use std::time::{SystemTime, UNIX_EPOCH};
394    let path = snapshot_cache_path()?;
395    let meta = std::fs::metadata(&path).ok()?;
396    let modified = meta.modified().ok()?;
397    let age = SystemTime::now().duration_since(modified).ok()?;
398    if age.as_secs() > SNAPSHOT_CACHE_TTL_SECS {
399        return None;
400    }
401    // Avoid stale-clock surprise: if file timestamp is far in the future, bail.
402    if modified.duration_since(UNIX_EPOCH).ok()?.as_secs() == 0 {
403        return None;
404    }
405    // Active session writes append to a .jsonl in place — file mtime advances,
406    // parent dir mtime does not. Drop the cache when any transcript is newer
407    // so mid-stream assistant turns reach context.json without a 300s lag.
408    if transcript_newer_than(modified) {
409        return None;
410    }
411    let bytes = std::fs::read(&path).ok()?;
412    serde_json::from_slice::<UsageSnapshot>(&bytes).ok()
413}
414
415#[cfg(not(target_arch = "wasm32"))]
416fn transcript_newer_than(threshold: std::time::SystemTime) -> bool {
417    let Ok(home) = std::env::var("HOME") else { return false };
418    let roots = [
419        std::path::PathBuf::from(&home).join(".claude").join("projects"),
420        std::path::PathBuf::from(&home).join(".codex").join("sessions"),
421    ];
422    for root in &roots {
423        if jsonl_newer_in_dir(root, threshold, 0) {
424            return true;
425        }
426    }
427    false
428}
429
430#[cfg(not(target_arch = "wasm32"))]
431fn jsonl_newer_in_dir(dir: &std::path::Path, threshold: std::time::SystemTime, depth: usize) -> bool {
432    // Transcripts live at <root>/<project>/<session>.jsonl — 4 levels is plenty
433    // and prevents pathological recursion if symlinks slip past file_type checks.
434    if depth > 4 {
435        return false;
436    }
437    let Ok(entries) = std::fs::read_dir(dir) else { return false };
438    for entry in entries.flatten() {
439        let Ok(ft) = entry.file_type() else { continue };
440        if ft.is_symlink() {
441            continue;
442        }
443        let path = entry.path();
444        if ft.is_dir() {
445            if jsonl_newer_in_dir(&path, threshold, depth + 1) {
446                return true;
447            }
448        } else if ft.is_file()
449            && path.extension().and_then(|s| s.to_str()) == Some("jsonl")
450        {
451            if let Ok(meta) = entry.metadata() {
452                if let Ok(m) = meta.modified() {
453                    if m > threshold {
454                        return true;
455                    }
456                }
457            }
458        }
459    }
460    false
461}
462
463#[cfg(not(target_arch = "wasm32"))]
464fn save_snapshot_cache(snapshot: &UsageSnapshot) {
465    let Some(path) = snapshot_cache_path() else { return };
466    if let Some(parent) = path.parent() {
467        let _ = std::fs::create_dir_all(parent);
468    }
469    if let Ok(bytes) = serde_json::to_vec(snapshot) {
470        // Best-effort; cache miss on next tick is acceptable.
471        let _ = std::fs::write(&path, bytes);
472    }
473}
474
475#[cfg(not(target_arch = "wasm32"))]
476pub fn collect_native() -> UsageSnapshot {
477    // Fast path: reuse a fresh on-disk snapshot to avoid re-scanning every
478    // transcript on each daemon tick. Invalidated at 300s or when any
479    // transcript is newer (see load_snapshot_cache).
480    if let Some(mut cached) = load_snapshot_cache() {
481        cached.accounts = collect_accounts();
482        return cached;
483    }
484
485    let mut snapshot = collect_rust();
486    if snapshot.source == "rust" {
487        // Persist the heavy collection (accounts are cheap + host-specific, so
488        // they're re-read each call below rather than cached).
489        save_snapshot_cache(&snapshot);
490    }
491    snapshot.accounts = collect_accounts();
492    snapshot
493}
494
495/// Reads all `~/.claude/auth-*.json` files and returns one `AccountInfo` per file.
496/// Marks the active account by matching the token stored in macOS Keychain under
497/// service "Claude Code-credentials".
498#[cfg(not(target_arch = "wasm32"))]
499fn collect_accounts() -> Vec<AccountInfo> {
500    use std::fs;
501    use serde_json;
502
503    let home = match std::env::var("HOME") {
504        Ok(h) => h,
505        Err(_) => return vec![],
506    };
507    let claude_dir = std::path::PathBuf::from(&home).join(".claude");
508
509    let read_dir = match fs::read_dir(&claude_dir) {
510        Ok(d) => d,
511        Err(_) => return vec![],
512    };
513
514    let paths: Vec<_> = read_dir
515        .filter_map(|e| e.ok())
516        .map(|e| e.path())
517        .filter(|p| {
518            p.file_name()
519                .and_then(|n| n.to_str())
520                .map(|n| n.starts_with("auth-") && n.ends_with(".json"))
521                .unwrap_or(false)
522        })
523        .collect();
524
525    // Read the active token prefix from keychain.
526    let active_token_prefix = active_token_prefix_from_keychain();
527
528    let mut accounts = Vec::new();
529    for path in &paths {
530        let stem = path
531            .file_stem()
532            .and_then(|s| s.to_str())
533            .unwrap_or("")
534            .trim_start_matches("auth-")
535            .to_string();
536
537        let Ok(content) = fs::read_to_string(path) else { continue };
538        let Ok(val) = serde_json::from_str::<serde_json::Value>(&content) else { continue };
539
540        let oauth = &val["claudeAiOauth"];
541        let subscription_type = oauth["subscriptionType"]
542            .as_str()
543            .unwrap_or("unknown")
544            .to_string();
545        let rate_limit_tier = oauth["rateLimitTier"]
546            .as_str()
547            .unwrap_or("")
548            .to_string();
549        let file_token = oauth["accessToken"].as_str().unwrap_or("");
550
551        let mut info = AccountInfo::from_tier(stem, subscription_type, rate_limit_tier);
552        if let Some(ref prefix) = active_token_prefix {
553            if !file_token.is_empty() && file_token.starts_with(prefix.as_str()) {
554                info.is_active = true;
555            }
556        }
557        accounts.push(info);
558    }
559
560    accounts.sort_by(|a, b| a.name.cmp(&b.name));
561
562    // If exactly one account exists, treat it as active regardless.
563    if accounts.len() == 1 {
564        accounts[0].is_active = true;
565    }
566
567    accounts
568}
569
570/// Returns the first 40 chars of the access token stored in the macOS Keychain
571/// under service "Claude Code-credentials", or None if unavailable.
572#[cfg(not(target_arch = "wasm32"))]
573fn active_token_prefix_from_keychain() -> Option<String> {
574    let output = std::process::Command::new("security")
575        .args(["find-generic-password", "-s", "Claude Code-credentials", "-w"])
576        .output()
577        .ok()?;
578
579    if !output.status.success() {
580        return None;
581    }
582
583    let raw = String::from_utf8_lossy(&output.stdout).trim().to_string();
584    // The stored value is the full auth JSON — parse it.
585    if let Ok(val) = serde_json::from_str::<serde_json::Value>(&raw) {
586        let token = val["claudeAiOauth"]["accessToken"].as_str()?;
587        return Some(token[..token.len().min(40)].to_string());
588    }
589    // Fallback: stored value might itself be just the token string.
590    if raw.starts_with("sk-ant") {
591        return Some(raw[..raw.len().min(40)].to_string());
592    }
593    None
594}
context_bar_core/usage_signal.rs

context_bar_core/
usage_signal.rs
