[codex-core] Fix inline compaction retry and model-switch preflight [ci changed_files]

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Cooper Gamble
2026-03-11 02:48:35 +00:00
co-authored by Codex
parent 9d20427fc2
commit d9916b5e6d
3 changed files with 274 additions and 32 deletions
+25 -5
View File
@@ -41,6 +41,7 @@ use crate::realtime_conversation::handle_start as handle_realtime_conversation_s
use crate::realtime_conversation::handle_text as handle_realtime_conversation_text;
use crate::rollout::session_index;
use crate::stream_events_utils::HandleOutputCtx;
use crate::stream_events_utils::PendingServerSideCompactionCheckpoint;
use crate::stream_events_utils::handle_non_tool_response_item;
use crate::stream_events_utils::handle_output_item_done;
use crate::stream_events_utils::last_assistant_message_from_item;
@@ -6097,10 +6098,6 @@ async fn maybe_run_previous_model_inline_compact(
turn_context: &Arc<TurnContext>,
total_usage_tokens: i64,
) -> CodexResult<bool> {
if inline_server_side_compaction_threshold(sess, turn_context).is_some() {
return Ok(false);
}
let Some(previous_turn_settings) = sess.previous_turn_settings().await else {
return Ok(false);
};
@@ -7154,6 +7151,9 @@ async fn try_run_sampling_request(
let mut needs_follow_up = false;
let mut last_agent_message: Option<String> = None;
let mut observed_server_side_compaction = false;
let mut pending_server_side_compaction_checkpoint: Option<
PendingServerSideCompactionCheckpoint,
> = None;
let mut active_item: Option<TurnItem> = None;
let mut should_emit_turn_diff = false;
let plan_mode = turn_context.collaboration_mode.mode == ModeKind::Plan;
@@ -7232,13 +7232,15 @@ async fn try_run_sampling_request(
turn_context: turn_context.clone(),
tool_runtime: tool_runtime.clone(),
cancellation_token: cancellation_token.child_token(),
history_before_turn: history_before_turn.clone(),
};
let output_result = handle_output_item_done(&mut ctx, item, previously_active_item)
.instrument(handle_responses)
.await?;
observed_server_side_compaction |= saw_server_side_compaction;
if let Some(pending_compaction) = output_result.pending_server_side_compaction {
pending_server_side_compaction_checkpoint = Some(pending_compaction);
}
if let Some(tool_future) = output_result.tool_future {
in_flight.push_back(tool_future);
}
@@ -7257,6 +7259,10 @@ async fn try_run_sampling_request(
.await
{
let mut turn_item = turn_item;
if matches!(turn_item, TurnItem::ContextCompaction(_)) {
active_item = Some(turn_item);
continue;
}
let mut seeded_parsed: Option<ParsedAssistantTextDelta> = None;
let mut seeded_item_id: Option<String> = None;
if matches!(turn_item, TurnItem::AgentMessage(_))
@@ -7337,6 +7343,20 @@ async fn try_run_sampling_request(
&mut assistant_message_stream_parsers,
)
.await;
if let Some(PendingServerSideCompactionCheckpoint { item, turn_item }) =
pending_server_side_compaction_checkpoint.take()
{
let turn_item = TurnItem::ContextCompaction(turn_item);
sess.emit_turn_item_started(&turn_context, &turn_item).await;
sess.apply_server_side_compaction_checkpoint(
turn_context.as_ref(),
item,
history_before_turn.as_slice(),
)
.await;
sess.emit_turn_item_completed(&turn_context, turn_item)
.await;
}
sess.update_token_usage_info(&turn_context, token_usage.as_ref())
.await;
should_emit_turn_diff = true;
+11 -27
View File
@@ -145,15 +145,20 @@ pub(crate) type InFlightFuture<'f> =
pub(crate) struct OutputItemResult {
pub last_agent_message: Option<String>,
pub needs_follow_up: bool,
pub pending_server_side_compaction: Option<PendingServerSideCompactionCheckpoint>,
pub tool_future: Option<InFlightFuture<'static>>,
}
pub(crate) struct PendingServerSideCompactionCheckpoint {
pub item: ResponseItem,
pub turn_item: ContextCompactionItem,
}
pub(crate) struct HandleOutputCtx {
pub sess: Arc<Session>,
pub turn_context: Arc<TurnContext>,
pub tool_runtime: ToolCallRuntime,
pub cancellation_token: CancellationToken,
pub history_before_turn: Arc<Vec<ResponseItem>>,
}
#[instrument(level = "trace", skip_all)]
@@ -166,39 +171,18 @@ pub(crate) async fn handle_output_item_done(
let plan_mode = ctx.turn_context.collaboration_mode.mode == ModeKind::Plan;
if matches!(item, ResponseItem::Compaction { .. }) {
let had_active_compaction = matches!(
previously_active_item.as_ref(),
Some(TurnItem::ContextCompaction(_))
);
let compaction_item = match previously_active_item {
Some(TurnItem::ContextCompaction(item)) => item,
_ => ContextCompactionItem::new(),
};
if !had_active_compaction {
ctx.sess
.emit_turn_item_started(
&ctx.turn_context,
&TurnItem::ContextCompaction(compaction_item.clone()),
)
.await;
}
debug!(
turn_id = %ctx.turn_context.sub_id,
"processing streamed server-side compaction item"
"buffering streamed server-side compaction item until response.completed"
);
ctx.sess
.apply_server_side_compaction_checkpoint(
ctx.turn_context.as_ref(),
item,
ctx.history_before_turn.as_slice(),
)
.await;
ctx.sess
.emit_turn_item_completed(
&ctx.turn_context,
TurnItem::ContextCompaction(compaction_item),
)
.await;
output.pending_server_side_compaction = Some(PendingServerSideCompactionCheckpoint {
item,
turn_item: compaction_item,
});
return Ok(output);
}
+238
View File
@@ -5,11 +5,15 @@ use std::path::PathBuf;
use anyhow::Result;
use codex_core::CodexAuth;
use codex_core::ModelProviderInfo;
use codex_core::WireApi;
use codex_core::compact::SUMMARY_PREFIX;
use codex_core::features::Feature;
use codex_protocol::items::TurnItem;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::protocol::ConversationStartParams;
use codex_protocol::protocol::ErrorEvent;
use codex_protocol::protocol::EventMsg;
@@ -60,6 +64,18 @@ fn summary_with_prefix(summary: &str) -> String {
format!("{SUMMARY_PREFIX}\n{summary}")
}
fn model_info_with_context_window(slug: &str, context_window: i64) -> ModelInfo {
let models_response: ModelsResponse =
serde_json::from_str(include_str!("../../models.json")).expect("valid models.json");
let mut model_info = models_response
.models
.into_iter()
.find(|model| model.slug == slug)
.unwrap_or_else(|| panic!("model `{slug}` missing from models.json"));
model_info.context_window = Some(context_window);
model_info
}
fn context_snapshot_options() -> ContextSnapshotOptions {
ContextSnapshotOptions::default()
.render_mode(ContextSnapshotRenderMode::KindWithTextPrefix { max_chars: 64 })
@@ -592,6 +608,93 @@ async fn auto_server_side_compaction_keeps_current_turn_inputs_for_follow_ups()
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auto_server_side_compaction_retries_without_committing_incomplete_checkpoint() -> Result<()>
{
skip_if_no_network!(Ok(()));
let compact_threshold = 120;
let first_turn_text = "inline compact retry turn one";
let inline_summary = summary_with_prefix("INLINE_SERVER_RETRY_SUMMARY");
let server = wiremock::MockServer::start().await;
let model_provider = ModelProviderInfo {
name: "openai".into(),
base_url: Some(format!("{}/v1", server.uri())),
env_key: Some("PATH".into()),
env_key_instructions: None,
experimental_bearer_token: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
env_http_headers: None,
request_max_retries: Some(0),
stream_max_retries: Some(1),
stream_idle_timeout_ms: Some(2000),
requires_openai_auth: false,
supports_websockets: false,
};
let responses_mock = responses::mount_response_sequence(
&server,
vec![
responses::sse_response(sse(vec![
responses::ev_assistant_message("m1", "FIRST_REMOTE_REPLY"),
responses::ev_completed_with_tokens("resp-1", 500),
])),
responses::sse_response(sse(vec![responses::ev_compaction(&inline_summary)])),
responses::sse_response(sse(vec![
responses::ev_assistant_message("m2", "AFTER_RETRY_REPLY"),
responses::ev_completed_with_tokens("resp-2", 40),
])),
responses::sse_response(sse(vec![
responses::ev_assistant_message("m3", "THIRD_TURN_REPLY"),
responses::ev_completed("resp-3"),
])),
],
)
.await;
let test = test_codex()
.with_config(move |config| {
config.model_provider = model_provider;
config
.features
.enable(Feature::ServerSideCompaction)
.expect("enable server-side compaction");
config.model_auto_compact_token_limit = Some(compact_threshold);
})
.build(&server)
.await?;
let codex = test.codex.clone();
submit_text_turn_and_wait(&codex, first_turn_text).await?;
submit_text_turn_and_wait(&codex, "inline compact retry turn two").await?;
submit_text_turn_and_wait(&codex, "inline compact retry turn three").await?;
let requests = responses_mock.requests();
assert_eq!(
requests.len(),
4,
"expected initial turn, incomplete inline attempt, retry, and third turn"
);
let third_turn_request = &requests[3];
assert!(
!third_turn_request.body_contains_text(&inline_summary),
"failed inline compaction should not rewrite local history before a retry succeeds"
);
assert!(
third_turn_request.body_contains_text(first_turn_text),
"failed inline compaction should leave pre-compaction history intact for later turns"
);
assert!(
third_turn_request.body_contains_text("AFTER_RETRY_REPLY"),
"retry response should be the history that later turns build on"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auto_server_side_compaction_stays_inline_with_custom_prompt() -> Result<()> {
skip_if_no_network!(Ok(()));
@@ -663,6 +766,141 @@ async fn auto_server_side_compaction_stays_inline_with_custom_prompt() -> Result
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn previous_model_preflight_compaction_still_runs_with_inline_server_side_compaction()
-> Result<()> {
skip_if_no_network!(Ok(()));
let server = wiremock::MockServer::start().await;
let previous_model = "gpt-5.2-codex";
let next_model = "gpt-5.1-codex-max";
let _models_mock = responses::mount_models_once(
&server,
ModelsResponse {
models: vec![
model_info_with_context_window(previous_model, 273_000),
model_info_with_context_window(next_model, 125_000),
],
},
)
.await;
let test = test_codex()
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_model(previous_model)
.with_config(|config| {
config
.features
.enable(Feature::ServerSideCompaction)
.expect("enable server-side compaction");
config.model_auto_compact_token_limit = Some(200);
})
.build(&server)
.await?;
let codex = test.codex.clone();
let initial_turn_request_mock = responses::mount_sse_once(
&server,
responses::sse(vec![
responses::ev_assistant_message("m1", "BEFORE_SWITCH_REPLY"),
responses::ev_completed_with_tokens("r1", 120_000),
]),
)
.await;
let post_compact_turn_request_mock = responses::mount_sse_once(
&server,
responses::sse(vec![
responses::ev_assistant_message("m2", "AFTER_SWITCH_REPLY"),
responses::ev_completed_with_tokens("r2", 80),
]),
)
.await;
let compact_mock = responses::mount_compact_user_history_with_summary_once(
&server,
&summary_with_prefix("REMOTE_SWITCH_SUMMARY"),
)
.await;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "BEFORE_SWITCH_USER".to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
codex
.submit(Op::OverrideTurnContext {
cwd: None,
approval_policy: None,
sandbox_policy: None,
windows_sandbox_level: None,
model: Some(next_model.to_string()),
effort: None,
summary: None,
service_tier: None,
collaboration_mode: None,
personality: None,
})
.await?;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "AFTER_SWITCH_USER".to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
assert_eq!(
initial_turn_request_mock.requests().len(),
1,
"expected initial turn request"
);
assert_eq!(
compact_mock.requests().len(),
1,
"expected previous-model preflight compaction before the smaller-model turn"
);
assert_eq!(
post_compact_turn_request_mock.requests().len(),
1,
"expected a single post-compaction smaller-model request"
);
let compact_request = compact_mock.single_request();
let compact_body = compact_request.body_json().to_string();
assert!(
!compact_body.contains("AFTER_SWITCH_USER"),
"preflight compaction should still exclude the incoming user message"
);
assert!(
!compact_body.contains("<model_switch>"),
"preflight compaction should still strip the incoming model switch item"
);
let post_compact_turn_request = post_compact_turn_request_mock.single_request();
assert!(
post_compact_turn_request.body_contains_text("REMOTE_SWITCH_SUMMARY"),
"smaller-model follow-up should use the previous-model compaction summary"
);
assert!(
post_compact_turn_request.body_contains_text("AFTER_SWITCH_USER"),
"smaller-model follow-up should still include the incoming user message"
);
assert!(
post_compact_turn_request.body_contains_text("<model_switch>"),
"smaller-model follow-up should still include the model switch item"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auto_server_side_compaction_reports_inline_compat_errors_without_fallback() -> Result<()> {
skip_if_no_network!(Ok(()));