feat: proxy context left after compaction (#6597)

codex-rs/core/src/codex.rs

@@ -96,6 +96,7 @@ use crate::protocol::StreamErrorEvent;
 use crate::protocol::Submission;
 use crate::protocol::TokenCountEvent;
 use crate::protocol::TokenUsage;
+use crate::protocol::TokenUsageInfo;
 use crate::protocol::TurnDiffEvent;
 use crate::protocol::WarningEvent;
 use crate::rollout::RolloutRecorder;
@@ -1082,6 +1083,36 @@ impl Session {
         self.send_token_count_event(turn_context).await;
     }
 
+    pub(crate) async fn override_last_token_usage_estimate(
+        &self,
+        turn_context: &TurnContext,
+        estimated_total_tokens: i64,
+    ) {
+        {
+            let mut state = self.state.lock().await;
+            let mut info = state.token_info().unwrap_or(TokenUsageInfo {
+                total_token_usage: TokenUsage::default(),
+                last_token_usage: TokenUsage::default(),
+                model_context_window: None,
+            });
+
+            info.last_token_usage = TokenUsage {
+                input_tokens: 0,
+                cached_input_tokens: 0,
+                output_tokens: 0,
+                reasoning_output_tokens: 0,
+                total_tokens: estimated_total_tokens.max(0),
+            };
+
+            if info.model_context_window.is_none() {
+                info.model_context_window = turn_context.client.get_model_context_window();
+            }
+
+            state.set_token_info(Some(info));
+        }
+        self.send_token_count_event(turn_context).await;
+    }
+
     pub(crate) async fn update_rate_limits(
         &self,
         turn_context: &TurnContext,

codex-rs/core/src/compact.rs

@@ -153,6 +153,15 @@ async fn run_compact_task_inner(
     new_history.extend(ghost_snapshots);
     sess.replace_history(new_history).await;
 
+    if let Some(estimated_tokens) = sess
+        .clone_history()
+        .await
+        .estimate_token_count(&turn_context)
+    {
+        sess.override_last_token_usage_estimate(&turn_context, estimated_tokens)
+            .await;
+    }
+
     let rollout_item = RolloutItem::Compacted(CompactedItem {
         message: summary_text.clone(),
     });

codex-rs/core/src/context_manager/history.rs

@@ -1,13 +1,14 @@
+use crate::codex::TurnContext;
+use crate::context_manager::normalize;
+use crate::context_manager::truncate::format_output_for_model_body;
+use crate::context_manager::truncate::globally_truncate_function_output_items;
 use codex_protocol::models::FunctionCallOutputPayload;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::TokenUsage;
 use codex_protocol::protocol::TokenUsageInfo;
+use codex_utils_tokenizer::Tokenizer;
 use std::ops::Deref;
 
-use crate::context_manager::normalize;
-use crate::context_manager::truncate::format_output_for_model_body;
-use crate::context_manager::truncate::globally_truncate_function_output_items;
-
 /// Transcript of conversation history
 #[derive(Debug, Clone, Default)]
 pub(crate) struct ContextManager {
@@ -28,6 +29,10 @@ impl ContextManager {
         self.token_info.clone()
     }
 
+    pub(crate) fn set_token_info(&mut self, info: Option<TokenUsageInfo>) {
+        self.token_info = info;
+    }
+
     pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
         match &mut self.token_info {
             Some(info) => info.fill_to_context_window(context_window),
@@ -68,6 +73,28 @@ impl ContextManager {
         history
     }
 
+    // Estimate the number of tokens in the history. Return None if no tokenizer
+    // is available. This does not consider the reasoning traces.
+    // /!\ The value is a lower bound estimate and does not represent the exact
+    // context length.
+    pub(crate) fn estimate_token_count(&self, turn_context: &TurnContext) -> Option<i64> {
+        let model = turn_context.client.get_model();
+        let tokenizer = Tokenizer::for_model(model.as_str()).ok()?;
+        let model_family = turn_context.client.get_model_family();
+
+        Some(
+            self.items
+                .iter()
+                .map(|item| {
+                    serde_json::to_string(&item)
+                        .map(|item| tokenizer.count(&item))
+                        .unwrap_or_default()
+                })
+                .sum::<i64>()
+                + tokenizer.count(model_family.base_instructions.as_str()),
+        )
+    }
+
     pub(crate) fn remove_first_item(&mut self) {
         if !self.items.is_empty() {
             // Remove the oldest item (front of the list). Items are ordered from

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

@@ -42,6 +42,10 @@ impl SessionState {
         self.history.replace(items);
     }
 
+    pub(crate) fn set_token_info(&mut self, info: Option<TokenUsageInfo>) {
+        self.history.set_token_info(info);
+    }
+
     // Token/rate limit helpers
     pub(crate) fn update_token_info_from_usage(
         &mut self,