mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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.
This commit is contained in:
@@ -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,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user