yapi_core/

compaction.rs

//! Compaction: summarize older context so a long session fits the model.
//!
//! Port of `packages/coding-agent/src/core/compaction/compaction.ts` and
//! `compaction/utils.ts` in pi `v1.0.0`. The functions here are pure apart from
//! the summary request. The agent session decides when to compact and records the
//! result.

use std::collections::BTreeSet;

use serde::Serialize;
use serde_json::{Value, json};
use tokio_util::sync::CancellationToken;
use yapi_ai::api::Apis;
use yapi_ai::errors::{is_retryable_assistant_error, retry_delay_ms};
use yapi_ai::registry::Auth;
use yapi_ai::stream::{CacheRetention, Request, StreamOptions};
use yapi_types::message::{
    AssistantMessage, Content, ContentBlock, Message, StopReason, SystemMessage, ThinkingLevel,
    Usage, UserMessage,
};
use yapi_types::model::Model;
use yapi_types::session::FileEntry;
use yapi_types::settings::Settings;

use crate::messages::convert_to_llm;
use crate::session::{ProjectedEntry, Projection, entry_messages};
use crate::time::{now_ms, uuid_v7};

/// When and how much to compact.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CompactionSettings {
    /// Automatic compaction is on.
    pub enabled: bool,
    /// Tokens kept free for the response. Compaction starts when less is left.
    pub reserve_tokens: u64,
    /// Recent tokens kept verbatim.
    pub keep_recent_tokens: u64,
}

impl CompactionSettings {
    /// The settings for `model`: its override, the ordinary setting, the default.
    pub fn resolve(settings: &Settings, model: Option<&Model>) -> CompactionSettings {
        let compaction = settings.compaction.as_ref();
        let key = model.map(|model| format!("{}/{}", model.provider, model.id));
        let overrides = key.as_ref().and_then(|key| {
            compaction
                .and_then(|c| c.model_overrides.as_ref())
                .and_then(|overrides| overrides.get(key))
        });
        CompactionSettings {
            enabled: compaction.and_then(|c| c.enabled).unwrap_or(true),
            reserve_tokens: overrides
                .and_then(|o| o.reserve_tokens)
                .or_else(|| compaction.and_then(|c| c.reserve_tokens))
                .unwrap_or(16384),
            keep_recent_tokens: overrides
                .and_then(|o| o.keep_recent_tokens)
                .or_else(|| compaction.and_then(|c| c.keep_recent_tokens))
                .unwrap_or(20000),
        }
    }
}

/// Retry policy for summary requests, from `settings.retry`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RetryPolicy {
    /// Retries are on.
    pub enabled: bool,
    /// Retries after the first attempt.
    pub max_retries: u32,
    /// Delay before the first retry, doubling each time.
    pub base_delay_ms: u64,
    /// Cap on one delay.
    pub max_delay_ms: u64,
}

impl RetryPolicy {
    /// The policy in `settings`, with pi's defaults.
    pub fn resolve(settings: &Settings) -> RetryPolicy {
        let retry = settings.retry.as_ref();
        RetryPolicy {
            enabled: retry.and_then(|r| r.enabled).unwrap_or(true),
            max_retries: retry.and_then(|r| r.max_retries).unwrap_or(3),
            base_delay_ms: retry.and_then(|r| r.base_delay_ms).unwrap_or(2000),
            max_delay_ms: retry
                .and_then(|r| r.max_agent_delay_ms)
                .unwrap_or(yapi_ai::errors::DEFAULT_MAX_AGENT_RETRY_DELAY_MS),
        }
    }

    /// The delay before retry `attempt` (1-based).
    pub fn delay_ms(&self, attempt: u32) -> u64 {
        retry_delay_ms(self.base_delay_ms, Some(self.max_delay_ms), attempt)
    }
}

/// Tokens a response's usage says the context held.
pub fn calculate_context_tokens(usage: &Usage) -> u64 {
    match usage.total_tokens {
        Some(total) if total > 0 => total,
        _ => usage.input + usage.output + usage.cache_read + usage.cache_write,
    }
}

fn assistant_usage(message: &Message) -> Option<&Usage> {
    match message {
        Message::Assistant(assistant)
            if !matches!(
                assistant.stop_reason,
                StopReason::Aborted | StopReason::Error
            ) && calculate_context_tokens(&assistant.usage) > 0 =>
        {
            Some(&assistant.usage)
        }
        _ => None,
    }
}

const ESTIMATED_IMAGE_CHARS: usize = 4800;

fn blocks_chars(blocks: &[ContentBlock]) -> usize {
    blocks
        .iter()
        .map(|block| match block {
            ContentBlock::Text(text) => yapi_types::js::len(&text.text),
            ContentBlock::Image(_) => ESTIMATED_IMAGE_CHARS,
            _ => 0,
        })
        .sum()
}

fn content_chars(content: &Content) -> usize {
    match content {
        Content::Text(text) => yapi_types::js::len(text),
        Content::Blocks(blocks) => blocks_chars(blocks),
    }
}

