core: reset context for token budget compaction

codex-rs/core/src/compact_token_budget.rs

@@ -0,0 +1,76 @@
+use std::sync::Arc;
+
+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::NewContextWindowMode;
+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;
+
+    run_compact_task_inner(&sess, &turn_context, 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>,
+) -> CodexResult<()> {
+    run_compact_task_inner(&sess, &turn_context, CompactionTrigger::Auto).await
+}
+
+async fn run_compact_task_inner(
+    sess: &Arc<Session>,
+    turn_context: &Arc<TurnContext>,
+    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(), NewContextWindowMode::ForceStart)
+        .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

@@ -251,6 +251,14 @@ use self::world_state::build_world_state_from_turn_context_item;
 #[cfg(test)]
 mod rollout_reconstruction_tests;
 
+#[derive(Clone, Copy, Debug)]
+pub(crate) enum NewContextWindowMode {
+    /// Start a fresh context window regardless of whether a new-context request is pending.
+    ForceStart,
+    /// Start a fresh context window only if one was explicitly requested.
+    StartIfRequested,
+}
+
 #[derive(Debug, PartialEq)]
 pub enum SteerInputError {
     NoActiveTurn(Vec<UserInput>),
@@ -3398,15 +3406,27 @@ impl Session {
         state.request_new_context_window();
     }
 
-    pub(crate) async fn maybe_start_new_context_window(
+    pub(crate) async fn start_new_context_window(
         &self,
         turn_context: &TurnContext,
+        mode: NewContextWindowMode,
     ) -> Option<u64> {
         let window = {
             let mut state = self.state.lock().await;
-            state.start_new_context_window_if_requested()
+            state.start_new_context_window(mode)
         };
         let (window_number, window_ids) = window?;
+        self.install_new_context_window(turn_context, window_number, window_ids)
+            .await;
+        Some(window_number)
+    }
+
+    async fn install_new_context_window(
+        &self,
+        turn_context: &TurnContext,
+        window_number: u64,
+        window_ids: AutoCompactWindowIds,
+    ) {
         let world_state = self.build_world_state(turn_context).await;
         let context_items = self
             .build_initial_context_with_world_state(turn_context, &world_state)
@@ -3435,7 +3455,6 @@ impl Session {
             state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact);
         }
         self.recompute_token_usage(turn_context).await;
-        Some(window_number)
     }
 
     pub(crate) async fn reference_context_item(&self) -> Option<TurnContextItem> {

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

@@ -39,6 +39,7 @@ use crate::responses_metadata::CodexResponsesMetadata;
 use crate::responses_metadata::CodexResponsesRequestKind;
 use crate::responses_retry::ResponsesStreamRequest;
 use crate::responses_retry::handle_retryable_response_stream_error;
+use crate::session::NewContextWindowMode;
 use crate::session::PreviousTurnSettings;
 use crate::session::TurnInput;
 use crate::session::session::Session;
@@ -341,7 +342,10 @@ pub(crate) async fn run_turn(
                 .await;
 
                 let started_new_context_window = sess
-                    .maybe_start_new_context_window(turn_context.as_ref())
+                    .start_new_context_window(
+                        turn_context.as_ref(),
+                        NewContextWindowMode::StartIfRequested,
+                    )
                     .await
                     .is_some();
                 if started_new_context_window && needs_follow_up {
@@ -357,6 +361,16 @@ pub(crate) async fn run_turn(
                     && token_limit_reached
                     && needs_follow_up
                 {
+                    if turn_context.config.features.enabled(Feature::TokenBudget)
+                        && model_needs_follow_up
+                    {
+                        // A token-budget reset clears the current-window history. Preserve
+                        // active tool-call continuation state first; pre-turn compaction will
+                        // reset the window before the next user turn if the budget is still
+                        // exhausted.
+                        can_drain_pending_input = false;
+                        continue;
+                    }
                     if let Err(err) = run_auto_compact(
                         &sess,
                         &turn_context,
@@ -989,6 +1003,17 @@ async fn run_auto_compact(
     reason: CompactionReason,
     phase: CompactionPhase,
 ) -> CodexResult<()> {
+    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),
+        )
+        .await?;
+        return Ok(());
+    }
+
     if should_use_remote_compact_task(turn_context.provider.info()) {
         if turn_context
             .config

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

@@ -12,6 +12,7 @@ use super::auto_compact_window::AutoCompactWindow;
 use super::auto_compact_window::AutoCompactWindowIds;
 use super::auto_compact_window::AutoCompactWindowSnapshot;
 use crate::context_manager::ContextManager;
+use crate::session::NewContextWindowMode;
 use crate::session::PreviousTurnSettings;
 use crate::session::session::SessionConfiguration;
 use crate::session::time_reminder::CurrentTimeReminderState;
@@ -176,10 +177,13 @@ impl SessionState {
         self.auto_compact_window.request_new_context_window();
     }
 
-    pub(crate) fn start_new_context_window_if_requested(
+    pub(crate) fn start_new_context_window(
         &mut self,
+        mode: NewContextWindowMode,
     ) -> Option<(u64, AutoCompactWindowIds)> {
-        if !self.auto_compact_window.take_new_context_window_request() {
+        if matches!(mode, NewContextWindowMode::StartIfRequested)
+            && !self.auto_compact_window.take_new_context_window_request()
+        {
             return None;
         }
 

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

@@ -4,20 +4,27 @@ use codex_config::types::McpServerTransportConfig;
 use codex_core::config::TokenBudgetConfig;
 use codex_features::Feature;
 use codex_model_provider_info::built_in_model_providers;
+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;
@@ -26,11 +33,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;
@@ -79,6 +88,62 @@ fn tool_names(request: &ResponsesRequest) -> Vec<String> {
         .collect()
 }
 
+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(()));
@@ -395,19 +460,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;
 
@@ -425,21 +489,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,
@@ -454,6 +519,141 @@ 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_preserves_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,
+        "active model follow-up should not run a compaction request before the continuation"
+    );
+    assert_eq!(
+        requests[1].function_call_output_content_and_success(call_id),
+        Some((
+            Some("You have 0 tokens left in this context window.".to_string()),
+            None
+        )),
+        "follow-up request should preserve active tool output after mid-turn 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_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, None);
+    assert_eq!(
+        follow_up_window_id, initial_window_id,
+        "mid-turn auto-compaction should not reset the token-budget context window while the model needs follow-up"
+    );
 
     Ok(())
 }