From a17da5e6e4a5a9b45396f0693b0a4d5b9df06318 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 6 Aug 2026 10:58:25 +0000 Subject: [PATCH] Fix first-turn model switching and rollback (#37260) ## Why A model switch on the first turn has no prior turn settings to compare, and rolling that turn back could leave its model instructions in history. ## What changed - Track the model that supplied the session's base instructions so first-turn switches inject the target model's instructions, including the selected personality. - Preserve inherited custom base instructions when starting from an empty fork. - Remove first-turn model-switch developer content during rollback so retries and cold resumes do not retain or duplicate stale instructions. ## Testing Add integration coverage for first-turn switches, personality changes, empty-prefix forks, rollback retries, and cold resume. GitOrigin-RevId: 9df532e9e337521ad3b315b77f8c98e039fb384e --- codex-rs/core/src/context_manager/history.rs | 26 ++- codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/session.rs | 7 +- codex-rs/core/src/session/tests.rs | 3 + codex-rs/core/src/session/world_state.rs | 31 ++- codex-rs/core/src/state/session.rs | 3 + codex-rs/core/tests/suite/model_switching.rs | 215 +++++++++++++++++++ 7 files changed, 273 insertions(+), 13 deletions(-) diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index a3cf288ee5..d900ec9813 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,4 +1,5 @@ use crate::context::ContextualUserFragment; +use crate::context::ModelSwitchInstructions; use crate::context::world_state::WorldState; use crate::context::world_state::WorldStateSnapshot; use crate::context_manager::normalize; @@ -245,7 +246,30 @@ impl ContextManager { cut_idx = self.trim_pre_turn_context_updates(&snapshot, first_instruction_turn_idx, cut_idx); - self.replace(snapshot[..cut_idx].to_vec()); + let mut retained_items = snapshot[..cut_idx].to_vec(); + if cut_idx == first_instruction_turn_idx + && let Some(first_turn_id) = snapshot[first_instruction_turn_idx].turn_id() + { + retained_items.retain_mut(|item| { + if item.turn_id() == Some(first_turn_id) + && let ResponseItem::Message { role, content, .. } = item + && role == "developer" + { + content.retain(|content| { + !matches!( + content, + ContentItem::InputText { text } + if ModelSwitchInstructions::matches_text(text) + ) + }); + !content.is_empty() + } else { + true + } + }); + } + + self.replace(retained_items); } pub(crate) fn update_token_info( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 96021f37a3..f69ff50476 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -729,6 +729,7 @@ impl Session { installation_id, auth_manager.clone(), models_manager.clone(), + model_info, exec_policy, tx_event.clone(), agent_status_tx.clone(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index e5de254b07..522055b7b2 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -515,6 +515,7 @@ impl Session { installation_id: String, auth_manager: Arc, models_manager: SharedModelsManager, + model_info: ModelInfo, exec_policy: Arc, tx_event: Sender, agent_status: watch::Sender, @@ -999,10 +1000,14 @@ impl Session { session_configuration.thread_name = thread_name.clone(); validate_config_lock_if_configured(&session_configuration).await?; export_config_lock_if_configured(&session_configuration, thread_id).await?; - let state = SessionState::new_with_auto_compact_window_ids( + let mut state = SessionState::new_with_auto_compact_window_ids( session_configuration.clone(), initial_auto_compact_window_ids, ); + state.base_instructions_model = (config.base_instructions.is_none() + && session_configuration.base_instructions + == model_info.get_model_instructions(config.personality)) + .then(|| model_info.slug.clone()); let managed_network_requirements_configured = config .config_layer_stack .requirements_toml() diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index be87b859c7..3a0ae6850a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5581,6 +5581,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + model_info, Arc::new(ExecPolicyManager::default()), tx_event, agent_status_tx, @@ -5995,6 +5996,7 @@ async fn make_session_with_config_and_rx( "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + model_info, Arc::new(ExecPolicyManager::default()), tx_event, agent_status_tx, @@ -6108,6 +6110,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + model_info, Arc::new(ExecPolicyManager::default()), tx_event, agent_status_tx, diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index a61a545d0e..14ac8fc656 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -37,19 +37,27 @@ impl Session { selected_capability_root_count = step_context.selected_capability_roots.len(), "building step world state" ); - let (previous_model, previous_context, base_instructions) = { - let state = self.state.lock().await; - ( - state - .previous_turn_settings() - .map(|previous| previous.model), - state.reference_context_item(), - state.session_configuration.base_instructions.clone(), - ) - }; let model_instructions = turn_context .model_info .get_model_instructions(turn_context.personality); + let (previous_model, previous_context, base_instructions) = { + let state = self.state.lock().await; + let base_instructions = state.session_configuration.base_instructions.clone(); + ( + state + .previous_turn_settings() + .map(|previous| previous.model) + .or_else(|| { + state + .base_instructions_model + .as_ref() + .filter(|_| base_instructions != model_instructions) + .cloned() + }), + state.reference_context_item(), + base_instructions, + ) + }; let personality_is_baked = turn_context.model_info.supports_personality() && base_instructions == model_instructions; let environment_subagents = if turn_context.config.include_environment_context { @@ -80,7 +88,8 @@ impl Session { turn_context.personality, previous_context .as_ref() - .map(|previous| previous.model.as_str()), + .map(|previous| previous.model.as_str()) + .or(previous_model.as_deref()), previous_context .as_ref() .and_then(|previous| previous.personality), diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 28a08ef4d6..f0b4e9ccc6 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -25,6 +25,8 @@ use codex_utils_output_truncation::TruncationPolicy; /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, + /// Model that generated the base instructions, when their provenance is known. + pub(crate) base_instructions_model: Option, pub(crate) history: ContextManager, pub(crate) latest_rate_limits: Option, pub(crate) server_reasoning_included: bool, @@ -62,6 +64,7 @@ impl SessionState { let history = ContextManager::new(); Self { session_configuration, + base_instructions_model: None, history, latest_rate_limits: None, server_reasoning_included: false, diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs index ce5f8fc9ba..4155f253ac 100644 --- a/codex-rs/core/tests/suite/model_switching.rs +++ b/codex-rs/core/tests/suite/model_switching.rs @@ -1,7 +1,10 @@ use anyhow::Result; use codex_config::types::Personality; +use codex_core::CodexThread; +use codex_core::ForkSnapshot; use codex_features::Feature; use codex_login::CodexAuth; +use codex_models_manager::bundled_models_response; use codex_models_manager::manager::RefreshStrategy; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; @@ -22,6 +25,7 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_image_generation_call; @@ -39,6 +43,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; +use test_case::test_case; use wiremock::MockServer; fn read_only_user_turn(test: &TestCodex, items: Vec, model: String) -> Op { @@ -67,6 +72,28 @@ fn read_only_user_turn(test: &TestCodex, items: Vec, model: String) - } } +async fn submit_model_turn( + thread: &CodexThread, + model: &str, + mut thread_settings: ThreadSettingsOverrides, +) -> Result<()> { + thread_settings.model = Some(model.to_string()); + thread + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "switch models".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings, + }) + .await?; + wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; + Ok(()) +} + fn test_model_info( slug: &str, display_name: &str, @@ -121,6 +148,194 @@ fn test_model_info( } } +#[test_case(None; "model only")] +#[test_case(Some(Personality::Pragmatic); "model and personality")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn first_turn_model_change_appends_model_instructions_developer_message( + personality: Option, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + let resp_mock = mount_sse_once(&server, sse_completed("resp-1")).await; + + let mut builder = test_codex().with_model("gpt-5.2").with_config(|config| { + config + .features + .enable(Feature::Personality) + .expect("test config should allow feature update"); + }); + let test = builder.build_with_auto_env(&server).await?; + let next_model = "gpt-5.4"; + + submit_model_turn( + &test.codex, + next_model, + ThreadSettingsOverrides { + personality, + ..Default::default() + }, + ) + .await?; + + let request = resp_mock.single_request(); + assert_eq!(request.body_json()["model"], next_model); + let developer_texts = request.message_input_texts("developer"); + let expected_instructions = bundled_models_response()? + .models + .into_iter() + .find(|model| model.slug == next_model) + .expect("expected target model in bundled catalog") + .get_model_instructions(personality.or(test.config.personality)); + assert!( + developer_texts.iter().any(|text| { + text.contains("") && text.contains(&expected_instructions) + }) + ); + assert!( + developer_texts + .iter() + .all(|text| !text.contains("")), + "model instructions already include the selected personality" + ); + + Ok(()) +} + +#[test_case(None; "model-generated base instructions")] +#[test_case(Some("inherited custom base instructions"); "custom base instructions")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn first_turn_after_empty_prefix_fork_preserves_inherited_base_instructions( + custom_base_instructions: Option<&'static str>, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + let resp_mock = mount_sse_once(&server, sse_completed("resp-fork")).await; + + let initial_model = "gpt-5.2"; + let mut builder = test_codex() + .with_model(initial_model) + .with_config(move |config| { + config.base_instructions = custom_base_instructions.map(str::to_string); + }); + let test = builder.build_with_auto_env(&server).await?; + test.codex.ensure_rollout_materialized().await; + test.codex.flush_rollout().await?; + + let mut fork_config = test.config.clone(); + fork_config.model = Some("gpt-5.4".to_string()); + fork_config.base_instructions = None; + let fork = test + .thread_manager + .fork_thread( + ForkSnapshot::TruncateBeforeNthUserMessage(0), + fork_config, + test.codex.rollout_path().expect("rollout path"), + /*thread_source*/ None, + /*parent_trace*/ None, + ) + .await?; + submit_model_turn( + &fork.thread, + initial_model, + ThreadSettingsOverrides::default(), + ) + .await?; + + let request = resp_mock.single_request(); + assert_eq!(request.body_json()["model"], initial_model); + if let Some(instructions) = custom_base_instructions { + assert_eq!(request.instructions_text(), instructions); + } + assert!( + request + .message_input_texts("developer") + .iter() + .all(|text| !text.contains("")), + "the inherited base instructions must not be replaced by catalog instructions" + ); + + Ok(()) +} + +#[derive(Clone, Copy)] +enum RollbackFollowup { + StartupModel, + SwitchedModel, + ColdResume, +} + +#[test_case(RollbackFollowup::StartupModel; "return to startup model")] +#[test_case(RollbackFollowup::SwitchedModel; "retry switched model")] +#[test_case(RollbackFollowup::ColdResume; "retry switched model after cold resume")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rollback_first_turn_model_change_removes_its_instructions( + followup: RollbackFollowup, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + let response_mock = mount_sse_sequence( + &server, + vec![sse_completed("resp-first"), sse_completed("resp-followup")], + ) + .await; + + let initial_model = "gpt-5.2"; + let switched_model = "gpt-5.4"; + let mut builder = test_codex().with_model(initial_model); + let test = builder.build_with_auto_env(&server).await?; + + submit_model_turn( + &test.codex, + switched_model, + ThreadSettingsOverrides::default(), + ) + .await?; + + test.codex + .submit(Op::ThreadRollback { num_turns: 1 }) + .await?; + wait_for_event(&test.codex, |ev| { + matches!(ev, EventMsg::ThreadRolledBack(_)) + }) + .await; + + let test = match followup { + RollbackFollowup::ColdResume => { + let mut resume_builder = test_codex().with_model(initial_model); + resume_builder.restart(&server, &test).await? + } + RollbackFollowup::StartupModel | RollbackFollowup::SwitchedModel => test, + }; + let followup_model = match followup { + RollbackFollowup::StartupModel => initial_model, + RollbackFollowup::SwitchedModel | RollbackFollowup::ColdResume => switched_model, + }; + submit_model_turn( + &test.codex, + followup_model, + ThreadSettingsOverrides::default(), + ) + .await?; + + let request = &response_mock.requests()[1]; + assert_eq!(request.body_json()["model"], followup_model); + let model_switch_count = request + .message_input_texts("developer") + .iter() + .filter(|text| text.contains("")) + .count(); + assert_eq!( + model_switch_count, + usize::from(followup_model == switched_model), + "rolled-back model instructions must not survive or be duplicated" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn model_change_appends_model_instructions_developer_message() -> Result<()> { skip_if_no_network!(Ok(()));