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

Remove async-trait from extension contributors (#27383)

Browse Source 
## Why

Extension contributors are registered behind `dyn Trait` objects, so
native `async fn`/RPITIT methods would make these traits
non-object-safe. Spell out the boxed, `Send` future contract directly so
`extension-api` no longer needs `async-trait` while retaining the
existing runtime model.

## What changed

- add a shared `ExtensionFuture` alias and use it for asynchronous
contributor methods
- migrate production and test implementations to return `Box::pin(async
move { ... })`
- remove `async-trait` dependencies where they are no longer used,
keeping it dev-only where unrelated test executors still require it

## Behavior

No behavior change is intended. Contributor futures remain boxed,
`Send`, dynamically dispatched, and lazily executed; cancellation and
callback ordering stay unchanged.

## Testing

- `just test -p codex-extension-api` (11 passed)
- affected extension crates (64 passed)
- targeted `codex-core` contributor tests (14 passed)
- `just fmt`
- `just bazel-lock-update`
- `just bazel-lock-check`

A broad local `codex-core` run compiled successfully but encountered
unrelated sandbox and missing test-binary fixture failures; CI will run
the full checks.
This commit is contained in:
jif
2026-06-10 13:31:09 +01:00
committed by GitHub
Unverified
parent ced1b8aa88
commit d2f6d23c6c
21 changed files with 736 additions and 597 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/tests.rs
+131 -96
View File
@@ -1922,26 +1922,27 @@ async fn record_token_usage_info_notifies_extension_contributors() {
records: Arc<std::sync::Mutex<Vec<RecordedTokenUsage>>>,
}
#[async_trait::async_trait]
impl codex_extension_api::TokenUsageContributor for TokenUsageRecorder {
async fn on_token_usage(
&self,
session_store: &codex_extension_api::ExtensionData,
thread_store: &codex_extension_api::ExtensionData,
turn_store: &codex_extension_api::ExtensionData,
token_usage: &TokenUsageInfo,
) {
self.records
.lock()
.expect("token usage records lock")
.push(RecordedTokenUsage {
session_level_id: session_store.level_id().to_string(),
thread_level_id: thread_store.level_id().to_string(),
turn_level_id: turn_store.level_id().to_string(),
token_usage: token_usage.clone(),
saw_session_store: session_store.get::<SessionTokenUsageMarker>().is_some(),
saw_thread_store: thread_store.get::<ThreadTokenUsageMarker>().is_some(),
});
fn on_token_usage<'a>(
&'a self,
session_store: &'a codex_extension_api::ExtensionData,
thread_store: &'a codex_extension_api::ExtensionData,
turn_store: &'a codex_extension_api::ExtensionData,
token_usage: &'a TokenUsageInfo,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
self.records
.lock()
.expect("token usage records lock")
.push(RecordedTokenUsage {
session_level_id: session_store.level_id().to_string(),
thread_level_id: thread_store.level_id().to_string(),
turn_level_id: turn_store.level_id().to_string(),
token_usage: token_usage.clone(),
saw_session_store: session_store.get::<SessionTokenUsageMarker>().is_some(),
saw_thread_store: thread_store.get::<ThreadTokenUsageMarker>().is_some(),
});
})
}
}
@@ -2040,25 +2041,32 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() {
records: Arc<std::sync::Mutex<Vec<RecordedTurnStart>>>,
}
#[async_trait::async_trait]
impl codex_extension_api::TurnLifecycleContributor for TurnStartRecorder {
async fn on_turn_start(&self, input: codex_extension_api::TurnStartInput<'_>) {
self.records
.lock()
.expect("turn start records lock")
.push(RecordedTurnStart {
session_level_id: input.session_store.level_id().to_string(),
thread_level_id: input.thread_store.level_id().to_string(),
turn_level_id: input.turn_store.level_id().to_string(),
turn_id: input.turn_id.to_string(),
collaboration_mode: input.collaboration_mode.clone(),
token_usage_at_turn_start: input.token_usage_at_turn_start.clone(),
saw_session_store: input
.session_store
.get::<SessionTurnStartMarker>()
.is_some(),
saw_thread_store: input.thread_store.get::<ThreadTurnStartMarker>().is_some(),
});
fn on_turn_start<'a>(
&'a self,
input: codex_extension_api::TurnStartInput<'a>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
self.records
.lock()
.expect("turn start records lock")
.push(RecordedTurnStart {
session_level_id: input.session_store.level_id().to_string(),
thread_level_id: input.thread_store.level_id().to_string(),
turn_level_id: input.turn_store.level_id().to_string(),
turn_id: input.turn_id.to_string(),
collaboration_mode: input.collaboration_mode.clone(),
token_usage_at_turn_start: input.token_usage_at_turn_start.clone(),
saw_session_store: input
.session_store
.get::<SessionTurnStartMarker>()
.is_some(),
saw_thread_store: input
.thread_store
.get::<ThreadTurnStartMarker>()
.is_some(),
});
})
}
}
@@ -2138,24 +2146,31 @@ async fn turn_error_lifecycle_exposes_error_and_stores() {
records: Arc<std::sync::Mutex<Vec<RecordedTurnError>>>,
}
#[async_trait::async_trait]
impl codex_extension_api::TurnLifecycleContributor for TurnErrorRecorder {
async fn on_turn_error(&self, input: codex_extension_api::TurnErrorInput<'_>) {
self.records
.lock()
.expect("turn error records lock")
.push(RecordedTurnError {
session_level_id: input.session_store.level_id().to_string(),
thread_level_id: input.thread_store.level_id().to_string(),
turn_level_id: input.turn_store.level_id().to_string(),
turn_id: input.turn_id.to_string(),
error: input.error,
saw_session_store: input
.session_store
.get::<SessionTurnErrorMarker>()
.is_some(),
saw_thread_store: input.thread_store.get::<ThreadTurnErrorMarker>().is_some(),
});
fn on_turn_error<'a>(
&'a self,
input: codex_extension_api::TurnErrorInput<'a>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
self.records
.lock()
.expect("turn error records lock")
.push(RecordedTurnError {
session_level_id: input.session_store.level_id().to_string(),
thread_level_id: input.thread_store.level_id().to_string(),
turn_level_id: input.turn_store.level_id().to_string(),
turn_id: input.turn_id.to_string(),
error: input.error,
saw_session_store: input
.session_store
.get::<SessionTurnErrorMarker>()
.is_some(),
saw_thread_store: input
.thread_store
.get::<ThreadTurnErrorMarker>()
.is_some(),
});
})
}
}
@@ -6411,16 +6426,20 @@ async fn submission_loop_channel_close_emits_thread_stop_lifecycle() {
expected_thread_id: ThreadId,
}
#[async_trait::async_trait]
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for ThreadStopRecorder {
async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
assert!(input.session_store.get::<SessionStopMarker>().is_some());
assert!(input.thread_store.get::<ThreadStopMarker>().is_some());
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
fn on_thread_stop<'a>(
&'a self,
input: codex_extension_api::ThreadStopInput<'a>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
assert!(input.session_store.get::<SessionStopMarker>().is_some());
assert!(input.thread_store.get::<ThreadStopMarker>().is_some());
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
})
}
}
@@ -6457,33 +6476,41 @@ async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lif
expected_turn_id: String,
}
#[async_trait::async_trait]
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for LifecycleRecorder {
async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
self.calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push("thread_stop");
fn on_thread_stop<'a>(
&'a self,
input: codex_extension_api::ThreadStopInput<'a>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
self.calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push("thread_stop");
})
}
}
#[async_trait::async_trait]
impl codex_extension_api::TurnLifecycleContributor for LifecycleRecorder {
async fn on_turn_abort(&self, input: codex_extension_api::TurnAbortInput<'_>) {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
assert_eq!(self.expected_turn_id, input.turn_store.level_id());
assert_eq!(TurnAbortReason::Interrupted, input.reason);
self.calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push("turn_abort");
fn on_turn_abort<'a>(
&'a self,
input: codex_extension_api::TurnAbortInput<'a>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
assert_eq!(self.expected_turn_id, input.turn_store.level_id());
assert_eq!(TurnAbortReason::Interrupted, input.reason);
self.calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push("turn_abort");
})
}
}
@@ -8698,15 +8725,19 @@ async fn task_finish_emits_thread_idle_lifecycle_after_active_turn_clears() {
expected_thread_id: ThreadId,
}
#[async_trait::async_trait]
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for ThreadIdleRecorder {
async fn on_thread_idle(&self, input: codex_extension_api::ThreadIdleInput<'_>) {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
self.idle_tx.send(()).await.expect("idle receiver open");
fn on_thread_idle<'a>(
&'a self,
input: codex_extension_api::ThreadIdleInput<'a>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
);
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
self.idle_tx.send(()).await.expect("idle receiver open");
})
}
}
@@ -8740,10 +8771,14 @@ async fn thread_idle_lifecycle_waits_for_trigger_turn_mailbox_work() {
calls: Arc<std::sync::atomic::AtomicUsize>,
}
#[async_trait::async_trait]
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for ThreadIdleRecorder {
async fn on_thread_idle(&self, _input: codex_extension_api::ThreadIdleInput<'_>) {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
fn on_thread_idle<'a>(
&'a self,
_input: codex_extension_api::ThreadIdleInput<'a>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
})
}
}
codex-rs/core/src/session/turn_tests.rs
+14 -13
View File
@@ -7,20 +7,21 @@ use std::sync::Arc;
struct RewriteAgentMessageContributor;
#[async_trait::async_trait]
impl TurnItemContributor for RewriteAgentMessageContributor {
async fn contribute(
&self,
_thread_store: &ExtensionData,
_turn_store: &ExtensionData,
item: &mut TurnItem,
) -> Result<(), String> {
if let TurnItem::AgentMessage(agent_message) = item {
agent_message.content = vec![AgentMessageContent::Text {
text: "plan contributed assistant text".to_string(),
}];
}
Ok(())
fn contribute<'a>(
&'a self,
_thread_store: &'a ExtensionData,
_turn_store: &'a ExtensionData,
item: &'a mut TurnItem,
) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> {
Box::pin(async move {
if let TurnItem::AgentMessage(agent_message) = item {
agent_message.content = vec![AgentMessageContent::Text {
text: "plan contributed assistant text".to_string(),
}];
}
Ok(())
})
}
}
codex-rs/core/src/stream_events_utils_tests.rs
+30 -28
View File
@@ -167,41 +167,43 @@ struct TestTurnItemContributor;
#[derive(Debug)]
struct TurnItemContributorRan;
#[async_trait::async_trait]
impl TurnItemContributor for TestTurnItemContributor {
async fn contribute(
&self,
_thread_store: &ExtensionData,
turn_store: &ExtensionData,
item: &mut TurnItem,
) -> Result<(), String> {
turn_store.insert(TurnItemContributorRan);
if let TurnItem::AgentMessage(agent_message) = item {
agent_message.memory_citation = Some(MemoryCitation {
entries: Vec::new(),
rollout_ids: Vec::new(),
});
}
Ok(())
fn contribute<'a>(
&'a self,
_thread_store: &'a ExtensionData,
turn_store: &'a ExtensionData,
item: &'a mut TurnItem,
) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> {
Box::pin(async move {
turn_store.insert(TurnItemContributorRan);
if let TurnItem::AgentMessage(agent_message) = item {
agent_message.memory_citation = Some(MemoryCitation {
entries: Vec::new(),
rollout_ids: Vec::new(),
});
}
Ok(())
})
}
}
struct RewriteAgentMessageContributor;
#[async_trait::async_trait]
impl TurnItemContributor for RewriteAgentMessageContributor {
async fn contribute(
&self,
_thread_store: &ExtensionData,
_turn_store: &ExtensionData,
item: &mut TurnItem,
) -> Result<(), String> {
if let TurnItem::AgentMessage(agent_message) = item {
agent_message.content = vec![AgentMessageContent::Text {
text: "contributed assistant text".to_string(),
}];
}
Ok(())
fn contribute<'a>(
&'a self,
_thread_store: &'a ExtensionData,
_turn_store: &'a ExtensionData,
item: &'a mut TurnItem,
) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> {
Box::pin(async move {
if let TurnItem::AgentMessage(agent_message) = item {
agent_message.content = vec![AgentMessageContent::Text {
text: "contributed assistant text".to_string(),
}];
}
Ok(())
})
}
}
codex-rs/core/src/thread_manager_tests.rs
+17 -13
View File
@@ -394,20 +394,24 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() {
observed: Arc<std::sync::Mutex<Option<(String, String)>>>,
}
#[async_trait::async_trait]
impl codex_extension_api::ThreadLifecycleContributor<Config> for InitialDataRecorder {
async fn on_thread_start(&self, input: codex_extension_api::ThreadStartInput<'_, Config>) {
let marker = input
.thread_store
.get::<InitialMarker>()
.expect("initial extension data should be available");
*self
.observed
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some((
input.thread_store.level_id().to_string(),
marker.0.to_string(),
));
fn on_thread_start<'a>(
&'a self,
input: codex_extension_api::ThreadStartInput<'a, Config>,
) -> codex_extension_api::ExtensionFuture<'a, ()> {
Box::pin(async move {
let marker = input
.thread_store
.get::<InitialMarker>()
.expect("initial extension data should be available");
*self
.observed
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some((
input.thread_store.level_id().to_string(),
marker.0.to_string(),
));
})
}
}
codex-rs/core/src/tools/handlers/extension_tools.rs
+10 -9
View File
@@ -388,16 +388,17 @@ mod tests {
struct RecordExtensionTurnItemContributor;
#[async_trait::async_trait]
impl TurnItemContributor for RecordExtensionTurnItemContributor {
async fn contribute(
&self,
_thread_store: &ExtensionData,
turn_store: &ExtensionData,
_item: &mut TurnItem,
) -> Result<(), String> {
turn_store.insert(ExtensionTurnItemContributorRan);
Ok(())
fn contribute<'a>(
&'a self,
_thread_store: &'a ExtensionData,
turn_store: &'a ExtensionData,
_item: &'a mut TurnItem,
) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> {
Box::pin(async move {
turn_store.insert(ExtensionTurnItemContributorRan);
Ok(())
})
}
}