mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
@@ -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(
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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(()));
|
||||
|
||||
Reference in New Issue
Block a user