mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00: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:
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.
|
||||
|
||||
Reference in new issue
Block a user