feat: add token usage contributor hook (#22485)

## Why

Extensions need a stable place to observe token accounting after Codex
folds model-provider usage into the session's cached `TokenUsageInfo`.
Without a contributor hook, extension-owned features that need last-turn
or cumulative token usage have to duplicate session plumbing or infer
state from client-facing `TokenCount` notifications.

## What changed

- Added `TokenUsageContributor` to `codex-extension-api`, passing
session/thread `ExtensionData`, `ThreadId`, turn id, and the current
`TokenUsageInfo`.
- Added registry builder/storage support for token-usage contributors.
- Invoked registered contributors from
`Session::record_token_usage_info` after the session token cache is
updated and before the client `TokenCount` notification is emitted.

## Testing

- Added `record_token_usage_info_notifies_extension_contributors`,
covering cumulative token usage updates and access to both extension
stores.

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

@@ -2858,8 +2858,23 @@ impl Session {
         token_usage: Option<&TokenUsage>,
     ) {
         if let Some(token_usage) = token_usage {
-            let mut state = self.state.lock().await;
-            state.update_token_info_from_usage(token_usage, turn_context.model_context_window());
+            let token_info = {
+                let mut state = self.state.lock().await;
+                state
+                    .update_token_info_from_usage(token_usage, turn_context.model_context_window());
+                state.token_info()
+            };
+            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,
+                        self.conversation_id,
+                        &turn_context.sub_id,
+                        token_info,
+                    );
+                }
+            }
         }
     }
 

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

@@ -1800,6 +1800,118 @@ async fn recompute_token_usage_updates_model_context_window() {
     assert_eq!(actual.model_context_window, Some(128_000));
 }
 
