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
This commit is contained in:
jif
2026-08-06 11:03:31 +00:00
committed by copyberry
parent 82b17bc724
commit a17da5e6e4
7 changed files with 273 additions and 13 deletions
+25 -1
View File
@@ -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(
+1
View File
@@ -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(),
+6 -1
View File
@@ -515,6 +515,7 @@ impl Session {
installation_id: String,
auth_manager: Arc<AuthManager>,
models_manager: SharedModelsManager,
model_info: ModelInfo,
exec_policy: Arc<ExecPolicyManager>,
tx_event: Sender<Event>,
agent_status: watch::Sender<AgentStatus>,
@@ -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()
+3
View File
@@ -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,
+20 -11
View File
@@ -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),
+3
View File
@@ -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<String>,
pub(crate) history: ContextManager,
pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
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,
@@ -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<UserInput>, model: String) -> Op {
@@ -67,6 +72,28 @@ fn read_only_user_turn(test: &TestCodex, items: Vec<UserInput>, 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<Personality>,
) -> 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("<model_switch>") && text.contains(&expected_instructions)
})
);
assert!(
developer_texts
.iter()
.all(|text| !text.contains("<personality_spec>")),
"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("<model_switch>")),
"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("<model_switch>"))
.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(()));