gyandoot-v2/crates/gyandoot-agent/src/context.rs

//! Production-grade context compression engine.
//!
//! Manages LLM conversation context by estimating token usage, pruning verbose
//! tool outputs, summarizing older messages, and assembling the final context
//! that gets sent to the provider.

use gyandoot_core::{ConversationContext, Message, MessageRole, ToolDefinition};
use tracing::debug;

// ─── ContextCompressor ──────────────────────────────────────────────────

/// Token-aware context window manager.
///
/// Strategies:
/// - Rough token estimation (chars/4) for fast checks.
/// - Summarization of older messages into a compact system message.
/// - Pruning of verbose tool outputs from mid-history.
/// - Keeping the most recent N turns intact.
pub struct ContextCompressor {
    /// Hard limit on context messages before summarization kicks in.
    max_messages_before_compress: usize,
    /// Minimum recent messages to keep intact.
    recent_context_count: usize,
}

impl ContextCompressor {
    /// Create a new compressor with sensible defaults.
    ///
    /// `max_tokens` is the model's context window size; compress triggers
    /// at 75% of this value.
    pub fn new(max_tokens: usize) -> Self {
        let max_messages = (max_tokens / 200).max(20);
        Self {
            max_messages_before_compress: max_messages,
            recent_context_count: (max_messages / 2).max(4),
        }
    }

    /// Create a compressor with explicit message thresholds.
    pub fn with_thresholds(max_messages: usize, recent_count: usize) -> Self {
        Self {
            max_messages_before_compress: max_messages,
            recent_context_count: recent_count.max(2),
        }
    }

    /// Estimate the number of tokens in a text string.
    ///
    /// Uses a simple char/4 heuristic. In production, swap this with a
    /// tiktoken-rs call for per-model accuracy.
    pub fn estimate_tokens(text: &str) -> usize {
        if text.is_empty() {
            return 1;
        }
        // ~4 chars per token on average for English text
        (text.len() / 4).max(1)
    }

    /// Estimate tokens for a slice of messages.
    pub fn estimate_messages_tokens(messages: &[Message]) -> usize {
        messages
            .iter()
            .map(|m| Self::estimate_tokens(&m.content) + 4) // +4 for role overhead
            .sum()
    }

    /// Compress a message list to fit within `max_tokens`.
    ///
    /// Strategy:
    /// 1. Always keep the system prompt (first message).
    /// 2. Prune verbose tool result bodies from middle messages.
    /// 3. Summarize older non-recent messages into a single condensed entry.
    /// 4. Keep the most recent N messages intact.
    pub async fn compress(
        &self,
        messages: Vec<Message>,
        max_tokens: usize,
    ) -> Vec<Message> {
        let total_tokens = Self::estimate_messages_tokens(&messages);
        let threshold = max_tokens * 3 / 4;

        if total_tokens <= threshold && messages.len() <= self.max_messages_before_compress {
            return messages;
        }

        debug!(
            "Compressing {} messages (estimated {} tokens, threshold {} tokens)",
            messages.len(),
            total_tokens,
            threshold,
        );

        // Step 1: Prune verbose tool results from middle messages
        let mut pruned = self.prune_tool_results(messages);

        // Step 2: If still over threshold, summarize
        let pruned_tokens = Self::estimate_messages_tokens(&pruned);
        if pruned_tokens > threshold || pruned.len() > self.max_messages_before_compress {
            pruned = self.summarize_middle(pruned);
        }

        debug!("Compressed to {} messages", pruned.len());
        pruned
    }

    /// Prune verbose tool outputs from message history, keeping only a
    /// compact summary for non-recent tool results.
    pub fn prune_tool_results(&self, messages: Vec<Message>) -> Vec<Message> {
        let keep_recent = self.recent_context_count;
        let mut result: Vec<Message> = Vec::with_capacity(messages.len());

        for (i, msg) in messages.iter().enumerate() {
            // Always keep the first message (system prompt)
            if i == 0 {
                result.push(msg.clone());
                continue;
            }

            // Keep recent messages intact
            if i >= messages.len().saturating_sub(keep_recent) {
                result.push(msg.clone());
                continue;
            }

            // Prune tool result content in middle messages
            if msg.role == MessageRole::Tool {
                let tool_result = msg.tool_result.as_deref().unwrap_or("");
                let truncated = if tool_result.len() > 200 {
                    format!(
                        "[Tool result: {} chars — truncated]\n...{}",
                        tool_result.len(),
                        &tool_result[..tool_result.len().min(200)]
                    )
                } else {
                    msg.content.clone()
                };
                result.push(Message::new(MessageRole::Tool, truncated));
            } else {
                result.push(msg.clone());
            }
        }

        result
    }