/// Estimated tokens of one message: characters / 4, rounded up, where an image
/// counts as 4800 characters.
pub fn estimate_tokens(message: &Message) -> u64 {
    let chars = match message {
        Message::System(system) => {
            let sections: usize = system
                .sections
                .iter()
                .flatten()
                .filter_map(|(_, text)| text.as_deref())
                .map(yapi_types::js::len)
                .sum();
            let tools = system
                .tools_added
                .as_ref()
                .and_then(|tools| yapi_types::json::to_string(tools).ok())
                .map_or(0, |json| yapi_types::js::len(&json));
            content_chars(&system.content) + sections + tools
        }
        Message::User(user) => content_chars(&user.content),
        Message::Assistant(assistant) => assistant
            .content
            .iter()
            .map(|block| match block {
                ContentBlock::Text(text) => yapi_types::js::len(&text.text),
                ContentBlock::Thinking(thinking) => yapi_types::js::len(&thinking.thinking),
                ContentBlock::ToolCall(call) => {
                    yapi_types::js::len(&call.name)
                        + yapi_types::js::len(&yapi_types::json::stringify(&call.arguments))
                }
                ContentBlock::Image(_) => 0,
            })
            .sum(),
        Message::Custom(custom) => content_chars(&custom.content),
        Message::ToolResult(result) => blocks_chars(&result.content),
        Message::BashExecution(bash) => {
            yapi_types::js::len(&bash.command) + yapi_types::js::len(&bash.output)
        }
        Message::BranchSummary(summary) => yapi_types::js::len(&summary.summary),
        Message::CompactionSummary(summary) => yapi_types::js::len(&summary.summary),
    };
    chars.div_ceil(4) as u64
}

/// Estimated context size: the last valid usage plus estimates for what follows.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ContextEstimate {
    /// Total tokens.
    pub tokens: u64,
    /// Tokens from the last usage.
    pub usage_tokens: u64,
    /// Estimated tokens after it.
    pub trailing_tokens: u64,
    /// Index of the message whose usage was used.
    pub last_usage_index: Option<usize>,
}

/// Estimates context tokens from the last valid usage and the messages after it.
pub fn estimate_context_tokens(messages: &[Message]) -> ContextEstimate {
    let usage = messages
        .iter()
        .enumerate()
        .rev()
        .find_map(|(index, message)| assistant_usage(message).map(|usage| (index, usage)));
    match usage {
        None => {
            let tokens = messages.iter().map(estimate_tokens).sum();
            ContextEstimate {
                tokens,
                usage_tokens: 0,
                trailing_tokens: tokens,
                last_usage_index: None,
            }
        }
        Some((index, usage)) => {
            let usage_tokens = calculate_context_tokens(usage);
            let trailing: u64 = messages[index + 1..].iter().map(estimate_tokens).sum();
            ContextEstimate {
                tokens: usage_tokens + trailing,
                usage_tokens,
                trailing_tokens: trailing,
                last_usage_index: Some(index),
            }
        }
    }
}

/// Estimates projected context without trusting usage recorded before a later
/// context edit or compaction.
pub fn estimate_projected_context_tokens(
    projection: &Projection<'_>,
    branch: &[&FileEntry],
) -> ContextEstimate {
    let estimate = estimate_context_tokens(&projection.messages);
    if let Some(usage_index) = estimate.last_usage_index {
        let mut next = 0;
        let mut usage_entry = None;
        for entry in &projection.entries {
            next += entry.messages.len();
            if usage_index < next {
                usage_entry = entry.source.meta().map(|meta| meta.id.as_str());
                break;
            }
        }
        let usage_position = usage_entry.and_then(|id| {
            branch
                .iter()
                .position(|entry| entry.meta().is_some_and(|meta| meta.id == id))
        });
        let invalidating = branch.iter().rposition(|entry| {
            matches!(entry, FileEntry::ContextEdit(_) | FileEntry::Compaction(_))
        });
        let after = match (usage_position, invalidating) {
            (Some(usage), Some(invalid)) => usage > invalid,
            (Some(_), None) => true,
            (None, _) => false,
        };
        if after {
            return estimate;
        }
    }
    let system = yapi_ai::transcript::current_system_message(&projection.messages);
    let mut tokens = system
        .map(|system| estimate_tokens(&Message::System(system)))
        .unwrap_or(0);
    tokens += projection
        .messages
        .iter()
        .filter(|message| !matches!(message, Message::System(_)))
        .map(estimate_tokens)
        .sum::<u64>();
    ContextEstimate {
        tokens,
        usage_tokens: 0,
        trailing_tokens: tokens,
        last_usage_index: None,
    }
}

/// Whether the context leaves less than the reserve free.
pub fn should_compact(tokens: u64, context_window: u64, settings: &CompactionSettings) -> bool {
    settings.enabled && (tokens as i128) > context_window as i128 - settings.reserve_tokens as i128
}

fn is_cut_point(message: &Message) -> bool {
    matches!(
        message,
        Message::User(_)
            | Message::Assistant(_)
            | Message::BashExecution(_)
            | Message::Custom(_)
            | Message::BranchSummary(_)
            | Message::CompactionSummary(_)
    )
}

fn is_turn_start(message: &Message) -> bool {
    matches!(
        message,
        Message::User(_)
            | Message::BashExecution(_)
            | Message::Custom(_)
            | Message::BranchSummary(_)
            | Message::CompactionSummary(_)
    )
}

fn is_compaction(entry: &ProjectedEntry<'_>) -> bool {
    matches!(entry.source, FileEntry::Compaction(_))
}

fn starts_turn(entry: &ProjectedEntry<'_>) -> bool {
    !is_compaction(entry) && entry.messages.iter().any(is_turn_start)
}

/// Where to cut the projected entries `start..end`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct CutPoint {
    first_kept: usize,
    turn_start: Option<usize>,
    split_turn: bool,
}

