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

ai_usagebar/context/
mod.rs
