From bee28e8a061c38c781c74299f2c1d90e3732da18 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Fri, 2 Oct 2026 20:42:17 +0000 Subject: [PATCH] Cap persisted command output in paginated history at 64 KiB (#50427) ## Why Paginated history previously persisted completed commands' full `aggregated_output`, allowing large command outputs to inflate stored history. ## What changed - Limit persisted command `aggregated_output` to 64 KiB, including a truncation marker. Preserve the beginning and end of the output without splitting UTF-8 characters. - Apply the shared persistence policy to rollout recording, thread-store appends, and legacy-to-paginated history migration. - Measure persisted payload bytes after truncation so persistence metrics reflect the stored representation. ## Testing Add coverage for the byte limit, UTF-8 boundaries, unchanged inputs that fit, and markers that exceed the budget. Verify capped output in paginated thread history and migrated legacy history, and verify persistence metrics use the truncated payload size. GitOrigin-RevId: 01d420d396ad34b42dc5dc208c0bf88d890639ee --- codex-rs/Cargo.lock | 1 + .../src/protocol/thread_history.rs | 2 +- .../session_importer.rs | 4 +- codex-rs/app-server/src/request_processors.rs | 8 +- codex-rs/rollout/Cargo.toml | 1 + codex-rs/rollout/src/lib.rs | 3 +- codex-rs/rollout/src/persistence_metrics.rs | 22 ++- .../rollout/src/persistence_metrics_tests.rs | 65 ++++++++ codex-rs/rollout/src/policy.rs | 142 +++++++++++++----- codex-rs/rollout/src/recorder_tests.rs | 6 - codex-rs/thread-store/src/in_memory.rs | 4 +- .../thread-store/src/local/live_writer.rs | 6 +- .../local/rollout_migration/canonicalizer.rs | 12 +- .../src/local/rollout_migration_tests.rs | 25 +-- .../thread_history_materialization_tests.rs | 71 +++++++++ codex-rs/utils/string/src/lib.rs | 1 + codex-rs/utils/string/src/truncate.rs | 26 +++- codex-rs/utils/string/src/truncate/tests.rs | 34 +++++ 18 files changed, 353 insertions(+), 80 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index b89b547895..cb7a7ba70c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4471,6 +4471,7 @@ dependencies = [ "codex-utils-absolute-path", "codex-utils-path", "codex-utils-process", + "codex-utils-string", "memchr", "pretty_assertions", "regex", diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 28e7e6ffad..e7128a59cc 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -356,7 +356,7 @@ impl ThreadHistoryBuilder { /// tracking used by running thread resume/rejoin. /// /// This function should handle all EventMsg variants that can be persisted in a rollout file. - /// See `should_persist_event_msg` in `codex-rs/core/rollout/policy.rs`. + /// See `persisted_rollout_item` in `codex-rs/rollout/src/policy.rs`. pub fn handle_event(&mut self, event: &EventMsg) { match event { EventMsg::UserMessage(payload) => self.handle_user_message(payload), diff --git a/codex-rs/app-server/src/external_agent_migration/session_importer.rs b/codex-rs/app-server/src/external_agent_migration/session_importer.rs index 85592dc7fb..33f28ed878 100644 --- a/codex-rs/app-server/src/external_agent_migration/session_importer.rs +++ b/codex-rs/app-server/src/external_agent_migration/session_importer.rs @@ -34,7 +34,7 @@ use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::RolloutItem; -use codex_rollout::is_persisted_rollout_item; +use codex_rollout::into_persisted_rollout_items; use codex_thread_store::AppendThreadItemsParams; use codex_thread_store::CreateThreadParams; use codex_thread_store::PersistContext; @@ -497,7 +497,7 @@ impl ExternalAgentSessionImporter { memory_mode, }, }; - rollout_items.retain(|item| is_persisted_rollout_item(item, ThreadHistoryMode::Legacy)); + rollout_items = into_persisted_rollout_items(rollout_items, ThreadHistoryMode::Legacy); let (created_at, updated_at) = rollout_items .iter() .filter_map(|item| match item { diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 38110ec115..9f01ed82bd 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -473,7 +473,7 @@ use codex_rmcp_client::perform_oauth_login_return_url; use codex_rollout::InitialHistory; use codex_rollout::ResumedHistory; use codex_rollout::RolloutItem; -use codex_rollout::is_persisted_rollout_item; +use codex_rollout::persisted_rollout_item; use codex_rollout::state_db::StateDbHandle; use codex_rollout::state_db::reconcile_rollout; use codex_state::ThreadMetadata; @@ -722,8 +722,10 @@ pub(crate) use self::thread_summary::thread_settings_from_config_snapshot; pub(crate) fn build_legacy_api_turns_from_rollout_items(items: &[RolloutItem]) -> Vec { let mut builder = ThreadHistoryBuilder::new(); for item in items { - if is_persisted_rollout_item(item, codex_protocol::protocol::ThreadHistoryMode::Legacy) { - builder.handle_rollout_item(item); + if let Some(item) = + persisted_rollout_item(item, codex_protocol::protocol::ThreadHistoryMode::Legacy) + { + builder.handle_rollout_item(item.as_ref()); } } builder.finish() diff --git a/codex-rs/rollout/Cargo.toml b/codex-rs/rollout/Cargo.toml index fb2b5b39e4..6ce9637f65 100644 --- a/codex-rs/rollout/Cargo.toml +++ b/codex-rs/rollout/Cargo.toml @@ -25,6 +25,7 @@ codex-protocol = { workspace = true } codex-state = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-path = { workspace = true } +codex-utils-string = { workspace = true } memchr = { workspace = true } regex = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 6bcbf57fde..6a068945eb 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -143,7 +143,8 @@ pub use model_context::ModelContextScanProgress; 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::into_persisted_rollout_items; +pub use policy::persisted_rollout_item; pub use policy::persisted_rollout_items; pub use policy::should_persist_response_item; pub use policy::should_persist_response_item_for_memories; diff --git a/codex-rs/rollout/src/persistence_metrics.rs b/codex-rs/rollout/src/persistence_metrics.rs index b21072c729..ad74af3312 100644 --- a/codex-rs/rollout/src/persistence_metrics.rs +++ b/codex-rs/rollout/src/persistence_metrics.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::io::Write; use std::sync::Arc; use std::sync::Mutex; @@ -11,7 +12,7 @@ use codex_protocol::protocol::ThreadHistoryMode; use crate::ResponseItemEnvelope; use crate::RolloutItem; -use crate::policy::is_persisted_rollout_item; +use crate::policy::persisted_rollout_item; const ITEM_BYTES_METRIC: &str = "codex.rollout.persistence.item_bytes"; const APPEND_METRIC: &str = "codex.rollout.persistence.append"; @@ -46,6 +47,7 @@ pub struct RolloutItemMeasurement { pub decision: PersistenceDecision, pub rollout_item_type: String, pub payload_bytes: Option, + pub persisted_payload_bytes: Option, } #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -106,22 +108,28 @@ pub fn measure_and_filter_rollout_items( }; for item in items { - let kept = is_persisted_rollout_item(item, history_mode); - let decision = if kept { + let projected = persisted_rollout_item(item, history_mode); + let decision = if projected.is_some() { PersistenceDecision::Kept } else { PersistenceDecision::Dropped }; let payload_bytes = serialized_len(item).ok(); + let persisted_payload_bytes = match projected.as_ref() { + Some(Cow::Borrowed(_)) => payload_bytes, + Some(Cow::Owned(item)) => serialized_len(item).ok(), + None => None, + }; add_to_totals(&mut measurement.pre_filter, payload_bytes); - if kept { - add_to_totals(&mut measurement.post_filter, payload_bytes); - persisted.push(item.clone()); + if let Some(projected) = projected { + add_to_totals(&mut measurement.post_filter, persisted_payload_bytes); + persisted.push(projected.into_owned()); } measurement.items.push(RolloutItemMeasurement { decision, rollout_item_type: rollout_item_type(item), payload_bytes, + persisted_payload_bytes, }); } @@ -199,7 +207,7 @@ fn finish_turn( 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); + add_to_totals(&mut totals.post_filter, item.persisted_payload_bytes); } } diff --git a/codex-rs/rollout/src/persistence_metrics_tests.rs b/codex-rs/rollout/src/persistence_metrics_tests.rs index 02a600a8e1..f6af5af10d 100644 --- a/codex-rs/rollout/src/persistence_metrics_tests.rs +++ b/codex-rs/rollout/src/persistence_metrics_tests.rs @@ -1,12 +1,18 @@ +use std::time::Duration; + use codex_protocol::ThreadId; +use codex_protocol::items::CommandExecutionItem; +use codex_protocol::items::CommandExecutionStatus; use codex_protocol::items::EnteredReviewModeItem; use codex_protocol::items::ExitedReviewModeItem; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; +use codex_protocol::parse_command::ParsedCommand; use codex_protocol::protocol::EnteredReviewModeEvent; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ExecCommandSource; use codex_protocol::protocol::ExitedReviewModeEvent; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::ReviewTarget; @@ -15,6 +21,7 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnStartedEvent; +use codex_utils_absolute_path::test_support::PathExt; use pretty_assertions::assert_eq; use super::CompletedTurnMeasurement; @@ -26,6 +33,7 @@ use super::measure_and_filter_rollout_items; use super::update_turn_measurements; use crate::ResponseItemEnvelope; use crate::RolloutItem; +use crate::policy::PERSISTED_COMMAND_OUTPUT_MAX_BYTES; fn retained_message(text: &str) -> RolloutItem { RolloutItem::ResponseItem(ResponseItemEnvelope::new(ResponseItem::Message { @@ -290,6 +298,63 @@ fn item_completion_persistence_depends_on_history_mode() { ); } +#[test] +fn projected_items_report_post_projection_bytes() { + let original_output = "x".repeat(PERSISTED_COMMAND_OUTPUT_MAX_BYTES * 2); + let item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent { + thread_id: ThreadId::default(), + turn_id: "turn".to_string(), + item: TurnItem::CommandExecution(CommandExecutionItem { + sandbox_type: None, + model_context: None, + id: "exec".to_string(), + plugin_id: None, + script_path: None, + process_id: None, + command: vec!["echo".to_string()], + cwd: std::env::temp_dir().abs().into(), + parsed_cmd: vec![ParsedCommand::Unknown { + cmd: "echo".to_string(), + }], + source: ExecCommandSource::Agent, + interaction_input: None, + status: CommandExecutionStatus::Completed, + aggregated_output: Some(original_output), + exit_code: Some(0), + duration: Some(Duration::ZERO), + }), + started_at_ms: Some(0), + completed_at_ms: 0, + })); + let original_bytes = serde_json::to_vec(&item) + .expect("serialize original command") + .len() as u64; + + let (persisted, measurement) = + measure_and_filter_rollout_items(std::slice::from_ref(&item), ThreadHistoryMode::Paginated); + let persisted_bytes = serde_json::to_vec(&persisted[0]) + .expect("serialize persisted command") + .len() as u64; + assert_eq!( + measurement.pre_filter, + super::RolloutSizeTotals { + items: 1, + payload_bytes: original_bytes, + } + ); + assert_eq!( + measurement.post_filter, + super::RolloutSizeTotals { + items: 1, + payload_bytes: persisted_bytes, + } + ); + assert!( + persisted_bytes < original_bytes, + "projection should reduce serialized bytes" + ); +} + #[test] fn review_mode_persistence_depends_on_history_mode() { let completed_items = vec![ diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index c82b5ce958..7eb21511d2 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -1,19 +1,40 @@ +use std::borrow::Cow; + use crate::RolloutItem; use crate::protocol::EventMsg; use codex_extension_items::ExtensionItem; +use codex_protocol::items::CommandExecutionItem; use codex_protocol::items::TurnItem; use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::SubAgentActivityKind; use codex_protocol::protocol::ThreadHistoryMode; +use codex_utils_string::truncate_middle_with_marker; -/// Whether a rollout `item` should be persisted in rollout files. -pub fn is_persisted_rollout_item(item: &RolloutItem, history_mode: ThreadHistoryMode) -> bool { +pub(super) const PERSISTED_COMMAND_OUTPUT_MAX_BYTES: usize = 64 * 1024; +const PERSISTED_COMMAND_OUTPUT_TRUNCATION_MARKER: &str = + "\n... command output truncated for persistence ...\n"; + +/// Returns the authoritative durable representation of a rollout item. +pub fn persisted_rollout_item( + item: &RolloutItem, + history_mode: ThreadHistoryMode, +) -> Option> { match item { - RolloutItem::ResponseItem(item) => should_persist_response_item(&item.item), + RolloutItem::ResponseItem(response_item) => { + should_persist_response_item(&response_item.item).then_some(Cow::Borrowed(item)) + } RolloutItem::InterAgentCommunication(_) - | RolloutItem::InterAgentCommunicationMetadata { .. } => true, - RolloutItem::EventMsg(ev) => should_persist_event_msg(ev, history_mode), - RolloutItem::RealtimeItem(_) => matches!(history_mode, ThreadHistoryMode::Paginated), + | RolloutItem::InterAgentCommunicationMetadata { .. } => Some(Cow::Borrowed(item)), + RolloutItem::EventMsg(ev) => { + persisted_event_msg(ev, history_mode).map(|event| match event { + Cow::Borrowed(_) => Cow::Borrowed(item), + Cow::Owned(event) => Cow::Owned(RolloutItem::EventMsg(event)), + }) + } + RolloutItem::RealtimeItem(_) => { + matches!(history_mode, ThreadHistoryMode::Paginated).then_some(Cow::Borrowed(item)) + } // Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns). RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) @@ -21,7 +42,7 @@ pub fn is_persisted_rollout_item(item: &RolloutItem, history_mode: ThreadHistory | RolloutItem::WorldState(_) | RolloutItem::RetainedContext(_) | RolloutItem::SecurityRiskScore(_) - | RolloutItem::SessionMeta(_) => true, + | RolloutItem::SessionMeta(_) => Some(Cow::Borrowed(item)), } } @@ -30,13 +51,25 @@ pub fn persisted_rollout_items( items: &[RolloutItem], history_mode: ThreadHistoryMode, ) -> Vec { - let mut persisted = Vec::new(); - for item in items { - if is_persisted_rollout_item(item, history_mode) { - persisted.push(item.clone()); - } - } - persisted + items + .iter() + .filter_map(|item| persisted_rollout_item(item, history_mode).map(Cow::into_owned)) + .collect() +} + +/// Return the owned rollout items that should be persisted for a live append. +pub fn into_persisted_rollout_items( + items: Vec, + history_mode: ThreadHistoryMode, +) -> Vec { + items + .into_iter() + .filter_map(|item| match persisted_rollout_item(&item, history_mode) { + Some(Cow::Borrowed(_)) => Some(item), + Some(Cow::Owned(item)) => Some(item), + None => None, + }) + .collect() } /// Whether a `ResponseItem` should be persisted in rollout files. @@ -89,34 +122,19 @@ pub fn should_persist_response_item_for_memories(item: &ResponseItem) -> bool { } } -/// Whether an `EventMsg` should be persisted in rollout files. -#[inline] -pub fn should_persist_event_msg(ev: &EventMsg, history_mode: ThreadHistoryMode) -> bool { +fn persisted_event_msg( + ev: &EventMsg, + history_mode: ThreadHistoryMode, +) -> Option> { match ev { - EventMsg::ItemCompleted(event) => { - // Paginated rollouts store TurnItems. - // Legacy rollouts keep only items with no lossless raw ResponseItem or legacy - // equivalent. - matches!(history_mode, ThreadHistoryMode::Paginated) - || matches!( - event.item, - TurnItem::FunctionCallOutput(_) - | TurnItem::Plan(_) - | TurnItem::Extension(ExtensionItem::Sleep(_)) - ) - || matches!( - &event.item, - TurnItem::SubAgentActivity(item) - if item.kind == SubAgentActivityKind::Completed - ) - } + EventMsg::ItemCompleted(event) => persisted_item_completed_event(ev, event, history_mode), EventMsg::TokenCount(_) | EventMsg::ThreadGoalUpdated(_) | EventMsg::ThreadRolledBack(_) | EventMsg::TurnAborted(_) | EventMsg::TurnStarted(_) | EventMsg::TurnComplete(_) - | EventMsg::ThreadSettingsApplied(_) => true, + | EventMsg::ThreadSettingsApplied(_) => Some(Cow::Borrowed(ev)), // Only persist these legacy events when the thread's history mode is Legacy. // New, paginated rollouts persist ItemCompleted events with TurnItems. @@ -131,12 +149,11 @@ pub fn should_persist_event_msg(ev: &EventMsg, history_mode: ThreadHistoryMode) | EventMsg::McpToolCallEnd(_) | EventMsg::WebSearchEnd(_) | EventMsg::ImageGenerationEnd(_) => { - matches!(history_mode, ThreadHistoryMode::Legacy) - } - EventMsg::SubAgentActivity(event) => { - matches!(history_mode, ThreadHistoryMode::Legacy) - && event.kind != SubAgentActivityKind::Completed + matches!(history_mode, ThreadHistoryMode::Legacy).then_some(Cow::Borrowed(ev)) } + EventMsg::SubAgentActivity(event) => (matches!(history_mode, ThreadHistoryMode::Legacy) + && event.kind != SubAgentActivityKind::Completed) + .then_some(Cow::Borrowed(ev)), // Transient, non-durable events. EventMsg::Error(_) @@ -201,6 +218,49 @@ pub fn should_persist_event_msg(ev: &EventMsg, history_mode: ThreadHistoryMode) | EventMsg::CollabAgentInteractionBegin(_) | EventMsg::CollabWaitingBegin(_) | EventMsg::CollabCloseBegin(_) - | EventMsg::CollabResumeBegin(_) => false, + | EventMsg::CollabResumeBegin(_) => None, + } +} + +/// Returns the persisted representation of an `ItemCompleted` event. +fn persisted_item_completed_event<'a>( + ev: &'a EventMsg, + event: &ItemCompletedEvent, + history_mode: ThreadHistoryMode, +) -> Option> { + match history_mode { + ThreadHistoryMode::Legacy => { + // Legacy rollouts keep only items with no lossless raw ResponseItem or legacy + // equivalent. + (matches!( + event.item, + TurnItem::FunctionCallOutput(_) + | TurnItem::Plan(_) + | TurnItem::Extension(ExtensionItem::Sleep(_)) + ) || matches!( + &event.item, + TurnItem::SubAgentActivity(item) + if item.kind == SubAgentActivityKind::Completed + )) + .then_some(Cow::Borrowed(ev)) + } + ThreadHistoryMode::Paginated => match &event.item { + TurnItem::CommandExecution(CommandExecutionItem { + aggregated_output: Some(aggregated_output), + .. + }) if aggregated_output.len() > PERSISTED_COMMAND_OUTPUT_MAX_BYTES => { + let aggregated_output = truncate_middle_with_marker( + aggregated_output, + PERSISTED_COMMAND_OUTPUT_MAX_BYTES, + PERSISTED_COMMAND_OUTPUT_TRUNCATION_MARKER, + ); + let mut event = event.clone(); + if let TurnItem::CommandExecution(command) = &mut event.item { + command.aggregated_output = Some(aggregated_output); + } + Some(Cow::Owned(EventMsg::ItemCompleted(event))) + } + _ => Some(Cow::Borrowed(ev)), + }, } } diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index af55bf6ddf..e57b8d99b5 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -458,12 +458,6 @@ async fn load_rollout_items_preserves_security_risk_scores() -> std::io::Result< sampled_at: None, }; let security_risk_item = RolloutItem::SecurityRiskScore(security_risk.clone()); - for history_mode in [ThreadHistoryMode::Legacy, ThreadHistoryMode::Paginated] { - assert!(crate::is_persisted_rollout_item( - &security_risk_item, - history_mode - )); - } let mut file = File::create(&rollout_path)?; for (ordinal, item) in [ diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 1e670fee77..152c1d0021 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -19,7 +19,7 @@ use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::RolloutItem; -use codex_rollout::persisted_rollout_items; +use codex_rollout::into_persisted_rollout_items; use codex_utils_absolute_path::AbsolutePathBuf; use crate::AppendThreadItemsParams; @@ -653,7 +653,7 @@ impl InMemoryThreadStore { } let mut state = self.state.lock().await; let history_mode = history_mode_from_state(&state, params.thread_id); - let persisted_items = persisted_rollout_items(params.items.as_slice(), history_mode); + let persisted_items = into_persisted_rollout_items(params.items, history_mode); if persisted_items.is_empty() { return Ok(()); } diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index f78dba4531..e1b49e82b7 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -8,7 +8,7 @@ use codex_rollout::RolloutConfig; use codex_rollout::RolloutItem; use codex_rollout::RolloutRecorder; use codex_rollout::RolloutRecorderParams; -use codex_rollout::is_persisted_rollout_item; +use codex_rollout::into_persisted_rollout_items; use tracing::warn; use super::LocalThreadStore; @@ -336,8 +336,8 @@ async fn write_and_project( let (recorder, rollout_id, history_mode) = live_writer_parts(store, thread_id).await?; let sync_rollout_path = matches!(&write_op, RolloutWriteOp::Persist | RolloutWriteOp::Flush); let write_op = match write_op { - RolloutWriteOp::AppendItems(mut items) => { - items.retain(|item| is_persisted_rollout_item(item, history_mode)); + RolloutWriteOp::AppendItems(items) => { + let items = into_persisted_rollout_items(items, history_mode); if items.is_empty() { return Ok(()); } diff --git a/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs b/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs index 8c446ccc5c..1575bcf090 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs @@ -22,6 +22,7 @@ use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnStartedEvent; use codex_rollout::RolloutItem; use codex_rollout::RolloutLine; +use std::borrow::Cow; use std::collections::HashSet; use tokio::fs::File; use tokio::io::AsyncWriteExt; @@ -278,10 +279,7 @@ impl LegacyRolloutCanonicalizer { } } else { let item = RolloutItem::EventMsg(event); - if codex_rollout::is_persisted_rollout_item(&item, ThreadHistoryMode::Paginated) - { - self.write_item(writer, ×tamp, item).await?; - } + self.write_item(writer, ×tamp, item).await?; } } item @ RolloutItem::InterAgentCommunication(_) => { @@ -466,6 +464,12 @@ impl LegacyRolloutCanonicalizer { timestamp: &str, item: RolloutItem, ) -> ThreadStoreResult<()> { + let item = match codex_rollout::persisted_rollout_item(&item, ThreadHistoryMode::Paginated) + { + Some(Cow::Borrowed(_)) => item, + Some(Cow::Owned(item)) => item, + None => return Ok(()), + }; let mut bytes = serde_json::to_vec(&RolloutLine { timestamp: timestamp.to_string(), ordinal: Some(self.next_ordinal), diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index c040d43a52..90e2f11512 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -188,7 +188,7 @@ fn rollout_response_item(item: ResponseItem) -> RolloutItem { RolloutItem::ResponseItem(item.into()) } -fn exec_completion(turn_id: &str, call_id: &str) -> RolloutItem { +fn exec_completion(turn_id: &str, call_id: &str, output: &str) -> RolloutItem { serde_json::from_value(json!({ "type": "event_msg", "payload": { @@ -199,12 +199,12 @@ fn exec_completion(turn_id: &str, call_id: &str) -> RolloutItem { "cwd": "file:///tmp", "parsed_cmd": [], "source": "agent", - "stdout": "ok", + "stdout": output, "stderr": "", - "aggregated_output": "ok", + "aggregated_output": output, "exit_code": 0, "duration": {"secs": 0, "nanos": 0}, - "formatted_output": "ok", + "formatted_output": output, "status": "completed" } })) @@ -417,7 +417,8 @@ async fn migration_publishes_canonical_projected_history_and_is_idempotent() { async fn migration_projects_explicit_and_implicit_legacy_completed_items() { let home = TempDir::new().expect("create Codex home"); let thread_id = ThreadId::new(); - let exec = exec_completion("explicit", "call-1"); + let output = format!("head\n{}\ntail", "x".repeat(128 * 1024)); + let exec = exec_completion("explicit", "call-1", &output); let reasoning = serde_json::from_value(json!({ "type": "event_msg", "payload": {"type": "agent_reasoning", "text": "summary"} @@ -471,7 +472,13 @@ async fn migration_projects_explicit_and_implicit_legacy_completed_items() { let reasoning: serde_json::Value = serde_json::from_slice(&items.items[1].item_json).expect("parse projected reasoning"); assert_eq!(command["type"], "commandExecution"); - assert_eq!(command["aggregatedOutput"], "ok"); + let aggregated_output = command["aggregatedOutput"] + .as_str() + .expect("migrated command output"); + assert_eq!(aggregated_output.len(), 64 * 1024); + assert!(aggregated_output.starts_with("head\n")); + assert!(aggregated_output.ends_with("\ntail")); + assert!(aggregated_output.contains("command output truncated for persistence")); assert_eq!(reasoning["type"], "reasoning"); assert_eq!(reasoning["summary"], json!(["summary"])); assert_eq!(reasoning["content"], json!(["raw"])); @@ -548,7 +555,7 @@ async fn migration_keeps_late_completions_in_their_original_turn() { user_message("old question"), started("current"), user_message("current question"), - exec_completion("old", "call-old"), + exec_completion("old", "call-old", "ok"), serde_json::from_value(json!({ "type": "event_msg", "payload": { @@ -989,9 +996,9 @@ async fn migration_keeps_late_completions_for_surviving_turns_across_rollback() started("remove"), user_message("remove question"), completed("old"), - exec_completion("old", "call-old"), + exec_completion("old", "call-old", "ok"), item_completed("old", "reason-old"), - exec_completion("remove", "call-remove"), + exec_completion("remove", "call-remove", "ok"), completed("remove"), RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns: 1, diff --git a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs index 07713ea56c..e7a663ffaf 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs @@ -8,13 +8,17 @@ use codex_app_server_protocol::ThreadItem; use codex_protocol::ThreadId; use codex_protocol::items::AgentMessageContent; use codex_protocol::items::AgentMessageItem; +use codex_protocol::items::CommandExecutionItem; +use codex_protocol::items::CommandExecutionStatus; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; +use codex_protocol::parse_command::ParsedCommand; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ExecCommandSource; use codex_protocol::protocol::HistoryPosition; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::RateLimitSnapshot; @@ -434,6 +438,73 @@ WHERE thread_id = ? assert_eq!(projection_state, (rollout_len, 5)); } +#[tokio::test] +async fn paginated_command_history_caps_aggregated_output() { + let home = TempDir::new().expect("temp dir"); + let store = projection_store(home.path()).await; + let thread_id = ThreadId::default(); + create_paginated_thread(&store, thread_id).await; + let output = format!("head\n{}\ntail", "x".repeat(128 * 1024)); + let command = CommandExecutionItem { + sandbox_type: None, + model_context: None, + id: "exec-1".to_string(), + plugin_id: None, + script_path: None, + process_id: None, + command: vec!["echo".to_string(), "hello".to_string()], + cwd: home.path().abs().into(), + parsed_cmd: vec![ParsedCommand::Unknown { + cmd: "echo hello".to_string(), + }], + source: ExecCommandSource::Agent, + interaction_input: None, + status: CommandExecutionStatus::Completed, + aggregated_output: Some(output), + exit_code: Some(0), + duration: Some(Duration::from_millis(12)), + }; + + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![ + turn_started("turn-1"), + completed_item(thread_id, "turn-1", TurnItem::CommandExecution(command)), + turn_completed("turn-1"), + ], + }) + .await + .expect("append command history"); + + let pool = codex_state::open_thread_history_db(&codex_state::SqliteConfig::new_for_testing( + home.path().abs(), + )) + .await + .expect("open thread history db"); + let item_json = sqlx::query_scalar::<_, String>( + "SELECT item_json FROM thread_items WHERE thread_id = ? AND item_id = ?", + ) + .bind(thread_id.to_string()) + .bind("exec-1") + .fetch_one(&pool) + .await + .expect("read projected command"); + let projected: ThreadItem = + serde_json::from_str(&item_json).expect("deserialize projected command"); + let ThreadItem::CommandExecution { + aggregated_output: Some(output), + .. + } = projected + else { + panic!("expected projected command output"); + }; + assert_eq!(output.len(), 64 * 1024); + assert!(output.starts_with("head\n")); + assert!(output.ends_with("\ntail")); + assert!(output.contains("command output truncated for persistence")); +} + #[tokio::test] async fn paginated_realtime_items_materialize_separately_in_rollout_order() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/utils/string/src/lib.rs b/codex-rs/utils/string/src/lib.rs index 1618739b08..cacbb9e2e5 100644 --- a/codex-rs/utils/string/src/lib.rs +++ b/codex-rs/utils/string/src/lib.rs @@ -7,6 +7,7 @@ pub use truncate::approx_bytes_for_tokens; pub use truncate::approx_token_count; pub use truncate::approx_tokens_from_byte_count; pub use truncate::truncate_middle_chars; +pub use truncate::truncate_middle_with_marker; pub use truncate::truncate_middle_with_token_budget; // Truncate a &str to a byte budget at a char boundary (prefix) diff --git a/codex-rs/utils/string/src/truncate.rs b/codex-rs/utils/string/src/truncate.rs index c0d1dd67d5..075d78fd7a 100644 --- a/codex-rs/utils/string/src/truncate.rs +++ b/codex-rs/utils/string/src/truncate.rs @@ -3,11 +3,35 @@ const APPROX_BYTES_PER_TOKEN: usize = 4; -/// Truncate a string to `max_bytes` using a character-count marker. +/// Retain at most `max_bytes` from a string's ends and add a character-count +/// marker outside that budget. +/// +/// Use [`truncate_middle_with_marker`] when the complete result must fit a hard +/// byte limit. pub fn truncate_middle_chars(s: &str, max_bytes: usize) -> String { truncate_with_byte_estimate(s, max_bytes, /*use_tokens*/ false) } +/// Truncate the middle of a UTF-8 string so the result, including `marker`, is +/// at most `max_bytes` bytes. +/// +/// Returns the original string when it already fits. If the marker itself does +/// not fit, returns the longest original prefix that does. +pub fn truncate_middle_with_marker(s: &str, max_bytes: usize, marker: &str) -> String { + if s.len() <= max_bytes { + return s.to_string(); + } + + let Some(content_budget) = max_bytes.checked_sub(marker.len()) else { + return s[..s.floor_char_boundary(max_bytes)].to_string(); + }; + let (left_budget, right_budget) = split_budget(content_budget); + let prefix_end = s.floor_char_boundary(left_budget); + let suffix_start = s.ceil_char_boundary(s.len().saturating_sub(right_budget)); + + assemble_truncated_output(&s[..prefix_end], &s[suffix_start..], marker) +} + /// Truncate the middle of a UTF-8 string to at most `max_tokens` approximate /// tokens, preserving the beginning and the end. Returns the possibly /// truncated string and `Some(original_token_count)` if truncation occurred; diff --git a/codex-rs/utils/string/src/truncate/tests.rs b/codex-rs/utils/string/src/truncate/tests.rs index a8b5fd3d13..eaa90257a6 100644 --- a/codex-rs/utils/string/src/truncate/tests.rs +++ b/codex-rs/utils/string/src/truncate/tests.rs @@ -1,5 +1,6 @@ use super::split_string; use super::truncate_middle_chars; +use super::truncate_middle_with_marker; use super::truncate_middle_with_token_budget; use pretty_assertions::assert_eq; @@ -129,3 +130,36 @@ fn truncate_middle_bytes_handles_utf8_content() { let out = truncate_middle_chars(s, /*max_bytes*/ 20); assert_eq!(out, "😀😀…21 chars truncated…with text\n"); } + +#[test] +fn truncate_middle_with_marker_enforces_the_total_byte_limit() { + let input = format!("head{}tail", "x".repeat(100)); + assert_eq!( + truncate_middle_with_marker(&input, /*max_bytes*/ 32, ""), + "headxxxxxxxxxxxxxtail" + ); +} + +#[test] +fn truncate_middle_with_marker_preserves_utf8_boundaries() { + assert_eq!( + truncate_middle_with_marker("head🦀🦀🦀🦀🦀tail", /*max_bytes*/ 18, "..."), + "head...🦀tail" + ); +} + +#[test] +fn truncate_middle_with_marker_returns_input_that_fits() { + assert_eq!( + truncate_middle_with_marker("short", /*max_bytes*/ 5, "..."), + "short" + ); +} + +#[test] +fn truncate_middle_with_marker_keeps_a_prefix_when_the_marker_does_not_fit() { + assert_eq!( + truncate_middle_with_marker("🦀abc", /*max_bytes*/ 4, "marker"), + "🦀" + ); +}