openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
jif
2026-07-22 11:17:28 +00:00
committed by copyberry
parent 84d2b203ed
commit fd51e50540
30 changed files with 70 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -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
-1
View File
@@ -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
-1
View File
@@ -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
-1
View File
@@ -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
+2 -14
View File
@@ -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,
&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<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
+2 -1
View File
@@ -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
-5
View File
@@ -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
-1
View File
@@ -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
+1 -2
View File
@@ -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
-1
View File
@@ -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
-2
View File
@@ -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
+1 -2
View File
@@ -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
+3 -3
View File
@@ -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
+1
View File
@@ -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
+1 -2
View File
@@ -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
-2
View File
@@ -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
+4 -11
View File
@@ -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
-2
View File
@@ -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
+5
View File
@@ -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
-2
View File
@@ -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
+1 -7
View File
@@ -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
-1
View File
@@ -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
+4 -10
View File
@@ -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()
}
codex-rs/ext/image-generation/src/extension.rs
-1
View File
@@ -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
-2
View File
@@ -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
+5 -28
View File
@@ -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::<Vec<_>>();
@@ -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::<Vec<_>>();
@@ -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);
codex-rs/ext/skills/src/extension.rs
+21 -9
View File
@@ -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
+4
View File
@@ -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
+13 -20
View File
@@ -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;
codex-rs/ext/web-search/src/extension.rs
+1 -4
View File
@@ -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| {
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::<Vec<_>>();