From 4dde907d27489046daf2e2bbe90e59eb6ffecad8 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 23 Jun 2026 15:49:30 -0700 Subject: [PATCH] refactor: extract context window token status (#29664) ## Why This PR keeps the mechanical helper extraction separate from the behavior change in #29665. The follow-up needs the token-window accounting from `turn.rs` in another call path, but reviewing that is much easier when the helper extraction is separate from the semantic change. ## What - Adds `session/context_window.rs` with `ContextWindowTokenStatus`. - Moves the existing auto-compaction token-status calculation out of `session/turn.rs`. - Replaces the duplicated inline remaining-token calculation in `turn.rs` with `tokens_until_compaction()`. This PR is intended to be behavior-preserving. The `get_context_remaining` behavior change is stacked separately in #29665. ## Testing - `just test -p codex-core auto_compact_body_after_prefix` --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/29664). * #29665 * __->__ #29664 --- codex-rs/core/src/session/context_window.rs | 76 ++++++++++++++++++ codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/turn.rs | 85 +++------------------ 3 files changed, 86 insertions(+), 76 deletions(-) create mode 100644 codex-rs/core/src/session/context_window.rs diff --git a/codex-rs/core/src/session/context_window.rs b/codex-rs/core/src/session/context_window.rs new file mode 100644 index 000000000..8dc3d002f --- /dev/null +++ b/codex-rs/core/src/session/context_window.rs @@ -0,0 +1,76 @@ +use super::session::Session; +use super::turn_context::TurnContext; +use codex_protocol::config_types::AutoCompactTokenLimitScope; + +#[derive(Debug)] +pub(crate) struct ContextWindowTokenStatus { + // Full active context usage, independent of the configured auto-compact scope. + pub(crate) active_context_tokens: i64, + // Usage counted against `model_auto_compact_token_limit` for the current scope. + pub(crate) auto_compact_scope_tokens: i64, + pub(crate) auto_compact_scope_limit: i64, + pub(crate) full_context_window_limit: Option, + pub(crate) auto_compact_window_prefill_tokens: Option, + pub(crate) full_context_window_limit_reached: bool, + pub(crate) token_limit_reached: bool, +} + +impl ContextWindowTokenStatus { + pub(crate) fn tokens_until_compaction(&self) -> i64 { + let full_context_remaining = self.full_context_window_limit.map_or(i64::MAX, |limit| { + limit.saturating_sub(self.active_context_tokens) + }); + self.auto_compact_scope_limit + .saturating_sub(self.auto_compact_scope_tokens) + .min(full_context_remaining) + .max(0) + } +} + +pub(crate) async fn context_window_token_status( + sess: &Session, + turn_context: &TurnContext, +) -> ContextWindowTokenStatus { + let active_context_tokens = sess.get_total_token_usage().await; + let mut auto_compact_window_prefill_tokens = None; + let (auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit) = + match turn_context.config.model_auto_compact_token_limit_scope { + AutoCompactTokenLimitScope::Total => ( + active_context_tokens, + turn_context + .model_info + .auto_compact_token_limit() + .unwrap_or(i64::MAX), + None, + ), + AutoCompactTokenLimitScope::BodyAfterPrefix => { + let window = sess.auto_compact_window_snapshot().await; + auto_compact_window_prefill_tokens = window.prefill_input_tokens; + let baseline = window.prefill_input_tokens.unwrap_or(active_context_tokens); + ( + active_context_tokens.saturating_sub(baseline), + turn_context + .config + .model_auto_compact_token_limit + .or_else(|| turn_context.model_info.auto_compact_token_limit()) + .unwrap_or(i64::MAX), + turn_context.model_context_window(), + ) + } + }; + let full_context_window_limit_reached = + full_context_window_limit.is_some_and(|full_context_window_limit| { + active_context_tokens >= full_context_window_limit + }); + let token_limit_reached = + auto_compact_scope_tokens >= auto_compact_scope_limit || full_context_window_limit_reached; + ContextWindowTokenStatus { + active_context_tokens, + auto_compact_scope_tokens, + auto_compact_scope_limit, + full_context_window_limit, + auto_compact_window_prefill_tokens, + full_context_window_limit_reached, + token_limit_reached, + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8b8044bad..63ebf6729 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -211,6 +211,7 @@ use codex_protocol::exec_output::StreamOutput; mod code_mode_warning; mod config_lock; +pub(crate) mod context_window; mod handlers; mod inject; mod input_queue; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 100089723..43ca2f5c9 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -293,8 +293,11 @@ pub(crate) async fn run_turn( let (has_pending_input, token_status, estimated_token_count) = async { let has_pending_input = sess.input_queue.has_pending_input(&sess.active_turn).await; - let token_status = - auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await; + let token_status = super::context_window::context_window_token_status( + sess.as_ref(), + turn_context.as_ref(), + ) + .await; let estimated_token_count = sess.get_estimated_token_count(turn_context.as_ref()).await; (has_pending_input, token_status, estimated_token_count) @@ -321,17 +324,7 @@ pub(crate) async fn run_turn( "post sampling token usage" ); - let tokens_after_sampling = token_status.active_context_tokens; - let full_context_remaining = token_status - .full_context_window_limit - .map_or(i64::MAX, |limit| { - limit.saturating_sub(tokens_after_sampling) - }); - let tokens_until_compaction = token_status - .auto_compact_scope_limit - .saturating_sub(token_status.auto_compact_scope_tokens) - .min(full_context_remaining) - .max(0); + let tokens_until_compaction = token_status.tokens_until_compaction(); super::token_budget::maybe_record( sess.as_ref(), turn_context.as_ref(), @@ -801,68 +794,6 @@ async fn track_turn_resolved_config_analytics( }); } -#[derive(Debug)] -struct AutoCompactTokenStatus { - // Full active context usage, independent of the configured auto-compact scope. - active_context_tokens: i64, - // Usage counted against `model_auto_compact_token_limit` for the current scope. - auto_compact_scope_tokens: i64, - auto_compact_scope_limit: i64, - full_context_window_limit: Option, - auto_compact_window_prefill_tokens: Option, - full_context_window_limit_reached: bool, - token_limit_reached: bool, -} - -async fn auto_compact_token_status( - sess: &Session, - turn_context: &TurnContext, -) -> AutoCompactTokenStatus { - let active_context_tokens = sess.get_total_token_usage().await; - let mut auto_compact_window_prefill_tokens = None; - let (auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit) = - match turn_context.config.model_auto_compact_token_limit_scope { - AutoCompactTokenLimitScope::Total => ( - active_context_tokens, - turn_context - .model_info - .auto_compact_token_limit() - .unwrap_or(i64::MAX), - None, - ), - AutoCompactTokenLimitScope::BodyAfterPrefix => { - let window = sess.auto_compact_window_snapshot().await; - auto_compact_window_prefill_tokens = window.prefill_input_tokens; - let baseline = window.prefill_input_tokens.unwrap_or(active_context_tokens); - ( - active_context_tokens.saturating_sub(baseline), - turn_context - .config - .model_auto_compact_token_limit - .or_else(|| turn_context.model_info.auto_compact_token_limit()) - .unwrap_or(i64::MAX), - turn_context.model_context_window(), - ) - } - }; - let full_context_window_limit_reached = - full_context_window_limit.is_some_and(|full_context_window_limit| { - active_context_tokens >= full_context_window_limit - }); - let token_limit_reached = - auto_compact_scope_tokens >= auto_compact_scope_limit || full_context_window_limit_reached; - - AutoCompactTokenStatus { - active_context_tokens, - auto_compact_scope_tokens, - auto_compact_scope_limit, - full_context_window_limit, - auto_compact_window_prefill_tokens, - full_context_window_limit_reached, - token_limit_reached, - } -} - #[instrument(level = "trace", skip_all)] async fn run_pre_sampling_compact( sess: &Arc, @@ -878,7 +809,9 @@ async fn run_pre_sampling_compact( } maybe_run_previous_model_inline_compact(sess, turn_context, client_session).await?; - let token_status = auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await; + let token_status = + super::context_window::context_window_token_status(sess.as_ref(), turn_context.as_ref()) + .await; // Compact if the configured auto-compaction budget or usable context window is exhausted. if token_status.token_limit_reached { let step_context = sess.capture_step_context(Arc::clone(turn_context)).await;