From 426b81a41bbb2cbcaa7e0eca3bf5d270fe69ac55 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Thu, 2 Apr 2026 18:52:32 -0700 Subject: [PATCH] [codex-analytics] add token usage metadata --- .../analytics/src/analytics_client_tests.rs | 72 +++++++++++++++++++ codex-rs/analytics/src/events.rs | 5 ++ codex-rs/analytics/src/reducer.rs | 34 +++++++++ .../app-server/src/bespoke_event_handling.rs | 18 ++++- .../app-server/tests/suite/v2/turn_start.rs | 5 ++ 5 files changed, 132 insertions(+), 2 deletions(-) diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index d9a080aa1b..3a0230113c 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -46,6 +46,9 @@ use codex_app_server_protocol::Thread; use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::ThreadStatus as AppServerThreadStatus; +use codex_app_server_protocol::ThreadTokenUsage; +use codex_app_server_protocol::ThreadTokenUsageUpdatedNotification; +use codex_app_server_protocol::TokenUsageBreakdown; use codex_app_server_protocol::Turn; use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnError as AppServerTurnError; @@ -173,6 +176,33 @@ fn sample_turn_started_notification(thread_id: &str, turn_id: &str) -> ServerNot }) } +fn sample_thread_token_usage_updated_notification( + thread_id: &str, + turn_id: &str, +) -> ServerNotification { + ServerNotification::ThreadTokenUsageUpdated(ThreadTokenUsageUpdatedNotification { + thread_id: thread_id.to_string(), + turn_id: turn_id.to_string(), + token_usage: ThreadTokenUsage { + total: TokenUsageBreakdown { + total_tokens: 500, + input_tokens: 200, + cached_input_tokens: 50, + output_tokens: 220, + reasoning_output_tokens: 30, + }, + last: TokenUsageBreakdown { + total_tokens: 321, + input_tokens: 123, + cached_input_tokens: 45, + output_tokens: 140, + reasoning_output_tokens: 13, + }, + model_context_window: Some(200_000), + }, + }) +} + fn sample_turn_completed_notification( thread_id: &str, turn_id: &str, @@ -223,6 +253,7 @@ async fn ingest_turn_prerequisites( include_initialize: bool, include_resolved_config: bool, include_started: bool, + include_token_usage: bool, ) { if include_initialize { reducer @@ -292,6 +323,17 @@ async fn ingest_turn_prerequisites( ) .await; } + + if include_token_usage { + reducer + .ingest( + AnalyticsFact::Notification(Box::new( + sample_thread_token_usage_updated_notification("thread-2", "turn-2"), + )), + out, + ) + .await; + } } fn expected_absolute_path(path: &PathBuf) -> String { @@ -887,6 +929,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), created_at: Some(455), completed_at: Some(456), @@ -928,6 +975,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, "created_at": 455, "completed_at": 456 @@ -947,6 +999,7 @@ async fn turn_lifecycle_emits_turn_event() { /*include_initialize*/ true, /*include_resolved_config*/ true, /*include_started*/ true, + /*include_token_usage*/ true, ) .await; reducer @@ -975,6 +1028,14 @@ async fn turn_lifecycle_emits_turn_event() { assert_eq!(payload["event_params"]["created_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] @@ -988,6 +1049,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 @@ -1017,6 +1079,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 @@ -1044,6 +1107,7 @@ async fn turn_completed_without_started_notification_emits_null_created_at() { /*include_initialize*/ true, /*include_resolved_config*/ true, /*include_started*/ false, + /*include_token_usage*/ false, ) .await; reducer @@ -1061,6 +1125,14 @@ async fn turn_completed_without_started_notification_emits_null_created_at() { let payload = serde_json::to_value(&out[0]).expect("serialize turn event"); assert_eq!(payload["event_params"]["created_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 { diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index fcb3d4d09f..3ae13b0a53 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -154,6 +154,11 @@ pub(crate) struct CodexTurnEventParams { pub(crate) subagent_tool_call_count: Option, pub(crate) web_search_count: Option, pub(crate) image_generation_count: Option, + pub(crate) input_tokens: Option, + pub(crate) cached_input_tokens: Option, + pub(crate) output_tokens: Option, + pub(crate) reasoning_output_tokens: Option, + pub(crate) total_tokens: Option, pub(crate) duration_ms: Option, pub(crate) created_at: Option, pub(crate) completed_at: Option, diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index 930ae8bc58..c32d043c90 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -34,6 +34,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; @@ -83,6 +84,7 @@ struct TurnState { num_input_images: Option, resolved_config: Option, created_at: Option, + token_usage: Option, completed: Option, } @@ -206,6 +208,7 @@ impl AnalyticsReducer { num_input_images: None, resolved_config: None, created_at: None, + token_usage: None, completed: None, }); turn_state.thread_id = Some(thread_id); @@ -358,6 +361,7 @@ impl AnalyticsReducer { num_input_images: None, resolved_config: None, created_at: None, + token_usage: None, completed: None, }); turn_state.connection_id = Some(connection_id); @@ -382,10 +386,23 @@ impl AnalyticsReducer { num_input_images: None, resolved_config: None, created_at: None, + token_usage: None, completed: None, }); turn_state.created_at = u64::try_from(notification.created_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, + created_at: None, + token_usage: None, + completed: None, + }); + turn_state.token_usage = Some(notification.token_usage.last); + } ServerNotification::TurnCompleted(notification) => { let turn_state = self.turns @@ -396,6 +413,7 @@ impl AnalyticsReducer { num_input_images: None, resolved_config: None, created_at: None, + token_usage: None, completed: None, }); turn_state.completed = Some(CompletedTurnState { @@ -509,6 +527,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, @@ -540,6 +559,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, created_at, completed_at: Some(completed.completed_at), diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index e7918324ea..18372b40da 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -1364,8 +1364,14 @@ pub(crate) async fn apply_bespoke_event_handling( .await; } EventMsg::TokenCount(token_count_event) => { - handle_token_count_event(conversation_id, event_turn_id, token_count_event, &outgoing) - .await; + handle_token_count_event( + conversation_id, + event_turn_id, + Some(&analytics_events_client), + token_count_event, + &outgoing, + ) + .await; } EventMsg::Error(ev) => { thread_watch_manager @@ -2208,6 +2214,7 @@ async fn handle_thread_rollback_failed( async fn handle_token_count_event( conversation_id: ThreadId, turn_id: String, + analytics_events_client: Option<&AnalyticsEventsClient>, token_count_event: TokenCountEvent, outgoing: &ThreadScopedOutgoingMessageSender, ) { @@ -2218,6 +2225,11 @@ async fn handle_token_count_event( turn_id, token_usage, }; + if let Some(analytics_events_client) = analytics_events_client { + analytics_events_client.track_notification( + ServerNotification::ThreadTokenUsageUpdated(notification.clone()), + ); + } outgoing .send_server_notification(ServerNotification::ThreadTokenUsageUpdated(notification)) .await; @@ -3550,6 +3562,7 @@ mod tests { handle_token_count_event( conversation_id, turn_id.clone(), + /*analytics_events_client*/ None, TokenCountEvent { info: Some(info), rate_limits: Some(rate_limits), @@ -3604,6 +3617,7 @@ mod tests { handle_token_count_event( conversation_id, turn_id.clone(), + /*analytics_events_client*/ None, TokenCountEvent { info: None, rate_limits: None, diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 5b612791d2..a8552b32d9 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -309,6 +309,11 @@ async fn turn_start_tracks_turn_event_analytics() -> Result<()> { assert!(event["event_params"]["created_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(()) }