core: reset context for token budget compaction

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