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

codex-rs/core/tests/suite/token_budget.rs

@@ -5,20 +5,27 @@ use codex_core::config::TokenBudgetConfig;
 use codex_features::Feature;
 use codex_model_provider_info::built_in_model_providers;
 use codex_protocol::config_types::AutoCompactTokenLimitScope;
+use codex_protocol::items::TurnItem;
 use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
 use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG;
 use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::HookEventName;
+use codex_protocol::protocol::HookRunStatus;
+use codex_protocol::protocol::ItemCompletedEvent;
+use codex_protocol::protocol::ItemStartedEvent;
 use codex_protocol::protocol::Op;
 use core_test_support::PathBufExt;
 use core_test_support::assert_regex_match;
 use core_test_support::context_snapshot;
 use core_test_support::context_snapshot::ContextSnapshotOptions;
+use core_test_support::hooks::trust_discovered_hooks;
 use core_test_support::responses::ResponsesRequest;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_completed_with_tokens;
 use core_test_support::responses::ev_function_call;
 use core_test_support::responses::ev_response_created;
+use core_test_support::responses::mount_compact_json_once;
 use core_test_support::responses::mount_sse_sequence;
 use core_test_support::responses::sse;
 use core_test_support::responses::start_mock_server;
@@ -27,11 +34,13 @@ use core_test_support::stdio_server_bin;
 use core_test_support::test_codex::local;
 use core_test_support::test_codex::test_codex;
 use core_test_support::wait_for_event;
+use core_test_support::wait_for_event_match;
 use core_test_support::wait_for_mcp_server;
 use pretty_assertions::assert_eq;
 use serde_json::Value;
 use serde_json::json;
 use std::collections::HashMap;
+use std::path::Path;
 use std::time::Duration;
 
 const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000;
@@ -96,6 +105,62 @@ fn ev_completed_with_usage(id: &str, input_tokens: i64, output_tokens: i64) -> V
     })
 }
 
