openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-20 12:47:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

rollout: own turn lifecycle replay

Browse Source
This commit is contained in:
Adam Perry
2026-06-23 06:24:13 +00:00
parent 24423f5712
commit 9c9326d6b3
5 changed files with 631 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
codex-rs/core/Cargo.toml
View File

@@ -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 }

32
codex-rs/core/src/thread_manager.rs
View File

@@ -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(),
},
};
}

4
codex-rs/rollout/src/lib.rs
View File

@@ -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;

227
codex-rs/rollout/src/turn_lifecycle.rs Normal file
View File

@@ -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<FinishedTurn>,
current_turn: Option<CurrentTurn>,
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;

386
codex-rs/rollout/src/turn_lifecycle_tests.rs Normal file
View File

@@ -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);
}