core: skip fallback after explicit context reset

This commit is contained in:
Bo Xie
2026-07-08 17:20:05 -07:00
parent 248115bed4
commit ba2bfb205b
3 changed files with 63 additions and 10 deletions

View File

@@ -19,6 +19,11 @@ use codex_protocol::items::TurnItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::TurnStartedEvent;
pub(crate) enum AutoCompactFallbackPolicy {
Run,
Skip,
}
/// Runs token-budget manual compaction as a normal compaction lifecycle.
///
/// Token-budget compaction skips model/server summarization and installs a fresh context window
@@ -61,6 +66,7 @@ pub(crate) async fn run_inline_auto_compact_task(
step_context: Arc<StepContext>,
client_session: &mut ModelClientSession,
initial_context_injection: InitialContextInjection,
fallback_policy: AutoCompactFallbackPolicy,
) -> CodexResult<()> {
let turn_context = &step_context.turn;
let world_state = match initial_context_injection {
@@ -69,13 +75,17 @@ pub(crate) async fn run_inline_auto_compact_task(
Arc::new(sess.build_world_state_for_step(&step_context).await)
}
};
let (auto_step_context, fallback_client_session) = match fallback_policy {
AutoCompactFallbackPolicy::Run => (Some(&step_context), Some(client_session)),
AutoCompactFallbackPolicy::Skip => (None, None),
};
run_compact_task_inner(
&sess,
turn_context,
world_state,
CompactionTrigger::Auto,
Some(&step_context),
Some(client_session),
auto_step_context,
fallback_client_session,
)
.await
}

View File

@@ -15,6 +15,7 @@ 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;
use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_remote_auto_compact_task_v2;
use crate::compact_token_budget::AutoCompactFallbackPolicy;
use crate::connectors;
use crate::context::ContextualUserFragment;
use crate::feedback_tags;
@@ -344,15 +345,24 @@ pub(crate) async fn run_turn(
.await;
// 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 needs_follow_up
&& (sess.take_new_context_window_request().await || token_limit_reached)
{
let new_context_window_requested = if needs_follow_up {
sess.take_new_context_window_request().await
} else {
false
};
if needs_follow_up && (new_context_window_requested || token_limit_reached) {
let fallback_policy = if new_context_window_requested {
AutoCompactFallbackPolicy::Skip
} else {
AutoCompactFallbackPolicy::Run
};
if let Err(err) = run_auto_compact(
&sess,
Arc::clone(&step_context),
/*fallback_step_context*/ None,
&mut client_session,
InitialContextInjection::BeforeLastUserMessage(Arc::clone(&world_state)),
fallback_policy,
CompactionReason::ContextLimit,
CompactionPhase::MidTurn,
)
@@ -815,6 +825,7 @@ async fn run_pre_sampling_compact(
/*fallback_step_context*/ None,
client_session,
InitialContextInjection::DoNotInject,
AutoCompactFallbackPolicy::Run,
CompactionReason::ContextLimit,
CompactionPhase::PreTurn,
)
@@ -892,6 +903,7 @@ async fn maybe_run_previous_model_inline_compact(
fallback_step_context,
client_session,
InitialContextInjection::DoNotInject,
AutoCompactFallbackPolicy::Run,
CompactionReason::CompHashChanged,
CompactionPhase::PreTurn,
)
@@ -939,6 +951,7 @@ async fn maybe_run_previous_model_inline_compact(
fallback_step_context,
client_session,
InitialContextInjection::DoNotInject,
AutoCompactFallbackPolicy::Run,
CompactionReason::ModelDownshift,
CompactionPhase::PreTurn,
)
@@ -958,6 +971,7 @@ async fn run_auto_compact(
fallback_step_context: Option<Arc<StepContext>>,
client_session: &mut ModelClientSession,
initial_context_injection: InitialContextInjection,
fallback_policy: AutoCompactFallbackPolicy,
reason: CompactionReason,
phase: CompactionPhase,
) -> CodexResult<()> {
@@ -970,6 +984,7 @@ async fn run_auto_compact(
step_context,
client_session,
initial_context_injection,
fallback_policy,
)
.await?;
return Ok(());

View File

@@ -1079,7 +1079,7 @@ async fn token_budget_auto_compact_fallback_runs_in_old_window_before_reset() ->
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
async fn new_context_tool_skips_fallback_and_starts_new_window_before_follow_up() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
@@ -1097,7 +1097,9 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
sse(vec![
ev_response_created("resp-1"),
ev_function_call(call_id, "new_context", "{}"),
ev_completed("resp-1"),
// Explicit rollover takes precedence even when this response also crosses the
// automatic compaction threshold.
ev_completed_with_tokens("resp-1", /*total_tokens*/ 9_500),
]),
sse(vec![
ev_response_created("resp-2"),
@@ -1112,13 +1114,19 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
],
)
.await;
let (extensions, fallback_tool_calls) = token_budget_auto_compact_fallback_extensions();
let test = test_codex()
.with_extensions(extensions)
.with_config(|config| {
config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
config.model_context_window = Some(10_000);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
config
.features
.enable(Feature::AutoCompactFallback)
.expect("auto compact fallback feature should be known");
})
.build(&server)
.await?;
@@ -1133,12 +1141,32 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
.any(|name| name == "new_context"),
"new_context should be exposed when token budget is enabled"
);
assert_eq!(
fallback_tool_calls.load(Ordering::SeqCst),
0,
"explicit new_context should skip fallback tool execution"
);
assert!(
!requests[1].body_contains_text(TOKEN_BUDGET_FALLBACK_PROMPT),
"the first new-window request should not include the fallback prompt"
);
let first_new_window_metadata: Value = serde_json::from_str(
&requests[1]
.header("x-codex-turn-metadata")
.expect("the first new-window request should include turn metadata"),
)
.expect("turn metadata should be valid json");
assert_eq!(
first_new_window_metadata["request_kind"].as_str(),
Some("turn"),
"explicit new_context should continue directly with a normal turn"
);
let thread_id = test.session_configured.thread_id;
let initial_token_budget = token_budget_contexts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
let (initial_first_window_id, _, initial_window_id) =
token_budget_window_ids(&initial_token_budget[0], thread_id);
let new_window_token_budget = token_budget_contexts(&requests[2]);
let new_window_token_budget = token_budget_contexts(&requests[1]);
assert_eq!(new_window_token_budget.len(), 1);
let (new_first_window_id, new_previous_window_id, new_window_id) =
token_budget_window_ids(&new_window_token_budget[0], thread_id);
@@ -1149,7 +1177,7 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
);
assert_ne!(new_window_id, initial_window_id);
assert!(
!requests[2].body_contains_text("request new context window"),
!requests[1].body_contains_text("request new context window"),
"new_context should drop the prior window history before continuing the turn"
);
assert_eq!(