Attribute turn metrics to the models used during the turn (#44656)

## Why Turn completion metrics used session-level telemetry, which could label usage with the wrong model after a model switch. A turn can also include compaction and responses from different models, so a single token total cannot attribute usage correctly. ## What changed - Accumulate token usage by model using each response's resolved step settings, and emit one histogram sample per model and token type per turn. - Preserve zero-valued samples when no usage is reported, without adding a sample for an unused selected model when only compaction ran. - Use turn telemetry for tool-call, memory, network-proxy, and running-process metrics. ## Testing Add regression tests for model switches, step-level model changes, compaction followed by a response or an aborted turn, and turns without responses. Verify repeated requests aggregate into one sample per model and earlier session usage is excluded. GitOrigin-RevId: 03ab1dcde7da0b16b3a500a9a05e4c86583ee585
2026-10-30 19:18:49 +00:00 · 2026-09-10 21:58:58 +00:00
parent 8570091e14
commit c8a8295e79
8 changed files with 352 additions and 46 deletions
@@ -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<Self>,
+        _session: Arc<Session>,
+        _turn_context: Arc<TurnContext>,
+        _input: Vec<TurnInput>,
+        _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<Str
    }
 }

+#[derive(Clone, Copy)]
+enum UsageScenario {
+    ResponseOnly,
+    ResponseAfterStepSwitch,
+    CompactThenResponse,
+    CompactThenStop,
+    NoResponse,
+}
+
+#[test_case::test_case(UsageScenario::ResponseOnly; "response_only")]
+#[test_case::test_case(UsageScenario::ResponseAfterStepSwitch; "response_after_step_switch")]
+#[test_case::test_case(UsageScenario::CompactThenResponse; "compact_then_response")]
+#[test_case::test_case(UsageScenario::CompactThenStop; "compact_then_stop")]
+#[test_case::test_case(UsageScenario::NoResponse; "no_response")]
+#[tokio::test]
+async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) {
+    let metrics = MetricsClient::new(
+        MetricsConfig::in_memory(
+            "test",
+            "codex-core",
+            env!("CARGO_PKG_VERSION"),
+            InMemoryMetricExporter::default(),
+        )
+        .with_runtime_reader(),
+    )
+    .expect("in-memory metrics client");
+    let (mut session, mut turn_context, _receiver) = make_session_and_context_with_rx().await;
+    let session_telemetry = session
+        .services
+        .session_telemetry
+        .clone()
+        .with_metrics(metrics.clone());
+    Arc::get_mut(&mut session)
+        .expect("session should be uniquely owned")
+        .services
+        .session_telemetry = session_telemetry.clone();
+    Arc::get_mut(&mut turn_context)
+        .expect("turn context should be uniquely owned")
+        .session_telemetry = session_telemetry;
+
+    let next_model = if turn_context.model_info().slug == "gpt-5.4" {
+        "gpt-5.2"
+    } else {
+        "gpt-5.4"
+    };
+    let previous_context = turn_context;
+    let turn_context = Arc::new(
+        previous_context
+            .with_model(next_model.to_string(), &session.services.models_manager)
+            .await,
+    );
+    // Earlier session usage must not leak into this turn's per-model samples.
+    session
+        .update_token_usage_info(
+            &previous_context,
+            Some(&TokenUsage {
+                input_tokens: 1_000,
+                total_tokens: 1_000,
+                ..TokenUsage::default()
+            }),
+        )
+        .await
+        .expect("earlier session usage should be recorded");
+    session
+        .spawn_task(Arc::clone(&turn_context), Vec::new(), PendingTask)
+        .await;
+    let mut expected_usage = Vec::new();
+    if matches!(
+        scenario,
+        UsageScenario::CompactThenResponse | UsageScenario::CompactThenStop
+    ) {
+        let compaction_usage = TokenUsage {
+            input_tokens: 100,
+            cached_input_tokens: 30,
+            cache_write_input_tokens: 20,
+            output_tokens: 70,
+            reasoning_output_tokens: 40,
+            total_tokens: 170,
+            codex_rollout_budget_units: None,
+        };
+        // Local pre-turn compaction records usage with the previous model's context.
+        session
+            .update_token_usage_info(&previous_context, Some(&compaction_usage))
+            .await
+            .expect("compaction usage should be recorded");
+        expected_usage.push((previous_context.model_info().slug.clone(), compaction_usage));
+    }
+    if matches!(
+        scenario,
+        UsageScenario::ResponseOnly
+            | UsageScenario::ResponseAfterStepSwitch
+            | UsageScenario::CompactThenResponse
+    ) {
+        let response_context = if matches!(scenario, UsageScenario::ResponseAfterStepSwitch) {
+            &previous_context
+        } else {
+            &turn_context
+        };
+        let response_usage = TokenUsage {
+            input_tokens: 10,
+            cached_input_tokens: 3,
+            cache_write_input_tokens: 2,
+            output_tokens: 7,
+            reasoning_output_tokens: 4,
+            total_tokens: 17,
+            codex_rollout_budget_units: None,
+        };
+        // Multiple requests for one model must still produce one turn histogram sample.
+        for _ in 0..2 {
+            session
+                .record_token_usage_info(
+                    &turn_context,
+                    &response_context.initial_settings,
+                    Some(&response_usage),
+                )
+                .await
+                .expect("response usage should be recorded");
+        }
+        let mut total_usage = response_usage.clone();
+        total_usage.add_assign(&response_usage);
+        expected_usage.push((response_context.model_info().slug.clone(), total_usage));
+    }
+    if matches!(scenario, UsageScenario::NoResponse) {
+        expected_usage.push((next_model.to_string(), TokenUsage::default()));
+    }
+
+    let task_result = if matches!(scenario, UsageScenario::CompactThenStop) {
+        // A PostCompact hook can stop the turn before the new model samples.
+        Err(CodexErr::TurnAborted)
+    } else {
+        Ok(None)
+    };
+    session.on_task_finished(turn_context, task_result).await;
+
+    let snapshot = metrics.snapshot().expect("runtime metrics snapshot");
+    let token_usage_metric = find_metric(&snapshot, TURN_TOKEN_USAGE_METRIC);
+    let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = token_usage_metric.data() else {
+        panic!("expected token usage histogram");
+    };
+    let mut token_usage_points = histogram
+        .data_points()
+        .map(|point| {
+            let mut attributes = attributes_to_map(point.attributes());
+            (
+                (
+                    attributes
+                        .remove("token_type")
+                        .expect("token usage metric should include a token type"),
+                    attributes
+                        .remove("model")
+                        .expect("token usage metric should include a model"),
+                ),
+                point.count(),
+                point.sum(),
+            )
+        })
+        .collect::<Vec<_>>();
+    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::<Vec<_>>();
+    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::<Vec<_>>();
+    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::<BTreeMap<_, _>>();
+    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();