From ba2bfb205bfeae0a8be4043b9c9ee1608ff91dbd Mon Sep 17 00:00:00 2001 From: Bo Xie Date: Wed, 8 Jul 2026 17:20:05 -0700 Subject: [PATCH] core: skip fallback after explicit context reset --- codex-rs/core/src/compact_token_budget.rs | 14 +++++++-- codex-rs/core/src/session/turn.rs | 21 +++++++++++-- codex-rs/core/tests/suite/token_budget.rs | 38 ++++++++++++++++++++--- 3 files changed, 63 insertions(+), 10 deletions(-) diff --git a/codex-rs/core/src/compact_token_budget.rs b/codex-rs/core/src/compact_token_budget.rs index f3bd3294d1..5527cd806b 100644 --- a/codex-rs/core/src/compact_token_budget.rs +++ b/codex-rs/core/src/compact_token_budget.rs @@ -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, 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 } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8c67167f6d..f693c8f675 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -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>, 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(()); diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 8d8c75b8af..8e0a72cf99 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -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!(