fn find_cut_point(
    entries: &[ProjectedEntry<'_>],
    start: usize,
    end: usize,
    keep_recent_tokens: u64,
) -> CutPoint {
    let cut_points: Vec<usize> = (start..end)
        .filter(|&index| {
            !is_compaction(&entries[index]) && entries[index].messages.iter().any(is_cut_point)
        })
        .collect();
    let Some(&first) = cut_points.first() else {
        return CutPoint {
            first_kept: start,
            turn_start: None,
            split_turn: false,
        };
    };
    let mut accumulated = 0;
    let mut exceeded = false;
    let mut cut = first;
    for index in (start..end).rev() {
        let tokens: u64 = entries[index].messages.iter().map(estimate_tokens).sum();
        if tokens == 0 {
            continue;
        }
        accumulated += tokens;
        if accumulated >= keep_recent_tokens {
            exceeded = true;
            cut = cut_points
                .iter()
                .copied()
                .find(|&candidate| candidate >= index)
                .unwrap_or(cut_points[cut_points.len() - 1]);
            break;
        }
    }

    // A recovery attempt and its omission edits are invisible after the last
    // visible input. Move past them only when the whole suffix is such an attempt.
    let suffix = &entries[(cut + 1).min(end)..end];
    let intrinsically_visible = |entry: &ProjectedEntry<'_>| {
        !matches!(entry.source, FileEntry::ContextEdit(_))
            && !entry_messages(entry.source).is_empty()
    };
    let omitted =
        |entry: &ProjectedEntry<'_>| intrinsically_visible(entry) && entry.messages.is_empty();
    let omitted_ids: Vec<&str> = suffix
        .iter()
        .filter(|entry| omitted(entry))
        .filter_map(|entry| entry.source.meta().map(|meta| meta.id.as_str()))
        .collect();
    let external_replacement = suffix.iter().any(|entry| match entry.source {
        FileEntry::ContextEdit(edit) => {
            edit.replacement.is_some() && !omitted_ids.contains(&edit.target_id.as_str())
        }
        _ => false,
    });
    let recovery_suffix = exceeded && !external_replacement && suffix.iter().any(|entry| {
        matches!(entry.source, FileEntry::Message(m) if matches!(m.message, Message::Assistant(_)))
            && omitted(entry)
    })
        && suffix.iter().all(|entry| {
            !is_compaction(entry) && (!intrinsically_visible(entry) || omitted(entry))
        });
    if recovery_suffix {
        cut += 1;
    }

    while cut > start {
        let previous = &entries[cut - 1];
        if is_compaction(previous) || !previous.messages.is_empty() {
            break;
        }
        cut -= 1;
    }
    let turn_start = if entries.get(cut).is_some_and(starts_turn) {
        None
    } else {
        (start..=cut.min(end.saturating_sub(1)))
            .rev()
            .find(|&index| starts_turn(&entries[index]))
    };
    CutPoint {
        first_kept: cut,
        turn_start,
        split_turn: turn_start.is_some(),
    }
}

/// Files read and changed by tool calls, for the summary.
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct FileOperations {
    /// Read paths.
    pub read: BTreeSet<String>,
    /// Written paths.
    pub written: BTreeSet<String>,
    /// Edited paths.
    pub edited: BTreeSet<String>,
}

impl FileOperations {
    fn add(&mut self, tool: &str, arguments: &serde_json::Map<String, Value>) {
        let Some(path) = arguments
            .get("path")
            .and_then(Value::as_str)
            .filter(|p| !p.is_empty())
        else {
            return;
        };
        let set = match tool {
            "read" => &mut self.read,
            "write" => &mut self.written,
            "edit" => &mut self.edited,
            _ => return,
        };
        set.insert(path.to_owned());
    }

    /// Adds the `readFiles` and `modifiedFiles` lists of an earlier summary's
    /// details.
    fn add_details(&mut self, details: &Value) {
        for (key, set) in [
            ("readFiles", &mut self.read),
            ("modifiedFiles", &mut self.edited),
        ] {
            let paths = details[key].as_array().into_iter().flatten();
            set.extend(paths.filter_map(Value::as_str).map(str::to_owned));
        }
    }

    /// Records the file tools a message called.
    pub fn extract(&mut self, message: &Message) {
        match message {
            Message::ToolResult(result) => {
                for call in result.nested_calls.iter().flat_map(|nested| &nested.calls) {
                    if let Some(arguments) = &call.arguments {
                        self.add(&call.name, arguments);
                    }
                }
            }
            Message::Assistant(assistant) => {
                for block in &assistant.content {
                    if let ContentBlock::ToolCall(call) = block {
                        self.add(&call.name, &call.arguments);
                    }
                }
            }
            _ => {}
        }
    }

    /// Paths only read, and paths changed, each sorted.
    pub fn lists(&self) -> (Vec<String>, Vec<String>) {
        let modified: BTreeSet<&String> = self.edited.iter().chain(&self.written).collect();
        let read = self
            .read
            .iter()
            .filter(|path| !modified.contains(path))
            .cloned()
            .collect();
        (read, modified.into_iter().cloned().collect())
    }
}

