Remove step-scoped data from extension contributors (#34734)

## What changed

- Remove the step-scoped `ExtensionData` argument from context, turn-input, and tool contributors.
- Pass the host's optional `McpResourceClient` through `ThreadStartInput` so extensions can retain session capabilities explicitly.
- Keep the MCP resource client in skills-owned session state for catalog loading and skill tools.

GitOrigin-RevId: bafa77bcd998aff408d6a396c5fd9ac268c4cce4

codex-rs/Cargo.lock

@@ -2998,6 +2998,7 @@ dependencies = [
  "codex-config",
  "codex-context-fragments",
  "codex-exec-server-protocol",
+ "codex-mcp",
  "codex-protocol",
  "codex-tools",
  "codex-utils-absolute-path",

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,14 +3223,8 @@ 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)
-        self.build_initial_context_with_world_state_and_mcp(
+            .await
-            turn_context,
-            world_state,
-            &mcp,
-            &step_store,
-        )
-        .await
     }
 
     pub(crate) async fn build_initial_context_with_world_state_and_mcp(
@@ -3239,7 +3232,6 @@ impl Session {
         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,7 +35,6 @@ impl StepContext {
         mcp: Arc<McpRuntimeSnapshot>,
         loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
     ) -> Self {
-        let extension_data = ExtensionData::new(turn.sub_id.clone());
         Self {
             turn,
             environments,
@@ -46,7 +42,6 @@ impl StepContext {
             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
+                .mcp_resource_client
-                .get::<McpResourceClient>()
+                .as_ref()
-                .expect("session store should contain an MCP resource client");
+                .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
+/// extension-private thread state and retain any session capabilities supplied
-/// value created by the host, not in these inputs.
+/// 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.

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

@@ -41,7 +41,6 @@ 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()))
     }
@@ -64,7 +63,6 @@ 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;
@@ -79,7 +77,6 @@ impl ToolContributor for AllContributors {
         &self,
         _session_store: &ExtensionData,
         _thread_store: &ExtensionData,
-        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
         Vec::new()
     }
@@ -161,7 +158,6 @@ 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,
@@ -223,13 +219,12 @@ 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, &step_store)
+                .contribute_thread_context(&session_store, &thread_store)
                 .await,
         );
     }
@@ -242,7 +237,6 @@ 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,7 +420,6 @@ 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

@@ -1134,6 +1134,7 @@ async fn installed_tools_with_start(
                 session_source: &session_source,
                 persistent_thread_state_available,
                 environments: &[],
+                mcp_resource_client: None,
                 session_store: &session_store,
                 thread_store: &thread_store,
             })
@@ -1143,9 +1144,7 @@ async fn installed_tools_with_start(
     registry
         .tool_contributors()
         .iter()
-        .flat_map(|contributor| {
+        .flat_map(|contributor| contributor.tools(&session_store, &thread_store))
-            contributor.tools(&session_store, &thread_store, &ExtensionData::new("step"))
-        })
         .collect()
 }
 
@@ -1189,6 +1188,7 @@ impl GoalExtensionHarness {
                     session_source: &session_source,
                     persistent_thread_state_available: true,
                     environments: &[],
+                    mcp_resource_client: None,
                     session_store: &session_store,
                     thread_store: &thread_store,
                 })
@@ -1207,13 +1207,7 @@ impl GoalExtensionHarness {
         self.registry
             .tool_contributors()
             .iter()
-            .flat_map(|contributor| {
+            .flat_map(|contributor| contributor.tools(&self.session_store, &self.thread_store))
-                contributor.tools(
-                    &self.session_store,
-                    &self.thread_store,
-                    &ExtensionData::new("step"),
-                )
-            })
             .collect()
     }
 

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

