mirror of
https://github.com/openai/codex.git
synced 2026-10-29 19:08:43 +00:00
[core] simplify inline server-side compaction handling [ci changed_files]
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
+37
-170
@@ -20,7 +20,6 @@ use crate::apps::render_apps_section;
|
||||
use crate::commit_attribution::commit_message_trailer_instruction;
|
||||
use crate::compact;
|
||||
use crate::compact::InitialContextInjection;
|
||||
use crate::compact::insert_initial_context_before_last_real_user_or_summary;
|
||||
use crate::compact::run_inline_auto_compact_task;
|
||||
use crate::compact::should_use_remote_compact_task;
|
||||
use crate::compact_remote::run_inline_remote_auto_compact_task;
|
||||
@@ -42,7 +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::ServerSideCompaction;
|
||||
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;
|
||||
@@ -3236,25 +3235,14 @@ impl Session {
|
||||
state.record_items(items.iter(), turn_context.truncation_policy);
|
||||
}
|
||||
|
||||
pub(crate) async fn apply_server_side_compaction_checkpoint(
|
||||
pub(crate) async fn apply_server_side_compaction(
|
||||
&self,
|
||||
turn_context: &TurnContext,
|
||||
item: ResponseItem,
|
||||
compaction_initial_context: &[ResponseItem],
|
||||
turn_start_context_items: &[ResponseItem],
|
||||
history_before_turn: &[ResponseItem],
|
||||
history_at_checkpoint: &[ResponseItem],
|
||||
) {
|
||||
// The server emits compaction as a streamed item before the response is fully complete.
|
||||
// Wait until `response.completed` to rewrite local history so later streamed items from the
|
||||
// same turn can still be appended in wire order before we collapse the checkpoint.
|
||||
let current_history = self.clone_history().await;
|
||||
let replacement_history = build_server_side_compaction_replacement_history(
|
||||
item.clone(),
|
||||
compaction_initial_context,
|
||||
turn_start_context_items,
|
||||
history_before_turn,
|
||||
history_at_checkpoint,
|
||||
current_history.raw_items(),
|
||||
);
|
||||
|
||||
@@ -3265,9 +3253,10 @@ impl Session {
|
||||
};
|
||||
self.replace_compacted_history(replacement_history, reference_context_item, compacted_item)
|
||||
.await;
|
||||
self.recompute_token_usage(turn_context).await;
|
||||
debug!(
|
||||
turn_id = %turn_context.sub_id,
|
||||
"applied local server-side compaction checkpoint"
|
||||
"applied local server-side compaction"
|
||||
);
|
||||
self.services.session_telemetry.counter(
|
||||
"codex.compaction_checkpoint_applied",
|
||||
@@ -5398,21 +5387,17 @@ pub(crate) async fn run_turn(
|
||||
// new user message are recorded. Estimate pending incoming items (context
|
||||
// diffs/full reinjection + user input) and trigger compaction preemptively
|
||||
// when they would push the thread over the compaction threshold.
|
||||
let mut pending_server_side_compaction =
|
||||
match run_pre_sampling_compact(&sess, &turn_context).await {
|
||||
Ok(pending) => pending,
|
||||
Err(err) => {
|
||||
error!("Failed to run pre-sampling compact: {err}");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let mut inline_compaction_trigger = match run_pre_sampling_compact(&sess, &turn_context).await {
|
||||
Ok(trigger) => trigger,
|
||||
Err(err) => {
|
||||
error!("Failed to run pre-sampling compact: {err}");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref());
|
||||
|
||||
let history_before_turn = sess.clone_history().await.raw_items().to_vec();
|
||||
let compaction_initial_context = sess.build_initial_context(turn_context.as_ref()).await;
|
||||
let turn_start_context_items = sess
|
||||
.record_context_updates_and_set_reference_context_item(turn_context.as_ref())
|
||||
sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref())
|
||||
.await;
|
||||
|
||||
let loaded_plugins = sess
|
||||
@@ -5689,7 +5674,6 @@ pub(crate) async fn run_turn(
|
||||
.map(|user_message| user_message.message())
|
||||
.collect::<Vec<String>>();
|
||||
let turn_metadata_header = turn_context.turn_metadata_state.current_header_value();
|
||||
let inline_compaction_for_request = pending_server_side_compaction;
|
||||
match run_sampling_request(
|
||||
Arc::clone(&sess),
|
||||
Arc::clone(&turn_context),
|
||||
@@ -5697,10 +5681,7 @@ pub(crate) async fn run_turn(
|
||||
&mut client_session,
|
||||
turn_metadata_header.as_deref(),
|
||||
sampling_request_input,
|
||||
&compaction_initial_context,
|
||||
&turn_start_context_items,
|
||||
&history_before_turn,
|
||||
inline_compaction_for_request.map(|pending| pending.threshold),
|
||||
inline_server_side_compaction_threshold(&sess, &turn_context),
|
||||
&turn_enabled_connectors,
|
||||
skills_outcome,
|
||||
&mut server_model_warning_emitted_for_turn,
|
||||
@@ -5709,7 +5690,7 @@ pub(crate) async fn run_turn(
|
||||
.await
|
||||
{
|
||||
Ok(sampling_request_output) => {
|
||||
if let Some(pending_compaction) = inline_compaction_for_request {
|
||||
if let Some(trigger) = inline_compaction_trigger.take() {
|
||||
let result = if sampling_request_output.observed_server_side_compaction {
|
||||
"applied"
|
||||
} else {
|
||||
@@ -5720,15 +5701,8 @@ pub(crate) async fn run_turn(
|
||||
} else {
|
||||
vec![("reason", "threshold_not_reached")]
|
||||
};
|
||||
record_compaction_metric(
|
||||
&sess,
|
||||
"server_side",
|
||||
pending_compaction.trigger,
|
||||
result,
|
||||
&extra_tags,
|
||||
);
|
||||
record_compaction_metric(&sess, "server_side", trigger, result, &extra_tags);
|
||||
}
|
||||
pending_server_side_compaction = None;
|
||||
let SamplingRequestResult {
|
||||
needs_follow_up,
|
||||
last_agent_message: sampling_request_last_agent_message,
|
||||
@@ -5752,9 +5726,7 @@ pub(crate) async fn run_turn(
|
||||
|
||||
// as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop.
|
||||
if token_limit_reached && needs_follow_up {
|
||||
if let Some(threshold) =
|
||||
inline_server_side_compaction_threshold(&sess, &turn_context)
|
||||
{
|
||||
if inline_server_side_compaction_threshold(&sess, &turn_context).is_some() {
|
||||
record_compaction_metric(
|
||||
&sess,
|
||||
"server_side",
|
||||
@@ -5762,10 +5734,7 @@ pub(crate) async fn run_turn(
|
||||
"requested",
|
||||
&[],
|
||||
);
|
||||
pending_server_side_compaction = Some(PendingServerSideCompaction {
|
||||
threshold,
|
||||
trigger: AutoCompactTrigger::AutoFollowUp,
|
||||
});
|
||||
inline_compaction_trigger = Some(AutoCompactTrigger::AutoFollowUp);
|
||||
} else {
|
||||
record_inline_compaction_skip(
|
||||
&sess,
|
||||
@@ -5952,70 +5921,13 @@ impl AutoCompactTrigger {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
struct PendingServerSideCompaction {
|
||||
threshold: i64,
|
||||
trigger: AutoCompactTrigger,
|
||||
}
|
||||
|
||||
fn build_server_side_compaction_replacement_history(
|
||||
compaction_item: ResponseItem,
|
||||
compaction_initial_context: &[ResponseItem],
|
||||
turn_start_context_items: &[ResponseItem],
|
||||
history_before_turn: &[ResponseItem],
|
||||
history_at_checkpoint: &[ResponseItem],
|
||||
current_history: &[ResponseItem],
|
||||
history_at_compaction: &[ResponseItem],
|
||||
) -> Vec<ResponseItem> {
|
||||
// Rebuild the active turn around the compaction checkpoint:
|
||||
// 1. keep the turn-local items that existed when compaction fired
|
||||
// 2. replace any prior same-turn compaction summary with the newest one
|
||||
// 3. re-append items that arrived later in the same streamed response
|
||||
// 4. reattach ghost snapshots at the end so undo state survives the rewrite
|
||||
let checkpoint_turn_items = history_at_checkpoint
|
||||
.strip_prefix(history_before_turn)
|
||||
.unwrap_or(history_at_checkpoint);
|
||||
let stripped_compaction_initial_context =
|
||||
checkpoint_turn_items.strip_prefix(compaction_initial_context);
|
||||
let stripped_turn_start_context_items =
|
||||
checkpoint_turn_items.strip_prefix(turn_start_context_items);
|
||||
let checkpoint_turn_items = match (
|
||||
stripped_compaction_initial_context,
|
||||
stripped_turn_start_context_items,
|
||||
) {
|
||||
(Some(after_compaction_initial_context), Some(after_turn_start_context_items)) => {
|
||||
if compaction_initial_context.len() >= turn_start_context_items.len() {
|
||||
after_compaction_initial_context
|
||||
} else {
|
||||
after_turn_start_context_items
|
||||
}
|
||||
}
|
||||
(Some(after_compaction_initial_context), None) => after_compaction_initial_context,
|
||||
(None, Some(after_turn_start_context_items)) => after_turn_start_context_items,
|
||||
(None, None) => checkpoint_turn_items,
|
||||
};
|
||||
let post_checkpoint_turn_items = current_history
|
||||
.strip_prefix(history_at_checkpoint)
|
||||
.unwrap_or_default();
|
||||
let mut replacement_history: Vec<ResponseItem> = checkpoint_turn_items
|
||||
.iter()
|
||||
.filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. }))
|
||||
.filter(|item| !matches!(item, ResponseItem::Compaction { .. }))
|
||||
.cloned()
|
||||
.collect();
|
||||
replacement_history.push(compaction_item);
|
||||
let mut replacement_history = vec![compaction_item];
|
||||
replacement_history.extend(
|
||||
post_checkpoint_turn_items
|
||||
.iter()
|
||||
.filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. }))
|
||||
.filter(|item| !matches!(item, ResponseItem::Compaction { .. }))
|
||||
.cloned(),
|
||||
);
|
||||
let mut replacement_history = insert_initial_context_before_last_real_user_or_summary(
|
||||
replacement_history,
|
||||
compaction_initial_context.to_vec(),
|
||||
);
|
||||
replacement_history.extend(
|
||||
current_history
|
||||
history_at_compaction
|
||||
.iter()
|
||||
.filter(|item| matches!(item, ResponseItem::GhostSnapshot { .. }))
|
||||
.cloned(),
|
||||
@@ -6087,14 +5999,17 @@ fn record_inline_compaction_skip(
|
||||
async fn run_pre_sampling_compact(
|
||||
sess: &Arc<Session>,
|
||||
turn_context: &Arc<TurnContext>,
|
||||
) -> CodexResult<Option<PendingServerSideCompaction>> {
|
||||
) -> CodexResult<Option<AutoCompactTrigger>> {
|
||||
let total_usage_tokens_before_compaction = sess.get_total_token_usage().await;
|
||||
maybe_run_previous_model_inline_compact(
|
||||
if maybe_run_previous_model_inline_compact(
|
||||
sess,
|
||||
turn_context,
|
||||
total_usage_tokens_before_compaction,
|
||||
)
|
||||
.await?;
|
||||
.await?
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let total_usage_tokens = sess.get_total_token_usage().await;
|
||||
let auto_compact_limit = turn_context
|
||||
.model_info
|
||||
@@ -6102,7 +6017,7 @@ async fn run_pre_sampling_compact(
|
||||
.unwrap_or(i64::MAX);
|
||||
// Compact if the total usage tokens are greater than the auto compact limit
|
||||
if total_usage_tokens >= auto_compact_limit {
|
||||
if let Some(threshold) = inline_server_side_compaction_threshold(sess, turn_context) {
|
||||
if inline_server_side_compaction_threshold(sess, turn_context).is_some() {
|
||||
record_compaction_metric(
|
||||
sess,
|
||||
"server_side",
|
||||
@@ -6110,10 +6025,7 @@ async fn run_pre_sampling_compact(
|
||||
"requested",
|
||||
&[],
|
||||
);
|
||||
return Ok(Some(PendingServerSideCompaction {
|
||||
threshold,
|
||||
trigger: AutoCompactTrigger::AutoPreTurn,
|
||||
}));
|
||||
return Ok(Some(AutoCompactTrigger::AutoPreTurn));
|
||||
}
|
||||
record_inline_compaction_skip(sess, turn_context, AutoCompactTrigger::AutoPreTurn);
|
||||
run_auto_compact(
|
||||
@@ -6138,13 +6050,6 @@ async fn maybe_run_previous_model_inline_compact(
|
||||
turn_context: &Arc<TurnContext>,
|
||||
total_usage_tokens: i64,
|
||||
) -> CodexResult<bool> {
|
||||
// Keep OpenAI auto-compaction on one path. If inline server-side compaction is eligible for
|
||||
// the current turn, let the normal pre-turn inline request handle it instead of running the
|
||||
// older previous-model client-side preflight flow first.
|
||||
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);
|
||||
};
|
||||
@@ -6401,9 +6306,6 @@ async fn run_sampling_request(
|
||||
client_session: &mut ModelClientSession,
|
||||
turn_metadata_header: Option<&str>,
|
||||
input: Vec<ResponseItem>,
|
||||
compaction_initial_context: &[ResponseItem],
|
||||
turn_start_context_items: &[ResponseItem],
|
||||
history_before_turn: &[ResponseItem],
|
||||
inline_compaction_threshold: Option<i64>,
|
||||
explicitly_enabled_connectors: &HashSet<String>,
|
||||
skills_outcome: Option<&SkillLoadOutcome>,
|
||||
@@ -6440,9 +6342,6 @@ async fn run_sampling_request(
|
||||
Arc::clone(&turn_diff_tracker),
|
||||
server_model_warning_emitted_for_turn,
|
||||
&prompt,
|
||||
compaction_initial_context,
|
||||
turn_start_context_items,
|
||||
history_before_turn,
|
||||
cancellation_token.child_token(),
|
||||
)
|
||||
.await
|
||||
@@ -7166,9 +7065,6 @@ async fn try_run_sampling_request(
|
||||
turn_diff_tracker: SharedTurnDiffTracker,
|
||||
server_model_warning_emitted_for_turn: &mut bool,
|
||||
prompt: &Prompt,
|
||||
compaction_initial_context: &[ResponseItem],
|
||||
turn_start_context_items: &[ResponseItem],
|
||||
history_before_turn: &[ResponseItem],
|
||||
cancellation_token: CancellationToken,
|
||||
) -> CodexResult<SamplingRequestResult> {
|
||||
feedback_tags!(
|
||||
@@ -7204,15 +7100,11 @@ 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;
|
||||
let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode);
|
||||
let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id));
|
||||
let history_before_turn = Arc::new(history_before_turn.to_vec());
|
||||
let receiving_span = trace_span!("receiving_stream");
|
||||
let outcome: CodexResult<SamplingRequestResult> = loop {
|
||||
let handle_responses = trace_span!(
|
||||
@@ -7291,8 +7183,14 @@ async fn try_run_sampling_request(
|
||||
.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(ServerSideCompaction { item, turn_item }) =
|
||||
output_result.server_side_compaction
|
||||
{
|
||||
sess.apply_server_side_compaction(turn_context.as_ref(), item)
|
||||
.await;
|
||||
sess.emit_turn_item_started(&turn_context, &turn_item).await;
|
||||
sess.emit_turn_item_completed(&turn_context, turn_item)
|
||||
.await;
|
||||
}
|
||||
if let Some(tool_future) = output_result.tool_future {
|
||||
in_flight.push_back(tool_future);
|
||||
@@ -7396,39 +7294,8 @@ async fn try_run_sampling_request(
|
||||
&mut assistant_message_stream_parsers,
|
||||
)
|
||||
.await;
|
||||
let mut applied_server_side_compaction_checkpoint = false;
|
||||
if let Some(PendingServerSideCompactionCheckpoint {
|
||||
history_at_checkpoint,
|
||||
item,
|
||||
turn_item,
|
||||
}) = pending_server_side_compaction_checkpoint.take()
|
||||
{
|
||||
sess.apply_server_side_compaction_checkpoint(
|
||||
turn_context.as_ref(),
|
||||
item,
|
||||
compaction_initial_context,
|
||||
turn_start_context_items,
|
||||
history_before_turn.as_slice(),
|
||||
history_at_checkpoint.as_slice(),
|
||||
)
|
||||
sess.update_token_usage_info(&turn_context, token_usage.as_ref())
|
||||
.await;
|
||||
sess.emit_turn_item_started(&turn_context, &turn_item).await;
|
||||
sess.emit_turn_item_completed(&turn_context, turn_item)
|
||||
.await;
|
||||
if let Some(token_usage) = token_usage.as_ref() {
|
||||
let mut state = sess.state.lock().await;
|
||||
state.update_token_info_from_usage(
|
||||
token_usage,
|
||||
turn_context.model_context_window(),
|
||||
);
|
||||
}
|
||||
sess.recompute_token_usage(&turn_context).await;
|
||||
applied_server_side_compaction_checkpoint = true;
|
||||
}
|
||||
if !applied_server_side_compaction_checkpoint {
|
||||
sess.update_token_usage_info(&turn_context, token_usage.as_ref())
|
||||
.await;
|
||||
}
|
||||
should_emit_turn_diff = true;
|
||||
|
||||
needs_follow_up |= sess.has_pending_input().await;
|
||||
|
||||
@@ -285,114 +285,38 @@ fn assistant_message_stream_parsers_seed_plan_parser_across_added_and_delta_boun
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() {
|
||||
fn build_server_side_compaction_replacement_history_replaces_prompt_history() {
|
||||
let prior_snapshot = ghost_snapshot("ghost-before");
|
||||
let same_turn_snapshot = ghost_snapshot("ghost-during");
|
||||
let history_before_turn = vec![user_message("earlier"), prior_snapshot.clone()];
|
||||
let turn_start_context_items = vec![
|
||||
developer_message("<model_switch>\nuse the new model"),
|
||||
environment_context_message("/fresh"),
|
||||
];
|
||||
let current_turn_user = user_message("current turn");
|
||||
let skill_injection = skill_message(
|
||||
"<skill>\n<name>demo</name>\n<path>/tmp/skills/demo/SKILL.md</path>\nbody\n</skill>",
|
||||
);
|
||||
let plugin_injection = developer_message("PLUGIN_HINT");
|
||||
let current_turn_tool_output = ResponseItem::FunctionCallOutput {
|
||||
call_id: "call-1".to_string(),
|
||||
output: FunctionCallOutputPayload::from_text("tool result".to_string()),
|
||||
};
|
||||
let current_history = vec![
|
||||
user_message("earlier"),
|
||||
prior_snapshot.clone(),
|
||||
turn_start_context_items[0].clone(),
|
||||
turn_start_context_items[1].clone(),
|
||||
current_turn_user.clone(),
|
||||
skill_injection.clone(),
|
||||
plugin_injection.clone(),
|
||||
current_turn_tool_output.clone(),
|
||||
developer_message("<model_switch>\nuse the new model"),
|
||||
environment_context_message("/fresh"),
|
||||
user_message("current turn"),
|
||||
skill_message(
|
||||
"<skill>\n<name>demo</name>\n<path>/tmp/skills/demo/SKILL.md</path>\nbody\n</skill>",
|
||||
),
|
||||
developer_message("PLUGIN_HINT"),
|
||||
ResponseItem::FunctionCallOutput {
|
||||
call_id: "call-1".to_string(),
|
||||
output: FunctionCallOutputPayload::from_text("tool result".to_string()),
|
||||
},
|
||||
same_turn_snapshot.clone(),
|
||||
ResponseItem::Compaction {
|
||||
encrypted_content: "INLINE_SUMMARY_1".to_string(),
|
||||
},
|
||||
];
|
||||
let compaction_item = ResponseItem::Compaction {
|
||||
encrypted_content: "INLINE_SUMMARY".to_string(),
|
||||
};
|
||||
|
||||
let replacement_history = build_server_side_compaction_replacement_history(
|
||||
compaction_item.clone(),
|
||||
&turn_start_context_items,
|
||||
&turn_start_context_items,
|
||||
&history_before_turn,
|
||||
¤t_history,
|
||||
¤t_history,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
replacement_history,
|
||||
vec![
|
||||
turn_start_context_items[0].clone(),
|
||||
turn_start_context_items[1].clone(),
|
||||
current_turn_user,
|
||||
skill_injection,
|
||||
plugin_injection,
|
||||
current_turn_tool_output,
|
||||
compaction_item,
|
||||
prior_snapshot,
|
||||
same_turn_snapshot,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_server_side_compaction_replacement_history_prefers_longer_initial_context_prefix() {
|
||||
let history_before_turn = vec![user_message("earlier")];
|
||||
let turn_start_context_items = vec![developer_message("<model_switch>\nuse the new model")];
|
||||
let compaction_initial_context = vec![
|
||||
turn_start_context_items[0].clone(),
|
||||
environment_context_message("/fresh"),
|
||||
];
|
||||
let current_turn_user = user_message("current turn");
|
||||
let prior_compaction = ResponseItem::Compaction {
|
||||
encrypted_content: "INLINE_SUMMARY_1".to_string(),
|
||||
};
|
||||
let new_compaction = ResponseItem::Compaction {
|
||||
encrypted_content: "INLINE_SUMMARY_2".to_string(),
|
||||
};
|
||||
let current_turn_tool_output = ResponseItem::FunctionCallOutput {
|
||||
call_id: "call-1".to_string(),
|
||||
output: FunctionCallOutputPayload::from_text("tool result".to_string()),
|
||||
};
|
||||
let history_at_checkpoint = vec![
|
||||
compaction_initial_context[0].clone(),
|
||||
compaction_initial_context[1].clone(),
|
||||
current_turn_user.clone(),
|
||||
prior_compaction,
|
||||
];
|
||||
let current_history = vec![
|
||||
history_at_checkpoint[0].clone(),
|
||||
history_at_checkpoint[1].clone(),
|
||||
history_at_checkpoint[2].clone(),
|
||||
history_at_checkpoint[3].clone(),
|
||||
current_turn_tool_output.clone(),
|
||||
];
|
||||
|
||||
let replacement_history = build_server_side_compaction_replacement_history(
|
||||
new_compaction.clone(),
|
||||
&compaction_initial_context,
|
||||
&turn_start_context_items,
|
||||
&history_before_turn,
|
||||
&history_at_checkpoint,
|
||||
¤t_history,
|
||||
);
|
||||
let replacement_history =
|
||||
build_server_side_compaction_replacement_history(compaction_item.clone(), ¤t_history);
|
||||
|
||||
assert_eq!(
|
||||
replacement_history,
|
||||
vec![
|
||||
compaction_initial_context[0].clone(),
|
||||
compaction_initial_context[1].clone(),
|
||||
current_turn_user,
|
||||
new_compaction,
|
||||
current_turn_tool_output,
|
||||
]
|
||||
vec![compaction_item, prior_snapshot, same_turn_snapshot]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -3920,7 +3844,7 @@ async fn task_finish_emits_turn_item_lifecycle_for_leftover_pending_user_input()
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn inline_compaction_output_item_buffers_checkpoint_without_committing_turn_item() {
|
||||
async fn inline_compaction_output_item_returns_immediate_compaction_apply_signal() {
|
||||
let (sess, tc, rx) = make_session_and_context_with_rx().await;
|
||||
while rx.try_recv().is_ok() {}
|
||||
|
||||
@@ -3956,7 +3880,7 @@ async fn inline_compaction_output_item_buffers_checkpoint_without_committing_tur
|
||||
.await
|
||||
.expect("handle output item");
|
||||
|
||||
assert!(output.pending_server_side_compaction.is_some());
|
||||
assert!(output.server_side_compaction.is_some());
|
||||
assert!(output.tool_future.is_none());
|
||||
assert!(!output.needs_follow_up);
|
||||
assert!(output.last_agent_message.is_none());
|
||||
@@ -3971,7 +3895,7 @@ async fn inline_compaction_output_item_buffers_checkpoint_without_committing_tur
|
||||
));
|
||||
assert!(
|
||||
rx.try_recv().is_err(),
|
||||
"expected no committed compaction item lifecycle before response.completed"
|
||||
"expected the caller to decide when to commit the compaction lifecycle"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -147,15 +147,11 @@ 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 server_side_compaction: Option<ServerSideCompaction>,
|
||||
pub tool_future: Option<InFlightFuture<'static>>,
|
||||
}
|
||||
|
||||
pub(crate) struct PendingServerSideCompactionCheckpoint {
|
||||
// Snapshot the raw history at the moment the compaction item streamed. We cannot build
|
||||
// replacement history yet because later same-turn output may still arrive before
|
||||
// `response.completed`.
|
||||
pub history_at_checkpoint: Vec<ResponseItem>,
|
||||
pub(crate) struct ServerSideCompaction {
|
||||
pub item: ResponseItem,
|
||||
pub turn_item: TurnItem,
|
||||
}
|
||||
@@ -181,12 +177,11 @@ pub(crate) async fn handle_output_item_done(
|
||||
Some(TurnItem::ContextCompaction(item)) => item,
|
||||
_ => ContextCompactionItem::new(),
|
||||
});
|
||||
// Preserve the raw wire event immediately, but defer the committed turn-item lifecycle
|
||||
// until `response.completed` so later streamed output from the same response is not
|
||||
// reordered around the local checkpoint rewrite.
|
||||
// Preserve the raw wire event immediately; the caller rewrites local history inline and
|
||||
// then continues appending later same-turn output after the checkpoint item.
|
||||
debug!(
|
||||
turn_id = %ctx.turn_context.sub_id,
|
||||
"emitting streamed server-side raw compaction item and buffering committed checkpoint until response.completed"
|
||||
"emitting streamed server-side raw compaction item for immediate local checkpoint apply"
|
||||
);
|
||||
ctx.sess
|
||||
.send_event(
|
||||
@@ -194,13 +189,7 @@ pub(crate) async fn handle_output_item_done(
|
||||
EventMsg::RawResponseItem(RawResponseItemEvent { item: item.clone() }),
|
||||
)
|
||||
.await;
|
||||
// Replacement history is derived on `response.completed` from this snapshot plus whatever
|
||||
// else the model streamed afterward.
|
||||
output.pending_server_side_compaction = Some(PendingServerSideCompactionCheckpoint {
|
||||
history_at_checkpoint: ctx.sess.clone_history().await.raw_items().to_vec(),
|
||||
item,
|
||||
turn_item,
|
||||
});
|
||||
output.server_side_compaction = Some(ServerSideCompaction { item, turn_item });
|
||||
return Ok(output);
|
||||
}
|
||||
|
||||
|
||||
@@ -9,9 +9,13 @@ use codex_core::ModelProviderInfo;
|
||||
use codex_core::WireApi;
|
||||
use codex_core::compact::SUMMARY_PREFIX;
|
||||
use codex_core::features::Feature;
|
||||
use codex_core::models_manager::manager::RefreshStrategy;
|
||||
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::AskForApproval;
|
||||
use codex_protocol::protocol::ConversationStartParams;
|
||||
use codex_protocol::protocol::ErrorEvent;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
@@ -22,11 +26,13 @@ use codex_protocol::protocol::RealtimeConversationRealtimeEvent;
|
||||
use codex_protocol::protocol::RealtimeEvent;
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
use codex_protocol::protocol::RolloutLine;
|
||||
use codex_protocol::protocol::SandboxPolicy;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use core_test_support::context_snapshot;
|
||||
use core_test_support::context_snapshot::ContextSnapshotOptions;
|
||||
use core_test_support::context_snapshot::ContextSnapshotRenderMode;
|
||||
use core_test_support::responses;
|
||||
use core_test_support::responses::mount_models_once;
|
||||
use core_test_support::responses::mount_sse_once;
|
||||
use core_test_support::responses::sse;
|
||||
use core_test_support::responses::start_websocket_server;
|
||||
@@ -62,6 +68,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 })
|
||||
@@ -488,6 +506,13 @@ async fn auto_server_side_compaction_uses_inline_context_management() -> Result<
|
||||
"expected three inline /responses requests"
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
requests[0].body_json().get("context_management"),
|
||||
Some(&json!([{
|
||||
"type": "compaction",
|
||||
"compact_threshold": compact_threshold,
|
||||
}])),
|
||||
);
|
||||
let inline_request = requests[1].body_json();
|
||||
assert_eq!(
|
||||
inline_request.get("context_management"),
|
||||
@@ -498,14 +523,21 @@ async fn auto_server_side_compaction_uses_inline_context_management() -> Result<
|
||||
);
|
||||
|
||||
let post_inline_request = &requests[2];
|
||||
assert_eq!(
|
||||
post_inline_request.body_json().get("context_management"),
|
||||
Some(&json!([{
|
||||
"type": "compaction",
|
||||
"compact_threshold": compact_threshold,
|
||||
}])),
|
||||
);
|
||||
let post_inline_body = post_inline_request.body_json().to_string();
|
||||
assert!(
|
||||
post_inline_request.body_contains_text(&inline_summary),
|
||||
"expected subsequent request to reuse inline compaction item"
|
||||
);
|
||||
assert!(
|
||||
post_inline_request.body_contains_text("<permissions instructions>"),
|
||||
"expected subsequent request to preserve canonical context after inline compaction"
|
||||
!post_inline_request.body_contains_text("<permissions instructions>"),
|
||||
"expected subsequent request to rely on the compaction item rather than duplicate canonical context"
|
||||
);
|
||||
assert!(
|
||||
!post_inline_request.body_contains_text(first_turn_text),
|
||||
@@ -533,7 +565,85 @@ async fn auto_server_side_compaction_uses_inline_context_management() -> Result<
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn auto_server_side_compaction_keeps_current_turn_inputs_for_follow_ups() -> Result<()> {
|
||||
async fn snapshot_request_shape_auto_server_side_compaction_history_layout() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
let compact_threshold = 120;
|
||||
let inline_summary = summary_with_prefix("INLINE_SERVER_SUMMARY");
|
||||
|
||||
let harness = TestCodexHarness::with_builder(
|
||||
test_codex()
|
||||
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
|
||||
.with_config(move |config| {
|
||||
config
|
||||
.features
|
||||
.enable(Feature::ServerSideCompaction)
|
||||
.expect("enable server-side compaction");
|
||||
config.model_auto_compact_token_limit = Some(compact_threshold);
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
let codex = harness.test().codex.clone();
|
||||
|
||||
let responses_mock = responses::mount_response_sequence(
|
||||
harness.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::ev_assistant_message("m2", "AFTER_INLINE_REPLY"),
|
||||
responses::ev_completed_with_tokens("resp-2", 80),
|
||||
])),
|
||||
responses::sse_response(sse(vec![responses::ev_completed("resp-3")])),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
submit_text_turn_and_wait(&codex, "inline compact turn one").await?;
|
||||
|
||||
codex
|
||||
.submit(Op::UserInput {
|
||||
items: vec![UserInput::Text {
|
||||
text: "inline compact turn two".to_string(),
|
||||
text_elements: Vec::new(),
|
||||
}],
|
||||
final_output_json_schema: None,
|
||||
})
|
||||
.await?;
|
||||
wait_for_event(&codex, |event| {
|
||||
matches!(event, EventMsg::ContextCompacted(_))
|
||||
})
|
||||
.await;
|
||||
wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await;
|
||||
|
||||
submit_text_turn_and_wait(&codex, "inline compact turn three").await?;
|
||||
|
||||
let requests = responses_mock.requests();
|
||||
assert_eq!(
|
||||
requests.len(),
|
||||
3,
|
||||
"expected three inline /responses requests"
|
||||
);
|
||||
|
||||
insta::assert_snapshot!(
|
||||
"auto_server_side_compaction_history_layout_shapes",
|
||||
format_labeled_requests_snapshot(
|
||||
"Inline server-side compaction: the compacting turn opts into context_management, and the next request reuses the streamed compaction checkpoint directly.",
|
||||
&[
|
||||
("Inline Compaction Request", &requests[1]),
|
||||
("Post-Inline-Compaction History Layout", &requests[2]),
|
||||
]
|
||||
)
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn auto_server_side_compaction_follow_ups_use_compaction_checkpoint() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
let compact_threshold = 120;
|
||||
@@ -592,26 +702,17 @@ async fn auto_server_side_compaction_keeps_current_turn_inputs_for_follow_ups()
|
||||
"expected same-turn follow-up to include the inline compaction item"
|
||||
);
|
||||
assert!(
|
||||
follow_up_request.body_contains_text("<permissions instructions>"),
|
||||
"expected same-turn follow-up to preserve canonical context after inline compaction"
|
||||
!follow_up_request.body_contains_text("<permissions instructions>"),
|
||||
"expected same-turn follow-up to rely on the compaction item rather than duplicate canonical context"
|
||||
);
|
||||
assert!(
|
||||
follow_up_request.body_contains_text(second_turn_text),
|
||||
"expected same-turn follow-up to retain the current turn user input"
|
||||
!follow_up_request.body_contains_text(second_turn_text),
|
||||
"expected same-turn follow-up to replace prior plaintext history with the compaction item"
|
||||
);
|
||||
assert!(
|
||||
!follow_up_request.body_contains_text(first_turn_text),
|
||||
"expected same-turn follow-up to drop pre-compaction history"
|
||||
);
|
||||
assert!(
|
||||
follow_up_body
|
||||
.find(second_turn_text)
|
||||
.expect("current turn text in follow-up request")
|
||||
< follow_up_body
|
||||
.find("INLINE_SERVER_SUMMARY")
|
||||
.expect("inline compaction marker in follow-up request"),
|
||||
"expected current-turn items to remain ahead of the inline compaction item"
|
||||
);
|
||||
assert!(
|
||||
follow_up_request
|
||||
.function_call_output_text("call-inline-mid-turn")
|
||||
@@ -2516,6 +2617,171 @@ async fn snapshot_request_shape_remote_pre_turn_compaction_strips_incoming_model
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn previous_model_preflight_remote_compaction_still_runs_with_inline_feature_enabled()
|
||||
-> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
let compact_threshold = 200;
|
||||
let previous_model = "gpt-5.2-codex";
|
||||
let next_model = "gpt-5.1-codex-max";
|
||||
let harness = TestCodexHarness::with_builder(
|
||||
test_codex()
|
||||
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
|
||||
.with_model(previous_model)
|
||||
.with_config(move |config| {
|
||||
config.model_auto_compact_token_limit = Some(compact_threshold);
|
||||
config
|
||||
.features
|
||||
.enable(Feature::ServerSideCompaction)
|
||||
.expect("enable server-side compaction");
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
let codex = harness.test().codex.clone();
|
||||
let models_mock = mount_models_once(
|
||||
harness.server(),
|
||||
ModelsResponse {
|
||||
models: vec![
|
||||
model_info_with_context_window(previous_model, 273_000),
|
||||
model_info_with_context_window(next_model, 125_000),
|
||||
],
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let _ = harness
|
||||
.test()
|
||||
.thread_manager
|
||||
.get_models_manager()
|
||||
.list_models(RefreshStrategy::Online)
|
||||
.await;
|
||||
|
||||
let initial_turn_request_mock = responses::mount_sse_once(
|
||||
harness.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(
|
||||
harness.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(
|
||||
harness.server(),
|
||||
&summary_with_prefix("REMOTE_SWITCH_SUMMARY"),
|
||||
)
|
||||
.await;
|
||||
|
||||
codex
|
||||
.submit(Op::UserTurn {
|
||||
items: vec![UserInput::Text {
|
||||
text: "BEFORE_SWITCH_USER".to_string(),
|
||||
text_elements: Vec::new(),
|
||||
}],
|
||||
final_output_json_schema: None,
|
||||
cwd: harness.test().cwd.path().to_path_buf(),
|
||||
approval_policy: AskForApproval::Never,
|
||||
sandbox_policy: SandboxPolicy::DangerFullAccess,
|
||||
model: previous_model.to_string(),
|
||||
effort: None,
|
||||
summary: None,
|
||||
service_tier: None,
|
||||
collaboration_mode: None,
|
||||
personality: None,
|
||||
})
|
||||
.await?;
|
||||
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
|
||||
|
||||
codex
|
||||
.submit(Op::UserTurn {
|
||||
items: vec![UserInput::Text {
|
||||
text: "AFTER_SWITCH_USER".to_string(),
|
||||
text_elements: Vec::new(),
|
||||
}],
|
||||
final_output_json_schema: None,
|
||||
cwd: harness.test().cwd.path().to_path_buf(),
|
||||
approval_policy: AskForApproval::Never,
|
||||
sandbox_policy: SandboxPolicy::DangerFullAccess,
|
||||
model: next_model.to_string(),
|
||||
effort: None,
|
||||
summary: None,
|
||||
service_tier: None,
|
||||
collaboration_mode: None,
|
||||
personality: None,
|
||||
})
|
||||
.await?;
|
||||
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
|
||||
|
||||
assert_eq!(
|
||||
compact_mock.requests().len(),
|
||||
1,
|
||||
"expected previous-model preflight to keep using remote compaction even when inline server-side compaction is enabled"
|
||||
);
|
||||
assert_eq!(models_mock.requests().len(), 1, "expected one models fetch");
|
||||
assert_eq!(
|
||||
initial_turn_request_mock.requests().len(),
|
||||
1,
|
||||
"expected initial turn request"
|
||||
);
|
||||
assert_eq!(
|
||||
post_compact_turn_request_mock.requests().len(),
|
||||
1,
|
||||
"expected post-compaction follow-up request"
|
||||
);
|
||||
let initial_turn_request = initial_turn_request_mock.single_request();
|
||||
let compact_request = compact_mock.single_request();
|
||||
let post_compact_turn_request = post_compact_turn_request_mock.single_request();
|
||||
assert_eq!(
|
||||
initial_turn_request
|
||||
.body_json()
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some(previous_model)
|
||||
);
|
||||
assert_eq!(
|
||||
compact_request
|
||||
.body_json()
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some(previous_model)
|
||||
);
|
||||
assert_eq!(
|
||||
post_compact_turn_request
|
||||
.body_json()
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some(next_model)
|
||||
);
|
||||
let compact_body = compact_request.body_json().to_string();
|
||||
assert!(
|
||||
!compact_body.contains("<model_switch>"),
|
||||
"expected previous-model preflight compact request to strip model-switch context"
|
||||
);
|
||||
let follow_up_body = post_compact_turn_request.body_json().to_string();
|
||||
assert!(
|
||||
follow_up_body.contains("<model_switch>"),
|
||||
"expected post-compaction next-model request to restore model-switch context"
|
||||
);
|
||||
assert_eq!(
|
||||
post_compact_turn_request
|
||||
.body_json()
|
||||
.get("context_management"),
|
||||
Some(&json!([{
|
||||
"type": "compaction",
|
||||
"compact_threshold": compact_threshold,
|
||||
}])),
|
||||
"expected eligible next-model requests to always include inline context management after preflight compaction"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
// TODO(ccunningham): Update once remote pre-turn compaction context-overflow handling includes
|
||||
// incoming user input and emits richer oversized-input messaging.
|
||||
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
---
|
||||
source: core/tests/suite/compact_remote.rs
|
||||
assertion_line: 558
|
||||
expression: "format_labeled_requests_snapshot(\"Inline server-side compaction: the compacting turn opts into context_management, and the next request reuses the streamed compaction checkpoint directly.\",\n&[(\"Inline Compaction Request\", &requests[1]),\n(\"Post-Inline-Compaction History Layout\", &requests[2]),])"
|
||||
---
|
||||
Scenario: Inline server-side compaction: the compacting turn opts into context_management, and the next request reuses the streamed compaction checkpoint directly.
|
||||
|
||||
## Inline Compaction Request
|
||||
00:message/developer:<PERMISSIONS_INSTRUCTIONS>
|
||||
01:message/user[2]:
|
||||
[01] <AGENTS_MD>
|
||||
[02] <ENVIRONMENT_CONTEXT:cwd=<CWD>>
|
||||
02:message/user:inline compact turn one
|
||||
03:message/assistant:FIRST_REMOTE_REPLY
|
||||
04:message/user:inline compact turn two
|
||||
|
||||
## Post-Inline-Compaction History Layout
|
||||
00:compaction:encrypted=true
|
||||
01:message/assistant:AFTER_INLINE_REPLY
|
||||
02:message/user:inline compact turn three
|
||||
Reference in New Issue
Block a user