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

context_bar_core/
usage_signal.rs
