diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 7c8151d20c..c280058670 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1055,6 +1055,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( diff --git a/codex-rs/core/src/codex/compact.rs b/codex-rs/core/src/codex/compact.rs index a8340cd000..eaeae64bf6 100644 --- a/codex-rs/core/src/codex/compact.rs +++ b/codex-rs/core/src/codex/compact.rs @@ -160,6 +160,8 @@ async fn run_compact_task_inner( .collect(); new_history.extend(ghost_snapshots); 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(), diff --git a/codex-rs/core/src/conversation_history.rs b/codex-rs/core/src/conversation_history.rs index 9c724df62c..51c4ab4ded 100644 --- a/codex-rs/core/src/conversation_history.rs +++ b/codex-rs/core/src/conversation_history.rs @@ -428,6 +428,12 @@ impl ConversationHistory { | ResponseItem::Other => item.clone(), } } + + pub(crate) fn reset_last_token_usage(&mut self) { + if let Some(info) = self.token_info.as_mut() { + info.last_token_usage = TokenUsage::default(); + } + } } pub(crate) fn format_output_for_model_body(content: &str) -> String { @@ -573,6 +579,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(); diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index a41d2b6342..85f98daa38 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -68,4 +68,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(); + } }