lean-ctx 3.9.19

Context Runtime for AI Agents with CCP. 79 MCP tools, 10 read modes, 95+ compression patterns, cross-session memory (CCP), persistent AI knowledge with temporal facts + contradiction detection, multi-agent context sharing, LITM-aware positioning, AAAK compact format, adaptive compression with Thompson Sampling bandits. Supports 24+ AI tools. Reduces LLM token consumption by up to 99%.
use super::baseline_recorder::{append_jsonl, read_jsonl};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::{io, path::PathBuf};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TreatmentCall {
    pub session_id: String,
    pub tool_name: String,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub model: String,
    pub timestamp: String,
    pub raw_cost: u64,
    pub compressed_tokens: u64,
    pub savings_tokens: u64,
    pub compression_ratio: f64,
}

impl TreatmentCall {
    pub fn new(
        session_id: &str,
        tool_name: &str,
        input_tokens: u64,
        original_output_tokens: u64,
        savings_tokens: u64,
        model: &str,
    ) -> Self {
        let compressed_tokens = original_output_tokens.saturating_sub(savings_tokens);
        Self {
            session_id: session_id.to_owned(),
            tool_name: tool_name.to_owned(),
            input_tokens,
            output_tokens: compressed_tokens,
            model: model.to_owned(),
            timestamp: Utc::now().to_rfc3339(),
            raw_cost: crate::core::value_gate::cost_tracker::calculate_cost(
                input_tokens,
                compressed_tokens,
                0,
                model,
            ),
            compressed_tokens,
            savings_tokens: original_output_tokens.saturating_sub(compressed_tokens),
            compression_ratio: if original_output_tokens == 0 {
                1.0
            } else {
                compressed_tokens as f64 / original_output_tokens as f64
            },
        }
    }
}

#[derive(Debug, Clone)]
pub struct TreatmentRecorder {
    directory: PathBuf,
}

impl TreatmentRecorder {
    pub fn new(directory: impl Into<PathBuf>) -> Self {
        Self {
            directory: directory.into(),
        }
    }
    pub fn path(&self) -> PathBuf {
        self.directory.join("treatment.jsonl")
    }
    pub fn record(&self, call: &TreatmentCall) -> io::Result<()> {
        append_jsonl(&self.path(), call)
    }
    pub fn read(&self) -> io::Result<Vec<TreatmentCall>> {
        read_jsonl(&self.path())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn treatment_recorder_includes_compression_data() {
        let dir = tempfile::tempdir().unwrap();
        let recorder = TreatmentRecorder::new(dir.path());
        recorder
            .record(&TreatmentCall::new("s1", "ctx_read", 10, 100, 60, "gpt-4o"))
            .unwrap();
        let entry = recorder.read().unwrap().pop().unwrap();
        assert_eq!(
            (
                entry.output_tokens,
                entry.compressed_tokens,
                entry.savings_tokens
            ),
            (40, 40, 60)
        );
        assert!((entry.compression_ratio - 0.4).abs() < f64::EPSILON);
    }
}