Preserve prewarmed reasoning effort across replay and early rollback (#44285)

## Why

Initial history replay and rollback before the first turn could clear the reasoning-effort baseline established by startup prewarm, preventing the first request from preserving that baseline when the selected effort changes.

## What changed

Preserve the reasoning-effort pin during history reconstruction. Clear it on rollback only after startup prewarm has been consumed. With reasoning-effort overrides enabled, a prewarmed `medium` baseline remains on the request while a subsequent `high` selection is sent as a `configuration_update`.

## Testing

Add regression coverage for prewarm completing before initial replay on resume and fork. Expand WebSocket tests to cover resumed and forked threads with overrides enabled and disabled, plus rollback before the first turn.

GitOrigin-RevId: 2e2d0e8a06b24f01b717458ca0576d63107c7697

codex-rs/core/src/session/handlers.rs

@@ -20,6 +20,7 @@ use crate::context::GuardianApprovedAction;
 use crate::context::NodeReplReviewEvidence;
 use crate::review_prompts::resolve_review_request;
 use crate::session::spawn_review_thread;
+use crate::state::ReasoningEffortPin;
 use crate::tasks::CompactTask;
 use crate::tasks::UserShellCommandMode;
 use crate::tasks::UserShellCommandTask;
@@ -335,6 +336,14 @@ pub async fn thread_rollback(sess: &Arc<Session>, sub_id: String, num_turns: u32
         .collect::<Vec<_>>();
     sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice())
         .await;
+    {
+        let mut state = sess.state.lock().await;
+        // Keep the baseline while startup prewarm is retained for the first turn,
+        // including when its task has not established the pin yet.
+        if state.startup_prewarm.is_none() {
+            state.reasoning_effort_pin = ReasoningEffortPin::Unset;
+        }
+    }
     sess.services
         .thread_extension_data
         .remove::<NodeReplReviewEvidence>();

codex-rs/core/src/session/mod.rs

@@ -1636,8 +1636,6 @@ impl Session {
             state
                 .history
                 .restore_review_context(Some(&retained_context), guardian_history.as_ref());
-            // The next send supplies the selected effort. Refresh its trusted override too.
-            state.reasoning_effort_pin = ReasoningEffortPin::Unset;
             if let Some(world_state) = world_state_baseline {
                 state.history.set_world_state_baseline(world_state);
             }

codex-rs/core/src/session/reasoning_effort.rs

@@ -1,7 +1,8 @@
 //! Cache-preserving effort updates and the request-effort baseline for a context window.
 //!
-//! Only trusted harness items establish overrides. Replay invalidates the runtime pin;
-//! successful compaction retires the overrides and allows a fresh request baseline.
+//! Only trusted harness items establish overrides. Replay and rollback preserve startup
+//! prewarm's baseline while it is retained; later rollback invalidates the runtime pin.
+//! Successful compaction retires the overrides and allows a fresh request baseline.
 
 use super::session::Session;
 use super::step_context::StepContext;
@@ -31,8 +32,8 @@ impl Session {
         let should_skip = {
             let mut state = self.state.lock().await;
             if matches!(state.reasoning_effort_pin, ReasoningEffortPin::Compacted) {
-                // Replay leaves the pin Unset. Only successful compaction allows the
-                // request baseline to establish the selection without another update.
+                // Only successful compaction allows the request baseline to establish
+                // the selection without another update.
                 state
                     .reasoning_effort_pin
                     .pin(&settings.model_info.slug, effort.clone());

codex-rs/core/src/session/reasoning_effort_tests.rs

@@ -1,11 +1,61 @@
-//! Request-effort lookup must leave live state untouched during compaction.
+//! Request-effort pins across initial replay and compaction lookups.
 
 use super::RequestEffortUsage;
+use crate::session::step_settings::ResolvedStepSettings;
 use crate::session::tests::make_session_and_context;
 use codex_features::Feature;
+use codex_history::InitialHistory;
+use codex_history::ResumedHistory;
+use codex_protocol::ThreadId;
 use codex_protocol::openai_models::ReasoningEffort;
 use pretty_assertions::assert_eq;
 use std::sync::Arc;
+use test_case::test_case;
+
+#[test_case(InitialHistory::Forked(Vec::new()); "fork")]
+#[test_case(InitialHistory::Resumed(ResumedHistory {
+    conversation_id: ThreadId::default(),
+    history: Arc::new(Vec::new()),
+    rollout_path: None,
+}); "resume")]
+#[tokio::test]
+async fn initial_replay_preserves_prewarmed_effort(history: InitialHistory) {
+    let (mut session, turn_context) = make_session_and_context().await;
+    session
+        .features
+        .enable(Feature::ReasoningEffortOverride)
+        .unwrap();
+    let mut model_info = Arc::clone(&turn_context.initial_settings.model_info);
+    Arc::make_mut(&mut model_info).use_responses_lite = true;
+    let mut selected = turn_context.initial_settings.selected().clone();
+    selected.collaboration_mode.settings.reasoning_effort = Some(ReasoningEffort::Medium);
+    let prewarm_settings = ResolvedStepSettings::new(
+        Arc::new(selected.clone()),
+        Arc::clone(&model_info),
+        /*fast_mode_enabled*/ false,
+    );
+    // Force the ordering where prewarm pins its request before initial replay finishes.
+    assert_eq!(
+        session
+            .reasoning_effort_for_request(&prewarm_settings, RequestEffortUsage::Sampling)
+            .await,
+        Some(ReasoningEffort::Medium),
+    );
+    session.record_initial_history(history).await;
+
+    selected.collaboration_mode.settings.reasoning_effort = Some(ReasoningEffort::High);
+    let turn_settings = ResolvedStepSettings::new(
+        Arc::new(selected),
+        model_info,
+        /*fast_mode_enabled*/ false,
+    );
+    assert_eq!(
+        session
+            .reasoning_effort_for_request(&turn_settings, RequestEffortUsage::Sampling)
+            .await,
+        Some(ReasoningEffort::Medium),
+    );
+}
 
 #[tokio::test]
 async fn compaction_effort_lookup_preserves_pin_for_fallback_models() {