mirror of
https://github.com/openai/codex.git
synced 2026-11-12 21:20:34 +00:00
codex: measure thread reads from transport output
This commit is contained in:
@@ -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),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}") =>
|
||||
|
||||
@@ -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:?}");
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user