Persist TurnContext after full-context reinjection

Author: charley-oai

codex-rs/core/src/codex.rs

@@ -2705,14 +2705,15 @@ impl Session {
             self.record_conversation_items(turn_context, &context_items)
                 .await;
         }
-        if previous_user_turn_model.is_none()
-            || (!should_inject_full_context && !context_items.is_empty())
+        if should_inject_full_context
+            || previous_user_turn_model.is_none()
+            || !context_items.is_empty()
         {
-            // Keep rollout TurnContext entries to:
+            // Keep rollout TurnContext entries for any turn that establishes or changes the
+            // persisted model-visible baseline:
             // - the first real user turn (to recover `previous_model` on resume)
+            // - full-context reinjection after the baseline was cleared
             // - steady-state turns that emitted explicit context diffs
-            // Full reinjection after compaction is tracked in runtime state only; resume will
-            // conservatively fall back to a missing baseline and reinject again if needed.
             self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item.clone())])
                 .await;
         }
@@ -5952,6 +5953,10 @@ mod tests {
     use crate::protocol::TokenCountEvent;
     use crate::protocol::TokenUsage;
     use crate::protocol::TokenUsageInfo;
+    use crate::protocol::UserMessageEvent;
+    use crate::rollout::policy::EventPersistenceMode;
+    use crate::rollout::recorder::RolloutRecorder;
+    use crate::rollout::recorder::RolloutRecorderParams;
     use crate::state::TaskKind;
     use crate::tasks::SessionTask;
     use crate::tasks::SessionTaskContext;
@@ -7978,6 +7983,83 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn record_context_updates_and_set_reference_context_item_persists_full_reinjection_to_rollout()
+     {
+        let (session, previous_context) = make_session_and_context().await;
+        let next_model = if previous_context.model_info.slug == "gpt-5.1" {
+            "gpt-5"
+        } else {
+            "gpt-5.1"
+        };
+        let turn_context = previous_context
+            .with_model(next_model.to_string(), &session.services.models_manager)
+            .await;
+        let config = session.get_config().await;
+        let recorder = RolloutRecorder::new(
+            config.as_ref(),
+            RolloutRecorderParams::new(
+                ThreadId::default(),
+                None,
+                SessionSource::Exec,
+                BaseInstructions::default(),
+                Vec::new(),
+                EventPersistenceMode::Limited,
+            ),
+            None,
+            None,
+        )
+        .await
+        .expect("create rollout recorder");
+        let rollout_path = recorder.rollout_path().to_path_buf();
+        {
+            let mut rollout = session.services.rollout.lock().await;
+            *rollout = Some(recorder);
+        }
+
+        session
+            .persist_rollout_items(&[RolloutItem::EventMsg(EventMsg::UserMessage(
+                UserMessageEvent {
+                    message: "seed rollout".to_string(),
+                    images: None,
+                    local_images: Vec::new(),
+                    text_elements: Vec::new(),
+                },
+            ))])
+            .await;
+        {
+            let mut state = session.state.lock().await;
+            state.set_reference_context_item(None);
+        }
+
+        session
+            .record_context_updates_and_set_reference_context_item(
+                &turn_context,
+                Some(previous_context.model_info.slug.as_str()),
+            )
+            .await;
+        session.ensure_rollout_materialized().await;
+        session.flush_rollout().await;
+
+        let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path)
+            .await
+            .expect("read rollout history")
+        else {
+            panic!("expected resumed rollout history");
+        };
+        let persisted_turn_context = resumed.history.iter().find_map(|item| match item {
+            RolloutItem::TurnContext(ctx) => Some(ctx.clone()),
+            _ => None,
+        });
+
+        assert_eq!(
+            serde_json::to_value(persisted_turn_context)
+                .expect("serialize persisted turn context item"),
+            serde_json::to_value(Some(turn_context.to_turn_context_item()))
+                .expect("serialize expected turn context item")
+        );
+    }
+
     #[tokio::test]
     async fn run_user_shell_command_does_not_set_reference_context_item() {
         let (session, _turn_context, rx) = make_session_and_context_with_rx().await;