[codex-analytics] add token usage metadata

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 {

codex-rs/analytics/src/events.rs

@@ -154,6 +154,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) created_at: Option<u64>,
     pub(crate) completed_at: Option<u64>,

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<usize>,
     resolved_config: Option<TurnResolvedConfigFact>,
     created_at: Option<u64>,
+    token_usage: Option<TokenUsageBreakdown>,
     completed: Option<CompletedTurnState>,
 }
 
@@ -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),

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
@@ -2216,6 +2222,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,
 ) {
@@ -2226,6 +2233,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;
@@ -3560,6 +3572,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),
@@ -3614,6 +3627,7 @@ mod tests {
         handle_token_count_event(
             conversation_id,
             turn_id.clone(),
+            /*analytics_events_client*/ None,
             TokenCountEvent {
                 info: None,
                 rate_limits: None,

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(())
 }