tiletopia/src-tauri/src/usage.rs

//! Reads claude-code session transcripts to report each session's **current
//! context-window occupancy** for the per-pane context indicator.
//!
//! claude writes one JSONL transcript per session at
//! `~/.claude/projects/<mangled-cwd>/<sessionId>.jsonl`. Every assistant line
//! carries `cwd`, `message.model`, and `message.usage`. The size of the prompt
//! sent on the most recent turn — `input_tokens + cache_read_input_tokens +
//! cache_creation_input_tokens` of the LAST assistant line — is a good proxy
//! for "how full is this session's context window right now".
//!
//! Windows-only: the transcripts live inside each WSL distro, reached via the
//! `\\wsl.localhost\<distro>\…` 9p share. Returns empty on non-Windows.

use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};

use parking_lot::Mutex;
use serde::Serialize;

use crate::pty::quiet_command;

/// Ignore sessions older than this, and cap the number returned — bounds the
/// scan cost on machines with a large transcript history.
const MAX_AGE_MS: i64 = 30 * 24 * 60 * 60 * 1000;
const MAX_SESSIONS: usize = 50;

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionContext {
    pub session_id: String,
    pub cwd: String,
    pub distro: String,
    pub last_active_ms: i64,
    /// Prompt size of the last assistant turn (input + both cache buckets) —
    /// the current context-window occupancy.
    pub context_tokens: u64,
    pub model: String,
}

/// Parsed-file cache entry, validated by (size, mtime) so we only re-parse the
/// one transcript that actually grew between polls.
struct CachedFile {
    size: u64,
    mtime_ms: i64,
    cwd: String,
    context_tokens: u64,
    model: String,
}

#[derive(Default)]
pub struct UsageCache {
    files: Mutex<HashMap<PathBuf, CachedFile>>,
    /// distro -> resolved `$HOME` (one wsl.exe probe per distro per process).
    homes: Mutex<HashMap<String, String>>,
}

/// Read each recent session's current context occupancy across the given WSL
/// distros (the distinct distros of currently-open WSL panes). Newest first,
/// capped to MAX_SESSIONS.
#[tauri::command]
pub async fn get_pane_context(
    distros: Vec<String>,
    cache: tauri::State<'_, UsageCache>,
) -> Result<Vec<SessionContext>, String> {
    if !cfg!(windows) {
        return Ok(Vec::new());
    }
    let cache = cache.inner();
    let mut out: Vec<SessionContext> = Vec::new();
    let mut seen = HashSet::new();
    for distro in distros.into_iter().filter(|d| !d.is_empty() && seen.insert(d.clone())) {
        match collect_distro(&distro, cache) {
            Ok(mut v) => out.append(&mut v),
            Err(e) => tracing::warn!("context scan for distro {distro} failed: {e}"),
        }
    }
    out.sort_by(|a, b| b.last_active_ms.cmp(&a.last_active_ms));
    out.truncate(MAX_SESSIONS);
    Ok(out)
}

fn collect_distro(distro: &str, cache: &UsageCache) -> Result<Vec<SessionContext>, String> {
    let home = resolve_home(distro, cache)?;
    let projects = projects_dir(distro, &home)
        .ok_or_else(|| format!("no ~/.claude/projects reachable for {distro}"))?;

    // Gather candidate transcripts (path, mtime), newest first.
    let now = now_ms();
    let mut candidates: Vec<(PathBuf, i64)> = Vec::new();
    for proj in std::fs::read_dir(&projects).map_err(|e| e.to_string())?.flatten() {
        let proj_path = proj.path();
        if !proj_path.is_dir() {
            continue;
        }
        let inner = match std::fs::read_dir(&proj_path) {
            Ok(it) => it,
            Err(_) => continue,
        };
        for f in inner.flatten() {
            let p = f.path();
            if p.extension().and_then(|e| e.to_str()) != Some("jsonl") {
                continue;
            }
            let mtime = std::fs::metadata(&p)
                .ok()
                .and_then(|m| m.modified().ok())
                .and_then(sys_to_ms)
                .unwrap_or(0);
            if now - mtime > MAX_AGE_MS {
                continue;
            }
            candidates.push((p, mtime));
        }
    }
    candidates.sort_by(|a, b| b.1.cmp(&a.1));
    candidates.truncate(MAX_SESSIONS);

    let mut out = Vec::new();
    for (path, mtime) in candidates {
        let (cwd, context_tokens, model) = match parse_or_cache(&path, cache) {
            Ok(v) => v,
            Err(e) => {
                tracing::debug!("skip transcript {}: {e}", path.display());
                continue;
            }
        };
        let session_id = path
            .file_stem()
            .map(|s| s.to_string_lossy().into_owned())
            .unwrap_or_default();
        out.push(SessionContext {
            session_id,
            cwd,
            distro: distro.to_string(),
            last_active_ms: mtime,
            context_tokens,
            model,
        });
    }
    Ok(out)
}

