[codex-analytics] add token usage metadata

This commit is contained in:
rhan-oai
2026-04-07 17:17:39 -07:00
parent 6ccf770583
commit 32a0c52acd
8 changed files with 141 additions and 0 deletions

View File

@@ -30,6 +30,7 @@ use crate::facts::SubAgentThreadStartedInput;
use crate::facts::TrackEventsContext;
use crate::facts::TurnResolvedConfigFact;
use crate::facts::TurnStatus;
use crate::facts::TurnTokenUsageFact;
use crate::reducer::AnalyticsReducer;
use crate::reducer::normalize_path_for_skill_id;
use crate::reducer::skill_id_for_local_skill;
@@ -66,6 +67,7 @@ use codex_protocol::config_types::ModeKind;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::TokenUsage;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::HashSet;
@@ -181,6 +183,20 @@ fn sample_turn_started_notification(thread_id: &str, turn_id: &str) -> ServerNot
})
}
fn sample_turn_token_usage_fact(thread_id: &str, turn_id: &str) -> TurnTokenUsageFact {
TurnTokenUsageFact {
thread_id: thread_id.to_string(),
turn_id: turn_id.to_string(),
token_usage: TokenUsage {
total_tokens: 321,
input_tokens: 123,
cached_input_tokens: 45,
output_tokens: 140,
reasoning_output_tokens: 13,
},
}
}
fn sample_turn_completed_notification(
thread_id: &str,
turn_id: &str,
@@ -232,6 +248,7 @@ async fn ingest_turn_prerequisites(
include_initialize: bool,
include_resolved_config: bool,
include_started: bool,
include_token_usage: bool,
) {
if include_initialize {
reducer
@@ -301,6 +318,17 @@ async fn ingest_turn_prerequisites(
)
.await;
}
if include_token_usage {
reducer
.ingest(
AnalyticsFact::Custom(CustomAnalyticsFact::TurnTokenUsage(Box::new(
sample_turn_token_usage_fact("thread-2", "turn-2"),
))),
out,
)
.await;
}
}
fn expected_absolute_path(path: &PathBuf) -> String {
@@ -1045,6 +1073,11 @@ fn turn_event_serializes_expected_shape() {
subagent_tool_call_count: None,
web_search_count: None,
image_generation_count: None,
input_tokens: None,
cached_input_tokens: None,
output_tokens: None,
reasoning_output_tokens: None,
total_tokens: None,
duration_ms: Some(1234),
started_at: Some(455),
completed_at: Some(456),
@@ -1086,6 +1119,11 @@ fn turn_event_serializes_expected_shape() {
"subagent_tool_call_count": null,
"web_search_count": null,
"image_generation_count": null,
"input_tokens": null,
"cached_input_tokens": null,
"output_tokens": null,
"reasoning_output_tokens": null,
"total_tokens": null,
"duration_ms": 1234,
"started_at": 455,
"completed_at": 456
@@ -1105,6 +1143,7 @@ async fn turn_lifecycle_emits_turn_event() {
/*include_initialize*/ true,
/*include_resolved_config*/ true,
/*include_started*/ true,
/*include_token_usage*/ true,
)
.await;
reducer
@@ -1133,6 +1172,14 @@ async fn turn_lifecycle_emits_turn_event() {
assert_eq!(payload["event_params"]["started_at"], json!(455));
assert_eq!(payload["event_params"]["completed_at"], json!(456));
assert_eq!(payload["event_params"]["duration_ms"], json!(1234));
assert_eq!(payload["event_params"]["input_tokens"], json!(123));
assert_eq!(payload["event_params"]["cached_input_tokens"], json!(45));
assert_eq!(payload["event_params"]["output_tokens"], json!(140));
assert_eq!(
payload["event_params"]["reasoning_output_tokens"],
json!(13)
);
assert_eq!(payload["event_params"]["total_tokens"], json!(321));
}
#[tokio::test]
@@ -1146,6 +1193,7 @@ async fn turn_does_not_emit_without_required_prerequisites() {
/*include_initialize*/ false,
/*include_resolved_config*/ true,
/*include_started*/ false,
/*include_token_usage*/ false,
)
.await;
reducer
@@ -1175,6 +1223,7 @@ async fn turn_does_not_emit_without_required_prerequisites() {
/*include_initialize*/ true,
/*include_resolved_config*/ false,
/*include_started*/ false,
/*include_token_usage*/ false,
)
.await;
reducer
@@ -1202,6 +1251,7 @@ async fn turn_lifecycle_emits_failed_turn_event() {
/*include_initialize*/ true,
/*include_resolved_config*/ true,
/*include_started*/ true,
/*include_token_usage*/ false,
)
.await;
reducer
@@ -1233,6 +1283,7 @@ async fn turn_lifecycle_emits_interrupted_turn_event_without_error() {
/*include_initialize*/ true,
/*include_resolved_config*/ true,
/*include_started*/ true,
/*include_token_usage*/ false,
)
.await;
reducer
@@ -1264,6 +1315,7 @@ async fn turn_completed_without_started_notification_emits_null_started_at() {
/*include_initialize*/ true,
/*include_resolved_config*/ true,
/*include_started*/ false,
/*include_token_usage*/ false,
)
.await;
reducer
@@ -1281,6 +1333,14 @@ async fn turn_completed_without_started_notification_emits_null_started_at() {
let payload = serde_json::to_value(&out[0]).expect("serialize turn event");
assert_eq!(payload["event_params"]["started_at"], json!(null));
assert_eq!(payload["event_params"]["duration_ms"], json!(1234));
assert_eq!(payload["event_params"]["input_tokens"], json!(null));
assert_eq!(payload["event_params"]["cached_input_tokens"], json!(null));
assert_eq!(payload["event_params"]["output_tokens"], json!(null));
assert_eq!(
payload["event_params"]["reasoning_output_tokens"],
json!(null)
);
assert_eq!(payload["event_params"]["total_tokens"], json!(null));
}
fn sample_plugin_metadata() -> PluginTelemetryMetadata {

View File

@@ -14,6 +14,7 @@ use crate::facts::SkillInvokedInput;
use crate::facts::SubAgentThreadStartedInput;
use crate::facts::TrackEventsContext;
use crate::facts::TurnResolvedConfigFact;
use crate::facts::TurnTokenUsageFact;
use crate::reducer::AnalyticsReducer;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::ClientResponse;
@@ -196,6 +197,12 @@ impl AnalyticsEventsClient {
));
}
pub fn track_turn_token_usage(&self, fact: TurnTokenUsageFact) {
self.record_fact(AnalyticsFact::Custom(CustomAnalyticsFact::TurnTokenUsage(
Box::new(fact),
)));
}
pub fn track_plugin_installed(&self, plugin: PluginTelemetryMetadata) {
self.record_fact(AnalyticsFact::Custom(
CustomAnalyticsFact::PluginStateChanged(PluginStateChangedInput {

View File

@@ -156,6 +156,11 @@ pub(crate) struct CodexTurnEventParams {
pub(crate) subagent_tool_call_count: Option<usize>,
pub(crate) web_search_count: Option<usize>,
pub(crate) image_generation_count: Option<usize>,
pub(crate) input_tokens: Option<i64>,
pub(crate) cached_input_tokens: Option<i64>,
pub(crate) output_tokens: Option<i64>,
pub(crate) reasoning_output_tokens: Option<i64>,
pub(crate) total_tokens: Option<i64>,
pub(crate) duration_ms: Option<u64>,
pub(crate) started_at: Option<u64>,
pub(crate) completed_at: Option<u64>,

View File

@@ -16,6 +16,7 @@ use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SkillScope;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::TokenUsage;
use serde::Serialize;
use std::path::PathBuf;
@@ -65,6 +66,13 @@ pub struct TurnResolvedConfigFact {
pub is_first_turn: bool,
}
#[derive(Clone)]
pub struct TurnTokenUsageFact {
pub turn_id: String,
pub thread_id: String,
pub token_usage: TokenUsage,
}
#[derive(Clone, Copy, Debug, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum TurnStatus {
@@ -133,6 +141,7 @@ pub(crate) enum AnalyticsFact {
pub(crate) enum CustomAnalyticsFact {
SubAgentThreadStarted(SubAgentThreadStartedInput),
TurnResolvedConfig(Box<TurnResolvedConfigFact>),
TurnTokenUsage(Box<TurnTokenUsageFact>),
SkillInvoked(SkillInvokedInput),
AppMentioned(AppMentionedInput),
AppUsed(AppUsedInput),

View File

@@ -12,6 +12,7 @@ pub use facts::SubAgentThreadStartedInput;
pub use facts::TrackEventsContext;
pub use facts::TurnResolvedConfigFact;
pub use facts::TurnStatus;
pub use facts::TurnTokenUsageFact;
pub use facts::build_track_events_context;
#[cfg(test)]

View File

@@ -30,6 +30,7 @@ use crate::facts::SkillInvokedInput;
use crate::facts::SubAgentThreadStartedInput;
use crate::facts::TurnResolvedConfigFact;
use crate::facts::TurnStatus;
use crate::facts::TurnTokenUsageFact;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::ClientResponse;
use codex_app_server_protocol::CodexErrorInfo;
@@ -46,6 +47,7 @@ use codex_protocol::config_types::ReasoningSummary;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SkillScope;
use codex_protocol::protocol::TokenUsage;
use sha1::Digest;
use std::collections::HashMap;
use std::path::Path;
@@ -85,6 +87,7 @@ struct TurnState {
num_input_images: Option<usize>,
resolved_config: Option<TurnResolvedConfigFact>,
started_at: Option<u64>,
token_usage: Option<TokenUsage>,
completed: Option<CompletedTurnState>,
}
@@ -129,6 +132,9 @@ impl AnalyticsReducer {
CustomAnalyticsFact::TurnResolvedConfig(input) => {
self.ingest_turn_resolved_config(*input, out);
}
CustomAnalyticsFact::TurnTokenUsage(input) => {
self.ingest_turn_token_usage(*input, out);
}
CustomAnalyticsFact::SkillInvoked(input) => {
self.ingest_skill_invoked(input, out).await;
}
@@ -221,6 +227,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);
@@ -229,6 +236,26 @@ impl AnalyticsReducer {
self.maybe_emit_turn_event(&turn_id, out);
}
fn ingest_turn_token_usage(
&mut self,
input: TurnTokenUsageFact,
out: &mut Vec<TrackEventRequest>,
) {
let turn_id = input.turn_id.clone();
let turn_state = self.turns.entry(turn_id.clone()).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.thread_id = Some(input.thread_id);
turn_state.token_usage = Some(input.token_usage);
self.maybe_emit_turn_event(&turn_id, out);
}
async fn ingest_skill_invoked(
&mut self,
input: SkillInvokedInput,
@@ -373,6 +400,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 +425,7 @@ impl AnalyticsReducer {
num_input_images: None,
resolved_config: None,
started_at: None,
token_usage: None,
completed: None,
});
turn_state.started_at = notification
@@ -414,6 +443,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 +562,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 +594,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),

View File

@@ -310,6 +310,11 @@ async fn turn_start_tracks_turn_event_analytics() -> Result<()> {
assert!(event["event_params"]["started_at"].as_u64().is_some());
assert!(event["event_params"]["completed_at"].as_u64().is_some());
assert!(event["event_params"]["duration_ms"].as_u64().is_some());
assert_eq!(event["event_params"]["input_tokens"], 0);
assert_eq!(event["event_params"]["cached_input_tokens"], 0);
assert_eq!(event["event_params"]["output_tokens"], 0);
assert_eq!(event["event_params"]["reasoning_output_tokens"], 0);
assert_eq!(event["event_params"]["total_tokens"], 0);
Ok(())
}

View File

@@ -30,6 +30,7 @@ use crate::hook_runtime::record_pending_input;
use crate::state::ActiveTurn;
use crate::state::RunningTask;
use crate::state::TaskKind;
use codex_analytics::TurnTokenUsageFact;
use codex_login::AuthManager;
use codex_models_manager::manager::ModelsManager;
use codex_otel::SessionTelemetry;
@@ -485,6 +486,13 @@ impl Session {
- token_usage_at_turn_start.total_tokens)
.max(0),
};
self.services
.analytics_events_client
.track_turn_token_usage(TurnTokenUsageFact {
turn_id: turn_context.sub_id.clone(),
thread_id: self.conversation_id.to_string(),
token_usage: turn_token_usage.clone(),
});
self.services.session_telemetry.histogram(
TURN_TOKEN_USAGE_METRIC,
turn_token_usage.total_tokens,