measure rollout load and ThreadItem projection sizes

This commit is contained in:
Tom Wiltzius
2026-06-24 19:52:03 -07:00
parent 96d8e34712
commit ccc788debf
7 changed files with 615 additions and 1 deletions
@@ -429,6 +429,7 @@ use codex_protocol::protocol::W3cTraceContext;
use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS;
use codex_protocol::user_input::UserInput as CoreInputItem;
use codex_rmcp_client::perform_oauth_login_return_url;
use codex_rollout::RolloutProjectionTelemetry;
use codex_rollout::is_persisted_rollout_item;
use codex_rollout::state_db::StateDbHandle;
use codex_rollout::state_db::reconcile_rollout;
@@ -619,3 +620,40 @@ pub(crate) fn build_api_turns_from_rollout_items(items: &[RolloutItem]) -> Vec<T
}
builder.finish()
}
fn record_rollout_projection(thread_id: ThreadId, turns: &[Turn]) {
let item_count = turns
.iter()
.map(|turn| turn.items.len() as u64)
.sum::<u64>();
RolloutProjectionTelemetry::new(thread_id).record_projected_turns(
turns,
turns.len() as u64,
item_count,
turns
.iter()
.filter(|turn| is_completed_user_assistant_turn(turn))
.map(|turn| (&turn.items, turn.items.len() as u64)),
);
}
fn is_completed_user_assistant_turn(turn: &Turn) -> bool {
turn.status == TurnStatus::Completed
&& turn
.items
.iter()
.any(|item| matches!(item, ThreadItem::UserMessage { .. }))
&& turn.items.iter().any(|item| {
matches!(
item,
ThreadItem::AgentMessage {
phase: None | Some(codex_protocol::models::MessagePhase::FinalAnswer),
..
}
)
})
}
#[cfg(test)]
#[path = "request_processors_projection_metrics_tests.rs"]
mod projection_metrics_tests;
@@ -2263,6 +2263,7 @@ impl ThreadRequestProcessor {
thread_from_stored_thread(stored_thread, fallback_provider, &self.config.cwd);
if include_turns && let Some(history) = history {
thread.turns = build_api_turns_from_rollout_items(&history.items);
record_rollout_projection(thread_id, &thread.turns);
}
Ok(Some(thread))
}
@@ -0,0 +1,55 @@
use codex_protocol::protocol::AgentMessageEvent;
use codex_protocol::protocol::AgentReasoningRawContentEvent;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::TurnStartedEvent;
use codex_protocol::protocol::UserMessageEvent;
use pretty_assertions::assert_eq;
use super::build_api_turns_from_rollout_items;
use super::is_completed_user_assistant_turn;
#[test]
fn cold_projection_coalesces_repeated_rollout_updates() {
let rollout_items = vec![
RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent {
turn_id: "turn-1".to_string(),
trace_id: None,
started_at: None,
model_context_window: None,
collaboration_mode_kind: Default::default(),
})),
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: "question".to_string(),
..Default::default()
})),
RolloutItem::EventMsg(EventMsg::AgentReasoningRawContent(
AgentReasoningRawContentEvent {
text: "first".to_string(),
},
)),
RolloutItem::EventMsg(EventMsg::AgentReasoningRawContent(
AgentReasoningRawContentEvent {
text: "second".to_string(),
},
)),
RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent {
message: "answer".to_string(),
phase: None,
memory_citation: None,
})),
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: "turn-1".to_string(),
last_agent_message: Some("answer".to_string()),
completed_at: None,
duration_ms: None,
time_to_first_token_ms: None,
})),
];
let turns = build_api_turns_from_rollout_items(&rollout_items);
assert_eq!(turns.len(), 1);
assert_eq!(turns[0].items.len(), 3);
assert!(is_completed_user_assistant_turn(&turns[0]));
}
+1
View File
@@ -59,6 +59,7 @@ pub use list::rollout_date_parts;
pub use metadata::builder_from_items;
pub use persistence_metrics::RolloutPersistenceBatchMeasurement;
pub use persistence_metrics::RolloutPersistenceTelemetry;
pub use persistence_metrics::RolloutProjectionTelemetry;
pub use persistence_metrics::measure_and_filter_rollout_items;
pub use policy::is_persisted_rollout_item;
pub use policy::persisted_rollout_items;
+321 -1
View File
@@ -1,5 +1,7 @@
use std::collections::HashSet;
use std::io::Write;
use std::sync::Arc;
use std::sync::LazyLock;
use std::sync::Mutex;
use codex_otel::MetricsClient;
@@ -8,15 +10,29 @@ use codex_protocol::items::TurnItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use serde::Serialize;
use crate::policy::is_persisted_rollout_item;
const ITEM_BYTES_METRIC: &str = "codex.rollout.persistence.item_bytes";
const APPEND_METRIC: &str = "codex.rollout.persistence.append";
const TURN_BYTES_METRIC: &str = "codex.rollout.persistence.turn_bytes";
const LOAD_THREAD_BYTES_METRIC: &str = "codex.rollout.persistence.load_thread_bytes";
const LOAD_THREAD_ITEMS_METRIC: &str = "codex.rollout.persistence.load_thread_items";
const LOAD_TURN_BYTES_METRIC: &str = "codex.rollout.persistence.load_turn_bytes";
const LOAD_TURN_ITEMS_METRIC: &str = "codex.rollout.persistence.load_turn_items";
const PROJECTION_THREAD_BYTES_METRIC: &str = "codex.rollout.persistence.projection_thread_bytes";
const PROJECTION_THREAD_ITEMS_METRIC: &str = "codex.rollout.persistence.projection_thread_items";
const PROJECTION_THREAD_TURNS_METRIC: &str = "codex.rollout.persistence.projection_thread_turns";
const PROJECTION_TURN_BYTES_METRIC: &str = "codex.rollout.persistence.projection_turn_bytes";
const PROJECTION_TURN_ITEMS_METRIC: &str = "codex.rollout.persistence.projection_turn_items";
const MEASUREMENT_ERROR_METRIC: &str = "codex.rollout.persistence.measurement_error";
const SAMPLE_DENOMINATOR: u64 = 100;
const SAMPLE_RATE_LABEL: &str = "0.01";
static MEASURED_ROLLOUT_LOAD_THREADS: LazyLock<Mutex<HashSet<ThreadId>>> =
LazyLock::new(|| Mutex::new(HashSet::new()));
static MEASURED_PROJECTED_THREADS: LazyLock<Mutex<HashSet<ThreadId>>> =
LazyLock::new(|| Mutex::new(HashSet::new()));
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PersistenceDecision {
@@ -200,7 +216,7 @@ fn add_item_to_turn(totals: &mut TurnSizeTotals, item: &RolloutItemMeasurement)
}
}
fn serialized_len(item: &RolloutItem) -> serde_json::Result<u64> {
fn serialized_len(item: &(impl Serialize + ?Sized)) -> serde_json::Result<u64> {
let mut writer = CountingWriter::default();
serde_json::to_writer(&mut writer, item)?;
Ok(writer.bytes)
@@ -391,6 +407,310 @@ impl RolloutPersistenceTelemetry {
}
}
/// Records detached per-thread and per-turn size proxies for rollout loads.
pub(crate) struct RolloutLoadTelemetry {
metrics: Option<MetricsClient>,
sampled: bool,
thread_id: ThreadId,
}
impl RolloutLoadTelemetry {
pub fn new(thread_id: ThreadId) -> Self {
let metrics = codex_otel::global();
let sampled = metrics.is_some() && is_thread_sampled(thread_id);
Self {
metrics,
sampled,
thread_id,
}
}
/// Records bytes for successfully loaded decompressed JSONL records, excluding newlines.
pub fn record_rollout_load(
&self,
totals: RolloutSizeTotals,
turn_totals: &[RolloutSizeTotals],
) {
let Some(metrics) = self.enabled_metrics() else {
return;
};
if !mark_thread_measured(&MEASURED_ROLLOUT_LOAD_THREADS, self.thread_id) {
return;
}
record_size_totals(
metrics,
totals,
LOAD_THREAD_BYTES_METRIC,
LOAD_THREAD_ITEMS_METRIC,
"rollout_jsonl",
"rollout_line_json_v1",
"rollout_load",
);
record_turn_size_totals(
metrics,
turn_totals,
LOAD_TURN_BYTES_METRIC,
LOAD_TURN_ITEMS_METRIC,
"rollout_jsonl",
"rollout_line_json_v1",
"rollout_load",
);
}
fn enabled_metrics(&self) -> Option<&MetricsClient> {
self.sampled.then_some(self.metrics.as_ref()).flatten()
}
}
/// Records detached size proxies for app-server ThreadItem projection reads.
pub struct RolloutProjectionTelemetry {
metrics: Option<MetricsClient>,
sampled: bool,
thread_id: ThreadId,
}
impl RolloutProjectionTelemetry {
pub fn new(thread_id: ThreadId) -> Self {
let metrics = codex_otel::global();
let sampled = metrics.is_some() && is_thread_sampled(thread_id);
Self {
metrics,
sampled,
thread_id,
}
}
/// Records compact JSON bytes for the final materialized app-server turns.
pub fn record_projected_turns<'a, T, U>(
&self,
turns: &T,
turn_count: u64,
item_count: u64,
forecast_turn_items: impl IntoIterator<Item = (&'a U, u64)>,
) where
T: Serialize + ?Sized,
U: Serialize + 'a,
{
let Some(metrics) = self.enabled_metrics() else {
return;
};
if MEASURED_PROJECTED_THREADS
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.contains(&self.thread_id)
{
return;
}
let payload_bytes = match serialized_len(turns) {
Ok(payload_bytes) => payload_bytes,
Err(_) => {
let _ = metrics.counter(
MEASUREMENT_ERROR_METRIC,
/*inc*/ 1,
&[
("phase", "projection_serialize"),
("representation", "thread_turns"),
("source", "app_server_thread_read"),
],
);
return;
}
};
let turn_totals = match forecast_turn_items
.into_iter()
.map(|(items, item_count)| {
serialized_len(items).map(|payload_bytes| RolloutSizeTotals {
items: item_count,
payload_bytes,
})
})
.collect::<serde_json::Result<Vec<_>>>()
{
Ok(turn_totals) => turn_totals,
Err(_) => {
let _ = metrics.counter(
MEASUREMENT_ERROR_METRIC,
/*inc*/ 1,
&[
("phase", "projection_turn_serialize"),
("representation", "thread_turns"),
("source", "app_server_thread_read"),
],
);
return;
}
};
if !mark_thread_measured(&MEASURED_PROJECTED_THREADS, self.thread_id) {
return;
}
record_size_totals(
metrics,
RolloutSizeTotals {
items: item_count,
payload_bytes,
},
PROJECTION_THREAD_BYTES_METRIC,
PROJECTION_THREAD_ITEMS_METRIC,
"thread_turns",
"app_server_turns_json_v2",
"app_server_thread_read",
);
record_turn_size_totals(
metrics,
&turn_totals,
PROJECTION_TURN_BYTES_METRIC,
PROJECTION_TURN_ITEMS_METRIC,
"thread_turns",
"app_server_turn_items_json_v2",
"app_server_thread_read",
);
let _ = metrics.histogram(
PROJECTION_THREAD_TURNS_METRIC,
saturating_i64(turn_count),
&[
("representation", "thread_turns"),
("encoding", "app_server_turns_json_v2"),
("sample_rate", SAMPLE_RATE_LABEL),
("source", "app_server_thread_read"),
],
);
}
fn enabled_metrics(&self) -> Option<&MetricsClient> {
self.sampled.then_some(self.metrics.as_ref()).flatten()
}
}
#[derive(Default)]
pub(crate) struct RolloutTurnSizeTracker {
current: Option<PendingRolloutTurn>,
completed: Vec<RolloutSizeTotals>,
}
#[derive(Default)]
struct PendingRolloutTurn {
totals: RolloutSizeTotals,
explicit: bool,
saw_user_message: bool,
saw_final_agent_message: bool,
}
impl RolloutTurnSizeTracker {
pub(crate) fn observe(&mut self, item: &RolloutItem, payload_bytes: u64) {
match item {
RolloutItem::EventMsg(EventMsg::TurnStarted(_)) => {
self.finish_implicit_turn();
self.current = Some(PendingRolloutTurn {
explicit: true,
..Default::default()
});
}
RolloutItem::EventMsg(EventMsg::UserMessage(_)) => {
if self
.current
.as_ref()
.is_some_and(|turn| !turn.explicit && turn.saw_user_message)
{
self.finish_implicit_turn();
}
self.current
.get_or_insert_with(PendingRolloutTurn::default)
.saw_user_message = true;
}
_ => {}
}
let Some(turn) = self.current.as_mut() else {
return;
};
turn.totals.items = turn.totals.items.saturating_add(1);
turn.totals.payload_bytes = turn.totals.payload_bytes.saturating_add(payload_bytes);
match item {
RolloutItem::EventMsg(EventMsg::AgentMessage(message))
if !matches!(
message.phase,
Some(codex_protocol::models::MessagePhase::Commentary)
) =>
{
turn.saw_final_agent_message = true;
}
RolloutItem::EventMsg(EventMsg::TurnComplete(_)) => self.finish_current_turn(),
RolloutItem::EventMsg(EventMsg::TurnAborted(_)) => {
self.current = None;
}
_ => {}
}
}
pub(crate) fn finish(mut self) -> Vec<RolloutSizeTotals> {
self.finish_implicit_turn();
self.completed
}
fn finish_implicit_turn(&mut self) {
if self.current.as_ref().is_some_and(|turn| !turn.explicit) {
self.finish_current_turn();
}
}
fn finish_current_turn(&mut self) {
if let Some(turn) = self.current.take()
&& turn.saw_user_message
&& turn.saw_final_agent_message
{
self.completed.push(turn.totals);
}
}
}
fn record_size_totals(
metrics: &MetricsClient,
totals: RolloutSizeTotals,
bytes_metric: &'static str,
items_metric: &'static str,
representation: &str,
encoding: &str,
source: &str,
) {
let tags = [
("representation", representation),
("encoding", encoding),
("sample_rate", SAMPLE_RATE_LABEL),
("source", source),
];
let _ = metrics.histogram(bytes_metric, saturating_i64(totals.payload_bytes), &tags);
let _ = metrics.histogram(items_metric, saturating_i64(totals.items), &tags);
}
fn record_turn_size_totals(
metrics: &MetricsClient,
turn_totals: &[RolloutSizeTotals],
bytes_metric: &'static str,
items_metric: &'static str,
representation: &str,
encoding: &str,
source: &str,
) {
let tags = [
("representation", representation),
("encoding", encoding),
("sample_rate", SAMPLE_RATE_LABEL),
("source", source),
];
for totals in turn_totals {
let _ = metrics.histogram(bytes_metric, saturating_i64(totals.payload_bytes), &tags);
let _ = metrics.histogram(items_metric, saturating_i64(totals.items), &tags);
}
}
fn mark_thread_measured(measured_threads: &Mutex<HashSet<ThreadId>>, thread_id: ThreadId) -> bool {
measured_threads
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(thread_id)
}
fn saturating_i64(value: u64) -> i64 {
value.try_into().unwrap_or(i64::MAX)
}
@@ -1,8 +1,13 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use codex_protocol::ThreadId;
use codex_protocol::items::TurnItem;
use codex_protocol::items::UserMessageItem;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::AgentMessageEvent;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ItemCompletedEvent;
use codex_protocol::protocol::RolloutItem;
@@ -10,14 +15,20 @@ use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::TurnStartedEvent;
use codex_protocol::protocol::UserMessageEvent;
use pretty_assertions::assert_eq;
use serde::Serialize;
use super::CompletedTurnMeasurement;
use super::RolloutProjectionTelemetry;
use super::RolloutSizeTotals;
use super::RolloutTurnSizeTracker;
use super::TurnMeasurementState;
use super::TurnOutcome;
use super::TurnSizeTotals;
use super::is_thread_sampled;
use super::measure_and_filter_rollout_items;
use super::serialized_len;
use super::update_turn_measurements;
fn retained_message(text: &str) -> RolloutItem {
@@ -258,3 +269,174 @@ fn filtered_item_completion_includes_its_nested_item_type() {
super::PersistenceDecision::Dropped
);
}
#[test]
fn projected_turn_size_matches_compact_json() {
let turns = vec![
serde_json::json!({"id": "turn-1", "items": [{"type": "userMessage"}]}),
serde_json::json!({"id": "turn-2", "items": []}),
];
assert_eq!(
serialized_len(&turns).expect("measure projected turns"),
serde_json::to_vec(&turns)
.expect("serialize projected turns")
.len() as u64
);
}
#[test]
fn rollout_turn_sizes_use_loaded_line_bytes_for_completed_user_turns() {
let items = [
RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent {
turn_id: "turn-1".to_string(),
trace_id: None,
started_at: None,
model_context_window: None,
collaboration_mode_kind: Default::default(),
})),
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: "question".to_string(),
..Default::default()
})),
RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent {
message: "answer".to_string(),
phase: None,
memory_citation: None,
})),
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: "turn-1".to_string(),
last_agent_message: Some("answer".to_string()),
completed_at: None,
duration_ms: None,
time_to_first_token_ms: None,
})),
];
let mut tracker = RolloutTurnSizeTracker::default();
for (index, item) in items.iter().enumerate() {
tracker.observe(item, (index + 1) as u64);
}
assert_eq!(
tracker.finish(),
vec![RolloutSizeTotals {
items: 4,
payload_bytes: 10,
}]
);
}
#[test]
fn rollout_turn_sizes_support_legacy_implicit_turns() {
let user = || {
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: "question".to_string(),
..Default::default()
}))
};
let answer = || {
RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent {
message: "answer".to_string(),
phase: None,
memory_citation: None,
}))
};
let mut tracker = RolloutTurnSizeTracker::default();
for item in [user(), answer(), user(), answer()] {
tracker.observe(&item, 5);
}
assert_eq!(
tracker.finish(),
vec![
RolloutSizeTotals {
items: 2,
payload_bytes: 10,
},
RolloutSizeTotals {
items: 2,
payload_bytes: 10,
},
]
);
}
#[test]
fn rollout_turn_sizes_exclude_incomplete_and_commentary_only_turns() {
let mut tracker = RolloutTurnSizeTracker::default();
let start = |turn_id: &str| {
RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent {
turn_id: turn_id.to_string(),
trace_id: None,
started_at: None,
model_context_window: None,
collaboration_mode_kind: Default::default(),
}))
};
let user = || {
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: "question".to_string(),
..Default::default()
}))
};
let commentary = RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent {
message: "working".to_string(),
phase: Some(codex_protocol::models::MessagePhase::Commentary),
memory_citation: None,
}));
let complete = RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: "turn-1".to_string(),
last_agent_message: None,
completed_at: None,
duration_ms: None,
time_to_first_token_ms: None,
}));
for item in [start("turn-1"), user(), commentary, complete] {
tracker.observe(&item, 1);
}
for item in [start("turn-2"), user()] {
tracker.observe(&item, 1);
}
assert_eq!(tracker.finish(), Vec::<RolloutSizeTotals>::new());
}
struct FailingSerialize {
serialization_attempted: Arc<AtomicBool>,
}
impl Serialize for FailingSerialize {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.serialization_attempted.store(true, Ordering::Relaxed);
Err(serde::ser::Error::custom(
"intentional serialization failure",
))
}
}
#[test]
fn exporter_disabled_path_does_not_serialize_projected_turns() {
let sampled_thread_id = (0..10_000_u128)
.find_map(|value| {
let thread_id = ThreadId::from_string(&format!("00000000-0000-0000-0000-{value:012x}"))
.expect("valid thread id");
is_thread_sampled(thread_id).then_some(thread_id)
})
.expect("sampled thread id");
let serialization_attempted = Arc::new(AtomicBool::new(false));
let turns = FailingSerialize {
serialization_attempted: Arc::clone(&serialization_attempted),
};
RolloutProjectionTelemetry::new(sampled_thread_id).record_projected_turns(
&turns,
/*turn_count*/ 1,
/*item_count*/ 1,
std::iter::empty::<(&serde_json::Value, u64)>(),
);
assert!(!serialization_attempted.load(Ordering::Relaxed));
}
+17
View File
@@ -44,6 +44,9 @@ use super::list::get_threads_in_root;
use super::list::parse_cursor;
use super::list::parse_timestamp_uuid_from_filename;
use super::metadata;
use super::persistence_metrics::RolloutLoadTelemetry;
use super::persistence_metrics::RolloutSizeTotals;
use super::persistence_metrics::RolloutTurnSizeTracker;
use super::session_index::find_thread_names_by_ids;
use crate::config::RolloutConfigView;
use crate::state_db;
@@ -903,6 +906,8 @@ impl RolloutRecorder {
let mut items: Vec<RolloutItem> = Vec::new();
let mut thread_id: Option<ThreadId> = None;
let mut parse_errors = 0usize;
let mut loaded_jsonl_bytes = 0_u64;
let mut loaded_turns = RolloutTurnSizeTracker::default();
let mut reader = compression::open_rollout_line_reader(path).await?;
let mut saw_non_empty_line = false;
while let Some(line) = reader.next_line().await? {
@@ -934,6 +939,9 @@ impl RolloutRecorder {
{
thread_id = Some(session_meta_line.meta.id);
}
let line_bytes = line.len() as u64;
loaded_jsonl_bytes = loaded_jsonl_bytes.saturating_add(line_bytes);
loaded_turns.observe(&item, line_bytes);
items.push(item);
}
Err(e) => {
@@ -952,6 +960,15 @@ impl RolloutRecorder {
thread_id,
parse_errors,
);
if let Some(thread_id) = thread_id {
RolloutLoadTelemetry::new(thread_id).record_rollout_load(
RolloutSizeTotals {
items: items.len() as u64,
payload_bytes: loaded_jsonl_bytes,
},
&loaded_turns.finish(),
);
}
Ok((items, thread_id, parse_errors))
}