    /// Summarize non-recent, non-system messages into a single condensed
    /// system message.
    fn summarize_middle(&self, messages: Vec<Message>) -> Vec<Message> {
        let keep_recent = self.recent_context_count;
        let mut compressed: Vec<Message> = Vec::with_capacity(self.max_messages_before_compress);

        // Always keep the system prompt
        if let Some(system) = messages.first() {
            compressed.push(system.clone());
        }

        // Messages to summarize: everything between system prompt and recent tail
        let summarize_end = messages.len().saturating_sub(keep_recent);
        let middle: Vec<&Message> = messages[1..summarize_end]
            .iter()
            .filter(|m| m.role != MessageRole::Tool)
            .collect();

        if !middle.is_empty() {
            let summary = self.build_summary(&middle);
            compressed.push(Message::new(
                MessageRole::System,
                format!(
                    "[Summarized conversation history — {}]\n\
                     (Context from earlier turns in this session)",
                    summary
                ),
            ));
        }

        // Append the most recent messages
        let keep_start = messages.len().saturating_sub(keep_recent).max(1);
        for msg in messages.iter().skip(keep_start) {
            compressed.push(msg.clone());
        }

        compressed
    }

    /// Build a text summary from a slice of middle messages.
    fn build_summary(&self, messages: &[&Message]) -> String {
        let count = messages.len();
        let user_msgs: Vec<&str> = messages
            .iter()
            .filter(|m| m.role == MessageRole::User)
            .map(|m| m.content.as_str())
            .collect();

        if user_msgs.is_empty() {
            format!("{} assistant messages", count)
        } else if user_msgs.len() <= 3 {
            format!(
                "Previous discussion: {}",
                user_msgs
                    .iter()
                    .map(|s| truncate_str(s, 80))
                    .collect::<Vec<_>>()
                    .join(" | ")
            )
        } else {
            format!(
                "{} turns covering: {}, ... (and {} more)",
                count,
                user_msgs[..3]
                    .iter()
                    .map(|s| truncate_str(s, 80))
                    .collect::<Vec<_>>()
                    .join(", "),
                user_msgs.len().saturating_sub(3)
            )
        }
    }

    /// Assemble the final `ConversationContext` from messages, system
    /// prompt, and tool definitions.
    pub fn build_context(
        &self,
        messages: &[Message],
        system_prompt: &str,
        tool_definitions: &[ToolDefinition],
    ) -> ConversationContext {
        ConversationContext {
            messages: messages.to_vec(),
            system_prompt: system_prompt.to_string(),
            tools: tool_definitions.to_vec(),
            max_tokens: 0, // caller sets this or provider defaults
        }
    }
}

/// Truncate a string to `max_chars`, appending "…" if shortened.
fn truncate_str(s: &str, max_chars: usize) -> String {
    if s.len() <= max_chars {
        s.to_string()
    } else {
        format!("{}…", &s[..max_chars.saturating_sub(1)])
    }
}

