refactor: extract context window token status

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

@@ -0,0 +1,78 @@
+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<i64>,
+    pub(crate) auto_compact_window_prefill_tokens: Option<i64>,
+    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)
+    }
+}
+
+impl Session {
+    pub(crate) async fn context_window_token_status(
+        &self,
+        turn_context: &TurnContext,
+    ) -> ContextWindowTokenStatus {
+        let active_context_tokens = self.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 = self.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,
+        }
+    }
+}

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

@@ -209,6 +209,7 @@ use codex_protocol::error::Result as CodexResult;
 use codex_protocol::exec_output::StreamOutput;
 
 mod config_lock;
+pub(crate) mod context_window;
 mod handlers;
 mod inject;
 mod input_queue;

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

@@ -294,8 +294,9 @@ 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 = sess
+                        .context_window_token_status(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)
@@ -322,17 +323,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(),
@@ -802,68 +793,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<i64>,
-    auto_compact_window_prefill_tokens: Option<i64>,
-    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<Session>,
@@ -879,7 +808,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 = sess
+        .context_window_token_status(turn_context.as_ref())
+        .await;
     // Compact if the configured auto-compaction budget or usable context window is exhausted.
     if token_status.token_limit_reached {
         run_auto_compact(