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 {