From 9c9326d6b3bf4ed8ffc4d9979815ae095f48ce01 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Tue, 23 Jun 2026 06:24:13 +0000 Subject: [PATCH] rollout: own turn lifecycle replay --- codex-rs/core/Cargo.toml | 2 +- codex-rs/core/src/thread_manager.rs | 32 +- codex-rs/rollout/src/lib.rs | 4 + codex-rs/rollout/src/turn_lifecycle.rs | 227 +++++++++++ codex-rs/rollout/src/turn_lifecycle_tests.rs | 386 +++++++++++++++++++ 5 files changed, 631 insertions(+), 20 deletions(-) create mode 100644 codex-rs/rollout/src/turn_lifecycle.rs create mode 100644 codex-rs/rollout/src/turn_lifecycle_tests.rs diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 4435a2eb63..ee42f551d4 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -26,7 +26,6 @@ clap = { workspace = true, features = ["derive"] } codex-analytics = { workspace = true } codex-agent-graph-store = { workspace = true } codex-api = { workspace = true } -codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } codex-code-mode = { workspace = true } @@ -135,6 +134,7 @@ codex-shell-escalation = { workspace = true } [dev-dependencies] assert_cmd = { workspace = true } assert_matches = { workspace = true } +codex-app-server-protocol = { workspace = true } codex-image-generation-extension = { workspace = true } codex-home = { workspace = true } codex-otel = { workspace = true } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 4846bfd890..a1d382455e 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -19,8 +19,6 @@ use crate::tasks::interrupted_turn_history_marker; use codex_agent_graph_store::AgentGraphStore; use codex_agent_graph_store::LocalAgentGraphStore; use codex_analytics::AnalyticsEventsClient; -use codex_app_server_protocol::ThreadHistoryBuilder; -use codex_app_server_protocol::TurnStatus; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; @@ -59,6 +57,8 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; +use codex_rollout::ExplicitTurnState; +use codex_rollout::RolloutTurnLifecycleTracker; use codex_rollout::state_db::StateDbHandle; use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LocalThreadStore; @@ -1768,28 +1768,22 @@ struct SnapshotTurnState { fn snapshot_turn_state(history: &InitialHistory) -> SnapshotTurnState { let rollout_items = history.get_rollout_items(); - let mut builder = ThreadHistoryBuilder::new(); + let mut tracker = RolloutTurnLifecycleTracker::new(); for item in rollout_items { - builder.handle_rollout_item(item); + tracker.handle_rollout_item(item); } - let active_turn_id = builder.active_turn_id_if_explicit(); - if builder.has_active_turn() && active_turn_id.is_some() { - let active_turn_snapshot = builder.active_turn_snapshot(); - if active_turn_snapshot - .as_ref() - .is_some_and(|turn| turn.status != TurnStatus::InProgress) - { - return SnapshotTurnState { + if let Some(active_turn) = tracker.current_explicit_turn() { + return match active_turn.state { + ExplicitTurnState::InProgress => SnapshotTurnState { + ends_mid_turn: true, + active_turn_id: Some(active_turn.turn_id.clone()), + active_turn_start_index: Some(active_turn.rollout_start_index), + }, + ExplicitTurnState::Terminal => SnapshotTurnState { ends_mid_turn: false, active_turn_id: None, active_turn_start_index: None, - }; - } - - return SnapshotTurnState { - ends_mid_turn: true, - active_turn_id, - active_turn_start_index: builder.active_turn_start_index(), + }, }; } diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 3c103750c1..2ddcbaa156 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -15,6 +15,7 @@ pub(crate) mod search; pub(crate) mod session_index; mod sqlite_metrics; pub mod state_db; +mod turn_lifecycle; pub(crate) use codex_protocol::protocol; @@ -76,6 +77,9 @@ pub use session_index::find_thread_names_by_ids; pub use session_index::remove_thread_name_entries; pub use state_db::StateDbHandle; pub use state_db::sqlite_telemetry_recorder; +pub use turn_lifecycle::CurrentExplicitTurn; +pub use turn_lifecycle::ExplicitTurnState; +pub use turn_lifecycle::RolloutTurnLifecycleTracker; #[cfg(test)] mod tests; diff --git a/codex-rs/rollout/src/turn_lifecycle.rs b/codex-rs/rollout/src/turn_lifecycle.rs new file mode 100644 index 0000000000..eb30dd8929 --- /dev/null +++ b/codex-rs/rollout/src/turn_lifecycle.rs @@ -0,0 +1,227 @@ +use codex_protocol::items::parse_hook_prompt_message; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; + +/// Whether the current explicit turn is still running or has reached a terminal state. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ExplicitTurnState { + InProgress, + Terminal, +} + +/// The explicit turn currently open at the end of the observed rollout. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct CurrentExplicitTurn { + /// Identifier carried by the `TurnStarted` event. + pub turn_id: String, + /// Zero-based index of that event in the raw rollout stream. + pub rollout_start_index: usize, + /// Whether the turn is still active or has reached a terminal boundary. + pub state: ExplicitTurnState, +} + +/// Tracks turn lifecycle boundaries in persisted rollout order. +/// +/// This intentionally models only lifecycle state. Implicit turns are retained as +/// placeholders so rollback counts and late explicit turn IDs remain correlated +/// without reconstructing an app-server presentation history. +#[derive(Debug, Default)] +pub struct RolloutTurnLifecycleTracker { + finished_turns: Vec, + current_turn: Option, + next_rollout_index: usize, +} + +#[derive(Debug)] +enum CurrentTurn { + Explicit(CurrentExplicitTurn), + Implicit(ImplicitTurnState), +} + +#[derive(Debug)] +enum FinishedTurn { + Explicit(String), + Implicit, +} + +#[derive(Debug)] +enum ImplicitTurnState { + CompactionOnly, + Materialized, +} + +impl RolloutTurnLifecycleTracker { + /// Create an empty lifecycle tracker. + pub fn new() -> Self { + Self::default() + } + + /// Observe one rollout item in persisted order. + pub fn handle_rollout_item(&mut self, item: &RolloutItem) { + let rollout_index = self.next_rollout_index; + self.next_rollout_index += 1; + + if matches!(item, RolloutItem::Compacted(_)) { + self.handle_compacted(); + return; + } + + if let RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::Message { + role, + content, + id, + .. + }) = item + && role == "user" + && parse_hook_prompt_message(id.as_ref(), content).is_some() + { + self.materialize_implicit_turn(); + return; + } + + let RolloutItem::EventMsg(event) = item else { + return; + }; + + match event { + EventMsg::UserMessage(_) => self.handle_implicit_user_turn(), + EventMsg::AgentMessage(event) if !event.message.is_empty() => { + self.materialize_implicit_turn(); + } + EventMsg::AgentReasoning(event) if !event.text.is_empty() => { + self.materialize_implicit_turn(); + } + EventMsg::AgentReasoningRawContent(event) if !event.text.is_empty() => { + self.materialize_implicit_turn(); + } + EventMsg::PatchApplyEnd(event) if event.turn_id.is_empty() => { + self.materialize_implicit_turn(); + } + EventMsg::ContextCompacted(_) + | EventMsg::EnteredReviewMode(_) + | EventMsg::ExitedReviewMode(_) + | EventMsg::McpToolCallEnd(_) + | EventMsg::WebSearchEnd(_) + | EventMsg::ImageGenerationEnd(_) + | EventMsg::SubAgentActivity(_) => self.materialize_implicit_turn(), + EventMsg::TurnStarted(event) => { + self.finish_current_turn(); + self.current_turn = Some(CurrentTurn::Explicit(CurrentExplicitTurn { + turn_id: event.turn_id.clone(), + rollout_start_index: rollout_index, + state: ExplicitTurnState::InProgress, + })); + } + EventMsg::TurnComplete(event) => self.handle_turn_complete(&event.turn_id), + EventMsg::TurnAborted(event) => self.handle_turn_aborted(event.turn_id.as_deref()), + EventMsg::Error(event) if event.affects_turn_status() => { + self.mark_current_explicit_turn_terminal(); + } + EventMsg::ThreadRolledBack(event) => { + self.finish_current_turn(); + let num_turns = usize::try_from(event.num_turns).unwrap_or(usize::MAX); + self.finished_turns + .truncate(self.finished_turns.len().saturating_sub(num_turns)); + } + _ => {} + } + } + + /// Return the explicitly opened turn that remains current, if any. + pub fn current_explicit_turn(&self) -> Option<&CurrentExplicitTurn> { + match self.current_turn.as_ref() { + Some(CurrentTurn::Explicit(turn)) => Some(turn), + Some(CurrentTurn::Implicit(_)) | None => None, + } + } + + fn handle_implicit_user_turn(&mut self) { + if matches!(self.current_turn.as_ref(), Some(CurrentTurn::Explicit(_))) { + return; + } + if matches!( + self.current_turn.as_ref(), + Some(CurrentTurn::Implicit(ImplicitTurnState::CompactionOnly)) + ) { + self.current_turn = Some(CurrentTurn::Implicit(ImplicitTurnState::Materialized)); + return; + } + if self.current_turn.is_some() { + self.finish_current_turn(); + } + self.current_turn = Some(CurrentTurn::Implicit(ImplicitTurnState::Materialized)); + } + + fn handle_compacted(&mut self) { + if self.current_turn.is_none() { + self.current_turn = Some(CurrentTurn::Implicit(ImplicitTurnState::CompactionOnly)); + } + } + + fn materialize_implicit_turn(&mut self) { + match self.current_turn.as_mut() { + Some(CurrentTurn::Explicit(_)) => {} + Some(CurrentTurn::Implicit(state)) => *state = ImplicitTurnState::Materialized, + None => { + self.current_turn = Some(CurrentTurn::Implicit(ImplicitTurnState::Materialized)); + } + } + } + + fn handle_turn_complete(&mut self, turn_id: &str) { + if self.current_explicit_turn_has_id(turn_id) { + self.finish_current_turn(); + return; + } + + if self.finished_turn_has_id(turn_id) { + return; + } + + self.finish_current_turn(); + } + + fn handle_turn_aborted(&mut self, turn_id: Option<&str>) { + if turn_id.is_some_and(|turn_id| self.current_explicit_turn_has_id(turn_id)) { + self.mark_current_explicit_turn_terminal(); + return; + } + + if turn_id.is_some_and(|turn_id| self.finished_turn_has_id(turn_id)) { + return; + } + + self.mark_current_explicit_turn_terminal(); + } + + fn mark_current_explicit_turn_terminal(&mut self) { + if let Some(CurrentTurn::Explicit(turn)) = self.current_turn.as_mut() { + turn.state = ExplicitTurnState::Terminal; + } + } + + fn current_explicit_turn_has_id(&self, turn_id: &str) -> bool { + self.current_explicit_turn() + .is_some_and(|turn| turn.turn_id == turn_id) + } + + fn finished_turn_has_id(&self, turn_id: &str) -> bool { + self.finished_turns + .iter() + .any(|turn| matches!(turn, FinishedTurn::Explicit(id) if id == turn_id)) + } + + fn finish_current_turn(&mut self) { + let Some(turn) = self.current_turn.take() else { + return; + }; + self.finished_turns.push(match turn { + CurrentTurn::Explicit(turn) => FinishedTurn::Explicit(turn.turn_id), + CurrentTurn::Implicit(_) => FinishedTurn::Implicit, + }); + } +} + +#[cfg(test)] +#[path = "turn_lifecycle_tests.rs"] +mod tests; diff --git a/codex-rs/rollout/src/turn_lifecycle_tests.rs b/codex-rs/rollout/src/turn_lifecycle_tests.rs new file mode 100644 index 0000000000..cc5340946e --- /dev/null +++ b/codex-rs/rollout/src/turn_lifecycle_tests.rs @@ -0,0 +1,386 @@ +use codex_protocol::items::HookPromptFragment; +use codex_protocol::items::build_hook_prompt_message; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::CompactedItem; +use codex_protocol::protocol::ErrorEvent; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::ThreadRolledBackEvent; +use codex_protocol::protocol::TurnAbortReason; +use codex_protocol::protocol::TurnAbortedEvent; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::protocol::UserMessageEvent; +use pretty_assertions::assert_eq; + +use super::CurrentExplicitTurn; +use super::ExplicitTurnState; +use super::RolloutTurnLifecycleTracker; + +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: Option<&str>) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnAborted(TurnAbortedEvent { + turn_id: turn_id.map(str::to_string), + reason: TurnAbortReason::Interrupted, + completed_at: None, + duration_ms: None, + })) +} + +fn user_message(message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: message.to_string(), + ..Default::default() + })) +} + +fn agent_message(message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent { + message: message.to_string(), + phase: None, + memory_citation: None, + })) +} + +fn compacted() -> RolloutItem { + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: None, + window_number: None, + first_window_id: None, + previous_window_id: None, + window_id: None, + }) +} + +fn hook_prompt() -> RolloutItem { + let fragments = [HookPromptFragment::from_single_hook( + "hook guidance", + "hook-run-1", + )]; + RolloutItem::ResponseItem(build_hook_prompt_message(&fragments).expect("hook prompt message")) +} + +fn rollback(num_turns: u32) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns, + })) +} + +fn observe(tracker: &mut RolloutTurnLifecycleTracker, items: &[RolloutItem]) { + for item in items { + tracker.handle_rollout_item(item); + } +} + +#[test] +fn records_raw_rollout_index_for_explicit_turn_start() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + RolloutItem::ResponseItem(ResponseItem::Other), + RolloutItem::ResponseItem(ResponseItem::Other), + turn_started("turn-a"), + ], + ); + + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-a".to_string(), + rollout_start_index: 2, + state: ExplicitTurnState::InProgress, + }) + ); +} + +#[test] +fn complete_closes_while_abort_and_error_retain_terminal_current_turn() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[turn_started("turn-a"), turn_aborted(Some("turn-a"))], + ); + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-a".to_string(), + rollout_start_index: 0, + state: ExplicitTurnState::Terminal, + }) + ); + + tracker.handle_rollout_item(&turn_complete("turn-a")); + assert_eq!(tracker.current_explicit_turn(), None); + + tracker.handle_rollout_item(&turn_started("turn-b")); + tracker.handle_rollout_item(&RolloutItem::EventMsg(EventMsg::Error(ErrorEvent { + message: "failed".to_string(), + codex_error_info: None, + }))); + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 3, + state: ExplicitTurnState::Terminal, + }) + ); +} + +#[test] +fn second_start_finishes_and_replaces_current_turn() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[turn_started("turn-a"), turn_started("turn-b")], + ); + + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 1, + state: ExplicitTurnState::InProgress, + }) + ); +} + +#[test] +fn late_historical_ids_do_not_affect_current_but_unknown_ids_do() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + turn_started("turn-a"), + turn_complete("turn-a"), + turn_started("turn-b"), + turn_complete("turn-a"), + turn_aborted(Some("turn-a")), + ], + ); + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 2, + state: ExplicitTurnState::InProgress, + }) + ); + + tracker.handle_rollout_item(&turn_aborted(Some("unknown"))); + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 2, + state: ExplicitTurnState::Terminal, + }) + ); + + tracker.handle_rollout_item(&turn_complete("unknown")); + assert_eq!(tracker.current_explicit_turn(), None); +} + +#[test] +fn rollback_zero_finishes_current_and_rolled_back_ids_become_unknown() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + turn_started("turn-a"), + rollback(/*num_turns*/ 0), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 2, + state: ExplicitTurnState::InProgress, + }) + ); + + observe( + &mut tracker, + &[ + rollback(/*num_turns*/ 1), + turn_started("turn-c"), + turn_complete("turn-b"), + ], + ); + assert_eq!(tracker.current_explicit_turn(), None); +} + +#[test] +fn implicit_turn_placeholders_keep_rollback_late_id_matching_aligned() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + turn_started("turn-a"), + turn_complete("turn-a"), + user_message("legacy turn"), + rollback(/*num_turns*/ 1), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 4, + state: ExplicitTurnState::InProgress, + }) + ); + + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + user_message("legacy turn"), + turn_started("turn-a"), + turn_complete("turn-a"), + rollback(/*num_turns*/ 1), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + assert_eq!(tracker.current_explicit_turn(), None); +} + +#[test] +fn non_user_materialized_turn_adds_an_implicit_rollback_placeholder() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + turn_started("turn-a"), + turn_complete("turn-a"), + agent_message("legacy response"), + rollback(/*num_turns*/ 1), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 4, + state: ExplicitTurnState::InProgress, + }) + ); +} + +#[test] +fn first_user_message_reuses_compaction_only_turn_and_second_starts_another() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + turn_started("turn-a"), + turn_complete("turn-a"), + compacted(), + user_message("first legacy turn"), + user_message("second legacy turn"), + rollback(/*num_turns*/ 3), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + + assert_eq!(tracker.current_explicit_turn(), None); +} + +#[test] +fn hook_prompt_adds_an_implicit_rollback_placeholder() { + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe( + &mut tracker, + &[ + turn_started("turn-a"), + turn_complete("turn-a"), + hook_prompt(), + rollback(/*num_turns*/ 1), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 4, + state: ExplicitTurnState::InProgress, + }) + ); +} + +#[test] +fn hook_prompt_materializes_compaction_slot_before_user_starts_another() { + let items_before_rollback = [ + turn_started("turn-a"), + turn_complete("turn-a"), + compacted(), + hook_prompt(), + user_message("legacy turn"), + ]; + + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe(&mut tracker, &items_before_rollback); + observe( + &mut tracker, + &[ + rollback(/*num_turns*/ 2), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + assert_eq!( + tracker.current_explicit_turn(), + Some(&CurrentExplicitTurn { + turn_id: "turn-b".to_string(), + rollout_start_index: 6, + state: ExplicitTurnState::InProgress, + }) + ); + + let mut tracker = RolloutTurnLifecycleTracker::new(); + observe(&mut tracker, &items_before_rollback); + observe( + &mut tracker, + &[ + rollback(/*num_turns*/ 3), + turn_started("turn-b"), + turn_complete("turn-a"), + ], + ); + assert_eq!(tracker.current_explicit_turn(), None); +}