core: reset context for token budget compaction (#29743)

## Why

When `Feature::TokenBudget` is enabled, compaction should behave like
`new_context`: start a fresh context window with the standard injected
context, without asking the server to summarize old history and without
carrying prior user or assistant messages into the next model request.

This is still a compaction operation from the client lifecycle
perspective. Manual `/compact` and auto-compaction should keep the same
observable side effects that clients and hooks expect, including compact
hooks and `TurnItem::ContextCompaction`.

## What changed

- Added `compact_token_budget` to run token-budget manual and inline
auto-compaction through a shared compaction lifecycle.
- Split pending `new_context` requests from forced context-window
startup: `take_new_context_window_request()` consumes pending requests,
and `start_new_context_window()` installs a fresh context window.
- Routed token-budget manual `/compact` and inline auto-compaction to
install a fresh context window locally instead of calling server/local
summarization.
- Preserved compact lifecycle side effects for token-budget compaction
by running pre/post compact hooks and emitting `ContextCompaction` item
start/completion events.
- Updated token-budget tests to assert fresh window IDs, absence of
server-side compaction calls, dropped prior transcript messages/tool
output after reset, and compact hook/item lifecycle behavior.

## Testing

- `just test -p codex-core
token_budget_context_uses_new_window_after_compaction`
- `just test -p codex-core token_budget_compaction_runs_compact_hooks`
- `just test -p codex-core
token_budget_mid_turn_auto_compaction_resets_before_active_follow_up`

---------

Co-authored-by: pakrym-oai <pakrym@openai.com>

codex-rs/core/src/compact_token_budget.rs

@@ -0,0 +1,90 @@
+use std::sync::Arc;
+
+use crate::compact::InitialContextInjection;
+use crate::context::world_state::WorldState;
+use crate::hook_runtime::PostCompactHookOutcome;
+use crate::hook_runtime::PreCompactHookOutcome;
+use crate::hook_runtime::run_post_compact_hooks;
+use crate::hook_runtime::run_pre_compact_hooks;
+use crate::session::session::Session;
+use crate::session::turn_context::TurnContext;
+use codex_analytics::CompactionTrigger;
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::Result as CodexResult;
+use codex_protocol::items::ContextCompactionItem;
+use codex_protocol::items::TurnItem;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::TurnStartedEvent;
+
+/// Runs token-budget manual compaction as a normal compaction lifecycle.
+///
+/// Token-budget compaction skips model/server summarization and installs a fresh context window
+/// instead. It is still modeled as compaction so compact hooks and `ContextCompaction` turn items
+/// observe the same lifecycle as local or remote compaction.
+pub(crate) async fn run_manual_compact_task(
+    sess: Arc<Session>,
+    turn_context: Arc<TurnContext>,
+) -> CodexResult<()> {
+    let start_event = EventMsg::TurnStarted(TurnStartedEvent {
+        turn_id: turn_context.sub_id.clone(),
+        trace_id: turn_context.trace_id.clone(),
+        started_at: turn_context.turn_timing_state.started_at_unix_secs().await,
+        model_context_window: turn_context.model_context_window(),
+        collaboration_mode_kind: turn_context.collaboration_mode.mode,
+    });
+    sess.send_event(&turn_context, start_event).await;
+
+    let world_state = Arc::new(
+        sess.build_world_state_for_environments(&turn_context, &turn_context.environments)
+            .await,
+    );
+    run_compact_task_inner(&sess, &turn_context, world_state, CompactionTrigger::Manual).await
+}
+
+/// Runs token-budget inline auto-compaction as a normal compaction lifecycle.
+///
+/// Token-budget compaction skips model/server summarization and installs a fresh context window
+/// instead. It is still modeled as compaction so compact hooks and `ContextCompaction` turn items
+/// observe the same lifecycle as local or remote compaction.
+pub(crate) async fn run_inline_auto_compact_task(
+    sess: Arc<Session>,
+    turn_context: Arc<TurnContext>,
+    initial_context_injection: InitialContextInjection,
+) -> CodexResult<()> {
+    let world_state = match initial_context_injection {
+        InitialContextInjection::BeforeLastUserMessage(world_state) => world_state,
+        InitialContextInjection::DoNotInject => Arc::new(
+            sess.build_world_state_for_environments(&turn_context, &turn_context.environments)
+                .await,
+        ),
+    };
+    run_compact_task_inner(&sess, &turn_context, world_state, CompactionTrigger::Auto).await
+}
+
+async fn run_compact_task_inner(
+    sess: &Arc<Session>,
+    turn_context: &Arc<TurnContext>,
+    world_state: Arc<WorldState>,
+    trigger: CompactionTrigger,
+) -> CodexResult<()> {
+    let pre_compact_outcome = run_pre_compact_hooks(sess, turn_context, trigger).await;
+    match pre_compact_outcome {
+        PreCompactHookOutcome::Continue => {}
+        PreCompactHookOutcome::Stopped => return Err(CodexErr::TurnAborted),
+    }
+
+    let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new());
+    sess.emit_turn_item_started(turn_context, &compaction_item)
+        .await;
+    sess.start_new_context_window(turn_context.as_ref(), world_state)
+        .await;
+    sess.emit_turn_item_completed(turn_context, compaction_item)
+        .await;
+
+    let post_compact_outcome = run_post_compact_hooks(sess, turn_context, trigger).await;
+    if let PostCompactHookOutcome::Stopped = post_compact_outcome {
+        return Err(CodexErr::TurnAborted);
+    }
+
+    Ok(())
+}

