Preserve token totals after /compact

Signed-off-by: Fahad <[email protected]>

Author: guidedways

codex-rs/core/src/codex.rs

@@ -995,6 +995,11 @@ impl Session {
         }
     }
 
+    async fn reset_last_token_usage(&self) {
+        let mut state = self.state.lock().await;
+        state.reset_last_token_usage();
+    }
+
     /// Record a user input item to conversation history and also persist a
     /// corresponding UserMessage EventMsg to rollout.
     async fn record_input_and_rollout_usermsg(

codex-rs/core/src/codex/compact.rs

@@ -154,6 +154,8 @@ async fn run_compact_task_inner(
     let initial_context = sess.build_initial_context(turn_context.as_ref());
     let new_history = build_compacted_history(initial_context, &user_messages, &summary_text);
     sess.replace_history(new_history).await;
+    sess.reset_last_token_usage().await;
+    sess.send_token_count_event(turn_context.as_ref()).await;
 
     let rollout_item = RolloutItem::Compacted(CompactedItem {
         message: summary_text.clone(),

codex-rs/core/src/conversation_history.rs

@@ -332,6 +332,12 @@ impl ConversationHistory {
             model_context_window,
         );
     }
+
+    pub(crate) fn reset_last_token_usage(&mut self) {
+        if let Some(info) = self.token_info.as_mut() {
+            info.last_token_usage = TokenUsage::default();
+        }
+    }
 }
 
 #[inline]
@@ -395,6 +401,24 @@ mod tests {
         }
     }
 
+    #[test]
+    fn replace_preserves_totals_and_reset_clears_last_usage() {
+        let mut history = ConversationHistory::new();
+        let mut usage = TokenUsage::default();
+        usage.total_tokens = 1280;
+        usage.input_tokens = 640;
+        history.update_token_info(&usage, Some(13_000));
+
+        history.replace(Vec::new());
+        history.reset_last_token_usage();
+
+        let info = history
+            .token_info()
+            .expect("token info should persist after replace");
+        assert_eq!(info.total_token_usage.total_tokens, 1280);
+        assert_eq!(info.last_token_usage.total_tokens, 0);
+    }
+
     #[test]
     fn filters_non_api_messages() {
         let mut h = ConversationHistory::default();

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

@@ -72,4 +72,8 @@ impl SessionState {
     pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
         self.history.set_token_usage_full(context_window);
     }
+
+    pub(crate) fn reset_last_token_usage(&mut self) {
+        self.history.reset_last_token_usage();
+    }
 }