From 61cbf3574eca870df6fa7f49648ec7e001901b5a Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 27 May 2026 15:49:08 +0200 Subject: [PATCH] Drop startup context when truncating forked rollouts (#24751) ## Summary - Change last-`n` fork truncation to start at the first fork-turn boundary instead of returning the full rollout when the fork history is shorter than the requested window. - Add coverage for the startup-prefix case in both rollout truncation tests and agent control spawn behavior. - Ensure bounded forked children still rebuild context after the cached prefix is truncated. ## Testing - Added unit coverage for truncation behavior when the parent history is under the requested fork-turn limit. - Added an agent control test covering bounded fork spawn behavior with startup context present. - Not run (not requested). --- codex-rs/core/src/agent/control_tests.rs | 102 ++++++++++++++++++ .../core/src/thread_rollout_truncation.rs | 16 +-- .../src/thread_rollout_truncation_tests.rs | 28 +++++ 3 files changed, 140 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 8209272f10..b5ba938ef7 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1113,6 +1113,108 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { .expect("parent shutdown should submit"); } +#[tokio::test] +async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_limit() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + let startup_turn_context = parent_thread.codex.session.new_default_turn().await; + parent_thread + .codex + .session + .record_conversation_items( + startup_turn_context.as_ref(), + &[ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "parent startup developer context".to_string(), + }], + phase: None, + }], + ) + .await; + parent_thread + .inject_user_message_without_turn("current parent task".to_string()) + .await; + let spawn_turn_context = parent_thread.codex.session.new_default_turn().await; + let parent_spawn_call_id = "spawn-call-last-n-under-limit".to_string(); + parent_thread + .codex + .session + .record_conversation_items( + spawn_turn_context.as_ref(), + &[spawn_agent_call(&parent_spawn_call_id)], + ) + .await; + parent_thread + .codex + .session + .ensure_rollout_materialized() + .await; + parent_thread + .codex + .session + .flush_rollout() + .await + .expect("parent rollout should flush"); + + let child_thread_id = harness + .control + .spawn_agent_with_metadata( + harness.config.clone(), + text_input("child task"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + fork_parent_spawn_call_id: Some(parent_spawn_call_id), + fork_mode: Some(SpawnAgentForkMode::LastNTurns(2)), + ..Default::default() + }, + ) + .await + .expect("bounded forked spawn should drop startup prefix") + .thread_id; + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + let history = child_thread.codex.session.clone_history().await; + assert!( + history_contains_text(history.raw_items(), "current parent task"), + "bounded fork should retain the requested recent parent turn" + ); + assert!( + !history_contains_text(history.raw_items(), "parent startup developer context"), + "bounded fork should drop parent startup context even when fewer turns exist than requested" + ); + assert!( + child_thread + .codex + .session + .reference_context_item() + .await + .is_none(), + "bounded forked child should still rebuild context after truncating the cached prefix" + ); + + let _ = harness + .control + .shutdown_live_agent(child_thread_id) + .await + .expect("child shutdown should submit"); + let _ = parent_thread + .submit(Op::Shutdown {}) + .await + .expect("parent shutdown should submit"); +} + #[tokio::test] async fn spawn_agent_fork_last_n_turns_strips_parent_usage_hints() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/thread_rollout_truncation.rs b/codex-rs/core/src/thread_rollout_truncation.rs index e20ee53d47..69e615c4fb 100644 --- a/codex-rs/core/src/thread_rollout_truncation.rs +++ b/codex-rs/core/src/thread_rollout_truncation.rs @@ -130,7 +130,8 @@ pub(crate) fn truncate_rollout_before_nth_user_message_from_start( /// 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. +/// 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, @@ -140,11 +141,14 @@ pub(crate) fn truncate_rollout_to_last_n_fork_turns( } 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]; + 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() } diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index df370a0546..e0a8837165 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -29,6 +29,17 @@ fn assistant_msg(text: &str) -> ResponseItem { } } +fn developer_msg(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + phase: None, + } +} + fn inter_agent_msg(text: &str, trigger_turn: bool) -> ResponseItem { let communication = InterAgentCommunication::new( AgentPath::root(), @@ -197,6 +208,23 @@ fn truncates_rollout_to_last_n_fork_turns_counts_trigger_turn_messages() { ); } +#[test] +fn truncates_rollout_to_last_n_fork_turns_drops_startup_prefix_even_when_under_limit() { + let rollout = vec![ + RolloutItem::ResponseItem(developer_msg("startup developer context")), + RolloutItem::ResponseItem(user_msg("current task")), + RolloutItem::ResponseItem(assistant_msg("answer")), + ]; + + let truncated = truncate_rollout_to_last_n_fork_turns(&rollout, /*n_from_end*/ 2); + let expected = rollout[1..].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![