From 68e045a631e108ae069660bc2ce1148f46d29130 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 13 May 2026 16:43:28 +0200 Subject: [PATCH] Make context contributors async (#22491) ## Summary - make ContextContributor return a boxed Send future - await context contributors during initial context assembly - update existing contributors and extension-api examples for the async contract ## Testing - cargo test -p codex-extension-api --examples - cargo test -p codex-git-attribution - cargo test -p codex-core build_initial_context_includes_git_attribution_from_extensions -- --nocapture - cargo test -p codex-core build_initial_context_omits_git_attribution_when_feature_is_disabled -- --nocapture - cargo test -p codex-core (fails in unrelated agent::control::tests::spawn_agent_fork_last_n_turns_keeps_only_recent_turns stack overflow) - just fix -p codex-extension-api - just fix -p codex-git-attribution - just fix -p codex-core - cargo clippy -p codex-extension-api --examples --- codex-rs/core/src/session/mod.rs | 14 ++++-- codex-rs/core/src/session/tests.rs | 34 +++++++------ .../examples/enabled_extensions.rs | 49 +++++++++++++++---- .../shared_state_extension.rs | 44 +++++++++-------- .../ext/extension-api/src/contributors.rs | 10 ++-- codex-rs/ext/git-attribution/src/lib.rs | 32 ++++++------ 6 files changed, 114 insertions(+), 69 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 3b8939ec5e..a1d598bff7 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2694,11 +2694,15 @@ impl Session { { developer_sections.push(plugin_instructions.render()); } - for contributor in self.services.extensions.context_contributors() { - for fragment in contributor.contribute( - &self.services.session_extension_data, - &self.services.thread_extension_data, - ) { + let context_contributors = self.services.extensions.context_contributors().to_vec(); + for contributor in context_contributors { + for fragment in contributor + .contribute( + &self.services.session_extension_data, + &self.services.thread_extension_data, + ) + .await + { match fragment.slot() { PromptSlot::DeveloperPolicy | PromptSlot::DeveloperCapabilities => { developer_sections.push(fragment.text().to_string()); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 45ac5ab9c5..5a4bc2d0e7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6352,21 +6352,25 @@ struct GitAttributionTestContributor; struct GitAttributionTestState; impl codex_extension_api::ContextContributor for GitAttributionTestContributor { - fn contribute( - &self, - _session_store: &codex_extension_api::ExtensionData, - thread_store: &codex_extension_api::ExtensionData, - ) -> Vec { - thread_store - .get::() - .is_some() - .then(|| { - codex_extension_api::PromptFragment::developer_policy( - "git attribution extension enabled", - ) - }) - .into_iter() - .collect() + fn contribute<'a>( + &'a self, + _session_store: &'a codex_extension_api::ExtensionData, + thread_store: &'a codex_extension_api::ExtensionData, + ) -> std::pin::Pin< + Box> + Send + 'a>, + > { + Box::pin(async move { + thread_store + .get::() + .is_some() + .then(|| { + codex_extension_api::PromptFragment::developer_policy( + "git attribution extension enabled", + ) + }) + .into_iter() + .collect() + }) } } diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs index f32af97c1a..9027c28241 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -1,6 +1,12 @@ #[path = "enabled_extensions/shared_state_extension.rs"] mod shared_state_extension; +use std::future::Future; +use std::pin::pin; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistryBuilder; use shared_state_extension::recorded_style_contributions; @@ -18,9 +24,21 @@ fn main() { let second_thread_store = ExtensionData::new("thread-2"); // 3. Reusing the same session store shares session state across threads. - let first_thread_fragments = contribute_prompt(®istry, &session_store, &first_thread_store); - contribute_prompt(®istry, &session_store, &first_thread_store); - contribute_prompt(®istry, &session_store, &second_thread_store); + let first_thread_fragments = block_on_ready(contribute_prompt( + ®istry, + &session_store, + &first_thread_store, + )); + block_on_ready(contribute_prompt( + ®istry, + &session_store, + &first_thread_store, + )); + block_on_ready(contribute_prompt( + ®istry, + &session_store, + &second_thread_store, + )); println!("first prompt fragments: {}", first_thread_fragments.len()); println!( @@ -49,14 +67,27 @@ fn main() { ); } -fn contribute_prompt( +async fn contribute_prompt( registry: &codex_extension_api::ExtensionRegistry<()>, session_store: &ExtensionData, thread_store: &ExtensionData, ) -> Vec { - registry - .context_contributors() - .iter() - .flat_map(|contributor| contributor.contribute(session_store, thread_store)) - .collect() + let mut fragments = Vec::new(); + for contributor in registry.context_contributors() { + fragments.extend(contributor.contribute(session_store, thread_store).await); + } + fragments +} + +fn block_on_ready(future: F) -> F::Output +where + F: Future, +{ + let waker = Waker::noop(); + let mut context = Context::from_waker(waker); + let mut future = pin!(future); + match future.as_mut().poll(&mut context) { + Poll::Ready(output) => output, + Poll::Pending => panic!("example context contributors should complete immediately"), + } } 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 25ab027c88..531f65b99e 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 @@ -17,17 +17,19 @@ pub fn install(registry: &mut ExtensionRegistryBuilder<()>) { struct StyleContributor; impl ContextContributor for StyleContributor { - fn contribute( - &self, - session_store: &ExtensionData, - thread_store: &ExtensionData, - ) -> Vec { - contribution_counts(session_store).record_style(); - contribution_counts(thread_store).record_style(); + fn contribute<'a>( + &'a self, + session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + contribution_counts(session_store).record_style(); + contribution_counts(thread_store).record_style(); - vec![PromptFragment::developer_policy( - "Prefer short answers unless the user asks for detail.", - )] + vec![PromptFragment::developer_policy( + "Prefer short answers unless the user asks for detail.", + )] + }) } } @@ -35,17 +37,19 @@ impl ContextContributor for StyleContributor { struct UsageContributor; impl ContextContributor for UsageContributor { - fn contribute( - &self, - session_store: &ExtensionData, - thread_store: &ExtensionData, - ) -> Vec { - contribution_counts(session_store).record_usage(); - contribution_counts(thread_store).record_usage(); + fn contribute<'a>( + &'a self, + session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + contribution_counts(session_store).record_usage(); + contribution_counts(thread_store).record_usage(); - vec![PromptFragment::developer_capability( - "This extension can contribute more than one prompt fragment.", - )] + vec![PromptFragment::developer_capability( + "This extension can contribute more than one prompt fragment.", + )] + }) } } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 7393f0b3db..aea59c8a50 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -26,11 +26,11 @@ pub use turn_lifecycle::TurnStopInput; /// Extension contribution that adds prompt fragments during prompt assembly. pub trait ContextContributor: Send + Sync { - fn contribute( - &self, - session_store: &ExtensionData, - thread_store: &ExtensionData, - ) -> Vec; + fn contribute<'a>( + &'a self, + session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + ) -> std::pin::Pin> + Send + 'a>>; } /// Contributor for host-owned thread lifecycle gates. diff --git a/codex-rs/ext/git-attribution/src/lib.rs b/codex-rs/ext/git-attribution/src/lib.rs index 1df7815d6e..03655296da 100644 --- a/codex-rs/ext/git-attribution/src/lib.rs +++ b/codex-rs/ext/git-attribution/src/lib.rs @@ -16,21 +16,23 @@ const DEFAULT_ATTRIBUTION_VALUE: &str = "Codex "; pub struct GitAttributionExtension; impl ContextContributor for GitAttributionExtension { - fn contribute( - &self, - _session_store: &ExtensionData, - thread_store: &ExtensionData, - ) -> Vec { - let Some(config_store) = thread_store.get::() else { - return Vec::new(); - }; - if !config_store.enabled { - return Vec::new(); - } - build_instruction(config_store.prompt.as_deref()) - .map(PromptFragment::developer_policy) - .into_iter() - .collect() + fn contribute<'a>( + &'a self, + _session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + let Some(config_store) = thread_store.get::() else { + return Vec::new(); + }; + if !config_store.enabled { + return Vec::new(); + } + build_instruction(config_store.prompt.as_deref()) + .map(PromptFragment::developer_policy) + .into_iter() + .collect() + }) } }