codex: measure thread reads from transport output

This commit is contained in:
Tom Wiltzius
2026-06-24 19:52:03 -07:00
parent 46f89fcce0
commit 2018e6c904
15 changed files with 278 additions and 187 deletions
+1
View File
@@ -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;
@@ -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<usize>,
}
#[derive(Debug)]
pub struct QueuedOutgoingMessage {
pub message: OutgoingMessage,
pub write_complete_tx: Option<oneshot::Sender<()>>,
pub write_complete_tx: Option<oneshot::Sender<OutgoingWriteComplete>>,
}
impl QueuedOutgoingMessage {
@@ -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<oneshot::Sender<()>>,
pub(super) write_complete_tx: Option<oneshot::Sender<crate::OutgoingWriteComplete>>,
}
#[derive(Clone)]
@@ -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::<usize>());
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 });
}
}
}
@@ -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)");
@@ -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<M, SinkError>(
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),
});
}
}
}
+5 -1
View File
@@ -681,7 +681,11 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
}
}
if let Some(write_complete_tx) = queued_message.write_complete_tx {
let _ = write_complete_tx.send(());
let _ = write_complete_tx.send(
codex_app_server_transport::OutgoingWriteComplete {
serialized_bytes: None,
},
);
}
}
}
+1 -1
View File
@@ -1218,7 +1218,7 @@ impl MessageProcessor {
self.thread_processor.thread_loaded_list(params).await
}
ClientRequest::ThreadRead { params, .. } => {
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
+69 -3
View File
@@ -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<oneshot::Sender<()>>,
write_complete_tx: Option<oneshot::Sender<OutgoingWriteComplete>>,
},
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<T>(
&self,
request_id: ConnectionRequestId,
response: T,
) -> oneshot::Receiver<OutgoingWriteComplete>
where
T: Into<ClientResponsePayload>,
{
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<oneshot::Sender<OutgoingWriteComplete>>,
) {
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<oneshot::Sender<OutgoingWriteComplete>>,
) {
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::<OutgoingEnvelope>(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::<OutgoingEnvelope>(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)
@@ -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<T
builder.finish()
}
fn record_rollout_projection(thread_id: ThreadId, turns: &[Turn]) {
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).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 {
@@ -19,6 +19,11 @@ 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,
@@ -667,11 +672,31 @@ impl ThreadRequestProcessor {
pub(crate) async fn thread_read(
&self,
request_id: &ConnectionRequestId,
params: ThreadReadParams,
) -> Result<Option<ClientResponsePayload>, 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<ThreadReadResponse, JSONRPCErrorError> {
) -> Result<ThreadReadView, JSONRPCErrorError> {
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<Thread, ThreadReadViewError> {
) -> Result<ThreadReadView, ThreadReadViewError> {
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<Option<Thread>, ThreadReadViewError> {
) -> Result<Option<ThreadReadView>, 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}") =>
+2 -1
View File
@@ -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<ConnectionId, OutboundConnectionState>,
connection_id: ConnectionId,
message: OutgoingMessage,
write_complete_tx: Option<tokio::sync::oneshot::Sender<()>>,
write_complete_tx: Option<tokio::sync::oneshot::Sender<OutgoingWriteComplete>>,
) -> bool {
let Some(connection_state) = connections.get(&connection_id) else {
warn!("dropping message for disconnected connection: {connection_id:?}");
+1
View File
@@ -59,6 +59,7 @@ pub use list::rollout_date_parts;
pub use metadata::builder_from_items;
pub use persistence_metrics::RolloutPersistenceBatchMeasurement;
pub use persistence_metrics::RolloutPersistenceTelemetry;
pub use persistence_metrics::RolloutProjectionMeasurement;
pub use persistence_metrics::RolloutProjectionTelemetry;
pub use persistence_metrics::measure_and_filter_rollout_items;
pub use policy::is_persisted_rollout_item;
+92 -86
View File
@@ -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<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 {
@@ -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<MetricsClient>,
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<Item = (&'a U, u64)>,
) where
T: Serialize + ?Sized,
U: Serialize + 'a,
{
let Some(metrics) = self.enabled_metrics() else {
return;
};
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;
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::<serde_json::Result<Vec<_>>>()
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<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>,
@@ -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::<RolloutSizeTotals>::new());
}
struct FailingSerialize {
serialization_attempted: Arc<AtomicBool>,
}
impl Serialize for FailingSerialize {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.serialization_attempted.store(true, Ordering::Relaxed);
Err(serde::ser::Error::custom(
"intentional serialization failure",
))
}
}
#[test]
fn exporter_disabled_path_does_not_serialize_projected_turns() {
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());
}