diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d67e79f79..315dc93c6 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2946,12 +2946,14 @@ impl Session { }; if let Some(token_info) = token_info.as_ref() { for contributor in self.services.extensions.token_usage_contributors() { - contributor.on_token_usage( - &self.services.session_extension_data, - &self.services.thread_extension_data, - turn_context.extension_data.as_ref(), - token_info, - ); + contributor + .on_token_usage( + &self.services.session_extension_data, + &self.services.thread_extension_data, + turn_context.extension_data.as_ref(), + token_info, + ) + .await; } } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index b0bc6330d..4511df458 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1874,8 +1874,9 @@ async fn record_token_usage_info_notifies_extension_contributors() { records: Arc>>, } + #[async_trait::async_trait] impl codex_extension_api::TokenUsageContributor for TokenUsageRecorder { - fn on_token_usage( + async fn on_token_usage( &self, session_store: &codex_extension_api::ExtensionData, thread_store: &codex_extension_api::ExtensionData, diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index b03846fd3..4aa28d104 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -88,9 +88,10 @@ pub trait ConfigContributor: Send + Sync { /// Implementations should keep this callback cheap. The host calls it after /// updating cached token usage and before emitting the corresponding client /// token-count notification. +#[async_trait::async_trait] pub trait TokenUsageContributor: Send + Sync { /// Called each time the host records token usage from a model response. - fn on_token_usage( + async fn on_token_usage( &self, _session_store: &ExtensionData, _thread_store: &ExtensionData, diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 97d715729..3204b9eb9 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -179,11 +179,12 @@ where } } +#[async_trait] impl TokenUsageContributor for GoalExtension where C: Send + Sync + 'static, { - fn on_token_usage( + async fn on_token_usage( &self, _session_store: &ExtensionData, thread_store: &ExtensionData, @@ -204,8 +205,6 @@ where // this recorded delta can be committed to the active persisted goal. // It also needs an event/input capability to emit ThreadGoalUpdated and // inject budget-limit steering when accounting changes goal status. - // TODO: if the storage/event path must await, TokenUsageContributor - // either needs to become async or receive a fire-and-forget host sink. } }