lifeloop-cli 0.1.0

//! Gemini CLI lifecycle telemetry reader.
//!
//! Parses Gemini telemetry JSONL emitting `gemini_cli.api_response` and
//! `gemini_cli.chat_compression` events. Extracts prompt-token totals,
//! model name, the native context-window field, and a compaction
//! signal.

use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};

use serde_json::Value;

use super::{
    EnvAlias, PressureObservation, TelemetryError, TelemetryResult, TokenUsage, compute_pct,
    file_mtime_epoch_s, general_context_window, general_host_model, home_dir, is_recent,
    number_key, resolve_env_string, resolve_env_u64, string_key,
};

const ADAPTER_ID: &str = "gemini";

const GEMINI_TELEMETRY_PATH_ALIASES: &[EnvAlias] = &[
    EnvAlias {
        lifeloop: "LIFELOOP_GEMINI_TELEMETRY_PATH",
        ccd_compat: "CCD_GEMINI_TELEMETRY_PATH",
    },
    EnvAlias {
        lifeloop: "LIFELOOP_GEMINI_TELEMETRY_OUTFILE",
        ccd_compat: "GEMINI_TELEMETRY_OUTFILE",
    },
];

const GEMINI_MODEL_ALIASES: &[EnvAlias] = &[EnvAlias {
    lifeloop: "LIFELOOP_GEMINI_MODEL",
    ccd_compat: "CCD_GEMINI_MODEL",
}];

const GEMINI_CONTEXT_WINDOW_ALIASES: &[EnvAlias] = &[EnvAlias {
    lifeloop: "LIFELOOP_GEMINI_CONTEXT_WINDOW_TOKENS",
    ccd_compat: "CCD_GEMINI_CONTEXT_WINDOW_TOKENS",
}];

#[derive(Default)]
struct GeminiSessionSnapshot {
    latest_prompt_tokens: Option<u64>,
    input_tokens: u64,
    output_tokens: u64,
    cache_read_input_tokens: u64,
    context_window: Option<u64>,
    compacted: bool,
    model_name: Option<String>,
}

pub fn current(repo_root: &Path) -> TelemetryResult<Option<PressureObservation>> {
    let Some((path, observed_at_epoch_s)) = telemetry_observation(repo_root)? else {
        return Ok(None);
    };
    let file = File::open(&path).map_err(TelemetryError::from)?;
    let mut bytes = Vec::new();
    use std::io::Read;
    BufReader::new(file)
        .read_to_end(&mut bytes)
        .map_err(TelemetryError::from)?;
    parse_telemetry_log(&bytes, observed_at_epoch_s)
}

/// Parse a Gemini telemetry log byte slice.
pub fn parse_telemetry_log(
    bytes: &[u8],
    observed_at_epoch_s: u64,
) -> TelemetryResult<Option<PressureObservation>> {
    let reader = BufReader::new(bytes);
    let mut sessions: HashMap<String, GeminiSessionSnapshot> = HashMap::new();
    let mut latest_session_key = None;

    for line in reader.lines() {
        let line = line.map_err(TelemetryError::from)?;
        if line.trim().is_empty() {
            continue;
        }
        let value: Value = match serde_json::from_str(&line) {
            Ok(v) => v,
            Err(_) => continue,
        };
        let event_name = event_name(&value).unwrap_or_else(|| event_name_from_line(&line));
        if event_name.is_empty() {
            continue;
        }
        let session_key =
            string_key(&value, &["session_id", "sessionId"]).unwrap_or_else(|| "global".to_owned());
        latest_session_key = Some(session_key.clone());
        let session = sessions.entry(session_key).or_default();
        apply_event(session, &event_name, &value);
    }

    let Some(session_key) = latest_session_key else {
        return Ok(None);
    };
    let Some(session) = sessions.remove(&session_key) else {
        return Ok(None);
    };
    if session.latest_prompt_tokens.is_none() && !session.compacted {
        return Ok(None);
    }

    Ok(Some(PressureObservation {
        adapter_id: ADAPTER_ID.into(),
        adapter_version: None,
        observed_at_epoch_s,
        model_name: session
            .model_name
            .or_else(|| resolve_env_string(GEMINI_MODEL_ALIASES))
            .or_else(general_host_model),
        total_tokens: session.latest_prompt_tokens,
        context_window_tokens: session.context_window,
        context_used_pct: session
            .latest_prompt_tokens
            .and_then(|t| compute_pct(t, session.context_window)),
        compaction_signal: session.compacted.then_some(true),
        usage: TokenUsage {
            input_tokens: session.input_tokens,
            output_tokens: session.output_tokens,
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: session.cache_read_input_tokens,
            blended_total_tokens: None,
        },
    }))
}

