roma-core 0.1.0

//! History compression strategies.
//!
//! Compression runs in two passes:
//!
//! 1. **Tag compression**: replace the `content` string of `ToolResult`
//!    blocks (typically the largest source of bloat) with a short stub
//!    once total tokens exceed the target.
//! 2. **Head trimming**: if still over budget, drop the oldest
//!    user/assistant pair. The very first message is preserved whenever
//!    possible because it often carries the task brief.
//!
//! The strategy is conservative: we never touch `ContentBlock::Text` on
//! assistant turns (which contain reasoning) or the final assistant
//! message of a tool call (which the next turn's tool_result must match
//! by tool_use_id).

use std::borrow::Cow;
use std::cmp::min;

use crate::tokenizer::Tokenizer;
use crate::types::{ContentBlock, Message, Role};

const COMPRESSED_TOOL_RESULT_STUB: &str = "[tool result compressed]";
const HEAD_TRIM_PROBE_MSGS: usize = 64;

/// Messages containing a stream-rule injection are never stubbed or
/// trimmed (spec C5/D6): long sessions are exactly where rules matter.
fn contains_rule_sentinel(text: &str) -> bool {
    text.contains(crate::rules::RULE_SENTINEL)
}

fn message_has_rule_injection(m: &Message) -> bool {
    m.content.iter().any(|b| match b {
        ContentBlock::Text { text } => contains_rule_sentinel(text),
        ContentBlock::ToolResult { content, .. } => contains_rule_sentinel(content),
        _ => false,
    })
}

/// Strategy for shrinking chat history under a target token count.
pub trait HistoryCompressor: Send + Sync {
    /// Compress `messages` in place until the total token count reported
    /// by `tokenizer` is at most `target_tokens`, or we run out of
    /// reversible compression moves.
    ///
    /// Returns the token count after compression (which may still exceed
    /// `target_tokens` if no further moves were safe).
    fn compress(
        &self,
        messages: &mut Vec<Message>,
        target_tokens: u32,
        tokenizer: &dyn Tokenizer,
    ) -> u32;
}

/// Default two-pass compressor.
#[derive(Debug, Default, Clone, Copy)]
pub struct DefaultHistoryCompressor;

impl HistoryCompressor for DefaultHistoryCompressor {
    fn compress(
        &self,
        messages: &mut Vec<Message>,
        target_tokens: u32,
        tokenizer: &dyn Tokenizer,
    ) -> u32 {
        let mut total = count_total(messages, tokenizer);
        if total <= target_tokens {
            return total;
        }

        // Pass 1: stub out ToolResult content blocks, oldest first.
        for i in 0..messages.len() {
            for block in messages[i].content.iter_mut() {
                if let ContentBlock::ToolResult { content, .. } = block
                    && content.as_str() != COMPRESSED_TOOL_RESULT_STUB
                    && !contains_rule_sentinel(content)
                {
                    *content = COMPRESSED_TOOL_RESULT_STUB.to_string();
                }
            }
            total = count_total(messages, tokenizer);
            if total <= target_tokens {
                return total;
            }
        }

        // Pass 2: trim from the head one *turn* at a time. A turn is a
        // sequence of messages anchored by a user-text message (not a
        // tool_result), running through any assistant + tool_result
        // interleaving up to the next user-text message. Removing a
        // whole turn preserves the Anthropic invariant that every
        // ToolUse block has a matching ToolResult later in history.
        let preserved_first = !messages.is_empty() && is_user_anchor(&messages[0]);
        let anchor_floor = if preserved_first { 1 } else { 0 };
        // `floor` advances past pinned turns so the search never revisits them.
        let mut floor = anchor_floor;
        let mut attempts = 0;
        while count_total(messages, tokenizer) > target_tokens
            && attempts < HEAD_TRIM_PROBE_MSGS
            && messages.len() > floor + 1
        {
            // Find the first user-anchor at or after floor.
            let Some(turn_start) = messages
                .iter()
                .enumerate()
                .skip(floor)
                .find(|(_, m)| is_user_anchor(m))
                .map(|(i, _)| i)
            else {
                break;
            };
            // Find the next user-anchor (exclusive end of this turn).
            let turn_end = messages
                .iter()
                .enumerate()
                .skip(turn_start + 1)
                .find(|(_, m)| is_user_anchor(m))
                .map_or(messages.len(), |(i, _)| i);
            // If the turn is the entire remainder, stop — we'd otherwise
            // delete the live conversation tail.
            if turn_end == messages.len() {
                break;
            }
            // C5: turns containing a rule injection are pinned — skip past
            // them instead of draining.
            if messages[turn_start..turn_end]
                .iter()
                .any(message_has_rule_injection)
            {
                floor = turn_end;
                continue;
            }
            messages.drain(turn_start..min(turn_end, messages.len()));
            attempts += 1;
        }

        count_total(messages, tokenizer)
    }
}

