diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 3a593e4551..fe998edca6 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -3243,6 +3243,7 @@ impl Session { compaction_initial_context: &[ResponseItem], turn_start_context_items: &[ResponseItem], history_before_turn: &[ResponseItem], + history_at_checkpoint: &[ResponseItem], ) { let current_history = self.clone_history().await; let replacement_history = build_server_side_compaction_replacement_history( @@ -3250,6 +3251,7 @@ impl Session { compaction_initial_context, turn_start_context_items, history_before_turn, + history_at_checkpoint, current_history.raw_items(), ); @@ -5961,21 +5963,32 @@ fn build_server_side_compaction_replacement_history( compaction_initial_context: &[ResponseItem], turn_start_context_items: &[ResponseItem], history_before_turn: &[ResponseItem], + history_at_checkpoint: &[ResponseItem], current_history: &[ResponseItem], ) -> Vec { - let current_turn_items = current_history + let checkpoint_turn_items = history_at_checkpoint .strip_prefix(history_before_turn) - .unwrap_or(current_history); - let current_turn_items = current_turn_items + .unwrap_or(history_at_checkpoint); + let checkpoint_turn_items = checkpoint_turn_items .strip_prefix(turn_start_context_items) - .unwrap_or(current_turn_items); - let mut replacement_history: Vec = current_turn_items + .unwrap_or(checkpoint_turn_items); + let post_checkpoint_turn_items = current_history + .strip_prefix(history_at_checkpoint) + .unwrap_or_default(); + let mut replacement_history: Vec = checkpoint_turn_items .iter() .filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. })) .filter(|item| !matches!(item, ResponseItem::Compaction { .. })) .cloned() .collect(); replacement_history.push(compaction_item); + replacement_history.extend( + post_checkpoint_turn_items + .iter() + .filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. })) + .filter(|item| !matches!(item, ResponseItem::Compaction { .. })) + .cloned(), + ); let mut replacement_history = insert_initial_context_before_last_real_user_or_summary( replacement_history, compaction_initial_context.to_vec(), @@ -7355,8 +7368,11 @@ async fn try_run_sampling_request( &mut assistant_message_stream_parsers, ) .await; - if let Some(PendingServerSideCompactionCheckpoint { item, turn_item }) = - pending_server_side_compaction_checkpoint.take() + if let Some(PendingServerSideCompactionCheckpoint { + history_at_checkpoint, + item, + turn_item, + }) = pending_server_side_compaction_checkpoint.take() { let turn_item = TurnItem::ContextCompaction(turn_item); sess.emit_turn_item_started(&turn_context, &turn_item).await; @@ -7366,6 +7382,7 @@ async fn try_run_sampling_request( compaction_initial_context, turn_start_context_items, history_before_turn.as_slice(), + history_at_checkpoint.as_slice(), ) .await; sess.emit_turn_item_completed(&turn_context, turn_item) diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index e655bccc83..1536ef8bac 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -317,6 +317,7 @@ fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() &turn_start_context_items, &history_before_turn, ¤t_history, + ¤t_history, ); assert_eq!( @@ -372,6 +373,7 @@ fn build_server_side_compaction_replacement_history_preserves_turn_scoped_inject &turn_start_context_items, &history_before_turn, ¤t_history, + ¤t_history, ); assert_eq!( @@ -426,6 +428,7 @@ fn build_server_side_compaction_replacement_history_replaces_prior_same_turn_sum &turn_start_context_items, &history_before_turn, ¤t_history, + ¤t_history, ); assert_eq!( @@ -476,6 +479,7 @@ fn build_server_side_compaction_replacement_history_replaces_prior_summary_with_ &turn_start_context_items, &history_before_turn, ¤t_history, + ¤t_history, ); assert_eq!( @@ -491,6 +495,66 @@ fn build_server_side_compaction_replacement_history_replaces_prior_summary_with_ ); } +#[test] +fn build_server_side_compaction_replacement_history_keeps_checkpoint_before_post_compaction_items() +{ + let prior_snapshot = ghost_snapshot("ghost-before"); + let same_turn_snapshot = ghost_snapshot("ghost-during"); + let history_before_turn = vec![user_message("earlier"), prior_snapshot.clone()]; + let turn_start_context_items = vec![ + developer_message("fresh permissions"), + environment_context_message("/fresh"), + ]; + let current_turn_user = user_message("current turn"); + let post_checkpoint_tool_call = ResponseItem::FunctionCall { + id: None, + call_id: "call-1".to_string(), + name: "test_tool".to_string(), + arguments: "{}".to_string(), + }; + let history_at_checkpoint = vec![ + user_message("earlier"), + prior_snapshot.clone(), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), + current_turn_user.clone(), + ]; + let current_history = vec![ + history_at_checkpoint[0].clone(), + history_at_checkpoint[1].clone(), + history_at_checkpoint[2].clone(), + history_at_checkpoint[3].clone(), + history_at_checkpoint[4].clone(), + post_checkpoint_tool_call.clone(), + same_turn_snapshot.clone(), + ]; + let compaction_item = ResponseItem::Compaction { + encrypted_content: "INLINE_SUMMARY".to_string(), + }; + + let replacement_history = build_server_side_compaction_replacement_history( + compaction_item.clone(), + &turn_start_context_items, + &turn_start_context_items, + &history_before_turn, + &history_at_checkpoint, + ¤t_history, + ); + + assert_eq!( + replacement_history, + vec![ + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), + current_turn_user, + compaction_item, + post_checkpoint_tool_call, + prior_snapshot, + same_turn_snapshot, + ] + ); +} + fn make_mcp_tool( server_name: &str, tool_name: &str, diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 1dbc2fd111..78dc7b997a 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -150,6 +150,7 @@ pub(crate) struct OutputItemResult { } pub(crate) struct PendingServerSideCompactionCheckpoint { + pub history_at_checkpoint: Vec, pub item: ResponseItem, pub turn_item: ContextCompactionItem, } @@ -180,6 +181,7 @@ pub(crate) async fn handle_output_item_done( "buffering streamed server-side compaction item until response.completed" ); output.pending_server_side_compaction = Some(PendingServerSideCompactionCheckpoint { + history_at_checkpoint: ctx.sess.clone_history().await.raw_items().to_vec(), item, turn_item: compaction_item, }); diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 4e10926b88..ae68c37bd9 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -512,6 +512,7 @@ async fn auto_server_side_compaction_uses_inline_context_management() -> Result< ); let post_inline_request = &requests[2]; + let post_inline_body = post_inline_request.body_json().to_string(); assert!( post_inline_request.body_contains_text(&inline_summary), "expected subsequent request to reuse inline compaction item" @@ -532,6 +533,15 @@ async fn auto_server_side_compaction_uses_inline_context_management() -> Result< post_inline_request.body_contains_text(third_turn_text), "expected next turn to append normally after inline compaction" ); + assert!( + post_inline_body + .find("INLINE_SERVER_SUMMARY") + .expect("inline compaction marker in subsequent request") + < post_inline_body + .find("AFTER_INLINE_REPLY") + .expect("post-compaction assistant reply in subsequent request"), + "expected post-compaction transcript items to remain after the inline compaction item" + ); Ok(()) } @@ -622,6 +632,15 @@ async fn auto_server_side_compaction_keeps_current_turn_inputs_for_follow_ups() .is_some(), "expected same-turn follow-up to include the tool output" ); + assert!( + follow_up_body + .find("INLINE_SERVER_SUMMARY") + .expect("inline compaction marker in follow-up request") + < follow_up_body + .find("call-inline-mid-turn") + .expect("post-compaction tool call in follow-up request"), + "expected the inline compaction item to remain ahead of post-compaction tool calls" + ); Ok(()) }