openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-13 11:47:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
Andrei Eternal
2026-09-10 21:45:08 +00:00
committed by copyberry
parent 8570091e14
commit c8a8295e79
8 changed files with 352 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
codex-rs/core/src/session/mod.rs
View File

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

12
codex-rs/core/src/session/tests.rs
View File

@@ -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");

6
codex-rs/core/src/session/turn.rs
View File

@@ -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;

2
codex-rs/core/src/state/mod.rs
View File

@@ -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;

2
codex-rs/core/src/state/turn.rs
View File

@@ -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<Arc<StepContext>>,

48
codex-rs/core/src/state/turn_token_usage.rs Normal file
View File

@@ -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<String, (SessionTelemetry, TokenUsage)>,
}
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::<Vec<_>>();
// 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],
);
}
}
}
}

57
codex-rs/core/src/tasks/mod.rs
View File

@@ -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),
&[],

257
codex-rs/core/src/tasks/mod_tests.rs
View File

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