diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 02d231c962..c8c00e1ec1 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -398,7 +398,6 @@ impl ModelClient { self.state .window_generation .store(window_generation, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); } pub(crate) fn advance_window_generation(&self) { diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index c0f62ec664..c6cd71dfaf 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -568,6 +568,11 @@ pub async fn thread_rollback(sess: &Arc, sub_id: String, num_turns: u32 sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice()) .await; sess.recompute_token_usage(turn_context.as_ref()).await; + sess.services.model_client.set_window_generation( + super::rollout_reconstruction::effective_window_generation_from_rollout( + replay_items.as_slice(), + ), + ); sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())]) .await; diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index a4c042af0c..d9c92fcda4 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -35,6 +35,81 @@ struct ActiveReplaySegment<'a> { base_replacement_history: Option<&'a [ResponseItem]>, } +#[derive(Debug, Default)] +struct WindowGenerationReplaySegment { + counts_as_user_turn: bool, + compaction_count: u64, +} + +fn finalize_window_generation_segment( + active_segment: WindowGenerationReplaySegment, + window_generation: &mut u64, + pending_rollback_turns: &mut usize, +) { + if *pending_rollback_turns > 0 { + if active_segment.counts_as_user_turn { + *pending_rollback_turns -= 1; + } + return; + } + + *window_generation = window_generation.saturating_add(active_segment.compaction_count); +} + +/// Replays rollout segments newest-to-oldest so compactions in rolled-back suffixes do not +/// contribute to the public context-window lineage generation. +pub(super) fn effective_window_generation_from_rollout(rollout_items: &[RolloutItem]) -> u64 { + let mut window_generation = 0u64; + let mut pending_rollback_turns = 0usize; + let mut active_segment: Option = None; + + for item in rollout_items.iter().rev() { + match item { + RolloutItem::Compacted(_) => { + let active_segment = + active_segment.get_or_insert_with(WindowGenerationReplaySegment::default); + active_segment.compaction_count = active_segment.compaction_count.saturating_add(1); + } + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { + pending_rollback_turns = pending_rollback_turns + .saturating_add(usize::try_from(rollback.num_turns).unwrap_or(usize::MAX)); + } + RolloutItem::EventMsg(EventMsg::TurnStarted(_)) => { + if let Some(active_segment) = active_segment.take() { + finalize_window_generation_segment( + active_segment, + &mut window_generation, + &mut pending_rollback_turns, + ); + } + } + RolloutItem::EventMsg(EventMsg::UserMessage(_)) => { + active_segment + .get_or_insert_with(WindowGenerationReplaySegment::default) + .counts_as_user_turn = true; + } + RolloutItem::ResponseItem(response_item) => { + let active_segment = + active_segment.get_or_insert_with(WindowGenerationReplaySegment::default); + active_segment.counts_as_user_turn |= is_user_turn_boundary(response_item); + } + RolloutItem::EventMsg(_) + | RolloutItem::TurnContext(_) + | RolloutItem::SessionMeta(_) => {} + } + } + + if let Some(active_segment) = active_segment { + finalize_window_generation_segment( + active_segment, + &mut window_generation, + &mut pending_rollback_turns, + ); + } + + window_generation +} + fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&str>) -> bool { active_turn_id .is_none_or(|turn_id| item_turn_id.is_none_or(|item_turn_id| item_turn_id == turn_id)) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 37a501fb38..6915665663 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -518,15 +518,15 @@ impl Session { InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, }; let window_generation = match &initial_history { - InitialHistory::Resumed(resumed_history) => u64::try_from( - resumed_history - .history - .iter() - .filter(|item| matches!(item, RolloutItem::Compacted(_))) - .count(), - ) - .unwrap_or(u64::MAX), - InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => 0, + InitialHistory::Resumed(resumed_history) => { + super::rollout_reconstruction::effective_window_generation_from_rollout( + &resumed_history.history, + ) + } + InitialHistory::Forked(history) => { + super::rollout_reconstruction::effective_window_generation_from_rollout(history) + } + InitialHistory::New | InitialHistory::Cleared => 0, }; // Kick off independent async setup tasks in parallel to reduce startup latency. // diff --git a/codex-rs/core/tests/suite/window_headers.rs b/codex-rs/core/tests/suite/window_headers.rs index 51bb132799..5115a007a8 100644 --- a/codex-rs/core/tests/suite/window_headers.rs +++ b/codex-rs/core/tests/suite/window_headers.rs @@ -164,6 +164,162 @@ async fn window_id_stays_stable_after_rollback_and_resume() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn window_id_rolls_back_across_compaction_and_persists_on_resume() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_assistant_message("msg-1", "first reply"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_assistant_message("msg-2", "summary"), + ev_completed("resp-2"), + ]), + sse(vec![ev_completed("resp-3")]), + sse(vec![ev_completed("resp-4")]), + sse(vec![ev_completed("resp-5")]), + ], + ) + .await; + + let mut builder = test_codex().with_config(|config| { + config.model_provider.name = "Non-OpenAI Model provider".to_string(); + config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string()); + }); + let initial = builder.build(&server).await?; + let initial_thread = Arc::clone(&initial.codex); + let rollout_path = initial + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + + submit_user_turn(&initial_thread, "before compact").await?; + submit_compact_turn(&initial_thread).await?; + submit_user_turn(&initial_thread, "discard me").await?; + initial_thread + .submit(Op::ThreadRollback { num_turns: 2 }) + .await?; + wait_for_event(&initial_thread, |event| { + matches!(event, EventMsg::ThreadRolledBack(_)) + }) + .await; + submit_user_turn(&initial_thread, "after rollback").await?; + shutdown_thread(&initial_thread).await?; + + let resumed = builder + .resume(&server, initial.home.clone(), rollout_path) + .await?; + submit_user_turn(&resumed.codex, "after resume").await?; + shutdown_thread(&resumed.codex).await?; + + let requests = request_log.requests(); + assert_eq!(requests.len(), 5, "expected five model requests"); + + let window_ids = requests.iter().map(window_id_parts).collect::>(); + let initial_thread_id = window_ids[0].0.clone(); + assert_eq!( + window_ids, + vec![ + (initial_thread_id.clone(), 0), + (initial_thread_id.clone(), 0), + (initial_thread_id.clone(), 1), + (initial_thread_id.clone(), 0), + (initial_thread_id, 0), + ] + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn forked_compacted_history_inherits_effective_generation_on_resume() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_assistant_message("msg-1", "first reply"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_assistant_message("msg-2", "summary"), + ev_completed("resp-2"), + ]), + sse(vec![ev_completed("resp-3")]), + sse(vec![ev_completed("resp-4")]), + ], + ) + .await; + + let mut builder = test_codex().with_config(|config| { + config.model_provider.name = "Non-OpenAI Model provider".to_string(); + config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string()); + }); + let initial = builder.build(&server).await?; + let initial_thread = Arc::clone(&initial.codex); + let rollout_path = initial + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + + submit_user_turn(&initial_thread, "before compact").await?; + submit_compact_turn(&initial_thread).await?; + shutdown_thread(&initial_thread).await?; + + let forked = initial + .thread_manager + .fork_thread( + /*snapshot*/ usize::MAX, + initial.config.clone(), + rollout_path, + /*thread_source*/ None, + /*persist_extended_history*/ false, + /*parent_trace*/ None, + ) + .await?; + let fork_rollout_path = forked + .session_configured + .rollout_path + .clone() + .expect("fork rollout path"); + submit_user_turn(&forked.thread, "after fork").await?; + shutdown_thread(&forked.thread).await?; + + let resumed = builder + .resume(&server, initial.home.clone(), fork_rollout_path) + .await?; + submit_user_turn(&resumed.codex, "after fork resume").await?; + shutdown_thread(&resumed.codex).await?; + + let requests = request_log.requests(); + assert_eq!(requests.len(), 4, "expected four model requests"); + + let window_ids = requests.iter().map(window_id_parts).collect::>(); + let initial_thread_id = window_ids[0].0.clone(); + let forked_thread_id = window_ids[2].0.clone(); + assert_ne!(forked_thread_id, initial_thread_id); + assert_eq!( + window_ids, + vec![ + (initial_thread_id.clone(), 0), + (initial_thread_id, 0), + (forked_thread_id.clone(), 1), + (forked_thread_id, 1), + ] + ); + + Ok(()) +} + async fn submit_user_turn(codex: &Arc, text: &str) -> Result<()> { codex .submit(Op::UserInput {