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"), + "🦀" + ); +}