mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
[codex] Instrument rollout persistence bytes (#29498)
- Add 1%-sampled rollout persistence metrics that report per-item and per-thread JSON byte totals before and after filtering when metrics export is enabled. - Tag each item with its exact response or event variant, including nested turn-item kinds for conditionally persisted completion events, so aggregate cloud-storage impact can be estimated by policy choice.
This commit is contained in:
Generated
+1
@@ -3979,6 +3979,7 @@ dependencies = [
|
||||
"chrono",
|
||||
"codex-git-utils",
|
||||
"codex-install-context",
|
||||
"codex-otel",
|
||||
"codex-protocol",
|
||||
"codex-rollout",
|
||||
"codex-state",
|
||||
|
||||
@@ -8,6 +8,7 @@ pub(crate) mod compression;
|
||||
pub(crate) mod config;
|
||||
pub(crate) mod list;
|
||||
pub(crate) mod metadata;
|
||||
mod persistence_metrics;
|
||||
pub(crate) mod policy;
|
||||
pub(crate) mod recorder;
|
||||
pub(crate) mod search;
|
||||
@@ -60,6 +61,9 @@ pub use list::read_session_meta_line;
|
||||
pub use list::read_thread_item_from_rollout;
|
||||
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::measure_and_filter_rollout_items;
|
||||
pub use policy::is_persisted_rollout_item;
|
||||
pub use policy::persisted_rollout_items;
|
||||
pub use policy::should_persist_response_item_for_memories;
|
||||
|
||||
@@ -0,0 +1,406 @@
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use codex_otel::MetricsClient;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::items::TurnItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
|
||||
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 MEASUREMENT_ERROR_METRIC: &str = "codex.rollout.persistence.measurement_error";
|
||||
const SAMPLE_DENOMINATOR: u64 = 100;
|
||||
const SAMPLE_RATE_LABEL: &str = "0.01";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PersistenceDecision {
|
||||
Kept,
|
||||
Dropped,
|
||||
}
|
||||
|
||||
impl PersistenceDecision {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Kept => "kept",
|
||||
Self::Dropped => "dropped",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct RolloutSizeTotals {
|
||||
pub items: u64,
|
||||
pub payload_bytes: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RolloutItemMeasurement {
|
||||
pub decision: PersistenceDecision,
|
||||
pub rollout_item_type: String,
|
||||
pub payload_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct RolloutPersistenceBatchMeasurement {
|
||||
pub pre_filter: RolloutSizeTotals,
|
||||
pub post_filter: RolloutSizeTotals,
|
||||
pub items: Vec<RolloutItemMeasurement>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
struct TurnSizeTotals {
|
||||
pre_filter: RolloutSizeTotals,
|
||||
post_filter: RolloutSizeTotals,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum TurnOutcome {
|
||||
Completed,
|
||||
Aborted,
|
||||
}
|
||||
|
||||
impl TurnOutcome {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Completed => "completed",
|
||||
Self::Aborted => "aborted",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
struct CompletedTurnMeasurement {
|
||||
totals: TurnSizeTotals,
|
||||
outcome: TurnOutcome,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, PartialEq, Eq)]
|
||||
struct TurnMeasurementState {
|
||||
pending: TurnSizeTotals,
|
||||
active: Option<TurnSizeTotals>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, PartialEq, Eq)]
|
||||
struct TurnMeasurementUpdate {
|
||||
completed: Vec<CompletedTurnMeasurement>,
|
||||
boundary_errors: Vec<&'static str>,
|
||||
}
|
||||
|
||||
/// Measures logical JSON sizes while applying the shared rollout persistence policy once.
|
||||
pub fn measure_and_filter_rollout_items(
|
||||
items: &[RolloutItem],
|
||||
) -> (Vec<RolloutItem>, RolloutPersistenceBatchMeasurement) {
|
||||
let mut persisted = Vec::new();
|
||||
let mut measurement = RolloutPersistenceBatchMeasurement {
|
||||
items: Vec::with_capacity(items.len()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
for item in items {
|
||||
let kept = is_persisted_rollout_item(item);
|
||||
let decision = if kept {
|
||||
PersistenceDecision::Kept
|
||||
} else {
|
||||
PersistenceDecision::Dropped
|
||||
};
|
||||
let payload_bytes = serialized_len(item).ok();
|
||||
add_to_totals(&mut measurement.pre_filter, payload_bytes);
|
||||
if kept {
|
||||
add_to_totals(&mut measurement.post_filter, payload_bytes);
|
||||
persisted.push(item.clone());
|
||||
}
|
||||
measurement.items.push(RolloutItemMeasurement {
|
||||
decision,
|
||||
rollout_item_type: rollout_item_type(item),
|
||||
payload_bytes,
|
||||
});
|
||||
}
|
||||
|
||||
(persisted, measurement)
|
||||
}
|
||||
|
||||
fn add_to_totals(totals: &mut RolloutSizeTotals, payload_bytes: Option<u64>) {
|
||||
totals.items = totals.items.saturating_add(1);
|
||||
if let Some(payload_bytes) = payload_bytes {
|
||||
totals.payload_bytes = totals.payload_bytes.saturating_add(payload_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
fn update_turn_measurements(
|
||||
state: &mut TurnMeasurementState,
|
||||
items: &[RolloutItem],
|
||||
measurement: &RolloutPersistenceBatchMeasurement,
|
||||
) -> TurnMeasurementUpdate {
|
||||
let mut update = TurnMeasurementUpdate::default();
|
||||
for (item, item_measurement) in items.iter().zip(&measurement.items) {
|
||||
match item {
|
||||
RolloutItem::EventMsg(EventMsg::TurnStarted(_)) => {
|
||||
if state.active.take().is_some() {
|
||||
update.boundary_errors.push("event.turn_started");
|
||||
}
|
||||
let mut totals = std::mem::take(&mut state.pending);
|
||||
add_item_to_turn(&mut totals, item_measurement);
|
||||
state.active = Some(totals);
|
||||
}
|
||||
RolloutItem::EventMsg(EventMsg::TurnComplete(_)) => {
|
||||
finish_turn(
|
||||
state,
|
||||
item_measurement,
|
||||
TurnOutcome::Completed,
|
||||
"event.turn_complete",
|
||||
&mut update,
|
||||
);
|
||||
}
|
||||
RolloutItem::EventMsg(EventMsg::TurnAborted(_)) => {
|
||||
finish_turn(
|
||||
state,
|
||||
item_measurement,
|
||||
TurnOutcome::Aborted,
|
||||
"event.turn_aborted",
|
||||
&mut update,
|
||||
);
|
||||
}
|
||||
_ => match state.active.as_mut() {
|
||||
Some(totals) => add_item_to_turn(totals, item_measurement),
|
||||
None => add_item_to_turn(&mut state.pending, item_measurement),
|
||||
},
|
||||
}
|
||||
}
|
||||
update
|
||||
}
|
||||
|
||||
fn finish_turn(
|
||||
state: &mut TurnMeasurementState,
|
||||
item: &RolloutItemMeasurement,
|
||||
outcome: TurnOutcome,
|
||||
boundary_type: &'static str,
|
||||
update: &mut TurnMeasurementUpdate,
|
||||
) {
|
||||
let Some(mut totals) = state.active.take() else {
|
||||
state.pending = TurnSizeTotals::default();
|
||||
update.boundary_errors.push(boundary_type);
|
||||
return;
|
||||
};
|
||||
add_item_to_turn(&mut totals, item);
|
||||
update
|
||||
.completed
|
||||
.push(CompletedTurnMeasurement { totals, outcome });
|
||||
}
|
||||
|
||||
fn add_item_to_turn(totals: &mut TurnSizeTotals, item: &RolloutItemMeasurement) {
|
||||
add_to_totals(&mut totals.pre_filter, item.payload_bytes);
|
||||
if item.decision == PersistenceDecision::Kept {
|
||||
add_to_totals(&mut totals.post_filter, item.payload_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
fn serialized_len(item: &RolloutItem) -> serde_json::Result<u64> {
|
||||
let mut writer = CountingWriter::default();
|
||||
serde_json::to_writer(&mut writer, item)?;
|
||||
Ok(writer.bytes)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct CountingWriter {
|
||||
bytes: u64,
|
||||
}
|
||||
|
||||
impl Write for CountingWriter {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
self.bytes = self.bytes.saturating_add(buf.len() as u64);
|
||||
Ok(buf.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn rollout_item_type(item: &RolloutItem) -> String {
|
||||
match item {
|
||||
RolloutItem::SessionMeta(_) => "session_meta".to_string(),
|
||||
RolloutItem::ResponseItem(item) => response_item_type(item).to_string(),
|
||||
RolloutItem::InterAgentCommunication(_) => "inter_agent_communication".to_string(),
|
||||
RolloutItem::Compacted(_) => "compacted".to_string(),
|
||||
RolloutItem::TurnContext(_) => "turn_context".to_string(),
|
||||
RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) => {
|
||||
format!("event.item_completed.{}", turn_item_type(&event.item))
|
||||
}
|
||||
RolloutItem::EventMsg(event) => format!("event.{event}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn turn_item_type(item: &TurnItem) -> &'static str {
|
||||
match item {
|
||||
TurnItem::UserMessage(_) => "user_message",
|
||||
TurnItem::HookPrompt(_) => "hook_prompt",
|
||||
TurnItem::AgentMessage(_) => "agent_message",
|
||||
TurnItem::Plan(_) => "plan",
|
||||
TurnItem::Reasoning(_) => "reasoning",
|
||||
TurnItem::WebSearch(_) => "web_search",
|
||||
TurnItem::ImageView(_) => "image_view",
|
||||
TurnItem::Sleep(_) => "sleep",
|
||||
TurnItem::ImageGeneration(_) => "image_generation",
|
||||
TurnItem::FileChange(_) => "file_change",
|
||||
TurnItem::McpToolCall(_) => "mcp_tool_call",
|
||||
TurnItem::ContextCompaction(_) => "context_compaction",
|
||||
}
|
||||
}
|
||||
|
||||
fn response_item_type(item: &ResponseItem) -> &'static str {
|
||||
match item {
|
||||
ResponseItem::Message { .. } => "response.message",
|
||||
ResponseItem::AgentMessage { .. } => "response.agent_message",
|
||||
ResponseItem::Reasoning { .. } => "response.reasoning",
|
||||
ResponseItem::LocalShellCall { .. } => "response.local_shell_call",
|
||||
ResponseItem::FunctionCall { .. } => "response.function_call",
|
||||
ResponseItem::ToolSearchCall { .. } => "response.tool_search_call",
|
||||
ResponseItem::FunctionCallOutput { .. } => "response.function_call_output",
|
||||
ResponseItem::ToolSearchOutput { .. } => "response.tool_search_output",
|
||||
ResponseItem::CustomToolCall { .. } => "response.custom_tool_call",
|
||||
ResponseItem::CustomToolCallOutput { .. } => "response.custom_tool_call_output",
|
||||
ResponseItem::WebSearchCall { .. } => "response.web_search_call",
|
||||
ResponseItem::ImageGenerationCall { .. } => "response.image_generation_call",
|
||||
ResponseItem::Compaction { .. } => "response.compaction",
|
||||
ResponseItem::CompactionTrigger { .. } => "response.compaction_trigger",
|
||||
ResponseItem::ContextCompaction { .. } => "response.context_compaction",
|
||||
ResponseItem::Other => "response.other",
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RolloutPersistenceTelemetry {
|
||||
metrics: Option<MetricsClient>,
|
||||
sampled: bool,
|
||||
turn_state: Option<Arc<Mutex<TurnMeasurementState>>>,
|
||||
}
|
||||
|
||||
impl RolloutPersistenceTelemetry {
|
||||
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,
|
||||
turn_state: sampled.then(|| Arc::new(Mutex::new(TurnMeasurementState::default()))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled_metrics().is_some()
|
||||
}
|
||||
|
||||
pub fn record_batch(
|
||||
&self,
|
||||
items: &[RolloutItem],
|
||||
measurement: &RolloutPersistenceBatchMeasurement,
|
||||
) {
|
||||
let Some(metrics) = self.enabled_metrics() else {
|
||||
return;
|
||||
};
|
||||
|
||||
for item in &measurement.items {
|
||||
if let Some(payload_bytes) = item.payload_bytes {
|
||||
let _ = metrics.histogram(
|
||||
ITEM_BYTES_METRIC,
|
||||
saturating_i64(payload_bytes),
|
||||
&[
|
||||
("decision", item.decision.as_str()),
|
||||
("rollout_item_type", item.rollout_item_type.as_str()),
|
||||
("encoding", "rollout_item_json_v1"),
|
||||
("sample_rate", SAMPLE_RATE_LABEL),
|
||||
],
|
||||
);
|
||||
} else {
|
||||
let _ = metrics.counter(
|
||||
MEASUREMENT_ERROR_METRIC,
|
||||
/*inc*/ 1,
|
||||
&[
|
||||
("rollout_item_type", item.rollout_item_type.as_str()),
|
||||
("phase", "serialize"),
|
||||
],
|
||||
);
|
||||
}
|
||||
}
|
||||
// Count successful input appends and the subset that remain storage operations after the
|
||||
// persistence policy removes filtered items.
|
||||
let _ = metrics.counter(
|
||||
APPEND_METRIC,
|
||||
/*inc*/ 1,
|
||||
&[("stage", "pre_filter"), ("sample_rate", SAMPLE_RATE_LABEL)],
|
||||
);
|
||||
if measurement.post_filter.items > 0 {
|
||||
let _ = metrics.counter(
|
||||
APPEND_METRIC,
|
||||
/*inc*/ 1,
|
||||
&[("stage", "post_filter"), ("sample_rate", SAMPLE_RATE_LABEL)],
|
||||
);
|
||||
}
|
||||
|
||||
let Some(turn_state) = self.turn_state.as_ref() else {
|
||||
return;
|
||||
};
|
||||
let turn_update = update_turn_measurements(
|
||||
&mut turn_state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner),
|
||||
items,
|
||||
measurement,
|
||||
);
|
||||
for boundary_type in turn_update.boundary_errors {
|
||||
let _ = metrics.counter(
|
||||
MEASUREMENT_ERROR_METRIC,
|
||||
/*inc*/ 1,
|
||||
&[
|
||||
("rollout_item_type", boundary_type),
|
||||
("phase", "turn_boundary"),
|
||||
],
|
||||
);
|
||||
}
|
||||
for turn in turn_update.completed {
|
||||
for (stage, totals) in [
|
||||
("pre_filter", turn.totals.pre_filter),
|
||||
("post_filter", turn.totals.post_filter),
|
||||
] {
|
||||
let _ = metrics.histogram(
|
||||
TURN_BYTES_METRIC,
|
||||
saturating_i64(totals.payload_bytes),
|
||||
&[
|
||||
("stage", stage),
|
||||
("outcome", turn.outcome.as_str()),
|
||||
("encoding", "rollout_item_json_v1"),
|
||||
("sample_rate", SAMPLE_RATE_LABEL),
|
||||
],
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn enabled_metrics(&self) -> Option<&MetricsClient> {
|
||||
self.sampled.then_some(self.metrics.as_ref()).flatten()
|
||||
}
|
||||
}
|
||||
|
||||
fn saturating_i64(value: u64) -> i64 {
|
||||
value.try_into().unwrap_or(i64::MAX)
|
||||
}
|
||||
|
||||
fn is_thread_sampled(thread_id: ThreadId) -> bool {
|
||||
let hash = thread_id
|
||||
.to_string()
|
||||
.bytes()
|
||||
.fold(0xcbf29ce484222325_u64, |hash, byte| {
|
||||
(hash ^ u64::from(byte)).wrapping_mul(0x100000001b3)
|
||||
});
|
||||
hash % SAMPLE_DENOMINATOR == 0
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "persistence_metrics_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,260 @@
|
||||
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::EventMsg;
|
||||
use codex_protocol::protocol::ItemCompletedEvent;
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
use codex_protocol::protocol::TurnAbortReason;
|
||||
use codex_protocol::protocol::TurnAbortedEvent;
|
||||
use codex_protocol::protocol::TurnCompleteEvent;
|
||||
use codex_protocol::protocol::TurnStartedEvent;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::CompletedTurnMeasurement;
|
||||
use super::TurnMeasurementState;
|
||||
use super::TurnOutcome;
|
||||
use super::TurnSizeTotals;
|
||||
use super::is_thread_sampled;
|
||||
use super::measure_and_filter_rollout_items;
|
||||
use super::update_turn_measurements;
|
||||
|
||||
fn retained_message(text: &str) -> RolloutItem {
|
||||
RolloutItem::ResponseItem(ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text: text.to_string(),
|
||||
}],
|
||||
phase: None,
|
||||
internal_chat_message_metadata_passthrough: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn turn_started(turn_id: &str) -> RolloutItem {
|
||||
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(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn turn_complete(turn_id: &str) -> RolloutItem {
|
||||
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
|
||||
turn_id: turn_id.to_string(),
|
||||
last_agent_message: None,
|
||||
completed_at: None,
|
||||
duration_ms: None,
|
||||
time_to_first_token_ms: None,
|
||||
}))
|
||||
}
|
||||
|
||||
fn turn_aborted(turn_id: &str) -> RolloutItem {
|
||||
RolloutItem::EventMsg(EventMsg::TurnAborted(TurnAbortedEvent {
|
||||
turn_id: Some(turn_id.to_string()),
|
||||
reason: TurnAbortReason::Interrupted,
|
||||
completed_at: None,
|
||||
duration_ms: None,
|
||||
}))
|
||||
}
|
||||
|
||||
fn update_for_batch(
|
||||
state: &mut TurnMeasurementState,
|
||||
items: &[RolloutItem],
|
||||
) -> super::TurnMeasurementUpdate {
|
||||
let (_, measurement) = measure_and_filter_rollout_items(items);
|
||||
update_turn_measurements(state, items, &measurement)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn thread_sampling_is_stable_and_selects_whole_threads() {
|
||||
let mut sampled = None;
|
||||
let mut unsampled = None;
|
||||
for value in 0..10_000_u128 {
|
||||
let thread_id = ThreadId::from_string(&format!("00000000-0000-0000-0000-{value:012x}"))
|
||||
.expect("valid thread id");
|
||||
if is_thread_sampled(thread_id) {
|
||||
sampled.get_or_insert(thread_id);
|
||||
} else {
|
||||
unsampled.get_or_insert(thread_id);
|
||||
}
|
||||
if sampled.is_some() && unsampled.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let sampled = sampled.expect("at least one sampled thread");
|
||||
let unsampled = unsampled.expect("at least one unsampled thread");
|
||||
assert!(is_thread_sampled(sampled));
|
||||
assert!(is_thread_sampled(sampled));
|
||||
assert!(!is_thread_sampled(unsampled));
|
||||
assert!(!is_thread_sampled(unsampled));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_batch_reports_exact_policy_counts_and_bytes() {
|
||||
let kept = retained_message("hello");
|
||||
let dropped = RolloutItem::ResponseItem(ResponseItem::Other);
|
||||
let items = vec![kept.clone(), dropped.clone()];
|
||||
|
||||
let (persisted, measurement) = measure_and_filter_rollout_items(&items);
|
||||
let kept_bytes = serde_json::to_vec(&kept)
|
||||
.expect("serialize kept item")
|
||||
.len() as u64;
|
||||
let dropped_bytes = serde_json::to_vec(&dropped)
|
||||
.expect("serialize dropped item")
|
||||
.len() as u64;
|
||||
|
||||
assert_eq!(
|
||||
serde_json::to_value(persisted).expect("serialize persisted items"),
|
||||
serde_json::to_value([kept]).expect("serialize expected items")
|
||||
);
|
||||
assert_eq!(measurement.pre_filter.items, 2);
|
||||
assert_eq!(
|
||||
measurement.pre_filter.payload_bytes,
|
||||
kept_bytes + dropped_bytes
|
||||
);
|
||||
assert_eq!(measurement.post_filter.items, 1);
|
||||
assert_eq!(measurement.post_filter.payload_bytes, kept_bytes);
|
||||
assert_eq!(measurement.items[0].payload_bytes, Some(kept_bytes));
|
||||
assert_eq!(measurement.items[1].payload_bytes, Some(dropped_bytes));
|
||||
assert_eq!(measurement.items[0].rollout_item_type, "response.message");
|
||||
assert_eq!(measurement.items[1].rollout_item_type, "response.other");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retained_items_are_byte_identical() {
|
||||
let item = retained_message("a moderately sized payload");
|
||||
let (persisted, measurement) = measure_and_filter_rollout_items(std::slice::from_ref(&item));
|
||||
|
||||
assert_eq!(
|
||||
serde_json::to_vec(&persisted[0]).expect("serialize persisted item"),
|
||||
serde_json::to_vec(&item).expect("serialize candidate item")
|
||||
);
|
||||
assert_eq!(
|
||||
measurement.post_filter.payload_bytes,
|
||||
measurement.items[0].payload_bytes.expect("payload bytes")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn turn_measurements_span_batches_and_include_items_before_start() {
|
||||
let first_turn = vec![
|
||||
retained_message("first prompt"),
|
||||
turn_started("turn-1"),
|
||||
retained_message("first response"),
|
||||
RolloutItem::ResponseItem(ResponseItem::Other),
|
||||
turn_complete("turn-1"),
|
||||
];
|
||||
let second_turn = vec![
|
||||
retained_message("second prompt"),
|
||||
turn_started("turn-2"),
|
||||
retained_message("second response"),
|
||||
turn_aborted("turn-2"),
|
||||
];
|
||||
let (_, first_expected) = measure_and_filter_rollout_items(&first_turn);
|
||||
let (_, second_expected) = measure_and_filter_rollout_items(&second_turn);
|
||||
let batches = [
|
||||
first_turn[..1].to_vec(),
|
||||
first_turn[1..3].to_vec(),
|
||||
vec![
|
||||
first_turn[3].clone(),
|
||||
first_turn[4].clone(),
|
||||
second_turn[0].clone(),
|
||||
],
|
||||
second_turn[1..].to_vec(),
|
||||
];
|
||||
|
||||
let mut state = TurnMeasurementState::default();
|
||||
let mut completed = Vec::new();
|
||||
let mut boundary_errors = Vec::new();
|
||||
for batch in batches {
|
||||
let update = update_for_batch(&mut state, &batch);
|
||||
completed.extend(update.completed);
|
||||
boundary_errors.extend(update.boundary_errors);
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
completed,
|
||||
vec![
|
||||
CompletedTurnMeasurement {
|
||||
totals: TurnSizeTotals {
|
||||
pre_filter: first_expected.pre_filter,
|
||||
post_filter: first_expected.post_filter,
|
||||
},
|
||||
outcome: TurnOutcome::Completed,
|
||||
},
|
||||
CompletedTurnMeasurement {
|
||||
totals: TurnSizeTotals {
|
||||
pre_filter: second_expected.pre_filter,
|
||||
post_filter: second_expected.post_filter,
|
||||
},
|
||||
outcome: TurnOutcome::Aborted,
|
||||
},
|
||||
]
|
||||
);
|
||||
assert_eq!(boundary_errors, Vec::<&str>::new());
|
||||
assert_eq!(state, TurnMeasurementState::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_turn_boundaries_reset_partial_measurements() {
|
||||
let mut state = TurnMeasurementState::default();
|
||||
let unmatched_completion = vec![retained_message("orphan"), turn_complete("turn-1")];
|
||||
let update = update_for_batch(&mut state, &unmatched_completion);
|
||||
|
||||
assert_eq!(update.completed, Vec::new());
|
||||
assert_eq!(update.boundary_errors, vec!["event.turn_complete"]);
|
||||
assert_eq!(state, TurnMeasurementState::default());
|
||||
|
||||
let replacement = vec![
|
||||
turn_started("turn-1"),
|
||||
retained_message("discarded partial turn"),
|
||||
turn_started("turn-2"),
|
||||
retained_message("retained turn"),
|
||||
turn_complete("turn-2"),
|
||||
];
|
||||
let (_, expected) = measure_and_filter_rollout_items(&replacement[2..]);
|
||||
let update = update_for_batch(&mut state, &replacement);
|
||||
|
||||
assert_eq!(
|
||||
update.completed,
|
||||
vec![CompletedTurnMeasurement {
|
||||
totals: TurnSizeTotals {
|
||||
pre_filter: expected.pre_filter,
|
||||
post_filter: expected.post_filter,
|
||||
},
|
||||
outcome: TurnOutcome::Completed,
|
||||
}]
|
||||
);
|
||||
assert_eq!(update.boundary_errors, vec!["event.turn_started"]);
|
||||
assert_eq!(state, TurnMeasurementState::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filtered_item_completion_includes_its_nested_item_type() {
|
||||
let item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
|
||||
thread_id: ThreadId::default(),
|
||||
turn_id: "turn".to_string(),
|
||||
item: TurnItem::UserMessage(UserMessageItem {
|
||||
id: "item".to_string(),
|
||||
client_id: None,
|
||||
content: Vec::new(),
|
||||
}),
|
||||
completed_at_ms: 0,
|
||||
}));
|
||||
|
||||
let (_, measurement) = measure_and_filter_rollout_items(&[item]);
|
||||
|
||||
assert_eq!(
|
||||
measurement.items[0].rollout_item_type,
|
||||
"event.item_completed.user_message"
|
||||
);
|
||||
assert_eq!(
|
||||
measurement.items[0].decision,
|
||||
super::PersistenceDecision::Dropped
|
||||
);
|
||||
}
|
||||
@@ -16,6 +16,7 @@ workspace = true
|
||||
chrono = { workspace = true, features = ["serde"] }
|
||||
codex-git-utils = { workspace = true }
|
||||
codex-install-context = { workspace = true }
|
||||
codex-otel = { workspace = true }
|
||||
codex-protocol = { workspace = true }
|
||||
codex-rollout = { workspace = true }
|
||||
codex-state = { workspace = true }
|
||||
|
||||
@@ -4,6 +4,8 @@ use std::sync::Arc;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
use codex_protocol::protocol::ThreadMemoryMode;
|
||||
use codex_rollout::RolloutPersistenceTelemetry;
|
||||
use codex_rollout::measure_and_filter_rollout_items;
|
||||
use codex_rollout::persisted_rollout_items;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::warn;
|
||||
@@ -32,6 +34,7 @@ pub struct LiveThread {
|
||||
thread_id: ThreadId,
|
||||
thread_store: Arc<dyn ThreadStore>,
|
||||
metadata_sync: Arc<Mutex<ThreadMetadataSync>>,
|
||||
persistence_telemetry: RolloutPersistenceTelemetry,
|
||||
}
|
||||
|
||||
/// Owns a live thread while session initialization is still fallible.
|
||||
@@ -95,6 +98,7 @@ impl LiveThread {
|
||||
thread_id,
|
||||
thread_store,
|
||||
metadata_sync: Arc::new(Mutex::new(metadata_sync)),
|
||||
persistence_telemetry: RolloutPersistenceTelemetry::new(thread_id),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -130,6 +134,7 @@ impl LiveThread {
|
||||
thread_id,
|
||||
thread_store,
|
||||
metadata_sync: Arc::new(Mutex::new(metadata_sync)),
|
||||
persistence_telemetry: RolloutPersistenceTelemetry::new(thread_id),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -139,16 +144,25 @@ impl LiveThread {
|
||||
fields(item_count = items.len())
|
||||
)]
|
||||
pub async fn append_items(&self, items: &[RolloutItem]) -> ThreadStoreResult<()> {
|
||||
let canonical_items = persisted_rollout_items(items);
|
||||
// Empty appends are intentionally ignored rather than represented as zero-sized batches.
|
||||
if items.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let (canonical_items, measurement) = if self.persistence_telemetry.is_enabled() {
|
||||
let (canonical_items, measurement) = measure_and_filter_rollout_items(items);
|
||||
(canonical_items, Some(measurement))
|
||||
} else {
|
||||
(persisted_rollout_items(items), None)
|
||||
};
|
||||
self.thread_store
|
||||
.append_items(AppendThreadItemsParams {
|
||||
thread_id: self.thread_id,
|
||||
items: items.to_vec(),
|
||||
})
|
||||
.await?;
|
||||
if let Some(measurement) = measurement.as_ref() {
|
||||
self.persistence_telemetry.record_batch(items, measurement);
|
||||
}
|
||||
if canonical_items.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -17,6 +17,8 @@ use crate::ResumeThreadParams;
|
||||
use crate::ThreadStoreError;
|
||||
use crate::ThreadStoreResult;
|
||||
|
||||
const ROLLOUT_SIZE_BYTES_METRIC: &str = "codex.rollout.size_bytes";
|
||||
|
||||
pub(super) async fn create_thread(
|
||||
store: &LocalThreadStore,
|
||||
params: CreateThreadParams,
|
||||
@@ -128,8 +130,15 @@ pub(super) async fn shutdown_thread(
|
||||
thread_id: ThreadId,
|
||||
) -> ThreadStoreResult<()> {
|
||||
let recorder = store.live_recorder(thread_id).await?;
|
||||
let rollout_path = recorder.rollout_path().to_path_buf();
|
||||
recorder.shutdown().await.map_err(thread_store_io_error)?;
|
||||
sync_materialized_rollout_path(store, thread_id).await?;
|
||||
if let Some(metrics) = codex_otel::global()
|
||||
&& let Ok(metadata) = tokio::fs::metadata(rollout_path).await
|
||||
{
|
||||
let size_bytes = i64::try_from(metadata.len()).unwrap_or(i64::MAX);
|
||||
let _ = metrics.histogram(ROLLOUT_SIZE_BYTES_METRIC, size_bytes, &[]);
|
||||
}
|
||||
store.live_recorders.lock().await.remove(&thread_id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user