Compare commits

..

1 Commits

Author SHA1 Message Date
rhan-oai
44463cb72b [codex-analytics] add token usage metadata 2026-04-04 20:59:01 -07:00
5 changed files with 132 additions and 2 deletions

View File

@@ -48,6 +48,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;
@@ -181,6 +184,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,
@@ -232,6 +262,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 +332,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 {
@@ -1045,6 +1087,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),
@@ -1086,6 +1133,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
@@ -1105,6 +1157,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 +1186,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]
@@ -1146,6 +1207,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 +1237,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 +1265,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
@@ -1219,6 +1283,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 {

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

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>,
created_at: Option<u64>,
token_usage: Option<TokenUsageBreakdown>,
completed: Option<CompletedTurnState>,
}
@@ -221,6 +223,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);
@@ -373,6 +376,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);
@@ -397,6 +401,7 @@ impl AnalyticsReducer {
num_input_images: None,
resolved_config: None,
created_at: None,
token_usage: None,
completed: None,
});
turn_state.created_at = notification
@@ -404,6 +409,18 @@ impl AnalyticsReducer {
.created_at
.and_then(|created_at| u64::try_from(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
@@ -414,6 +431,7 @@ impl AnalyticsReducer {
num_input_images: None,
resolved_config: None,
created_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,
created_at,
completed_at: Some(completed.completed_at),

View File

@@ -1367,8 +1367,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
@@ -2217,6 +2223,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,
) {
@@ -2227,6 +2234,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;
@@ -3582,6 +3594,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),
@@ -3636,6 +3649,7 @@ mod tests {
handle_token_count_event(
conversation_id,
turn_id.clone(),
/*analytics_events_client*/ None,
TokenCountEvent {
info: None,
rate_limits: None,

View File

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