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 recent_sessions: Vec<SessionRecord>,
142    #[serde(default)]
143    pub active_sessions: Vec<ActiveSession>,
144    #[serde(default)]
145    pub session_5h_resets_at: Option<String>,
146    #[serde(default)]
147    pub week_7d_resets_at: Option<String>,
148}
149
150#[derive(Clone, Debug, Default, Serialize, Deserialize)]
151pub struct TimeBucket {
152    #[serde(default, alias = "week", alias = "month")]
153    pub date: String,
154    #[serde(default)]
155    pub tokens: u64,
156    #[serde(default)]
157    pub sessions: u64,
158    // Token-category split + estimated USD cost (cost view).
159    #[serde(default)]
160    pub input: u64,
161    #[serde(default)]
162    pub output: u64,
163    #[serde(default)]
164    pub cache_creation: u64,
165    #[serde(default)]
166    pub cache_read: u64,
167    #[serde(default)]
168    pub cost: f64,
169}
170
171#[derive(Clone, Debug, Default, Serialize, Deserialize)]
172pub struct NamedBucket {
173    #[serde(default, alias = "project")]
174    pub model: String,
175    #[serde(default)]
176    pub tokens: u64,
177    #[serde(default)]
178    pub sessions: u64,
179    #[serde(default)]
180    pub input: u64,
181    #[serde(default)]
182    pub output: u64,
183    #[serde(default)]
184    pub cache_creation: u64,
185    #[serde(default)]
186    pub cache_read: u64,
187    #[serde(default)]
188    pub cost: f64,
189}
190
191#[derive(Clone, Debug, Default, Serialize, Deserialize)]
192pub struct SessionRecord {
193    #[serde(default)]
194    pub id: String,
195    #[serde(default)]
196    pub started_at: String,
197    #[serde(default)]
198    pub ended_at: String,
199    #[serde(default)]
200    pub duration_minutes: f64,
201    #[serde(default)]
202    pub tokens: u64,
203    #[serde(default)]
204    pub model: String,
205    #[serde(default)]
206    pub project: String,
207    #[serde(default)]
208    pub input: u64,
209    #[serde(default)]
210    pub output: u64,
211    #[serde(default)]
212    pub cache_creation: u64,
213    #[serde(default)]
214    pub cache_read: u64,
215    #[serde(default)]
216    pub cost: f64,
217}
218
219/// One (day × project) row — the `better-ccusage daily --instances` cross-tab.
220#[derive(Clone, Debug, Default, Serialize, Deserialize)]
221pub struct DailyInstance {
222    #[serde(default)]
223    pub date: String,
224    #[serde(default)]
225    pub project: String,
226    #[serde(default)]
227    pub models: Vec<String>,
228    #[serde(default)]
229    pub tokens: u64,
230    #[serde(default)]
231    pub sessions: u64,
232    #[serde(default)]
233    pub input: u64,
234    #[serde(default)]
235    pub output: u64,
236    #[serde(default)]
237    pub cache_creation: u64,
238    #[serde(default)]
239    pub cache_read: u64,
240    #[serde(default)]
241    pub cost: f64,
242}
243
244#[derive(Clone, Debug, Default, Serialize, Deserialize)]
245pub struct ToolSummary {
246    #[serde(default)]
247    pub name: String,
248    #[serde(default)]
249    pub sessions_7d: u64,
250    #[serde(default)]
251    pub sessions_today: u64,
252    #[serde(default)]
253    pub tokens_7d: u64,
254    #[serde(default)]
255    pub tokens_today: u64,
256    #[serde(default)]
257    pub last_used: Option<String>,
258    #[serde(default)]
259    pub last_model: Option<String>,
260}
261
262/// Subscription account read from `~/.claude/auth-*.json`.
263#[derive(Clone, Debug, Default, Serialize, Deserialize)]
264pub struct AccountInfo {
265    /// Filename stem, e.g. "hasan" from auth-hasan.json.
266    pub name: String,
267    /// "pro", "max", "free", etc.
268    pub subscription_type: String,
269    /// Raw tier string from the auth file.
270    pub rate_limit_tier: String,
271    /// Rolling 5-hour message limit (0 = unknown).
272    pub limit_5h_messages: u32,
273    /// Rolling 7-day message limit (0 = unknown).
274    pub limit_7d_messages: u32,
275    /// Whether this is the currently active account (matched via keychain).
276    pub is_active: bool,
277}
278
279#[cfg(not(target_arch = "wasm32"))]
280impl AccountInfo {
281    fn from_tier(name: String, subscription_type: String, rate_limit_tier: String) -> Self {
282        let (limit_5h_messages, limit_7d_messages) = match rate_limit_tier.as_str() {
283            t if t.contains("max_20x") => (900, 4500),
284            t if t.contains("max_5x") => (225, 1125),
285            t if t.contains("max") => (225, 1125),
286            _ => (45, 225),
287        };
288        Self { name, subscription_type, rate_limit_tier, limit_5h_messages, limit_7d_messages, is_active: false }
289    }
290}
291
292#[derive(Clone, Debug, Default, Serialize, Deserialize)]
293pub struct UsageSnapshot {
294    #[serde(default)]
295    pub claude: AgentUsage,
296    #[serde(default)]
297    pub codex: AgentUsage,
298    #[serde(default)]
299    pub others: Vec<ToolSummary>,
300    #[serde(default)]
301    pub accounts: Vec<AccountInfo>,
302    #[serde(default)]
303    pub collected_at: Option<String>,
304    #[serde(default)]
305    pub source: String,
306    /// Where the cost rate table came from: "live", "cache", or "fallback".
307    #[serde(default)]
308    pub pricing_source: Option<String>,
309    /// Always true: costs are API-equivalent estimates, not billed amounts.
310    #[serde(default)]
311    pub pricing_is_estimate: bool,
312}
313
314impl UsageSnapshot {
315    pub fn unavailable(reason: impl Into<String>) -> Self {
316        Self {
317            source: reason.into(),
318            ..Default::default()
319        }
320    }
321}
322
323#[cfg(target_arch = "wasm32")]
324const SCRIPT: &str = include_str!("usage_signal.py");
325
326#[cfg(target_arch = "wasm32")]
327pub fn collect(worktree: &zed::Worktree) -> UsageSnapshot {
328    let Some(python) = worktree
329        .which("python3")
330        .or_else(|| worktree.which("python"))
331    else {
332        return UsageSnapshot::unavailable("python3 not found on PATH");
333    };
334
335    let mut command = Command::new(python);
336    command = command.arg("-c").arg(SCRIPT);
337    command = command.envs(worktree.shell_env());
338
339    let output = match command.output() {
340        Ok(value) => value,
341        Err(error) => {
342            return UsageSnapshot::unavailable(format!("python spawn failed: {error}"));
343        }
344    };
345
346    if output.status != Some(0) {
347        let stderr = String::from_utf8_lossy(&output.stderr);
348        return UsageSnapshot::unavailable(format!(
349            "usage_signal.py exited with status {:?}: {}",
350            output.status,
351            stderr.trim()
352        ));
353    }
354
355    match serde_json::from_slice::<UsageSnapshot>(&output.stdout) {
356        Ok(snapshot) => snapshot,
357        Err(error) => UsageSnapshot::unavailable(format!("usage parse failed: {error}")),
358    }
359}
360
361/// Pure-Rust native snapshot — the engine, no `python3`. Reads `~/.claude` +
362/// `~/.codex` transcripts directly, prices with the live/cached LiteLLM table,
363/// applies the statusline + usage-API + Codex rate-limit overlays, and probes
364/// other AI tools. `accounts` is filled in by [`collect_native`].
365#[cfg(not(target_arch = "wasm32"))]
366fn collect_rust() -> UsageSnapshot {
367    use crate::aggregate::iso_utc;
368
369    let home = match std::env::var("HOME") {
370        Ok(h) => std::path::PathBuf::from(h),
371        Err(_) => return UsageSnapshot::unavailable("HOME not set"),
372    };
373    let now = std::time::SystemTime::now()
374        .duration_since(std::time::UNIX_EPOCH)
375        .map(|d| d.as_secs() as f64)
376        .unwrap_or(0.0);
377
378    let (table, pricing_source) = crate::pricing::load_pricing();
379    let claude = crate::collect::collect_claude_enriched(&home, now, &table);
380    let codex = crate::collect::collect_codex_enriched(&home, now, &table);
381    let others = crate::others::collect_others(&home, now);
382
383    UsageSnapshot {
384        claude,
385        codex,
386        others,
387        accounts: Vec::new(),
388        collected_at: Some(iso_utc(now)),
389        source: "rust".to_string(),
390        pricing_source: Some(pricing_source),
391        pricing_is_estimate: true,
392    }
393}
394
395/// Where we cache the full Python-emitted snapshot. Distinct from
396/// `usage_api_cache.json` (which Python uses for upstream API responses).
397#[cfg(not(target_arch = "wasm32"))]
398fn snapshot_cache_path() -> Option<std::path::PathBuf> {
399    let home = std::env::var("HOME").ok()?;
400    Some(std::path::PathBuf::from(home).join(".context-bar").join("usage.cache.json"))
401}
402
403/// TTL for the Rust-side snapshot cache. Matches the Python `CACHE_TTL_OK`
404/// upstream API window so we never spawn Python more often than the data
405/// changes anyway.
406#[cfg(not(target_arch = "wasm32"))]
407const SNAPSHOT_CACHE_TTL_SECS: u64 = 300;
408
409#[cfg(not(target_arch = "wasm32"))]
410fn load_snapshot_cache() -> Option<UsageSnapshot> {
411    use std::time::{SystemTime, UNIX_EPOCH};
412    let path = snapshot_cache_path()?;
413    let meta = std::fs::metadata(&path).ok()?;
414    let modified = meta.modified().ok()?;
415    let age = SystemTime::now().duration_since(modified).ok()?;
416    if age.as_secs() > SNAPSHOT_CACHE_TTL_SECS {
417        return None;
418    }
419    // Avoid stale-clock surprise: if file timestamp is far in the future, bail.
420    if modified.duration_since(UNIX_EPOCH).ok()?.as_secs() == 0 {
421        return None;
422    }
423    // Active session writes append to a .jsonl in place — file mtime advances,
424    // parent dir mtime does not. Drop the cache when any transcript is newer
425    // so mid-stream assistant turns reach context.json without a 300s lag.
426    if transcript_newer_than(modified) {
427        return None;
428    }
429    let bytes = std::fs::read(&path).ok()?;
430    serde_json::from_slice::<UsageSnapshot>(&bytes).ok()
431}
432
433#[cfg(not(target_arch = "wasm32"))]
434fn transcript_newer_than(threshold: std::time::SystemTime) -> bool {
435    let Ok(home) = std::env::var("HOME") else { return false };
436    let roots = [
437        std::path::PathBuf::from(&home).join(".claude").join("projects"),
438        std::path::PathBuf::from(&home).join(".codex").join("sessions"),
439    ];
440    for root in &roots {
441        if jsonl_newer_in_dir(root, threshold, 0) {
442            return true;
443        }
444    }
445    false
446}
447
448#[cfg(not(target_arch = "wasm32"))]
449fn jsonl_newer_in_dir(dir: &std::path::Path, threshold: std::time::SystemTime, depth: usize) -> bool {
450    // Transcripts live at <root>/<project>/<session>.jsonl — 4 levels is plenty
451    // and prevents pathological recursion if symlinks slip past file_type checks.
452    if depth > 4 {
453        return false;
454    }
455    let Ok(entries) = std::fs::read_dir(dir) else { return false };
456    for entry in entries.flatten() {
457        let Ok(ft) = entry.file_type() else { continue };
458        if ft.is_symlink() {
459            continue;
460        }
461        let path = entry.path();
462        if ft.is_dir() {
463            if jsonl_newer_in_dir(&path, threshold, depth + 1) {
464                return true;
465            }
466        } else if ft.is_file()
467            && path.extension().and_then(|s| s.to_str()) == Some("jsonl")
468        {
469            if let Ok(meta) = entry.metadata() {
470                if let Ok(m) = meta.modified() {
471                    if m > threshold {
472                        return true;
473                    }
474                }
475            }
476        }
477    }
478    false
479}
480
481#[cfg(not(target_arch = "wasm32"))]
482fn save_snapshot_cache(snapshot: &UsageSnapshot) {
483    let Some(path) = snapshot_cache_path() else { return };
484    if let Some(parent) = path.parent() {
485        let _ = std::fs::create_dir_all(parent);
486    }
487    if let Ok(bytes) = serde_json::to_vec(snapshot) {
488        // Best-effort; cache miss on next tick is acceptable.
489        let _ = std::fs::write(&path, bytes);
490    }
491}
492
493#[cfg(not(target_arch = "wasm32"))]
494pub fn collect_native() -> UsageSnapshot {
495    // Fast path: reuse a fresh on-disk snapshot to avoid re-scanning every
496    // transcript on each daemon tick. Invalidated at 300s or when any
497    // transcript is newer (see load_snapshot_cache).
498    if let Some(mut cached) = load_snapshot_cache() {
499        cached.accounts = collect_accounts();
500        return cached;
501    }
502
503    let mut snapshot = collect_rust();
504    if snapshot.source == "rust" {
505        // Persist the heavy collection (accounts are cheap + host-specific, so
506        // they're re-read each call below rather than cached).
507        save_snapshot_cache(&snapshot);
508    }
509    snapshot.accounts = collect_accounts();
510    snapshot
511}
512
513/// Reads all `~/.claude/auth-*.json` files and returns one `AccountInfo` per file.
514/// Marks the active account by matching the token stored in macOS Keychain under
515/// service "Claude Code-credentials".
516#[cfg(not(target_arch = "wasm32"))]
517fn collect_accounts() -> Vec<AccountInfo> {
518    use std::fs;
519    use serde_json;
520
521    let home = match std::env::var("HOME") {
522        Ok(h) => h,
523        Err(_) => return vec![],
524    };
525    let claude_dir = std::path::PathBuf::from(&home).join(".claude");
526
527    let read_dir = match fs::read_dir(&claude_dir) {
528        Ok(d) => d,
529        Err(_) => return vec![],
530    };
531
532    let paths: Vec<_> = read_dir
533        .filter_map(|e| e.ok())
534        .map(|e| e.path())
535        .filter(|p| {
536            p.file_name()
537                .and_then(|n| n.to_str())
538                .map(|n| n.starts_with("auth-") && n.ends_with(".json"))
539                .unwrap_or(false)
540        })
541        .collect();
542
543    // Read the active token prefix from keychain.
544    let active_token_prefix = active_token_prefix_from_keychain();
545
546    let mut accounts = Vec::new();
547    for path in &paths {
548        let stem = path
549            .file_stem()
550            .and_then(|s| s.to_str())
551            .unwrap_or("")
552            .trim_start_matches("auth-")
553            .to_string();
554
555        let Ok(content) = fs::read_to_string(path) else { continue };
556        let Ok(val) = serde_json::from_str::<serde_json::Value>(&content) else { continue };
557
558        let oauth = &val["claudeAiOauth"];
559        let subscription_type = oauth["subscriptionType"]
560            .as_str()
561            .unwrap_or("unknown")
562            .to_string();
563        let rate_limit_tier = oauth["rateLimitTier"]
564            .as_str()
565            .unwrap_or("")
566            .to_string();
567        let file_token = oauth["accessToken"].as_str().unwrap_or("");
568
569        let mut info = AccountInfo::from_tier(stem, subscription_type, rate_limit_tier);
570        if let Some(ref prefix) = active_token_prefix {
571            if !file_token.is_empty() && file_token.starts_with(prefix.as_str()) {
572                info.is_active = true;
573            }
574        }
575        accounts.push(info);
576    }
577
578    accounts.sort_by(|a, b| a.name.cmp(&b.name));
579
580    // If exactly one account exists, treat it as active regardless.
581    if accounts.len() == 1 {
582        accounts[0].is_active = true;
583    }
584
585    accounts
586}
587
588/// Returns the first 40 chars of the access token stored in the macOS Keychain
589/// under service "Claude Code-credentials", or None if unavailable.
590#[cfg(not(target_arch = "wasm32"))]
591fn active_token_prefix_from_keychain() -> Option<String> {
592    let output = std::process::Command::new("security")
593        .args(["find-generic-password", "-s", "Claude Code-credentials", "-w"])
594        .output()
595        .ok()?;
596
597    if !output.status.success() {
598        return None;
599    }
600
601    let raw = String::from_utf8_lossy(&output.stdout).trim().to_string();
602    // The stored value is the full auth JSON — parse it.
603    if let Ok(val) = serde_json::from_str::<serde_json::Value>(&raw) {
604        let token = val["claudeAiOauth"]["accessToken"].as_str()?;
605        return Some(token[..token.len().min(40)].to_string());
606    }
607    // Fallback: stored value might itself be just the token string.
608    if raw.starts_with("sk-ant") {
609        return Some(raw[..raw.len().min(40)].to_string());
610    }
611    None
612}
context_bar_core/usage_signal.rs

context_bar_core/
usage_signal.rs
