Remove step-scoped data from extension contributors

codex-rs/core/src/compact.rs

@@ -97,7 +97,6 @@ pub(crate) async fn build_compaction_initial_context(
                     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)))

codex-rs/core/src/compact_tests.rs

@@ -24,7 +24,6 @@ async fn process_compacted_history_with_test_session(
             &turn_context,
             world_state.as_ref(),
             step_context.mcp.as_ref(),
-            &step_context.extension_data,
         )
         .await;
     let initial_context_injection = InitialContextInjection::BeforeLastUserMessage {

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

@@ -114,7 +114,6 @@ 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

@@ -3182,7 +3182,6 @@ 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
@@ -3224,22 +3223,15 @@ impl Session {
         world_state: &WorldState,
     ) -> Vec<ResponseItem> {
         let mcp = self.services.latest_mcp_runtime();
-        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
+        self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp)
+            .await
     }
 
     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);
@@ -3354,7 +3346,6 @@ impl Session {
                 .contribute_thread_context(
                     &self.services.session_extension_data,
                     &self.services.thread_extension_data,
-                    step_store,
                 )
                 .await
             {
@@ -3374,7 +3365,6 @@ 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
@@ -3549,7 +3539,6 @@ impl Session {
                 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();
@@ -3607,7 +3596,6 @@ impl Session {
                     turn_context,
                     world_state.as_ref(),
                     step_context.mcp.as_ref(),
-                    &step_context.extension_data,
                 )
                 .await;
             let snapshot = world_state.snapshot();

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

@@ -1041,13 +1041,14 @@ impl Session {
             )));
             let session_extension_data =
                 codex_extension_api::ExtensionData::new(session_id.to_string());
-            session_extension_data.insert(McpResourceClient::new(Arc::clone(&mcp_runtime)));
+            let mcp_resource_client = Arc::new(McpResourceClient::new(Arc::clone(&mcp_runtime)));
             for contributor in extensions.thread_lifecycle_contributors() {
                 contributor.on_thread_start(codex_extension_api::ThreadStartInput {
                     config: config.as_ref(),
                     session_source: &session_configuration.session_source,
                     persistent_thread_state_available: state_db_ctx.is_some(),
                     environments: session_configuration.environment_selections(),
+                    mcp_resource_client: Some(Arc::clone(&mcp_resource_client)),
                     session_store: &session_extension_data,
                     thread_store: &thread_extension_data,
                 }).await;

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

@@ -6,7 +6,6 @@ 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;
 
@@ -23,8 +22,6 @@ 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>>,
 }
@@ -38,15 +35,13 @@ impl StepContext {
         mcp: Arc<McpRuntimeSnapshot>,
         loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
     ) -> Self {
-        let extension_data = ExtensionData::new(turn.sub_id.clone());
         Self {
             turn,
             environments,
             selected_capability_roots,
             executor_capability_discovery,
             mcp,
             mcp_tool_snapshot: OnceCell::new(),
-            extension_data,
             loaded_agents_md,
         }
     }

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

@@ -8370,7 +8370,6 @@ 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/turn.rs

@@ -734,7 +734,6 @@ 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
@@ -1351,7 +1350,7 @@ pub(crate) async fn built_tools(
         ToolRouterParams {
             tool_runtimes: mcp_tool_runtimes,
             tool_suggest_candidates,
-            extension_tool_executors: extension_tool_executors(sess, step_context),
+            extension_tool_executors: extension_tool_executors(sess),
             dynamic_tools: turn_context.dynamic_tools.as_slice(),
         },
         &sess.services.tool_search_handler_cache,

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

@@ -124,7 +124,6 @@ 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,7 +246,6 @@ 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
@@ -257,7 +256,6 @@ 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

@@ -44,7 +44,6 @@ 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)]
     }
@@ -392,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, step_context.as_ref()),
+            extension_tool_executors: extension_tool_executors(&session),
             dynamic_tools: turn.dynamic_tools.as_slice(),
         },
         &Default::default(),

codex-rs/core/tests/suite/mcp_tool_exposure.rs

@@ -51,9 +51,9 @@ impl ThreadLifecycleContributor<Config> for McpResourceClientCapture {
     ) -> ExtensionFuture<'a, ()> {
         Box::pin(async move {
             let client = input
-                .session_store
-                .get::<McpResourceClient>()
-                .expect("session store should contain an MCP resource client");
+                .mcp_resource_client
+                .as_ref()
+                .expect("host should supply an MCP resource client");
             *self
                 .client
                 .lock()

codex-rs/ext/extension-api/Cargo.toml

@@ -17,6 +17,7 @@ workspace = true
 codex-config = { workspace = true }
 codex-context-fragments = { workspace = true }
 codex-exec-server-protocol = { workspace = true }
+codex-mcp = { workspace = true }
 codex-protocol = { workspace = true }
 codex-tools = { workspace = true }
 codex-utils-absolute-path = { workspace = true }

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

@@ -73,11 +73,10 @@ 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, &step_store)
+                .contribute_thread_context(session_store, thread_store)
                 .await,
         );
     }

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

@@ -21,7 +21,6 @@ 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();
@@ -42,7 +41,6 @@ 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,18 +78,16 @@ 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.
+    /// Returns thread-scoped context using the supplied extension state.
     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()
         })
     }
@@ -120,8 +118,8 @@ pub trait ContextContributor: Send + Sync {
 /// Contributor for host-owned thread lifecycle gates.
 ///
 /// Implementations should use these callbacks to seed, rehydrate, or flush
-/// extension-private thread state. Heavy dependencies belong on the extension
-/// value created by the host, not in these inputs.
+/// extension-private thread state and retain any session capabilities supplied
+/// by the host. Other heavy dependencies belong on the extension value.
 pub trait ThreadLifecycleContributor<C: Sync>: Send + Sync {
     /// Called after host startup has initialized the thread-scoped store.
     fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> {
@@ -208,16 +206,12 @@ 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>>>;
 }
 
@@ -277,12 +271,11 @@ pub trait SkillInvocationContributor: Send + Sync {
 
 /// Extension contribution that exposes native tools owned by a feature.
 pub trait ToolContributor: Send + Sync {
-    /// Returns native tools bound to the supplied sampling-step capabilities.
+    /// Returns native tools bound to the supplied extension state.
     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,8 +15,6 @@ 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/thread_lifecycle.rs

@@ -1,4 +1,7 @@
+use std::sync::Arc;
+
 use crate::ExtensionData;
+use codex_mcp::McpResourceClient;
 use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::TurnEnvironmentSelection;
 
@@ -20,6 +23,8 @@ pub struct ThreadStartInput<'a, C> {
     pub persistent_thread_state_available: bool,
     /// Execution environments selected for this thread.
     pub environments: &'a [TurnEnvironmentSelection],
+    /// MCP resource access supplied by the host for this session.
+    pub mcp_resource_client: Option<Arc<McpResourceClient>>,
     /// Store scoped to the host session runtime.
     pub session_store: &'a ExtensionData,
     /// Store scoped to this thread runtime.

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

@@ -20,8 +20,6 @@ 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.