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() {

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

@@ -30,8 +30,9 @@ use codex_protocol::protocol::TurnContextItem;
 use codex_utils_output_truncation::TruncationPolicy;
 use tokio_util::task::AbortOnDropHandle;
 
-/// Runtime request effort. A successful compaction can establish a fresh baseline, while
-/// startup and replay must explicitly re-establish the selected effort in surviving history.
+/// Runtime request effort, initially unset and established by prewarm or sampling.
+/// Rollback clears it after startup prewarm is consumed; successful compaction allows
+/// a fresh baseline without an override.
 pub(crate) enum ReasoningEffortPin {
     Unset,
     Compacted,

codex-rs/core/tests/suite/reasoning_effort_override.rs

@@ -1,5 +1,6 @@
 //! Trusted reasoning-effort updates follow surviving history and the next turn's selected settings.
 
+use codex_core::ForkSnapshot;
 use codex_core::RecoverTurnRequest;
 use codex_core::StartIfIdleSubmission;
 use codex_core::SuspendTurnOutcome;
@@ -372,11 +373,27 @@ async fn reasoning_effort_override_normalizes_ultra_before_comparing_updates(
     Ok(())
 }
 
-#[test_case(true; "feature enabled")]
-#[test_case(false; "feature disabled")]
+#[derive(Clone, Copy)]
+enum PrewarmStartup {
+    New,
+    Resume,
+    Fork,
+    ResumeThenRollback,
+    ForkThenRollback,
+}
+
+#[test_case(true, PrewarmStartup::New; "new thread feature enabled")]
+#[test_case(false, PrewarmStartup::New; "new thread feature disabled")]
+#[test_case(true, PrewarmStartup::Resume; "resumed thread feature enabled")]
+#[test_case(false, PrewarmStartup::Resume; "resumed thread feature disabled")]
+#[test_case(true, PrewarmStartup::Fork; "forked thread feature enabled")]
+#[test_case(false, PrewarmStartup::Fork; "forked thread feature disabled")]
+#[test_case(true, PrewarmStartup::ResumeThenRollback; "resumed thread rollback before first turn")]
+#[test_case(true, PrewarmStartup::ForkThenRollback; "forked thread rollback before first turn")]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn reasoning_effort_override_websocket_prewarm_preserves_baseline(
     feature_enabled: bool,
+    startup: PrewarmStartup,
 ) -> anyhow::Result<()> {
     skip_if_no_network!(Ok(()));
     let server = responses::start_mock_server().await;
@@ -392,17 +409,65 @@ async fn reasoning_effort_override_websocket_prewarm_preserves_baseline(
     ]])
     .await;
     let base_url = format!("{}/v1", websocket.uri());
-    let test = override_builder()
-        .with_config(move |config| {
-            config.model_provider.base_url = Some(base_url);
-            config.model_provider.supports_websockets = true;
-            config
-                .features
-                .set_enabled(Feature::ReasoningEffortOverride, feature_enabled)
-                .expect("configure reasoning effort overrides");
-        })
-        .build_with_auto_env(&server)
-        .await?;
+    let configure_prewarm = move |config: &mut Config| {
+        config.model_provider.base_url = Some(base_url.clone());
+        config.model_provider.supports_websockets = true;
+        config
+            .features
+            .set_enabled(Feature::ReasoningEffortOverride, feature_enabled)
+            .expect("configure reasoning effort overrides");
+    };
+    let mut builder = override_builder().with_config(configure_prewarm.clone());
+    let test = if !matches!(startup, PrewarmStartup::New) {
+        let previous_mock = responses::mount_sse_once(
+            &server,
+            responses::sse(vec![responses::ev_completed("previous")]),
+        )
+        .await;
+        let mut previous = override_builder()
+            .with_config(|config| {
+                config
+                    .features
+                    .disable(Feature::ReasoningEffortOverride)
+                    .expect("disable reasoning effort overrides for source history");
+                config.model_provider.supports_websockets = false;
+            })
+            .build_with_auto_env(&server)
+            .await?;
+        previous.submit_text_turn("previous turn").await?;
+        assert_eq!(
+            previous_mock
+                .single_request()
+                .message_input_texts("user")
+                .last()
+                .map(String::as_str),
+            Some("previous turn"),
+        );
+        if matches!(
+            startup,
+            PrewarmStartup::Fork | PrewarmStartup::ForkThenRollback
+        ) {
+            previous.codex.shutdown_and_wait().await?;
+            let mut config = previous.config.clone();
+            configure_prewarm(&mut config);
+            let forked = previous
+                .thread_manager
+                .fork_thread(
+                    ForkSnapshot::Interrupted,
+                    codex_core::StartThreadOptions::new(config.clone()),
+                    previous.codex.rollout_path().expect("rollout path"),
+                )
+                .await?;
+            previous.codex = forked.thread;
+            previous.session_configured = forked.session_configured;
+            previous.config = config;
+            previous
+        } else {
+            builder.restart(&server, &previous).await?
+        }
+    } else {
+        builder.build_with_auto_env(&server).await?
+    };
     let warmup = tokio::time::timeout(
         std::time::Duration::from_secs(/*secs*/ 10),
         websocket.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0),
@@ -410,6 +475,23 @@ async fn reasoning_effort_override_websocket_prewarm_preserves_baseline(
     .await?;
     assert_eq!(warmup.body_json()["generate"], false);
     assert_eq!(warmup.body_json()["reasoning"]["effort"], "medium");
+    if matches!(
+        startup,
+        PrewarmStartup::ResumeThenRollback | PrewarmStartup::ForkThenRollback
+    ) {
+        // Observing the warmup request guarantees that Medium is pinned before rollback.
+        test.codex
+            .submit(Op::ThreadRollback { num_turns: 1 })
+            .await?;
+        let rollback = wait_for_event(&test.codex, |event| {
+            matches!(event, EventMsg::ThreadRolledBack(_) | EventMsg::Error(_))
+        })
+        .await;
+        assert!(
+            matches!(rollback, EventMsg::ThreadRolledBack(_)),
+            "rollback failed: {rollback:?}",
+        );
+    }
     submit_thread_settings(
         &test.codex,
         ThreadSettingsOverrides {