From 083c1962f918b38d7d74f6c1ffec93483bf22062 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 13 May 2026 14:32:23 +0200 Subject: [PATCH] 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 | 19 ++- codex-rs/core/src/session/tests.rs | 112 ++++++++++++++++++ .../ext/extension-api/src/contributors.rs | 20 ++++ codex-rs/ext/extension-api/src/lib.rs | 1 + codex-rs/ext/extension-api/src/registry.rs | 15 +++ 5 files changed, 165 insertions(+), 2 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index fa63715272..770d356ef1 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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, + ); + } + } } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 5fcb652f8c..2964b0ad5f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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>>, + } + + 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::().is_some(), + saw_thread_store: thread_store.get::().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::::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::>(); + assert_eq!(expected, actual); +} + #[tokio::test] async fn record_initial_history_reconstructs_forked_transcript() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 0aee629a43..961b294263 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/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. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 57ecefd964..65e167c771 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/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; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index f6a4dd6667..f2818cacaa 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/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 { thread_lifecycle_contributors: Vec>>, turn_lifecycle_contributors: Vec>, + token_usage_contributors: Vec>, context_contributors: Vec>, tool_contributors: Vec>, turn_item_contributors: Vec>, @@ -24,6 +26,7 @@ impl Default for ExtensionRegistryBuilder { 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 ExtensionRegistryBuilder { self.turn_lifecycle_contributors.push(contributor); } + /// Registers one token-usage contributor. + pub fn token_usage_contributor(&mut self, contributor: Arc) { + self.token_usage_contributors.push(contributor); + } + /// Registers one prompt contributor. pub fn prompt_contributor(&mut self, contributor: Arc) { self.context_contributors.push(contributor); @@ -76,6 +84,7 @@ impl ExtensionRegistryBuilder { 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 ExtensionRegistryBuilder { pub struct ExtensionRegistry { thread_lifecycle_contributors: Vec>>, turn_lifecycle_contributors: Vec>, + token_usage_contributors: Vec>, context_contributors: Vec>, tool_contributors: Vec>, turn_item_contributors: Vec>, @@ -105,6 +115,11 @@ impl ExtensionRegistry { &self.turn_lifecycle_contributors } + /// Returns the registered token-usage contributors. + pub fn token_usage_contributors(&self) -> &[Arc] { + &self.token_usage_contributors + } + /// Claims the first rendered approval-review prompt accepted by an /// installed contributor. pub fn approval_review<'a>(