/// Probe `$HOME` inside the distro once and cache it. `sh -c` (not a login
/// shell) so rc-file output can't contaminate stdout.
fn resolve_home(distro: &str, cache: &UsageCache) -> Result<String, String> {
    if let Some(h) = cache.homes.lock().get(distro) {
        return Ok(h.clone());
    }
    let out = quiet_command("wsl.exe")
        .args(["-d", distro, "--", "sh", "-c", "printf %s \"$HOME\""])
        .output()
        .map_err(|e| format!("wsl.exe $HOME probe: {e}"))?;
    if !out.status.success() {
        return Err(format!("wsl.exe $HOME probe exited {:?}", out.status.code()));
    }
    let home = String::from_utf8_lossy(&out.stdout).trim().to_string();
    if home.is_empty() {
        return Err("empty $HOME".into());
    }
    cache.homes.lock().insert(distro.to_string(), home.clone());
    Ok(home)
}

/// `~/.claude/projects` as a Windows UNC path into the distro. Tries the newer
/// `\\wsl.localhost\` share first, then the legacy `\\wsl$\` alias.
fn projects_dir(distro: &str, home: &str) -> Option<PathBuf> {
    let home_rel = home.trim_start_matches('/');
    for prefix in [
        format!(r"\\wsl.localhost\{distro}"),
        format!(r"\\wsl$\{distro}"),
    ] {
        let p = PathBuf::from(prefix)
            .join(home_rel)
            .join(".claude")
            .join("projects");
        if p.is_dir() {
            return Some(p);
        }
    }
    None
}

fn parse_or_cache(
    path: &Path,
    cache: &UsageCache,
) -> Result<(String, u64, String), String> {
    let meta = std::fs::metadata(path).map_err(|e| e.to_string())?;
    let size = meta.len();
    let mtime = meta.modified().ok().and_then(sys_to_ms).unwrap_or(0);
    if let Some(c) = cache.files.lock().get(path) {
        if c.size == size && c.mtime_ms == mtime {
            return Ok((c.cwd.clone(), c.context_tokens, c.model.clone()));
        }
    }
    let (cwd, context_tokens, model) = parse_file(path)?;
    cache.files.lock().insert(
        path.to_path_buf(),
        CachedFile {
            size,
            mtime_ms: mtime,
            cwd: cwd.clone(),
            context_tokens,
            model: model.clone(),
        },
    );
    Ok((cwd, context_tokens, model))
}

/// Returns (cwd, context_tokens, model) where context_tokens is the prompt size
/// of the LAST assistant turn — the current context-window occupancy.
fn parse_file(path: &Path) -> Result<(String, u64, String), String> {
    let content = std::fs::read_to_string(path).map_err(|e| e.to_string())?;
    let mut cwd = String::new();
    let mut context_tokens = 0u64;
    let mut model = String::new();

    for line in content.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let v: serde_json::Value = match serde_json::from_str(line) {
            Ok(v) => v,
            Err(_) => continue, // tolerate a truncated final line / stray text
        };
        if cwd.is_empty() {
            if let Some(c) = v.get("cwd").and_then(|x| x.as_str()) {
                cwd = c.to_string();
            }
        }
        if v.get("type").and_then(|x| x.as_str()) != Some("assistant") {
            continue;
        }
        let msg = match v.get("message") {
            Some(m) => m,
            None => continue,
        };
        let usage = match msg.get("usage") {
            Some(u) => u,
            None => continue,
        };
        let tok = |k: &str| usage.get(k).and_then(|x| x.as_u64()).unwrap_or(0);
        // Overwrite each turn so we end up with the LAST assistant line's values.
        context_tokens = tok("input_tokens")
            + tok("cache_read_input_tokens")
            + tok("cache_creation_input_tokens");
        model = msg
            .get("model")
            .and_then(|x| x.as_str())
            .unwrap_or("unknown")
            .to_string();
    }

    Ok((cwd, context_tokens, model))
}

fn now_ms() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0)
}

fn sys_to_ms(t: SystemTime) -> Option<i64> {
    t.duration_since(UNIX_EPOCH).ok().map(|d| d.as_millis() as i64)
}