ai_usagebar/context/
mod.rs

1//! Best-effort, local-only Claude Code context-window discovery.
2//!
3//! Claude Code transcripts are an undocumented JSONL surface, so this module
4//! deliberately treats every line as optional data. It reads only bounded
5//! tails, ignores unknown/corrupt records, never follows discovered symlinks,
6//! and reports unknown usage instead of fabricating zero. A compaction record
7//! after the latest usage invalidates that reading until another assistant
8//! response supplies the new context size.
9
10pub mod activity;
11
12use std::fs::{self, File};
13use std::io::{Read, Seek, SeekFrom};
14use std::path::{Path, PathBuf};
15use std::time::SystemTime;
16
17use chrono::{DateTime, Utc};
18use serde_json::Value;
19
20use crate::config::ContextConfig;
21use crate::error::{AppError, Result};
22
23/// The overlay is a recent-session monitor, not an unbounded history browser.
24pub const MAX_SESSIONS: usize = 100;
25const MAX_WALK_ENTRIES: usize = 10_000;
26const MAX_TAIL_BYTES: usize = 2 * 1024 * 1024;
27const MAX_LINE_BYTES: usize = 512 * 1024;
28const MAX_DISPLAY_CHARS: usize = 120;
29
30#[derive(Debug, Clone)]
31pub struct ContextScan {
32    pub sessions: Vec<ContextSession>,
33    /// JSONL candidates observed before the recent-session cap was applied.
34    pub discovered: usize,
35    /// Unreadable entries and malformed/oversized tail records ignored.
36    pub skipped: usize,
37    /// True when the directory-entry safety cap stopped traversal early.
38    pub walk_capped: bool,
39}
40
41#[derive(Debug, Clone)]
42pub struct ContextSession {
43    pub session_id: String,
44    pub title: Option<String>,
45    pub project: String,
46    pub model: Option<String>,
47    pub modified_at: DateTime<Utc>,
48    pub usage: ContextUsage,
49}
50
51impl ContextSession {
52    pub fn display_name(&self) -> String {
53        self.title
54            .clone()
55            .unwrap_or_else(|| format!("session {}", short_id(&self.session_id)))
56    }
57}
58
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub enum ContextUsage {
61    Available {
62        /// Matches Claude Code's input-only context percentage: fresh input +
63        /// cache creation + cache reads from the most recent API response.
64        input_tokens: u64,
65        window_tokens: Option<u64>,
66        percent: Option<u16>,
67    },
68    /// A compact boundary landed after the last assistant usage record.
69    Compacted,
70    /// No usable assistant usage was present in the bounded transcript tail.
71    Unknown,
72}
73
74#[derive(Debug)]
75struct Candidate {
76    path: PathBuf,
77    modified: SystemTime,
78}
79
80#[derive(Debug, Default)]
81struct Discovery {
82    candidates: Vec<Candidate>,
83    skipped: usize,
84    walk_capped: bool,
85}
86
87#[derive(Debug, Default)]
88struct TailData {
89    session_id: Option<String>,
90    title: Option<String>,
91    cwd: Option<String>,
92    model: Option<String>,
93    usage: TailUsage,
94    skipped: usize,
95}
96
97#[derive(Debug, Default)]
98enum TailUsage {
99    Tokens(u64),
100    Compacted,
101    #[default]
102    Unknown,
103}
104
105/// Resolve Claude Code's conventional local transcript directory without
106/// reading it. Kept separate so tests always inject their own root.
107pub fn default_projects_path() -> Result<PathBuf> {
108    Ok(crate::cache::home_dir()?.join(".claude").join("projects"))
109}
110
111/// Scan the most recently modified top-level session transcripts below
112/// `root`. Subagent sidechains are excluded, and discovered symlinks are never
113/// followed. The caller should run this blocking filesystem work on Tokio's
114/// blocking pool.
115pub fn scan_dir(root: &Path, config: &ContextConfig) -> Result<ContextScan> {
116    let mut discovery = discover(root)?;
117    discovery.candidates.sort_by(|a, b| {
118        b.modified
119            .cmp(&a.modified)
120            .then_with(|| a.path.cmp(&b.path))
121    });
122    let discovered = discovery.candidates.len();
123    discovery.candidates.truncate(MAX_SESSIONS);
124
125    let mut sessions = Vec::with_capacity(discovery.candidates.len());
126    for candidate in discovery.candidates {
127        match parse_candidate(root, candidate, config) {
128            Ok((session, skipped)) => {
129                discovery.skipped += skipped;
130                sessions.push(session);
131            }
132            Err(_) => discovery.skipped += 1,
133        }
134    }
135
136    Ok(ContextScan {
137        sessions,
138        discovered,
139        skipped: discovery.skipped,
140        walk_capped: discovery.walk_capped,
141    })
142}
143
144fn discover(root: &Path) -> Result<Discovery> {
145    // Opening the configured root is the one error that should reach the user;
146    // failures inside it are isolated to the affected entry.
147    fs::read_dir(root).map_err(|e| AppError::io_at(root, e))?;
148
149    let mut result = Discovery::default();
150    let mut stack = vec![root.to_path_buf()];
151    let mut visited = 0usize;
152
153    'walk: while let Some(dir) = stack.pop() {
154        let entries = match fs::read_dir(&dir) {
155            Ok(entries) => entries,
156            Err(_) => {
157                result.skipped += 1;
158                continue;
159            }
160        };
161        for entry in entries {
162            if visited >= MAX_WALK_ENTRIES {
163                result.walk_capped = true;
164                break 'walk;
165            }
166            visited += 1;
167
168            let entry = match entry {
169                Ok(entry) => entry,
170                Err(_) => {
171                    result.skipped += 1;
172                    continue;
173                }
174            };
175            let file_type = match entry.file_type() {
176                Ok(kind) => kind,
177                Err(_) => {
178                    result.skipped += 1;
179                    continue;
180                }
181            };
182            if file_type.is_symlink() {
183                continue;
184            }
185            if file_type.is_dir() {
186                if entry.file_name() != "subagents" {
187                    stack.push(entry.path());
188                }
189                continue;
190            }
191            if !file_type.is_file()
192                || entry.path().extension().and_then(|ext| ext.to_str()) != Some("jsonl")
193            {
194                continue;
195            }
196            let metadata = match entry.metadata() {
197                Ok(metadata) => metadata,
198                Err(_) => {
199                    result.skipped += 1;
200                    continue;
201                }
202            };
203            result.candidates.push(Candidate {
204                path: entry.path(),
205                modified: metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH),
206            });
207        }
208    }
209    Ok(result)
210}
211
212fn parse_candidate(
213    root: &Path,
214    candidate: Candidate,
215    config: &ContextConfig,
216) -> Result<(ContextSession, usize)> {
217    let tail = read_tail(&candidate.path)?;
218    let parsed = parse_tail(&tail);
219    let file_id = candidate
220        .path
221        .file_stem()
222        .and_then(|stem| stem.to_str())
223        .unwrap_or("unknown-session");
224    let session_id = parsed
225        .session_id
226        .filter(|id| !id.is_empty())
227        .unwrap_or_else(|| sanitize_display(file_id));
228    let project = parsed
229        .cwd
230        .as_deref()
231        .and_then(project_from_cwd)
232        .or_else(|| project_from_path(root, &candidate.path))
233        .unwrap_or_else(|| "unknown project".into());
234    let window_tokens = config.window_tokens_for(parsed.model.as_deref());
235    let usage = match parsed.usage {
236        TailUsage::Tokens(input_tokens) => ContextUsage::Available {
237            input_tokens,
238            window_tokens,
239            percent: window_tokens.map(|window| percent(input_tokens, window)),
240        },
241        TailUsage::Compacted => ContextUsage::Compacted,
242        TailUsage::Unknown => ContextUsage::Unknown,
243    };
244
245    Ok((
246        ContextSession {
247            session_id,
248            title: parsed.title.filter(|title| !title.is_empty()),
249            project,
250            model: parsed.model,
251            modified_at: DateTime::<Utc>::from(candidate.modified),
252            usage,
253        },
254        parsed.skipped,
255    ))
256}
257
258fn read_tail(path: &Path) -> Result<Vec<u8>> {
259    let mut file = File::open(path).map_err(|e| AppError::io_at(path, e))?;
260    let len = file.metadata().map_err(|e| AppError::io_at(path, e))?.len();
261    let read_len = len.min(MAX_TAIL_BYTES as u64);
262    let start = len.saturating_sub(read_len);
263    file.seek(SeekFrom::Start(start))
264        .map_err(|e| AppError::io_at(path, e))?;
265    let mut bytes = vec![0; read_len as usize];
266    file.read_exact(&mut bytes)
267        .map_err(|e| AppError::io_at(path, e))?;
268
269    // The first bytes of a bounded tail may be the middle of a JSON record.
270    // Drop that fragment rather than feeding it to the tolerant parser.
271    if start > 0 {
272        if let Some(newline) = bytes.iter().position(|byte| *byte == b'\n') {
273            bytes.drain(..=newline);
274        } else {
275            bytes.clear();
276        }
277    }
278    Ok(bytes)
279}
280
281fn parse_tail(bytes: &[u8]) -> TailData {
282    let mut out = TailData::default();
283    for raw in bytes.split(|byte| *byte == b'\n') {
284        let raw = raw.strip_suffix(b"\r").unwrap_or(raw);
285        if raw.is_empty() {
286            continue;
287        }
288        if raw.len() > MAX_LINE_BYTES {
289            out.skipped += 1;
290            continue;
291        }
292        let value: Value = match serde_json::from_slice(raw) {
293            Ok(value) => value,
294            Err(_) => {
295                out.skipped += 1;
296                continue;
297            }
298        };
299
300        capture_string(&value, "sessionId", &mut out.session_id);
301        capture_string(&value, "cwd", &mut out.cwd);
302        match value.get("type").and_then(Value::as_str) {
303            Some("custom-title") => {
304                capture_string(&value, "customTitle", &mut out.title);
305                capture_string(&value, "title", &mut out.title);
306            }
307            Some("assistant") => {
308                if let Some(model) = value.pointer("/message/model").and_then(Value::as_str) {
309                    out.model = Some(sanitize_display(model));
310                }
311                if let Some(tokens) = input_tokens(&value) {
312                    out.usage = TailUsage::Tokens(tokens);
313                }
314            }
315            Some("system")
316                if value.get("subtype").and_then(Value::as_str) == Some("compact_boundary") =>
317            {
318                out.usage = TailUsage::Compacted;
319            }
320            _ => {}
321        }
322    }
323    out
324}
325
326fn capture_string(value: &Value, key: &str, target: &mut Option<String>) {
327    if let Some(value) = value.get(key).and_then(Value::as_str) {
328        let value = sanitize_display(value);
329        if !value.is_empty() {
330            *target = Some(value);
331        }
332    }
333}
334
335fn input_tokens(value: &Value) -> Option<u64> {
336    let usage = value.pointer("/message/usage")?;
337    let input = usage.get("input_tokens")?.as_u64()?;
338    let cache_creation = optional_u64(usage, "cache_creation_input_tokens")?;
339    let cache_read = optional_u64(usage, "cache_read_input_tokens")?;
340    input.checked_add(cache_creation)?.checked_add(cache_read)
341}
342
343fn optional_u64(value: &Value, key: &str) -> Option<u64> {
344    match value.get(key) {
345        None | Some(Value::Null) => Some(0),
346        Some(value) => value.as_u64(),
347    }
348}
349
350fn percent(input_tokens: u64, window_tokens: u64) -> u16 {
351    debug_assert!(window_tokens > 0);
352    let value = (u128::from(input_tokens) * 100) / u128::from(window_tokens);
353    value.min(u128::from(u16::MAX)) as u16
354}
355
356fn project_from_cwd(cwd: &str) -> Option<String> {
357    Path::new(cwd)
358        .file_name()
359        .map(|name| sanitize_display(&name.to_string_lossy()))
360        .filter(|name| !name.is_empty())
361}
362
363fn project_from_path(root: &Path, transcript: &Path) -> Option<String> {
364    let relative = transcript.strip_prefix(root).ok()?;
365    let first = relative.components().next()?.as_os_str();
366    let project = sanitize_display(&first.to_string_lossy());
367    (!project.is_empty()).then_some(project)
368}
369
370fn short_id(id: &str) -> String {
371    id.chars().take(8).collect()
372}
373
374/// Strip terminal control characters and cap untrusted transcript labels
375/// before they enter ratatui's backing buffer.
376pub fn sanitize_display(value: &str) -> String {
377    value
378        .chars()
379        .filter_map(|ch| {
380            if ch.is_control() {
381                if ch.is_whitespace() { Some(' ') } else { None }
382            } else {
383                Some(ch)
384            }
385        })
386        .take(MAX_DISPLAY_CHARS)
387        .collect::<String>()
388        .trim()
389        .to_string()
390}
391
392#[cfg(test)]
393mod tests {
394    use std::io::Write;
395
396    use serde_json::json;
397    use tempfile::TempDir;
398
399    use super::*;
400
401    fn assistant(
402        session: &str,
403        cwd: &str,
404        model: &str,
405        input: u64,
406        create: u64,
407        read: u64,
408    ) -> String {
409        json!({
410            "type": "assistant",
411            "sessionId": session,
412            "cwd": cwd,
413            "message": {
414                "model": model,
415                "usage": {
416                    "input_tokens": input,
417                    "cache_creation_input_tokens": create,
418                    "cache_read_input_tokens": read
419                }
420            }
421        })
422        .to_string()
423    }
424
425    fn write_session(root: &Path, project: &str, name: &str, lines: &[String]) -> PathBuf {
426        let dir = root.join(project);
427        fs::create_dir_all(&dir).unwrap();
428        let path = dir.join(format!("{name}.jsonl"));
429        let mut file = File::create(&path).unwrap();
430        for line in lines {
431            writeln!(file, "{line}").unwrap();
432        }
433        path
434    }
435
436    #[test]
437    fn latest_assistant_usage_matches_claude_input_only_formula() {
438        let dir = TempDir::new().unwrap();
439        write_session(
440            dir.path(),
441            "-work-project",
442            "abc-123",
443            &[
444                assistant("abc-123", "/work/project", "claude-test", 1, 2, 3),
445                assistant("abc-123", "/work/project", "claude-test", 10, 20, 30),
446            ],
447        );
448        let mut config = ContextConfig {
449            context_window_tokens: Some(200),
450            ..ContextConfig::default()
451        };
452        config
453            .model_context_window_tokens
454            .insert("claude-test".into(), 100);
455
456        let scan = scan_dir(dir.path(), &config).unwrap();
457        assert_eq!(scan.sessions.len(), 1);
458        let session = &scan.sessions[0];
459        assert_eq!(session.session_id, "abc-123");
460        assert_eq!(session.project, "project");
461        assert_eq!(session.model.as_deref(), Some("claude-test"));
462        assert_eq!(
463            session.usage,
464            ContextUsage::Available {
465                input_tokens: 60,
466                window_tokens: Some(100),
467                percent: Some(60),
468            }
469        );
470    }
471
472    #[test]
473    fn compact_boundary_invalidates_the_previous_reading() {
474        let dir = TempDir::new().unwrap();
475        let compact = json!({"type": "system", "subtype": "compact_boundary"}).to_string();
476        write_session(
477            dir.path(),
478            "project",
479            "session",
480            &[
481                assistant("session", "/work/project", "claude-test", 90, 0, 0),
482                compact,
483            ],
484        );
485
486        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
487        assert_eq!(scan.sessions[0].usage, ContextUsage::Compacted);
488    }
489
490    #[test]
491    fn assistant_after_compaction_supplies_the_new_reading() {
492        let dir = TempDir::new().unwrap();
493        let compact = json!({"type": "system", "subtype": "compact_boundary"}).to_string();
494        write_session(
495            dir.path(),
496            "project",
497            "session",
498            &[
499                assistant("session", "/work/project", "claude-test", 90, 0, 0),
500                compact,
501                assistant("session", "/work/project", "claude-test", 12, 0, 0),
502            ],
503        );
504
505        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
506        assert_eq!(
507            scan.sessions[0].usage,
508            ContextUsage::Available {
509                input_tokens: 12,
510                window_tokens: None,
511                percent: None,
512            }
513        );
514    }
515
516    #[test]
517    fn corrupt_or_truncated_records_do_not_hide_the_last_good_usage() {
518        let dir = TempDir::new().unwrap();
519        let path = write_session(
520            dir.path(),
521            "project",
522            "session",
523            &[assistant(
524                "session",
525                "/work/project",
526                "claude-test",
527                42,
528                0,
529                0,
530            )],
531        );
532        let mut file = fs::OpenOptions::new().append(true).open(path).unwrap();
533        write!(file, "{{\"type\":\"assistant\"").unwrap();
534
535        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
536        assert!(scan.skipped >= 1);
537        assert_eq!(
538            scan.sessions[0].usage,
539            ContextUsage::Available {
540                input_tokens: 42,
541                window_tokens: None,
542                percent: None,
543            }
544        );
545    }
546
547    #[test]
548    fn an_old_usage_record_outside_the_tail_cap_is_not_loaded_unboundedly() {
549        let dir = TempDir::new().unwrap();
550        let path = write_session(
551            dir.path(),
552            "project",
553            "session",
554            &[assistant(
555                "session",
556                "/work/project",
557                "claude-test",
558                42,
559                0,
560                0,
561            )],
562        );
563        let mut file = fs::OpenOptions::new().append(true).open(path).unwrap();
564        file.write_all(&vec![b'x'; MAX_TAIL_BYTES + 1024]).unwrap();
565        writeln!(file).unwrap();
566
567        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
568        assert_eq!(scan.sessions[0].usage, ContextUsage::Unknown);
569    }
570
571    #[test]
572    fn subagent_transcripts_are_excluded() {
573        let dir = TempDir::new().unwrap();
574        write_session(
575            dir.path(),
576            "project",
577            "main",
578            &[assistant("main", "/work/project", "claude-test", 1, 0, 0)],
579        );
580        write_session(
581            &dir.path().join("project"),
582            "subagents",
583            "agent",
584            &[assistant("agent", "/work/project", "claude-test", 99, 0, 0)],
585        );
586
587        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
588        assert_eq!(scan.sessions.len(), 1);
589        assert_eq!(scan.sessions[0].session_id, "main");
590    }
591
592    #[cfg(unix)]
593    #[test]
594    fn discovered_symlinks_are_not_followed() {
595        use std::os::unix::fs::symlink;
596
597        let dir = TempDir::new().unwrap();
598        let outside = TempDir::new().unwrap();
599        let outside_file = write_session(
600            outside.path(),
601            "project",
602            "outside",
603            &[assistant("outside", "/outside", "claude-test", 99, 0, 0)],
604        );
605        fs::create_dir_all(dir.path().join("project")).unwrap();
606        symlink(outside_file, dir.path().join("project/linked.jsonl")).unwrap();
607
608        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
609        assert!(scan.sessions.is_empty());
610    }
611
612    #[test]
613    fn custom_titles_and_paths_are_sanitized_before_display() {
614        let dir = TempDir::new().unwrap();
615        let title =
616            json!({"type": "custom-title", "customTitle": "build\u{1b}[2J\nrelease"}).to_string();
617        write_session(
618            dir.path(),
619            "project",
620            "session",
621            &[
622                title,
623                assistant("session", "/work/proj\nname", "claude-test", 1, 0, 0),
624            ],
625        );
626
627        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
628        let session = &scan.sessions[0];
629        assert_eq!(session.title.as_deref(), Some("build[2J release"));
630        assert_eq!(session.project, "proj name");
631    }
632
633    #[test]
634    fn invalid_usage_numbers_are_unknown_not_zero() {
635        let dir = TempDir::new().unwrap();
636        let invalid = json!({
637            "type": "assistant",
638            "message": {"usage": {"input_tokens": "12"}}
639        })
640        .to_string();
641        write_session(dir.path(), "project", "session", &[invalid]);
642
643        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
644        assert_eq!(scan.sessions[0].usage, ContextUsage::Unknown);
645    }
646
647    #[test]
648    fn recent_session_result_is_capped() {
649        let dir = TempDir::new().unwrap();
650        for i in 0..=MAX_SESSIONS {
651            write_session(
652                dir.path(),
653                "project",
654                &format!("session-{i}"),
655                &[assistant(
656                    &format!("session-{i}"),
657                    "/work/project",
658                    "claude-test",
659                    i as u64,
660                    0,
661                    0,
662                )],
663            );
664        }
665
666        let scan = scan_dir(dir.path(), &ContextConfig::default()).unwrap();
667        assert_eq!(scan.discovered, MAX_SESSIONS + 1);
668        assert_eq!(scan.sessions.len(), MAX_SESSIONS);
669    }
670
671    #[test]
672    fn missing_root_is_an_actionable_error() {
673        let dir = TempDir::new().unwrap();
674        let missing = dir.path().join("missing");
675        let error = scan_dir(&missing, &ContextConfig::default())
676            .unwrap_err()
677            .to_string();
678        assert!(error.contains(missing.to_string_lossy().as_ref()));
679    }
680}
ai_usagebar/context/mod.rs

ai_usagebar/context/
mod.rs
