core/src/thread_rollout_truncation.rs
182 lines
//! Helpers for truncating rollouts based on "user turn" boundaries.//!//! In core, "user turns" are detected by scanning `ResponseItem::Message` items and//! interpreting them via `event_mapping::parse_turn_item(...)`.use crate::context_manager::is_user_turn_boundary;use crate::event_mapping;use codex_protocol::items::TurnItem;use codex_protocol::models::ResponseItem;use codex_protocol::protocol::EventMsg;use codex_protocol::protocol::InitialHistory;use codex_protocol::protocol::InterAgentCommunication;use codex_protocol::protocol::RolloutItem;pub(crate) fn initial_history_has_prior_user_turns(conversation_history: &InitialHistory) -> bool { conversation_history.scan_rollout_items(rollout_item_is_user_turn_boundary)}fn rollout_item_is_user_turn_boundary(item: &RolloutItem) -> bool { match item { RolloutItem::ResponseItem(item) => is_user_turn_boundary(item), RolloutItem::InterAgentCommunication(_) => true, _ => false, }}/// Return the indices of user message boundaries in a rollout.////// A user message boundary is a `RolloutItem::ResponseItem(ResponseItem::Message { .. })`/// whose parsed turn item is `TurnItem::UserMessage`.////// Rollouts can contain `ThreadRolledBack` markers. Those markers indicate that the/// last N user turns were removed from the effective thread history; we apply them here so/// indexing uses the post-rollback history rather than the raw stream.pub(crate) fn user_message_positions_in_rollout(items: &[RolloutItem]) -> Vec<usize> { let mut user_positions = Vec::new(); for (idx, item) in items.iter().enumerate() { match item { RolloutItem::ResponseItem(item @ ResponseItem::Message { .. }) if matches!( event_mapping::parse_turn_item(item), Some(TurnItem::UserMessage(_)) ) => { user_positions.push(idx); } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { let num_turns = usize::try_from(rollback.num_turns).unwrap_or(usize::MAX); let new_len = user_positions.len().saturating_sub(num_turns); user_positions.truncate(new_len); } _ => {} } } user_positions}/// Return the indices of fork-turn boundaries in a rollout.////// A fork-turn boundary is either:/// - a real user message boundary, or/// - an inter-agent communication whose `trigger_turn` is `true`, or/// - a legacy assistant inter-agent envelope with the same flag.////// Like `user_message_positions_in_rollout`, this applies `ThreadRolledBack` markers so indexing/// reflects the effective post-rollback history. Rollback counts instruction turns, so a rollback/// removes the stale suffix starting at the earliest rolled-back instruction-turn boundary instead/// of simply truncating the mixed fork-boundary list.pub(crate) fn fork_turn_positions_in_rollout(items: &[RolloutItem]) -> Vec<usize> { let mut rollback_turn_positions = Vec::new(); let mut fork_turn_positions = Vec::new(); for (idx, item) in items.iter().enumerate() { match item { RolloutItem::ResponseItem(item) => { if is_user_turn_boundary(item) { rollback_turn_positions.push(idx); } if is_real_user_message_boundary(item) || is_trigger_turn_boundary(item) { fork_turn_positions.push(idx); } } RolloutItem::InterAgentCommunication(communication) => { rollback_turn_positions.push(idx); if communication.trigger_turn { fork_turn_positions.push(idx); } } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { let num_turns = usize::try_from(rollback.num_turns).unwrap_or(usize::MAX); if num_turns == 0 { continue; } let Some(rollback_start_idx) = rollback_turn_positions .len() .checked_sub(num_turns) .map(|rollback_start| rollback_turn_positions[rollback_start]) .or_else(|| rollback_turn_positions.first().copied()) else { continue; }; let new_rollback_len = rollback_turn_positions.len().saturating_sub(num_turns); rollback_turn_positions.truncate(new_rollback_len); fork_turn_positions.retain(|position| *position < rollback_start_idx); } _ => {} } } fork_turn_positions}/// Return a prefix of `items` obtained by cutting strictly before the nth user message.////// The boundary index is 0-based from the start of `items` (so `n_from_start = 0` returns/// a prefix that excludes the first user message and everything after it).////// If `n_from_start` is `usize::MAX`, this returns the full rollout (no truncation)./// If fewer than or equal to `n_from_start` user messages exist, this returns the full/// rollout unchanged.pub(crate) fn truncate_rollout_before_nth_user_message_from_start( items: &[RolloutItem], n_from_start: usize,) -> Vec<RolloutItem> { if n_from_start == usize::MAX { return items.to_vec(); } let user_positions = user_message_positions_in_rollout(items); // If fewer than or equal to n user messages exist, keep the full rollout. if user_positions.len() <= n_from_start { return items.to_vec(); } // Cut strictly before the nth user message (do not keep the nth itself). let cut_idx = user_positions[n_from_start]; items[..cut_idx].to_vec()}/// Return a suffix of `items` that keeps the last `n_from_end` fork turns.////// If fewer than or equal to `n_from_end` fork turns exist, this keeps from the first fork-turn/// boundary and still drops pre-turn startup context.pub(crate) fn truncate_rollout_to_last_n_fork_turns( items: &[RolloutItem], n_from_end: usize,) -> Vec<RolloutItem> { if n_from_end == 0 { return Vec::new(); } let fork_turn_positions = fork_turn_positions_in_rollout(items); let Some(keep_idx) = fork_turn_positions .len() .checked_sub(n_from_end) .map(|position| fork_turn_positions[position]) .or_else(|| fork_turn_positions.first().copied()) else { return Vec::new(); }; items[keep_idx..].to_vec()}fn is_real_user_message_boundary(item: &ResponseItem) -> bool { matches!( event_mapping::parse_turn_item(item), Some(TurnItem::UserMessage(_)) )}fn is_trigger_turn_boundary(item: &ResponseItem) -> bool { let ResponseItem::Message { role, content, .. } = item else { return false; }; role == "assistant" && InterAgentCommunication::from_message_content(content) .is_some_and(|communication| communication.trigger_turn)}#[cfg(test)]#[path = "thread_rollout_truncation_tests.rs"]mod tests;