// ─── Tests ──────────────────────────────────────────────────────────────

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

    fn make_msg(role: MessageRole, content: impl Into<String>) -> Message {
        Message::new(role, content)
    }

    #[test]
    fn test_estimate_tokens() {
        assert_eq!(ContextCompressor::estimate_tokens("abcd"), 1);
        assert_eq!(ContextCompressor::estimate_tokens(""), 1);
        let long = "a".repeat(400);
        assert_eq!(ContextCompressor::estimate_tokens(&long), 100);
    }

    #[test]
    fn test_estimate_messages_tokens() {
        let msgs = vec![
            make_msg(MessageRole::User, "hello"),
            make_msg(MessageRole::Assistant, "world"),
        ];
        let tokens = ContextCompressor::estimate_messages_tokens(&msgs);
        // Each content ~1 token, each +4 overhead = ~10 total
        assert!(tokens >= 10);
    }

    #[tokio::test]
    async fn test_no_compress_below_threshold() {
        let compressor = ContextCompressor::with_thresholds(20, 5);
        let messages: Vec<Message> = (0..5)
            .map(|i| make_msg(MessageRole::User, format!("msg {i}")))
            .collect();

        let result = compressor.compress(messages.clone(), 9999).await;
        assert_eq!(result.len(), 5);
    }

    #[tokio::test]
    async fn test_compress_above_threshold() {
        let compressor = ContextCompressor::with_thresholds(10, 4);
        let mut messages = vec![make_msg(MessageRole::System, "You are an AI.")];
        for i in 0..20 {
            messages.push(make_msg(MessageRole::User, format!("user msg {i}")));
            messages.push(make_msg(MessageRole::Assistant, format!("response {i}")));
        }

        let result = compressor.compress(messages, 9999).await;
        assert!(result.len() < 20);
        assert_eq!(result[0].content, "You are an AI.");
    }

    #[tokio::test]
    async fn test_compress_token_budget() {
        let compressor = ContextCompressor::with_thresholds(50, 8);
        let mut messages = vec![make_msg(MessageRole::System, "sys")];
        for i in 0..30 {
            messages.push(make_msg(MessageRole::User, format!("user msg {i}")));
            messages.push(make_msg(MessageRole::Assistant, format!("response {i}")));
        }
        // total messages: 1 + 60 = 61
        // estimated tokens: ~ 61 * (3/4 + 4) ≈ 290
        let result = compressor.compress(messages, 500).await;
        assert!(!result.is_empty());
        assert_eq!(result[0].content, "sys");
    }

    #[test]
    fn test_prune_tool_results_short() {
        let compressor = ContextCompressor::with_thresholds(20, 5);
        let messages = vec![
            make_msg(MessageRole::System, "sys"),
            make_msg(MessageRole::User, "hello"),
            Message::new(MessageRole::Tool, "short result").with_tool_result("short result".into()),
        ];
        let pruned = compressor.prune_tool_results(messages);
        assert_eq!(pruned.len(), 3);
        assert_eq!(pruned[2].content, "short result");
    }

    #[test]
    fn test_prune_tool_results_long() {
        let compressor = ContextCompressor::with_thresholds(20, 3);
        let long_result = "x".repeat(500);
        let mut messages = vec![
            make_msg(MessageRole::System, "sys"),
            make_msg(MessageRole::User, "hello"),
        ];
        // Add middle messages so tool result falls outside recent zone
        // total = 10, recent = 3, recent_start = 7
        for i in 0..5 {
            messages.push(make_msg(MessageRole::User, format!("middle {i}")));
        }
        messages.push(
            Message::new(MessageRole::Tool, &long_result)
                .with_tool_result(long_result.clone()),
        );
        messages.push(make_msg(MessageRole::Assistant, "response 1"));
        messages.push(make_msg(MessageRole::Assistant, "response 2"));
        messages.push(make_msg(MessageRole::Assistant, "recent response"));

        let pruned = compressor.prune_tool_results(messages);
        // Find the tool message
        let tool_msg = pruned.iter().find(|m| m.role == MessageRole::Tool).unwrap();
        // Long tool content should be truncated
        assert!(
            tool_msg.content.contains("[Tool result:"),
            "Expected truncated tool result, got: {}...",
            &tool_msg.content[..tool_msg.content.len().min(100)]
        );
    }

    #[test]
    fn test_truncate_str() {
        assert_eq!(truncate_str("hello", 10), "hello");
        let long = "a".repeat(100);
        let truncated = truncate_str(&long, 10);
        // 9 chars + '…' (multi-byte) = 10 characters
        assert_eq!(truncated.chars().count(), 10);
        assert!(truncated.ends_with('…'));
    }

    #[test]
    fn test_build_summary_empty() {
        let compressor = ContextCompressor::with_thresholds(20, 5);
        let summary = compressor.build_summary(&[]);
        assert!(summary.contains("0"));
    }

    #[test]
    fn test_build_summary_few_messages() {
        let compressor = ContextCompressor::with_thresholds(20, 5);
        let msg = make_msg(MessageRole::User, "hello there");
        let msgs = vec![&msg];
        let summary = compressor.build_summary(&msgs);
        assert!(summary.contains("hello there"));
    }

    #[test]
    fn test_build_context() {
        let compressor = ContextCompressor::with_thresholds(20, 5);
        let messages = vec![make_msg(MessageRole::User, "hi")];
        let ctx = compressor.build_context(&messages, "system prompt", &[]);
        assert_eq!(ctx.messages.len(), 1);
        assert_eq!(ctx.system_prompt, "system prompt");
    }
}