/// The file lists as tags appended to a summary.
pub fn format_file_operations(read: &[String], modified: &[String]) -> String {
    let mut sections = Vec::new();
    if !read.is_empty() {
        sections.push(format!("<read-files>\n{}\n</read-files>", read.join("\n")));
    }
    if !modified.is_empty() {
        sections.push(format!(
            "<modified-files>\n{}\n</modified-files>",
            modified.join("\n")
        ));
    }
    if sections.is_empty() {
        String::new()
    } else {
        format!("\n\n{}", sections.join("\n\n"))
    }
}

const TOOL_RESULT_MAX_CHARS: usize = 2000;

fn truncate_for_summary(text: &str, max: usize) -> String {
    let length = yapi_types::js::len(text);
    if length <= max {
        return text.to_owned();
    }
    format!(
        "{}\n\n[... {} more characters truncated]",
        yapi_types::js::slice(text, 0, max),
        length - max
    )
}

/// The conversation as plain text, so the summarizer does not continue it.
/// Takes messages already converted for the model.
pub fn serialize_conversation(messages: &[Message]) -> String {
    let mut parts = Vec::new();
    for message in messages {
        match message {
            Message::User(user) => {
                let text = user.content.text("");
                if !text.is_empty() {
                    parts.push(format!("[User]: {text}"));
                }
            }
            Message::Assistant(assistant) => {
                let mut thinking = Vec::new();
                let mut calls = Vec::new();
                for block in &assistant.content {
                    match block {
                        ContentBlock::Thinking(block) => thinking.push(block.thinking.as_str()),
                        ContentBlock::ToolCall(call) => {
                            let arguments: Vec<String> = call
                                .arguments
                                .iter()
                                .map(|(key, value)| {
                                    format!("{key}={}", yapi_types::json::stringify(value))
                                })
                                .collect();
                            calls.push(format!("{}({})", call.name, arguments.join(", ")));
                        }
                        _ => {}
                    }
                }
                if !thinking.is_empty() {
                    parts.push(format!("[Assistant thinking]: {}", thinking.join("\n")));
                }
                if assistant
                    .content
                    .iter()
                    .any(|block| matches!(block, ContentBlock::Text(_)))
                {
                    parts.push(format!(
                        "[Assistant]: {}",
                        yapi_types::message::blocks_text(&assistant.content, "\n")
                    ));
                }
                if !calls.is_empty() {
                    parts.push(format!("[Assistant tool calls]: {}", calls.join("; ")));
                }
            }
            Message::ToolResult(result) => {
                let text = yapi_types::message::blocks_text(&result.content, "");
                if !text.is_empty() {
                    parts.push(format!(
                        "[Tool result]: {}",
                        truncate_for_summary(&text, TOOL_RESULT_MAX_CHARS)
                    ));
                }
            }
            _ => {}
        }
    }
    parts.join("\n\n")
}

/// System prompt of summary requests.
pub const SUMMARIZATION_SYSTEM_PROMPT: &str = "You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified.\n\nDo NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.";

const SUMMARIZATION_PROMPT: &str = "The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.\n\nUse this EXACT format:\n\n## Goal\n[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]\n\n## Constraints & Preferences\n- [Any constraints, preferences, or requirements mentioned by user]\n- [Or \"(none)\" if none were mentioned]\n\n## Progress\n### Done\n- [x] [Completed tasks/changes]\n\n### In Progress\n- [ ] [Current work]\n\n### Blocked\n- [Issues preventing progress, if any]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale]\n\n## Next Steps\n1. [Ordered list of what should happen next]\n\n## Critical Context\n- [Any data, examples, or references needed to continue]\n- [Or \"(none)\" if not applicable]\n\nKeep each section concise. Preserve exact file paths, function names, and error messages.";

const UPDATE_SUMMARIZATION_INSTRUCTIONS: &str = "Update the existing structured summary with new information. RULES:\n- PRESERVE all existing information from the previous summary\n- ADD new progress, decisions, and context from the new messages\n- UPDATE the Progress section: move items from \"In Progress\" to \"Done\" when completed\n- UPDATE \"Next Steps\" based on what was accomplished\n- PRESERVE exact file paths, function names, and error messages\n- If something is no longer relevant, you may remove it\n\nUse this EXACT format:\n\n## Goal\n[Preserve existing goals, add new ones if the task expanded]\n\n## Constraints & Preferences\n- [Preserve existing, add new ones discovered]\n\n## Progress\n### Done\n- [x] [Include previously done items AND newly completed items]\n\n### In Progress\n- [ ] [Current work - update based on progress]\n\n### Blocked\n- [Current blockers - remove if resolved]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale] (preserve all previous, add new)\n\n## Next Steps\n1. [Update based on current state]\n\n## Critical Context\n- [Preserve important context, add new if needed]\n\nKeep each section concise. Preserve exact file paths, function names, and error messages.";

const TURN_PREFIX_SUMMARIZATION_PROMPT: &str = "The messages above are earlier context from an ongoing conversation. Later messages are stored separately and do not need to be reconstructed.\n\nCreate a concise checkpoint of the user's request and the progress shown above. This checkpoint will be placed before the later messages so the conversation can continue with the necessary context.\n\n## Original Request\n[What did the user ask for?]\n\n## Progress So Far\n- [Key decisions and work completed in these messages]\n\n## Context Needed to Continue\n- [Information from these messages needed to understand the later work]\n\nOnly summarize information explicitly present above. Do not infer or recreate later messages.";

