From db2a273b1226985cb8e29091b5de3f981db7809a Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Thu, 9 Jul 2026 14:36:03 -0700 Subject: [PATCH] Split step capture from environment refresh --- codex-rs/core/src/codex_thread.rs | 2 + codex-rs/core/src/compact_remote.rs | 5 ++- codex-rs/core/src/compact_remote_v2.rs | 5 ++- codex-rs/core/src/compact_token_budget.rs | 5 ++- codex-rs/core/src/prompt_debug.rs | 5 ++- codex-rs/core/src/session/mod.rs | 43 +++----------------- codex-rs/core/src/session/step_context.rs | 39 ++++++++++++++++++ codex-rs/core/src/session/tests.rs | 6 ++- codex-rs/core/src/session/turn.rs | 24 +++++++++-- codex-rs/core/src/session_startup_prewarm.rs | 1 + 10 files changed, 89 insertions(+), 46 deletions(-) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 257a32733b..b9e0a7e882 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -481,6 +481,7 @@ impl CodexThread { .codex .session .capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&self.codex.session) .await; self.codex .session @@ -610,6 +611,7 @@ impl CodexThread { self.codex .session .capture_step_context(turn_context) + .refresh_env(&self.codex.session) .await .mcp .clone() diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 12233197f2..3e4fdfa93f 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -75,7 +75,10 @@ pub(crate) async fn run_remote_compact_task( turn_context: Arc, ) -> CodexResult<()> { // Standalone compaction is its own request boundary, so it captures a fresh step. - let step_context = sess.capture_step_context(Arc::clone(&turn_context)).await; + let step_context = sess + .capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&sess) + .await; let start_event = EventMsg::TurnStarted(TurnStartedEvent { turn_id: turn_context.sub_id.clone(), trace_id: turn_context.trace_id.clone(), diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index bf9527f34e..c70876dc18 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -86,7 +86,10 @@ pub(crate) async fn run_remote_compact_task( turn_context: Arc, ) -> CodexResult<()> { // Standalone compaction is its own request boundary, so it captures a fresh step. - let step_context = sess.capture_step_context(Arc::clone(&turn_context)).await; + let step_context = sess + .capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&sess) + .await; let start_event = EventMsg::TurnStarted(TurnStartedEvent { turn_id: turn_context.sub_id.clone(), trace_id: turn_context.trace_id.clone(), diff --git a/codex-rs/core/src/compact_token_budget.rs b/codex-rs/core/src/compact_token_budget.rs index 0d71e7ab2f..8d51c6b3d3 100644 --- a/codex-rs/core/src/compact_token_budget.rs +++ b/codex-rs/core/src/compact_token_budget.rs @@ -36,7 +36,10 @@ pub(crate) async fn run_manual_compact_task( sess.send_event(&turn_context, start_event).await; // Manual compaction runs outside run_turn, so it captures its own current step. - let step_context = sess.capture_step_context(Arc::clone(&turn_context)).await; + let step_context = sess + .capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&sess) + .await; let world_state = Arc::new(sess.build_world_state_for_step(&step_context).await); run_compact_task_inner(&sess, &turn_context, world_state, CompactionTrigger::Manual).await } diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index cac9bb41ae..27307cf2be 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -78,7 +78,10 @@ pub(crate) async fn build_prompt_input_from_session( ) -> CodexResult> { let turn_context = sess.new_default_turn().await; // Prompt debugging builds a standalone request without entering run_turn. - let step_context = sess.capture_step_context(Arc::clone(&turn_context)).await; + let step_context = sess + .capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&sess) + .await; sess.record_context_updates_and_set_reference_context_item(step_context.as_ref()) .await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9339413b33..bf16ef5e58 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2880,49 +2880,18 @@ impl Session { world_state } - /// Captures one request-scoped view of dynamic state. - /// - /// This may refresh filesystem-derived state. Normal turns should call it only from - /// `run_turn` and pass the result down; standalone request or history boundaries may capture - /// their own step. - #[tracing::instrument(name = "step_context.capture", level = "info", skip_all)] - pub(crate) async fn capture_step_context( + /// Captures one lightweight request-scoped view before refreshing dynamic state. + pub(crate) fn capture_step_context( self: &Arc, turn_context: Arc, ) -> Arc { - let deferred_executor_enabled = turn_context - .config - .features - .enabled(Feature::DeferredExecutor); - // Keep the old turn-frozen environment view unless deferred executors are enabled. - let environments = if deferred_executor_enabled { - self.services.turn_environments.snapshot().await - } else { - turn_context.environments.clone() - }; - if deferred_executor_enabled { - self.services - .agents_md_manager - .refresh(&turn_context.config, &environments) - .await; - } - let loaded_agents_md = self.services.agents_md_manager.get_loaded().await; - let selected_capability_roots = self - .resolve_selected_capability_roots_for_step(&environments) - .await; - let mcp = self - .mcp_runtime_for_step( - turn_context.as_ref(), - &environments, - &selected_capability_roots, - ) - .await; + let environments = turn_context.environments.clone(); Arc::new(StepContext::new( turn_context, environments, - selected_capability_roots, - mcp, - loaded_agents_md, + Vec::new(), + self.services.latest_mcp_runtime(), + /*loaded_agents_md*/ None, )) } diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 1d7316d914..dc74ae65d9 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -3,8 +3,10 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::McpRuntimeSnapshot; +use crate::session::session::Session; use crate::session::turn_context::TurnContext; use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_features::Feature; use codex_mcp::ToolInfo; use tokio::sync::OnceCell; @@ -46,4 +48,41 @@ impl StepContext { .get_or_init(|| self.mcp.manager().list_all_tools()) .await } + + #[tracing::instrument(name = "step_context.capture", level = "info", skip_all)] + pub(crate) async fn refresh_env(self: Arc, session: &Arc) -> Arc { + let deferred_executor_enabled = + self.turn.config.features.enabled(Feature::DeferredExecutor); + // Keep the old turn-frozen environment view unless deferred executors are enabled. + let environments = if deferred_executor_enabled { + session.services.turn_environments.snapshot().await + } else { + self.turn.environments.clone() + }; + if deferred_executor_enabled { + session + .services + .agents_md_manager + .refresh(&self.turn.config, &environments) + .await; + } + let loaded_agents_md = session.services.agents_md_manager.get_loaded().await; + let selected_capability_roots = session + .resolve_selected_capability_roots_for_step(&environments) + .await; + let mcp = session + .mcp_runtime_for_step( + self.turn.as_ref(), + &environments, + &selected_capability_roots, + ) + .await; + Arc::new(Self::new( + Arc::clone(&self.turn), + environments, + selected_capability_roots, + mcp, + loaded_agents_md, + )) + } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index c1f36ad2b6..ae0afb1e3c 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7752,6 +7752,7 @@ async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { let old_runtime = session.services.latest_mcp_runtime(); let step_context = session .capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&session) .await; assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); let old_token = session.mcp_startup_cancellation_token().await; @@ -7901,7 +7902,10 @@ async fn built_tools_uses_the_step_mcp_runtime() -> anyhow::Result<()> { let (session, turn_context) = make_session_and_context().await; let session = Arc::new(session); let turn_context = Arc::new(turn_context); - let step_context = session.capture_step_context(turn_context).await; + let step_context = session + .capture_step_context(turn_context) + .refresh_env(&session) + .await; let mut refresh_config = step_context.turn.config.as_ref().clone(); refresh_config.mcp_servers.set(HashMap::from([( diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8811e2f0aa..465b183636 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -165,7 +165,10 @@ pub(crate) async fn run_turn( } // run_turn owns the step used to seed context and make the first sampling request. - let first_step_context = sess.capture_step_context(Arc::clone(&turn_context)).await; + let first_step_context = sess + .capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&sess) + .await; // Keep the exact model-visible state used by this turn and its inline compactions. let (mut world_state, display_roots) = tokio::join!( sess.record_context_updates_and_set_reference_context_item(first_step_context.as_ref()), @@ -247,7 +250,11 @@ pub(crate) async fn run_turn( // Capture once so context, advertised tools, and tool calls share one request view. let step_context = match next_step_context.take() { Some(step_context) => step_context, - None => sess.capture_step_context(Arc::clone(&turn_context)).await, + None => { + sess.capture_step_context(Arc::clone(&turn_context)) + .refresh_env(&sess) + .await + } }; let sampling_request_result: CodexResult<_> = async { super::time_reminder::maybe_record_current_time_reminder( @@ -807,7 +814,10 @@ async fn run_pre_sampling_compact( // Compact if the configured auto-compaction budget or usable context window is exhausted. if token_status.token_limit_reached { // Pre-turn compaction runs before run_turn creates the normal sampling step. - let step_context = sess.capture_step_context(Arc::clone(turn_context)).await; + let step_context = sess + .capture_step_context(Arc::clone(turn_context)) + .refresh_env(sess) + .await; run_auto_compact( sess, step_context, @@ -849,7 +859,11 @@ async fn capture_current_model_fallback_step_context( { return None; } - Some(sess.capture_step_context(Arc::clone(turn_context)).await) + Some( + sess.capture_step_context(Arc::clone(turn_context)) + .refresh_env(sess) + .await, + ) } /// Runs pre-sampling compaction against the previous model when its compaction compatibility @@ -878,6 +892,7 @@ async fn maybe_run_previous_model_inline_compact( if should_compact_for_comp_hash_change { let step_context = sess .capture_step_context(Arc::clone(&previous_model_turn_context)) + .refresh_env(sess) .await; let fallback_step_context = capture_current_model_fallback_step_context( sess, @@ -925,6 +940,7 @@ async fn maybe_run_previous_model_inline_compact( if should_run { let step_context = sess .capture_step_context(Arc::clone(&previous_model_turn_context)) + .refresh_env(sess) .await; let fallback_step_context = capture_current_model_fallback_step_context( sess, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index a5a4c5e437..09fa59083c 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -269,6 +269,7 @@ async fn schedule_startup_prewarm_inner( // Startup prewarm runs before run_turn and needs its own tool-building snapshot. let step_context = session .capture_step_context(Arc::clone(&startup_turn_context)) + .refresh_env(&session) .await; let startup_router = built_tools( session.as_ref(),