From de414dc785e3372f87ca474046127888698f38d0 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Mon, 22 Jun 2026 23:38:43 -0700 Subject: [PATCH] core: require explicit environment snapshots --- codex-rs/core/src/compact_tests.rs | 6 +++++- codex-rs/core/src/session/mod.rs | 12 +++++------- codex-rs/core/src/session/tests.rs | 16 +++++++++++++--- .../core/src/session/tests/guardian_tests.rs | 6 +++++- codex-rs/core/src/thread_manager_tests.rs | 4 +++- .../core/src/thread_rollout_truncation_tests.rs | 4 +++- 6 files changed, 34 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index b17361a0ea..8dd4dd4778 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -14,7 +14,11 @@ async fn process_compacted_history_with_test_session( session .set_previous_turn_settings(previous_turn_settings.cloned()) .await; - let world_state = Arc::new(session.build_world_state(&turn_context).await); + let world_state = Arc::new( + session + .build_world_state_for_environments(&turn_context, &turn_context.environments) + .await, + ); let initial_context = session .build_initial_context_with_world_state(&turn_context, world_state.as_ref()) .await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 339e0444be..a51cadd576 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3030,7 +3030,7 @@ impl Session { items } - async fn build_world_state_for_environments( + pub(crate) async fn build_world_state_for_environments( &self, turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, @@ -3050,11 +3050,6 @@ impl Session { ) } - pub(crate) async fn build_world_state(&self, turn_context: &TurnContext) -> WorldState { - self.build_world_state_for_environments(turn_context, &turn_context.environments) - .await - } - pub(crate) async fn build_initial_context_with_world_state( &self, turn_context: &TurnContext, @@ -3471,7 +3466,10 @@ impl Session { let turn_context_item = turn_context.to_turn_context_item(); let turn_context_changed = reference_context_item.as_ref() != Some(&turn_context_item); let should_inject_full_context = reference_context_item.is_none(); - let world_state = Arc::new(self.build_world_state(turn_context).await); + let world_state = Arc::new( + self.build_world_state_for_environments(turn_context, &turn_context.environments) + .await, + ); let mut context_items = if should_inject_full_context { let context_items = self .build_initial_context_with_world_state(turn_context, world_state.as_ref()) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 5f3f79ab57..c29dd576b6 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5021,7 +5021,9 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { } async fn build_initial_context(session: &Session, turn_context: &TurnContext) -> Vec { - let world_state = session.build_world_state(turn_context).await; + let world_state = session + .build_world_state_for_environments(turn_context, &turn_context.environments) + .await; session .build_initial_context_with_world_state(turn_context, &world_state) .await @@ -7853,7 +7855,11 @@ async fn record_context_updates_includes_turn_context_fragments_on_steady_state_ }); let mut previous_context_item = turn_context.to_turn_context_item(); previous_context_item.turn_id = Some("previous-turn-id".to_string()); - let world_state = Arc::new(session.build_world_state(&turn_context).await); + let world_state = Arc::new( + session + .build_world_state_for_environments(&turn_context, &turn_context.environments) + .await, + ); { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item)); @@ -8527,7 +8533,11 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline .with_model(next_model.to_string(), &session.services.models_manager) .await; let previous_context_item = previous_context.to_turn_context_item(); - let world_state = Arc::new(session.build_world_state(&previous_context).await); + let world_state = Arc::new( + session + .build_world_state_for_environments(&previous_context, &previous_context.environments) + .await, + ); { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item.clone())); diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index b4668b358e..d465f975c0 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -523,7 +523,11 @@ 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 world_state = Arc::new(session.build_world_state(&turn_context).await); + let world_state = Arc::new( + session + .build_world_state_for_environments(&turn_context, &turn_context.environments) + .await, + ); let initial_context_injection = InitialContextInjection::BeforeLastUserMessage(world_state); let (refreshed, _) = crate::compact_remote::process_compacted_history( diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 9092d63689..ffc0a4121a 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -225,7 +225,9 @@ fn out_of_range_truncation_drops_pre_user_active_turn_prefix() { #[tokio::test] async fn ignores_session_prefix_messages_when_truncating() { let (session, turn_context) = make_session_and_context().await; - let world_state = session.build_world_state(&turn_context).await; + let world_state = session + .build_world_state_for_environments(&turn_context, &turn_context.environments) + .await; let mut items = session .build_initial_context_with_world_state(&turn_context, &world_state) .await; diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index ca3cabe60c..513bfee1a7 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -166,7 +166,9 @@ fn truncates_rollout_from_start_applies_thread_rollback_markers() { #[tokio::test] async fn ignores_session_prefix_messages_when_truncating_rollout_from_start() { let (session, turn_context) = make_session_and_context().await; - let world_state = session.build_world_state(&turn_context).await; + let world_state = session + .build_world_state_for_environments(&turn_context, &turn_context.environments) + .await; let mut items = session .build_initial_context_with_world_state(&turn_context, &world_state) .await;