/// What compaction will summarize and keep, serialized as pi's
/// `CompactionPreparation`.
#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct Preparation {
    /// First entry kept verbatim.
    pub first_kept_entry_id: String,
    /// Messages summarized and dropped.
    pub messages_to_summarize: Vec<Message>,
    /// The start of a turn split by the cut, summarized separately.
    pub turn_prefix_messages: Vec<Message>,
    /// The cut falls inside a turn.
    pub is_split_turn: bool,
    /// Estimated context tokens before compaction.
    pub tokens_before: u64,
    /// The previous compaction's summary, to update.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub previous_summary: Option<String>,
    /// Files touched by the summarized messages.
    pub file_ops: FileOperations,
    /// The settings used.
    pub settings: CompactionSettings,
}

/// Plans a compaction of a branch, or `None` when there is nothing to summarize.
pub fn prepare_compaction(
    branch: &[&FileEntry],
    settings: CompactionSettings,
) -> Option<Preparation> {
    if matches!(branch.last(), Some(FileEntry::Compaction(_))) {
        return None;
    }
    let projection = crate::session::build_projection(branch);
    let entries = &projection.entries;
    // The newest compaction is projected first. Older ones contribute nothing.
    let previous = entries
        .iter()
        .position(|entry| is_compaction(entry) && !entry.messages.is_empty());
    let (previous_summary, start) = match previous {
        Some(index) => match entries[index].source {
            FileEntry::Compaction(compaction) => (Some(compaction.summary.clone()), index + 1),
            _ => (None, 0),
        },
        None => (None, 0),
    };
    let end = entries.len();
    let tokens_before = estimate_projected_context_tokens(&projection, branch).tokens;
    let cut = find_cut_point(entries, start, end, settings.keep_recent_tokens);
    let first_kept_entry_id = entries
        .get(cut.first_kept)?
        .source
        .meta()
        .map(|meta| meta.id.clone())
        .filter(|id| !id.is_empty())?;
    let history_end = if cut.split_turn {
        cut.turn_start.unwrap_or(cut.first_kept)
    } else {
        cut.first_kept
    };
    let summarizable = |range: std::ops::Range<usize>| -> Vec<Message> {
        entries[range]
            .iter()
            .filter(|entry| !is_compaction(entry))
            .flat_map(|entry| {
                entry
                    .messages
                    .iter()
                    .filter(|message| !matches!(message, Message::System(_)))
                    .cloned()
            })
            .collect()
    };
    let messages_to_summarize = summarizable(start..history_end.max(start));
    let turn_prefix_messages = match (cut.split_turn, cut.turn_start) {
        (true, Some(turn_start)) => summarizable(turn_start..cut.first_kept),
        _ => Vec::new(),
    };
    if messages_to_summarize.is_empty() && turn_prefix_messages.is_empty() {
        return None;
    }
    let mut file_ops = FileOperations::default();
    if let Some(index) = previous
        && let FileEntry::Compaction(compaction) = entries[index].source
        && compaction.from_hook != Some(true)
        && let Some(details) = &compaction.details
    {
        file_ops.add_details(details);
    }
    for message in messages_to_summarize.iter().chain(&turn_prefix_messages) {
        file_ops.extract(message);
    }
    Some(Preparation {
        first_kept_entry_id,
        messages_to_summarize,
        turn_prefix_messages,
        is_split_turn: cut.split_turn,
        tokens_before,
        previous_summary,
        file_ops,
        settings,
    })
}

/// What summary requests need besides the messages.
pub struct Summarizer<'a> {
    /// The model to summarize with.
    pub model: &'a Model,
    /// Wire APIs.
    pub apis: Apis,
    /// Credentials for the model.
    pub auth: Auth,
    /// The session's thinking level, used when the model reasons.
    pub thinking_level: ThinkingLevel,
    /// Retries for transient failures.
    pub retry: RetryPolicy,
    /// The settings every request carries, as the session's requests carry
    /// them, or the error an invalid one raises.
    pub request: Result<StreamOptions, String>,
    /// Cancels the requests.
    pub cancel: CancellationToken,
    /// Hears about retries.
    pub on_retry: Box<dyn Fn(SummaryRetry) + Send + Sync + 'a>,
    /// Observers of the requests.
    pub hooks: yapi_ai::stream::RequestHooks,
}

/// pi-ai's `RetryCallbacks`: what a summarizer reports while it retries.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SummaryRetry {
    /// A failed request is retried after `delay_ms`.
    Scheduled {
        /// The retry's number, from 1.
        attempt: u32,
        /// How many retries the policy allows.
        max_attempts: u32,
        /// The wait before it.
        delay_ms: u64,
        /// Why the request failed.
        error_message: String,
    },
    /// The retried request starts.
    AttemptStart,
    /// Retrying ended, in success or not.
    Finished,
}

/// A summary and the usage of the requests that made it.
#[derive(Clone, Debug, PartialEq)]
pub struct Summary {
    /// The text.
    pub text: String,
    /// Usage of the summary requests.
    pub usage: Usage,
}

/// Why a summary response cannot become a checkpoint, if it cannot: it
/// failed, reached the token cap or called a tool.
pub(crate) fn check_summary(response: &AssistantMessage, label: &str) -> Result<(), String> {
    match response.stop_reason {
        StopReason::Error => Err(format!(
            "{label} failed: {}",
            response
                .error_message
                .as_deref()
                .filter(|m| !m.is_empty())
                .unwrap_or("Unknown error")
        )),
        StopReason::Length => Err(format!(
            "{label} failed: generation hit the token cap and the summary is incomplete"
        )),
        _ if response
            .content
            .iter()
            .any(|block| matches!(block, ContentBlock::ToolCall(_))) =>
        {
            Err(format!("{label} attempted to call a tool"))
        }
        _ => Ok(()),
    }
}