codex-rs/core/src/lib.rs

@@ -21,6 +21,7 @@ pub use turn_metadata::detached_memory_responses_metadata;
 mod codex_thread;
 mod compact_remote;
 mod compact_remote_v2;
+mod compact_token_budget;
 mod config_lock;
 pub use codex_thread::BackgroundTerminalInfo;
 pub use codex_thread::CodexThread;

codex-rs/core/src/session/mod.rs

@@ -3425,16 +3425,21 @@ impl Session {
         state.request_new_context_window();
     }
 
-    pub(crate) async fn maybe_start_new_context_window(
+    pub(crate) async fn take_new_context_window_request(&self) -> bool {
+        let mut state = self.state.lock().await;
+        state.take_new_context_window_request()
+    }
+
+    pub(crate) async fn start_new_context_window(
         &self,
         turn_context: &TurnContext,
         world_state: Arc<WorldState>,
-    ) -> Option<u64> {
+    ) -> u64 {
         let window = {
             let mut state = self.state.lock().await;
-            state.start_new_context_window_if_requested()
+            state.start_new_context_window()
         };
-        let (window_number, window_ids) = window?;
+        let (window_number, window_ids) = window;
         let context_items = self
             .build_initial_context_with_world_state(turn_context, world_state.as_ref())
             .await;
@@ -3462,7 +3467,7 @@ impl Session {
             state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact);
         }
         self.recompute_token_usage(turn_context).await;
-        Some(window_number)
+        window_number
     }
 
     pub(crate) async fn reference_context_item(&self) -> Option<TurnContextItem> {

codex-rs/core/src/session/turn.rs

@@ -331,22 +331,14 @@ pub(crate) async fn run_turn(
                 )
                 .await;
 
-                let started_new_context_window = sess
-                    .maybe_start_new_context_window(turn_context.as_ref(), Arc::clone(&world_state))
-                    .await
-                    .is_some();
-                if started_new_context_window && needs_follow_up {
-                    can_drain_pending_input = !model_needs_follow_up;
-                    continue;
-                }
-
                 // 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 turn_context
+                let auto_compact_needed = turn_context
                     .config
                     .features
                     .enabled(Feature::AutoCompaction)
-                    && token_limit_reached
-                    && needs_follow_up
+                    && token_limit_reached;
+                if needs_follow_up
+                    && (sess.take_new_context_window_request().await || auto_compact_needed)
                 {
                     if let Err(err) = run_auto_compact(
                         &sess,
@@ -928,6 +920,18 @@ async fn run_auto_compact(
     phase: CompactionPhase,
 ) -> CodexResult<()> {
     let turn_context = &step_context.turn;
+    if turn_context.config.features.enabled(Feature::TokenBudget) {
+        // Compaction is the reset request, so force a new context window
+        // instead of consuming a pending `new_context` tool request.
+        crate::compact_token_budget::run_inline_auto_compact_task(
+            Arc::clone(sess),
+            Arc::clone(turn_context),
+            initial_context_injection,
+        )
+        .await?;
+        return Ok(());
+    }
+
     if should_use_remote_compact_task(turn_context.provider.info()) {
         if turn_context
             .config

codex-rs/core/src/state/session.rs

@@ -187,16 +187,14 @@ impl SessionState {
         self.auto_compact_window.request_new_context_window();
     }
 
-    pub(crate) fn start_new_context_window_if_requested(
-        &mut self,
-    ) -> Option<(u64, AutoCompactWindowIds)> {
-        if !self.auto_compact_window.take_new_context_window_request() {
-            return None;
-        }
+    pub(crate) fn take_new_context_window_request(&mut self) -> bool {
+        self.auto_compact_window.take_new_context_window_request()
+    }
 
+    pub(crate) fn start_new_context_window(&mut self) -> (u64, AutoCompactWindowIds) {
         let window = self.auto_compact_window.advance();
         self.auto_compact_window.clear_prefill();
-        Some(window)
+        window
     }
 
     pub(crate) fn token_info(&self) -> Option<TokenUsageInfo> {

codex-rs/core/src/tasks/compact.rs

@@ -7,6 +7,7 @@ use super::emit_compact_metric;
 use crate::session::TurnInput;
 use crate::session::turn_context::TurnContext;
 use crate::state::TaskKind;
+use codex_features::Feature;
 use codex_protocol::error::CodexErr;
 use codex_protocol::user_input::UserInput;
 use tokio_util::sync::CancellationToken;
@@ -31,6 +32,11 @@ impl SessionTask for CompactTask {
         _cancellation_token: CancellationToken,
     ) -> SessionTaskResult {
         let session = session.clone_session();
+        if ctx.config.features.enabled(Feature::TokenBudget) {
+            crate::compact_token_budget::run_manual_compact_task(session, ctx).await?;
+            return Ok(None);
+        }
+
         let result = if crate::compact::should_use_remote_compact_task(ctx.provider.info()) {
             if ctx
                 .config