mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Remove async-trait from extension contributors (#27383)
## 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
GitHub
+199
-184
@@ -1,12 +1,12 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::Weak;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use codex_analytics::AnalyticsEventsClient;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_extension_api::ConfigContributor;
|
||||
use codex_extension_api::ExtensionData;
|
||||
use codex_extension_api::ExtensionEventSink;
|
||||
use codex_extension_api::ExtensionFuture;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::ThreadIdleInput;
|
||||
use codex_extension_api::ThreadLifecycleContributor;
|
||||
@@ -95,76 +95,83 @@ impl<C> GoalExtension<C> {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> ThreadLifecycleContributor<C> for GoalExtension<C>
|
||||
where
|
||||
C: Send + Sync + 'static,
|
||||
{
|
||||
async fn on_thread_start(&self, input: ThreadStartInput<'_, C>) {
|
||||
let enabled = (self.goals_enabled)(input.config);
|
||||
let tools_available_for_thread = input.persistent_thread_state_available
|
||||
&& !matches!(
|
||||
input.session_source,
|
||||
SessionSource::SubAgent(SubAgentSource::Review)
|
||||
);
|
||||
input
|
||||
.thread_store
|
||||
.insert(GoalExtensionConfig::from_enabled(enabled));
|
||||
let accounting_state = input
|
||||
.thread_store
|
||||
.get_or_init::<GoalAccountingState>(GoalAccountingState::default);
|
||||
let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else {
|
||||
return;
|
||||
};
|
||||
let runtime = input.thread_store.get_or_init::<GoalRuntimeHandle>(|| {
|
||||
GoalRuntimeHandle::new(
|
||||
thread_id,
|
||||
Arc::clone(&self.state_dbs),
|
||||
self.event_emitter.clone(),
|
||||
self.metrics.clone(),
|
||||
self.thread_manager.clone(),
|
||||
accounting_state,
|
||||
GoalRuntimeConfig {
|
||||
analytics: self.analytics.clone(),
|
||||
enabled,
|
||||
tools_available_for_thread,
|
||||
},
|
||||
)
|
||||
});
|
||||
runtime.set_enabled(enabled);
|
||||
self.goal_service.register_runtime(&runtime);
|
||||
fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let enabled = (self.goals_enabled)(input.config);
|
||||
let tools_available_for_thread = input.persistent_thread_state_available
|
||||
&& !matches!(
|
||||
input.session_source,
|
||||
SessionSource::SubAgent(SubAgentSource::Review)
|
||||
);
|
||||
input
|
||||
.thread_store
|
||||
.insert(GoalExtensionConfig::from_enabled(enabled));
|
||||
let accounting_state = input
|
||||
.thread_store
|
||||
.get_or_init::<GoalAccountingState>(GoalAccountingState::default);
|
||||
let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else {
|
||||
return;
|
||||
};
|
||||
let runtime = input.thread_store.get_or_init::<GoalRuntimeHandle>(|| {
|
||||
GoalRuntimeHandle::new(
|
||||
thread_id,
|
||||
Arc::clone(&self.state_dbs),
|
||||
self.event_emitter.clone(),
|
||||
self.metrics.clone(),
|
||||
self.thread_manager.clone(),
|
||||
accounting_state,
|
||||
GoalRuntimeConfig {
|
||||
analytics: self.analytics.clone(),
|
||||
enabled,
|
||||
tools_available_for_thread,
|
||||
},
|
||||
)
|
||||
});
|
||||
runtime.set_enabled(enabled);
|
||||
self.goal_service.register_runtime(&runtime);
|
||||
})
|
||||
}
|
||||
|
||||
async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
|
||||
if let Err(err) = runtime.restore_after_resume().await {
|
||||
tracing::warn!(
|
||||
"failed to restore goal runtime after thread resume for {}: {err}",
|
||||
runtime.thread_id()
|
||||
);
|
||||
}
|
||||
if let Err(err) = runtime.restore_after_resume().await {
|
||||
tracing::warn!(
|
||||
"failed to restore goal runtime after thread resume for {}: {err}",
|
||||
runtime.thread_id()
|
||||
);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn on_thread_idle(&self, input: ThreadIdleInput<'_>) {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
|
||||
if let Err(err) = runtime.continue_if_idle().await {
|
||||
tracing::warn!(
|
||||
"failed to continue active goal for idle thread {}: {err}",
|
||||
runtime.thread_id()
|
||||
);
|
||||
}
|
||||
if let Err(err) = runtime.continue_if_idle().await {
|
||||
tracing::warn!(
|
||||
"failed to continue active goal for idle thread {}: {err}",
|
||||
runtime.thread_id()
|
||||
);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn on_thread_stop(&self, input: ThreadStopInput<'_>) {
|
||||
if let Some(runtime) = goal_runtime_handle(input.thread_store) {
|
||||
self.goal_service.unregister_runtime(&runtime);
|
||||
}
|
||||
fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
if let Some(runtime) = goal_runtime_handle(input.thread_store) {
|
||||
self.goal_service.unregister_runtime(&runtime);
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -187,153 +194,161 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> TurnLifecycleContributor for GoalExtension<C>
|
||||
where
|
||||
C: Send + Sync + 'static,
|
||||
{
|
||||
async fn on_turn_start(&self, input: TurnStartInput<'_>) {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let accounting = runtime.accounting_state();
|
||||
accounting.start_turn(
|
||||
input.turn_id,
|
||||
input.collaboration_mode.mode,
|
||||
input.token_usage_at_turn_start,
|
||||
);
|
||||
if matches!(
|
||||
input.collaboration_mode.mode,
|
||||
codex_protocol::config_types::ModeKind::Plan
|
||||
) {
|
||||
accounting.clear_current_turn_goal();
|
||||
return;
|
||||
}
|
||||
let Ok(goal) = self
|
||||
.state_dbs
|
||||
.thread_goals()
|
||||
.get_thread_goal(runtime.thread_id())
|
||||
.await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if let Some(goal) = goal
|
||||
&& matches!(
|
||||
goal.status,
|
||||
codex_state::ThreadGoalStatus::Active
|
||||
| codex_state::ThreadGoalStatus::BudgetLimited
|
||||
)
|
||||
{
|
||||
accounting.mark_turn_goal_active(input.turn_id, goal.goal_id);
|
||||
}
|
||||
let accounting = runtime.accounting_state();
|
||||
accounting.start_turn(
|
||||
input.turn_id,
|
||||
input.collaboration_mode.mode,
|
||||
input.token_usage_at_turn_start,
|
||||
);
|
||||
if matches!(
|
||||
input.collaboration_mode.mode,
|
||||
codex_protocol::config_types::ModeKind::Plan
|
||||
) {
|
||||
accounting.clear_current_turn_goal();
|
||||
return;
|
||||
}
|
||||
let Ok(goal) = self
|
||||
.state_dbs
|
||||
.thread_goals()
|
||||
.get_thread_goal(runtime.thread_id())
|
||||
.await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if let Some(goal) = goal
|
||||
&& matches!(
|
||||
goal.status,
|
||||
codex_state::ThreadGoalStatus::Active
|
||||
| codex_state::ThreadGoalStatus::BudgetLimited
|
||||
)
|
||||
{
|
||||
accounting.mark_turn_goal_active(input.turn_id, goal.goal_id);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn on_turn_stop(&self, input: TurnStopInput<'_>) {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let turn_id = input.turn_store.level_id();
|
||||
if let Err(err) = runtime
|
||||
.account_active_goal_progress(
|
||||
turn_id,
|
||||
&format!("{turn_id}:turn-stop"),
|
||||
codex_state::GoalAccountingMode::ActiveOnly,
|
||||
BudgetLimitedGoalDisposition::ClearActive,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"failed to account active goal progress at turn stop for {turn_id}: {err}"
|
||||
);
|
||||
return;
|
||||
}
|
||||
runtime.accounting_state().finish_turn(turn_id);
|
||||
let turn_id = input.turn_store.level_id();
|
||||
if let Err(err) = runtime
|
||||
.account_active_goal_progress(
|
||||
turn_id,
|
||||
&format!("{turn_id}:turn-stop"),
|
||||
codex_state::GoalAccountingMode::ActiveOnly,
|
||||
BudgetLimitedGoalDisposition::ClearActive,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"failed to account active goal progress at turn stop for {turn_id}: {err}"
|
||||
);
|
||||
return;
|
||||
}
|
||||
runtime.accounting_state().finish_turn(turn_id);
|
||||
})
|
||||
}
|
||||
|
||||
async fn on_turn_abort(&self, input: TurnAbortInput<'_>) {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let turn_id = input.turn_store.level_id();
|
||||
if let Err(err) = runtime
|
||||
.account_active_goal_progress(
|
||||
turn_id,
|
||||
&format!("{turn_id}:turn-abort"),
|
||||
codex_state::GoalAccountingMode::ActiveOnly,
|
||||
BudgetLimitedGoalDisposition::ClearActive,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"failed to account active goal progress after turn abort for {turn_id}: {err}"
|
||||
);
|
||||
return;
|
||||
}
|
||||
runtime.accounting_state().finish_turn(turn_id);
|
||||
let turn_id = input.turn_store.level_id();
|
||||
if let Err(err) = runtime
|
||||
.account_active_goal_progress(
|
||||
turn_id,
|
||||
&format!("{turn_id}:turn-abort"),
|
||||
codex_state::GoalAccountingMode::ActiveOnly,
|
||||
BudgetLimitedGoalDisposition::ClearActive,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"failed to account active goal progress after turn abort for {turn_id}: {err}"
|
||||
);
|
||||
return;
|
||||
}
|
||||
runtime.accounting_state().finish_turn(turn_id);
|
||||
})
|
||||
}
|
||||
|
||||
async fn on_turn_error(&self, input: TurnErrorInput<'_>) {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
fn on_turn_error<'a>(&'a self, input: TurnErrorInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
|
||||
return;
|
||||
};
|
||||
|
||||
let reason = match input.error {
|
||||
CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit,
|
||||
// The turn has ended because the error was non-retryable or its
|
||||
// retries were exhausted. Block the goal to prevent automatic
|
||||
// continuation from looping and consuming tokens, as can happen
|
||||
// with compaction errors.
|
||||
_ => ActiveGoalStopReason::TurnError,
|
||||
};
|
||||
if let Err(err) = runtime
|
||||
.stop_active_goal_for_turn(input.turn_id, reason)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
error = ?input.error,
|
||||
"failed to stop active goal after turn error: {err}"
|
||||
);
|
||||
}
|
||||
let reason = match input.error {
|
||||
CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit,
|
||||
// The turn has ended because the error was non-retryable or its
|
||||
// retries were exhausted. Block the goal to prevent automatic
|
||||
// continuation from looping and consuming tokens, as can happen
|
||||
// with compaction errors.
|
||||
_ => ActiveGoalStopReason::TurnError,
|
||||
};
|
||||
if let Err(err) = runtime
|
||||
.stop_active_goal_for_turn(input.turn_id, reason)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
error = ?input.error,
|
||||
"failed to stop active goal after turn error: {err}"
|
||||
);
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<C> TokenUsageContributor for GoalExtension<C>
|
||||
where
|
||||
C: Send + Sync + 'static,
|
||||
{
|
||||
async fn on_token_usage(
|
||||
&self,
|
||||
_session_store: &ExtensionData,
|
||||
thread_store: &ExtensionData,
|
||||
turn_store: &ExtensionData,
|
||||
token_usage: &TokenUsageInfo,
|
||||
) {
|
||||
let Some(runtime) = goal_runtime_handle(thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
fn on_token_usage<'a>(
|
||||
&'a self,
|
||||
_session_store: &'a ExtensionData,
|
||||
thread_store: &'a ExtensionData,
|
||||
turn_store: &'a ExtensionData,
|
||||
token_usage: &'a TokenUsageInfo,
|
||||
) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(runtime) = goal_runtime_handle(thread_store) else {
|
||||
return;
|
||||
};
|
||||
if !runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(_recorded) = runtime
|
||||
.accounting_state()
|
||||
.record_token_usage(turn_store.level_id(), &token_usage.total_token_usage)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some(_recorded) = runtime
|
||||
.accounting_state()
|
||||
.record_token_usage(turn_store.level_id(), &token_usage.total_token_usage)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user