diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index d30e5a3eaf..abed28fd7a 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2072,13 +2072,13 @@ impl Session { InitialHistory::New => { // Defer initial context insertion until the first real turn starts so // turn/start overrides can be merged before we write model-visible context. - self.set_previous_turn_settings(None).await; + self.reset_reference_turn_context_state().await; } InitialHistory::Resumed(resumed_history) => { let rollout_items = resumed_history.history; - let previous_turn_settings = self - .apply_rollout_reconstruction(&turn_context, &rollout_items) + self.apply_rollout_reconstruction(&turn_context, &rollout_items) .await; + let previous_turn_settings = self.previous_turn_settings().await; // If resuming, warn when the last recorded model differs from the current one. let curr: &str = turn_context.model_info.slug.as_str(); @@ -2153,19 +2153,15 @@ impl Session { &self, turn_context: &TurnContext, rollout_items: &[RolloutItem], - ) -> Option { + ) { let reconstructed_rollout = self .reconstruct_history_from_rollout(turn_context, rollout_items) .await; - let previous_turn_settings = reconstructed_rollout.previous_turn_settings.clone(); - self.replace_history( + let mut state = self.state.lock().await; + state.replace_history_with_reference_turn_context_state( reconstructed_rollout.history, - reconstructed_rollout.reference_context_item, - ) - .await; - self.set_previous_turn_settings(previous_turn_settings.clone()) - .await; - previous_turn_settings + reconstructed_rollout.reference_turn_context_state, + ); } fn last_token_info_from_rollout(rollout_items: &[RolloutItem]) -> Option { @@ -2180,12 +2176,14 @@ impl Session { state.previous_turn_settings() } - pub(crate) async fn set_previous_turn_settings( - &self, - previous_turn_settings: Option, - ) { + pub(crate) async fn record_regular_turn_context(&self, turn_context_item: TurnContextItem) { let mut state = self.state.lock().await; - state.set_previous_turn_settings(previous_turn_settings); + state.record_regular_turn_context(turn_context_item); + } + + async fn reset_reference_turn_context_state(&self) { + let mut state = self.state.lock().await; + state.reset_reference_turn_context_state(); } fn maybe_refresh_shell_snapshot_for_cwd( @@ -3623,10 +3621,10 @@ impl Session { self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item.clone())]) .await; - // Advance the in-memory diff baseline even when this turn emitted no model-visible - // context items. This keeps later runtime diffing aligned with the current turn state. - let mut state = self.state.lock().await; - state.set_reference_context_item(Some(turn_context_item)); + // Advance the in-memory turn-context tracker even when this turn emitted no model-visible + // context items. Regular turns become both the latest turn-settings source and the active + // model-visible reference baseline for subsequent diffing. + self.record_regular_turn_context(turn_context_item).await; } pub(crate) async fn update_token_usage_info( @@ -5642,15 +5640,6 @@ pub(crate) async fn run_turn( let response_item: ResponseItem = initial_input_for_turn.clone().into(); sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item) .await; - // Track the previous-turn baseline from the regular user-turn path only so - // standalone tasks (compact/shell/review/undo) cannot suppress future - // model/realtime injections. - sess.set_previous_turn_settings(Some(PreviousTurnSettings { - model: turn_context.model_info.slug.clone(), - realtime_active: Some(turn_context.realtime_active), - })) - .await; - if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) .await; diff --git a/codex-rs/core/src/codex/rollout_reconstruction.rs b/codex-rs/core/src/codex/rollout_reconstruction.rs index 11ddbc1928..f769e3ced6 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction.rs @@ -1,36 +1,19 @@ use super::*; +use crate::context_manager::ReferenceTurnContextState; // Return value of `Session::reconstruct_history_from_rollout`, bundling the rebuilt history with // the resume/fork hydration metadata derived from the same replay. #[derive(Debug)] pub(super) struct RolloutReconstruction { pub(super) history: Vec, - pub(super) previous_turn_settings: Option, - pub(super) reference_context_item: Option, -} - -#[derive(Debug, Default)] -enum TurnReferenceContextItem { - /// No `TurnContextItem` has been seen for this replay span yet. - /// - /// This differs from `Cleared`: `NeverSet` means there is no evidence this turn ever - /// established a baseline, while `Cleared` means a baseline existed and a later compaction - /// invalidated it. Only the latter must emit an explicit clearing segment for resume/fork - /// hydration. - #[default] - NeverSet, - /// A previously established baseline was invalidated by later compaction. - Cleared, - /// The latest baseline established by this replay span. - Latest(Box), + pub(super) reference_turn_context_state: ReferenceTurnContextState, } #[derive(Debug, Default)] struct ActiveReplaySegment<'a> { turn_id: Option, counts_as_user_turn: bool, - previous_turn_settings: Option, - reference_context_item: TurnReferenceContextItem, + reference_turn_context_state: ReferenceTurnContextState, base_replacement_history: Option<&'a [ResponseItem]>, } @@ -42,8 +25,7 @@ fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&s fn finalize_active_segment<'a>( active_segment: ActiveReplaySegment<'a>, base_replacement_history: &mut Option<&'a [ResponseItem]>, - previous_turn_settings: &mut Option, - reference_context_item: &mut TurnReferenceContextItem, + reference_turn_context_state: &mut ReferenceTurnContextState, pending_rollback_turns: &mut usize, ) { // Thread rollback drops the newest surviving real user-message boundaries. In replay, that @@ -64,21 +46,56 @@ fn finalize_active_segment<'a>( *base_replacement_history = Some(segment_base_replacement_history); } - // `previous_turn_settings` come from the newest surviving user turn that established them. - if previous_turn_settings.is_none() && active_segment.counts_as_user_turn { - *previous_turn_settings = active_segment.previous_turn_settings; + merge_surviving_segment_turn_context_state( + reference_turn_context_state, + active_segment.reference_turn_context_state, + active_segment.counts_as_user_turn, + ); +} + +/// Merge one surviving replay segment's turn-context bookkeeping into the aggregate +/// `reference_turn_context_state` we are reconstructing for the newest surviving history tail. +/// +/// `segment_turn_context_state` is the per-segment state collected while replaying one reverse +/// segment. `reference_turn_context_state` is the cross-segment accumulator for the surviving +/// transcript after rollback has skipped newer user turns. +fn merge_surviving_segment_turn_context_state( + reference_turn_context_state: &mut ReferenceTurnContextState, + segment_turn_context_state: ReferenceTurnContextState, + counts_as_user_turn: bool, +) { + // Only real user turns should backfill "previous turn settings". Standalone task turns may + // carry lifecycle events, but they must not become the latest real turn context. + if counts_as_user_turn + && reference_turn_context_state + .latest_turn_context_item() + .is_none() + && let Some(turn_context_item) = segment_turn_context_state.latest_turn_context_item() + { + reference_turn_context_state.set_latest_turn_context_item(Some(turn_context_item)); } - // `reference_context_item` comes from the newest surviving user turn baseline, or - // from a surviving compaction that explicitly cleared that baseline. - if matches!(reference_context_item, TurnReferenceContextItem::NeverSet) - && (active_segment.counts_as_user_turn - || matches!( - active_segment.reference_context_item, - TurnReferenceContextItem::Cleared - )) + // A compaction seen in this segment hides older reference baselines, but it must not erase a + // newer stored reference baseline we already captured from a later surviving user turn. + if segment_turn_context_state.compacted_since_model_saw_reference_turn_context() + && reference_turn_context_state + .stored_reference_turn_context_item() + .is_none() { - *reference_context_item = active_segment.reference_context_item; + reference_turn_context_state.note_compaction(); + } + + // The model-visible reference baseline comes from the newest surviving user turn that both + // carries a stored baseline and has not been hidden by a later surviving compaction. + if counts_as_user_turn + && !reference_turn_context_state.compacted_since_model_saw_reference_turn_context() + && reference_turn_context_state + .stored_reference_turn_context_item() + .is_none() + && let Some(turn_context_item) = + segment_turn_context_state.stored_reference_turn_context_item() + { + reference_turn_context_state.set_reference_context_item(Some(turn_context_item)); } } @@ -90,12 +107,11 @@ impl Session { ) -> RolloutReconstruction { // Replay metadata should already match the shape of the future lazy reverse loader, even // while history materialization still uses an eager bridge. Scan newest-to-oldest, - // stopping once a surviving replacement-history checkpoint and the required resume metadata - // are both known; then replay only the buffered surviving tail forward to preserve exact - // history semantics. + // stopping once a surviving replacement-history checkpoint and the latest surviving turn + // context are both known; then replay only the buffered surviving tail forward to + // preserve exact history semantics. let mut base_replacement_history: Option<&[ResponseItem]> = None; - let mut previous_turn_settings = None; - let mut reference_context_item = TurnReferenceContextItem::NeverSet; + let mut reference_turn_context_state = ReferenceTurnContextState::default(); // Rollback is "drop the newest N user turns". While scanning in reverse, that becomes // "skip the next N user-turn segments we finalize". let mut pending_rollback_turns = 0usize; @@ -113,12 +129,9 @@ impl Session { active_segment.get_or_insert_with(ActiveReplaySegment::default); // Looking backward, compaction clears any older baseline unless a newer // `TurnContextItem` in this same segment has already re-established it. - if matches!( - active_segment.reference_context_item, - TurnReferenceContextItem::NeverSet - ) { - active_segment.reference_context_item = TurnReferenceContextItem::Cleared; - } + active_segment + .reference_turn_context_state + .note_compaction_during_reverse_replay(); if active_segment.base_replacement_history.is_none() && let Some(replacement_history) = &compacted.replacement_history { @@ -170,17 +183,9 @@ impl Session { active_segment.turn_id.as_deref(), ctx.turn_id.as_deref(), ) { - active_segment.previous_turn_settings = Some(PreviousTurnSettings { - model: ctx.model.clone(), - realtime_active: ctx.realtime_active, - }); - if matches!( - active_segment.reference_context_item, - TurnReferenceContextItem::NeverSet - ) { - active_segment.reference_context_item = - TurnReferenceContextItem::Latest(Box::new(ctx.clone())); - } + active_segment + .reference_turn_context_state + .note_turn_context_during_reverse_replay(ctx); } } RolloutItem::EventMsg(EventMsg::TurnStarted(event)) => { @@ -195,8 +200,7 @@ impl Session { finalize_active_segment( active_segment, &mut base_replacement_history, - &mut previous_turn_settings, - &mut reference_context_item, + &mut reference_turn_context_state, &mut pending_rollback_turns, ); } @@ -207,12 +211,13 @@ impl Session { } if base_replacement_history.is_some() - && previous_turn_settings.is_some() - && !matches!(reference_context_item, TurnReferenceContextItem::NeverSet) + && pending_rollback_turns == 0 + && reference_turn_context_state + .latest_turn_context_item() + .is_some() { - // At this point we have both eager resume metadata values and the replacement- - // history base for the surviving tail, so older rollout items cannot affect this - // result. + // At this point the replay-derived metadata and replacement-history base for the + // surviving tail are both fixed, so older rollout items cannot affect this result. break; } } @@ -221,8 +226,7 @@ impl Session { finalize_active_segment( active_segment, &mut base_replacement_history, - &mut previous_turn_settings, - &mut reference_context_item, + &mut reference_turn_context_state, &mut pending_rollback_turns, ); } @@ -276,22 +280,12 @@ impl Session { } } - let reference_context_item = match reference_context_item { - TurnReferenceContextItem::NeverSet | TurnReferenceContextItem::Cleared => None, - TurnReferenceContextItem::Latest(turn_reference_context_item) => { - Some(*turn_reference_context_item) - } - }; - let reference_context_item = if saw_legacy_compaction_without_replacement_history { - None - } else { - reference_context_item - }; - + if saw_legacy_compaction_without_replacement_history { + reference_turn_context_state.note_compaction(); + } RolloutReconstruction { history: history.raw_items().to_vec(), - previous_turn_settings, - reference_context_item, + reference_turn_context_state, } } } diff --git a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs index 6cc99a2907..c0173517e9 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs @@ -228,18 +228,34 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com vec![turn_one_user, turn_one_assistant] ); assert_eq!( - reconstructed.previous_turn_settings, + reconstructed + .reference_turn_context_state + .previous_turn_settings(), Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), realtime_active: Some(turn_context.realtime_active), }) ); assert_eq!( - serde_json::to_value(reconstructed.reference_context_item) - .expect("serialize reconstructed reference context item"), - serde_json::to_value(Some(first_context_item)) + serde_json::to_value( + reconstructed + .reference_turn_context_state + .reference_context_item(), + ) + .expect("serialize reconstructed reference context item"), + serde_json::to_value(Some(first_context_item.clone())) .expect("serialize expected reference context item") ); + assert_eq!( + serde_json::to_value( + reconstructed + .reference_turn_context_state + .latest_turn_context_item(), + ) + .expect("serialize surviving turn context item"), + serde_json::to_value(Some(first_context_item)) + .expect("serialize expected surviving turn context item") + ); } #[tokio::test] @@ -310,20 +326,166 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_inc vec![turn_one_user, turn_one_assistant] ); assert_eq!( - reconstructed.previous_turn_settings, + reconstructed + .reference_turn_context_state + .previous_turn_settings(), Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), realtime_active: Some(turn_context.realtime_active), }) ); assert_eq!( - serde_json::to_value(reconstructed.reference_context_item) - .expect("serialize reconstructed reference context item"), + serde_json::to_value( + reconstructed + .reference_turn_context_state + .reference_context_item(), + ) + .expect("serialize reconstructed reference context item"), serde_json::to_value(Some(first_context_item)) .expect("serialize expected reference context item") ); } +#[tokio::test] +async fn reconstruct_history_rollback_backfills_surviving_turn_context_from_older_turn() { + let (session, turn_context) = make_session_and_context().await; + let first_context_item = turn_context.to_turn_context_item(); + let first_turn_id = first_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let second_turn_id = "second-turn-without-context".to_string(); + let rolled_back_turn_id = "rolled-back-turn".to_string(); + let turn_one_user = user_message("turn 1 user"); + let turn_one_assistant = assistant_message("turn 1 assistant"); + let turn_two_user = user_message("turn 2 user"); + let turn_two_assistant = assistant_message("turn 2 assistant"); + + let rollout_items = vec![ + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: first_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 1 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(first_context_item.clone()), + RolloutItem::ResponseItem(turn_one_user.clone()), + RolloutItem::ResponseItem(turn_one_assistant.clone()), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: first_turn_id, + last_agent_message: None, + }, + )), + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: second_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 2 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::ResponseItem(turn_two_user.clone()), + RolloutItem::ResponseItem(turn_two_assistant.clone()), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: second_turn_id, + last_agent_message: None, + }, + )), + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: rolled_back_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 3 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(TurnContextItem { + turn_id: Some(rolled_back_turn_id.clone()), + model: "rolled-back-model".to_string(), + ..first_context_item.clone() + }), + RolloutItem::ResponseItem(user_message("turn 3 user")), + RolloutItem::ResponseItem(assistant_message("turn 3 assistant")), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: rolled_back_turn_id, + last_agent_message: None, + }, + )), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack( + codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 }, + )), + ]; + + let reconstructed = session + .reconstruct_history_from_rollout(&turn_context, &rollout_items) + .await; + + assert_eq!( + reconstructed.history, + vec![ + turn_one_user, + turn_one_assistant, + turn_two_user, + turn_two_assistant + ] + ); + assert_eq!( + reconstructed + .reference_turn_context_state + .previous_turn_settings(), + Some(PreviousTurnSettings { + model: turn_context.model_info.slug.clone(), + realtime_active: Some(turn_context.realtime_active), + }) + ); + assert_eq!( + serde_json::to_value( + reconstructed + .reference_turn_context_state + .reference_context_item(), + ) + .expect("serialize reconstructed reference context item"), + serde_json::to_value(Some(first_context_item.clone())) + .expect("serialize expected reference context item") + ); + assert_eq!( + serde_json::to_value( + reconstructed + .reference_turn_context_state + .latest_turn_context_item(), + ) + .expect("serialize surviving turn context item"), + serde_json::to_value(Some(first_context_item)) + .expect("serialize expected surviving turn context item") + ); +} + #[tokio::test] async fn reconstruct_history_rollback_skips_non_user_turns_for_history_and_metadata() { let (session, turn_context) = make_session_and_context().await; @@ -416,15 +578,21 @@ async fn reconstruct_history_rollback_skips_non_user_turns_for_history_and_metad vec![turn_one_user, turn_one_assistant] ); assert_eq!( - reconstructed.previous_turn_settings, + reconstructed + .reference_turn_context_state + .previous_turn_settings(), Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), realtime_active: Some(turn_context.realtime_active), }) ); assert_eq!( - serde_json::to_value(reconstructed.reference_context_item) - .expect("serialize reconstructed reference context item"), + serde_json::to_value( + reconstructed + .reference_turn_context_state + .reference_context_item(), + ) + .expect("serialize reconstructed reference context item"), serde_json::to_value(Some(first_context_item)) .expect("serialize expected reference context item") ); @@ -473,8 +641,18 @@ async fn reconstruct_history_rollback_clears_history_and_metadata_when_exceeding .await; assert_eq!(reconstructed.history, Vec::new()); - assert_eq!(reconstructed.previous_turn_settings, None); - assert!(reconstructed.reference_context_item.is_none()); + assert_eq!( + reconstructed + .reference_turn_context_state + .previous_turn_settings(), + None + ); + assert!( + reconstructed + .reference_turn_context_state + .reference_context_item() + .is_none() + ); } #[tokio::test] @@ -685,7 +863,12 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_ user_message("legacy summary"), ] ); - assert!(reconstructed.reference_context_item.is_none()); + assert!( + reconstructed + .reference_turn_context_state + .reference_context_item() + .is_none() + ); } #[tokio::test] @@ -731,7 +914,12 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear .reconstruct_history_from_rollout(&turn_context, &rollout_items) .await; - assert!(reconstructed.reference_context_item.is_none()); + assert!( + reconstructed + .reference_turn_context_state + .reference_context_item() + .is_none() + ); } #[tokio::test] diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index fd1bb576b4..b8043ed8cc 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -118,6 +118,39 @@ fn user_message(text: &str) -> ResponseItem { } } +fn turn_context_item_with_previous_turn_settings( + turn_context: &TurnContext, + previous_turn_settings: PreviousTurnSettings, +) -> TurnContextItem { + let mut turn_context_item = turn_context.to_turn_context_item(); + turn_context_item.model = previous_turn_settings.model; + turn_context_item.realtime_active = previous_turn_settings.realtime_active; + turn_context_item +} + +async fn seed_previous_turn_settings( + session: &Session, + turn_context: &TurnContext, + previous_turn_settings: PreviousTurnSettings, +) { + session + .record_regular_turn_context(turn_context_item_with_previous_turn_settings( + turn_context, + previous_turn_settings, + )) + .await; +} + +async fn seed_previous_turn_settings_without_reference( + session: &Session, + turn_context: &TurnContext, + previous_turn_settings: PreviousTurnSettings, +) { + seed_previous_turn_settings(session, turn_context, previous_turn_settings).await; + let mut state = session.state.lock().await; + state.set_reference_context_item(None); +} + fn assistant_message(text: &str) -> ResponseItem { ResponseItem::Message { id: None, @@ -1090,11 +1123,6 @@ async fn thread_rollback_drops_last_turn_from_history() { .map(RolloutItem::ResponseItem) .collect(); sess.persist_rollout_items(&rollout_items).await; - sess.set_previous_turn_settings(Some(PreviousTurnSettings { - model: "stale-model".to_string(), - realtime_active: Some(tc.realtime_active), - })) - .await; { let mut state = sess.state.lock().await; state.set_reference_context_item(Some(tc.to_turn_context_item())); @@ -1252,11 +1280,6 @@ async fn thread_rollback_recomputes_previous_turn_settings_and_reference_context Some(first_context_item.clone()), ) .await; - sess.set_previous_turn_settings(Some(PreviousTurnSettings { - model: "stale-model".to_string(), - realtime_active: None, - })) - .await; handlers::thread_rollback(&sess, "sub-1".to_string(), 1).await; let rollback_event = wait_for_thread_rolled_back(&rx).await; @@ -3328,7 +3351,7 @@ async fn record_model_warning_appends_user_message() { #[tokio::test] async fn spawn_task_does_not_update_previous_turn_settings_for_non_run_turn_tasks() { let (sess, tc, _rx) = make_session_and_context_with_rx().await; - sess.set_previous_turn_settings(None).await; + sess.reset_reference_turn_context_state().await; let input = vec![UserInput::Text { text: "hello".to_string(), text_elements: Vec::new(), @@ -3512,9 +3535,7 @@ async fn build_settings_update_items_uses_previous_turn_settings_for_realtime_en .await; current_context.realtime_active = false; - session - .set_previous_turn_settings(Some(previous_turn_settings)) - .await; + seed_previous_turn_settings(&session, &previous_context, previous_turn_settings).await; let update_items = session .build_settings_update_items(Some(&previous_context_item), ¤t_context) .await; @@ -3675,8 +3696,7 @@ async fn build_initial_context_uses_previous_turn_settings_for_realtime_end() { realtime_active: Some(true), }; - session - .set_previous_turn_settings(Some(previous_turn_settings)) + seed_previous_turn_settings_without_reference(&session, &turn_context, previous_turn_settings) .await; let initial_context = session.build_initial_context(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); @@ -3697,8 +3717,7 @@ async fn build_initial_context_restates_realtime_start_when_reference_context_is realtime_active: Some(true), }; - session - .set_previous_turn_settings(Some(previous_turn_settings)) + seed_previous_turn_settings_without_reference(&session, &turn_context, previous_turn_settings) .await; let initial_context = session.build_initial_context(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); @@ -3853,8 +3872,7 @@ async fn build_initial_context_prepends_model_switch_message() { realtime_active: None, }; - session - .set_previous_turn_settings(Some(previous_turn_settings)) + seed_previous_turn_settings_without_reference(&session, &turn_context, previous_turn_settings) .await; let initial_context = session.build_initial_context(&turn_context).await; @@ -3917,12 +3935,15 @@ async fn record_context_updates_and_set_reference_context_item_persists_full_rei state.set_reference_context_item(None); } - session - .set_previous_turn_settings(Some(PreviousTurnSettings { + seed_previous_turn_settings_without_reference( + &session, + &previous_context, + PreviousTurnSettings { model: previous_context.model_info.slug.clone(), realtime_active: Some(previous_context.realtime_active), - })) - .await; + }, + ) + .await; session .record_context_updates_and_set_reference_context_item(&turn_context) .await; diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 92e889d647..4ed5fe0f79 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -6,9 +6,15 @@ async fn process_compacted_history_with_test_session( previous_turn_settings: Option<&PreviousTurnSettings>, ) -> (Vec, Vec) { let (session, turn_context) = crate::codex::make_session_and_context().await; - session - .set_previous_turn_settings(previous_turn_settings.cloned()) - .await; + if let Some(previous_turn_settings) = previous_turn_settings { + let mut previous_turn_context_item = turn_context.to_turn_context_item(); + previous_turn_context_item.model = previous_turn_settings.model.clone(); + previous_turn_context_item.realtime_active = previous_turn_settings.realtime_active; + session + .record_regular_turn_context(previous_turn_context_item) + .await; + session.replace_history(Vec::new(), None).await; + } let initial_context = session.build_initial_context(&turn_context).await; let refreshed = crate::compact_remote::process_compacted_history( &session, diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 19dbcf7293..ad8dd81cb2 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,3 +1,4 @@ +use crate::codex::PreviousTurnSettings; use crate::codex::TurnContext; use crate::context_manager::normalize; use crate::event_mapping::is_contextual_user_message_content; @@ -32,15 +33,111 @@ pub(crate) struct ContextManager { /// The oldest items are at the beginning of the vector. items: Vec, token_info: Option, - /// Reference context snapshot used for diffing and producing model-visible - /// settings update items. + reference_turn_context_state: ReferenceTurnContextState, +} + +/// Session-owned bookkeeping for turn-context state that survives history replay, +/// rollback, and compaction. +/// +/// This intentionally tracks both the latest real turn context we know about and the +/// model-visible reference baseline, because those diverge when compaction hides the +/// baseline without erasing the last real turn's settings. +#[derive(Debug, Clone, Default)] +pub(crate) struct ReferenceTurnContextState { + /// The most recent real turn context we reconstructed or recorded, even if a later + /// compaction means the model can no longer rely on it as the active baseline. /// - /// This is the baseline for the next regular model turn, and may already - /// match the current turn after context updates are persisted. + /// This drives `previous_turn_settings()` and other rollback bookkeeping, which + /// intentionally survive compaction until a newer real turn replaces them. + latest_turn_context_item: Option, + /// The last turn context item that established the model's reference baseline. /// - /// When this is `None`, settings diffing treats the next turn as having no - /// baseline and emits a full reinjection of context state. - reference_context_item: Option, + /// Unlike `latest_turn_context_item`, this is only model-visible when + /// `compacted_since_model_saw_reference_turn_context` is false. + reference_turn_context_item: Option, + /// Whether compaction has crossed the current reference baseline without a later + /// reinjection or real turn context re-establishing it. + /// + /// When this is true, `reference_context_item()` must return `None` even if + /// `reference_turn_context_item` still retains the last stored baseline for replay or + /// rollback bookkeeping. + compacted_since_model_saw_reference_turn_context: bool, +} + +impl ReferenceTurnContextState { + pub(crate) fn reset(&mut self) { + *self = Self::default(); + } + + pub(crate) fn note_compaction(&mut self) { + self.compacted_since_model_saw_reference_turn_context = true; + } + + pub(crate) fn note_compaction_during_reverse_replay(&mut self) { + if self.reference_turn_context_item.is_none() { + self.compacted_since_model_saw_reference_turn_context = true; + } + } + + pub(crate) fn note_turn_context_during_reverse_replay( + &mut self, + turn_context_item: &TurnContextItem, + ) { + if self.latest_turn_context_item.is_none() { + self.latest_turn_context_item = Some(turn_context_item.clone()); + } + if self.reference_turn_context_item.is_none() { + self.reference_turn_context_item = Some(turn_context_item.clone()); + } + } + + pub(crate) fn set_latest_turn_context_item(&mut self, item: Option) { + self.latest_turn_context_item = item; + } + + pub(crate) fn record_regular_turn_context(&mut self, turn_context_item: TurnContextItem) { + self.latest_turn_context_item = Some(turn_context_item.clone()); + self.reference_turn_context_item = Some(turn_context_item); + self.compacted_since_model_saw_reference_turn_context = false; + } + + pub(crate) fn set_reference_context_item(&mut self, item: Option) { + if let Some(item) = item { + self.reference_turn_context_item = Some(item); + self.compacted_since_model_saw_reference_turn_context = false; + } else { + self.note_compaction(); + } + } + + pub(crate) fn latest_turn_context_item(&self) -> Option { + self.latest_turn_context_item.clone() + } + + pub(crate) fn stored_reference_turn_context_item(&self) -> Option { + self.reference_turn_context_item.clone() + } + + pub(crate) fn compacted_since_model_saw_reference_turn_context(&self) -> bool { + self.compacted_since_model_saw_reference_turn_context + } + + pub(crate) fn previous_turn_settings(&self) -> Option { + self.latest_turn_context_item + .as_ref() + .map(|turn_context_item| PreviousTurnSettings { + model: turn_context_item.model.clone(), + realtime_active: turn_context_item.realtime_active, + }) + } + + pub(crate) fn reference_context_item(&self) -> Option { + if self.compacted_since_model_saw_reference_turn_context { + None + } else { + self.reference_turn_context_item.clone() + } + } } #[derive(Debug, Clone, Copy, Default)] @@ -56,7 +153,7 @@ impl ContextManager { Self { items: Vec::new(), token_info: TokenUsageInfo::new_or_append(&None, &None, None), - reference_context_item: None, + reference_turn_context_state: ReferenceTurnContextState::default(), } } @@ -68,12 +165,33 @@ impl ContextManager { self.token_info = info; } + pub(crate) fn set_reference_turn_context_state( + &mut self, + reference_turn_context_state: ReferenceTurnContextState, + ) { + self.reference_turn_context_state = reference_turn_context_state; + } + + pub(crate) fn reset_reference_turn_context_state(&mut self) { + self.reference_turn_context_state.reset(); + } + + pub(crate) fn record_regular_turn_context(&mut self, turn_context_item: TurnContextItem) { + self.reference_turn_context_state + .record_regular_turn_context(turn_context_item); + } + + pub(crate) fn previous_turn_settings(&self) -> Option { + self.reference_turn_context_state.previous_turn_settings() + } + pub(crate) fn set_reference_context_item(&mut self, item: Option) { - self.reference_context_item = item; + self.reference_turn_context_state + .set_reference_context_item(item); } pub(crate) fn reference_context_item(&self) -> Option { - self.reference_context_item.clone() + self.reference_turn_context_state.reference_context_item() } pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 853f8af5ac..b3c2a0ff89 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -3,6 +3,7 @@ mod normalize; pub(crate) mod updates; pub(crate) use history::ContextManager; +pub(crate) use history::ReferenceTurnContextState; pub(crate) use history::TotalTokenUsageBreakdown; pub(crate) use history::estimate_response_item_model_visible_bytes; pub(crate) use history::is_codex_generated_item; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 40faa4b856..484110197a 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -9,6 +9,7 @@ use tokio::task::JoinHandle; use crate::codex::PreviousTurnSettings; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; +use crate::context_manager::ReferenceTurnContextState; use crate::error::Result as CodexResult; use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; @@ -26,10 +27,6 @@ pub(crate) struct SessionState { pub(crate) server_reasoning_included: bool, pub(crate) dependency_env: HashMap, pub(crate) mcp_dependency_prompted: HashSet, - /// Settings used by the latest regular user turn, used for turn-to-turn - /// model/realtime handling on subsequent regular turns (including full-context - /// reinjection after resume or `/compact`). - previous_turn_settings: Option, /// Startup regular task pre-created during session initialization. pub(crate) startup_regular_task: Option>>, pub(crate) active_connector_selection: HashSet, @@ -48,7 +45,6 @@ impl SessionState { server_reasoning_included: false, dependency_env: HashMap::new(), mcp_dependency_prompted: HashSet::new(), - previous_turn_settings: None, startup_regular_task: None, active_connector_selection: HashSet::new(), pending_session_start_source: None, @@ -66,13 +62,7 @@ impl SessionState { } pub(crate) fn previous_turn_settings(&self) -> Option { - self.previous_turn_settings.clone() - } - pub(crate) fn set_previous_turn_settings( - &mut self, - previous_turn_settings: Option, - ) { - self.previous_turn_settings = previous_turn_settings; + self.history.previous_turn_settings() } pub(crate) fn clone_history(&self) -> ContextManager { @@ -89,10 +79,28 @@ impl SessionState { .set_reference_context_item(reference_context_item); } + pub(crate) fn replace_history_with_reference_turn_context_state( + &mut self, + items: Vec, + reference_turn_context_state: ReferenceTurnContextState, + ) { + self.history.replace(items); + self.history + .set_reference_turn_context_state(reference_turn_context_state); + } + + pub(crate) fn reset_reference_turn_context_state(&mut self) { + self.history.reset_reference_turn_context_state(); + } + pub(crate) fn set_token_info(&mut self, info: Option) { self.history.set_token_info(info); } + pub(crate) fn record_regular_turn_context(&mut self, turn_context_item: TurnContextItem) { + self.history.record_regular_turn_context(turn_context_item); + } + pub(crate) fn set_reference_context_item(&mut self, item: Option) { self.history.set_reference_context_item(item); }