chore: make token usage async (#23305)

Make the `TokenUsageContributor` async. This will be required for future
extension and it's basically free

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;
                 }
             }
         }

codex-rs/core/src/session/tests.rs

@@ -1874,8 +1874,9 @@ 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 {
-        fn on_token_usage(
+        async fn on_token_usage(
             &self,
             session_store: &codex_extension_api::ExtensionData,
             thread_store: &codex_extension_api::ExtensionData,

codex-rs/ext/extension-api/src/contributors.rs

@@ -88,9 +88,10 @@ pub trait ConfigContributor<C>: 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,

codex-rs/ext/goal/src/extension.rs

@@ -179,11 +179,12 @@ where
     }
 }
 
+#[async_trait]
 impl<C> TokenUsageContributor for GoalExtension<C>
 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.
     }
 }