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:
jif-oai
2026-05-13 14:32:23 +02:00
committed by GitHub
parent fcc2a92743
commit 083c1962f9
5 changed files with 165 additions and 2 deletions
+17 -2
View File
@@ -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,
);
}
}
}
}