/// A user message anchors a new turn iff it carries text content (not
/// just `tool_result` blocks). This is the boundary at which trimming is
/// safe: dropping `[turn_start..turn_end)` removes a self-contained
/// chunk of conversation including any internal tool_use/tool_result
/// pairs.
fn is_user_anchor(m: &Message) -> bool {
    if m.role != Role::User {
        return false;
    }
    // A tool-result carrier is interior to its turn, never a boundary —
    // even when it also carries text (memory snapshots and stream-rule
    // injections are appended to tool-result messages). Treating it as an
    // anchor would let a drain separate tool_result from its tool_use.
    if m.content
        .iter()
        .any(|b| matches!(b, ContentBlock::ToolResult { .. }))
    {
        return false;
    }
    m.content
        .iter()
        .any(|b| matches!(b, ContentBlock::Text { .. }))
}

/// Count tokens across all blocks of all messages, approximating a
/// provider-side payload. We count each block's textual payload and add
/// a small per-message overhead.
fn count_total(messages: &[Message], tokenizer: &dyn Tokenizer) -> u32 {
    let mut total: u32 = 0;
    for m in messages {
        total = total.saturating_add(4); // per-message overhead
        for b in &m.content {
            total = total.saturating_add(tokenizer.count(&block_text(b)));
        }
    }
    total
}

