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
This commit is contained in:
felixxia-oai
2026-09-08 10:41:15 +00:00
committed by copyberry
parent df522cae16
commit 31ccaf40c2
7 changed files with 504 additions and 14 deletions
+3
View File
@@ -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);
}
+30 -4
View File
@@ -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<ReasoningEffort> {
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,
+2 -1
View File
@@ -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,
+4 -1
View File
@@ -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,
+1
View File
@@ -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;
+36
View File
@@ -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<ReasoningEffort> {
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<PreviousTurnSettings>,
/// 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<SessionStartupPrewarmHandle>,
/// 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(),
@@ -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::<Vec<_>>();
assert_eq!(
mocks
requests
.iter()
.map(|mock| effort_updates(&mock.single_request()))
.map(|request| request.body_json()["reasoning"]["effort"].clone())
.collect::<Vec<_>>(),
vec![initial_update["reasoning"]["effort"].clone(); requests.len()],
);
assert_eq!(
requests.iter().map(effort_updates).collect::<Vec<_>>(),
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<_>>(),
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::<Vec<_>>();
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::<Vec<_>>(), [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::<Vec<_>>(),
[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<Config> 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(())
}