openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

codex: simplify thread read size metrics

Browse Source
This commit is contained in:
Tom Wiltzius
2026-06-24 20:05:25 -07:00
parent 7981353f76
commit 7b83805eab
10 changed files with 231 additions and 648 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
codex-rs/app-server/src/lib.rs
View File

@@ -110,6 +110,7 @@ mod request_processors;
mod request_serialization;
mod server_request_error;
mod skills_watcher;
mod thread_read_metrics;
mod thread_state;
mod thread_status;
mod transport;

42
codex-rs/app-server/src/request_processors.rs
View File

@@ -429,8 +429,6 @@ 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::RolloutProjectionMeasurement;
use codex_rollout::RolloutProjectionTelemetry;
use codex_rollout::is_persisted_rollout_item;
use codex_rollout::state_db::StateDbHandle;
use codex_rollout::state_db::reconcile_rollout;
@@ -621,43 +619,3 @@ pub(crate) fn build_api_turns_from_rollout_items(items: &[RolloutItem]) -> Vec<T
}
builder.finish()
}
fn prepare_rollout_projection_measurement(
thread_id: ThreadId,
turns: &[Turn],
) -> Option<RolloutProjectionMeasurement> {
let item_count = turns
.iter()
.map(|turn| turn.items.len() as u64)
.sum::<u64>();
RolloutProjectionTelemetry::new(thread_id).prepare_response_measurement(
turns.len() as u64,
item_count,
turns
.iter()
.filter(|turn| is_completed_user_assistant_turn(turn))
.map(|turn| turn.items.len() as u64)
.collect(),
)
}
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;

116
codex-rs/app-server/src/request_processors/thread_processor.rs
View File

