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

Add step-scoped data to extension contributors (#34525)

Browse Source 
## 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
This commit is contained in:
jif
2026-07-21 11:53:56 +00:00
committed by copyberry
parent 2d85e6d3a6
commit c44c4de7b4
28 changed files with 195 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/compact.rs
+18 -11
View File
@@ -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<WorldState>),
BeforeLastUserMessage {
world_state: Arc<WorldState>,
step_context: Arc<StepContext>,
},
DoNotInject,
}
pub(crate) async fn build_compaction_initial_context(
sess: &Session,
turn_context: &TurnContext,
initial_context_injection: &InitialContextInjection,
) -> (Vec<ResponseItem>, Option<Arc<WorldState>>) {
// 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())
}
};
codex-rs/core/src/compact_remote.rs
+4 -10
View File
@@ -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<ResponseItem>,
initial_context_injection: &InitialContextInjection,
) -> (Vec<ResponseItem>, Option<Arc<WorldState>>) {
@@ -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);
(
codex-rs/core/src/compact_remote_v2.rs
+4 -8
View File
@@ -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())
}
};
codex-rs/core/src/compact_tests.rs
+13 -5
View File
@@ -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,
)
codex-rs/core/src/compact_token_budget.rs
+6 -6
View File
@@ -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<StepContext>,
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<Session>,
turn_context: &Arc<TurnContext>,
step_context: &Arc<StepContext>,
world_state: Arc<WorldState>,
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;
codex-rs/core/src/guardian/tests.rs
+1
View File
@@ -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<codex_extension_api::PromptFragment>> {
Box::pin(async move {
if thread_store
codex-rs/core/src/session/mod.rs
+26 -7
View File
@@ -3152,8 +3152,9 @@ impl Session {
async fn build_turn_context_contribution_items(
&self,
turn_context: &TurnContext,
step_context: &StepContext,
) -> Vec<ResponseItem> {
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<ResponseItem> {
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<ResponseItem> {
let mut developer_sections = Vec::<String>::with_capacity(8);
let mut contextual_user_sections = Vec::<String>::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<WorldState>,
) -> 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,
);
}
codex-rs/core/src/session/step_context.rs
+5
View File
@@ -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<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>>,
}
@@ -35,6 +38,7 @@ impl StepContext {
mcp: Arc<McpRuntimeSnapshot>,
loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
) -> 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,
}
}
codex-rs/core/src/session/tests.rs
+6 -3
View File
@@ -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<dyn std::future::Future<Output = Vec<codex_extension_api::PromptFragment>> + Send + 'a>,
> {
codex-rs/core/src/session/tests/guardian_tests.rs
+6 -3
View File
@@ -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,
codex-rs/core/src/session/turn.rs
+9 -4
View File
@@ -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<Session>,
turn_context: &TurnContext,
step_context: &StepContext,
user_input: &[UserInput],
cancellation_token: &CancellationToken,
) -> Option<Vec<ResponseItem>> {
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,
codex-rs/core/src/session/world_state.rs
+1
View File
@@ -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
{
codex-rs/core/src/tools/router.rs
+2
View File
@@ -246,6 +246,7 @@ 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
@@ -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()
codex-rs/core/src/tools/router_tests.rs
+2 -1
View File
@@ -43,6 +43,7 @@ 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)]
}
@@ -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(),
codex-rs/ext/extension-api/examples/enabled_extensions.rs
+2 -1
View File
@@ -73,10 +73,11 @@ 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)
.contribute_thread_context(session_store, thread_store, &step_store)
.await,
);
}
codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs
+2
View File
@@ -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<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + 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<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
+9 -1
View File
@@ -78,15 +78,18 @@ 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.
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()
})
}
@@ -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<Box<dyn ContextualUserFragment + Send>>>;
}
@@ -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<Arc<dyn ToolExecutor<ToolCall>>>;
}
codex-rs/ext/extension-api/src/contributors/context.rs
+2
View File
@@ -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<i64>,
}
codex-rs/ext/extension-api/src/contributors/world_state.rs
+2
View File
@@ -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.
codex-rs/ext/extension-api/tests/registry.rs
+7 -1
View File
@@ -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<PromptFragment>> {
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<dyn ContextualUserFragment + Send>>> {
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<Arc<dyn ToolExecutor<ToolCall>>> {
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<PromptFragment>> {
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,
codex-rs/ext/goal/src/extension.rs
+1
View File
@@ -420,6 +420,7 @@ 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
+10 -2
View File
@@ -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()
}
codex-rs/ext/image-generation/src/extension.rs
+1
View File
@@ -86,6 +86,7 @@ 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,6 +52,7 @@ 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 {
@@ -100,6 +101,7 @@ 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
+29 -6
View File
@@ -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::<Vec<_>>();
@@ -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::<Vec<_>>();
@@ -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);
codex-rs/ext/skills/src/extension.rs
+3
View File
@@ -113,6 +113,7 @@ 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 {
@@ -214,6 +215,7 @@ 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();
@@ -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<dyn ContextualUserFragment + Send>>> {
Box::pin(async move {
let Some(thread_state) = thread_store.get::<SkillsThreadState>() else {
codex-rs/ext/skills/tests/skills_extension.rs
+18 -4
View File
@@ -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;
codex-rs/ext/web-search/src/extension.rs
+4 -1
View File
@@ -120,6 +120,7 @@ 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();
@@ -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::<Vec<_>>();