atman_runtime/

compaction.rs

use crate::message::{Message, MessagePart, MessageRole};

pub const KEEP_RECENT_MESSAGES: usize = 10;
pub const KEEP_RECENT_USER_TURNS: usize = 5;
const KEEP_RECENT_TOKEN_FRACTION: f64 = 0.05;

pub fn estimate_tokens_for_message(msg: &Message) -> u64 {
    let mut chars = 0usize;
    for part in &msg.parts {
        chars += match part {
            MessagePart::CompactSummary { summary, .. } => summary.len(),
            MessagePart::Text { text } => text.len(),
            MessagePart::Thinking { thinking, .. } => thinking.len(),
            MessagePart::ToolResult { content, .. } => content.len(),
            MessagePart::Image { .. } => 512,
            MessagePart::ToolUse { name, input, .. } => name.len() + input.to_string().len(),
        };
    }
    chars = chars.saturating_add(estimate_role_overhead(msg.role));
    (chars as f64 / 3.5).ceil() as u64
}

fn estimate_role_overhead(role: MessageRole) -> usize {
    match role {
        MessageRole::System => 12,
        MessageRole::User => 8,
        MessageRole::Assistant => 8,
        MessageRole::Tool => 16,
    }
}

pub fn estimate_tokens_for_messages(messages: &[Message]) -> u64 {
    messages.iter().map(estimate_tokens_for_message).sum()
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompactRange {
    pub start: usize,
    pub end: usize,
    pub tokens_saved_estimate: u64,
}

pub fn is_plan_related(msg: &Message) -> bool {
    for part in &msg.parts {
        match part {
            MessagePart::ToolUse { name, .. } if name.starts_with("plan.") => return true,
            MessagePart::ToolResult { content, .. } if content.starts_with("# Plan:") => {
                return true;
            }
            _ => {}
        }
    }
    false
}

pub fn is_compaction_summary(msg: &Message) -> bool {
    if !matches!(msg.role, MessageRole::System) {
        return false;
    }
    msg.parts
        .iter()
        .any(|part| matches!(part, MessagePart::CompactSummary { .. }))
}

fn find_kth_recent_user(messages: &[Message], k: usize) -> usize {
    let mut user_count = 0;
    for (index, message) in messages.iter().enumerate().rev() {
        if message.role == MessageRole::User {
            user_count += 1;
            if user_count == k {
                return index;
            }
        }
    }
    0
}

pub fn find_compact_range(messages: &[Message], budget: u64) -> Option<CompactRange> {
    let total = estimate_tokens_for_messages(messages);
    if total <= budget || messages.len() < 4 {
        return None;
    }

    let start = messages
        .iter()
        .rposition(is_compaction_summary)
        .unwrap_or(0);
    let keep_recent_tokens = (budget as f64 * KEEP_RECENT_TOKEN_FRACTION).ceil() as u64;
    let mut recent_tokens = 0u64;
    let mut token_end = messages.len();
    for (index, message) in messages.iter().enumerate().rev() {
        recent_tokens = recent_tokens.saturating_add(estimate_tokens_for_message(message));
        token_end = index;
        if recent_tokens >= keep_recent_tokens {
            break;
        }
    }
    let message_end = messages.len().saturating_sub(KEEP_RECENT_MESSAGES);
    let end = message_end
        .min(token_end)
        .min(find_kth_recent_user(messages, KEEP_RECENT_USER_TURNS));
    if end < start + 2 {
        return None;
    }

    let tokens_saved_estimate = messages[start..end]
        .iter()
        .map(estimate_tokens_for_message)
        .sum();
    Some(CompactRange {
        start,
        end,
        tokens_saved_estimate,
    })
}

pub fn estimate_compacted_message_tokens(
    messages: &[Message],
    range: &CompactRange,
    summary: &str,
) -> u64 {
    let turn_id = messages
        .get(range.start)
        .map(|m| m.turn_id.clone())
        .unwrap_or_else(crate::event::TurnId::now);
    let after = replace_range_with_summary(messages, range, summary.to_string(), turn_id);
    estimate_tokens_for_messages(&after)
}

pub fn filter_orphan_tool_messages(messages: &mut Vec<Message>) {
    let use_ids: std::collections::HashSet<String> = messages
        .iter()
        .flat_map(|m| {
            m.parts.iter().filter_map(|p| match p {
                MessagePart::ToolUse { id, .. } => Some(id.clone()),
                _ => None,
            })
        })
        .collect();
    let mut seen_results: std::collections::HashSet<String> = std::collections::HashSet::new();
    messages.retain(|m| {
        for p in &m.parts {
            if let MessagePart::ToolResult { tool_use_id, .. } = p {
                if !use_ids.contains(tool_use_id) {
                    return false;
                }
                if !seen_results.insert(tool_use_id.clone()) {
                    return false;
                }
            }
        }
        true
    });
}

pub fn find_compact_summaries(messages: &[Message]) -> Vec<CompactSummary> {
    let mut out = Vec::new();
    for (idx, msg) in messages.iter().enumerate() {
        if let Some(summary) = compact_summary(msg) {
            out.push(CompactSummary {
                message_index: idx,
                seq_start: summary.seq_start,
                seq_end: summary.seq_end,
                count: summary.count,
            });
        }
    }
    out
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompactSummary {
    pub message_index: usize,
    pub seq_start: u64,
    pub seq_end: u64,
    pub count: usize,
}

struct CompactSummaryPart {
    seq_start: u64,
    seq_end: u64,
    count: usize,
}

fn extract_anchor(messages: &[Message]) -> Option<(String, &[Message])> {
    let first = messages.first()?;
    let summary = first.parts.iter().find_map(|part| match part {
        MessagePart::CompactSummary { summary, .. } => Some(summary.clone()),
        _ => None,
    })?;
    Some((summary, &messages[1..]))
}

fn compact_summary(msg: &Message) -> Option<CompactSummaryPart> {
    if msg.role != MessageRole::System {
        return None;
    }
    msg.parts.iter().find_map(|part| match part {
        MessagePart::CompactSummary {
            seq_start,
            seq_end,
            count,
            ..
        } => Some(CompactSummaryPart {
            seq_start: *seq_start,
            seq_end: *seq_end,
            count: *count,
        }),
        _ => None,
    })
}

pub async fn maybe_auto_compact(
    session: &crate::session::Session,
    model: &str,
    providers: &crate::provider::ProviderRegistry,
) {
    let _compact_guard = session.acquire_compact_lock().await;
    maybe_auto_compact_locked(session, model, providers).await;
}

pub fn spawn_auto_compact(
    session: std::sync::Arc<crate::session::Session>,
    model: String,
    providers: crate::provider::ProviderRegistry,
) {
    tokio::task::spawn_blocking(move || {
        let Ok(rt) = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
        else {
            session.push_system_note("compaction skipped: background runtime init failed".into());
            return;
        };
        rt.block_on(async move {
            maybe_auto_compact(&session, &model, &providers).await;
        });
    });
}

pub async fn start_auto_compact(
    session: std::sync::Arc<crate::session::Session>,
    model: String,
    providers: crate::provider::ProviderRegistry,
) {
    let compact_guard = session.acquire_compact_lock_owned().await;
    tokio::task::spawn_blocking(move || {
        let Ok(rt) = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
        else {
            drop(compact_guard);
            session.push_system_note("compaction skipped: background runtime init failed".into());
            return;
        };
        rt.block_on(async move {
            maybe_auto_compact_locked(&session, &model, &providers).await;
            drop(compact_guard);
        });
    });
}

async fn maybe_auto_compact_locked(
    session: &crate::session::Session,
    model: &str,
    providers: &crate::provider::ProviderRegistry,
) {
    let forced = session.take_manual_compact_request();
    let info = crate::model_registry::model_info(model);
    let trigger = info.compaction_trigger_threshold();
    let target = info.compaction_target_after();
    let msgs = session.messages();
    let provider_tokens = session.last_input_tokens();
    let current = if provider_tokens > 0 {
        provider_tokens
    } else {
        estimate_tokens_for_messages(&msgs)
    };
    if !forced && current <= trigger {
        return;
    }
    if !forced && !session.approval_cooldown_ok_for_compact() {
        return;
    }
    let Some(range) = find_compact_range(&msgs, target) else {
        session.emit_compact_warning(
            model,
            current,
            trigger,
            info.context_budget,
            "no compactible span — history too short or already fully compacted",
        );
        return;
    };
    let _ = session
        .stream_tx()
        .send(crate::stream::StreamFrame::CompactionSummary {
            phase: crate::stream::CompactionPhase::Running,
            range_start: range.start,
            range_end: range.end.saturating_sub(1),
            summary: String::new(),
            before_tokens: current,
            after_tokens: 0,
            compacted_count: range.end - range.start,
        });
    let send_failed = |session: &crate::session::Session, reason: &str| {
        let _ = session
            .stream_tx()
            .send(crate::stream::StreamFrame::CompactionSummary {
                phase: crate::stream::CompactionPhase::Failed,
                range_start: range.start,
                range_end: range.end.saturating_sub(1),
                summary: reason.to_string(),
                before_tokens: current,
                after_tokens: current,
                compacted_count: range.end - range.start,
            });
    };
    let mut filtered: Vec<Message> = msgs[range.start..range.end].to_vec();
    filter_orphan_tool_messages(&mut filtered);
    let (anchor, new_messages) = extract_anchor(&filtered)
        .map(|(anchor, remaining)| (Some(anchor), remaining.to_vec()))
        .unwrap_or_else(|| (None, filtered.clone()));
    let summary =
        match generate_llm_summary(anchor.as_deref(), &new_messages, model, providers).await {
            Ok(text) => text,
            Err(err) => {
                session.emit_compact_warning(
                    model,
                    current,
                    trigger,
                    info.context_budget,
                    &format!("LLM summary failed: {err}. Degraded to placeholder."),
                );
                format!(
                    "[atman: compacted {} messages, LLM summary unavailable at {}]",
                    range.end - range.start,
                    chrono::Utc::now().to_rfc3339()
                )
            }
        };
    let final_summary =
        match request_review_if_enabled(session, forced, &filtered, &range, current, summary).await
        {
            ReviewOutcome::Commit(s) => s,
            ReviewOutcome::Rejected => {
                send_failed(
                    session,
                    "compaction rejected by user; keeping full transcript",
                );
                session.push_system_note(
                    "compaction rejected by user; keeping full transcript".into(),
                );
                return;
            }
        };
    let after_tokens = estimate_compacted_message_tokens(&msgs, &range, &final_summary);
    if after_tokens >= current {
        send_failed(
            session,
            &format!(
                "compaction skipped: summary would not shrink transcript ({} >= {} tokens)",
                after_tokens, current
            ),
        );
        session.push_system_note(format!(
            "compaction skipped: summary would not shrink transcript ({} >= {} tokens)",
            after_tokens, current
        ));
        return;
    }
    match session.compact_messages(final_summary, range, current) {
        Some(result) => {
            session.push_system_note(format!(
                "auto-compacted {}..{} — {} → {} tokens",
                result.compacted_start,
                result.compacted_end,
                result.before_tokens,
                result.after_tokens
            ));
        }
        None => {
            session.emit_compact_warning(
                model,
                current,
                trigger,
                info.context_budget,
                "no compactible span — history too short or already fully compacted",
            );
        }
    }
}

enum ReviewOutcome {
    Commit(String),
    Rejected,
}

async fn request_review_if_enabled(
    session: &crate::session::Session,
    forced: bool,
    slice: &[Message],
    range: &CompactRange,
    tokens_before: u64,
    summary: String,
) -> ReviewOutcome {
    if !session.compact_review_mode().should_review(forced) {
        return ReviewOutcome::Commit(summary);
    }
    let reviews = session.compact_reviews();
    if reviews.subscriber_count() == 0 {
        return ReviewOutcome::Commit(summary);
    }
    let pending = crate::session::PendingCompactReview {
        review_id: uuid::Uuid::now_v7().to_string(),
        summary: summary.clone(),
        slice_preview: format_slice_for_preview(slice),
        slice_count: slice.len(),
        range_start: range.start,
        range_end: range.end,
        tokens_before,
        emitted_at: chrono::Utc::now(),
    };
    let rx = reviews.request(pending);
    match rx.await {
        Ok(crate::session::CompactReviewDecision::AcceptAsIs) => ReviewOutcome::Commit(summary),
        Ok(crate::session::CompactReviewDecision::AcceptEdited { summary: edited }) => {
            ReviewOutcome::Commit(edited)
        }
        Ok(crate::session::CompactReviewDecision::Reject) | Err(_) => ReviewOutcome::Rejected,
    }
}

fn format_slice_for_preview(slice: &[Message]) -> String {
    let mut out = String::new();
    for (i, msg) in slice.iter().enumerate() {
        let role = msg.role.as_str();
        let body = serialize_message_for_summary(msg);
        let truncated: String = body.chars().take(400).collect();
        out.push_str(&format!("[{i}] {role}: {truncated}\n"));
    }
    out.chars().take(16_000).collect()
}

const SUMMARY_SYSTEM_PROMPT: &str =
    "You are an anchored context summarization assistant for coding sessions.";

const SUMMARY_INSTRUCTIONS: &str = r#"You are an anchored context summarization assistant.

Below is:
1. <current-anchor>: the existing handoff state, which is authoritative and must be preserved.
2. <new-messages>: only the messages that arrived since the anchor was written.

Merge the NEW facts from <new-messages> INTO the current anchor, producing an upgraded full anchor.

STRUCTURAL RULES (data model, not optional style):
- ## Objective: unchanged unless the new messages show the user explicitly redirected.
- ### Completed: ONLY ADD newly completed items. Never remove or re-evaluate an existing completed item. If a completed item is now in question, add it to ### Active or ### Blocked instead. NEVER delete from Completed.
- ### Active: update based on new messages; move newly-done items to Completed.
- ### Blocked: update based on new messages; remove resolved ones.
- ## Decisions: only add new decisions. Never remove old ones.
- ## Next Move: replace based on current end state.
- Keep every section, even when empty.
- Preserve exact file paths, symbols, commands, error strings, identifiers.

Output exactly this Markdown structure:
## Objective
## Important Details
## Work State
### Completed
### Active
### Blocked
## Decisions
## Next Move
## Relevant Files

Do not mention the summary process or that context was compacted.
Respond in the same language as the conversation."#;

async fn generate_llm_summary(
    anchor: Option<&str>,
    slice: &[Message],
    model: &str,
    providers: &crate::provider::ProviderRegistry,
) -> Result<String, crate::error::RuntimeError> {
    let provider = providers.resolve(model).ok_or_else(|| {
        crate::error::RuntimeError::ToolFailed(format!("no provider for {model}"))
    })?;
    let payload = format_slice_for_summary(slice);
    let (messages, dump_user) = if let Some(anchor) = anchor {
        let anchor_user = format!("<current-anchor>\n{anchor}\n</current-anchor>");
        let new_user =
            format!("<new-messages>\n{payload}\n</new-messages>\n\n{SUMMARY_INSTRUCTIONS}");
        (
            vec![
                Message::user_text(crate::event::TurnId::now(), anchor_user.clone()),
                Message::user_text(crate::event::TurnId::now(), new_user.clone()),
            ],
            format!("{anchor_user}\n\n{new_user}"),
        )
    } else {
        let user = format!(
            "<conversation_history>\n{payload}\n</conversation_history>\n\n{SUMMARY_INSTRUCTIONS}"
        );
        (
            vec![Message::user_text(
                crate::event::TurnId::now(),
                user.clone(),
            )],
            user,
        )
    };
    if let Ok(dir) = std::env::var("ATMAN_COMPACT_DUMP") {
        let _ = std::fs::write(
            format!("{dir}/compact_request.txt"),
            format!("=== SYSTEM ===\n{SUMMARY_SYSTEM_PROMPT}\n\n=== USER ===\n{dump_user}"),
        );
    }
    let req = crate::provider::LlmRequest {
        model: model.into(),
        messages,
        system: Some(SUMMARY_SYSTEM_PROMPT.into()),
        input: crate::value::Value::Unit,
        schema: None,
        cache_prompt: false,
        tools: Vec::new(),
        thinking_enabled: false,
        stall_timeout_secs: 0,
    };
    let outcome = provider.call(req).await?;
    let text = outcome.text_concat();
    if text.trim().is_empty() {
        return Err(crate::error::RuntimeError::ToolFailed(
            "empty summary from provider".into(),
        ));
    }
    Ok(text)
}

fn format_slice_for_summary(slice: &[Message]) -> String {
    let mut out = String::new();
    for (i, msg) in slice.iter().enumerate() {
        let role = msg.role.as_str();
        let body = serialize_message_for_summary(msg);
        let truncated: String = body.chars().take(4000).collect();
        out.push_str(&format!("[{i}] {role}: {truncated}\n\n"));
    }
    out.chars().take(120_000).collect()
}

fn serialize_message_for_summary(msg: &Message) -> String {
    let mut parts = Vec::new();
    for part in &msg.parts {
        match part {
            MessagePart::CompactSummary { summary, .. } => {
                parts.push(summary.clone());
            }
            MessagePart::Text { text } => {
                parts.push(text.clone());
            }
            MessagePart::Thinking { thinking, .. } => {
                let truncated: String = thinking.chars().take(1000).collect();
                parts.push(format!("[thinking: {truncated}]"));
            }
            MessagePart::ToolUse { name, input, .. } => {
                let input_str = if input.is_null() {
                    String::new()
                } else {
                    input.to_string()
                };
                let truncated: String = input_str.chars().take(2000).collect();
                parts.push(format!("[tool_call: {name}({truncated})]"));
            }
            MessagePart::ToolResult {
                content,
                is_error,
                tool_use_id,
            } => {
                let truncated: String = content.chars().take(3000).collect();
                let marker = if *is_error { "ERROR" } else { "ok" };
                let id_short: String = tool_use_id.chars().take(12).collect();
                parts.push(format!("[tool_result {id_short}… {marker}: {truncated}]"));
            }
            MessagePart::Image { .. } => {
                parts.push("[image]".into());
            }
        }
    }
    parts.join(" ")
}

pub fn replace_range_with_summary(
    messages: &[Message],
    range: &CompactRange,
    summary: String,
    turn_id: crate::event::TurnId,
) -> Vec<Message> {
    let mut out = Vec::with_capacity(1 + messages.len().saturating_sub(range.end));
    out.push(Message::system_compact_summary(
        turn_id,
        summary,
        range.start as u64,
        range.end.saturating_sub(1) as u64,
        range.end - range.start,
    ));
    out.extend_from_slice(&messages[range.end..]);
    out
}

/// Result of compacting a messages_handle in place.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HandleCompactResult {
    pub before_tokens: u64,
    pub after_tokens: u64,
    pub compacted_start: usize,
    pub compacted_end: usize,
}

/// Compact a messages_handle in place (data-layer primitive, operates on any
/// FlowRun's segment). Returns `None` if under `budget` or no compactable
/// range. Caller should hold the FlowRun's `compact_lock`.
pub fn compact_messages_on_handle(
    handle: &std::sync::Arc<std::sync::Mutex<Vec<Message>>>,
    summary: String,
    budget: u64,
) -> Option<HandleCompactResult> {
    let mut msgs = handle.lock().unwrap();
    let before_tokens = estimate_tokens_for_messages(&msgs);
    let range = find_compact_range(&msgs, budget)?;
    let turn_id = msgs
        .get(range.start)
        .map(|m| m.turn_id.clone())
        .unwrap_or_else(crate::event::TurnId::now);
    let after = replace_range_with_summary(&msgs, &range, summary, turn_id);
    let after_tokens = estimate_tokens_for_messages(&after);
    if after_tokens >= before_tokens {
        return None;
    }
    let result = HandleCompactResult {
        before_tokens,
        after_tokens,
        compacted_start: range.start,
        compacted_end: range.end.saturating_sub(1),
    };
    *msgs = after;
    Some(result)
}

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

    fn user(text: &str) -> Message {
        Message::user_text(TurnId::now(), text)
    }
    fn assistant(text: &str) -> Message {
        Message::assistant_text(TurnId::now(), text)
    }
    fn system(text: &str) -> Message {
        Message::system_text(TurnId::now(), text)
    }

    #[test]
    fn summary_instructions_keep_decisions_before_next_move() {
        let objective = SUMMARY_INSTRUCTIONS
            .find("## Objective")
            .expect("objective");
        let decisions = SUMMARY_INSTRUCTIONS
            .find("## Decisions")
            .expect("decisions");
        let next_move = SUMMARY_INSTRUCTIONS
            .find("## Next Move")
            .expect("next move");
        assert!(objective < decisions);
        assert!(decisions < next_move);
    }

    #[test]
    fn estimate_scales_with_char_length() {
        let short = user("hi");
        let long = user(&"x".repeat(3500));
        assert!(estimate_tokens_for_message(&long) > estimate_tokens_for_message(&short) * 100);
    }

    #[test]
    fn find_compact_returns_none_when_under_budget() {
        let msgs = vec![user("a"), assistant("b"), user("c"), assistant("d")];
        assert!(find_compact_range(&msgs, 1000).is_none());
    }

    #[test]
    fn find_compact_returns_none_for_short_history() {
        let msgs = vec![user(&"x".repeat(9000))];
        assert!(find_compact_range(&msgs, 100).is_none());
    }

    #[test]
    fn find_kth_recent_user_handles_exact_excess_and_mixed_history() {
        let exact = vec![
            user("u0"),
            assistant("a0"),
            system("s0"),
            tool_result("call-0", "result", false),
            user("u1"),
            assistant("a1"),
            user("u2"),
            system("s1"),
            user("u3"),
            assistant("a3"),
            user("u4"),
        ];
        // With exactly five users, the fifth recent user is the first message.
        assert_eq!(find_kth_recent_user(&exact, KEEP_RECENT_USER_TURNS), 0);

        let mut excess = exact.clone();
        excess.push(user("u5"));
        assert_eq!(find_kth_recent_user(&excess, KEEP_RECENT_USER_TURNS), 4);

        let too_few = vec![user("a"), assistant("b"), assistant("c")];
        assert_eq!(find_kth_recent_user(&too_few, KEEP_RECENT_USER_TURNS), 0);
    }

    #[test]
    fn find_compact_range_preserves_minimum_recent_messages_without_anchor() {
        let mut msgs = vec![system("head")];
        msgs.extend((0..25).map(|index| assistant(&format!("old {index}"))));
        msgs.extend(
            (0..6).flat_map(|index| [user(&format!("user {index}")), assistant("assistant")]),
        );
        let range = find_compact_range(&msgs, 1).expect("range");
        assert_eq!(range.start, 0);
        assert_eq!(range.end, msgs.len() - KEEP_RECENT_MESSAGES);
    }

    #[test]
    fn find_compact_range_preserves_minimum_window_for_large_tool_result() {
        let mut msgs = vec![user(&"h".repeat(500_000))];
        for index in 1..31 {
            if matches!(index, 20 | 22 | 24 | 26 | 28 | 30) {
                msgs.push(user(&"u".repeat(800)));
            } else {
                msgs.push(assistant(&"a".repeat(800)));
            }
        }
        msgs.push(tool_result("call-large", &"t".repeat(22_000), false));

        let budget = 120_000;
        let minimum_recent_tokens = (budget as f64 * KEEP_RECENT_TOKEN_FRACTION).ceil() as u64;
        assert!(estimate_tokens_for_message(msgs.last().unwrap()) > minimum_recent_tokens);

        let range = find_compact_range(&msgs, budget).expect("range");
        assert!(
            range.end <= msgs.len() - KEEP_RECENT_MESSAGES,
            "range was {range:?}"
        );
        assert!(msgs.len() - range.end >= KEEP_RECENT_MESSAGES);
        assert!(estimate_tokens_for_messages(&msgs[range.end..]) >= minimum_recent_tokens);
    }

    #[test]
    fn find_compact_range_handles_four_to_twenty_one_message_histories() {
        for len in 4..=21 {
            let msgs = (0..len)
                .map(|_| user(&"x".repeat(5000)))
                .collect::<Vec<_>>();
            assert_eq!(
                find_compact_range(&msgs, 1).is_some(),
                len >= KEEP_RECENT_MESSAGES + 2,
                "len={len}"
            );
        }
    }

    #[test]
    fn find_compact_range_recent_users_limit_mixed_history() {
        let msgs = vec![
            system("head"),
            assistant("a0"),
            user("u0"),
            tool_result("call-0", "r0", false),
            assistant("a1"),
            user("u1"),
            assistant("a2"),
            tool_result("call-1", "r1", false),
            user("u2"),
            assistant("a3"),
            system("note"),
            user("u3"),
            tool_result("call-2", "r2", false),
            assistant("a4"),
            user("u4"),
            assistant("a5"),
            tool_result("call-3", "r3", false),
            user("u5"),
            assistant("a6"),
            system("tail"),
            assistant("a7"),
        ];

        let range = find_compact_range(&msgs, 1).expect("range");
        assert_eq!(range.end, 5);
        assert_eq!(msgs[range.end].role, MessageRole::User);
    }

    #[test]
    fn find_compact_range_preserves_recent_user_turns_after_anchor() {
        let mut msgs = vec![system("head"), compaction_summary("summary")];
        msgs.extend((0..6).flat_map(|index| {
            [
                user(&format!("user {index}")),
                assistant("assistant"),
                assistant("tool fragment"),
            ]
        }));
        msgs.extend((0..12).map(|_| assistant("recent fragment")));
        let range = find_compact_range(&msgs, 1).expect("range");
        assert_eq!(range.start, 1);
        assert_eq!(range.end, 5);
        assert_eq!(msgs[range.end].role, MessageRole::User);
    }

    #[test]
    fn find_compact_range_returns_none_when_end_cannot_cover_two_messages() {
        let msgs = vec![
            system("head"),
            compaction_summary("summary"),
            assistant("tail"),
            user("tail"),
        ];
        assert!(find_compact_range(&msgs, 1).is_none());
    }

    #[test]
    fn extract_anchor_removes_leading_compact_summary() {
        let messages = vec![compaction_summary("anchor"), user("new")];
        let (anchor, remaining) = extract_anchor(&messages).expect("anchor");
        assert_eq!(anchor, "anchor");
        assert_eq!(remaining, &messages[1..]);
    }

    #[test]
    fn extract_anchor_returns_none_without_leading_summary() {
        let messages = vec![user("new")];
        assert!(extract_anchor(&messages).is_none());
    }

    #[test]
    fn compact_messages_on_handle_replaces_range_in_place() {
        let mut messages = vec![system("head")];
        messages.extend((0..9).map(|index| assistant(&format!("old {index}"))));
        messages.extend(
            (0..6).flat_map(|index| [user(&format!("user {index}")), assistant(&"x".repeat(4000))]),
        );
        messages.extend((0..10).map(|index| assistant(&format!("recent {index}"))));
        messages.push(user("tail"));
        let handle: std::sync::Arc<std::sync::Mutex<Vec<Message>>> =
            std::sync::Arc::new(std::sync::Mutex::new(messages));
        // The recent-message and recent-user limits meet at the fifth recent user.
        // The compacted prefix is replaced by one summary while the tail remains.
        let result = compact_messages_on_handle(&handle, "gist".into(), 100);
        let result = result.expect("should compact");
        assert!(result.after_tokens < result.before_tokens);
        let msgs = handle.lock().unwrap();
        assert_eq!(result.compacted_start, 0);
        assert_eq!(result.compacted_end, 13);
        assert!(is_compaction_summary(&msgs[0]));
        assert_eq!(msgs.last().unwrap().text_concat(), "tail");
    }

    #[test]
    fn compact_messages_on_handle_none_when_under_budget() {
        let handle: std::sync::Arc<std::sync::Mutex<Vec<Message>>> =
            std::sync::Arc::new(std::sync::Mutex::new(vec![user("short")]));
        assert!(compact_messages_on_handle(&handle, "g".into(), 100_000).is_none());
    }

    #[test]
    fn compact_messages_on_handle_none_when_summary_would_not_shrink() {
        // 4 messages, all tiny → find_compact_range returns a range but the
        // summary message itself is comparable in size, so after >= before.
        // Construct a case where find_compact_range returns Some but shrink
        // check rejects it: make the range cover near-empty messages so the
        // summary overhead exceeds the savings.
        let handle: std::sync::Arc<std::sync::Mutex<Vec<Message>>> =
            std::sync::Arc::new(std::sync::Mutex::new(vec![
                system("h"),
                user("."),
                assistant("."),
                user("."),
                assistant("."),
                user("t"),
            ]));
        // budget=1 forces a range, but messages are so small the summary won't help
        let result = compact_messages_on_handle(&handle, "x".into(), 1);
        // Either no range found (len 6 but tiny), or shrink rejected.
        // The key invariant: handle is unchanged if None.
        let before_len = handle.lock().unwrap().len();
        if result.is_none() {
            assert_eq!(handle.lock().unwrap().len(), before_len);
        }
    }

    #[test]
    fn replace_range_puts_summary_system_message_in_place() {
        let msgs = vec![
            system("head"),
            user("m1"),
            assistant("m2"),
            user("m3"),
            assistant("m4"),
            user("tail"),
        ];
        let range = CompactRange {
            start: 1,
            end: 5,
            tokens_saved_estimate: 100,
        };
        let out = replace_range_with_summary(
            &msgs,
            &range,
            "gist: talked about m1..m4".into(),
            TurnId::now(),
        );
        assert_eq!(out.len(), 2, "summary + tail");
        assert_eq!(out[0].role, MessageRole::System);
        assert!(out[0].text_concat().contains("gist: talked about"));
        assert!(matches!(
            out[0].parts.as_slice(),
            [MessagePart::CompactSummary {
                seq_start: 1,
                seq_end: 4,
                count: 4,
                ..
            }]
        ));
        assert_eq!(out[1].role, MessageRole::User);
        assert_eq!(out[1].text_concat(), "tail");
    }

    #[test]
    fn find_compact_range_anchors_on_latest_structured_summary() {
        let mut msgs = vec![
            system("head"),
            Message::system_compact_summary(TurnId::now(), "old", 0, 1, 2),
        ];
        msgs.extend(
            (0..6).flat_map(|index| [user(&format!("user {index}")), assistant("assistant")]),
        );
        msgs.extend((0..12).map(|index| assistant(&format!("recent {index}"))));
        let range = find_compact_range(&msgs, 1).expect("range");
        assert_eq!(range.start, 1);
        assert_eq!(range.end, 4);
    }

    fn assistant_with_tool_use(text: &str, tool_name: &str, input: serde_json::Value) -> Message {
        Message {
            role: MessageRole::Assistant,
            parts: vec![
                MessagePart::Text { text: text.into() },
                MessagePart::ToolUse {
                    id: "call_test".into(),
                    name: tool_name.into(),
                    input,
                },
            ],
            turn_id: TurnId::now(),
        }
    }

    fn tool_result(id: &str, content: &str, is_error: bool) -> Message {
        Message {
            role: MessageRole::Tool,
            parts: vec![MessagePart::ToolResult {
                tool_use_id: id.into(),
                content: content.into(),
                is_error,
            }],
            turn_id: TurnId::now(),
        }
    }

    fn thinking(text: &str) -> Message {
        Message {
            role: MessageRole::Assistant,
            parts: vec![
                MessagePart::Thinking {
                    thinking: text.into(),
                    signature: None,
                },
                MessagePart::Text {
                    text: "after thinking".into(),
                },
            ],
            turn_id: TurnId::now(),
        }
    }

    #[test]
    fn format_slice_for_summary_includes_tool_use() {
        let slice = vec![
            user("read the file"),
            assistant_with_tool_use(
                "let me check",
                "fs.read",
                serde_json::json!({"path": "/tmp/foo.rs"}),
            ),
            tool_result("call_test", "fn main() {}", false),
        ];
        let out = format_slice_for_summary(&slice);
        assert!(out.contains("fs.read"), "missing tool name: {out}");
        assert!(out.contains("/tmp/foo.rs"), "missing tool input: {out}");
        assert!(
            out.contains("fn main()"),
            "missing tool_result content: {out}"
        );
        assert!(out.contains("tool_call"), "missing tool_call marker: {out}");
        assert!(
            out.contains("tool_result"),
            "missing tool_result marker: {out}"
        );
    }

    #[test]
    fn format_slice_for_summary_includes_thinking() {
        let slice = vec![thinking("I should consider the edge case")];
        let out = format_slice_for_summary(&slice);
        assert!(out.contains("thinking"), "missing thinking marker: {out}");
        assert!(out.contains("edge case"), "missing thinking content: {out}");
    }

    #[test]
    fn format_slice_for_summary_marks_error_tool_results() {
        let slice = vec![tool_result("call_1", "permission denied", true)];
        let out = format_slice_for_summary(&slice);
        assert!(out.contains("ERROR"), "missing ERROR marker: {out}");
    }

    #[test]
    fn format_slice_for_summary_truncates_long_tool_input() {
        let long_input = serde_json::json!({"content": "x".repeat(5000)});
        let slice = vec![assistant_with_tool_use("check", "fs.write", long_input)];
        let out = format_slice_for_summary(&slice);
        let tool_call_line = out
            .lines()
            .find(|l| l.contains("tool_call"))
            .unwrap_or_else(|| panic!("no tool_call line in {out}"));
        assert!(
            tool_call_line.chars().count() < 2200,
            "tool_call line not truncated: {tool_call_line}"
        );
    }

    fn compaction_summary(text: &str) -> Message {
        Message::system_compact_summary(TurnId::now(), text, 1, 5, 5)
    }

    #[test]
    fn is_compaction_summary_detects_structured_variant() {
        assert!(is_compaction_summary(&compaction_summary("gist")));
        assert!(!is_compaction_summary(&system("plain system msg")));
        assert!(!is_compaction_summary(&user("user msg")));
    }

    #[test]
    fn find_compact_range_spans_across_compaction_summaries() {
        let mut msgs = vec![
            system("head"),
            user(&"x".repeat(3000)),
            assistant(&"y".repeat(3000)),
            user(&"z".repeat(3000)),
            compaction_summary("first compaction summary"),
        ];
        msgs.extend(
            (0..6).flat_map(|index| [user(&format!("user {index}")), assistant(&"x".repeat(3000))]),
        );
        msgs.extend((0..10).map(|index| assistant(&format!("recent {index}"))));
        let range = find_compact_range(&msgs, 500).expect("expected range across summary");
        assert_eq!(
            range.start, 4,
            "range should anchor at the structured summary"
        );
        assert!(
            range.end > 4,
            "range should include later work, got {range:?}"
        );
        assert!(
            range.end - range.start >= 3,
            "range must cover >= 3 msgs, got {}",
            range.end - range.start
        );
    }

    #[test]
    fn find_compact_starts_from_summary() {
        let mut msgs = vec![user("a"), assistant("b"), compaction_summary("summary 1")];
        msgs.extend(
            (0..6).flat_map(|index| [user(&format!("user {index}")), assistant("assistant")]),
        );
        msgs.extend((0..12).map(|index| assistant(&format!("recent {index}"))));
        let range = find_compact_range(&msgs, 10).expect("expected range");
        assert_eq!(
            range.start, 2,
            "should start from the compact summary anchor"
        );
        assert_eq!(range.end, 5, "the fifth recent user is retained");
    }

    #[test]
    fn find_compact_range_includes_older_compaction_summaries() {
        let mut msgs = vec![compaction_summary("summary 0")];
        msgs.extend((0..6).flat_map(|index| {
            [
                user(&format!("old user {index}")),
                assistant(&"x".repeat(2000)),
            ]
        }));
        msgs.push(compaction_summary("summary 1"));
        msgs.extend((0..6).flat_map(|index| {
            [
                user(&format!("new user {index}")),
                assistant(&"z".repeat(2000)),
            ]
        }));
        msgs.extend((0..10).map(|index| assistant(&format!("tail {index}"))));
        let range = find_compact_range(&msgs, 500).expect("expected range");
        assert_eq!(range.start, 13, "should compact from the latest summary");
        assert!(
            range.end > range.start,
            "should include work after the latest summary"
        );
    }

    #[test]
    fn compacted_message_tokens_detects_growth() {
        let msgs = vec![compaction_summary("summary 0"), user("a"), assistant("b")];
        let range = CompactRange {
            start: 1,
            end: 3,
            tokens_saved_estimate: 0,
        };
        let before = estimate_tokens_for_messages(&msgs);
        let after = estimate_compacted_message_tokens(
            &msgs,
            &range,
            "a very long summary that expands the transcript a lot",
        );
        assert!(after > before, "expected growth to be detectable");
    }

    #[test]
    fn find_compact_starts_from_zero_without_summary() {
        let mut msgs = (0..26)
            .map(|index| assistant(&format!("old {index}")))
            .collect::<Vec<_>>();
        msgs.extend(
            (0..6).flat_map(|index| [user(&format!("user {index}")), assistant("assistant")]),
        );
        let range = find_compact_range(&msgs, 10).expect("expected range");
        assert_eq!(range.start, 0, "should start from 0 without summary");
        assert_eq!(range.end, 28, "the recent-message limit is retained");
    }

    #[test]
    fn filter_orphan_tool_messages_removes_orphan_results() {
        use crate::message::{Message, MessagePart, MessageRole};
        let turn = TurnId::now();
        let msgs = vec![
            Message {
                role: MessageRole::Tool,
                parts: vec![MessagePart::ToolResult {
                    tool_use_id: "orphan".into(),
                    content: "no matching use".into(),
                    is_error: false,
                }],
                turn_id: turn.clone(),
            },
            Message {
                role: MessageRole::Assistant,
                parts: vec![MessagePart::ToolUse {
                    id: "call_1".into(),
                    name: "fs.read".into(),
                    input: serde_json::json!({}),
                }],
                turn_id: turn.clone(),
            },
            Message {
                role: MessageRole::Tool,
                parts: vec![MessagePart::ToolResult {
                    tool_use_id: "call_1".into(),
                    content: "ok".into(),
                    is_error: false,
                }],
                turn_id: turn,
            },
        ];
        let mut filtered = msgs;
        filter_orphan_tool_messages(&mut filtered);
        assert_eq!(filtered.len(), 2, "orphan result should be removed");
    }
}