From 2018e6c904dec195bc517d3d0240e46131bd6b02 Mon Sep 17 00:00:00 2001 From: Tom Wiltzius Date: Tue, 23 Jun 2026 10:53:10 -0700 Subject: [PATCH] codex: measure thread reads from transport output --- codex-rs/app-server-transport/src/lib.rs | 1 + .../src/outgoing_message.rs | 9 +- .../src/transport/remote_control/mod.rs | 2 +- .../src/transport/remote_control/websocket.rs | 4 +- .../src/transport/stdio.rs | 6 +- .../src/transport/websocket.rs | 6 +- codex-rs/app-server/src/in_process.rs | 6 +- codex-rs/app-server/src/message_processor.rs | 2 +- codex-rs/app-server/src/outgoing_message.rs | 72 ++++++- codex-rs/app-server/src/request_processors.rs | 14 +- .../request_processors/thread_processor.rs | 107 +++++++---- codex-rs/app-server/src/transport.rs | 3 +- codex-rs/rollout/src/lib.rs | 1 + codex-rs/rollout/src/persistence_metrics.rs | 178 +++++++++--------- .../rollout/src/persistence_metrics_tests.rs | 54 +----- 15 files changed, 278 insertions(+), 187 deletions(-) diff --git a/codex-rs/app-server-transport/src/lib.rs b/codex-rs/app-server-transport/src/lib.rs index 9fb49f591d..43ece2b05b 100644 --- a/codex-rs/app-server-transport/src/lib.rs +++ b/codex-rs/app-server-transport/src/lib.rs @@ -5,6 +5,7 @@ pub use outgoing_message::ConnectionId; pub use outgoing_message::OutgoingError; pub use outgoing_message::OutgoingMessage; pub use outgoing_message::OutgoingResponse; +pub use outgoing_message::OutgoingWriteComplete; pub use outgoing_message::QueuedOutgoingMessage; pub use transport::AppServerStartupLock; pub use transport::AppServerTransport; diff --git a/codex-rs/app-server-transport/src/outgoing_message.rs b/codex-rs/app-server-transport/src/outgoing_message.rs index ff56b9fef9..8f8c0ab755 100644 --- a/codex-rs/app-server-transport/src/outgoing_message.rs +++ b/codex-rs/app-server-transport/src/outgoing_message.rs @@ -42,10 +42,17 @@ pub struct OutgoingError { pub id: RequestId, } +/// Details available after an outgoing message has been written by a transport. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct OutgoingWriteComplete { + /// Length of the JSON text written by the transport, when serialization occurred. + pub serialized_bytes: Option, +} + #[derive(Debug)] pub struct QueuedOutgoingMessage { pub message: OutgoingMessage, - pub write_complete_tx: Option>, + pub write_complete_tx: Option>, } impl QueuedOutgoingMessage { diff --git a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs index d0e934d90b..cb524dd4dc 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs @@ -98,7 +98,7 @@ pub(super) struct QueuedServerEnvelope { pub(super) event: ServerEvent, pub(super) client_id: ClientId, pub(super) stream_id: StreamId, - pub(super) write_complete_tx: Option>, + pub(super) write_complete_tx: Option>, } #[derive(Clone)] diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index e34cd34091..e14a6e9a0e 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -1047,6 +1047,8 @@ impl RemoteControlWebsocket { (payloads, queued_server_envelope.write_complete_tx) }; + let serialized_bytes = + (!payloads.is_empty()).then(|| payloads.iter().map(String::len).sum::()); for payload in payloads { tokio::select! { _ = shutdown_token.cancelled() => return Ok(()), @@ -1058,7 +1060,7 @@ impl RemoteControlWebsocket { } } if let Some(write_complete_tx) = write_complete_tx { - let _ = write_complete_tx.send(()); + let _ = write_complete_tx.send(crate::OutgoingWriteComplete { serialized_bytes }); } } } diff --git a/codex-rs/app-server-transport/src/transport/stdio.rs b/codex-rs/app-server-transport/src/transport/stdio.rs index 2d30296cd0..82d81a9ae6 100644 --- a/codex-rs/app-server-transport/src/transport/stdio.rs +++ b/codex-rs/app-server-transport/src/transport/stdio.rs @@ -4,6 +4,7 @@ use super::TransportEvent; use super::forward_incoming_message; use super::next_connection_id; use super::serialize_outgoing_message; +use crate::outgoing_message::OutgoingWriteComplete; use crate::outgoing_message::QueuedOutgoingMessage; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::JSONRPCMessage; @@ -85,13 +86,16 @@ pub async fn start_stdio_connection( let Some(mut json) = serialize_outgoing_message(queued_message.message) else { continue; }; + let serialized_bytes = json.len(); json.push('\n'); if let Err(err) = stdout.write_all(json.as_bytes()).await { error!("Failed to write to stdout: {err}"); break; } if let Some(write_complete_tx) = queued_message.write_complete_tx { - let _ = write_complete_tx.send(()); + let _ = write_complete_tx.send(OutgoingWriteComplete { + serialized_bytes: Some(serialized_bytes), + }); } } info!("stdout writer exited (channel closed)"); diff --git a/codex-rs/app-server-transport/src/transport/websocket.rs b/codex-rs/app-server-transport/src/transport/websocket.rs index 5677faae67..c59a0f0924 100644 --- a/codex-rs/app-server-transport/src/transport/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/websocket.rs @@ -8,6 +8,7 @@ use super::forward_incoming_message; use super::next_connection_id; use super::serialize_outgoing_message; use crate::outgoing_message::ConnectionId; +use crate::outgoing_message::OutgoingWriteComplete; use crate::outgoing_message::QueuedOutgoingMessage; use axum::Router; use axum::body::Body; @@ -316,11 +317,14 @@ async fn run_websocket_outbound_loop( let Some(json) = serialize_outgoing_message(queued_message.message) else { continue; }; + let serialized_bytes = json.len(); if websocket_writer.send(M::text(json)).await.is_err() { break; } if let Some(write_complete_tx) = queued_message.write_complete_tx { - let _ = write_complete_tx.send(()); + let _ = write_complete_tx.send(OutgoingWriteComplete { + serialized_bytes: Some(serialized_bytes), + }); } } } diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 938d539a9f..e05e87a3a1 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -681,7 +681,11 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult { - self.thread_processor.thread_read(params).await + self.thread_processor.thread_read(&request_id, params).await } ClientRequest::ThreadTurnsList { params, .. } => { self.thread_processor.thread_turns_list(params).await diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 9ce815bafe..b9811b3c68 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -31,6 +31,7 @@ pub(crate) use codex_app_server_transport::ConnectionId; pub(crate) use codex_app_server_transport::OutgoingError; pub(crate) use codex_app_server_transport::OutgoingMessage; pub(crate) use codex_app_server_transport::OutgoingResponse; +pub(crate) use codex_app_server_transport::OutgoingWriteComplete; pub(crate) use codex_app_server_transport::QueuedOutgoingMessage; #[cfg(test)] @@ -85,7 +86,7 @@ pub(crate) enum OutgoingEnvelope { ToConnection { connection_id: ConnectionId, message: OutgoingMessage, - write_complete_tx: Option>, + write_complete_tx: Option>, }, Broadcast { message: OutgoingMessage, @@ -507,10 +508,34 @@ impl OutgoingMessageSender { self.send_response_as(request_id, response.into()).await; } + pub(crate) async fn send_response_with_write_complete( + &self, + request_id: ConnectionRequestId, + response: T, + ) -> oneshot::Receiver + where + T: Into, + { + let (write_complete_tx, write_complete_rx) = oneshot::channel(); + self.send_response_as_inner(request_id, response.into(), Some(write_complete_tx)) + .await; + write_complete_rx + } + pub(crate) async fn send_response_as( &self, request_id: ConnectionRequestId, response: ClientResponsePayload, + ) { + self.send_response_as_inner(request_id, response, None) + .await; + } + + async fn send_response_as_inner( + &self, + request_id: ConnectionRequestId, + response: ClientResponsePayload, + write_complete_tx: Option>, ) { let connection_id = request_id.connection_id; let request_id_for_analytics = request_id.request_id.clone(); @@ -536,6 +561,7 @@ impl OutgoingMessageSender { connection_id, outgoing_message, "response", + write_complete_tx, ) .await; } @@ -655,6 +681,7 @@ impl OutgoingMessageSender { request_id.connection_id, outgoing_message, "error", + None, ) .await; } @@ -665,11 +692,12 @@ impl OutgoingMessageSender { connection_id: ConnectionId, message: OutgoingMessage, message_kind: &'static str, + write_complete_tx: Option>, ) { let send_fut = self.sender.send(OutgoingEnvelope::ToConnection { connection_id, message, - write_complete_tx: None, + write_complete_tx, }); let send_result = if let Some(request_context) = request_context { send_fut.instrument(request_context.span()).await @@ -1050,6 +1078,42 @@ mod tests { } } + #[tokio::test] + async fn send_response_with_write_complete_returns_serialized_size() { + let (tx, mut rx) = mpsc::channel::(4); + let outgoing = + OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled()); + let request_id = ConnectionRequestId { + connection_id: ConnectionId(42), + request_id: RequestId::Integer(7), + }; + + let write_complete_rx = outgoing + .send_response_with_write_complete( + request_id, + ClientResponsePayload::ThreadArchive( + codex_app_server_protocol::ThreadArchiveResponse {}, + ), + ) + .await; + let envelope = rx.recv().await.expect("response envelope"); + let OutgoingEnvelope::ToConnection { + write_complete_tx: Some(write_complete_tx), + .. + } = envelope + else { + panic!("expected targeted response with write completion sender"); + }; + let expected = OutgoingWriteComplete { + serialized_bytes: Some(321), + }; + write_complete_tx + .send(expected) + .expect("write completion receiver"); + + assert_eq!(write_complete_rx.await, Ok(expected)); + } + #[tokio::test] async fn send_response_clears_registered_request_context() { let (tx, _rx) = mpsc::channel::(4); @@ -1152,7 +1216,9 @@ mod tests { assert!(matches!(message, OutgoingMessage::AppServerNotification(_))); write_complete_tx .expect("write completion sender should be attached") - .send(()) + .send(OutgoingWriteComplete { + serialized_bytes: Some(123), + }) .expect("receiver should still be waiting"); timeout(Duration::from_secs(1), send_task) diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 16c4de83b1..ef411b4a20 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -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::RolloutProjectionMeasurement; use codex_rollout::RolloutProjectionTelemetry; use codex_rollout::is_persisted_rollout_item; use codex_rollout::state_db::StateDbHandle; @@ -621,20 +622,23 @@ 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).record_projected_turns( - turns, + 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, turn.items.len() as u64)), - ); + .map(|turn| turn.items.len() as u64) + .collect(), + ) } fn is_completed_user_assistant_turn(turn: &Turn) -> bool { 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 553677d4cc..ec4131b378 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -19,6 +19,11 @@ struct ThreadListFilters { parent_thread_id: Option, } +struct ThreadReadView { + thread: Thread, + projection_measurement: Option, +} + fn collect_resume_override_mismatches( request: &ThreadResumeParams, config_snapshot: &ThreadConfigSnapshot, @@ -667,11 +672,31 @@ impl ThreadRequestProcessor { pub(crate) async fn thread_read( &self, + request_id: &ConnectionRequestId, params: ThreadReadParams, ) -> Result, JSONRPCErrorError> { - self.thread_read_response_inner(params) - .await - .map(|response| Some(response.into())) + let view = self.thread_read_response_inner(params).await?; + let response = ThreadReadResponse { + thread: view.thread, + }; + 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; + } + Ok(None) } pub(crate) async fn thread_turns_list( @@ -2153,7 +2178,7 @@ impl ThreadRequestProcessor { async fn thread_read_response_inner( &self, params: ThreadReadParams, - ) -> Result { + ) -> Result { let ThreadReadParams { thread_id, include_turns, @@ -2162,11 +2187,9 @@ impl ThreadRequestProcessor { let thread_uuid = ThreadId::from_string(&thread_id) .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; - let thread = self - .read_thread_view(thread_uuid, include_turns) + self.read_thread_view(thread_uuid, include_turns) .await - .map_err(thread_read_view_error)?; - Ok(ThreadReadResponse { thread }) + .map_err(thread_read_view_error) } /// Builds the API view for `thread/read` from persisted metadata plus optional live state. @@ -2174,50 +2197,59 @@ 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 thread = if include_turns { + let mut view = 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?; - self.load_live_thread_view( - thread_id, - include_turns, - loaded_thread, - persisted_thread, - ) - .await? - } else if let Some(thread) = self + .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 .load_persisted_thread_for_read(thread_id, include_turns) .await? { // Unloaded thread with turns: load metadata and history together // from the ThreadStore. - thread + view } else { return Err(ThreadReadViewError::InvalidRequest(format!( "thread not loaded: {thread_id}" ))); } - } else if let Some(thread) = self + } else if let Some(view) = self .load_persisted_thread_for_read(thread_id, include_turns) .await? { // Persisted metadata-only read: no live thread state is needed. - thread + view } 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. - self.load_live_thread_view( - thread_id, - include_turns, - loaded_thread, - /*persisted_thread*/ None, - ) - .await? + ThreadReadView { + thread: self + .load_live_thread_view( + thread_id, + include_turns, + loaded_thread, + /*persisted_thread*/ None, + ) + .await?, + projection_measurement: None, + } } else { return Err(ThreadReadViewError::InvalidRequest(format!( "thread not loaded: {thread_id}" @@ -2232,22 +2264,22 @@ impl ThreadRequestProcessor { let thread_status = self .thread_watch_manager - .loaded_status_for_thread(&thread.id) + .loaded_status_for_thread(&view.thread.id) .await; set_thread_status_and_interrupt_stale_turns( - &mut thread, + &mut view.thread, thread_status, has_live_in_progress_turn, ); - Ok(thread) + Ok(view) } 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 @@ -2261,11 +2293,16 @@ 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); - record_rollout_projection(thread_id, &thread.turns); + projection_measurement = + prepare_rollout_projection_measurement(thread_id, &thread.turns); } - Ok(Some(thread)) + Ok(Some(ThreadReadView { + thread, + projection_measurement, + })) } Err(ThreadStoreError::InvalidRequest { message }) if message == format!("no rollout found for thread id {thread_id}") => diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index 9c40653764..63d7eaf6a5 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -17,6 +17,7 @@ pub(crate) use codex_app_server_transport::CHANNEL_CAPACITY; pub(crate) use codex_app_server_transport::ConnectionId; pub(crate) use codex_app_server_transport::ConnectionOrigin; pub(crate) use codex_app_server_transport::OutgoingMessage; +pub(crate) use codex_app_server_transport::OutgoingWriteComplete; pub(crate) use codex_app_server_transport::QueuedOutgoingMessage; pub(crate) use codex_app_server_transport::RemoteControlEnableError; pub(crate) use codex_app_server_transport::RemoteControlHandle; @@ -135,7 +136,7 @@ async fn send_message_to_connection( connections: &mut HashMap, connection_id: ConnectionId, message: OutgoingMessage, - write_complete_tx: Option>, + write_complete_tx: Option>, ) -> bool { let Some(connection_state) = connections.get(&connection_id) else { warn!("dropping message for disconnected connection: {connection_id:?}"); diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 3419185459..8d96dd7257 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -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::RolloutProjectionMeasurement; pub use persistence_metrics::RolloutProjectionTelemetry; pub use persistence_metrics::measure_and_filter_rollout_items; pub use policy::is_persisted_rollout_item; diff --git a/codex-rs/rollout/src/persistence_metrics.rs b/codex-rs/rollout/src/persistence_metrics.rs index 264a17ebc3..f916b1f768 100644 --- a/codex-rs/rollout/src/persistence_metrics.rs +++ b/codex-rs/rollout/src/persistence_metrics.rs @@ -21,11 +21,14 @@ const LOAD_THREAD_BYTES_METRIC: &str = "codex.rollout.persistence.load_thread_by 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 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_TURN_BYTES_METRIC: &str = "codex.rollout.persistence.projection_turn_bytes"; -const PROJECTION_TURN_ITEMS_METRIC: &str = "codex.rollout.persistence.projection_turn_items"; +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"; @@ -33,6 +36,8 @@ 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 { @@ -462,7 +467,7 @@ impl RolloutLoadTelemetry { } } -/// Records detached size proxies for app-server ThreadItem projection reads. +/// Prepares detached size proxies for app-server ThreadItem projection reads. pub struct RolloutProjectionTelemetry { metrics: Option, sampled: bool, @@ -480,100 +485,35 @@ impl RolloutProjectionTelemetry { } } - /// Records compact JSON bytes for the final materialized app-server turns. - pub fn record_projected_turns<'a, T, U>( + /// Prepares counts to record with the existing serialized `thread/read` response. + pub fn prepare_response_measurement( &self, - turns: &T, turn_count: u64, item_count: u64, - forecast_turn_items: impl IntoIterator, - ) where - T: Serialize + ?Sized, - U: Serialize + 'a, - { - let Some(metrics) = self.enabled_metrics() else { - return; - }; + 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; + return None; } - 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::>>() + if !IN_FLIGHT_PROJECTED_THREADS + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(self.thread_id) { - 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; + return None; } - record_size_totals( + Some(RolloutProjectionMeasurement { 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"), - ], - ); + thread_id: self.thread_id, + turn_count, + item_count, + completed_turn_item_counts, + }) } fn enabled_metrics(&self) -> Option<&MetricsClient> { @@ -581,6 +521,72 @@ impl RolloutProjectionTelemetry { } } +/// 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, diff --git a/codex-rs/rollout/src/persistence_metrics_tests.rs b/codex-rs/rollout/src/persistence_metrics_tests.rs index 117abd54be..fc8b64834a 100644 --- a/codex-rs/rollout/src/persistence_metrics_tests.rs +++ b/codex-rs/rollout/src/persistence_metrics_tests.rs @@ -1,7 +1,3 @@ -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; @@ -17,7 +13,6 @@ 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; @@ -28,7 +23,6 @@ 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 { @@ -270,21 +264,6 @@ fn filtered_item_completion_includes_its_nested_item_type() { ); } -#[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 = [ @@ -401,24 +380,8 @@ fn rollout_turn_sizes_exclude_incomplete_and_commentary_only_turns() { assert_eq!(tracker.finish(), Vec::::new()); } -struct FailingSerialize { - serialization_attempted: Arc, -} - -impl Serialize for FailingSerialize { - fn serialize(&self, _serializer: S) -> Result - 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() { +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}")) @@ -426,17 +389,8 @@ fn exporter_disabled_path_does_not_serialize_projected_turns() { 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), - }; + let measurement = RolloutProjectionTelemetry::new(sampled_thread_id) + .prepare_response_measurement(/*turn_count*/ 1, /*item_count*/ 1, vec![1]); - 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)); + assert!(measurement.is_none()); }