From 5d3f8752fc49662356832373db448f2b4bb21c34 Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Wed, 9 Sep 2026 19:16:49 +0000 Subject: [PATCH] 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 | 9 ++ codex-rs/core/src/session/mod.rs | 2 - codex-rs/core/src/session/reasoning_effort.rs | 9 +- .../src/session/reasoning_effort_tests.rs | 52 ++++++++- codex-rs/core/src/state/session.rs | 5 +- .../tests/suite/reasoning_effort_override.rs | 108 +++++++++++++++--- 6 files changed, 163 insertions(+), 22 deletions(-) diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 33adc834c2..cf173b4688 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/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, sub_id: String, num_turns: u32 .collect::>(); 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::(); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 62f8766329..d813225a80 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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); } diff --git a/codex-rs/core/src/session/reasoning_effort.rs b/codex-rs/core/src/session/reasoning_effort.rs index fa585a51e8..5529a0d43f 100644 --- a/codex-rs/core/src/session/reasoning_effort.rs +++ b/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()); diff --git a/codex-rs/core/src/session/reasoning_effort_tests.rs b/codex-rs/core/src/session/reasoning_effort_tests.rs index 18c976e54a..3ce2739c65 100644 --- a/codex-rs/core/src/session/reasoning_effort_tests.rs +++ b/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() { diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 175c1149ff..fa14b1da0f 100644 --- a/codex-rs/core/src/state/session.rs +++ b/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, diff --git a/codex-rs/core/tests/suite/reasoning_effort_override.rs b/codex-rs/core/tests/suite/reasoning_effort_override.rs index 8a895c7130..e6420e0a14 100644 --- a/codex-rs/core/tests/suite/reasoning_effort_override.rs +++ b/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 {