impl Summarizer<'_> {
    async fn complete_once(
        &self,
        system: &str,
        prompt: &str,
        max_tokens: u64,
        session_id: &str,
    ) -> AssistantMessage {
        let messages = vec![
            Message::System(SystemMessage {
                content: Content::Text(system.to_owned()),
                sections: None,
                timestamp: 0,
                tools_added: None,
                tools_removed: None,
            }),
            Message::User(UserMessage {
                content: Content::Blocks(vec![ContentBlock::text(prompt)]),
                timestamp: now_ms(),
            }),
        ];
        let reasoning = (self.model.reasoning && self.thinking_level != ThinkingLevel::Off)
            .then_some(self.thinking_level);
        let options = StreamOptions {
            reasoning,
            max_tokens: Some(max_tokens),
            session_id: Some(session_id.to_owned()),
            cache_retention: Some(CacheRetention::None),
            cancel: self.cancel.clone(),
            hooks: self.hooks.clone(),
            ..self.request.clone().unwrap_or_default()
        };
        let mut request = Request {
            model: self.model.clone(),
            messages,
            options,
        };
        let checked = self.request.as_ref().map(|_| ()).map_err(Clone::clone);
        let stream = match checked.and_then(|()| self.auth.clone().apply(&mut request)) {
            Ok(()) => {
                yapi_ai::stream::RequestHooks::prepare(&mut request).await;
                self.apis.stream(request)
            }
            Err(message) => {
                yapi_ai::api::failed_stream(&request.model, &request.options.cancel, message)
            }
        };
        match stream.result().await {
            Some(message) => message,
            None => {
                let mut message = yapi_ai::stream::new_output(self.model, now_ms());
                message.stop_reason = StopReason::Error;
                message.error_message = Some("Stream ended without a result".into());
                message
            }
        }
    }

    /// One summary request under `system` with retries on transient
    /// failures, reported as pi-ai's `retryAssistantCall` reports them.
    pub(crate) async fn complete(
        &self,
        system: &str,
        prompt: &str,
        max_tokens: u64,
    ) -> AssistantMessage {
        let session_id = uuid_v7();
        let max_attempts = if self.retry.enabled {
            self.retry.max_retries
        } else {
            0
        };
        let notify = &self.on_retry;
        let mut attempt = 0;
        loop {
            let response = self
                .complete_once(system, prompt, max_tokens, &session_id)
                .await;
            if response.stop_reason != StopReason::Error
                || attempt >= max_attempts
                || !is_retryable_assistant_error(&response)
            {
                if attempt > 0 {
                    notify(SummaryRetry::Finished);
                }
                return response;
            }
            attempt += 1;
            let delay_ms = self.retry.delay_ms(attempt);
            notify(SummaryRetry::Scheduled {
                attempt,
                max_attempts,
                delay_ms,
                error_message: response
                    .error_message
                    .clone()
                    .filter(|message| !message.is_empty())
                    .unwrap_or_else(|| "Unknown error".into()),
            });
            tokio::select! {
                () = self.cancel.cancelled() => {
                    notify(SummaryRetry::Finished);
                    let mut response = response;
                    response.stop_reason = StopReason::Aborted;
                    response.error_message = None;
                    return response;
                }
                () = tokio::time::sleep(std::time::Duration::from_millis(delay_ms)) => {}
            }
            notify(SummaryRetry::AttemptStart);
        }
    }

    async fn summarize(
        &self,
        prompt: &str,
        max_tokens: u64,
        label: &str,
    ) -> Result<Summary, String> {
        // As in pi, a request started after the abort fails with its kind,
        // while one aborted mid-stream counts as finished. The caller then
        // sees the abort.
        if self.cancel.is_cancelled() {
            return Err(format!("{label} failed: This operation was aborted"));
        }
        let response = self
            .complete(SUMMARIZATION_SYSTEM_PROMPT, prompt, max_tokens)
            .await;
        if response.stop_reason == StopReason::Aborted {
            return Ok(Summary {
                text: yapi_types::message::blocks_text(&response.content, "\n"),
                usage: response.usage,
            });
        }
        check_summary(&response, label)?;
        Ok(Summary {
            text: yapi_types::message::blocks_text(&response.content, "\n"),
            usage: response.usage,
        })
    }

    fn max_tokens(&self, fraction: f64, reserve_tokens: u64) -> u64 {
        let budget = (fraction * reserve_tokens as f64).floor() as u64;
        if self.model.max_tokens > 0 {
            budget.min(self.model.max_tokens)
        } else {
            budget
        }
    }

    /// Summarizes `messages`, updating `previous` when given.
    pub async fn generate_summary(
        &self,
        messages: &[Message],
        reserve_tokens: u64,
        custom_instructions: Option<&str>,
        previous: Option<&str>,
    ) -> Result<Summary, String> {
        let mut base = if previous.is_some() {
            format!(
                "The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.\n\n{UPDATE_SUMMARIZATION_INSTRUCTIONS}"
            )
        } else {
            SUMMARIZATION_PROMPT.to_owned()
        };
        if let Some(instructions) = custom_instructions.filter(|text| !text.is_empty()) {
            base = format!("{base}\n\nAdditional focus: {instructions}");
        }
        let conversation = serialize_conversation(&convert_to_llm(messages.to_vec()));
        let mut prompt = format!("<conversation>\n{conversation}\n</conversation>\n\n");
        if let Some(previous) = previous {
            prompt += &format!("<previous-summary>\n{previous}\n</previous-summary>\n\n");
        }
        prompt += &base;
        self.summarize(
            &prompt,
            self.max_tokens(0.8, reserve_tokens),
            "Summarization",
        )
        .await
    }

    async fn turn_prefix_summary(
        &self,
        messages: &[Message],
        reserve_tokens: u64,
    ) -> Result<Summary, String> {
        let conversation = serialize_conversation(&convert_to_llm(messages.to_vec()));
        let prompt = format!(
            "# Conversation\n{conversation}\n\n# Instructions\n{TURN_PREFIX_SUMMARIZATION_PROMPT}"
        );
        self.summarize(
            &prompt,
            self.max_tokens(0.5, reserve_tokens),
            "Turn prefix summarization",
        )
        .await
    }

    /// Generates the compaction summary for a preparation.
    pub async fn compact(
        &self,
        preparation: &Preparation,
        custom_instructions: Option<&str>,
    ) -> Result<yapi_types::event::CompactionResult, String> {
        let reserve = preparation.settings.reserve_tokens;
        let previous = preparation.previous_summary.as_deref();
        let (mut summary, usage) =
            if preparation.is_split_turn && !preparation.turn_prefix_messages.is_empty() {
                let history = if preparation.messages_to_summarize.is_empty() {
                    None
                } else {
                    Some(
                        self.generate_summary(
                            &preparation.messages_to_summarize,
                            reserve,
                            custom_instructions,
                            previous,
                        )
                        .await?,
                    )
                };
                let prefix = self
                    .turn_prefix_summary(&preparation.turn_prefix_messages, reserve)
                    .await?;
                let history_text = match &history {
                    Some(history) => history.text.clone(),
                    None => previous.unwrap_or("No prior history.").to_owned(),
                };
                let usage = match history {
                    Some(history) => history.usage.combine(&prefix.usage),
                    None => prefix.usage,
                };
                (
                    format!(
                        "{history_text}\n\n---\n\n**Turn Context (split turn):**\n\n{}",
                        prefix.text
                    ),
                    usage,
                )
            } else {
                let result = self
                    .generate_summary(
                        &preparation.messages_to_summarize,
                        reserve,
                        custom_instructions,
                        previous,
                    )
                    .await?;
                (result.text, result.usage)
            };
        let (read, modified) = preparation.file_ops.lists();
        summary += &format_file_operations(&read, &modified);
        Ok(yapi_types::event::CompactionResult {
            summary,
            first_kept_entry_id: preparation.first_kept_entry_id.clone(),
            tokens_before: preparation.tokens_before,
            estimated_tokens_after: None,
            usage: Some(usage),
            details: Some(json!({"readFiles": read, "modifiedFiles": modified})),
        })
    }
}

