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 authored and GitHub committed 2026-05-13 14:32:23 +02:00
1 parent fcc2a92743
commit 083c1962f9
5 files changed
+165 -2

No files matched your search

@@ -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.