+fn python_hook_command(script_path: &Path) -> String {
+    format!("python3 \"{}\"", script_path.display())
+}
+
+fn write_token_budget_compact_hooks(home: &Path) {
+    let script_path = home.join("token_budget_compact_hook.py");
+    std::fs::write(
+        &script_path,
+        "import json\nimport sys\njson.load(sys.stdin)\n",
+    )
+    .expect("write compact hook script");
+    let hooks = json!({
+        "hooks": {
+            "PreCompact": [{
+                "matcher": "manual",
+                "hooks": [{
+                    "type": "command",
+                    "command": python_hook_command(&script_path),
+                }]
+            }],
+            "PostCompact": [{
+                "matcher": "manual",
+                "hooks": [{
+                    "type": "command",
+                    "command": python_hook_command(&script_path),
+                }]
+            }]
+        }
+    });
+    std::fs::write(home.join("hooks.json"), hooks.to_string()).expect("write hooks.json");
+}
+
+async fn assert_context_compaction_item_lifecycle(codex: &std::sync::Arc<codex_core::CodexThread>) {
+    let mut saw_compaction_started = false;
+    let mut saw_compaction_completed = false;
+
+    loop {
+        let event = codex.next_event().await.expect("next event");
+        match event.msg {
+            EventMsg::ItemStarted(ItemStartedEvent {
+                item: TurnItem::ContextCompaction(_),
+                ..
+            }) => saw_compaction_started = true,
+            EventMsg::ItemCompleted(ItemCompletedEvent {
+                item: TurnItem::ContextCompaction(_),
+                ..
+            }) => saw_compaction_completed = true,
+            EventMsg::TurnComplete(_) => break,
+            _ => {}
+        }
+    }
+
+    assert!(saw_compaction_started, "compaction item should start");
+    assert!(saw_compaction_completed, "compaction item should complete");
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> {
     skip_if_no_network!(Ok(()));
@@ -540,19 +605,18 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
     let responses = mount_sse_sequence(
         &server,
         vec![
-            sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
             sse(vec![
-                ev_response_created("resp-compact"),
-                ev_assistant_message("msg-compact", "compact summary"),
-                ev_completed("resp-compact"),
+                ev_response_created("resp-1"),
+                ev_assistant_message("msg-1", "assistant before compact"),
+                ev_completed("resp-1"),
             ]),
             sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
         ],
     )
     .await;
+    let compact = mount_compact_json_once(&server, json!({ "output": [] })).await;
 
     let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
-    model_provider.name = "OpenAI-compatible test provider".to_string();
     model_provider.base_url = Some(format!("{}/v1", server.uri()));
     model_provider.supports_websockets = false;
 
@@ -570,21 +634,22 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
 
     test.submit_turn("before compact").await?;
     test.codex.submit(Op::Compact).await?;
-    wait_for_event(&test.codex, |event| {
-        matches!(event, EventMsg::TurnComplete(_))
-    })
-    .await;
+    assert_context_compaction_item_lifecycle(&test.codex).await;
     test.submit_turn("after compact").await?;
 
     let requests = responses.requests();
-    assert_eq!(requests.len(), 3);
+    assert_eq!(requests.len(), 2);
+    assert!(
+        compact.requests().is_empty(),
+        "token budget compaction should not call server-side compaction"
+    );
 
     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_previous_window_id, initial_window_id) =
         token_budget_window_ids(&initial_token_budget[0], thread_id);
-    let post_compaction_token_budget = token_budget_contexts(&requests[2]);
+    let post_compaction_token_budget = token_budget_contexts(&requests[1]);
     assert_eq!(post_compaction_token_budget.len(), 1);
     let (
         post_compaction_first_window_id,
@@ -599,6 +664,147 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
         Some(initial_window_id.as_str())
     );
     assert_ne!(post_compaction_window_id, initial_window_id);
+    assert!(
+        !requests[1].body_contains_text("before compact"),
+        "token budget compaction should drop prior user messages"
+    );
+    assert!(
+        !requests[1].body_contains_text("assistant before compact"),
+        "token budget compaction should drop prior assistant messages"
+    );
+    assert!(
+        requests[1].body_contains_text("after compact"),
+        "follow-up should still include the new turn input"
+    );
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn token_budget_compaction_runs_compact_hooks() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let test = test_codex()
+        .with_pre_build_hook(write_token_budget_compact_hooks)
+        .with_config(|config| {
+            config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
+            config
+                .features
+                .enable(Feature::TokenBudget)
+                .expect("test config should allow token budget");
+            trust_discovered_hooks(config);
+        })
+        .build(&server)
+        .await?;
+
+    test.codex.submit(Op::Compact).await?;
+
+    let pre_compact = wait_for_event_match(&test.codex, |event| match event {
+        EventMsg::HookCompleted(completed)
+            if completed.run.event_name == HookEventName::PreCompact =>
+        {
+            Some(completed.clone())
+        }
+        _ => None,
+    })
+    .await;
+    assert_eq!(pre_compact.run.status, HookRunStatus::Completed);
+
+    let post_compact = wait_for_event_match(&test.codex, |event| match event {
+        EventMsg::HookCompleted(completed)
+            if completed.run.event_name == HookEventName::PostCompact =>
+        {
+            Some(completed.clone())
+        }
+        _ => None,
+    })
+    .await;
+    assert_eq!(post_compact.run.status, HookRunStatus::Completed);
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let call_id = "remaining-call";
+    let responses = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_function_call(call_id, "get_context_remaining", "{}"),
+                ev_completed_with_tokens("resp-1", /*total_tokens*/ 9_500),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_assistant_message("msg-2", "done"),
+                ev_completed("resp-2"),
+            ]),
+        ],
+    )
+    .await;
+    let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
+    model_provider.name = "OpenAI (test)".into();
+    model_provider.base_url = Some(format!("{}/v1", server.uri()));
+    model_provider.supports_websockets = false;
+    let test = test_codex()
+        .with_config(move |config| {
+            config.model_provider = model_provider;
+            config.model_context_window = Some(10_000);
+            config
+                .features
+                .enable(Feature::TokenBudget)
+                .expect("test config should allow token budget");
+        })
+        .build(&server)
+        .await?;
+
+    test.submit_turn("trigger mid-turn auto compaction").await?;
+
+    let requests = responses.requests();
+    assert_eq!(
+        requests.len(),
+        2,
+        "token-budget auto-compaction should reset locally before the continuation"
+    );
+    assert!(
+        !requests[1].input().iter().any(|item| {
+            item.get("type").and_then(Value::as_str) == Some("function_call_output")
+                && item.get("call_id").and_then(Value::as_str) == Some(call_id)
+        }),
+        "fresh token-budget windows should drop active tool output with the prior history"
+    );
+
+    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 follow_up_token_budget = token_budget_contexts(&requests[1]);
+    assert_eq!(follow_up_token_budget.len(), 1);
+    let (follow_up_first_window_id, follow_up_previous_window_id, follow_up_window_id) =
+        token_budget_window_ids(&follow_up_token_budget[0], thread_id);
+    assert_eq!(follow_up_first_window_id, initial_first_window_id);
+    assert_eq!(
+        follow_up_previous_window_id.as_deref(),
+        Some(initial_window_id.as_str())
+    );
+    assert!(
+        !requests[1].body_contains_text("trigger mid-turn auto compaction"),
+        "fresh token-budget windows should drop prior user messages"
+    );
+    assert_ne!(
+        follow_up_window_id, initial_window_id,
+        "mid-turn token-budget auto-compaction should reset the context window"
+    );
 
     Ok(())
 }