diff --git a/codex-rs/core/src/thread_rollout_truncation.rs b/codex-rs/core/src/thread_rollout_truncation.rs index 5fa1881ffa..d27f954993 100644 --- a/codex-rs/core/src/thread_rollout_truncation.rs +++ b/codex-rs/core/src/thread_rollout_truncation.rs @@ -7,8 +7,21 @@ use crate::event_mapping; use codex_protocol::items::TurnItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::RolloutItem; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ForkTurnBoundaryKind { + User, + TriggerTurnEnvelope, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct ForkTurnBoundary { + idx: usize, + kind: ForkTurnBoundaryKind, +} + /// Return the indices of user message boundaries in a rollout. /// /// A user message boundary is a `RolloutItem::ResponseItem(ResponseItem::Message { .. })` @@ -18,26 +31,29 @@ use codex_protocol::protocol::RolloutItem; /// 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 { - 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 + fork_turn_boundaries_in_rollout(items) + .into_iter() + .filter_map(|boundary| match boundary.kind { + ForkTurnBoundaryKind::User => Some(boundary.idx), + ForkTurnBoundaryKind::TriggerTurnEnvelope => None, + }) + .collect() +} + +/// Return the indices of fork-turn boundaries in a rollout. +/// +/// A fork-turn boundary is either: +/// - a real user message boundary, or +/// - an assistant inter-agent envelope whose parsed `trigger_turn` is `true`. +/// +/// Rollbacks are applied to the effective instruction-turn stack rather than to user-only +/// boundaries, so a rollback can correctly remove trigger-turn inter-agent envelopes even when +/// there are no real user messages in the rolled-back suffix. +pub(crate) fn fork_turn_positions_in_rollout(items: &[RolloutItem]) -> Vec { + fork_turn_boundaries_in_rollout(items) + .into_iter() + .map(|boundary| boundary.idx) + .collect() } /// Return a prefix of `items` obtained by cutting strictly before the nth user message. @@ -68,6 +84,70 @@ pub(crate) fn truncate_rollout_before_nth_user_message_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 returns the full rollout. +pub(crate) fn truncate_rollout_to_last_n_fork_turns( + items: &[RolloutItem], + n_from_end: usize, +) -> Vec { + if n_from_end == 0 { + return Vec::new(); + } + + let fork_turn_positions = fork_turn_positions_in_rollout(items); + if fork_turn_positions.len() <= n_from_end { + return items.to_vec(); + } + + let keep_idx = fork_turn_positions[fork_turn_positions.len() - n_from_end]; + items[keep_idx..].to_vec() +} + +fn fork_turn_boundaries_in_rollout(items: &[RolloutItem]) -> Vec { + let mut boundaries = Vec::new(); + for (idx, item) in items.iter().enumerate() { + match item { + RolloutItem::ResponseItem(item) if is_real_user_message_boundary(item) => { + boundaries.push(ForkTurnBoundary { + idx, + kind: ForkTurnBoundaryKind::User, + }); + } + RolloutItem::ResponseItem(item) if is_trigger_turn_boundary(item) => { + boundaries.push(ForkTurnBoundary { + idx, + kind: ForkTurnBoundaryKind::TriggerTurnEnvelope, + }); + } + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { + let num_turns = usize::try_from(rollback.num_turns).unwrap_or(usize::MAX); + let new_len = boundaries.len().saturating_sub(num_turns); + boundaries.truncate(new_len); + } + _ => {} + } + } + boundaries +} + +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; diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index d251b722d9..a2146b60af 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -1,7 +1,9 @@ use super::*; use crate::codex::make_session_and_context; +use codex_protocol::AgentPath; use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemReasoningSummary; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::ThreadRolledBackEvent; use pretty_assertions::assert_eq; @@ -29,6 +31,17 @@ fn assistant_msg(text: &str) -> ResponseItem { } } +fn inter_agent_msg(text: &str, trigger_turn: bool) -> ResponseItem { + let communication = InterAgentCommunication::new( + AgentPath::root(), + AgentPath::try_from("/root/worker").expect("agent path"), + Vec::new(), + text.to_string(), + trigger_turn, + ); + communication.to_response_input_item().into() +} + #[test] fn truncates_rollout_from_start_before_nth_user_only() { let items = [ @@ -157,3 +170,97 @@ async fn ignores_session_prefix_messages_when_truncating_rollout_from_start() { serde_json::to_value(&expected).unwrap() ); } + +#[test] +fn truncates_rollout_to_last_n_fork_turns_counts_trigger_turn_messages() { + let rollout = vec![ + RolloutItem::ResponseItem(user_msg("u1")), + RolloutItem::ResponseItem(assistant_msg("a1")), + RolloutItem::ResponseItem(inter_agent_msg( + "queued message", + /*trigger_turn*/ false, + )), + RolloutItem::ResponseItem(assistant_msg("a2")), + RolloutItem::ResponseItem(inter_agent_msg( + "triggered task", + /*trigger_turn*/ true, + )), + RolloutItem::ResponseItem(assistant_msg("a3")), + RolloutItem::ResponseItem(user_msg("u2")), + RolloutItem::ResponseItem(assistant_msg("a4")), + ]; + + let truncated = truncate_rollout_to_last_n_fork_turns(&rollout, 2); + let expected = rollout[4..].to_vec(); + + assert_eq!( + serde_json::to_value(&truncated).unwrap(), + serde_json::to_value(&expected).unwrap() + ); +} + +#[test] +fn truncates_rollout_to_last_n_fork_turns_applies_thread_rollback_markers() { + let rollout = vec![ + RolloutItem::ResponseItem(user_msg("u1")), + RolloutItem::ResponseItem(assistant_msg("a1")), + RolloutItem::ResponseItem(inter_agent_msg( + "triggered task", + /*trigger_turn*/ true, + )), + RolloutItem::ResponseItem(assistant_msg("a2")), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + })), + RolloutItem::ResponseItem(user_msg("u2")), + RolloutItem::ResponseItem(assistant_msg("a3")), + ]; + + let truncated = truncate_rollout_to_last_n_fork_turns(&rollout, 2); + + assert_eq!( + serde_json::to_value(&truncated).unwrap(), + serde_json::to_value(&rollout).unwrap() + ); +} + +#[test] +fn truncates_rollout_to_last_n_fork_turns_discards_rolled_back_trigger_turn_only_suffix() { + let rollout = vec![ + RolloutItem::ResponseItem(inter_agent_msg("task A", /*trigger_turn*/ true)), + RolloutItem::ResponseItem(assistant_msg("working on task A")), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + })), + RolloutItem::ResponseItem(inter_agent_msg("task B", /*trigger_turn*/ true)), + RolloutItem::ResponseItem(assistant_msg("working on task B")), + ]; + + let truncated = truncate_rollout_to_last_n_fork_turns(&rollout, 2); + let expected = rollout[3..].to_vec(); + + assert_eq!( + serde_json::to_value(&truncated).unwrap(), + serde_json::to_value(&expected).unwrap() + ); +} + +#[test] +fn truncates_rollout_to_last_n_fork_turns_keeps_full_rollout_when_n_is_large() { + let rollout = vec![ + RolloutItem::ResponseItem(user_msg("u1")), + RolloutItem::ResponseItem(assistant_msg("a1")), + RolloutItem::ResponseItem(inter_agent_msg( + "triggered task", + /*trigger_turn*/ true, + )), + RolloutItem::ResponseItem(assistant_msg("a2")), + ]; + + let truncated = truncate_rollout_to_last_n_fork_turns(&rollout, 10); + + assert_eq!( + serde_json::to_value(&truncated).unwrap(), + serde_json::to_value(&rollout).unwrap() + ); +}