const BRANCH_SUMMARY_PREAMBLE: &str = "The user explored a different conversation branch before returning here.\nSummary of that exploration:\n\n";

const BRANCH_SUMMARY_PROMPT: &str = "Create a structured summary of this conversation branch for context when returning later.\n\nUse this EXACT format:\n\n## Goal\n[What was the user trying to accomplish in this branch?]\n\n## Constraints & Preferences\n- [Any constraints, preferences, or requirements mentioned]\n- [Or \"(none)\" if none were mentioned]\n\n## Progress\n### Done\n- [x] [Completed tasks/changes]\n\n### In Progress\n- [ ] [Work that was started but not finished]\n\n### Blocked\n- [Issues preventing progress, if any]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale]\n\n## Next Steps\n1. [What should happen next to continue this work]\n\nKeep each section concise. Preserve exact file paths, function names, and error messages.";

/// The message an entry contributes to a branch summary. Tool results are
/// skipped because the calls already contain their context.
fn branch_entry_message(entry: &FileEntry) -> Option<Message> {
    match entry {
        FileEntry::Message(entry) if matches!(entry.message, Message::ToolResult(_)) => None,
        // A compaction's summary comes after the system message it may include.
        FileEntry::Message(_)
        | FileEntry::CustomMessage(_)
        | FileEntry::BranchSummary(_)
        | FileEntry::Compaction(_) => entry_messages(entry).pop(),
        _ => None,
    }
}

/// Messages of the abandoned branch, newest first until `token_budget` (0 for
/// none), and the files touched across all of it.
pub fn prepare_branch_entries(
    entries: &[FileEntry],
    token_budget: u64,
) -> (Vec<Message>, FileOperations) {
    let mut file_ops = FileOperations::default();
    for entry in entries {
        if let FileEntry::BranchSummary(summary) = entry
            && summary.from_hook != Some(true)
            && let Some(details) = &summary.details
        {
            file_ops.add_details(details);
        }
    }
    let mut messages: Vec<Message> = Vec::new();
    let mut total = 0;
    for entry in entries.iter().rev() {
        let Some(message) = branch_entry_message(entry) else {
            continue;
        };
        file_ops.extract(&message);
        let tokens = estimate_tokens(&message);
        if token_budget > 0 && total + tokens > token_budget {
            // Summaries are worth keeping when there is still some room.
            if matches!(
                entry,
                FileEntry::Compaction(_) | FileEntry::BranchSummary(_)
            ) && (total as f64) < token_budget as f64 * 0.9
            {
                messages.insert(0, message);
            }
            break;
        }
        messages.insert(0, message);
        total += tokens;
    }
    (messages, file_ops)
}

