ralph/

compaction.rs

use crate::errors::Result;
use crate::prompts::compaction_prompt;
use crate::providers::{LlmProvider, Message};
use std::sync::OnceLock;
use tiktoken_rs::CoreBPE;

// ── Token counting ─────────────────────────────────────────────────────────────

/// Get the cl100k_base BPE, initialised once.
fn bpe() -> &'static CoreBPE {
    static BPE: OnceLock<CoreBPE> = OnceLock::new();
    BPE.get_or_init(|| tiktoken_rs::cl100k_base().expect("tiktoken cl100k_base failed to load"))
}

/// Count tokens in a single string using cl100k_base.
pub fn count_tokens(text: &str) -> u64 {
    bpe().encode_ordinary(text).len() as u64
}

/// Extract all text from a message, including tool use inputs and tool results.
/// `as_text()` only returns `ContentPart::Text` nodes, missing the bulk of
/// tool-heavy conversations (file reads, command output, etc.).
fn message_text(msg: &Message) -> String {
    use crate::providers::{ContentPart, MessageContent};
    match &msg.content {
        MessageContent::Text(s) => s.clone(),
        MessageContent::Parts(parts) => {
            let mut buf = String::new();
            for part in parts {
                match part {
                    ContentPart::Text { text } => buf.push_str(text),
                    ContentPart::ToolUse { name, input, .. } => {
                        buf.push_str(name);
                        buf.push_str(&input.to_string());
                    }
                    ContentPart::ToolResult { content, .. } => {
                        buf.push_str(content);
                    }
                }
            }
            buf
        }
    }
}

/// Count total tokens across all messages.
pub fn estimate_tokens(messages: &[Message]) -> u64 {
    // 4 overhead tokens per message (role tag etc.) + content tokens
    messages
        .iter()
        .map(|m| count_tokens(&message_text(m)) + 4)
        .sum()
}

/// Returns true if compaction should fire.
/// Fires when `estimate_tokens(messages) + tool_overhead_tokens >= compact_at_tokens`.
pub fn should_compact(
    messages: &[Message],
    compact_at_tokens: u64,
    tool_overhead_tokens: u64,
) -> bool {
    if compact_at_tokens == 0 {
        return true;
    }
    let used = estimate_tokens(messages) + tool_overhead_tokens;
    used >= compact_at_tokens
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Trim `window` from position 1 onward (preserving index 0, the summary
/// message) until the total token count is at or below `target_tokens`.
fn trim_to_fit(window: &mut Vec<Message>, target_tokens: u64) {
    while window.len() > 1 && estimate_tokens(window) > target_tokens {
        window.remove(1);
    }
}

/// Summarize a slice of messages into a single string using the LLM.
/// The slice is front-trimmed to fit within the model's context window
/// before the compaction prompt is appended.
async fn summarize_chunk(
    chunk: &[Message],
    provider: &dyn LlmProvider,
    context_window: u64,
) -> Result<String> {
    let prompt_tokens = count_tokens(compaction_prompt()) + 4;
    let available = (context_window * 90 / 100).saturating_sub(prompt_tokens);

    let mut start = 0usize;
    while start < chunk.len() {
        let tokens: u64 = chunk[start..]
            .iter()
            .map(|m| count_tokens(&message_text(m)) + 4)
            .sum();
        if tokens <= available {
            break;
        }
        start += 1;
    }

    let mut msgs = chunk[start..].to_vec();
    msgs.push(Message::user(compaction_prompt()));

    let response = provider.chat(&msgs, &[]).await?;
    Ok(response
        .text
        .unwrap_or_else(|| "Summary unavailable.".to_string()))
}

// ── Primary compaction ────────────────────────────────────────────────────────

/// Compact the message history via an LLM-generated summary.
///
/// Strategy:
/// 1. Front-trim messages so the compaction request fits in the context window.
/// 2. Summarise with the LLM.
/// 3. Build new_window = [summary] + last `keep_recent_turns` turns verbatim.
/// 4. Trim new_window to ≤ `compact_to_tokens` so future turns have room.
///
/// Returns: (new_window, original_archive)
pub async fn compact(
    messages: &[Message],
    provider: &dyn LlmProvider,
    keep_recent_turns: usize,
    compact_to_tokens: u64,
) -> Result<(Vec<Message>, Vec<Message>)> {
    let context_window = provider.context_window();
    let summary_text = summarize_chunk(messages, provider, context_window).await?;

    let recent_start = messages.len().saturating_sub(keep_recent_turns * 2);
    let mut new_window = Vec::new();
    new_window.push(Message::user(format!(
        "## Compacted Session Summary\n\n{}\n\n---\nContinuing from where we left off.",
        summary_text
    )));
    new_window.extend_from_slice(&messages[recent_start..]);

    trim_to_fit(&mut new_window, compact_to_tokens);

    Ok((new_window, messages.to_vec()))
}

// ── Fallback: split-and-double-compact ───────────────────────────────────────

/// Hierarchical compaction used when primary `compact()` fails.
///
/// Strategy:
/// 1. Isolate the recent verbatim tail (last `keep_recent_turns` turns).
/// 2. Split the remaining older messages into two halves.
/// 3. Summarise each half independently with the LLM (two API calls).
/// 4. Combine both summaries + recent tail into the new window.
/// 5. Trim new_window to ≤ 50 % of context_window.
///
/// Because each half is at most ~50 % of the full history, the summarisation
/// requests are well within the context window even when the full history
/// is not.
pub async fn compact_split(
    messages: &[Message],
    provider: &dyn LlmProvider,
    keep_recent_turns: usize,
    compact_to_tokens: u64,
) -> Result<(Vec<Message>, Vec<Message>)> {
    let context_window = provider.context_window();
    let recent_start = messages.len().saturating_sub(keep_recent_turns * 2);
    let recent = &messages[recent_start..];
    let old = &messages[..recent_start];

    // If there is nothing to split, fall through to a single-pass summary.
    let (summary_1, summary_2) = if old.len() < 2 {
        let s = summarize_chunk(messages, provider, context_window).await?;
        (s, String::new())
    } else {
        let mid = old.len() / 2;
        let (s1, s2) = tokio::try_join!(
            summarize_chunk(&old[..mid], provider, context_window),
            summarize_chunk(&old[mid..], provider, context_window),
        )?;
        (s1, s2)
    };

    let combined = if summary_2.is_empty() {
        summary_1
    } else {
        format!(
            "## Part 1 (earliest history)\n\n{}\n\n## Part 2 (continued)\n\n{}",
            summary_1, summary_2
        )
    };

    let mut new_window = Vec::new();
    new_window.push(Message::user(format!(
        "## Compacted Session Summary\n\n{}\n\n---\nContinuing from where we left off.",
        combined
    )));
    new_window.extend_from_slice(recent);

    trim_to_fit(&mut new_window, compact_to_tokens);

    Ok((new_window, messages.to_vec()))
}

// ── Archive ───────────────────────────────────────────────────────────────────

/// Append the pre-compaction history to the archive on disk.
pub fn archive(session_dir: &std::path::Path, original: &[Message]) -> Result<()> {
    let archive_path = session_dir.join("compacted.json");
    let mut existing: Vec<Vec<Message>> = if archive_path.exists() {
        let s = std::fs::read_to_string(&archive_path).unwrap_or_default();
        serde_json::from_str(&s).unwrap_or_default()
    } else {
        Vec::new()
    };
    existing.push(original.to_vec());
    let content = serde_json::to_string_pretty(&existing)?;
    std::fs::write(&archive_path, content)?;
    Ok(())
}