From b680bfccee31face7c58efcd80d6d9192d221e8d Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Mon, 22 Jun 2026 19:18:07 -0700 Subject: [PATCH] core: use current step state after compaction --- codex-rs/core/src/compact.rs | 29 +++- codex-rs/core/src/compact_remote.rs | 24 ++-- codex-rs/core/src/compact_remote_v2.rs | 3 +- codex-rs/core/src/compact_tests.rs | 2 +- codex-rs/core/src/session/mod.rs | 53 ++++++-- .../core/src/session/tests/guardian_tests.rs | 2 +- codex-rs/core/src/session/turn.rs | 5 +- codex-rs/core/tests/suite/remote_env.rs | 128 ++++++++++++++++++ 8 files changed, 215 insertions(+), 31 deletions(-) diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index c60bd155e8..7bb79a93b0 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -4,6 +4,7 @@ use std::time::Instant; use crate::Prompt; use crate::client::ModelClientSession; use crate::client_common::ResponseEvent; +use crate::context::world_state::WorldState; use crate::hook_runtime::PostCompactHookOutcome; use crate::hook_runtime::PreCompactHookOutcome; use crate::hook_runtime::run_post_compact_hooks; @@ -66,6 +67,23 @@ pub(crate) enum InitialContextInjection { DoNotInject, } +pub(crate) async fn build_compaction_initial_context( + sess: &Session, + turn_context: &TurnContext, + initial_context_injection: InitialContextInjection, +) -> (Vec, Option) { + // Return the rendered state with its items so history and its baseline stay identical. + match initial_context_injection { + InitialContextInjection::BeforeLastUserMessage => { + let (items, world_state) = sess + .build_initial_context_for_current_step(turn_context) + .await; + (items, Some(world_state)) + } + InitialContextInjection::DoNotInject => (Vec::new(), None), + } +} + pub(crate) fn should_use_remote_compact_task(provider: &ModelProviderInfo) -> bool { provider.supports_remote_compaction() } @@ -309,11 +327,13 @@ async fn run_compact_task_inner_impl( } let (window_number, window_ids) = sess.advance_auto_compact_window().await; - if matches!( + let (initial_context, world_state_baseline) = build_compaction_initial_context( + sess.as_ref(), + turn_context.as_ref(), initial_context_injection, - InitialContextInjection::BeforeLastUserMessage - ) { - let initial_context = sess.build_initial_context(turn_context.as_ref()).await; + ) + .await; + if !initial_context.is_empty() { new_history = insert_initial_context_before_last_real_user_or_summary(new_history, initial_context); } @@ -333,6 +353,7 @@ async fn run_compact_task_inner_impl( turn_context.as_ref(), new_history, reference_context_item, + world_state_baseline, compacted_item, ) .await; diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 3d23140684..2a433bafa6 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -6,8 +6,10 @@ use crate::client::CompactConversationRequestSettings; use crate::compact::CompactionAnalyticsAttempt; use crate::compact::CompactionAnalyticsDetails; use crate::compact::InitialContextInjection; +use crate::compact::build_compaction_initial_context; use crate::compact::compaction_status_from_result; use crate::compact::insert_initial_context_before_last_real_user_or_summary; +use crate::context::world_state::WorldState; use crate::context_manager::ContextManager; use crate::hook_runtime::PostCompactHookOutcome; use crate::hook_runtime::PreCompactHookOutcome; @@ -237,7 +239,7 @@ async fn run_remote_compact_task_inner_impl( window_id, CodexResponsesRequestKind::Compaction(compaction_metadata), ); - let mut new_history = sess + let new_history = sess .services .model_client .compact_conversation_history( @@ -259,7 +261,7 @@ async fn run_remote_compact_task_inner_impl( ) .await?; let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; - new_history = process_compacted_history( + let (new_history, world_state_baseline) = process_compacted_history( sess.as_ref(), turn_context.as_ref(), new_history, @@ -290,6 +292,7 @@ async fn run_remote_compact_task_inner_impl( turn_context.as_ref(), new_history, reference_context_item, + world_state_baseline, compacted_item, ) .await; @@ -305,21 +308,18 @@ pub(crate) async fn process_compacted_history( turn_context: &TurnContext, mut compacted_history: Vec, initial_context_injection: InitialContextInjection, -) -> Vec { +) -> (Vec, Option) { // Mid-turn compaction is the only path that must inject initial context above the last user // message in the replacement history. Pre-turn compaction instead injects context after the // compaction item, but mid-turn compaction keeps the compaction item last for model training. - let initial_context = if matches!( - initial_context_injection, - InitialContextInjection::BeforeLastUserMessage - ) { - sess.build_initial_context(turn_context).await - } else { - Vec::new() - }; + let (initial_context, world_state_baseline) = + build_compaction_initial_context(sess, turn_context, initial_context_injection).await; compacted_history.retain(should_keep_compacted_history_item); - insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context) + ( + insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context), + world_state_baseline, + ) } /// Returns whether an item from remote compaction output should be preserved. diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 54d139a0e1..8d2d6a7cd5 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -293,7 +293,7 @@ async fn run_remote_compact_task_inner_impl( build_v2_compacted_history(&prompt_input, compaction_output); analytics_details.retained_image_count = Some(retained_images); let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; - let new_history = process_compacted_history( + let (new_history, world_state_baseline) = process_compacted_history( sess.as_ref(), turn_context.as_ref(), compacted_history, @@ -321,6 +321,7 @@ async fn run_remote_compact_task_inner_impl( turn_context.as_ref(), new_history, reference_context_item, + world_state_baseline, compacted_item, ) .await; diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index bb6adf643d..ac6fec73ea 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -14,7 +14,7 @@ async fn process_compacted_history_with_test_session( .set_previous_turn_settings(previous_turn_settings.cloned()) .await; let initial_context = session.build_initial_context(&turn_context).await; - let refreshed = crate::compact_remote::process_compacted_history( + let (refreshed, _) = crate::compact_remote::process_compacted_history( &session, &turn_context, compacted_history, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 1bd5bae8bf..9ca9bea5e0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -32,6 +32,7 @@ use crate::context::NetworkRuleSaved; use crate::context::PermissionsInstructions; use crate::context::PersonalitySpecInstructions; use crate::context::RecommendedPluginsInstructions; +use crate::context::world_state::EnvironmentsState; use crate::context::world_state::WorldState; use crate::current_time::TimeProvider; use crate::default_skill_metadata_budget; @@ -238,6 +239,7 @@ use self::session::AppServerClientMetadata; use self::session::Session; use self::session::SessionConfiguration; pub(crate) use self::session::SessionSettingsUpdate; +use self::step_context::StepContext; #[cfg(test)] use self::turn::AssistantMessageStreamParsers; #[cfg(test)] @@ -2785,6 +2787,19 @@ impl Session { self.record_conversation_items(turn_context, &items).await; } + pub(crate) async fn capture_step_context(&self, turn_context: &TurnContext) -> StepContext { + let environments = if turn_context + .config + .features + .enabled(Feature::DeferredExecutor) + { + self.services.turn_environments.snapshot().await + } else { + turn_context.environments.clone() + }; + StepContext { environments } + } + pub(crate) async fn record_inter_agent_communication( &self, turn_context: &TurnContext, @@ -2884,6 +2899,7 @@ impl Session { turn_context: &TurnContext, items: Vec, reference_context_item: Option, + world_state_baseline: Option, compacted_item: CompactedItem, ) { let items = if turn_context.config.features.enabled(Feature::ItemIds) { @@ -2895,11 +2911,6 @@ impl Session { replacement_history: Some(items.clone()), ..compacted_item }; - let world_state_baseline = if reference_context_item.is_some() { - Some(self.build_world_state(turn_context).await) - } else { - None - }; { let mut state = self.state.lock().await; state.replace_history(items, reference_context_item.clone()); @@ -3031,6 +3042,7 @@ impl Session { items } + #[cfg(test)] pub(crate) async fn build_initial_context( &self, turn_context: &TurnContext, @@ -3040,6 +3052,32 @@ impl Session { .await } + pub(crate) async fn build_initial_context_for_current_step( + &self, + turn_context: &TurnContext, + ) -> (Vec, WorldState) { + let step_context = self.capture_step_context(turn_context).await; + let mut world_state = WorldState::default(); + if turn_context.config.include_environment_context { + let environment_subagents = self + .services + .agent_control + .format_environment_context_subagents(self.thread_id) + .await; + world_state.add_section( + EnvironmentsState::from_turn_context_with_environments( + turn_context, + &step_context.environments, + ) + .with_subagents(environment_subagents), + ); + } + let items = self + .build_initial_context_with_world_state(turn_context, &world_state) + .await; + (items, world_state) + } + async fn build_world_state(&self, turn_context: &TurnContext) -> WorldState { let environment_subagents = if turn_context.config.include_environment_context { self.services @@ -3407,9 +3445,8 @@ impl Session { state.start_new_context_window_if_requested() }; let (window_number, window_ids) = window?; - let world_state = self.build_world_state(turn_context).await; - let context_items = self - .build_initial_context_with_world_state(turn_context, &world_state) + let (context_items, world_state) = self + .build_initial_context_for_current_step(turn_context) .await; let turn_context_item = turn_context.to_turn_context_item(); let replacement_history = context_items; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index e44b12ef72..dc0979676c 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -524,7 +524,7 @@ async fn process_compacted_history_preserves_separate_guardian_developer_message turn_context.session_source = guardian_source; turn_context.developer_instructions = Some(guardian_policy.clone()); - let refreshed = crate::compact_remote::process_compacted_history( + let (refreshed, _) = crate::compact_remote::process_compacted_history( &session, &turn_context, vec![ diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 81f9f52356..40aa96cea9 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -42,7 +42,6 @@ use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::PreviousTurnSettings; use crate::session::TurnInput; use crate::session::session::Session; -use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::HandleOutputCtx; use crate::stream_events_utils::TurnItemContributorPolicy; @@ -247,9 +246,7 @@ pub(crate) async fn run_turn( .features .enabled(Feature::DeferredExecutor) { - let step_context = StepContext { - environments: sess.services.turn_environments.snapshot().await, - }; + let step_context = sess.capture_step_context(turn_context.as_ref()).await; sess.record_step_environment_context_if_changed( turn_context.as_ref(), &step_context, diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 2df03f92a8..63cca20b68 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -1,6 +1,7 @@ use anyhow::Context; use anyhow::Result; use codex_config::types::ApprovalsReviewer; +use codex_core::compact::SUMMARIZATION_PROMPT; use codex_core::config::Constrained; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; @@ -38,6 +39,7 @@ use core_test_support::get_remote_test_env; use core_test_support::responses::ev_apply_patch_custom_tool_call; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once; @@ -683,6 +685,132 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn deferred_executor_compaction_reinjects_current_environment_once() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let server = start_mock_server().await; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + "wait-for-startup", + "request_user_input", + &json!({ + "questions": [{ + "id": "continue", + "header": "Continue", + "question": "Continue after startup?", + "options": [{ + "label": "Yes (Recommended)", + "description": "Continue the test." + }, { + "label": "No", + "description": "Stop the test." + }] + }] + }) + .to_string(), + ), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 96), + ]), + sse(vec![ + ev_assistant_message("msg-compact", "AUTO_COMPACT_SUMMARY"), + ev_completed_with_tokens("resp-compact", /*total_tokens*/ 10), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + let mut builder = test_codex() + .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) + .with_config(|config| { + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!( + config + .features + .enable(Feature::DefaultModeRequestUserInput) + .is_ok() + ); + config.model_provider.name = "OpenAI (test)".to_string(); + config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string()); + config.model_context_window = Some(100); + config.model_auto_compact_token_limit = Some(90); + }); + let test = timeout(Duration::from_secs(5), builder.build(&server)) + .await + .context("thread startup should not wait for the remote environment")??; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "wait for the environment".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + + serve_environment_info(listener).await; + test.codex + .submit(Op::UserInputAnswer { + id: request.turn_id, + response: RequestUserInputResponse { + answers: HashMap::from([( + "continue".to_string(), + RequestUserInputAnswer { + answers: vec!["Yes (Recommended)".to_string()], + }, + )]), + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + let initial_context = requests[0].message_input_texts("user"); + assert!( + initial_context + .iter() + .any(|text| text.contains("starting")) + ); + + let post_compaction_context = requests[2].message_input_texts("user"); + assert_eq!( + post_compaction_context + .iter() + .filter(|text| text.contains("starting")) + .count(), + 0 + ); + assert_eq!( + post_compaction_context + .iter() + .filter(|text| text.contains("zsh")) + .count(), + 1 + ); + + Ok(()) +} + fn absolute_path(path: PathBuf) -> AbsolutePathBuf { AbsolutePathBuf::try_from(path).expect("path should be absolute") }