Add step-scoped data to extension contributors (#34525)

## What changed

- Add an `ExtensionData` store to each `StepContext`.
- Pass the step store to context, world-state, turn-input, and tool contributors so extensions can use capabilities bound to the current sampling step.
- Preserve the captured step context when rebuilding initial context during compaction.

GitOrigin-RevId: 51f3105d3a793f60534f122f9be8247c131c2ca5

codex-rs/core/src/compact.rs

@@ -15,6 +15,7 @@ use crate::responses_metadata::CompactionTurnMetadata;
 #[cfg(test)]
 use crate::session::PreviousTurnSettings;
 use crate::session::session::Session;
+use crate::session::step_context::StepContext;
 use crate::session::turn::get_last_assistant_message_from_turn;
 use crate::session::turn_context::TurnContext;
 use crate::util::backoff;
@@ -64,20 +65,30 @@ const COMPACT_USER_MESSAGE_MAX_TOKENS: usize = 20_000;
 /// initial context into the replacement history just above the last real user message.
 #[derive(Debug)]
 pub(crate) enum InitialContextInjection {
-    BeforeLastUserMessage(Arc<WorldState>),
+    BeforeLastUserMessage {
+        world_state: Arc<WorldState>,
+        step_context: Arc<StepContext>,
+    },
     DoNotInject,
 }
 
 pub(crate) async fn build_compaction_initial_context(
     sess: &Session,
-    turn_context: &TurnContext,
     initial_context_injection: &InitialContextInjection,
 ) -> (Vec<ResponseItem>, Option<Arc<WorldState>>) {
     // Return the rendered state with its items so history and its baseline stay identical.
     match initial_context_injection {
-        InitialContextInjection::BeforeLastUserMessage(world_state) => {
+        InitialContextInjection::BeforeLastUserMessage {
+            world_state,
+            step_context,
+        } => {
             let items = sess
-                .build_initial_context_with_world_state(turn_context, world_state.as_ref())
+                .build_initial_context_with_world_state_and_mcp(
+                    step_context.turn.as_ref(),
+                    world_state.as_ref(),
+                    step_context.mcp.as_ref(),
+                    &step_context.extension_data,
+                )
                 .await;
             (items, Some(Arc::clone(world_state)))
         }
@@ -334,19 +345,15 @@ async fn run_compact_task_inner_impl(
     }
     let (window_number, window_ids) = sess.advance_auto_compact_window().await;
 
-    let (initial_context, world_state_baseline) = build_compaction_initial_context(
-        sess.as_ref(),
-        turn_context.as_ref(),
-        &initial_context_injection,
-    )
-    .await;
+    let (initial_context, world_state_baseline) =
+        build_compaction_initial_context(sess.as_ref(), &initial_context_injection).await;
     if !initial_context.is_empty() {
         new_history =
             insert_initial_context_before_last_real_user_or_summary(new_history, initial_context);
     }
     let reference_context_item = match initial_context_injection {
         InitialContextInjection::DoNotInject => None,
-        InitialContextInjection::BeforeLastUserMessage(_) => {
+        InitialContextInjection::BeforeLastUserMessage { .. } => {
             Some(turn_context.to_turn_context_item())
         }
     };

codex-rs/core/src/compact_remote.rs

@@ -260,17 +260,12 @@ async fn run_remote_compact_task_inner_impl(
         trace_input_history,
     } = attempt;
     let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await;
-    let (new_history, world_state_baseline) = process_compacted_history(
-        sess.as_ref(),
-        compaction_turn_context.as_ref(),
-        new_history,
-        &initial_context_injection,
-    )
-    .await;
+    let (new_history, world_state_baseline) =
+        process_compacted_history(sess.as_ref(), new_history, &initial_context_injection).await;
 
     let reference_context_item = match initial_context_injection {
         InitialContextInjection::DoNotInject => None,
-        InitialContextInjection::BeforeLastUserMessage(_) => {
+        InitialContextInjection::BeforeLastUserMessage { .. } => {
             Some(compaction_turn_context.to_turn_context_item())
         }
     };
@@ -308,7 +303,6 @@ async fn run_remote_compact_task_inner_impl(
 
 pub(crate) async fn process_compacted_history(
     sess: &Session,
-    turn_context: &TurnContext,
     mut compacted_history: Vec<ResponseItem>,
     initial_context_injection: &InitialContextInjection,
 ) -> (Vec<ResponseItem>, Option<Arc<WorldState>>) {
@@ -316,7 +310,7 @@ pub(crate) async fn process_compacted_history(
     // 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, world_state_baseline) =
-        build_compaction_initial_context(sess, turn_context, initial_context_injection).await;
+        build_compaction_initial_context(sess, initial_context_injection).await;
 
     compacted_history.retain(should_keep_compacted_history_item);
     (

codex-rs/core/src/compact_remote_v2.rs

@@ -283,17 +283,13 @@ 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, world_state_baseline) = process_compacted_history(
-        sess.as_ref(),
-        compaction_turn_context.as_ref(),
-        compacted_history,
-        &initial_context_injection,
-    )
-    .await;
+    let (new_history, world_state_baseline) =
+        process_compacted_history(sess.as_ref(), compacted_history, &initial_context_injection)
+            .await;
 
     let reference_context_item = match initial_context_injection {
         InitialContextInjection::DoNotInject => None,
-        InitialContextInjection::BeforeLastUserMessage(_) => {
+        InitialContextInjection::BeforeLastUserMessage { .. } => {
             Some(compaction_turn_context.to_turn_context_item())
         }
     };

codex-rs/core/src/compact_tests.rs

@@ -1,5 +1,4 @@
 use super::*;
-use crate::session::tests::build_world_state_from_turn_context;
 use codex_model_provider_info::ModelProviderInfo;
 use codex_model_provider_info::WireApi;
 use codex_protocol::ResponseItemId;
@@ -17,14 +16,23 @@ async fn process_compacted_history_with_test_session(
     session
         .set_previous_turn_settings(previous_turn_settings.cloned())
         .await;
-    let world_state = Arc::new(build_world_state_from_turn_context(&session, &turn_context).await);
+    let step_context =
+        crate::session::step_context::StepContext::for_test(Arc::clone(&turn_context));
+    let world_state = Arc::new(session.build_world_state_for_step(&step_context).await);
     let initial_context = session
-        .build_initial_context_with_world_state(&turn_context, world_state.as_ref())
+        .build_initial_context_with_world_state_and_mcp(
+            &turn_context,
+            world_state.as_ref(),
+            step_context.mcp.as_ref(),
+            &step_context.extension_data,
+        )
         .await;
-    let initial_context_injection = InitialContextInjection::BeforeLastUserMessage(world_state);
+    let initial_context_injection = InitialContextInjection::BeforeLastUserMessage {
+        world_state,
+        step_context,
+    };
     let (refreshed, _) = crate::compact_remote::process_compacted_history(
         &session,
-        &turn_context,
         compacted_history,
         &initial_context_injection,
     )

codex-rs/core/src/compact_token_budget.rs

@@ -38,7 +38,7 @@ pub(crate) async fn run_manual_compact_task(
     // 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 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
+    run_compact_task_inner(&sess, &step_context, world_state, CompactionTrigger::Manual).await
 }
 
 /// Runs token-budget inline auto-compaction as a normal compaction lifecycle.
@@ -51,22 +51,22 @@ pub(crate) async fn run_inline_auto_compact_task(
     step_context: Arc<StepContext>,
     initial_context_injection: InitialContextInjection,
 ) -> CodexResult<()> {
-    let turn_context = &step_context.turn;
     let world_state = match initial_context_injection {
-        InitialContextInjection::BeforeLastUserMessage(world_state) => world_state,
+        InitialContextInjection::BeforeLastUserMessage { world_state, .. } => world_state,
         InitialContextInjection::DoNotInject => {
             Arc::new(sess.build_world_state_for_step(&step_context).await)
         }
     };
-    run_compact_task_inner(&sess, turn_context, world_state, CompactionTrigger::Auto).await
+    run_compact_task_inner(&sess, &step_context, world_state, CompactionTrigger::Auto).await
 }
 
 async fn run_compact_task_inner(
     sess: &Arc<Session>,
-    turn_context: &Arc<TurnContext>,
+    step_context: &Arc<StepContext>,
     world_state: Arc<WorldState>,
     trigger: CompactionTrigger,
 ) -> CodexResult<()> {
+    let turn_context = &step_context.turn;
     let pre_compact_outcome = run_pre_compact_hooks(sess, turn_context, trigger).await;
     match pre_compact_outcome {
         PreCompactHookOutcome::Continue => {}
@@ -76,7 +76,7 @@ async fn run_compact_task_inner(
     let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new());
     sess.emit_turn_item_started(turn_context, &compaction_item)
         .await;
-    sess.start_new_context_window(turn_context.as_ref(), world_state)
+    sess.start_new_context_window(step_context, world_state)
         .await;
     sess.emit_turn_item_completed(turn_context, compaction_item)
         .await;

codex-rs/core/src/guardian/tests.rs

@@ -114,6 +114,7 @@ impl codex_extension_api::ContextContributor for GuardianMemoryContextProbe {
         &'a self,
         _session_store: &'a codex_extension_api::ExtensionData,
         thread_store: &'a codex_extension_api::ExtensionData,
+        _step_store: &'a codex_extension_api::ExtensionData,
     ) -> codex_extension_api::ExtensionFuture<'a, Vec<codex_extension_api::PromptFragment>> {
         Box::pin(async move {
             if thread_store

codex-rs/core/src/session/mod.rs

@@ -3152,8 +3152,9 @@ impl Session {
 
     async fn build_turn_context_contribution_items(
         &self,
-        turn_context: &TurnContext,
+        step_context: &StepContext,
     ) -> Vec<ResponseItem> {
+        let turn_context = step_context.turn.as_ref();
         let mut developer_sections = Vec::new();
         let mut contextual_user_sections = Vec::new();
         let mut separate_developer_sections = Vec::new();
@@ -3167,6 +3168,7 @@ impl Session {
                     session_store: &self.services.session_extension_data,
                     thread_store: &self.services.thread_extension_data,
                     turn_store: turn_context.extension_data.as_ref(),
+                    step_store: &step_context.extension_data,
                     model_context_window: turn_context.model_context_window(),
                 })
                 .await
@@ -3201,21 +3203,29 @@ impl Session {
         items
     }
 
+    #[cfg(test)]
     pub(crate) async fn build_initial_context_with_world_state(
         &self,
         turn_context: &TurnContext,
         world_state: &WorldState,
     ) -> Vec<ResponseItem> {
         let mcp = self.services.latest_mcp_runtime();
-        self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp)
-            .await
+        let step_store = codex_extension_api::ExtensionData::new(turn_context.sub_id.clone());
+        self.build_initial_context_with_world_state_and_mcp(
+            turn_context,
+            world_state,
+            &mcp,
+            &step_store,
+        )
+        .await
     }
 
-    async fn build_initial_context_with_world_state_and_mcp(
+    pub(crate) async fn build_initial_context_with_world_state_and_mcp(
         &self,
         turn_context: &TurnContext,
         world_state: &WorldState,
         mcp: &McpRuntimeSnapshot,
+        step_store: &codex_extension_api::ExtensionData,
     ) -> Vec<ResponseItem> {
         let mut developer_sections = Vec::<String>::with_capacity(8);
         let mut contextual_user_sections = Vec::<String>::with_capacity(2);
@@ -3330,6 +3340,7 @@ impl Session {
                 .contribute_thread_context(
                     &self.services.session_extension_data,
                     &self.services.thread_extension_data,
+                    step_store,
                 )
                 .await
             {
@@ -3349,6 +3360,7 @@ impl Session {
                     session_store: &self.services.session_extension_data,
                     thread_store: &self.services.thread_extension_data,
                     turn_store: turn_context.extension_data.as_ref(),
+                    step_store,
                     model_context_window: turn_context.model_context_window(),
                 })
                 .await
@@ -3509,16 +3521,22 @@ impl Session {
 
     pub(crate) async fn start_new_context_window(
         &self,
-        turn_context: &TurnContext,
+        step_context: &StepContext,
         world_state: Arc<WorldState>,
     ) -> u64 {
+        let turn_context = step_context.turn.as_ref();
         let window = {
             let mut state = self.state.lock().await;
             state.start_new_context_window()
         };
         let (window_number, window_ids) = window;
         let context_items = self
-            .build_initial_context_with_world_state(turn_context, world_state.as_ref())
+            .build_initial_context_with_world_state_and_mcp(
+                turn_context,
+                world_state.as_ref(),
+                step_context.mcp.as_ref(),
+                &step_context.extension_data,
+            )
             .await;
         let turn_context_item = turn_context.to_turn_context_item();
         self.replace_compacted_history(
@@ -3579,6 +3597,7 @@ impl Session {
                     turn_context,
                     world_state.as_ref(),
                     step_context.mcp.as_ref(),
+                    &step_context.extension_data,
                 )
                 .await;
             let snapshot = world_state.snapshot();
@@ -3611,7 +3630,7 @@ impl Session {
         };
         if !should_inject_full_context && turn_context_changed {
             context_items.extend(
-                self.build_turn_context_contribution_items(turn_context)
+                self.build_turn_context_contribution_items(step_context)
                     .await,
             );
         }

codex-rs/core/src/session/step_context.rs

@@ -6,6 +6,7 @@ use crate::session::McpRuntimeSnapshot;
 use crate::session::turn_context::TurnContext;
 use codex_exec_server::ExecutorCapabilityDiscoverySnapshot;
 use codex_exec_server::ResolvedSelectedCapabilityRoot;
+use codex_extension_api::ExtensionData;
 use codex_mcp::ToolInfo;
 use tokio::sync::OnceCell;
 
@@ -22,6 +23,8 @@ pub(crate) struct StepContext {
     pub(crate) mcp: Arc<McpRuntimeSnapshot>,
     /// The fixed MCP tool list used for this exact sampling request.
     mcp_tool_snapshot: OnceCell<Vec<ToolInfo>>,
+    /// Extension capabilities bound to this exact sampling step.
+    pub(crate) extension_data: ExtensionData,
     /// The canonical AGENTS.md value observed with this environment snapshot.
     pub(crate) loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
 }
@@ -35,6 +38,7 @@ impl StepContext {
         mcp: Arc<McpRuntimeSnapshot>,
         loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
     ) -> Self {
+        let extension_data = ExtensionData::new(turn.sub_id.clone());
         Self {
             turn,
             environments,
@@ -42,6 +46,7 @@ impl StepContext {
             executor_capability_discovery,
             mcp,
             mcp_tool_snapshot: OnceCell::new(),
+            extension_data,
             loaded_agents_md,
         }
     }

codex-rs/core/src/session/tests.rs

@@ -2815,11 +2815,13 @@ async fn start_new_context_window_assigns_and_persists_item_ids() {
     .await;
     let rollout_path =
         attach_thread_persistence(Arc::get_mut(&mut session).expect("unique session")).await;
-    let world_state =
-        Arc::new(build_world_state_from_turn_context(session.as_ref(), &turn_context).await);
+    let step_context = session
+        .capture_step_context(Arc::clone(&turn_context))
+        .await;
+    let world_state = Arc::new(session.build_world_state_for_step(&step_context).await);
 
     session
-        .start_new_context_window(turn_context.as_ref(), world_state)
+        .start_new_context_window(&step_context, world_state)
         .await;
 
     let live_history = session.clone_history().await;
@@ -8274,6 +8276,7 @@ impl codex_extension_api::ContextContributor for PromptExtensionTestContributor
         &'a self,
         _session_store: &'a codex_extension_api::ExtensionData,
         thread_store: &'a codex_extension_api::ExtensionData,
+        _step_store: &'a codex_extension_api::ExtensionData,
     ) -> std::pin::Pin<
         Box<dyn std::future::Future<Output = Vec<codex_extension_api::PromptFragment>> + Send + 'a>,
     > {

codex-rs/core/src/session/tests/guardian_tests.rs

@@ -539,12 +539,15 @@ 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 turn_context = Arc::new(turn_context);
-    let world_state = Arc::new(build_world_state_from_turn_context(&session, &turn_context).await);
-    let initial_context_injection = InitialContextInjection::BeforeLastUserMessage(world_state);
+    let step_context = StepContext::for_test(Arc::clone(&turn_context));
+    let world_state = Arc::new(session.build_world_state_for_step(&step_context).await);
+    let initial_context_injection = InitialContextInjection::BeforeLastUserMessage {
+        world_state,
+        step_context,
+    };
 
     let (refreshed, _) = crate::compact_remote::process_compacted_history(
         &session,
-        &turn_context,
         vec![
             ResponseItem::Message {
                 id: None,

codex-rs/core/src/session/turn.rs

@@ -357,7 +357,10 @@ pub(crate) async fn run_turn(
                         Arc::clone(&step_context),
                         /*fallback_step_context*/ None,
                         &mut client_session,
-                        InitialContextInjection::BeforeLastUserMessage(Arc::clone(&world_state)),
+                        InitialContextInjection::BeforeLastUserMessage {
+                            world_state: Arc::clone(&world_state),
+                            step_context: Arc::clone(&step_context),
+                        },
                         CompactionReason::ContextLimit,
                         CompactionPhase::MidTurn,
                     )
@@ -584,7 +587,7 @@ async fn build_skills_and_plugins(
     let skills_outcome = turn_context.turn_skills.snapshot.outcome();
     let connector_slug_counts = build_connector_slug_counts(&available_connectors);
     let extension_injection_items =
-        build_extension_turn_input_items(sess, turn_context, &user_input, cancellation_token)
+        build_extension_turn_input_items(sess, step_context, &user_input, cancellation_token)
             .await?;
     let skill_name_counts_lower =
         build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1;
@@ -687,10 +690,11 @@ async fn build_skills_and_plugins(
 )]
 async fn build_extension_turn_input_items(
     sess: &Arc<Session>,
-    turn_context: &TurnContext,
+    step_context: &StepContext,
     user_input: &[UserInput],
     cancellation_token: &CancellationToken,
 ) -> Option<Vec<ResponseItem>> {
+    let turn_context = step_context.turn.as_ref();
     let contributors = sess.services.extensions.turn_input_contributors().to_vec();
     if contributors.is_empty() {
         return Some(Vec::new());
@@ -725,6 +729,7 @@ async fn build_extension_turn_input_items(
                 &sess.services.session_extension_data,
                 &sess.services.thread_extension_data,
                 turn_context.extension_data.as_ref(),
+                &step_context.extension_data,
             )
             .or_cancel(cancellation_token)
             .await
@@ -1341,7 +1346,7 @@ pub(crate) async fn built_tools(
         ToolRouterParams {
             tool_runtimes: mcp_tool_runtimes,
             tool_suggest_candidates,
-            extension_tool_executors: extension_tool_executors(sess),
+            extension_tool_executors: extension_tool_executors(sess, step_context),
             dynamic_tools: turn_context.dynamic_tools.as_slice(),
         },
         &sess.services.tool_search_handler_cache,

codex-rs/core/src/session/world_state.rs

@@ -124,6 +124,7 @@ impl Session {
                     session_store: &self.services.session_extension_data,
                     thread_store: &self.services.thread_extension_data,
                     turn_store: turn_context.extension_data.as_ref(),
+                    step_store: &step_context.extension_data,
                 })
                 .await
             {

codex-rs/core/src/tools/router.rs

@@ -246,6 +246,7 @@ impl ToolRouter {
 #[instrument(level = "trace", skip_all)]
 pub(crate) fn extension_tool_executors(
     session: &Session,
+    step_context: &StepContext,
 ) -> Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>> {
     session
         .services
@@ -256,6 +257,7 @@ pub(crate) fn extension_tool_executors(
             contributor.tools(
                 &session.services.session_extension_data,
                 &session.services.thread_extension_data,
+                &step_context.extension_data,
             )
         })
         .collect()

codex-rs/core/src/tools/router_tests.rs

@@ -43,6 +43,7 @@ impl codex_extension_api::ToolContributor for ExtensionEchoContributor {
         &self,
         _session_store: &ExtensionData,
         _thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn ToolExecutor<ExtensionToolCall>>> {
         vec![Arc::new(ExtensionEchoExecutor)]
     }
@@ -390,7 +391,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow
         ToolRouterParams {
             tool_suggest_candidates: None,
             tool_runtimes: Vec::new(),
-            extension_tool_executors: extension_tool_executors(&session),
+            extension_tool_executors: extension_tool_executors(&session, step_context.as_ref()),
             dynamic_tools: turn.dynamic_tools.as_slice(),
         },
         &Default::default(),