fn block_text(b: &ContentBlock) -> Cow<'_, str> {
    match b {
        ContentBlock::Text { text } => Cow::Borrowed(text),
        ContentBlock::Thinking { thinking, .. } => Cow::Borrowed(thinking),
        // `input` is a serde_json::Value; `as_str()` returns None for
        // Objects/Arrays, so we use `to_string()` to get the JSON text.
        ContentBlock::ToolUse { input, .. } => Cow::Owned(input.to_string()),
        ContentBlock::ToolResult { content, .. } => Cow::Borrowed(content),
    }
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::tokenizer::ApproxTokenizer;
    use crate::types::{Message, Role};

    fn tool_result(id: &str, content: &str) -> Message {
        Message {
            role: Role::User,
            content: vec![ContentBlock::ToolResult {
                tool_use_id: id.into(),
                content: content.into(),
                is_error: Some(false),
            }],
        }
    }

    #[test]
    fn noop_when_already_under_budget() {
        let mut msgs = vec![Message::user_text("hi")];
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;
        let n = c.compress(&mut msgs, 1000, &t);
        assert!(n < 100);
        assert_eq!(msgs.len(), 1);
    }

    #[test]
    fn stubs_tool_results_when_over_budget() {
        let big = "x".repeat(4000); // ~1000 tokens under approx.
        let mut msgs = vec![
            Message::user_text("hi"),
            tool_result("t1", &big),
            Message::assistant_text("ok"),
        ];
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;
        let _ = c.compress(&mut msgs, 50, &t);
        // The tool result content must have been replaced with the stub.
        match &msgs[1].content[0] {
            ContentBlock::ToolResult { content, .. } => {
                assert_eq!(content, COMPRESSED_TOOL_RESULT_STUB);
            }
            _ => panic!("expected tool result"),
        }
    }

    #[test]
    fn trims_head_preserving_first_message() {
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;
        let mut msgs = vec![
            Message::user_text("<brief>"),
            Message::user_text("a".repeat(4000)),
            Message::assistant_text("b".repeat(4000)),
            Message::user_text("c".repeat(4000)),
            Message::assistant_text("final"),
        ];
        let before = msgs.len();
        c.compress(&mut msgs, 20, &t);
        assert!(msgs.len() < before);
        // First message (brief) preserved.
        match &msgs[0].content[0] {
            ContentBlock::Text { text } => assert_eq!(text, "<brief>"),
            _ => panic!("first block should be the brief"),
        }
    }

    #[test]
    fn stops_when_only_two_messages_remain() {
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;
        // Extreme case: tight budget but only 2 messages; the compressor
        // must stop trimming rather than empty the transcript.
        let mut msgs = vec![
            Message::user_text("a".repeat(1000)),
            Message::assistant_text("b".repeat(1000)),
        ];
        c.compress(&mut msgs, 1, &t);
        assert!(!msgs.is_empty());
    }

    /// Regression for P0-3: head trimming must not orphan a ToolUse
    /// without its ToolResult, or vice versa. After compression the set
    /// of `tool_use` ids and the set of `tool_use_id`s referenced by
    /// `tool_result`s must be equal.
    #[test]
    fn trim_preserves_tool_use_tool_result_pairing() {
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;

        fn tool_use_msg(id: &str, name: &str) -> Message {
            Message {
                role: Role::Assistant,
                content: vec![ContentBlock::ToolUse {
                    id: id.into(),
                    name: name.into(),
                    input: serde_json::json!({"path": "x".repeat(2000)}),
                }],
            }
        }
        fn tool_result_msg(id: &str, body: &str) -> Message {
            Message {
                role: Role::User,
                content: vec![ContentBlock::ToolResult {
                    tool_use_id: id.into(),
                    content: body.into(),
                    is_error: Some(false),
                }],
            }
        }

        let mut msgs = vec![
            // Brief — preserved.
            Message::user_text("<brief>"),
            // Turn 1 — tool round trip.
            Message::user_text("turn1 user"),
            tool_use_msg("t1", "file_read"),
            tool_result_msg("t1", &"r".repeat(4000)),
            Message::assistant_text("turn1 done"),
            // Turn 2 — tool round trip.
            Message::user_text("turn2 user"),
            tool_use_msg("t2", "file_write"),
            tool_result_msg("t2", &"r".repeat(4000)),
            Message::assistant_text("turn2 done"),
            // Tail — newest user prompt.
            Message::user_text("turn3 user"),
        ];
        c.compress(&mut msgs, 50, &t);

        let mut tool_use_ids = std::collections::BTreeSet::new();
        let mut tool_result_ids = std::collections::BTreeSet::new();
        for m in &msgs {
            for b in &m.content {
                match b {
                    ContentBlock::ToolUse { id, .. } => {
                        tool_use_ids.insert(id.clone());
                    }
                    ContentBlock::ToolResult { tool_use_id, .. } => {
                        tool_result_ids.insert(tool_use_id.clone());
                    }
                    _ => {}
                }
            }
        }
        assert_eq!(
            tool_use_ids, tool_result_ids,
            "tool_use ids and tool_result ids must remain in sync after trim"
        );
        // The very last user message (live tail) must also survive.
        match &msgs.last().unwrap().content[0] {
            ContentBlock::Text { text } => assert_eq!(text, "turn3 user"),
            _ => panic!("expected the live tail to be preserved"),
        }
    }

    #[test]
    fn tool_result_with_rule_note_is_never_stubbed() {
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;
        let big = "x".repeat(4000);
        let mut msgs = vec![
            Message::user_text("hi"),
            tool_result("t1", &big),
            tool_result(
                "t2",
                &format!("{big}\n\n[stream rule `r` reminder]\nnote body"),
            ),
            Message::assistant_text("ok"),
        ];
        let _ = c.compress(&mut msgs, 50, &t);
        match &msgs[1].content[0] {
            ContentBlock::ToolResult { content, .. } => {
                assert_eq!(content, COMPRESSED_TOOL_RESULT_STUB)
            }
            _ => panic!("expected tool result"),
        }
        match &msgs[2].content[0] {
            ContentBlock::ToolResult { content, .. } => {
                assert!(
                    content.contains("[stream rule `r` reminder]"),
                    "sentinel-bearing tool result must not be stubbed"
                )
            }
            _ => panic!("expected tool result"),
        }
    }

    #[test]
    fn turn_with_rule_injection_survives_head_trim() {
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;
        let mut msgs = vec![
            Message::user_text("<brief>"),
            // Pinned turn: the user message carries a rule injection.
            Message::user_text("do the thing\n\n[stream rule `r` fired]\n\nrule body"),
            Message::assistant_text("a".repeat(4000)),
            // Unpinned padding turns.
            Message::user_text("b".repeat(4000)),
            Message::assistant_text("c".repeat(4000)),
            Message::user_text("d".repeat(4000)),
            Message::assistant_text("e".repeat(4000)),
            // Live tail.
            Message::user_text("latest question"),
        ];
        c.compress(&mut msgs, 50, &t);
        assert!(
            msgs.iter()
                .any(|m| m.text_content().contains("[stream rule `r` fired]")),
            "the turn containing the rule injection must survive compression"
        );
        assert_eq!(msgs.last().unwrap().text_content(), "latest question");
    }

    /// Regression: a tool-result message that also carries a text block
    /// (stream-rule injections and memory snapshots are appended to
    /// tool-result messages) must never be treated as a turn anchor.
    /// Otherwise pass 2 can drain the turn *ending* at that pseudo-anchor,
    /// orphaning the ToolResult from its ToolUse — every subsequent
    /// provider request then 400s, permanently corrupting the session.
    #[test]
    fn tool_result_with_appended_text_is_never_an_anchor() {
        let t = ApproxTokenizer;
        let c = DefaultHistoryCompressor;

        fn tool_use_msg(id: &str, name: &str) -> Message {
            Message {
                role: Role::Assistant,
                content: vec![ContentBlock::ToolUse {
                    id: id.into(),
                    name: name.into(),
                    input: serde_json::json!({"path": "x".repeat(2000)}),
                }],
            }
        }

        let mut msgs = vec![
            // Brief — preserved.
            Message::user_text("<brief>"),
            // Turn with a real tool round trip; the tool-result message
            // also carries a stream-rule injection as a Text block.
            Message::user_text("turn1 user"),
            tool_use_msg("t1", "file_read"),
            Message {
                role: Role::User,
                content: vec![
                    ContentBlock::ToolResult {
                        tool_use_id: "t1".into(),
                        content: "r".repeat(4000),
                        is_error: Some(false),
                    },
                    ContentBlock::Text {
                        text: "\n\n[stream rule `r` fired]\n\nbody".into(),
                    },
                ],
            },
            Message::assistant_text("turn1 done"),
            // Unpinned padding turns.
            Message::user_text("b".repeat(4000)),
            Message::assistant_text("c".repeat(4000)),
            Message::user_text("d".repeat(4000)),
            Message::assistant_text("e".repeat(4000)),
            // Live tail.
            Message::user_text("latest question"),
        ];
        c.compress(&mut msgs, 50, &t);

        let mut tool_use_ids = std::collections::BTreeSet::new();
        let mut tool_result_ids = std::collections::BTreeSet::new();
        for m in &msgs {
            for b in &m.content {
                match b {
                    ContentBlock::ToolUse { id, .. } => {
                        tool_use_ids.insert(id.clone());
                    }
                    ContentBlock::ToolResult { tool_use_id, .. } => {
                        tool_result_ids.insert(tool_use_id.clone());
                    }
                    _ => {}
                }
            }
        }
        assert_eq!(
            tool_use_ids, tool_result_ids,
            "tool_use ids and tool_result ids must remain in sync after trim"
        );
        assert!(
            msgs.iter()
                .any(|m| m.text_content().contains("[stream rule `r` fired]")),
            "the rule injection must survive compression"
        );
    }
}