Skip to main content

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