diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4824e1c79d..777c7834ef 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -4533,7 +4533,7 @@ impl Session { token_usage: Option<&TokenUsage>, ) -> CodexResult<()> { let result = self - .record_token_usage_info(turn_context, token_usage) + .record_token_usage_info(turn_context, &turn_context.initial_settings, token_usage) .await; self.send_token_count_event(turn_context).await; result @@ -4576,6 +4576,7 @@ impl Session { pub(crate) async fn record_token_usage_info( &self, turn_context: &TurnContext, + settings: &ResolvedStepSettings, token_usage: Option<&TokenUsage>, ) -> CodexResult<()> { if let Some(token_usage) = token_usage { @@ -4591,6 +4592,17 @@ impl Session { } state.token_info() }; + let turn_state = self + .input_queue + .turn_state_for_sub_id(&self.active_turn, &turn_context.sub_id) + .await; + if let Some(turn_state) = turn_state { + turn_state.lock().await.token_usage_by_model.record( + settings.selected_collaboration_mode().model(), + settings.telemetry(&turn_context.session_telemetry), + token_usage, + ); + } let budget_result = self.record_rollout_budget_usage(token_usage); if let Some(token_info) = token_info.as_ref() { for contributor in self.services.extensions.token_usage_contributors() { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6efd613754..40ab7b5436 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3121,11 +3121,19 @@ async fn record_token_usage_info_notifies_extension_contributors() { }; session - .record_token_usage_info(&turn_context, Some(&first_usage)) + .record_token_usage_info( + &turn_context, + &turn_context.initial_settings, + Some(&first_usage), + ) .await .expect("first usage should be recorded"); session - .record_token_usage_info(&turn_context, Some(&second_usage)) + .record_token_usage_info( + &turn_context, + &turn_context.initial_settings, + Some(&second_usage), + ) .await .expect("second usage should be recorded"); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8b1089860c..4fb1aac62f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2811,7 +2811,11 @@ async fn try_run_sampling_request( ) .await; let budget_result = sess - .record_token_usage_info(&turn_context, token_usage.as_ref()) + .record_token_usage_info( + &turn_context, + &step_context.settings, + token_usage.as_ref(), + ) .await; should_emit_token_count = true; should_emit_turn_diff = true; diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index b5100bbcef..07ce5c366f 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -3,6 +3,7 @@ mod auto_compact_window; mod service; mod session; mod turn; +mod turn_token_usage; pub(crate) use crate::tools::ExecutedToolCalls; pub(crate) use additional_context::AdditionalContextStore; @@ -18,3 +19,4 @@ pub(crate) use turn::PendingRequestPermissions; pub(crate) use turn::RunningTask; pub(crate) use turn::TaskKind; pub(crate) use turn::TurnState; +pub(crate) use turn_token_usage::TurnTokenUsage; diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index b03dbf737f..324b39ae46 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -17,6 +17,7 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use rmcp::model::RequestId; use tokio::sync::oneshot; +use super::TurnTokenUsage; use crate::agent::control::AgentExecutionGuard; use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::session::TurnInputQueue; @@ -101,6 +102,7 @@ pub(crate) struct TurnState { pub(crate) tool_calls: u64, pub(crate) has_memory_citation: bool, pub(crate) token_usage_at_turn_start: TokenUsage, + pub(crate) token_usage_by_model: TurnTokenUsage, /// The last step captured for execution or selected from a speculative fallback. /// Remains absent until a step is captured; standalone local compaction has no step. pub(crate) last_known_step_context: Option>, diff --git a/codex-rs/core/src/state/turn_token_usage.rs b/codex-rs/core/src/state/turn_token_usage.rs new file mode 100644 index 0000000000..455c78ccd5 --- /dev/null +++ b/codex-rs/core/src/state/turn_token_usage.rs @@ -0,0 +1,48 @@ +//! Turn token metrics grouped by the model that actually produced each response. +//! Multiple responses from one model contribute a single histogram sample per turn. + +use std::collections::BTreeMap; + +use codex_otel::SessionTelemetry; +use codex_otel::TURN_TOKEN_USAGE_METRIC; +use codex_protocol::protocol::TokenUsage; + +#[derive(Default)] +pub(crate) struct TurnTokenUsage { + by_model: BTreeMap, +} + +impl TurnTokenUsage { + pub(crate) fn record(&mut self, model: &str, telemetry: SessionTelemetry, usage: &TokenUsage) { + let (_, total) = self + .by_model + .entry(model.to_owned()) + .or_insert_with(|| (telemetry, TokenUsage::default())); + total.add_assign(usage); + } + + pub(crate) fn emit(self, fallback: &SessionTelemetry, tmp_mem: (&str, &str)) { + let mut usage = self.by_model.into_values().collect::>(); + // Preserve zero-valued completion samples for turns without reported usage. + // Do not create a sample for the selected model if only compaction ran. + if usage.is_empty() { + usage.push((fallback.clone(), TokenUsage::default())); + } + for (telemetry, usage) in usage { + for (token_type, value) in [ + ("total", usage.total_tokens), + ("input", usage.input_tokens), + ("cached_input", usage.cached_input()), + ("cache_write_input", usage.cache_write_input_tokens), + ("output", usage.output_tokens), + ("reasoning_output", usage.reasoning_output_tokens), + ] { + telemetry.histogram( + TURN_TOKEN_USAGE_METRIC, + value.max(0), + &[("token_type", token_type), tmp_mem], + ); + } + } + } +} diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index f3033d488b..8891462dfc 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -44,7 +44,6 @@ use codex_otel::SessionTelemetry; use codex_otel::TURN_E2E_DURATION_METRIC; use codex_otel::TURN_MEMORY_METRIC; use codex_otel::TURN_NETWORK_PROXY_METRIC; -use codex_otel::TURN_TOKEN_USAGE_METRIC; use codex_otel::TURN_TOOL_CALL_METRIC; use codex_otel::TURN_UNIFIED_EXEC_RUNNING_PROCESSES_METRIC; use codex_protocol::models::ResponseItem; @@ -631,12 +630,18 @@ impl Session { .input_queue .take_pending_input_for_turn_state(turn_state.as_ref()) .await; - let (turn_had_memory_citation, turn_tool_calls, token_usage_at_turn_start) = { - let ts = turn_state.lock().await; + let ( + turn_had_memory_citation, + turn_tool_calls, + token_usage_at_turn_start, + token_usage_by_model, + ) = { + let mut ts = turn_state.lock().await; ( ts.has_memory_citation, ts.tool_calls, ts.token_usage_at_turn_start.clone(), + std::mem::take(&mut ts.token_usage_by_model), ) }; run_hooks_and_record_inputs( @@ -647,6 +652,7 @@ impl Session { PersistContext::Standard, ) .await; + let turn_telemetry = &turn_context.session_telemetry; // Emit token usage metrics. { // TODO(jif): drop this @@ -673,12 +679,8 @@ impl Session { } None => false, }; - emit_turn_network_proxy_metric( - &self.services.session_telemetry, - network_proxy_active, - tmp_mem, - ); - self.services.session_telemetry.histogram( + emit_turn_network_proxy_metric(turn_telemetry, network_proxy_active, tmp_mem); + turn_telemetry.histogram( TURN_TOOL_CALL_METRIC, i64::try_from(turn_tool_calls).unwrap_or(i64::MAX), &[tmp_mem], @@ -739,46 +741,17 @@ impl Session { .track_turn_token_usage(TurnTokenUsageFact { turn_id: turn_context.sub_id.clone(), thread_id: self.thread_id.to_string(), - token_usage: turn_token_usage.clone(), + token_usage: turn_token_usage, }); - self.services.session_telemetry.histogram( - TURN_TOKEN_USAGE_METRIC, - turn_token_usage.total_tokens, - &[("token_type", "total"), tmp_mem], - ); - self.services.session_telemetry.histogram( - TURN_TOKEN_USAGE_METRIC, - turn_token_usage.input_tokens, - &[("token_type", "input"), tmp_mem], - ); - self.services.session_telemetry.histogram( - TURN_TOKEN_USAGE_METRIC, - turn_token_usage.cached_input(), - &[("token_type", "cached_input"), tmp_mem], - ); - self.services.session_telemetry.histogram( - TURN_TOKEN_USAGE_METRIC, - turn_token_usage.cache_write_input_tokens, - &[("token_type", "cache_write_input"), tmp_mem], - ); - self.services.session_telemetry.histogram( - TURN_TOKEN_USAGE_METRIC, - turn_token_usage.output_tokens, - &[("token_type", "output"), tmp_mem], - ); - self.services.session_telemetry.histogram( - TURN_TOKEN_USAGE_METRIC, - turn_token_usage.reasoning_output_tokens, - &[("token_type", "reasoning_output"), tmp_mem], - ); + token_usage_by_model.emit(turn_telemetry, tmp_mem); } emit_turn_memory_metric( - &self.services.session_telemetry, + turn_telemetry, turn_context.config.features.enabled(Feature::MemoryTool), turn_context.config.memories.use_memories, turn_had_memory_citation, ); - self.services.session_telemetry.counter( + turn_telemetry.counter( TURN_UNIFIED_EXEC_RUNNING_PROCESSES_METRIC, i64::try_from(self.list_background_terminals().await.len()).unwrap_or(i64::MAX), &[], diff --git a/codex-rs/core/src/tasks/mod_tests.rs b/codex-rs/core/src/tasks/mod_tests.rs index e426b0b710..94c2373519 100644 --- a/codex-rs/core/src/tasks/mod_tests.rs +++ b/codex-rs/core/src/tasks/mod_tests.rs @@ -1,14 +1,26 @@ +use super::SessionTask; +use super::SessionTaskResult; use super::TASK_COMPACT_METRIC; use super::emit_compact_metric; use super::emit_turn_memory_metric; use super::emit_turn_network_proxy_metric; +use crate::session::TurnInput; +use crate::session::session::Session; +use crate::session::tests::make_session_and_context_with_rx; +use crate::session::turn_context::TurnContext; +use crate::state::TaskKind; use codex_otel::MetricsClient; use codex_otel::MetricsConfig; use codex_otel::SessionTelemetry; use codex_otel::TURN_MEMORY_METRIC; use codex_otel::TURN_NETWORK_PROXY_METRIC; +use codex_otel::TURN_TOKEN_USAGE_METRIC; +use codex_otel::TURN_TOOL_CALL_METRIC; +use codex_otel::TURN_UNIFIED_EXEC_RUNNING_PROCESSES_METRIC; use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TokenUsage; use opentelemetry::KeyValue; use opentelemetry_sdk::metrics::InMemoryMetricExporter; use opentelemetry_sdk::metrics::data::AggregatedMetrics; @@ -17,6 +29,30 @@ use opentelemetry_sdk::metrics::data::MetricData; use opentelemetry_sdk::metrics::data::ResourceMetrics; use pretty_assertions::assert_eq; use std::collections::BTreeMap; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; + +struct PendingTask; + +impl SessionTask for PendingTask { + fn kind(&self) -> TaskKind { + TaskKind::Regular + } + + fn span_name(&self) -> &'static str { + "session_task.pending" + } + + async fn run( + self: Arc, + _session: Arc, + _turn_context: Arc, + _input: Vec, + _cancellation_token: CancellationToken, + ) -> SessionTaskResult { + std::future::pending().await + } +} fn test_session_telemetry() -> SessionTelemetry { let exporter = InMemoryMetricExporter::default(); @@ -75,6 +111,227 @@ fn metric_point(resource_metrics: &ResourceMetrics, name: &str) -> (BTreeMap>(); + let mut expected_points = expected_usage + .into_iter() + .flat_map(|(model, usage)| { + [ + ("cache_write_input", usage.cache_write_input_tokens), + ("cached_input", usage.cached_input()), + ("input", usage.input_tokens), + ("output", usage.output_tokens), + ("reasoning_output", usage.reasoning_output_tokens), + ("total", usage.total_tokens), + ] + .map(|(token_type, value)| ((token_type.to_string(), model.clone()), 1, value as f64)) + }) + .collect::>(); + token_usage_points.sort_by(|left, right| left.0.cmp(&right.0)); + expected_points.sort_by(|left, right| left.0.cmp(&right.0)); + assert_eq!(token_usage_points, expected_points); + + let tool_call_metric = find_metric(&snapshot, TURN_TOOL_CALL_METRIC); + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = tool_call_metric.data() else { + panic!("expected tool call histogram"); + }; + let tool_call_models = histogram + .data_points() + .map(|point| { + attributes_to_map(point.attributes()) + .remove("model") + .expect("tool call metric should include a model") + }) + .collect::>(); + assert_eq!(tool_call_models, vec![next_model.to_string()]); + + let counter_models = [ + TURN_MEMORY_METRIC, + TURN_NETWORK_PROXY_METRIC, + TURN_UNIFIED_EXEC_RUNNING_PROCESSES_METRIC, + ] + .into_iter() + .map(|name| { + let (mut attributes, _value) = metric_point(&snapshot, name); + ( + name.to_string(), + attributes + .remove("model") + .expect("turn counter should include a model"), + ) + }) + .collect::>(); + assert_eq!( + counter_models, + BTreeMap::from([ + (TURN_MEMORY_METRIC.to_string(), next_model.to_string()), + ( + TURN_NETWORK_PROXY_METRIC.to_string(), + next_model.to_string(), + ), + ( + TURN_UNIFIED_EXEC_RUNNING_PROCESSES_METRIC.to_string(), + next_model.to_string(), + ), + ]) + ); +} + #[test] fn emit_turn_network_proxy_metric_records_active_turn() { let session_telemetry = test_session_telemetry();