+#[tokio::test]
+async fn record_token_usage_info_notifies_extension_contributors() {
+    struct SessionTokenUsageMarker;
+    struct ThreadTokenUsageMarker;
+
+    #[derive(Debug, PartialEq, Eq)]
+    struct RecordedTokenUsage {
+        thread_id: ThreadId,
+        turn_id: String,
+        token_usage: TokenUsageInfo,
+        saw_session_store: bool,
+        saw_thread_store: bool,
+    }
+
+    struct TokenUsageRecorder {
+        records: Arc<std::sync::Mutex<Vec<RecordedTokenUsage>>>,
+    }
+
+    impl codex_extension_api::TokenUsageContributor for TokenUsageRecorder {
+        fn on_token_usage(
+            &self,
+            session_store: &codex_extension_api::ExtensionData,
+            thread_store: &codex_extension_api::ExtensionData,
+            thread_id: ThreadId,
+            turn_id: &str,
+            token_usage: &TokenUsageInfo,
+        ) {
+            self.records
+                .lock()
+                .expect("token usage records lock")
+                .push(RecordedTokenUsage {
+                    thread_id,
+                    turn_id: turn_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(),
+                });
+        }
+    }
+
+    let (mut session, turn_context) = make_session_and_context().await;
+    let records = Arc::new(std::sync::Mutex::new(Vec::new()));
+    let mut builder = codex_extension_api::ExtensionRegistryBuilder::<crate::config::Config>::new();
+    builder.token_usage_contributor(Arc::new(TokenUsageRecorder {
+        records: Arc::clone(&records),
+    }));
+    session.services.extensions = Arc::new(builder.build());
+    session
+        .services
+        .session_extension_data
+        .insert(SessionTokenUsageMarker);
+    session
+        .services
+        .thread_extension_data
+        .insert(ThreadTokenUsageMarker);
+
+    let first_usage = TokenUsage {
+        input_tokens: 10,
+        cached_input_tokens: 2,
+        output_tokens: 20,
+        reasoning_output_tokens: 3,
+        total_tokens: 33,
+    };
+    let second_usage = TokenUsage {
+        input_tokens: 7,
+        cached_input_tokens: 1,
+        output_tokens: 8,
+        reasoning_output_tokens: 5,
+        total_tokens: 20,
+    };
+
+    session
+        .record_token_usage_info(&turn_context, Some(&first_usage))
+        .await;
+    session
+        .record_token_usage_info(&turn_context, Some(&second_usage))
+        .await;
+
+    let mut expected_total_usage = first_usage.clone();
+    expected_total_usage.add_assign(&second_usage);
+    let expected = vec![
+        RecordedTokenUsage {
+            thread_id: session.conversation_id,
+            turn_id: turn_context.sub_id.clone(),
+            token_usage: TokenUsageInfo {
+                total_token_usage: first_usage.clone(),
+                last_token_usage: first_usage,
+                model_context_window: turn_context.model_context_window(),
+            },
+            saw_session_store: true,
+            saw_thread_store: true,
+        },
+        RecordedTokenUsage {
+            thread_id: session.conversation_id,
+            turn_id: turn_context.sub_id.clone(),
+            token_usage: TokenUsageInfo {
+                total_token_usage: expected_total_usage,
+                last_token_usage: second_usage,
+                model_context_window: turn_context.model_context_window(),
+            },
+            saw_session_store: true,
+            saw_thread_store: true,
+        },
+    ];
+    let actual = records
+        .lock()
+        .expect("token usage records lock")
+        .drain(..)
+        .collect::<Vec<_>>();
+    assert_eq!(expected, actual);
+}
+
 #[tokio::test]
 async fn record_initial_history_reconstructs_forked_transcript() {
     let (session, turn_context) = make_session_and_context().await;

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

@@ -1,8 +1,10 @@
 use std::future::Future;
 use std::sync::Arc;
 
+use codex_protocol::ThreadId;
 use codex_protocol::items::TurnItem;
 use codex_protocol::protocol::ReviewDecision;
+use codex_protocol::protocol::TokenUsageInfo;
 
 use crate::ExtensionData;
 
@@ -66,6 +68,24 @@ pub trait TurnLifecycleContributor: Send + Sync {
     fn on_turn_abort(&self, _input: TurnAbortInput<'_>) {}
 }
 
+/// Contributor for token usage checkpoints reported by the model provider.
+///
+/// Implementations should keep this callback cheap. The host calls it after
+/// updating cached token usage and before emitting the corresponding client
+/// token-count notification.
+pub trait TokenUsageContributor: Send + Sync {
+    /// Called each time the host records token usage from a model response.
+    fn on_token_usage(
+        &self,
+        _session_store: &ExtensionData,
+        _thread_store: &ExtensionData,
+        _thread_id: ThreadId,
+        _turn_id: &str,
+        _token_usage: &TokenUsageInfo,
+    ) {
+    }
+}
+
 /// Extension contribution that exposes native tools owned by a feature.
 pub trait ToolContributor: Send + Sync {
     /// Returns the native tools visible for the supplied extension stores.

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

@@ -25,6 +25,7 @@ pub use contributors::ThreadLifecycleContributor;
 pub use contributors::ThreadResumeInput;
 pub use contributors::ThreadStartInput;
 pub use contributors::ThreadStopInput;
+pub use contributors::TokenUsageContributor;
 pub use contributors::ToolContributor;
 pub use contributors::TurnAbortInput;
 pub use contributors::TurnItemContributionFuture;

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

@@ -5,6 +5,7 @@ use crate::ApprovalReviewFuture;
 use crate::ContextContributor;
 use crate::ExtensionData;
 use crate::ThreadLifecycleContributor;
+use crate::TokenUsageContributor;
 use crate::ToolContributor;
 use crate::TurnItemContributor;
 use crate::TurnLifecycleContributor;
@@ -13,6 +14,7 @@ use crate::TurnLifecycleContributor;
 pub struct ExtensionRegistryBuilder<C> {
     thread_lifecycle_contributors: Vec<Arc<dyn ThreadLifecycleContributor<C>>>,
     turn_lifecycle_contributors: Vec<Arc<dyn TurnLifecycleContributor>>,
+    token_usage_contributors: Vec<Arc<dyn TokenUsageContributor>>,
     context_contributors: Vec<Arc<dyn ContextContributor>>,
     tool_contributors: Vec<Arc<dyn ToolContributor>>,
     turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
@@ -24,6 +26,7 @@ impl<C> Default for ExtensionRegistryBuilder<C> {
         Self {
             thread_lifecycle_contributors: Vec::new(),
             turn_lifecycle_contributors: Vec::new(),
+            token_usage_contributors: Vec::new(),
             approval_review_contributors: Vec::new(),
             context_contributors: Vec::new(),
             tool_contributors: Vec::new(),
@@ -56,6 +59,11 @@ impl<C> ExtensionRegistryBuilder<C> {
         self.turn_lifecycle_contributors.push(contributor);
     }
 
+    /// Registers one token-usage contributor.
+    pub fn token_usage_contributor(&mut self, contributor: Arc<dyn TokenUsageContributor>) {
+        self.token_usage_contributors.push(contributor);
+    }
+
     /// Registers one prompt contributor.
     pub fn prompt_contributor(&mut self, contributor: Arc<dyn ContextContributor>) {
         self.context_contributors.push(contributor);
@@ -76,6 +84,7 @@ impl<C> ExtensionRegistryBuilder<C> {
         ExtensionRegistry {
             thread_lifecycle_contributors: self.thread_lifecycle_contributors,
             turn_lifecycle_contributors: self.turn_lifecycle_contributors,
+            token_usage_contributors: self.token_usage_contributors,
             approval_review_contributors: self.approval_review_contributors,
             context_contributors: self.context_contributors,
             tool_contributors: self.tool_contributors,
@@ -88,6 +97,7 @@ impl<C> ExtensionRegistryBuilder<C> {
 pub struct ExtensionRegistry<C> {
     thread_lifecycle_contributors: Vec<Arc<dyn ThreadLifecycleContributor<C>>>,
     turn_lifecycle_contributors: Vec<Arc<dyn TurnLifecycleContributor>>,
+    token_usage_contributors: Vec<Arc<dyn TokenUsageContributor>>,
     context_contributors: Vec<Arc<dyn ContextContributor>>,
     tool_contributors: Vec<Arc<dyn ToolContributor>>,
     turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
@@ -105,6 +115,11 @@ impl<C> ExtensionRegistry<C> {
         &self.turn_lifecycle_contributors
     }
 
+    /// Returns the registered token-usage contributors.
+    pub fn token_usage_contributors(&self) -> &[Arc<dyn TokenUsageContributor>] {
+        &self.token_usage_contributors
+    }
+
     /// Claims the first rendered approval-review prompt accepted by an
     /// installed contributor.
     pub fn approval_review<'a>(