mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
measure rollout load and ThreadItem projection sizes
This commit is contained in:
@@ -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::RolloutProjectionTelemetry;
|
||||
use codex_rollout::is_persisted_rollout_item;
|
||||
use codex_rollout::state_db::StateDbHandle;
|
||||
use codex_rollout::state_db::reconcile_rollout;
|
||||
@@ -619,3 +620,40 @@ pub(crate) fn build_api_turns_from_rollout_items(items: &[RolloutItem]) -> Vec<T
|
||||
}
|
||||
builder.finish()
|
||||
}
|
||||
|
||||
fn record_rollout_projection(thread_id: ThreadId, turns: &[Turn]) {
|
||||
let item_count = turns
|
||||
.iter()
|
||||
.map(|turn| turn.items.len() as u64)
|
||||
.sum::<u64>();
|
||||
RolloutProjectionTelemetry::new(thread_id).record_projected_turns(
|
||||
turns,
|
||||
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)),
|
||||
);
|
||||
}
|
||||
|
||||
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;
|
||||
|
||||
@@ -2263,6 +2263,7 @@ impl ThreadRequestProcessor {
|
||||
thread_from_stored_thread(stored_thread, fallback_provider, &self.config.cwd);
|
||||
if include_turns && let Some(history) = history {
|
||||
thread.turns = build_api_turns_from_rollout_items(&history.items);
|
||||
record_rollout_projection(thread_id, &thread.turns);
|
||||
}
|
||||
Ok(Some(thread))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
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]));
|
||||
}
|
||||
@@ -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::RolloutProjectionTelemetry;
|
||||
pub use persistence_metrics::measure_and_filter_rollout_items;
|
||||
pub use policy::is_persisted_rollout_item;
|
||||
pub use policy::persisted_rollout_items;
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
use std::sync::LazyLock;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use codex_otel::MetricsClient;
|
||||
@@ -8,15 +10,29 @@ 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 PROJECTION_THREAD_BYTES_METRIC: &str = "codex.rollout.persistence.projection_thread_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 MEASUREMENT_ERROR_METRIC: &str = "codex.rollout.persistence.measurement_error";
|
||||
const SAMPLE_DENOMINATOR: u64 = 100;
|
||||
const SAMPLE_RATE_LABEL: &str = "0.01";
|
||||
static MEASURED_ROLLOUT_LOAD_THREADS: LazyLock<Mutex<HashSet<ThreadId>>> =
|
||||
LazyLock::new(|| Mutex::new(HashSet::new()));
|
||||
static MEASURED_PROJECTED_THREADS: LazyLock<Mutex<HashSet<ThreadId>>> =
|
||||
LazyLock::new(|| Mutex::new(HashSet::new()));
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PersistenceDecision {
|
||||
@@ -200,7 +216,7 @@ fn add_item_to_turn(totals: &mut TurnSizeTotals, item: &RolloutItemMeasurement)
|
||||
}
|
||||
}
|
||||
|
||||
fn serialized_len(item: &RolloutItem) -> serde_json::Result<u64> {
|
||||
fn serialized_len(item: &(impl Serialize + ?Sized)) -> serde_json::Result<u64> {
|
||||
let mut writer = CountingWriter::default();
|
||||
serde_json::to_writer(&mut writer, item)?;
|
||||
Ok(writer.bytes)
|
||||
@@ -391,6 +407,310 @@ impl RolloutPersistenceTelemetry {
|
||||
}
|
||||
}
|
||||
|
||||
/// Records detached per-thread and per-turn size proxies for rollout loads.
|
||||
pub(crate) struct RolloutLoadTelemetry {
|
||||
metrics: Option<MetricsClient>,
|
||||
sampled: bool,
|
||||
thread_id: ThreadId,
|
||||
}
|
||||
|
||||
impl RolloutLoadTelemetry {
|
||||
pub fn new(thread_id: ThreadId) -> Self {
|
||||
let metrics = codex_otel::global();
|
||||
let sampled = metrics.is_some() && is_thread_sampled(thread_id);
|
||||
Self {
|
||||
metrics,
|
||||
sampled,
|
||||
thread_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// Records bytes for successfully loaded decompressed JSONL records, excluding newlines.
|
||||
pub fn record_rollout_load(
|
||||
&self,
|
||||
totals: RolloutSizeTotals,
|
||||
turn_totals: &[RolloutSizeTotals],
|
||||
) {
|
||||
let Some(metrics) = self.enabled_metrics() else {
|
||||
return;
|
||||
};
|
||||
if !mark_thread_measured(&MEASURED_ROLLOUT_LOAD_THREADS, self.thread_id) {
|
||||
return;
|
||||
}
|
||||
record_size_totals(
|
||||
metrics,
|
||||
totals,
|
||||
LOAD_THREAD_BYTES_METRIC,
|
||||
LOAD_THREAD_ITEMS_METRIC,
|
||||
"rollout_jsonl",
|
||||
"rollout_line_json_v1",
|
||||
"rollout_load",
|
||||
);
|
||||
record_turn_size_totals(
|
||||
metrics,
|
||||
turn_totals,
|
||||
LOAD_TURN_BYTES_METRIC,
|
||||
LOAD_TURN_ITEMS_METRIC,
|
||||
"rollout_jsonl",
|
||||
"rollout_line_json_v1",
|
||||
"rollout_load",
|
||||
);
|
||||
}
|
||||
|
||||
fn enabled_metrics(&self) -> Option<&MetricsClient> {
|
||||
self.sampled.then_some(self.metrics.as_ref()).flatten()
|
||||
}
|
||||
}
|
||||
|
||||
/// Records detached size proxies for app-server ThreadItem projection reads.
|
||||
pub struct RolloutProjectionTelemetry {
|
||||
metrics: Option<MetricsClient>,
|
||||
sampled: bool,
|
||||
thread_id: ThreadId,
|
||||
}
|
||||
|
||||
impl RolloutProjectionTelemetry {
|
||||
pub fn new(thread_id: ThreadId) -> Self {
|
||||
let metrics = codex_otel::global();
|
||||
let sampled = metrics.is_some() && is_thread_sampled(thread_id);
|
||||
Self {
|
||||
metrics,
|
||||
sampled,
|
||||
thread_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// Records compact JSON bytes for the final materialized app-server turns.
|
||||
pub fn record_projected_turns<'a, T, U>(
|
||||
&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;
|
||||
};
|
||||
if MEASURED_PROJECTED_THREADS
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.contains(&self.thread_id)
|
||||
{
|
||||
return;
|
||||
}
|
||||
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<_>>>()
|
||||
{
|
||||
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;
|
||||
}
|
||||
record_size_totals(
|
||||
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"),
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
fn enabled_metrics(&self) -> Option<&MetricsClient> {
|
||||
self.sampled.then_some(self.metrics.as_ref()).flatten()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct RolloutTurnSizeTracker {
|
||||
current: Option<PendingRolloutTurn>,
|
||||
completed: Vec<RolloutSizeTotals>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct PendingRolloutTurn {
|
||||
totals: RolloutSizeTotals,
|
||||
explicit: bool,
|
||||
saw_user_message: bool,
|
||||
saw_final_agent_message: bool,
|
||||
}
|
||||
|
||||
impl RolloutTurnSizeTracker {
|
||||
pub(crate) fn observe(&mut self, item: &RolloutItem, payload_bytes: u64) {
|
||||
match item {
|
||||
RolloutItem::EventMsg(EventMsg::TurnStarted(_)) => {
|
||||
self.finish_implicit_turn();
|
||||
self.current = Some(PendingRolloutTurn {
|
||||
explicit: true,
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
RolloutItem::EventMsg(EventMsg::UserMessage(_)) => {
|
||||
if self
|
||||
.current
|
||||
.as_ref()
|
||||
.is_some_and(|turn| !turn.explicit && turn.saw_user_message)
|
||||
{
|
||||
self.finish_implicit_turn();
|
||||
}
|
||||
self.current
|
||||
.get_or_insert_with(PendingRolloutTurn::default)
|
||||
.saw_user_message = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let Some(turn) = self.current.as_mut() else {
|
||||
return;
|
||||
};
|
||||
turn.totals.items = turn.totals.items.saturating_add(1);
|
||||
turn.totals.payload_bytes = turn.totals.payload_bytes.saturating_add(payload_bytes);
|
||||
|
||||
match item {
|
||||
RolloutItem::EventMsg(EventMsg::AgentMessage(message))
|
||||
if !matches!(
|
||||
message.phase,
|
||||
Some(codex_protocol::models::MessagePhase::Commentary)
|
||||
) =>
|
||||
{
|
||||
turn.saw_final_agent_message = true;
|
||||
}
|
||||
RolloutItem::EventMsg(EventMsg::TurnComplete(_)) => self.finish_current_turn(),
|
||||
RolloutItem::EventMsg(EventMsg::TurnAborted(_)) => {
|
||||
self.current = None;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn finish(mut self) -> Vec<RolloutSizeTotals> {
|
||||
self.finish_implicit_turn();
|
||||
self.completed
|
||||
}
|
||||
|
||||
fn finish_implicit_turn(&mut self) {
|
||||
if self.current.as_ref().is_some_and(|turn| !turn.explicit) {
|
||||
self.finish_current_turn();
|
||||
}
|
||||
}
|
||||
|
||||
fn finish_current_turn(&mut self) {
|
||||
if let Some(turn) = self.current.take()
|
||||
&& turn.saw_user_message
|
||||
&& turn.saw_final_agent_message
|
||||
{
|
||||
self.completed.push(turn.totals);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn record_size_totals(
|
||||
metrics: &MetricsClient,
|
||||
totals: RolloutSizeTotals,
|
||||
bytes_metric: &'static str,
|
||||
items_metric: &'static str,
|
||||
representation: &str,
|
||||
encoding: &str,
|
||||
source: &str,
|
||||
) {
|
||||
let tags = [
|
||||
("representation", representation),
|
||||
("encoding", encoding),
|
||||
("sample_rate", SAMPLE_RATE_LABEL),
|
||||
("source", source),
|
||||
];
|
||||
let _ = metrics.histogram(bytes_metric, saturating_i64(totals.payload_bytes), &tags);
|
||||
let _ = metrics.histogram(items_metric, saturating_i64(totals.items), &tags);
|
||||
}
|
||||
|
||||
fn record_turn_size_totals(
|
||||
metrics: &MetricsClient,
|
||||
turn_totals: &[RolloutSizeTotals],
|
||||
bytes_metric: &'static str,
|
||||
items_metric: &'static str,
|
||||
representation: &str,
|
||||
encoding: &str,
|
||||
source: &str,
|
||||
) {
|
||||
let tags = [
|
||||
("representation", representation),
|
||||
("encoding", encoding),
|
||||
("sample_rate", SAMPLE_RATE_LABEL),
|
||||
("source", source),
|
||||
];
|
||||
for totals in turn_totals {
|
||||
let _ = metrics.histogram(bytes_metric, saturating_i64(totals.payload_bytes), &tags);
|
||||
let _ = metrics.histogram(items_metric, saturating_i64(totals.items), &tags);
|
||||
}
|
||||
}
|
||||
|
||||
fn mark_thread_measured(measured_threads: &Mutex<HashSet<ThreadId>>, thread_id: ThreadId) -> bool {
|
||||
measured_threads
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.insert(thread_id)
|
||||
}
|
||||
|
||||
fn saturating_i64(value: u64) -> i64 {
|
||||
value.try_into().unwrap_or(i64::MAX)
|
||||
}
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
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;
|
||||
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;
|
||||
@@ -10,14 +15,20 @@ 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 serde::Serialize;
|
||||
|
||||
use super::CompletedTurnMeasurement;
|
||||
use super::RolloutProjectionTelemetry;
|
||||
use super::RolloutSizeTotals;
|
||||
use super::RolloutTurnSizeTracker;
|
||||
use super::TurnMeasurementState;
|
||||
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 {
|
||||
@@ -258,3 +269,174 @@ fn filtered_item_completion_includes_its_nested_item_type() {
|
||||
super::PersistenceDecision::Dropped
|
||||
);
|
||||
}
|
||||
|
||||
#[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 = [
|
||||
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, 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, 1);
|
||||
}
|
||||
for item in [start("turn-2"), user()] {
|
||||
tracker.observe(&item, 1);
|
||||
}
|
||||
|
||||
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() {
|
||||
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 serialization_attempted = Arc::new(AtomicBool::new(false));
|
||||
let turns = FailingSerialize {
|
||||
serialization_attempted: Arc::clone(&serialization_attempted),
|
||||
};
|
||||
|
||||
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));
|
||||
}
|
||||
|
||||
@@ -44,6 +44,9 @@ 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;
|
||||
@@ -903,6 +906,8 @@ impl RolloutRecorder {
|
||||
let mut items: Vec<RolloutItem> = Vec::new();
|
||||
let mut thread_id: Option<ThreadId> = None;
|
||||
let mut parse_errors = 0usize;
|
||||
let mut loaded_jsonl_bytes = 0_u64;
|
||||
let mut loaded_turns = RolloutTurnSizeTracker::default();
|
||||
let mut reader = compression::open_rollout_line_reader(path).await?;
|
||||
let mut saw_non_empty_line = false;
|
||||
while let Some(line) = reader.next_line().await? {
|
||||
@@ -934,6 +939,9 @@ 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) => {
|
||||
@@ -952,6 +960,15 @@ 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))
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user