clawgarden-agent 0.7.9

//! Token budget-based context manager
//!
//! Selects history to pass to LLM within token budget.
//! Prioritizes recent messages, includes initial messages if budget allows.

/// Token estimation factor (1 token ≈ 3 bytes, mixed English+Korean)
/// Korean may be overestimated but conservative (safe) direction.
const BYTES_PER_TOKEN: usize = 3;

/// Context manager
pub struct ContextManager {
    /// Token budget
    budget_tokens: usize,
}

impl ContextManager {
    pub fn new(budget_tokens: usize) -> Self {
        Self { budget_tokens }
    }

    pub fn with_default_budget() -> Self {
        let config = clawgarden_proto::AppConfig::load();
        Self::new(config.agent.context_budget_tokens)
    }

    /// Estimate token count for text.
    pub fn estimate_tokens(text: &str) -> usize {
        if text.is_empty() {
            return 0;
        }
        (text.len() / BYTES_PER_TOKEN).max(1)
    }

    /// Select messages from history within token budget.
    ///
    /// Strategy:
    /// 1. Include recent messages first (most recent first)
    /// 2. Include initial messages (first 2) if budget allows
    /// 3. Return in chronological order (oldest first)
    ///
    /// `user_message` is deducted separately from budget, so long user messages
    /// automatically reduce history.
    pub fn select_history(&self, history: &[String], user_message: &str) -> Vec<String> {
        if history.is_empty() {
            return Vec::new();
        }

        // Deduct budget used by the user message
        let user_tokens = Self::estimate_tokens(user_message);
        let mut remaining = self.budget_tokens.saturating_sub(user_tokens);

        // 1. Include recent messages first (most recent first)
        let mut selected_indices: Vec<usize> = Vec::new();

        for i in (0..history.len()).rev() {
            let tokens = Self::estimate_tokens(&history[i]);
            if tokens <= remaining {
                selected_indices.push(i);
                remaining -= tokens;
            } else {
                break;
            }
        }

        // 2. Include initial messages (first 2) if budget allows
        let initial_count = 2.min(history.len());
        for i in 0..initial_count {
            if !selected_indices.contains(&i) {
                let tokens = Self::estimate_tokens(&history[i]);
                if tokens <= remaining {
                    selected_indices.push(i);
                    remaining -= tokens;
                }
            }
        }

        // 3. Sort chronologically
        selected_indices.sort();

        selected_indices
            .into_iter()
            .map(|i| history[i].clone())
            .collect()
    }
}

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

    fn make_history(count: usize, words_per_msg: usize) -> Vec<String> {
        (0..count)
            .map(|i| {
                let words: Vec<String> = (0..words_per_msg)
                    .map(|j| format!("word{}_{}", i, j))
                    .collect();
                format!("[agent_{}]: {}", i, words.join(" "))
            })
            .collect()
    }

    #[test]
    fn test_estimate_tokens_empty() {
        assert_eq!(ContextManager::estimate_tokens(""), 0);
    }

    #[test]
    fn test_estimate_tokens_nonzero() {
        let tokens = ContextManager::estimate_tokens("Hello world");
        assert!(tokens > 0);
    }

    #[test]
    fn test_select_history_empty() {
        let mgr = ContextManager::with_default_budget();
        let result = mgr.select_history(&[], "hello");
        assert!(result.is_empty());
    }

    #[test]
    fn test_select_history_fits_all() {
        let mgr = ContextManager::new(10000);
        let history = make_history(5, 5);
        let result = mgr.select_history(&history, "test");
        assert_eq!(result.len(), 5, "enough budget includes all");
    }

    #[test]
    fn test_select_history_truncates_recent_first() {
        // Small budget: only recent messages included
        let mgr = ContextManager::new(200);
        let history = make_history(20, 10);
        let result = mgr.select_history(&history, "test");
        assert!(result.len() < 20, "small budget truncates");
        // Last message should be included
        assert!(result.last().unwrap().contains("agent_19"));
    }

    #[test]
    fn test_select_history_preserves_initial() {
        // Enough budget: initial messages also included
        let mgr = ContextManager::new(10000);
        let history = make_history(20, 5);
        let result = mgr.select_history(&history, "test");
        assert!(
            result.first().unwrap().contains("agent_0"),
            "initial messages preserved"
        );
    }

    #[test]
    fn test_user_message_deducts_budget() {
        // Large user message consumes budget, reducing history
        let mgr = ContextManager::new(500);
        let history = make_history(50, 10);
        let big_user_msg = "a".repeat(3000); // ~1000 tokens
        let result = mgr.select_history(&history, &big_user_msg);
        assert!(result.len() < 50, "large user message reduces history");
    }
}