/// What summarizing a branch produced.
#[derive(Clone, Debug, PartialEq)]
pub enum BranchSummary {
    /// A summary to record, with its usage and file lists.
    Done {
        /// The text.
        summary: String,
        /// Usage of the request, if one was made.
        usage: Option<Usage>,
        /// Files only read.
        read_files: Vec<String>,
        /// Files changed.
        modified_files: Vec<String>,
    },
    /// The request was cancelled.
    Aborted,
}

impl Summarizer<'_> {
    /// Summarizes abandoned branch entries (chronological). `replace` makes
    /// `custom_instructions` the whole instruction instead of an addition.
    pub async fn branch_summary(
        &self,
        entries: &[FileEntry],
        custom_instructions: Option<&str>,
        replace: bool,
        reserve_tokens: u64,
    ) -> Result<BranchSummary, String> {
        let window = if self.model.context_window > 0 {
            self.model.context_window
        } else {
            128_000
        };
        let (messages, file_ops) =
            prepare_branch_entries(entries, window.saturating_sub(reserve_tokens));
        if messages.is_empty() {
            return Ok(BranchSummary::Done {
                summary: "No content to summarize".into(),
                usage: None,
                read_files: Vec::new(),
                modified_files: Vec::new(),
            });
        }
        let conversation = serialize_conversation(&convert_to_llm(messages));
        let instructions = match custom_instructions.filter(|text| !text.is_empty()) {
            Some(custom) if replace => custom.to_owned(),
            Some(custom) => format!("{BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: {custom}"),
            None => BRANCH_SUMMARY_PROMPT.to_owned(),
        };
        let prompt = format!("<conversation>\n{conversation}\n</conversation>\n\n{instructions}");
        let max_tokens = if self.model.max_tokens > 0 {
            self.model.max_tokens.min(4096)
        } else {
            4096
        };
        let response = self
            .complete(SUMMARIZATION_SYSTEM_PROMPT, &prompt, max_tokens)
            .await;
        if response.stop_reason == StopReason::Aborted {
            return Ok(BranchSummary::Aborted);
        }
        check_summary(&response, "Branch summarization")?;
        let (read_files, modified_files) = file_ops.lists();
        let summary = format!(
            "{BRANCH_SUMMARY_PREAMBLE}{}{}",
            yapi_types::message::blocks_text(&response.content, "\n"),
            format_file_operations(&read_files, &modified_files)
        );
        Ok(BranchSummary::Done {
            summary,
            usage: Some(response.usage),
            read_files,
            modified_files,
        })
    }
}

/// Entries from `old_leaf` back to the deepest common ancestor with `target`,
/// oldest first, and that ancestor.
pub fn collect_branch_entries(
    session: &crate::session::SessionManager,
    old_leaf: Option<&str>,
    target: &str,
) -> (Vec<FileEntry>, Option<String>) {
    let Some(old_leaf) = old_leaf else {
        return (Vec::new(), None);
    };
    let ids = |path: Vec<&FileEntry>| -> Vec<String> {
        path.iter()
            .filter_map(|entry| entry.meta().map(|meta| meta.id.clone()))
            .collect()
    };
    let old_path = ids(session.branch_path(Some(old_leaf)));
    let target_path = ids(session.branch_path(Some(target)));
    let common = target_path
        .iter()
        .rev()
        .find(|id| old_path.contains(id))
        .cloned();
    let mut entries = Vec::new();
    let mut current = Some(old_leaf.to_owned());
    while let Some(id) = current.filter(|id| Some(id) != common.as_ref()) {
        let Some(entry) = session.entry(&id) else {
            break;
        };
        current = entry.meta().and_then(|meta| meta.parent_id.clone());
        entries.push(entry.clone());
    }
    entries.reverse();
    (entries, common)
}

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

    #[test]
    fn serializes_conversations() {
        let messages: Vec<Message> = serde_json::from_value(json!([
            {"role":"user","content":"hi","timestamp":1},
            {"role":"assistant","content":[
                {"type":"thinking","thinking":"plan"},
                {"type":"text","text":"Reading."},
                {"type":"toolCall","id":"c","name":"read","arguments":{"path":"a.txt","limit":null}}
            ],"api":"x","provider":"p","model":"m","usage":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"toolUse","timestamp":2},
            {"role":"toolResult","toolCallId":"c","toolName":"read","content":[{"type":"text","text":"x"}],"isError":false,"timestamp":3}
        ]))
        .unwrap();
        assert_eq!(
            serialize_conversation(&messages),
            "[User]: hi\n\n[Assistant thinking]: plan\n\n[Assistant]: Reading.\n\n[Assistant tool calls]: read(path=\"a.txt\", limit=null)\n\n[Tool result]: x"
        );
        let mut ops = FileOperations::default();
        for message in &messages {
            ops.extract(message);
        }
        assert_eq!(ops.lists(), (vec!["a.txt".to_owned()], vec![]));
    }

    #[test]
    fn thresholds() {
        let settings = CompactionSettings {
            enabled: true,
            reserve_tokens: 100,
            keep_recent_tokens: 10,
        };
        assert!(should_compact(950, 1000, &settings));
        assert!(!should_compact(900, 1000, &settings));
        assert!(should_compact(1, 50, &settings));
    }
}