From fd51e505401b7b2da958edc269e4d7280be86bd5 Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 22 Jul 2026 11:09:55 +0000 Subject: [PATCH] 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 | 1 + codex-rs/core/src/compact.rs | 1 - codex-rs/core/src/compact_tests.rs | 1 - codex-rs/core/src/guardian/tests.rs | 1 - codex-rs/core/src/session/mod.rs | 16 ++------- codex-rs/core/src/session/session.rs | 3 +- codex-rs/core/src/session/step_context.rs | 5 --- codex-rs/core/src/session/tests.rs | 1 - codex-rs/core/src/session/turn.rs | 3 +- codex-rs/core/src/session/world_state.rs | 1 - codex-rs/core/src/tools/router.rs | 2 -- codex-rs/core/src/tools/router_tests.rs | 3 +- .../core/tests/suite/mcp_tool_exposure.rs | 6 ++-- codex-rs/ext/extension-api/Cargo.toml | 1 + .../examples/enabled_extensions.rs | 3 +- .../shared_state_extension.rs | 2 -- .../ext/extension-api/src/contributors.rs | 15 +++------ .../extension-api/src/contributors/context.rs | 2 -- .../src/contributors/thread_lifecycle.rs | 5 +++ .../src/contributors/world_state.rs | 2 -- codex-rs/ext/extension-api/tests/registry.rs | 8 +---- codex-rs/ext/goal/src/extension.rs | 1 - .../ext/goal/tests/goal_extension_backend.rs | 14 +++----- .../ext/image-generation/src/extension.rs | 1 - codex-rs/ext/memories/src/extension.rs | 2 -- codex-rs/ext/memories/src/tests.rs | 33 +++---------------- codex-rs/ext/skills/src/extension.rs | 30 ++++++++++++----- codex-rs/ext/skills/src/state.rs | 4 +++ codex-rs/ext/skills/tests/skills_extension.rs | 33 ++++++++----------- codex-rs/ext/web-search/src/extension.rs | 5 +-- 30 files changed, 70 insertions(+), 135 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index bd0a291eaa77..f9639fb53d90 100644 --- a/codex-rs/Cargo.lock +++ b/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", diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 4e777fb9ff58..633695dc6a17 100644 --- a/codex-rs/core/src/compact.rs +++ b/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))) diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 5608b123ef04..f34d71f96b86 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/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 { diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 6b8ea2524a8a..f376f3bf623d 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/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> { Box::pin(async move { if thread_store diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c77aa110956b..5fc6c628d9b2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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 { 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( @@ -3239,7 +3232,6 @@ impl Session { turn_context: &TurnContext, world_state: &WorldState, mcp: &McpRuntimeSnapshot, - step_store: &codex_extension_api::ExtensionData, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::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(); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index b6394f8ed0a9..910ade583bd9 100644 --- a/codex-rs/core/src/session/session.rs +++ b/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; diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 479824681935..78895950a444 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/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, /// The fixed MCP tool list used for this exact sampling request. mcp_tool_snapshot: OnceCell>, - /// 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>, } @@ -38,7 +35,6 @@ impl StepContext { mcp: Arc, loaded_agents_md: Option>, ) -> 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, } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index fe15ddac5bd8..8a3321ccd09e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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> + Send + 'a>, > { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index a696c1332230..106e3bdf72ff 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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, diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index df3c4d09d595..28ca49148607 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/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 { diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 43b90c59058a..32e139cae734 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/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>> { 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() diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 1dc8328d0939..09f1be8ae1f9 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/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>> { 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(), diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 7fae5eb6337a..b9788907f752 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -51,9 +51,9 @@ impl ThreadLifecycleContributor for McpResourceClientCapture { ) -> ExtensionFuture<'a, ()> { Box::pin(async move { let client = input - .session_store - .get::() - .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() diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index f465bef36fde..06944b471cf5 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/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 } diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs index 4c332fc6e028..45b178bc8131 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -73,11 +73,10 @@ async fn contribute_prompt( thread_store: &ExtensionData, ) -> Vec { 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, ); } diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs index 4dc7fe874edb..414a67215bb8 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs +++ b/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> + 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> + Send + 'a>> { Box::pin(async move { contribution_counts(session_store).record_usage(); diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 3f40fb67e566..a58605e952ab 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -78,18 +78,16 @@ pub trait McpServerContributor: 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> { 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: 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>>; } @@ -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>>; } diff --git a/codex-rs/ext/extension-api/src/contributors/context.rs b/codex-rs/ext/extension-api/src/contributors/context.rs index 95e33f4ad36c..fb6239ec4b43 100644 --- a/codex-rs/ext/extension-api/src/contributors/context.rs +++ b/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, } diff --git a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs index 23852a2df0b4..4fa6e6721c4f 100644 --- a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs +++ b/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>, /// Store scoped to the host session runtime. pub session_store: &'a ExtensionData, /// Store scoped to this thread runtime. diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index ae60de7ade50..75fae6d4a83d 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/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. diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index 0f01ad0f2878..de9cdf15393b 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/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> { 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::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>> { 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> { 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, diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index f9ae24fbc638..4fa1081db9e1 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -420,7 +420,6 @@ where &self, _session_store: &ExtensionData, thread_store: &ExtensionData, - _step_store: &ExtensionData, ) -> Vec>> { let Some(runtime) = goal_runtime_handle(thread_store) else { return Vec::new(); diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 37af7986d4a4..206216058067 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/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| { - contributor.tools(&session_store, &thread_store, &ExtensionData::new("step")) - }) + .flat_map(|contributor| contributor.tools(&session_store, &thread_store)) .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| { - contributor.tools( - &self.session_store, - &self.thread_store, - &ExtensionData::new("step"), - ) - }) + .flat_map(|contributor| contributor.tools(&self.session_store, &self.thread_store)) .collect() } diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index b4a598ddf333..38a24eed7c63 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/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>> { let Some(config) = thread_store.get::() else { return Vec::new(); diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs index e61a10df201f..464ce24576c2 100644 --- a/codex-rs/ext/memories/src/extension.rs +++ b/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> + Send + 'a>> { Box::pin(async move { let Some(config) = thread_store.get::() else { @@ -101,7 +100,6 @@ impl ToolContributor for MemoriesExtension { &self, _session_store: &ExtensionData, thread_store: &ExtensionData, - _step_store: &ExtensionData, ) -> Vec>> { let Some(config) = thread_store.get::() else { return Vec::new(); diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index e9ab1dade396..cd13c062e75b 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/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( - &ExtensionData::new("session"), - &thread_store, - &ExtensionData::new("step"), - ) + .tools(&ExtensionData::new("session"), &thread_store) .is_empty() ); } @@ -81,11 +76,7 @@ fn tools_are_not_contributed_when_dedicated_tools_disabled() { assert!( extension - .tools( - &ExtensionData::new("session"), - &thread_store, - &ExtensionData::new("step"), - ) + .tools(&ExtensionData::new("session"), &thread_store) .is_empty() ); } @@ -101,11 +92,7 @@ fn tools_are_contributed_when_enabled_with_dedicated_tools() { }); let tool_names = extension - .tools( - &ExtensionData::new("session"), - &thread_store, - &ExtensionData::new("step"), - ) + .tools(&ExtensionData::new("session"), &thread_store) .into_iter() .map(|tool| tool.tool_name()) .collect::>(); @@ -136,13 +123,7 @@ fn install_registers_dedicated_tool_contributor() { let tool_names = registry .tool_contributors() .iter() - .flat_map(|contributor| { - contributor.tools( - &ExtensionData::new("session"), - &thread_store, - &ExtensionData::new("step"), - ) - }) + .flat_map(|contributor| contributor.tools(&ExtensionData::new("session"), &thread_store)) .map(|tool| tool.tool_name()) .collect::>(); @@ -200,11 +181,7 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() { }); let fragments = extension - .contribute_thread_context( - &ExtensionData::new("session"), - &thread_store, - &ExtensionData::new("step"), - ) + .contribute_thread_context(&ExtensionData::new("session"), &thread_store) .await; assert_eq!(fragments.len(), 1); diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index 78743878f0b6..e70b85fa5784 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/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> + Send + 'a>> { Box::pin(async move { let Some(thread_state) = thread_store.get::() 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::(), + mcp_resources: session_store + .get::() + .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::(), + mcp_resources: input + .session_store + .get::() + .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>> { let Some(thread_state) = thread_store.get::() else { return Vec::new(); @@ -235,7 +241,9 @@ where skill_tools( self.providers.clone(), - session_store.get::(), + session_store + .get::() + .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::pin(async move { let Some(thread_state) = thread_store.get::() else { @@ -286,6 +293,9 @@ where }; let config = thread_state.config(); + let mcp_resources = session_store + .get::() + .and_then(|state| state.mcp_resources.clone()); let host_snapshot = turn_store.get::(); let host_catalog_in_world_state = turn_store.get::().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::(), + mcp_resources, executor_capability_discovery: None, }; let host_query = query.clone(); @@ -467,7 +477,9 @@ impl SkillsExtension { package: entry.id.clone(), resource: entry.main_prompt.clone(), host_snapshot, - mcp_resources: session_store.get::(), + mcp_resources: session_store + .get::() + .and_then(|state| state.mcp_resources.clone()), }, ) .await diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index 9f91f9ea2a65..5c026321d9ce 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/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>, +} + pub(crate) struct SkillsThreadState { config: Mutex, orchestrator_skills_available: bool, diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index b121b3b9b6d3..75e6036dfaf3 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/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( - &session_store, - &thread_store, - &ExtensionData::new("step"), - ); + let tools = registry.tool_contributors()[0].tools(&session_store, &thread_store); 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; diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index c57350666434..d59a34649dd0 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/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>> { let Some(config) = thread_store.get::() else { return Vec::new(); @@ -208,9 +207,7 @@ mod tests { let tool_names = registry .tool_contributors() .iter() - .flat_map(|contributor| { - contributor.tools(&session_store, &thread_store, &ExtensionData::new("step")) - }) + .flat_map(|contributor| contributor.tools(&session_store, &thread_store)) .map(|tool| (tool.tool_name(), tool.supports_parallel_tool_calls())) .collect::>();