From c44c4de7b410993dacb2e88c7084c9c968bc963a Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 21 Jul 2026 11:47:24 +0000 Subject: [PATCH] Add step-scoped data to extension contributors (#34525) ## What changed - Add an `ExtensionData` store to each `StepContext`. - Pass the step store to context, world-state, turn-input, and tool contributors so extensions can use capabilities bound to the current sampling step. - Preserve the captured step context when rebuilding initial context during compaction. GitOrigin-RevId: 51f3105d3a793f60534f122f9be8247c131c2ca5 --- codex-rs/core/src/compact.rs | 29 +++++++++------ codex-rs/core/src/compact_remote.rs | 14 +++----- codex-rs/core/src/compact_remote_v2.rs | 12 +++---- codex-rs/core/src/compact_tests.rs | 18 +++++++--- codex-rs/core/src/compact_token_budget.rs | 12 +++---- codex-rs/core/src/guardian/tests.rs | 1 + codex-rs/core/src/session/mod.rs | 33 +++++++++++++---- codex-rs/core/src/session/step_context.rs | 5 +++ codex-rs/core/src/session/tests.rs | 9 +++-- .../core/src/session/tests/guardian_tests.rs | 9 +++-- codex-rs/core/src/session/turn.rs | 13 ++++--- 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 +- .../examples/enabled_extensions.rs | 3 +- .../shared_state_extension.rs | 2 ++ .../ext/extension-api/src/contributors.rs | 10 +++++- .../extension-api/src/contributors/context.rs | 2 ++ .../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 | 12 +++++-- .../ext/image-generation/src/extension.rs | 1 + codex-rs/ext/memories/src/extension.rs | 2 ++ codex-rs/ext/memories/src/tests.rs | 35 +++++++++++++++---- codex-rs/ext/skills/src/extension.rs | 3 ++ codex-rs/ext/skills/tests/skills_extension.rs | 22 +++++++++--- codex-rs/ext/web-search/src/extension.rs | 5 ++- 28 files changed, 195 insertions(+), 74 deletions(-) diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index a082c7d20a..23ee7a722a 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -15,6 +15,7 @@ use crate::responses_metadata::CompactionTurnMetadata; #[cfg(test)] use crate::session::PreviousTurnSettings; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::get_last_assistant_message_from_turn; use crate::session::turn_context::TurnContext; use crate::util::backoff; @@ -64,20 +65,30 @@ const COMPACT_USER_MESSAGE_MAX_TOKENS: usize = 20_000; /// initial context into the replacement history just above the last real user message. #[derive(Debug)] pub(crate) enum InitialContextInjection { - BeforeLastUserMessage(Arc), + BeforeLastUserMessage { + world_state: Arc, + step_context: Arc, + }, DoNotInject, } pub(crate) async fn build_compaction_initial_context( sess: &Session, - turn_context: &TurnContext, initial_context_injection: &InitialContextInjection, ) -> (Vec, Option>) { // Return the rendered state with its items so history and its baseline stay identical. match initial_context_injection { - InitialContextInjection::BeforeLastUserMessage(world_state) => { + InitialContextInjection::BeforeLastUserMessage { + world_state, + step_context, + } => { let items = sess - .build_initial_context_with_world_state(turn_context, world_state.as_ref()) + .build_initial_context_with_world_state_and_mcp( + step_context.turn.as_ref(), + world_state.as_ref(), + step_context.mcp.as_ref(), + &step_context.extension_data, + ) .await; (items, Some(Arc::clone(world_state))) } @@ -334,19 +345,15 @@ async fn run_compact_task_inner_impl( } let (window_number, window_ids) = sess.advance_auto_compact_window().await; - let (initial_context, world_state_baseline) = build_compaction_initial_context( - sess.as_ref(), - turn_context.as_ref(), - &initial_context_injection, - ) - .await; + let (initial_context, world_state_baseline) = + build_compaction_initial_context(sess.as_ref(), &initial_context_injection).await; if !initial_context.is_empty() { new_history = insert_initial_context_before_last_real_user_or_summary(new_history, initial_context); } let reference_context_item = match initial_context_injection { InitialContextInjection::DoNotInject => None, - InitialContextInjection::BeforeLastUserMessage(_) => { + InitialContextInjection::BeforeLastUserMessage { .. } => { Some(turn_context.to_turn_context_item()) } }; diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 16718f3719..07fc669c97 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -260,17 +260,12 @@ async fn run_remote_compact_task_inner_impl( trace_input_history, } = attempt; let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; - let (new_history, world_state_baseline) = process_compacted_history( - sess.as_ref(), - compaction_turn_context.as_ref(), - new_history, - &initial_context_injection, - ) - .await; + let (new_history, world_state_baseline) = + process_compacted_history(sess.as_ref(), new_history, &initial_context_injection).await; let reference_context_item = match initial_context_injection { InitialContextInjection::DoNotInject => None, - InitialContextInjection::BeforeLastUserMessage(_) => { + InitialContextInjection::BeforeLastUserMessage { .. } => { Some(compaction_turn_context.to_turn_context_item()) } }; @@ -308,7 +303,6 @@ async fn run_remote_compact_task_inner_impl( pub(crate) async fn process_compacted_history( sess: &Session, - turn_context: &TurnContext, mut compacted_history: Vec, initial_context_injection: &InitialContextInjection, ) -> (Vec, Option>) { @@ -316,7 +310,7 @@ pub(crate) async fn process_compacted_history( // message in the replacement history. Pre-turn compaction instead injects context after the // compaction item, but mid-turn compaction keeps the compaction item last for model training. let (initial_context, world_state_baseline) = - build_compaction_initial_context(sess, turn_context, initial_context_injection).await; + build_compaction_initial_context(sess, initial_context_injection).await; compacted_history.retain(should_keep_compacted_history_item); ( diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 9a39f481df..b6340732e2 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -283,17 +283,13 @@ async fn run_remote_compact_task_inner_impl( build_v2_compacted_history(&prompt_input, compaction_output); analytics_details.retained_image_count = Some(retained_images); let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; - let (new_history, world_state_baseline) = process_compacted_history( - sess.as_ref(), - compaction_turn_context.as_ref(), - compacted_history, - &initial_context_injection, - ) - .await; + let (new_history, world_state_baseline) = + process_compacted_history(sess.as_ref(), compacted_history, &initial_context_injection) + .await; let reference_context_item = match initial_context_injection { InitialContextInjection::DoNotInject => None, - InitialContextInjection::BeforeLastUserMessage(_) => { + InitialContextInjection::BeforeLastUserMessage { .. } => { Some(compaction_turn_context.to_turn_context_item()) } }; diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 35cfd2574f..5608b123ef 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -1,5 +1,4 @@ use super::*; -use crate::session::tests::build_world_state_from_turn_context; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_protocol::ResponseItemId; @@ -17,14 +16,23 @@ async fn process_compacted_history_with_test_session( session .set_previous_turn_settings(previous_turn_settings.cloned()) .await; - let world_state = Arc::new(build_world_state_from_turn_context(&session, &turn_context).await); + let step_context = + crate::session::step_context::StepContext::for_test(Arc::clone(&turn_context)); + let world_state = Arc::new(session.build_world_state_for_step(&step_context).await); let initial_context = session - .build_initial_context_with_world_state(&turn_context, world_state.as_ref()) + .build_initial_context_with_world_state_and_mcp( + &turn_context, + world_state.as_ref(), + step_context.mcp.as_ref(), + &step_context.extension_data, + ) .await; - let initial_context_injection = InitialContextInjection::BeforeLastUserMessage(world_state); + let initial_context_injection = InitialContextInjection::BeforeLastUserMessage { + world_state, + step_context, + }; let (refreshed, _) = crate::compact_remote::process_compacted_history( &session, - &turn_context, compacted_history, &initial_context_injection, ) diff --git a/codex-rs/core/src/compact_token_budget.rs b/codex-rs/core/src/compact_token_budget.rs index 0d71e7ab2f..f8b5763c95 100644 --- a/codex-rs/core/src/compact_token_budget.rs +++ b/codex-rs/core/src/compact_token_budget.rs @@ -38,7 +38,7 @@ pub(crate) async fn run_manual_compact_task( // Manual compaction runs outside run_turn, so it captures its own current step. let step_context = sess.capture_step_context(Arc::clone(&turn_context)).await; let world_state = Arc::new(sess.build_world_state_for_step(&step_context).await); - run_compact_task_inner(&sess, &turn_context, world_state, CompactionTrigger::Manual).await + run_compact_task_inner(&sess, &step_context, world_state, CompactionTrigger::Manual).await } /// Runs token-budget inline auto-compaction as a normal compaction lifecycle. @@ -51,22 +51,22 @@ pub(crate) async fn run_inline_auto_compact_task( step_context: Arc, initial_context_injection: InitialContextInjection, ) -> CodexResult<()> { - let turn_context = &step_context.turn; let world_state = match initial_context_injection { - InitialContextInjection::BeforeLastUserMessage(world_state) => world_state, + InitialContextInjection::BeforeLastUserMessage { world_state, .. } => world_state, InitialContextInjection::DoNotInject => { Arc::new(sess.build_world_state_for_step(&step_context).await) } }; - run_compact_task_inner(&sess, turn_context, world_state, CompactionTrigger::Auto).await + run_compact_task_inner(&sess, &step_context, world_state, CompactionTrigger::Auto).await } async fn run_compact_task_inner( sess: &Arc, - turn_context: &Arc, + step_context: &Arc, world_state: Arc, trigger: CompactionTrigger, ) -> CodexResult<()> { + let turn_context = &step_context.turn; let pre_compact_outcome = run_pre_compact_hooks(sess, turn_context, trigger).await; match pre_compact_outcome { PreCompactHookOutcome::Continue => {} @@ -76,7 +76,7 @@ async fn run_compact_task_inner( let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new()); sess.emit_turn_item_started(turn_context, &compaction_item) .await; - sess.start_new_context_window(turn_context.as_ref(), world_state) + sess.start_new_context_window(step_context, world_state) .await; sess.emit_turn_item_completed(turn_context, compaction_item) .await; diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index c622d68541..117dd512c9 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -114,6 +114,7 @@ impl codex_extension_api::ContextContributor for GuardianMemoryContextProbe { &'a self, _session_store: &'a codex_extension_api::ExtensionData, thread_store: &'a codex_extension_api::ExtensionData, + _step_store: &'a codex_extension_api::ExtensionData, ) -> codex_extension_api::ExtensionFuture<'a, Vec> { 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 bc686eb7a7..cca0385e01 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3152,8 +3152,9 @@ impl Session { async fn build_turn_context_contribution_items( &self, - turn_context: &TurnContext, + step_context: &StepContext, ) -> Vec { + let turn_context = step_context.turn.as_ref(); let mut developer_sections = Vec::new(); let mut contextual_user_sections = Vec::new(); let mut separate_developer_sections = Vec::new(); @@ -3167,6 +3168,7 @@ impl Session { session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), + step_store: &step_context.extension_data, model_context_window: turn_context.model_context_window(), }) .await @@ -3201,21 +3203,29 @@ impl Session { items } + #[cfg(test)] pub(crate) async fn build_initial_context_with_world_state( &self, turn_context: &TurnContext, world_state: &WorldState, ) -> Vec { let mcp = self.services.latest_mcp_runtime(); - self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp) - .await + 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 } - async fn build_initial_context_with_world_state_and_mcp( + pub(crate) async fn build_initial_context_with_world_state_and_mcp( &self, turn_context: &TurnContext, world_state: &WorldState, mcp: &McpRuntimeSnapshot, + step_store: &codex_extension_api::ExtensionData, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); @@ -3330,6 +3340,7 @@ impl Session { .contribute_thread_context( &self.services.session_extension_data, &self.services.thread_extension_data, + step_store, ) .await { @@ -3349,6 +3360,7 @@ impl Session { session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), + step_store, model_context_window: turn_context.model_context_window(), }) .await @@ -3509,16 +3521,22 @@ impl Session { pub(crate) async fn start_new_context_window( &self, - turn_context: &TurnContext, + step_context: &StepContext, world_state: Arc, ) -> u64 { + let turn_context = step_context.turn.as_ref(); let window = { let mut state = self.state.lock().await; state.start_new_context_window() }; let (window_number, window_ids) = window; let context_items = self - .build_initial_context_with_world_state(turn_context, world_state.as_ref()) + .build_initial_context_with_world_state_and_mcp( + turn_context, + world_state.as_ref(), + step_context.mcp.as_ref(), + &step_context.extension_data, + ) .await; let turn_context_item = turn_context.to_turn_context_item(); self.replace_compacted_history( @@ -3579,6 +3597,7 @@ impl Session { turn_context, world_state.as_ref(), step_context.mcp.as_ref(), + &step_context.extension_data, ) .await; let snapshot = world_state.snapshot(); @@ -3611,7 +3630,7 @@ impl Session { }; if !should_inject_full_context && turn_context_changed { context_items.extend( - self.build_turn_context_contribution_items(turn_context) + self.build_turn_context_contribution_items(step_context) .await, ); } diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 78895950a4..4798246819 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -6,6 +6,7 @@ use crate::session::McpRuntimeSnapshot; use crate::session::turn_context::TurnContext; use codex_exec_server::ExecutorCapabilityDiscoverySnapshot; use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_extension_api::ExtensionData; use codex_mcp::ToolInfo; use tokio::sync::OnceCell; @@ -22,6 +23,8 @@ pub(crate) struct StepContext { pub(crate) mcp: Arc, /// 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>, } @@ -35,6 +38,7 @@ impl StepContext { mcp: Arc, loaded_agents_md: Option>, ) -> Self { + let extension_data = ExtensionData::new(turn.sub_id.clone()); Self { turn, environments, @@ -42,6 +46,7 @@ impl StepContext { executor_capability_discovery, mcp, mcp_tool_snapshot: OnceCell::new(), + extension_data, loaded_agents_md, } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index aab520684e..dab8637501 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2815,11 +2815,13 @@ async fn start_new_context_window_assigns_and_persists_item_ids() { .await; let rollout_path = attach_thread_persistence(Arc::get_mut(&mut session).expect("unique session")).await; - let world_state = - Arc::new(build_world_state_from_turn_context(session.as_ref(), &turn_context).await); + let step_context = session + .capture_step_context(Arc::clone(&turn_context)) + .await; + let world_state = Arc::new(session.build_world_state_for_step(&step_context).await); session - .start_new_context_window(turn_context.as_ref(), world_state) + .start_new_context_window(&step_context, world_state) .await; let live_history = session.clone_history().await; @@ -8274,6 +8276,7 @@ impl codex_extension_api::ContextContributor for PromptExtensionTestContributor &'a self, _session_store: &'a codex_extension_api::ExtensionData, thread_store: &'a codex_extension_api::ExtensionData, + _step_store: &'a codex_extension_api::ExtensionData, ) -> std::pin::Pin< Box> + Send + 'a>, > { diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index d2cee609f1..d6c103a951 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -539,12 +539,15 @@ async fn process_compacted_history_preserves_separate_guardian_developer_message turn_context.session_source = guardian_source; turn_context.developer_instructions = Some(guardian_policy.clone()); let turn_context = Arc::new(turn_context); - let world_state = Arc::new(build_world_state_from_turn_context(&session, &turn_context).await); - let initial_context_injection = InitialContextInjection::BeforeLastUserMessage(world_state); + let step_context = StepContext::for_test(Arc::clone(&turn_context)); + let world_state = Arc::new(session.build_world_state_for_step(&step_context).await); + let initial_context_injection = InitialContextInjection::BeforeLastUserMessage { + world_state, + step_context, + }; let (refreshed, _) = crate::compact_remote::process_compacted_history( &session, - &turn_context, vec![ ResponseItem::Message { id: None, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 32751315b6..58706d52f5 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -357,7 +357,10 @@ pub(crate) async fn run_turn( Arc::clone(&step_context), /*fallback_step_context*/ None, &mut client_session, - InitialContextInjection::BeforeLastUserMessage(Arc::clone(&world_state)), + InitialContextInjection::BeforeLastUserMessage { + world_state: Arc::clone(&world_state), + step_context: Arc::clone(&step_context), + }, CompactionReason::ContextLimit, CompactionPhase::MidTurn, ) @@ -584,7 +587,7 @@ async fn build_skills_and_plugins( let skills_outcome = turn_context.turn_skills.snapshot.outcome(); let connector_slug_counts = build_connector_slug_counts(&available_connectors); let extension_injection_items = - build_extension_turn_input_items(sess, turn_context, &user_input, cancellation_token) + build_extension_turn_input_items(sess, step_context, &user_input, cancellation_token) .await?; let skill_name_counts_lower = build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1; @@ -687,10 +690,11 @@ async fn build_skills_and_plugins( )] async fn build_extension_turn_input_items( sess: &Arc, - turn_context: &TurnContext, + step_context: &StepContext, user_input: &[UserInput], cancellation_token: &CancellationToken, ) -> Option> { + let turn_context = step_context.turn.as_ref(); let contributors = sess.services.extensions.turn_input_contributors().to_vec(); if contributors.is_empty() { return Some(Vec::new()); @@ -725,6 +729,7 @@ async fn build_extension_turn_input_items( &sess.services.session_extension_data, &sess.services.thread_extension_data, turn_context.extension_data.as_ref(), + &step_context.extension_data, ) .or_cancel(cancellation_token) .await @@ -1341,7 +1346,7 @@ pub(crate) async fn built_tools( ToolRouterParams { tool_runtimes: mcp_tool_runtimes, tool_suggest_candidates, - extension_tool_executors: extension_tool_executors(sess), + extension_tool_executors: extension_tool_executors(sess, step_context), dynamic_tools: turn_context.dynamic_tools.as_slice(), }, &sess.services.tool_search_handler_cache, diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 28ca491486..df3c4d09d5 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -124,6 +124,7 @@ impl Session { session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), + step_store: &step_context.extension_data, }) .await { diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 32e139cae7..43b90c5905 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -246,6 +246,7 @@ impl ToolRouter { #[instrument(level = "trace", skip_all)] pub(crate) fn extension_tool_executors( session: &Session, + step_context: &StepContext, ) -> Vec>> { session .services @@ -256,6 +257,7 @@ pub(crate) fn extension_tool_executors( contributor.tools( &session.services.session_extension_data, &session.services.thread_extension_data, + &step_context.extension_data, ) }) .collect() diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index b918035ccc..52a0208e72 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -43,6 +43,7 @@ impl codex_extension_api::ToolContributor for ExtensionEchoContributor { &self, _session_store: &ExtensionData, _thread_store: &ExtensionData, + _step_store: &ExtensionData, ) -> Vec>> { vec![Arc::new(ExtensionEchoExecutor)] } @@ -390,7 +391,7 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow ToolRouterParams { tool_suggest_candidates: None, tool_runtimes: Vec::new(), - extension_tool_executors: extension_tool_executors(&session), + extension_tool_executors: extension_tool_executors(&session, step_context.as_ref()), dynamic_tools: turn.dynamic_tools.as_slice(), }, &Default::default(), diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs index 45b178bc81..4c332fc6e0 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -73,10 +73,11 @@ 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) + .contribute_thread_context(session_store, thread_store, &step_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 414a67215b..4dc7fe874e 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,6 +21,7 @@ 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(); @@ -41,6 +42,7 @@ impl ContextContributor for UsageContributor { &'a self, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, + _step_store: &'a ExtensionData, ) -> std::pin::Pin> + 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 b8cb2cecc0..3f40fb67e5 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -78,15 +78,18 @@ 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. 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() }) } @@ -205,12 +208,16 @@ pub trait TurnLifecycleContributor: Send + Sync { /// host, not in this input. pub trait TurnInputContributor: Send + Sync { /// Returns additional contextual fragments for one submitted turn. + /// + /// `step_store` contains host capabilities bound to the sampling step that + /// will consume these fragments. fn contribute<'a>( &'a self, input: TurnInputContext, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, turn_store: &'a ExtensionData, + step_store: &'a ExtensionData, ) -> ExtensionFuture<'a, Vec>>; } @@ -270,11 +277,12 @@ pub trait SkillInvocationContributor: Send + Sync { /// Extension contribution that exposes native tools owned by a feature. pub trait ToolContributor: Send + Sync { - /// Returns the native tools visible for the supplied extension stores. + /// Returns native tools bound to the supplied sampling-step capabilities. fn tools( &self, session_store: &ExtensionData, thread_store: &ExtensionData, + step_store: &ExtensionData, ) -> Vec>>; } diff --git a/codex-rs/ext/extension-api/src/contributors/context.rs b/codex-rs/ext/extension-api/src/contributors/context.rs index fb6239ec4b..95e33f4ad3 100644 --- a/codex-rs/ext/extension-api/src/contributors/context.rs +++ b/codex-rs/ext/extension-api/src/contributors/context.rs @@ -15,6 +15,8 @@ pub struct TurnContextContributionInput<'a> { pub thread_store: &'a ExtensionData, /// Store scoped to this turn. pub turn_store: &'a ExtensionData, + /// Store scoped to this sampling step. + pub step_store: &'a ExtensionData, /// Effective model context window for this turn, when known. pub model_context_window: Option, } 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 75fae6d4a8..ae60de7ade 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -20,6 +20,8 @@ pub struct WorldStateContributionInput<'a> { pub session_store: &'a ExtensionData, pub thread_store: &'a ExtensionData, pub turn_store: &'a ExtensionData, + /// Host capabilities bound to this exact sampling step. + pub step_store: &'a ExtensionData, } /// What the harness knows about the previous value of one extension-owned section. diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index de9cdf1539..0f01ad0f28 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -41,6 +41,7 @@ impl ContextContributor for AllContributors { &'a self, _session_store: &'a ExtensionData, _thread_store: &'a ExtensionData, + _step_store: &'a ExtensionData, ) -> ExtensionFuture<'a, Vec> { Box::pin(std::future::ready(Vec::new())) } @@ -63,6 +64,7 @@ impl TurnInputContributor for AllContributors { _session_store: &'a ExtensionData, _thread_store: &'a ExtensionData, _turn_store: &'a ExtensionData, + _step_store: &'a ExtensionData, ) -> ExtensionFuture<'a, Vec>> { Box::pin(async move { let _self = self; @@ -77,6 +79,7 @@ impl ToolContributor for AllContributors { &self, _session_store: &ExtensionData, _thread_store: &ExtensionData, + _step_store: &ExtensionData, ) -> Vec>> { Vec::new() } @@ -158,6 +161,7 @@ impl ContextContributor for NamedContextContributor { &'a self, _session_store: &'a ExtensionData, _thread_store: &'a ExtensionData, + _step_store: &'a ExtensionData, ) -> ExtensionFuture<'a, Vec> { Box::pin(std::future::ready(vec![PromptFragment::developer_policy( self.0, @@ -219,12 +223,13 @@ async fn contributors_preserve_registration_order() { let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); let turn_store = ExtensionData::new("turn"); + let step_store = ExtensionData::new("step"); let mut fragments = Vec::new(); for contributor in registry.context_contributors() { fragments.extend( contributor - .contribute_thread_context(&session_store, &thread_store) + .contribute_thread_context(&session_store, &thread_store, &step_store) .await, ); } @@ -237,6 +242,7 @@ async fn contributors_preserve_registration_order() { session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &step_store, model_context_window: Some(123), }) .await, diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 4fa1081db9..f9ae24fbc6 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -420,6 +420,7 @@ 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 6f5e81ec95..37af7986d4 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1143,7 +1143,9 @@ async fn installed_tools_with_start( registry .tool_contributors() .iter() - .flat_map(|contributor| contributor.tools(&session_store, &thread_store)) + .flat_map(|contributor| { + contributor.tools(&session_store, &thread_store, &ExtensionData::new("step")) + }) .collect() } @@ -1205,7 +1207,13 @@ impl GoalExtensionHarness { self.registry .tool_contributors() .iter() - .flat_map(|contributor| contributor.tools(&self.session_store, &self.thread_store)) + .flat_map(|contributor| { + contributor.tools( + &self.session_store, + &self.thread_store, + &ExtensionData::new("step"), + ) + }) .collect() } diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 38a24eed7c..b4a598ddf3 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -86,6 +86,7 @@ impl ToolContributor for ImageGenerationExtension { &self, _session_store: &ExtensionData, thread_store: &ExtensionData, + _step_store: &ExtensionData, ) -> Vec>> { 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 464ce24576..e61a10df20 100644 --- a/codex-rs/ext/memories/src/extension.rs +++ b/codex-rs/ext/memories/src/extension.rs @@ -52,6 +52,7 @@ impl ContextContributor for MemoriesExtension { &'a self, _session_store: &'a ExtensionData, thread_store: &'a ExtensionData, + _step_store: &'a ExtensionData, ) -> std::pin::Pin> + Send + 'a>> { Box::pin(async move { let Some(config) = thread_store.get::() else { @@ -100,6 +101,7 @@ 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 4394ea9704..e9ab1dade3 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -41,7 +41,8 @@ fn tools_are_not_contributed_without_thread_config() { extension .tools( &ExtensionData::new("session"), - &ExtensionData::new("thread") + &ExtensionData::new("thread"), + &ExtensionData::new("step") ) .is_empty() ); @@ -59,7 +60,11 @@ fn tools_are_not_contributed_when_disabled() { assert!( extension - .tools(&ExtensionData::new("session"), &thread_store) + .tools( + &ExtensionData::new("session"), + &thread_store, + &ExtensionData::new("step"), + ) .is_empty() ); } @@ -76,7 +81,11 @@ fn tools_are_not_contributed_when_dedicated_tools_disabled() { assert!( extension - .tools(&ExtensionData::new("session"), &thread_store) + .tools( + &ExtensionData::new("session"), + &thread_store, + &ExtensionData::new("step"), + ) .is_empty() ); } @@ -92,7 +101,11 @@ fn tools_are_contributed_when_enabled_with_dedicated_tools() { }); let tool_names = extension - .tools(&ExtensionData::new("session"), &thread_store) + .tools( + &ExtensionData::new("session"), + &thread_store, + &ExtensionData::new("step"), + ) .into_iter() .map(|tool| tool.tool_name()) .collect::>(); @@ -123,7 +136,13 @@ fn install_registers_dedicated_tool_contributor() { let tool_names = registry .tool_contributors() .iter() - .flat_map(|contributor| contributor.tools(&ExtensionData::new("session"), &thread_store)) + .flat_map(|contributor| { + contributor.tools( + &ExtensionData::new("session"), + &thread_store, + &ExtensionData::new("step"), + ) + }) .map(|tool| tool.tool_name()) .collect::>(); @@ -181,7 +200,11 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() { }); let fragments = extension - .contribute_thread_context(&ExtensionData::new("session"), &thread_store) + .contribute_thread_context( + &ExtensionData::new("session"), + &thread_store, + &ExtensionData::new("step"), + ) .await; assert_eq!(fragments.len(), 1); diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index 115d1662e1..e34cac46f9 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -113,6 +113,7 @@ where &'a self, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, + _step_store: &'a ExtensionData, ) -> std::pin::Pin> + Send + 'a>> { Box::pin(async move { let Some(thread_state) = thread_store.get::() else { @@ -214,6 +215,7 @@ where &self, session_store: &ExtensionData, thread_store: &ExtensionData, + _step_store: &ExtensionData, ) -> Vec>> { let Some(thread_state) = thread_store.get::() else { return Vec::new(); @@ -269,6 +271,7 @@ where session_store: &'a ExtensionData, thread_store: &'a ExtensionData, turn_store: &'a ExtensionData, + _step_store: &'a ExtensionData, ) -> ExtensionFuture<'a, Vec>> { Box::pin(async move { let Some(thread_state) = thread_store.get::() else { diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index c6bc1e5e4f..3a30b2bd9c 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -121,6 +121,7 @@ async fn installed_extension_uses_host_service_snapshot() -> TestResult { &session_store, &thread_store, &turn_store, + &ExtensionData::new("step"), ) .await; @@ -192,7 +193,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach .await; let prompt_fragments = registry.context_contributors()[0] - .contribute_thread_context(&session_store, &thread_store) + .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step")) .await; assert!(prompt_fragments.is_empty()); @@ -212,6 +213,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &ExtensionData::new("step"), }) .await; assert_eq!(1, available_sections.len()); @@ -239,6 +241,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach &session_store, &thread_store, &turn_store, + &ExtensionData::new("step"), ) .await; @@ -265,6 +268,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &unavailable_turn_store, + step_store: &ExtensionData::new("step"), }) .await; let unavailable_snapshot = unavailable_sections[0].snapshot().clone(); @@ -288,6 +292,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &restored_turn_store, + step_store: &ExtensionData::new("step"), }) .await; let restored_snapshot = restored_sections[0].snapshot().clone(); @@ -316,6 +321,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &listing_disabled_turn_store, + step_store: &ExtensionData::new("step"), }) .await; let listing_disabled_fragment = listing_disabled_sections[0] @@ -380,7 +386,7 @@ async fn default_context_truncates_catalog_descriptions() -> TestResult { .await; let fragments = registry.context_contributors()[0] - .contribute_thread_context(&session_store, &thread_store) + .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step")) .await; assert_eq!(1, fragments.len()); let rendered = fragments[0].text(); @@ -429,7 +435,11 @@ async fn skills_list_truncates_catalog_descriptions_in_tool_output() -> TestResu }) .await; - let tools = registry.tool_contributors()[0].tools(&session_store, &thread_store); + let tools = registry.tool_contributors()[0].tools( + &session_store, + &thread_store, + &ExtensionData::new("step"), + ); let list_tool = tools .iter() .find(|tool| tool.tool_name().name == "list") @@ -503,7 +513,7 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult { .await; let initial_fragments = registry.context_contributors()[0] - .contribute_thread_context(&session_store, &thread_store) + .contribute_thread_context(&session_store, &thread_store, &ExtensionData::new("step")) .await; assert!(initial_fragments.is_empty()); let EventMsg::Warning(warning) = event_rx.try_recv()?.msg else { @@ -528,6 +538,7 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult { &session_store, &thread_store, &ExtensionData::new(turn_id), + &ExtensionData::new("step"), ) .await; assert!(fragments.is_empty()); @@ -607,6 +618,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &ExtensionData::new("step"), }) .await; let fragments = registry.turn_input_contributors()[0] @@ -622,6 +634,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te &session_store, &thread_store, &turn_store, + &ExtensionData::new("step"), ) .await; @@ -697,6 +710,7 @@ async fn prompt_hidden_skill_can_still_be_invoked() -> TestResult { &session_store, &thread_store, &ExtensionData::new("turn-1"), + &ExtensionData::new("step"), ) .await; diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index d59a34649d..c573506664 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -120,6 +120,7 @@ impl ToolContributor for WebSearchExtension { &self, session_store: &ExtensionData, thread_store: &ExtensionData, + _step_store: &ExtensionData, ) -> Vec>> { let Some(config) = thread_store.get::() else { return Vec::new(); @@ -207,7 +208,9 @@ mod tests { let tool_names = registry .tool_contributors() .iter() - .flat_map(|contributor| contributor.tools(&session_store, &thread_store)) + .flat_map(|contributor| { + contributor.tools(&session_store, &thread_store, &ExtensionData::new("step")) + }) .map(|tool| (tool.tool_name(), tool.supports_parallel_tool_calls())) .collect::>();