From 31ccaf40c2298bb3286c8fe274e1e21c498bb70a Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Tue, 8 Sep 2026 10:24:16 +0000 Subject: [PATCH] Pin request reasoning effort while configuration overrides are active (#43795) ## Why Reasoning-effort changes should preserve the request baseline while trusted `configuration_update` items carry the selected effort. ## What changed - Pin request-level reasoning effort for the current model when reasoning-effort overrides are available, including during WebSocket prewarm. - Re-establish the selected effort after rollback, resume, or a model switch, refreshing the trusted override even when the selection is unchanged. - Continue using the selected request effort when overrides are unavailable. ## Testing Add regression coverage for WebSocket prewarm and reconnect, rollback, resume, and model switches. Update transition tests to assert that request-level effort stays pinned while configuration updates change. GitOrigin-RevId: 89cfee6289543230347b0488d1388190b4da7216 --- codex-rs/core/src/session/mod.rs | 3 + codex-rs/core/src/session/reasoning_effort.rs | 34 +- codex-rs/core/src/session/turn.rs | 3 +- codex-rs/core/src/session_startup_prewarm.rs | 5 +- codex-rs/core/src/state/mod.rs | 1 + codex-rs/core/src/state/session.rs | 36 ++ .../tests/suite/reasoning_effort_override.rs | 436 +++++++++++++++++- 7 files changed, 504 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b4e98168fa..5922f61e61 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -299,6 +299,7 @@ use crate::state::AcceptedUserInputResponse; use crate::state::AutoCompactWindowIds; use crate::state::AutoCompactWindowSnapshot; use crate::state::PendingRequestPermissions; +use crate::state::ReasoningEffortPin; use crate::state::SessionServices; use crate::state::SessionState; #[cfg(test)] @@ -1605,6 +1606,8 @@ 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 c11d010fa7..568b7a14ad 100644 --- a/codex-rs/core/src/session/reasoning_effort.rs +++ b/codex-rs/core/src/session/reasoning_effort.rs @@ -1,11 +1,12 @@ -//! Gated reasoning-effort updates appended after accepted input. +//! Cache-preserving effort updates and the original sampling request effort. //! -//! Only trusted harness items establish an effort. Updates append to surviving -//! history without replacing the request-level reasoning effort. +//! Only trusted harness items establish overrides. Replay invalidates the runtime +//! pin so the next send re-establishes the currently selected effort. use super::session::Session; use super::step_context::StepContext; use super::step_settings::ResolvedStepSettings; +use crate::state::ReasoningEffortPin; use codex_features::Feature; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; @@ -43,7 +44,10 @@ impl Session { _ => None, } }); - established_effort == Some(&effort) + state + .reasoning_effort_pin + .get(&settings.model_info.slug) + .is_some_and(|pinned| established_effort.unwrap_or(&pinned) == &effort) }; if should_skip { return; @@ -64,6 +68,28 @@ impl Session { .await; } + /// Reuses the original request effort while updates carry the selected effort. + pub(crate) async fn reasoning_effort_for_request( + &self, + settings: &ResolvedStepSettings, + ) -> Option { + let selected_effort = settings.reasoning_effort().cloned(); + if !self.enabled(Feature::ReasoningEffortOverride) { + return selected_effort; + } + let effort = self.effort_for_configuration_update(settings).await; + let mut state = self.state.lock().await; + let Some(effort) = effort else { + state.reasoning_effort_pin = ReasoningEffortPin::Unset; + return selected_effort; + }; + Some( + state + .reasoning_effort_pin + .pin(&settings.model_info.slug, effort), + ) + } + async fn effort_for_configuration_update( &self, settings: &ResolvedStepSettings, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 7f31667988..93216d7c83 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2313,7 +2313,8 @@ async fn try_run_sampling_request( prompt, &step_context.settings.model_info, &step_context.session_telemetry, - step_context.settings.reasoning_effort().cloned(), + sess.reasoning_effort_for_request(&step_context.settings) + .await, step_context.settings.reasoning_summary, step_context.settings.service_tier.clone(), responses_metadata, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 546962fea1..02ce6abf1e 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -318,12 +318,15 @@ async fn schedule_startup_prewarm_inner( .await; let mut client_session = session.services.model_client.new_session(); let websocket_warmup_started_at = Instant::now(); + // Prewarm establishes the request baseline before the first turn can change effort. client_session .prewarm_websocket( &startup_prompt, &step_context.settings.model_info, &step_context.session_telemetry, - step_context.settings.reasoning_effort().cloned(), + session + .reasoning_effort_for_request(&step_context.settings) + .await, step_context.settings.reasoning_summary, step_context.settings.service_tier.clone(), &responses_metadata, diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 218b25ba92..71d9616e52 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -9,6 +9,7 @@ pub(crate) use additional_context::AdditionalContextStore; pub(crate) use auto_compact_window::AutoCompactWindowIds; pub(crate) use auto_compact_window::AutoCompactWindowSnapshot; pub(crate) use service::SessionServices; +pub(crate) use session::ReasoningEffortPin; pub(crate) use session::SessionState; pub(crate) use turn::AcceptedUserInputResponse; pub(crate) use turn::ActiveTurn; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 5333661a6f..901a5a5aca 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -3,6 +3,7 @@ use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::models::BaseInstructionsProvenance; use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::ReasoningEffort; use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; @@ -29,6 +30,38 @@ use codex_protocol::protocol::TurnContextItem; use codex_utils_output_truncation::TruncationPolicy; use tokio_util::task::AbortOnDropHandle; +/// Runtime request effort; startup and replay must re-establish the selected effort. +pub(crate) enum ReasoningEffortPin { + Unset, + Active { + model: String, + effort: ReasoningEffort, + }, +} + +impl ReasoningEffortPin { + pub(crate) fn get(&self, model: &str) -> Option { + match self { + Self::Active { + model: pinned_model, + effort, + } if pinned_model == model => Some(effort.clone()), + Self::Unset | Self::Active { .. } => None, + } + } + + pub(crate) fn pin(&mut self, model: &str, effort: ReasoningEffort) -> ReasoningEffort { + if let Some(pinned) = self.get(model) { + return pinned; + } + *self = Self::Active { + model: model.to_owned(), + effort: effort.clone(), + }; + effort + } +} + /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, @@ -46,6 +79,8 @@ pub(crate) struct SessionState { previous_turn_settings: Option, /// Runtime accounting state for the active auto-compaction window. auto_compact_window: AutoCompactWindow, + /// Original request effort for the current model while configuration updates remain active. + pub(crate) reasoning_effort_pin: ReasoningEffortPin, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, /// Retained after completion so later turns do not repeat speculative captures. @@ -84,6 +119,7 @@ impl SessionState { additional_context: AdditionalContextStore::default(), previous_turn_settings: None, auto_compact_window: AutoCompactWindow::new_with_ids(auto_compact_window_ids), + reasoning_effort_pin: ReasoningEffortPin::Unset, startup_prewarm: None, shell_snapshot_prewarm: None, current_time_reminder: CurrentTimeReminderState::default(), diff --git a/codex-rs/core/tests/suite/reasoning_effort_override.rs b/codex-rs/core/tests/suite/reasoning_effort_override.rs index 8fa7be8dcf..fc834a9356 100644 --- a/codex-rs/core/tests/suite/reasoning_effort_override.rs +++ b/codex-rs/core/tests/suite/reasoning_effort_override.rs @@ -1,8 +1,15 @@ //! Trusted reasoning-effort updates follow surviving history and the next turn's selected settings. +use codex_core::config::Config; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ThreadIdleInput; +use codex_extension_api::ThreadLifecycleContributor; use codex_features::Feature; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::openai_models::ReasoningEffortPreset; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadSettingsOverrides; use core_test_support::responses; use core_test_support::responses::ResponsesRequest; @@ -10,9 +17,14 @@ use core_test_support::skip_if_no_network; use core_test_support::submit_thread_settings; use core_test_support::test_codex::TestCodexBuilder; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::Value; +use std::sync::Arc; +use std::time::Duration; use test_case::test_case; +use tokio::sync::Notify; +use tokio::time::timeout; fn override_builder() -> TestCodexBuilder { test_codex() @@ -91,11 +103,19 @@ async fn reasoning_effort_override_persistent_transitions( .await?; test.submit_text_turn("continue").await?; } + let requests = mocks + .iter() + .map(responses::ResponseMock::single_request) + .collect::>(); assert_eq!( - mocks + requests .iter() - .map(|mock| effort_updates(&mock.single_request())) + .map(|request| request.body_json()["reasoning"]["effort"].clone()) .collect::>(), + vec![initial_update["reasoning"]["effort"].clone(); requests.len()], + ); + assert_eq!( + requests.iter().map(effort_updates).collect::>(), vec![ vec![initial_update.clone()], vec![initial_update.clone(), changed_update.clone()], @@ -160,12 +180,7 @@ async fn reasoning_effort_override_preserves_prefix_and_only_appends_on_change() .iter() .map(|request| request.body_json()["reasoning"]["effort"].clone()) .collect::>(), - vec![ - Value::from("medium"), - Value::from("high"), - Value::from("high"), - Value::from("low"), - ], + vec![Value::from("medium"); requests.len()], ); let cache_keys = requests .iter() @@ -253,6 +268,171 @@ async fn reasoning_effort_override_normalizes_ultra_before_comparing_updates( Ok(()) } +#[test_case(true; "feature enabled")] +#[test_case(false; "feature disabled")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn reasoning_effort_override_websocket_prewarm_preserves_baseline( + feature_enabled: bool, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let websocket = responses::start_websocket_server(vec![vec![ + vec![ + responses::ev_response_created("warmup"), + responses::ev_completed("warmup"), + ], + vec![ + responses::ev_response_created("first"), + responses::ev_completed("first"), + ], + ]]) + .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 warmup = tokio::time::timeout( + std::time::Duration::from_secs(/*secs*/ 10), + websocket.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0), + ) + .await?; + assert_eq!(warmup.body_json()["generate"], false); + assert_eq!(warmup.body_json()["reasoning"]["effort"], "medium"); + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + effort: Some(Some(ReasoningEffort::High)), + ..Default::default() + }, + ) + .await?; + test.submit_text_turn("first message after changing effort") + .await?; + + let connection = websocket.single_connection(); + assert_eq!(connection.len(), 2); + let first = connection[1].body_json(); + let updates = first["input"] + .as_array() + .expect("first turn input") + .iter() + .filter(|item| item["type"] == "configuration_update") + .cloned() + .collect::>(); + assert_eq!( + (first["reasoning"]["effort"].clone(), updates), + if feature_enabled { + ( + Value::from("medium"), + vec![effort_update(ReasoningEffort::High)], + ) + } else { + (Value::from("high"), Vec::new()) + }, + ); + websocket.shutdown().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn reasoning_effort_override_websocket_appends_then_replays_after_reconnect() +-> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let websocket = responses::start_websocket_server(vec![ + vec![ + vec![ + responses::ev_response_created("warmup"), + responses::ev_completed("warmup"), + ], + vec![ + responses::ev_response_created("first"), + responses::ev_completed("first"), + ], + vec![ + responses::ev_response_created("changed"), + responses::ev_completed("changed"), + ], + ], + vec![vec![ + responses::ev_response_created("replayed"), + responses::ev_completed("replayed"), + ]], + ]) + .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; + }) + .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), + ) + .await?; + assert_eq!(warmup.body_json()["generate"], false); + test.submit_text_turn("first message").await?; + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + effort: Some(Some(ReasoningEffort::High)), + ..Default::default() + }, + ) + .await?; + test.submit_text_turn("changed effort").await?; + test.submit_text_turn("unchanged effort after reconnect") + .await?; + + let connections = websocket.connections(); + assert_eq!(connections.iter().map(Vec::len).collect::>(), [3, 1]); + let first = connections[0][1].body_json(); + let changed = connections[0][2].body_json(); + let replayed = connections[1][0].body_json(); + assert_eq!(changed["previous_response_id"], "first"); + assert_eq!(replayed.get("previous_response_id"), None); + assert_eq!( + responses::strip_response_item_ids_from_json(responses::strip_metadata_from_json( + changed["input"].clone() + )), + serde_json::json!([ + message("user", "changed effort"), + effort_update(ReasoningEffort::High), + ]) + ); + assert_eq!( + replayed["input"] + .as_array() + .expect("replayed input") + .iter() + .filter(|item| item["type"] == "configuration_update") + .cloned() + .collect::>(), + [effort_update(ReasoningEffort::High)] + ); + assert_eq!( + [&first, &changed, &replayed].map(|body| body["reasoning"]["effort"].clone()), + [ + Value::from("medium"), + Value::from("medium"), + Value::from("medium") + ] + ); + websocket.shutdown().await; + Ok(()) +} + #[derive(Clone, Copy)] enum OverrideUnavailable { FeatureDisabled, @@ -314,3 +494,243 @@ async fn reasoning_effort_override_unavailable_uses_request_effort( ); Ok(()) } + +#[test_case(1, ReasoningEffort::High; "later turn keep changed selection")] +#[test_case(1, ReasoningEffort::Medium; "later turn match restored settings")] +#[test_case(2, ReasoningEffort::High; "through first turn")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn reasoning_effort_override_rollback_reestablishes_selected_effort( + num_turns: u32, + effort: ReasoningEffort, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + #[derive(Default)] + struct ThreadIdle { + ready: Notify, + } + + impl ThreadLifecycleContributor for ThreadIdle { + fn on_thread_idle<'a>(&'a self, _input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + self.ready.notify_one(); + }) + } + } + + let idle = Arc::new(ThreadIdle::default()); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(idle.clone()); + let server = responses::start_mock_server().await; + let first = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("first")]), + ) + .await; + let removed = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("removed")]), + ) + .await; + let replacement = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("replacement")]), + ) + .await; + let test = override_builder() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&server) + .await?; + test.submit_text_turn("first turn").await?; + timeout(Duration::from_secs(/*secs*/ 10), idle.ready.notified()).await?; + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + effort: Some(Some(ReasoningEffort::High)), + ..Default::default() + }, + ) + .await?; + test.submit_text_turn("removed turn").await?; + // Completion can be delivered before the active slot is cleared. Wait for + // actual idle so this test exercises effort replay, not rollback admission. + timeout(Duration::from_secs(/*secs*/ 10), idle.ready.notified()).await?; + test.codex.submit(Op::ThreadRollback { num_turns }).await?; + let rollback = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::ThreadRolledBack(_) | EventMsg::Error(_)) + }) + .await; + let EventMsg::ThreadRolledBack(rollback) = rollback else { + panic!("rollback failed: {rollback:?}"); + }; + assert_eq!(rollback.num_turns, num_turns); + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + effort: Some(Some(effort.clone())), + ..Default::default() + }, + ) + .await?; + test.submit_text_turn("replacement turn").await?; + timeout(Duration::from_secs(/*secs*/ 10), idle.ready.notified()).await?; + + assert_eq!( + effort_updates(&first.single_request()), + [effort_update(ReasoningEffort::Medium)] + ); + assert_eq!( + effort_updates(&removed.single_request()), + [ + effort_update(ReasoningEffort::Medium), + effort_update(ReasoningEffort::High) + ] + ); + let request = replacement.single_request(); + assert_eq!(request.body_json()["reasoning"]["effort"], effort.as_str()); + let mut expected = if num_turns == 1 { + vec![effort_update(ReasoningEffort::Medium)] + } else { + Vec::new() + }; + expected.push(effort_update(effort.clone())); + assert_eq!(effort_updates(&request), expected); + let input = request.input(); + assert!( + !input + .iter() + .any(|item| item["content"][0]["text"] == "removed turn") + ); + let tail = responses::strip_response_item_ids_from_json(responses::strip_metadata_from_json( + Value::Array(input[input.len() - 2..].to_vec()), + )); + assert_eq!( + tail, + serde_json::json!([ + message("user", "replacement turn"), + effort_update(effort.clone()), + ]) + ); + Ok(()) +} + +#[test_case(ReasoningEffort::Medium; "same selection")] +#[test_case(ReasoningEffort::High; "changed selection")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn reasoning_effort_override_model_switch_reestablishes_selected_effort( + effort: ReasoningEffort, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let first = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("first")]), + ) + .await; + let second = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("second")]), + ) + .await; + let test = override_builder() + .with_model_info_override("gpt-5.4", |model| model.comp_hash = None) + .with_model_info_override("gpt-5.5", |model| { + model.comp_hash = None; + model.use_responses_lite = true; + }) + .with_model("gpt-5.4") + .build_with_auto_env(&server) + .await?; + test.submit_text_turn("first model").await?; + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + model: Some("gpt-5.5".to_string()), + effort: Some(Some(effort.clone())), + ..Default::default() + }, + ) + .await?; + test.submit_text_turn("second model").await?; + assert_eq!( + effort_updates(&first.single_request()), + [effort_update(ReasoningEffort::Medium)] + ); + let request = second.single_request(); + assert_eq!( + effort_updates(&request), + [ + effort_update(ReasoningEffort::Medium), + effort_update(effort.clone()) + ] + ); + assert_eq!( + serde_json::json!([ + request.body_json()["model"], + request.body_json()["reasoning"]["effort"] + ]), + serde_json::json!(["gpt-5.5", effort]) + ); + Ok(()) +} + +#[test_case(ReasoningEffort::High; "unchanged effort")] +#[test_case(ReasoningEffort::Medium; "changed effort")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn reasoning_effort_override_resume_refreshes_selected_effort( + effort: ReasoningEffort, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let first = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("first")]), + ) + .await; + let second = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("second")]), + ) + .await; + let resumed = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("resumed")]), + ) + .await; + let test = override_builder().build_with_auto_env(&server).await?; + test.submit_text_turn("medium turn").await?; + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + effort: Some(Some(ReasoningEffort::High)), + ..Default::default() + }, + ) + .await?; + test.submit_text_turn("high turn").await?; + let selected = effort.clone(); + let restarted = override_builder() + .with_config(move |config| { + config.model_reasoning_effort = Some(selected); + }) + .restart(&server, &test) + .await?; + restarted.submit_text_turn("after resume").await?; + + assert_eq!( + [first.single_request(), second.single_request()] + .map(|request| request.body_json()["reasoning"]["effort"].clone()), + [Value::from("medium"), Value::from("medium")] + ); + let request = resumed.single_request(); + assert_eq!(request.body_json()["reasoning"]["effort"], effort.as_str()); + // Even an unchanged selection must refresh the override when replay invalidates the pin. + assert_eq!( + effort_updates(&request), + [ + effort_update(ReasoningEffort::Medium), + effort_update(ReasoningEffort::High), + effort_update(effort.clone()) + ] + ); + Ok(()) +}