fn telemetry_observation(repo_root: &Path) -> TelemetryResult<Option<(PathBuf, u64)>> {
    let Some(path) = telemetry_path(repo_root)? else {
        return Ok(None);
    };
    let Some(observed_at_epoch_s) = file_mtime_epoch_s(&path)? else {
        return Ok(None);
    };
    if resolve_env_string(GEMINI_TELEMETRY_PATH_ALIASES).is_none()
        && !is_recent(observed_at_epoch_s)?
    {
        return Ok(None);
    }
    Ok(Some((path, observed_at_epoch_s)))
}

fn telemetry_path(repo_root: &Path) -> TelemetryResult<Option<PathBuf>> {
    if let Some(path) = resolve_env_string(GEMINI_TELEMETRY_PATH_ALIASES) {
        let path = PathBuf::from(path);
        return Ok(path.is_file().then_some(path));
    }
    let candidates = [
        repo_root.join(".gemini/telemetry.log"),
        repo_root.join("telemetry.log"),
        home_dir()?.join(".gemini/telemetry.log"),
    ];
    Ok(candidates.into_iter().find(|p| p.is_file()))
}

fn apply_event(session: &mut GeminiSessionSnapshot, event_name: &str, value: &Value) {
    if event_name.contains("gemini_cli.api_response") {
        apply_api_response(session, value);
    }
    if event_name.contains("gemini_cli.chat_compression") {
        session.compacted = true;
    }
}

fn apply_api_response(session: &mut GeminiSessionSnapshot, value: &Value) {
    if let Some(prompt_tokens) = prompt_tokens(value) {
        session.latest_prompt_tokens = Some(prompt_tokens);
    }
    if let Some(input_tokens) = input_tokens(value) {
        session.input_tokens = session.input_tokens.saturating_add(input_tokens);
    }
    if let Some(output_tokens) = output_tokens(value) {
        session.output_tokens = session.output_tokens.saturating_add(output_tokens);
    }
    if let Some(cache_read_tokens) = cache_read_input_tokens(value) {
        session.cache_read_input_tokens = session
            .cache_read_input_tokens
            .saturating_add(cache_read_tokens);
    }
    session.context_window = number_key(
        value,
        &[
            "model_context_window",
            "context_window",
            "contextWindow",
            "limit_context",
        ],
    )
    .or_else(|| resolve_env_u64(GEMINI_CONTEXT_WINDOW_ALIASES))
    .or_else(general_context_window);
    if let Some(model_name) = string_key(value, &["model", "model_name", "modelName", "model_id"]) {
        session.model_name = Some(model_name);
    }
}

fn event_name(value: &Value) -> Option<String> {
    string_key(value, &["name", "event_name", "eventName"])
}

fn event_name_from_line(line: &str) -> String {
    if line.contains("gemini_cli.api_response") {
        "gemini_cli.api_response".to_owned()
    } else if line.contains("gemini_cli.chat_compression") {
        "gemini_cli.chat_compression".to_owned()
    } else {
        String::new()
    }
}

fn prompt_tokens(value: &Value) -> Option<u64> {
    let input_tokens = input_tokens(value).unwrap_or(0);
    let cached_tokens = cache_read_input_tokens(value).unwrap_or(0);
    let total = input_tokens.saturating_add(cached_tokens);
    (total > 0).then_some(total)
}

fn input_tokens(value: &Value) -> Option<u64> {
    number_key(
        value,
        &[
            "input_token_count",
            "inputTokens",
            "prompt_tokens",
            "promptTokenCount",
        ],
    )
}

fn output_tokens(value: &Value) -> Option<u64> {
    number_key(
        value,
        &[
            "output_token_count",
            "outputTokens",
            "completion_tokens",
            "completionTokenCount",
        ],
    )
}

fn cache_read_input_tokens(value: &Value) -> Option<u64> {
    number_key(
        value,
        &[
            "cached_content_token_count",
            "cacheReadInputTokens",
            "cachedTokens",
        ],
    )
}