@@ -86,7 +86,6 @@ 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,7 +52,6 @@ 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 {
@@ -101,7 +100,6 @@ 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

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

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

@@ -24,7 +24,6 @@ use codex_extension_api::TurnInputContext;
 use codex_extension_api::TurnInputContributor;
 use codex_extension_api::WorldStateContributionInput;
 use codex_extension_api::WorldStateSectionContribution;
-use codex_mcp::McpResourceClient;
 use codex_otel::MetricsClient;
 use codex_protocol::openai_models::ModelInfo;
 use codex_protocol::protocol::Event;
@@ -51,6 +50,7 @@ use crate::selection::collect_explicit_skill_mentions;
 use crate::shadow_selection_experiment::ShadowSelectionExperiment;
 use crate::sources::SkillProviders;
 use crate::state::ExecutorSkillsStepState;
+use crate::state::SkillsSessionState;
 use crate::state::SkillsThreadState;
 use crate::state::SkillsTurnState;
 use crate::tools::skill_tools;
@@ -70,6 +70,9 @@ where
 {
     fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> {
         Box::pin(async move {
+            input.session_store.insert(SkillsSessionState {
+                mcp_resources: input.mcp_resource_client.clone(),
+            });
             let orchestrator_skills_available = !input
                 .environments
                 .iter()
@@ -114,7 +117,6 @@ 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 {
@@ -133,7 +135,9 @@ where
                         include_host_skills: false,
                         include_bundled_skills: config.bundled_skills_enabled,
                         include_orchestrator_skills: thread_state.orchestrator_skills_enabled(),
-                        mcp_resources: session_store.get::<McpResourceClient>(),
+                        mcp_resources: session_store
+                            .get::<SkillsSessionState>()
+                            .and_then(|state| state.mcp_resources.clone()),
                         executor_capability_discovery: None,
                     },
                     &thread_state,
@@ -176,7 +180,10 @@ where
                         include_host_skills: false,
                         include_bundled_skills: config.bundled_skills_enabled,
                         include_orchestrator_skills: false,
-                        mcp_resources: input.session_store.get::<McpResourceClient>(),
+                        mcp_resources: input
+                            .session_store
+                            .get::<SkillsSessionState>()
+                            .and_then(|state| state.mcp_resources.clone()),
                         executor_capability_discovery: input.executor_capability_discovery.cloned(),
                     },
                 )
@@ -222,7 +229,6 @@ 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();
@@ -235,7 +241,9 @@ where
 
         skill_tools(
             self.providers.clone(),
-            session_store.get::<McpResourceClient>(),
+            session_store
+                .get::<SkillsSessionState>()
+                .and_then(|state| state.mcp_resources.clone()),
             thread_state,
             Arc::clone(&self.shadow_selection),
         )
@@ -278,7 +286,6 @@ 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 {
@@ -286,6 +293,9 @@ where
             };
 
             let config = thread_state.config();
+            let mcp_resources = session_store
+                .get::<SkillsSessionState>()
+                .and_then(|state| state.mcp_resources.clone());
             let host_snapshot = turn_store.get::<HostSkillsSnapshot>();
             let host_catalog_in_world_state =
                 turn_store.get::<HostSkillsCatalogInWorldState>().is_some();
@@ -296,7 +306,7 @@ where
                 include_host_skills: !host_catalog_in_world_state,
                 include_bundled_skills: config.bundled_skills_enabled,
                 include_orchestrator_skills: thread_state.orchestrator_skills_enabled(),
-                mcp_resources: session_store.get::<McpResourceClient>(),
+                mcp_resources,
                 executor_capability_discovery: None,
             };
             let host_query = query.clone();
@@ -467,7 +477,9 @@ impl<C> SkillsExtension<C> {
                     package: entry.id.clone(),
                     resource: entry.main_prompt.clone(),
                     host_snapshot,
-                    mcp_resources: session_store.get::<McpResourceClient>(),
+                    mcp_resources: session_store
+                        .get::<SkillsSessionState>()
+                        .and_then(|state| state.mcp_resources.clone()),
                 },
             )
             .await

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

@@ -26,6 +26,10 @@ use crate::sources::SkillProviders;
 const MAX_CACHED_ORCHESTRATOR_RESOURCES: usize = 100;
 const MAX_CACHED_ORCHESTRATOR_CONTENT_BYTES: usize = 8 * 1024 * 1024;
 
+pub(crate) struct SkillsSessionState {
+    pub(crate) mcp_resources: Option<Arc<McpResourceClient>>,
+}
+
 pub(crate) struct SkillsThreadState {
     config: Mutex<SkillsExtensionConfig>,
     orchestrator_skills_available: bool,

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

@@ -85,6 +85,7 @@ async fn installed_extension_uses_host_service_snapshot() -> TestResult {
             session_source: &session_source,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
@@ -122,7 +123,6 @@ async fn installed_extension_uses_host_service_snapshot() -> TestResult {
             &session_store,
             &thread_store,
             &turn_store,
-            &ExtensionData::new("step"),
         )
         .await;
 
@@ -188,13 +188,14 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
             session_source: &session_source,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
         .await;
 
     let prompt_fragments = registry.context_contributors()[0]
-        .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step"))
+        .contribute_thread_context(&session_store, &thread_store)
         .await;
     assert!(prompt_fragments.is_empty());
 
@@ -214,7 +215,6 @@ 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());
@@ -242,7 +242,6 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
             &session_store,
             &thread_store,
             &turn_store,
-            &ExtensionData::new("step"),
         )
         .await;
 
@@ -269,7 +268,6 @@ 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();
@@ -293,7 +291,6 @@ 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();
@@ -322,7 +319,6 @@ 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]
@@ -381,13 +377,14 @@ async fn default_context_truncates_catalog_descriptions() -> TestResult {
             session_source: &session_source,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
         .await;
 
     let fragments = registry.context_contributors()[0]
-        .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step"))
+        .contribute_thread_context(&session_store, &thread_store)
         .await;
     assert_eq!(1, fragments.len());
     let rendered = fragments[0].text();
@@ -513,16 +510,13 @@ async fn skills_list_truncates_catalog_descriptions_in_tool_output() -> TestResu
             session_source: &session_source,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
         .await;
 
-    let tools = registry.tool_contributors()[0].tools(
+    let tools = registry.tool_contributors()[0].tools(&session_store, &thread_store);
-        &session_store,
-        &thread_store,
-        &ExtensionData::new("step"),
-    );
     let list_tool = tools
         .iter()
         .find(|tool| tool.tool_name().name == "list")
@@ -590,13 +584,14 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult {
             session_source: &session_source,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
         .await;
 
     let initial_fragments = registry.context_contributors()[0]
-        .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step"))
+        .contribute_thread_context(&session_store, &thread_store)
         .await;
     assert!(initial_fragments.is_empty());
     let EventMsg::Warning(warning) = event_rx.try_recv()?.msg else {
@@ -621,7 +616,6 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult {
                 &session_store,
                 &thread_store,
                 &ExtensionData::new(turn_id),
-                &ExtensionData::new("step"),
             )
             .await;
         assert!(fragments.is_empty());
@@ -681,6 +675,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te
             session_source: &session_source,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
@@ -701,7 +696,6 @@ 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]
@@ -717,7 +711,6 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te
             &session_store,
             &thread_store,
             &turn_store,
-            &ExtensionData::new("step"),
         )
         .await;
 
@@ -789,6 +782,7 @@ async fn model_context_window_scales_executor_catalog_but_not_thread_catalog() -
             session_source: &SessionSource::Cli,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
@@ -798,7 +792,7 @@ async fn model_context_window_scales_executor_catalog_but_not_thread_catalog() -
     thread_store.insert(model_info);
 
     let thread_fragments = registry.context_contributors()[0]
-        .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step"))
+        .contribute_thread_context(&session_store, &thread_store)
         .await;
     assert_eq!(1, thread_fragments.len());
     assert!(thread_fragments[0].text().contains("skill-39"));
@@ -826,7 +820,6 @@ async fn model_context_window_scales_executor_catalog_but_not_thread_catalog() -
             session_store: &session_store,
             thread_store: &thread_store,
             turn_store: &turn_store,
-            step_store: &ExtensionData::new("step"),
         })
         .await;
     let fragment = sections[0]
@@ -878,6 +871,7 @@ async fn prompt_hidden_skill_can_still_be_invoked() -> TestResult {
             session_source: &session_source,
             persistent_thread_state_available: true,
             environments: &[],
+            mcp_resource_client: None,
             session_store: &session_store,
             thread_store: &thread_store,
         })
@@ -896,7 +890,6 @@ 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,7 +120,6 @@ 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();
@@ -208,9 +207,7 @@ mod tests {
         let tool_names = registry
             .tool_contributors()
             .iter()
-            .flat_map(|contributor| {
+            .flat_map(|contributor| contributor.tools(&session_store, &thread_store))
-                contributor.tools(&session_store, &thread_store, &ExtensionData::new("step"))
-            })
             .map(|tool| (tool.tool_name(), tool.supports_parallel_tool_calls()))
             .collect::<Vec<_>>();