[codex-analytics] add token usage metadata

This commit is contained in:
rhan-oai
2026-04-06 12:26:33 -07:00
parent 031eb880e7
commit 472c15a74b
5 changed files with 132 additions and 2 deletions

View File

@@ -36,6 +36,7 @@ use codex_app_server_protocol::CodexErrorInfo;
use codex_app_server_protocol::InitializeParams;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ServerNotification;
use codex_app_server_protocol::TokenUsageBreakdown;
use codex_app_server_protocol::UserInput;
use codex_git_utils::collect_git_info;
use codex_git_utils::get_git_repo_root;
@@ -85,6 +86,7 @@ struct TurnState {
num_input_images: Option<usize>,
resolved_config: Option<TurnResolvedConfigFact>,
started_at: Option<u64>,
token_usage: Option<TokenUsageBreakdown>,
completed: Option<CompletedTurnState>,
}
@@ -221,6 +223,7 @@ impl AnalyticsReducer {
num_input_images: None,
resolved_config: None,
started_at: None,
token_usage: None,
completed: None,
});
turn_state.thread_id = Some(thread_id);
@@ -373,6 +376,7 @@ impl AnalyticsReducer {
num_input_images: None,
resolved_config: None,
started_at: None,
token_usage: None,
completed: None,
});
turn_state.connection_id = Some(connection_id);
@@ -397,6 +401,7 @@ impl AnalyticsReducer {
num_input_images: None,
resolved_config: None,
started_at: None,
token_usage: None,
completed: None,
});
turn_state.started_at = notification
@@ -404,6 +409,18 @@ impl AnalyticsReducer {
.started_at
.and_then(|started_at| u64::try_from(started_at).ok());
}
ServerNotification::ThreadTokenUsageUpdated(notification) => {
let turn_state = self.turns.entry(notification.turn_id).or_insert(TurnState {
connection_id: None,
thread_id: None,
num_input_images: None,
resolved_config: None,
started_at: None,
token_usage: None,
completed: None,
});
turn_state.token_usage = Some(notification.token_usage.last);
}
ServerNotification::TurnCompleted(notification) => {
let turn_state =
self.turns
@@ -414,6 +431,7 @@ impl AnalyticsReducer {
num_input_images: None,
resolved_config: None,
started_at: None,
token_usage: None,
completed: None,
});
turn_state.completed = Some(CompletedTurnState {
@@ -532,6 +550,7 @@ fn codex_turn_event_params(
personality,
is_first_turn,
} = resolved_config;
let token_usage = turn_state.token_usage.clone();
CodexTurnEventParams {
thread_id,
turn_id,
@@ -563,6 +582,21 @@ fn codex_turn_event_params(
subagent_tool_call_count: None,
web_search_count: None,
image_generation_count: None,
input_tokens: token_usage
.as_ref()
.map(|token_usage| token_usage.input_tokens),
cached_input_tokens: token_usage
.as_ref()
.map(|token_usage| token_usage.cached_input_tokens),
output_tokens: token_usage
.as_ref()
.map(|token_usage| token_usage.output_tokens),
reasoning_output_tokens: token_usage
.as_ref()
.map(|token_usage| token_usage.reasoning_output_tokens),
total_tokens: token_usage
.as_ref()
.map(|token_usage| token_usage.total_tokens),
duration_ms: completed.duration_ms,
started_at,
completed_at: Some(completed.completed_at),