From 8da9b522f5a9eb097af9f74b7227eb02085d86b8 Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Mon, 9 Mar 2026 01:12:22 +0000 Subject: [PATCH] [codex-core] Preserve inline compaction turn prompt state [ci changed_files] Co-authored-by: Codex --- codex-rs/core/src/codex.rs | 36 +++++--- codex-rs/core/src/codex_tests.rs | 96 +++++++++++++++++---- codex-rs/core/tests/suite/compact_remote.rs | 95 ++++++++++++++++++++ 3 files changed, 196 insertions(+), 31 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index cf27ff4c56..3a593e4551 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -3240,13 +3240,15 @@ impl Session { &self, turn_context: &TurnContext, item: ResponseItem, + compaction_initial_context: &[ResponseItem], + turn_start_context_items: &[ResponseItem], history_before_turn: &[ResponseItem], ) { let current_history = self.clone_history().await; - let initial_context = self.build_initial_context(turn_context).await; let replacement_history = build_server_side_compaction_replacement_history( item.clone(), - initial_context, + compaction_initial_context, + turn_start_context_items, history_before_turn, current_history.raw_items(), ); @@ -5406,7 +5408,9 @@ pub(crate) async fn run_turn( let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); let history_before_turn = sess.clone_history().await.raw_items().to_vec(); - sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref()) + let compaction_initial_context = sess.build_initial_context(turn_context.as_ref()).await; + let turn_start_context_items = sess + .record_context_updates_and_set_reference_context_item(turn_context.as_ref()) .await; let loaded_plugins = sess @@ -5691,6 +5695,8 @@ pub(crate) async fn run_turn( &mut client_session, turn_metadata_header.as_deref(), sampling_request_input, + &compaction_initial_context, + &turn_start_context_items, &history_before_turn, inline_compaction_for_request.map(|pending| pending.threshold), &turn_enabled_connectors, @@ -5952,31 +5958,27 @@ struct PendingServerSideCompaction { fn build_server_side_compaction_replacement_history( compaction_item: ResponseItem, - initial_context: Vec, + compaction_initial_context: &[ResponseItem], + turn_start_context_items: &[ResponseItem], history_before_turn: &[ResponseItem], current_history: &[ResponseItem], ) -> Vec { let current_turn_items = current_history .strip_prefix(history_before_turn) - .or_else(|| current_history.strip_prefix(initial_context.as_slice())) .unwrap_or(current_history); + let current_turn_items = current_turn_items + .strip_prefix(turn_start_context_items) + .unwrap_or(current_turn_items); let mut replacement_history: Vec = current_turn_items .iter() .filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. })) .filter(|item| !matches!(item, ResponseItem::Compaction { .. })) - .filter(|item| match item { - ResponseItem::Message { role, .. } if role == "developer" => false, - ResponseItem::Message { role, content, .. } if role == "user" => { - !crate::event_mapping::is_contextual_user_message_content(content) - } - _ => true, - }) .cloned() .collect(); replacement_history.push(compaction_item); let mut replacement_history = insert_initial_context_before_last_real_user_or_summary( replacement_history, - initial_context, + compaction_initial_context.to_vec(), ); replacement_history.extend( current_history @@ -6358,6 +6360,8 @@ async fn run_sampling_request( client_session: &mut ModelClientSession, turn_metadata_header: Option<&str>, input: Vec, + compaction_initial_context: &[ResponseItem], + turn_start_context_items: &[ResponseItem], history_before_turn: &[ResponseItem], inline_compaction_threshold: Option, explicitly_enabled_connectors: &HashSet, @@ -6395,6 +6399,8 @@ async fn run_sampling_request( Arc::clone(&turn_diff_tracker), server_model_warning_emitted_for_turn, &prompt, + compaction_initial_context, + turn_start_context_items, history_before_turn, cancellation_token.child_token(), ) @@ -7119,6 +7125,8 @@ async fn try_run_sampling_request( turn_diff_tracker: SharedTurnDiffTracker, server_model_warning_emitted_for_turn: &mut bool, prompt: &Prompt, + compaction_initial_context: &[ResponseItem], + turn_start_context_items: &[ResponseItem], history_before_turn: &[ResponseItem], cancellation_token: CancellationToken, ) -> CodexResult { @@ -7355,6 +7363,8 @@ async fn try_run_sampling_request( sess.apply_server_side_compaction_checkpoint( turn_context.as_ref(), item, + compaction_initial_context, + turn_start_context_items, history_before_turn.as_slice(), ) .await; diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index ade9511f9c..e655bccc83 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -289,7 +289,7 @@ fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() 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 initial_context = vec![ + let turn_start_context_items = vec![ developer_message("fresh permissions"), environment_context_message("/fresh"), ]; @@ -301,8 +301,8 @@ fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() let current_history = vec![ user_message("earlier"), prior_snapshot.clone(), - developer_message("stale permissions diff"), - environment_context_message("/stale"), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), current_turn_user.clone(), current_turn_tool_output.clone(), same_turn_snapshot.clone(), @@ -313,7 +313,8 @@ fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() let replacement_history = build_server_side_compaction_replacement_history( compaction_item.clone(), - initial_context.clone(), + &turn_start_context_items, + &turn_start_context_items, &history_before_turn, ¤t_history, ); @@ -321,8 +322,8 @@ fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() assert_eq!( replacement_history, vec![ - initial_context[0].clone(), - initial_context[1].clone(), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), current_turn_user, current_turn_tool_output, compaction_item, @@ -332,12 +333,69 @@ fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() ); } +#[test] +fn build_server_side_compaction_replacement_history_preserves_turn_scoped_injections() { + 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("\nuse the new model"), + environment_context_message("/fresh"), + ]; + let current_turn_user = user_message("current turn"); + let skill_injection = skill_message( + "\ndemo\n/tmp/skills/demo/SKILL.md\nbody\n", + ); + let plugin_injection = developer_message("PLUGIN_HINT"); + let current_turn_tool_output = ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text("tool result".to_string()), + }; + let current_history = vec![ + user_message("earlier"), + prior_snapshot.clone(), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), + current_turn_user.clone(), + skill_injection.clone(), + plugin_injection.clone(), + current_turn_tool_output.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, + ¤t_history, + ); + + assert_eq!( + replacement_history, + vec![ + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), + current_turn_user, + skill_injection, + plugin_injection, + current_turn_tool_output, + compaction_item, + prior_snapshot, + same_turn_snapshot, + ] + ); +} + #[test] fn build_server_side_compaction_replacement_history_replaces_prior_same_turn_summary() { 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 initial_context = vec![ + let turn_start_context_items = vec![ developer_message("fresh permissions"), environment_context_message("/fresh"), ]; @@ -353,8 +411,8 @@ fn build_server_side_compaction_replacement_history_replaces_prior_same_turn_sum encrypted_content: "INLINE_SUMMARY_2".to_string(), }; let current_history = vec![ - initial_context[0].clone(), - initial_context[1].clone(), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), current_turn_user.clone(), current_turn_tool_output.clone(), prior_compaction, @@ -364,7 +422,8 @@ fn build_server_side_compaction_replacement_history_replaces_prior_same_turn_sum let replacement_history = build_server_side_compaction_replacement_history( new_compaction.clone(), - initial_context.clone(), + &turn_start_context_items, + &turn_start_context_items, &history_before_turn, ¤t_history, ); @@ -372,8 +431,8 @@ fn build_server_side_compaction_replacement_history_replaces_prior_same_turn_sum assert_eq!( replacement_history, vec![ - initial_context[0].clone(), - initial_context[1].clone(), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), current_turn_user, current_turn_tool_output, new_compaction, @@ -387,7 +446,7 @@ fn build_server_side_compaction_replacement_history_replaces_prior_same_turn_sum fn build_server_side_compaction_replacement_history_replaces_prior_summary_with_empty_history() { let same_turn_snapshot = ghost_snapshot("ghost-during"); let history_before_turn = Vec::new(); - let initial_context = vec![ + let turn_start_context_items = vec![ developer_message("fresh permissions"), environment_context_message("/fresh"), ]; @@ -403,8 +462,8 @@ fn build_server_side_compaction_replacement_history_replaces_prior_summary_with_ encrypted_content: "INLINE_SUMMARY_2".to_string(), }; let current_history = vec![ - initial_context[0].clone(), - initial_context[1].clone(), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), prior_compaction, current_turn_user.clone(), current_turn_tool_output.clone(), @@ -413,7 +472,8 @@ fn build_server_side_compaction_replacement_history_replaces_prior_summary_with_ let replacement_history = build_server_side_compaction_replacement_history( new_compaction.clone(), - initial_context.clone(), + &turn_start_context_items, + &turn_start_context_items, &history_before_turn, ¤t_history, ); @@ -421,8 +481,8 @@ fn build_server_side_compaction_replacement_history_replaces_prior_summary_with_ assert_eq!( replacement_history, vec![ - initial_context[0].clone(), - initial_context[1].clone(), + turn_start_context_items[0].clone(), + turn_start_context_items[1].clone(), current_turn_user, current_turn_tool_output, new_compaction, diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 63a2ef6b64..4e10926b88 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -626,6 +626,101 @@ async fn auto_server_side_compaction_keeps_current_turn_inputs_for_follow_ups() Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn auto_server_side_compaction_follow_up_preserves_model_switch_updates() -> Result<()> { + skip_if_no_network!(Ok(())); + + let compact_threshold = 120; + let previous_model = "gpt-5.2-codex"; + let next_model = "gpt-5.1-codex-max"; + let inline_summary = summary_with_prefix("INLINE_SWITCH_SUMMARY"); + let server = wiremock::MockServer::start().await; + let _models_mock = responses::mount_models_once( + &server, + ModelsResponse { + models: vec![ + model_info_with_context_window(previous_model, 273_000), + model_info_with_context_window(next_model, 273_000), + ], + }, + ) + .await; + let responses_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_assistant_message("m1", "FIRST_REMOTE_REPLY"), + responses::ev_completed_with_tokens("resp-1", 500), + ]), + responses::sse(vec![ + responses::ev_compaction(&inline_summary), + responses::ev_function_call("call-inline-model-switch", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("resp-2", 80), + ]), + responses::sse(vec![ + responses::ev_assistant_message("m3", "AFTER_INLINE_SWITCH_TOOL_REPLY"), + responses::ev_completed("resp-3"), + ]), + ], + ) + .await; + + let test = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model(previous_model) + .with_config(move |config| { + config + .features + .enable(Feature::ServerSideCompaction) + .expect("enable server-side compaction"); + config.model_auto_compact_token_limit = Some(compact_threshold); + }) + .build(&server) + .await?; + let codex = test.codex.clone(); + + submit_text_turn_and_wait(&codex, "BEFORE_SWITCH_USER").await?; + + codex + .submit(Op::OverrideTurnContext { + cwd: None, + approval_policy: None, + sandbox_policy: None, + windows_sandbox_level: None, + model: Some(next_model.to_string()), + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + submit_text_turn_and_wait(&codex, "AFTER_SWITCH_USER").await?; + + let requests = responses_mock.requests(); + assert_eq!( + requests.len(), + 3, + "expected initial request, switched-model inline compaction request, and same-turn follow-up" + ); + + let follow_up_request = &requests[2]; + assert!( + follow_up_request.body_contains_text(&inline_summary), + "expected same-turn follow-up to include the inline compaction item" + ); + assert!( + follow_up_request.body_contains_text("AFTER_SWITCH_USER"), + "expected same-turn follow-up to retain the switched-model user input" + ); + assert!( + follow_up_request.body_contains_text(""), + "expected same-turn follow-up to preserve the original model switch update" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn auto_server_side_compaction_retries_without_committing_incomplete_checkpoint() -> Result<()> {