@@ -1,5 +1,6 @@
use super::*;
use crate::error_code::method_not_found;
use crate::thread_read_metrics::ThreadReadMeasurement;
use codex_app_server_protocol::SelectedCapabilityRoot;
use codex_extension_api::ExtensionDataInit;
use codex_protocol::config_types::MultiAgentMode;
@@ -19,11 +20,6 @@ struct ThreadListFilters {
parent_thread_id: Option<ThreadId>,
}
struct ThreadReadView {
thread: Thread,
projection_measurement: Option<RolloutProjectionMeasurement>,
}
fn collect_resume_override_mismatches(
request: &ThreadResumeParams,
config_snapshot: &ThreadConfigSnapshot,
@@ -675,27 +671,22 @@ impl ThreadRequestProcessor {
request_id: &ConnectionRequestId,
params: ThreadReadParams,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
let view = self.thread_read_response_inner(params).await?;
let response = ThreadReadResponse {
thread: view.thread,
let response = self.thread_read_response_inner(params).await?;
let Some(measurement) = ThreadReadMeasurement::prepare(&response.thread) else {
return Ok(Some(response.into()));
};
if let Some(measurement) = view.projection_measurement {
let write_complete_rx = self
.outgoing
.send_response_with_write_complete(request_id.clone(), response)
.await;
tokio::spawn(async move {
if let Ok(write_complete) = write_complete_rx.await
&& let Some(serialized_bytes) = write_complete.serialized_bytes
{
measurement.record_serialized_response(serialized_bytes);
}
});
} else {
self.outgoing
.send_response(request_id.clone(), response)
.await;
}
let write_complete_rx = self
.outgoing
.send_response_with_write_complete(request_id.clone(), response)
.await;
tokio::spawn(async move {
if let Ok(write_complete) = write_complete_rx.await
&& let Some(serialized_bytes) = write_complete.serialized_bytes
{
measurement.record_serialized_response(serialized_bytes);
}
});
Ok(None)
}
@@ -2178,7 +2169,7 @@ impl ThreadRequestProcessor {
async fn thread_read_response_inner(
&self,
params: ThreadReadParams,
) -> Result<ThreadReadView, JSONRPCErrorError> {
) -> Result<ThreadReadResponse, JSONRPCErrorError> {
let ThreadReadParams {
thread_id,
include_turns,
@@ -2187,9 +2178,11 @@ impl ThreadRequestProcessor {
let thread_uuid = ThreadId::from_string(&thread_id)
.map_err(|err| invalid_request(format!("invalid thread id: {err}")))?;
self.read_thread_view(thread_uuid, include_turns)
let thread = self
.read_thread_view(thread_uuid, include_turns)
.await
.map_err(thread_read_view_error)
.map_err(thread_read_view_error)?;
Ok(ThreadReadResponse { thread })
}
/// Builds the API view for `thread/read` from persisted metadata plus optional live state.
@@ -2197,59 +2190,50 @@ impl ThreadRequestProcessor {
&self,
thread_id: ThreadId,
include_turns: bool,
) -> Result<ThreadReadView, ThreadReadViewError> {
) -> Result<Thread, ThreadReadViewError> {
let loaded_thread = self.thread_manager.get_thread(thread_id).await.ok();
let mut view = if include_turns {
let mut thread = if include_turns {
if let Some(loaded_thread) = loaded_thread.as_ref() {
// Loaded thread with turns: use persisted metadata when it exists,
// but reconstruct turns from the live ThreadStore history.
let persisted_thread = self
.load_persisted_thread_for_read(thread_id, /*include_turns*/ false)
.await?
.map(|view| view.thread);
ThreadReadView {
thread: self
.load_live_thread_view(
thread_id,
include_turns,
loaded_thread,
persisted_thread,
)
.await?,
projection_measurement: None,
}
} else if let Some(view) = self
.await?;
self.load_live_thread_view(
thread_id,
include_turns,
loaded_thread,
persisted_thread,
)
.await?
} else if let Some(thread) = self
.load_persisted_thread_for_read(thread_id, include_turns)
.await?
{
// Unloaded thread with turns: load metadata and history together
// from the ThreadStore.
view
thread
} else {
return Err(ThreadReadViewError::InvalidRequest(format!(
"thread not loaded: {thread_id}"
)));
}
} else if let Some(view) = self
} else if let Some(thread) = self
.load_persisted_thread_for_read(thread_id, include_turns)
.await?
{
// Persisted metadata-only read: no live thread state is needed.
view
thread
} else if let Some(loaded_thread) = loaded_thread.as_ref() {
// Loaded metadata-only read before persistence is materialized: build
// the response from the live thread snapshot.
ThreadReadView {
thread: self
.load_live_thread_view(
thread_id,
include_turns,
loaded_thread,
/*persisted_thread*/ None,
)
.await?,
projection_measurement: None,
}
self.load_live_thread_view(
thread_id,
include_turns,
loaded_thread,
/*persisted_thread*/ None,
)
.await?
} else {
return Err(ThreadReadViewError::InvalidRequest(format!(
"thread not loaded: {thread_id}"
@@ -2264,22 +2248,22 @@ impl ThreadRequestProcessor {
let thread_status = self
.thread_watch_manager
.loaded_status_for_thread(&view.thread.id)
.loaded_status_for_thread(&thread.id)
.await;
set_thread_status_and_interrupt_stale_turns(
&mut view.thread,
&mut thread,
thread_status,
has_live_in_progress_turn,
);
Ok(view)
Ok(thread)
}
async fn load_persisted_thread_for_read(
&self,
thread_id: ThreadId,
include_turns: bool,
) -> Result<Option<ThreadReadView>, ThreadReadViewError> {
) -> Result<Option<Thread>, ThreadReadViewError> {
let fallback_provider = self.config.model_provider_id.as_str();
match self
.thread_store
@@ -2293,16 +2277,10 @@ impl ThreadRequestProcessor {
Ok(stored_thread) => {
let (mut thread, history) =
thread_from_stored_thread(stored_thread, fallback_provider, &self.config.cwd);
let mut projection_measurement = None;
if include_turns && let Some(history) = history {
thread.turns = build_api_turns_from_rollout_items(&history.items);
projection_measurement =
prepare_rollout_projection_measurement(thread_id, &thread.turns);
}
Ok(Some(ThreadReadView {
thread,
projection_measurement,
}))
Ok(Some(thread))
}
Err(ThreadStoreError::InvalidRequest { message })
if message == format!("no rollout found for thread id {thread_id}") =>

55
codex-rs/app-server/src/request_processors_projection_metrics_tests.rs
View File

@@ -1,55 +0,0 @@
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]));
}

120
codex-rs/app-server/src/thread_read_metrics.rs Normal file
View File

@@ -0,0 +1,120 @@
use std::collections::HashSet;
use std::sync::LazyLock;
use std::sync::Mutex;
use codex_app_server_protocol::Thread;
use codex_app_server_protocol::ThreadItem;
use codex_app_server_protocol::ThreadStatus;
use codex_app_server_protocol::Turn;
use codex_app_server_protocol::TurnStatus;
use codex_otel::MetricsClient;
use codex_protocol::models::MessagePhase;
const RESPONSE_BYTES_METRIC: &str = "codex.app_server.thread_read.response_bytes";
const COMPLETED_TURNS_METRIC: &str = "codex.app_server.thread_read.completed_turns";
const COMPLETED_TURN_ITEMS_METRIC: &str = "codex.app_server.thread_read.completed_turn_items";
const SAMPLE_DENOMINATOR: u64 = 100;
const SAMPLE_RATE_LABEL: &str = "0.01";
static MEASURED_THREADS: LazyLock<Mutex<HashSet<String>>> =
LazyLock::new(|| Mutex::new(HashSet::new()));
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
struct ThreadReadCounts {
completed_turns: u64,
completed_turn_items: u64,
}
pub(crate) struct ThreadReadMeasurement {
metrics: MetricsClient,
counts: ThreadReadCounts,
}
impl ThreadReadMeasurement {
pub(crate) fn prepare(thread: &Thread) -> Option<Self> {
if !matches!(thread.status, ThreadStatus::NotLoaded) || !is_thread_sampled(&thread.id) {
return None;
}
let metrics = codex_otel::global()?;
let counts = count_completed_turns(&thread.turns);
if counts.completed_turns == 0
|| !MEASURED_THREADS
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(thread.id.clone())
{
return None;
}
Some(Self { metrics, counts })
}
pub(crate) fn record_serialized_response(self, response_bytes: usize) {
let tags = [
("encoding", "app_server_transport_json"),
("sample_rate", SAMPLE_RATE_LABEL),
("source", "cold_thread_read"),
];
let _ = self.metrics.histogram(
RESPONSE_BYTES_METRIC,
saturating_i64(response_bytes as u64),
&tags,
);
let _ = self.metrics.histogram(
COMPLETED_TURNS_METRIC,
saturating_i64(self.counts.completed_turns),
&tags,
);
let _ = self.metrics.histogram(
COMPLETED_TURN_ITEMS_METRIC,
saturating_i64(self.counts.completed_turn_items),
&tags,
);
}
}
fn count_completed_turns(turns: &[Turn]) -> ThreadReadCounts {
let mut counts = ThreadReadCounts::default();
for turn in turns
.iter()
.filter(|turn| is_completed_user_assistant_turn(turn))
{
counts.completed_turns = counts.completed_turns.saturating_add(1);
counts.completed_turn_items = counts
.completed_turn_items
.saturating_add(turn.items.len() as u64);
}
counts
}
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(MessagePhase::FinalAnswer),
..
}
)
})
}
fn is_thread_sampled(thread_id: &str) -> bool {
let hash = thread_id
.bytes()
.fold(0xcbf29ce484222325_u64, |hash, byte| {
(hash ^ u64::from(byte)).wrapping_mul(0x100000001b3)
});
hash % SAMPLE_DENOMINATOR == 0
}
fn saturating_i64(value: u64) -> i64 {
value.try_into().unwrap_or(i64::MAX)
}
#[cfg(test)]
#[path = "thread_read_metrics_tests.rs"]
mod tests;

62
codex-rs/app-server/src/thread_read_metrics_tests.rs Normal file
View File

@@ -0,0 +1,62 @@
use codex_app_server_protocol::ThreadItem;
use codex_app_server_protocol::Turn;
use codex_app_server_protocol::TurnItemsView;
use codex_app_server_protocol::TurnStatus;
use codex_protocol::models::MessagePhase;
use pretty_assertions::assert_eq;
use super::ThreadReadCounts;
use super::count_completed_turns;
fn turn(status: TurnStatus, items: Vec<ThreadItem>) -> Turn {
Turn {
id: "turn".to_string(),
items,
items_view: TurnItemsView::Full,
status,
error: None,
started_at: None,
completed_at: None,
duration_ms: None,
}
}
#[test]
fn counts_items_in_completed_user_assistant_turns() {
let user_message = || ThreadItem::UserMessage {
id: "user".to_string(),
client_id: None,
content: Vec::new(),
};
let agent_message = |phase| ThreadItem::AgentMessage {
id: "agent".to_string(),
text: "answer".to_string(),
phase,
memory_citation: None,
};
let turns = vec![
turn(
TurnStatus::Completed,
vec![user_message(), agent_message(None)],
),
turn(
TurnStatus::Completed,
vec![
user_message(),
agent_message(Some(MessagePhase::Commentary)),
],
),
turn(
TurnStatus::InProgress,
vec![user_message(), agent_message(None)],
),
];
assert_eq!(
count_completed_turns(&turns),
ThreadReadCounts {
completed_turns: 1,
completed_turn_items: 2,
}
);
}

2
codex-rs/rollout/src/lib.rs
View File

@@ -59,8 +59,6 @@ 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::RolloutProjectionMeasurement;
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;

328
codex-rs/rollout/src/persistence_metrics.rs
View File

@@ -1,7 +1,5 @@
use std::collections::HashSet;
use std::io::Write;
use std::sync::Arc;
use std::sync::LazyLock;
use std::sync::Mutex;
use codex_otel::MetricsClient;
@@ -10,34 +8,15 @@ 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 THREAD_READ_RESPONSE_BYTES_METRIC: &str =
"codex.rollout.persistence.thread_read_response_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_COMPLETED_TURNS_METRIC: &str =
"codex.rollout.persistence.projection_completed_turns";
const PROJECTION_COMPLETED_TURN_ITEMS_METRIC: &str =
"codex.rollout.persistence.projection_completed_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()));
static IN_FLIGHT_PROJECTED_THREADS: LazyLock<Mutex<HashSet<ThreadId>>> =
LazyLock::new(|| Mutex::new(HashSet::new()));
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PersistenceDecision {
@@ -221,7 +200,7 @@ fn add_item_to_turn(totals: &mut TurnSizeTotals, item: &RolloutItemMeasurement)
}
}
fn serialized_len(item: &(impl Serialize + ?Sized)) -> serde_json::Result<u64> {
fn serialized_len(item: &RolloutItem) -> serde_json::Result<u64> {
let mut writer = CountingWriter::default();
serde_json::to_writer(&mut writer, item)?;
Ok(writer.bytes)
@@ -412,311 +391,6 @@ 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()
}
}
/// Prepares 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,
}
}
/// Prepares counts to record with the existing serialized `thread/read` response.
pub fn prepare_response_measurement(
&self,
turn_count: u64,
item_count: u64,
completed_turn_item_counts: Vec<u64>,
) -> Option<RolloutProjectionMeasurement> {
let metrics = self.enabled_metrics()?.clone();
if MEASURED_PROJECTED_THREADS
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.contains(&self.thread_id)
{
return None;
}
if !IN_FLIGHT_PROJECTED_THREADS
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(self.thread_id)
{
return None;
}
Some(RolloutProjectionMeasurement {
metrics,
thread_id: self.thread_id,
turn_count,
item_count,
completed_turn_item_counts,
})
}
fn enabled_metrics(&self) -> Option<&MetricsClient> {
self.sampled.then_some(self.metrics.as_ref()).flatten()
}
}
/// Counts associated with one sampled, cold `thread/read` response.
pub struct RolloutProjectionMeasurement {
metrics: MetricsClient,
thread_id: ThreadId,
turn_count: u64,
item_count: u64,
completed_turn_item_counts: Vec<u64>,
}
impl RolloutProjectionMeasurement {
/// Records bytes already produced by the app-server transport serializer.
pub fn record_serialized_response(self, response_bytes: usize) {
if !mark_thread_measured(&MEASURED_PROJECTED_THREADS, self.thread_id) {
return;
}
let response_tags = [
("representation", "thread_read_response"),
("encoding", "app_server_transport_json"),
("sample_rate", SAMPLE_RATE_LABEL),
("source", "app_server_thread_read"),
];
let _ = self.metrics.histogram(
THREAD_READ_RESPONSE_BYTES_METRIC,
saturating_i64(response_bytes as u64),
&response_tags,
);
let projection_tags = [
("representation", "thread_items"),
("encoding", "app_server_protocol_v2"),
("sample_rate", SAMPLE_RATE_LABEL),
("source", "app_server_thread_read"),
];
let _ = self.metrics.histogram(
PROJECTION_THREAD_ITEMS_METRIC,
saturating_i64(self.item_count),
&projection_tags,
);
let _ = self.metrics.histogram(
PROJECTION_THREAD_TURNS_METRIC,
saturating_i64(self.turn_count),
&projection_tags,
);
let _ = self.metrics.histogram(
PROJECTION_COMPLETED_TURNS_METRIC,
saturating_i64(self.completed_turn_item_counts.len() as u64),
&projection_tags,
);
for item_count in &self.completed_turn_item_counts {
let _ = self.metrics.histogram(
PROJECTION_COMPLETED_TURN_ITEMS_METRIC,
saturating_i64(*item_count),
&projection_tags,
);
}
}
}
impl Drop for RolloutProjectionMeasurement {
fn drop(&mut self) {
IN_FLIGHT_PROJECTED_THREADS
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(&self.thread_id);
}
}
#[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)
}

