agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make context contributors async (#22491)

Browse Source 
## 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
This commit is contained in:
jif-oai
2026-05-13 16:43:28 +02:00
committed by GitHub
Unverified
parent 1dcc89f1d4
commit 68e045a631
6 changed files with 114 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/mod.rs
+9 -5
View File
@@ -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());
codex-rs/core/src/session/tests.rs
+19 -15
View File
@@ -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<codex_extension_api::PromptFragment> {
thread_store
.get::<GitAttributionTestState>()
.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<dyn std::future::Future<Output = Vec<codex_extension_api::PromptFragment>> + Send + 'a>,
> {
Box::pin(async move {
thread_store
.get::<GitAttributionTestState>()
.is_some()
.then(|| {
codex_extension_api::PromptFragment::developer_policy(
"git attribution extension enabled",
)
})
.into_iter()
.collect()
})
}
}
codex-rs/ext/extension-api/examples/enabled_extensions.rs
+40 -9
View File
@@ -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(&registry, &session_store, &first_thread_store);
contribute_prompt(&registry, &session_store, &first_thread_store);
contribute_prompt(&registry, &session_store, &second_thread_store);
let first_thread_fragments = block_on_ready(contribute_prompt(
&registry,
&session_store,
&first_thread_store,
));
block_on_ready(contribute_prompt(
&registry,
&session_store,
&first_thread_store,
));
block_on_ready(contribute_prompt(
&registry,
&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<codex_extension_api::PromptFragment> {
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<F>(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"),
}
}
codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs
+24 -20
View File
@@ -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<PromptFragment> {
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<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + 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<PromptFragment> {
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<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + 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.",
)]
})
}
}
codex-rs/ext/extension-api/src/contributors.rs
+5 -5
View File
@@ -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<PromptFragment>;
fn contribute<'a>(
&'a self,
session_store: &'a ExtensionData,
thread_store: &'a ExtensionData,
) -> std::pin::Pin<Box<dyn Future<Output = Vec<PromptFragment>> + Send + 'a>>;
}
/// Contributor for host-owned thread lifecycle gates.
codex-rs/ext/git-attribution/src/lib.rs
+17 -15
View File
@@ -16,21 +16,23 @@ const DEFAULT_ATTRIBUTION_VALUE: &str = "Codex <noreply@openai.com>";
pub struct GitAttributionExtension;
impl ContextContributor for GitAttributionExtension {
fn contribute(
&self,
_session_store: &ExtensionData,
thread_store: &ExtensionData,
) -> Vec<PromptFragment> {
let Some(config_store) = thread_store.get::<GitAttributionConfig>() 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<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
Box::pin(async move {
let Some(config_store) = thread_store.get::<GitAttributionConfig>() 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()
})
}
}