From 7b83805eabb94e5f8f130eb8a25d15feafeb5bd5 Mon Sep 17 00:00:00 2001 From: Tom Wiltzius Date: Wed, 24 Jun 2026 20:05:25 -0700 Subject: [PATCH] codex: simplify thread read size metrics --- codex-rs/app-server/src/lib.rs | 1 + codex-rs/app-server/src/request_processors.rs | 42 --- .../request_processors/thread_processor.rs | 116 +++---- ...est_processors_projection_metrics_tests.rs | 55 --- .../app-server/src/thread_read_metrics.rs | 120 +++++++ .../src/thread_read_metrics_tests.rs | 62 ++++ codex-rs/rollout/src/lib.rs | 2 - codex-rs/rollout/src/persistence_metrics.rs | 328 +----------------- .../rollout/src/persistence_metrics_tests.rs | 136 -------- codex-rs/rollout/src/recorder.rs | 17 - 10 files changed, 231 insertions(+), 648 deletions(-) delete mode 100644 codex-rs/app-server/src/request_processors_projection_metrics_tests.rs create mode 100644 codex-rs/app-server/src/thread_read_metrics.rs create mode 100644 codex-rs/app-server/src/thread_read_metrics_tests.rs diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index beed35848b..4d2791e254 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -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; diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index ef411b4a20..500ecc4734 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -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 Option { - let item_count = turns - .iter() - .map(|turn| turn.items.len() as u64) - .sum::(); - 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; diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index ec4131b378..b41fd6814c 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -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, } -struct ThreadReadView { - thread: Thread, - projection_measurement: Option, -} - fn collect_resume_override_mismatches( request: &ThreadResumeParams, config_snapshot: &ThreadConfigSnapshot, @@ -675,27 +671,22 @@ impl ThreadRequestProcessor { request_id: &ConnectionRequestId, params: ThreadReadParams, ) -> Result, 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 { + ) -> Result { 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 { + ) -> Result { 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, ThreadReadViewError> { + ) -> Result, 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}") => diff --git a/codex-rs/app-server/src/request_processors_projection_metrics_tests.rs b/codex-rs/app-server/src/request_processors_projection_metrics_tests.rs deleted file mode 100644 index 4475855782..0000000000 --- a/codex-rs/app-server/src/request_processors_projection_metrics_tests.rs +++ /dev/null @@ -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])); -} diff --git a/codex-rs/app-server/src/thread_read_metrics.rs b/codex-rs/app-server/src/thread_read_metrics.rs new file mode 100644 index 0000000000..947468223f --- /dev/null +++ b/codex-rs/app-server/src/thread_read_metrics.rs @@ -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>> = + 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 { + 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; diff --git a/codex-rs/app-server/src/thread_read_metrics_tests.rs b/codex-rs/app-server/src/thread_read_metrics_tests.rs new file mode 100644 index 0000000000..8e65bb9587 --- /dev/null +++ b/codex-rs/app-server/src/thread_read_metrics_tests.rs @@ -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) -> 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, + } + ); +} diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 8d96dd7257..3c103750c1 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -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; diff --git a/codex-rs/rollout/src/persistence_metrics.rs b/codex-rs/rollout/src/persistence_metrics.rs index f916b1f768..104d99fcf2 100644 --- a/codex-rs/rollout/src/persistence_metrics.rs +++ b/codex-rs/rollout/src/persistence_metrics.rs @@ -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>> = - LazyLock::new(|| Mutex::new(HashSet::new())); -static MEASURED_PROJECTED_THREADS: LazyLock>> = - LazyLock::new(|| Mutex::new(HashSet::new())); -static IN_FLIGHT_PROJECTED_THREADS: LazyLock>> = - 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 { +fn serialized_len(item: &RolloutItem) -> serde_json::Result { 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, - 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, - 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, - ) -> Option { - 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, -} - -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, - completed: Vec, -} - -#[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 { - 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>, 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) } diff --git a/codex-rs/rollout/src/persistence_metrics_tests.rs b/codex-rs/rollout/src/persistence_metrics_tests.rs index fc8b64834a..963805f43d 100644 --- a/codex-rs/rollout/src/persistence_metrics_tests.rs +++ b/codex-rs/rollout/src/persistence_metrics_tests.rs @@ -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::::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()); -} diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 46a68bdf36..90edb2c490 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -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 = Vec::new(); let mut thread_id: Option = 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)) }