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(
+    let (initial_context, world_state_baseline) =
-        sess.as_ref(),
+        build_compaction_initial_context(sess.as_ref(), &initial_context_injection).await;
-        turn_context.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(
+    let (new_history, world_state_baseline) =
-        sess.as_ref(),
+        process_compacted_history(sess.as_ref(), new_history, &initial_context_injection).await;
-        compaction_turn_context.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(
+    let (new_history, world_state_baseline) =
-        sess.as_ref(),
+        process_compacted_history(sess.as_ref(), compacted_history, &initial_context_injection)
-        compaction_turn_context.as_ref(),
+            .await;
-        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)
+        let step_store = codex_extension_api::ExtensionData::new(turn_context.sub_id.clone());
-            .await
+        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 =
+    let step_context = session
-        Arc::new(build_world_state_from_turn_context(session.as_ref(), &turn_context).await);
+        .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 step_context = StepContext::for_test(Arc::clone(&turn_context));
-    let initial_context_injection = InitialContextInjection::BeforeLastUserMessage(world_state);
+    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(),

codex-rs/ext/extension-api/examples/enabled_extensions.rs

@@ -73,10 +73,11 @@ async fn contribute_prompt(
     thread_store: &ExtensionData,
 ) -> Vec<codex_extension_api::PromptFragment> {
     let mut fragments = Vec::new();
+    let step_store = ExtensionData::new("step");
     for contributor in registry.context_contributors() {
         fragments.extend(
             contributor
-                .contribute_thread_context(session_store, thread_store)
+                .contribute_thread_context(session_store, thread_store, &step_store)
                 .await,
         );
     }

codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs

@@ -21,6 +21,7 @@ impl ContextContributor for StyleContributor {
         &'a self,
         session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
         Box::pin(async move {
             contribution_counts(session_store).record_style();
@@ -41,6 +42,7 @@ impl ContextContributor for UsageContributor {
         &'a self,
         session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
         Box::pin(async move {
             contribution_counts(session_store).record_usage();

codex-rs/ext/extension-api/src/contributors.rs

@@ -78,15 +78,18 @@ pub trait McpServerContributor<C: Sync>: Send + Sync {
 /// fragment: thread/session context for stable inputs, and turn context for
 /// fragments that depend on turn-local host state.
 pub trait ContextContributor: Send + Sync {
+    /// Returns thread-scoped context using capabilities from the current sampling step.
     fn contribute_thread_context<'a>(
         &'a self,
         session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
+        step_store: &'a ExtensionData,
     ) -> ExtensionFuture<'a, Vec<PromptFragment>> {
         Box::pin(async move {
             let _self = self;
             let _session_store = session_store;
             let _thread_store = thread_store;
+            let _step_store = step_store;
             Vec::new()
         })
     }
@@ -205,12 +208,16 @@ pub trait TurnLifecycleContributor: Send + Sync {
 /// host, not in this input.
 pub trait TurnInputContributor: Send + Sync {
     /// Returns additional contextual fragments for one submitted turn.
+    ///
+    /// `step_store` contains host capabilities bound to the sampling step that
+    /// will consume these fragments.
     fn contribute<'a>(
         &'a self,
         input: TurnInputContext,
         session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
         turn_store: &'a ExtensionData,
+        step_store: &'a ExtensionData,
     ) -> ExtensionFuture<'a, Vec<Box<dyn ContextualUserFragment + Send>>>;
 }
 
@@ -270,11 +277,12 @@ pub trait SkillInvocationContributor: Send + Sync {
 
 /// Extension contribution that exposes native tools owned by a feature.
 pub trait ToolContributor: Send + Sync {
-    /// Returns the native tools visible for the supplied extension stores.
+    /// Returns native tools bound to the supplied sampling-step capabilities.
     fn tools(
         &self,
         session_store: &ExtensionData,
         thread_store: &ExtensionData,
+        step_store: &ExtensionData,
     ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>>;
 }
 

codex-rs/ext/extension-api/src/contributors/context.rs

@@ -15,6 +15,8 @@ pub struct TurnContextContributionInput<'a> {
     pub thread_store: &'a ExtensionData,
     /// Store scoped to this turn.
     pub turn_store: &'a ExtensionData,
+    /// Store scoped to this sampling step.
+    pub step_store: &'a ExtensionData,
     /// Effective model context window for this turn, when known.
     pub model_context_window: Option<i64>,
 }

codex-rs/ext/extension-api/src/contributors/world_state.rs

@@ -20,6 +20,8 @@ pub struct WorldStateContributionInput<'a> {
     pub session_store: &'a ExtensionData,
     pub thread_store: &'a ExtensionData,
     pub turn_store: &'a ExtensionData,
+    /// Host capabilities bound to this exact sampling step.
+    pub step_store: &'a ExtensionData,
 }
 
 /// What the harness knows about the previous value of one extension-owned section.

codex-rs/ext/extension-api/tests/registry.rs

@@ -41,6 +41,7 @@ impl ContextContributor for AllContributors {
         &'a self,
         _session_store: &'a ExtensionData,
         _thread_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> ExtensionFuture<'a, Vec<PromptFragment>> {
         Box::pin(std::future::ready(Vec::new()))
     }
@@ -63,6 +64,7 @@ impl TurnInputContributor for AllContributors {
         _session_store: &'a ExtensionData,
         _thread_store: &'a ExtensionData,
         _turn_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> ExtensionFuture<'a, Vec<Box<dyn ContextualUserFragment + Send>>> {
         Box::pin(async move {
             let _self = self;
@@ -77,6 +79,7 @@ impl ToolContributor for AllContributors {
         &self,
         _session_store: &ExtensionData,
         _thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
         Vec::new()
     }
@@ -158,6 +161,7 @@ impl ContextContributor for NamedContextContributor {
         &'a self,
         _session_store: &'a ExtensionData,
         _thread_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> ExtensionFuture<'a, Vec<PromptFragment>> {
         Box::pin(std::future::ready(vec![PromptFragment::developer_policy(
             self.0,
@@ -219,12 +223,13 @@ async fn contributors_preserve_registration_order() {
     let session_store = ExtensionData::new("session");
     let thread_store = ExtensionData::new("thread");
     let turn_store = ExtensionData::new("turn");
+    let step_store = ExtensionData::new("step");
 
     let mut fragments = Vec::new();
     for contributor in registry.context_contributors() {
         fragments.extend(
             contributor
-                .contribute_thread_context(&session_store, &thread_store)
+                .contribute_thread_context(&session_store, &thread_store, &step_store)
                 .await,
         );
     }
@@ -237,6 +242,7 @@ async fn contributors_preserve_registration_order() {
                     session_store: &session_store,
                     thread_store: &thread_store,
                     turn_store: &turn_store,
+                    step_store: &step_store,
                     model_context_window: Some(123),
                 })
                 .await,

codex-rs/ext/goal/src/extension.rs

@@ -420,6 +420,7 @@ where
         &self,
         _session_store: &ExtensionData,
         thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {
         let Some(runtime) = goal_runtime_handle(thread_store) else {
             return Vec::new();

codex-rs/ext/goal/tests/goal_extension_backend.rs

@@ -1143,7 +1143,9 @@ async fn installed_tools_with_start(
     registry
         .tool_contributors()
         .iter()
-        .flat_map(|contributor| contributor.tools(&session_store, &thread_store))
+        .flat_map(|contributor| {
+            contributor.tools(&session_store, &thread_store, &ExtensionData::new("step"))
+        })
         .collect()
 }
 
@@ -1205,7 +1207,13 @@ impl GoalExtensionHarness {
         self.registry
             .tool_contributors()
             .iter()
-            .flat_map(|contributor| contributor.tools(&self.session_store, &self.thread_store))
+            .flat_map(|contributor| {
+                contributor.tools(
+                    &self.session_store,
+                    &self.thread_store,
+                    &ExtensionData::new("step"),
+                )
+            })
             .collect()
     }
 

codex-rs/ext/image-generation/src/extension.rs

@@ -86,6 +86,7 @@ impl ToolContributor for ImageGenerationExtension {
         &self,
         _session_store: &ExtensionData,
         thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
         let Some(config) = thread_store.get::<ImageGenerationExtensionConfig>() else {
             return Vec::new();

codex-rs/ext/memories/src/extension.rs

@@ -52,6 +52,7 @@ impl ContextContributor for MemoriesExtension {
         &'a self,
         _session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
         Box::pin(async move {
             let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
@@ -100,6 +101,7 @@ impl ToolContributor for MemoriesExtension {
         &self,
         _session_store: &ExtensionData,
         thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {
         let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
             return Vec::new();

codex-rs/ext/memories/src/tests.rs

@@ -41,7 +41,8 @@ fn tools_are_not_contributed_without_thread_config() {
         extension
             .tools(
                 &ExtensionData::new("session"),
-                &ExtensionData::new("thread")
+                &ExtensionData::new("thread"),
+                &ExtensionData::new("step")
             )
             .is_empty()
     );
@@ -59,7 +60,11 @@ fn tools_are_not_contributed_when_disabled() {
 
     assert!(
         extension
-            .tools(&ExtensionData::new("session"), &thread_store)
+            .tools(
+                &ExtensionData::new("session"),
+                &thread_store,
+                &ExtensionData::new("step"),
+            )
             .is_empty()
     );
 }
@@ -76,7 +81,11 @@ fn tools_are_not_contributed_when_dedicated_tools_disabled() {
 
     assert!(
         extension
-            .tools(&ExtensionData::new("session"), &thread_store)
+            .tools(
+                &ExtensionData::new("session"),
+                &thread_store,
+                &ExtensionData::new("step"),
+            )
             .is_empty()
     );
 }
@@ -92,7 +101,11 @@ fn tools_are_contributed_when_enabled_with_dedicated_tools() {
     });
 
     let tool_names = extension
-        .tools(&ExtensionData::new("session"), &thread_store)
+        .tools(
+            &ExtensionData::new("session"),
+            &thread_store,
+            &ExtensionData::new("step"),
+        )
         .into_iter()
         .map(|tool| tool.tool_name())
         .collect::<Vec<_>>();
@@ -123,7 +136,13 @@ fn install_registers_dedicated_tool_contributor() {
     let tool_names = registry
         .tool_contributors()
         .iter()
-        .flat_map(|contributor| contributor.tools(&ExtensionData::new("session"), &thread_store))
+        .flat_map(|contributor| {
+            contributor.tools(
+                &ExtensionData::new("session"),
+                &thread_store,
+                &ExtensionData::new("step"),
+            )
+        })
         .map(|tool| tool.tool_name())
         .collect::<Vec<_>>();
 
@@ -181,7 +200,11 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() {
     });
 
     let fragments = extension
-        .contribute_thread_context(&ExtensionData::new("session"), &thread_store)
+        .contribute_thread_context(
+            &ExtensionData::new("session"),
+            &thread_store,
+            &ExtensionData::new("step"),
+        )
         .await;
 
     assert_eq!(fragments.len(), 1);

codex-rs/ext/skills/src/extension.rs

@@ -113,6 +113,7 @@ where
         &'a self,
         session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
         Box::pin(async move {
             let Some(thread_state) = thread_store.get::<SkillsThreadState>() else {
@@ -214,6 +215,7 @@ where
         &self,
         session_store: &ExtensionData,
         thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
         let Some(thread_state) = thread_store.get::<SkillsThreadState>() else {
             return Vec::new();
@@ -269,6 +271,7 @@ where
         session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
         turn_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> ExtensionFuture<'a, Vec<Box<dyn ContextualUserFragment + Send>>> {
         Box::pin(async move {
             let Some(thread_state) = thread_store.get::<SkillsThreadState>() else {

codex-rs/ext/skills/tests/skills_extension.rs

@@ -121,6 +121,7 @@ async fn installed_extension_uses_host_service_snapshot() -> TestResult {
             &session_store,
             &thread_store,
             &turn_store,
+            &ExtensionData::new("step"),
         )
         .await;
 
@@ -192,7 +193,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
         .await;
 
     let prompt_fragments = registry.context_contributors()[0]
-        .contribute_thread_context(&session_store, &thread_store)
+        .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step"))
         .await;
     assert!(prompt_fragments.is_empty());
 
@@ -212,6 +213,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
             session_store: &session_store,
             thread_store: &thread_store,
             turn_store: &turn_store,
+            step_store: &ExtensionData::new("step"),
         })
         .await;
     assert_eq!(1, available_sections.len());
@@ -239,6 +241,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
             &session_store,
             &thread_store,
             &turn_store,
+            &ExtensionData::new("step"),
         )
         .await;
 
@@ -265,6 +268,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
             session_store: &session_store,
             thread_store: &thread_store,
             turn_store: &unavailable_turn_store,
+            step_store: &ExtensionData::new("step"),
         })
         .await;
     let unavailable_snapshot = unavailable_sections[0].snapshot().clone();
@@ -288,6 +292,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
             session_store: &session_store,
             thread_store: &thread_store,
             turn_store: &restored_turn_store,
+            step_store: &ExtensionData::new("step"),
         })
         .await;
     let restored_snapshot = restored_sections[0].snapshot().clone();
@@ -316,6 +321,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
             session_store: &session_store,
             thread_store: &thread_store,
             turn_store: &listing_disabled_turn_store,
+            step_store: &ExtensionData::new("step"),
         })
         .await;
     let listing_disabled_fragment = listing_disabled_sections[0]
@@ -380,7 +386,7 @@ async fn default_context_truncates_catalog_descriptions() -> TestResult {
         .await;
 
     let fragments = registry.context_contributors()[0]
-        .contribute_thread_context(&session_store, &thread_store)
+        .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step"))
         .await;
     assert_eq!(1, fragments.len());
     let rendered = fragments[0].text();
@@ -429,7 +435,11 @@ async fn skills_list_truncates_catalog_descriptions_in_tool_output() -> TestResu
         })
         .await;
 
-    let tools = registry.tool_contributors()[0].tools(&session_store, &thread_store);
+    let tools = registry.tool_contributors()[0].tools(
+        &session_store,
+        &thread_store,
+        &ExtensionData::new("step"),
+    );
     let list_tool = tools
         .iter()
         .find(|tool| tool.tool_name().name == "list")
@@ -503,7 +513,7 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult {
         .await;
 
     let initial_fragments = registry.context_contributors()[0]
-        .contribute_thread_context(&session_store, &thread_store)
+        .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step"))
         .await;
     assert!(initial_fragments.is_empty());
     let EventMsg::Warning(warning) = event_rx.try_recv()?.msg else {
@@ -528,6 +538,7 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult {
                 &session_store,
                 &thread_store,
                 &ExtensionData::new(turn_id),
+                &ExtensionData::new("step"),
             )
             .await;
         assert!(fragments.is_empty());
@@ -607,6 +618,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te
             session_store: &session_store,
             thread_store: &thread_store,
             turn_store: &turn_store,
+            step_store: &ExtensionData::new("step"),
         })
         .await;
     let fragments = registry.turn_input_contributors()[0]
@@ -622,6 +634,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te
             &session_store,
             &thread_store,
             &turn_store,
+            &ExtensionData::new("step"),
         )
         .await;
 
@@ -697,6 +710,7 @@ async fn prompt_hidden_skill_can_still_be_invoked() -> TestResult {
             &session_store,
             &thread_store,
             &ExtensionData::new("turn-1"),
+            &ExtensionData::new("step"),
         )
         .await;
 

codex-rs/ext/web-search/src/extension.rs

@@ -120,6 +120,7 @@ impl ToolContributor for WebSearchExtension {
         &self,
         session_store: &ExtensionData,
         thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {
         let Some(config) = thread_store.get::<WebSearchExtensionConfig>() else {
             return Vec::new();
@@ -207,7 +208,9 @@ mod tests {
         let tool_names = registry
             .tool_contributors()
             .iter()
-            .flat_map(|contributor| contributor.tools(&session_store, &thread_store))
+            .flat_map(|contributor| {
+                contributor.tools(&session_store, &thread_store, &ExtensionData::new("step"))
+            })
             .map(|tool| (tool.tool_name(), tool.supports_parallel_tool_calls()))
             .collect::<Vec<_>>();