136
codex-rs/rollout/src/persistence_metrics_tests.rs
View File

@@ -3,7 +3,6 @@ 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;
@@ -11,13 +10,9 @@ 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 super::CompletedTurnMeasurement;
use super::RolloutProjectionTelemetry;
use super::RolloutSizeTotals;
use super::RolloutTurnSizeTracker;
use super::TurnMeasurementState;
use super::TurnOutcome;
use super::TurnSizeTotals;
@@ -263,134 +258,3 @@ fn filtered_item_completion_includes_its_nested_item_type() {
super::PersistenceDecision::Dropped
);
}
#[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, /*payload_bytes*/ 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, /*payload_bytes*/ 1);
}
for item in [start("turn-2"), user()] {
tracker.observe(&item, /*payload_bytes*/ 1);
}
assert_eq!(tracker.finish(), Vec::<RolloutSizeTotals>::new());
}
#[test]
fn exporter_disabled_path_does_not_prepare_projection_measurement() {
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 measurement = RolloutProjectionTelemetry::new(sampled_thread_id)
.prepare_response_measurement(/*turn_count*/ 1, /*item_count*/ 1, vec![1]);
assert!(measurement.is_none());
}

17
codex-rs/rollout/src/recorder.rs
View File

@@ -44,9 +44,6 @@ 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;
@@ -906,8 +903,6 @@ 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? {
@@ -939,9 +934,6 @@ 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) => {
@@ -960,15 +952,6 @@ 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))
}