chore: consolidate compaction token usage (#6894)

This commit is contained in:
jif-oai 2025-11-19 11:26:01 +00:00 committed by GitHub
parent 44c747837a
commit 3e9e1d993d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 10 additions and 23 deletions

View file

@ -1085,11 +1085,14 @@ 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,
) {
pub(crate) async fn recompute_token_usage(&self, turn_context: &TurnContext) {
let Some(estimated_total_tokens) = self
.clone_history()
.await
.estimate_token_count(turn_context)
else {
return;
};
{
let mut state = self.state.lock().await;
let mut info = state.token_info().unwrap_or(TokenUsageInfo {

View file

@ -171,15 +171,7 @@ async fn run_compact_task_inner(
.collect();
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;
}
sess.recompute_token_usage(&turn_context).await;
let rollout_item = RolloutItem::Compacted(CompactedItem {
message: summary_text.clone(),

View file

@ -66,15 +66,7 @@ async fn run_remote_compact_task_inner_impl(
new_history.extend(ghost_snapshots);
}
sess.replace_history(new_history.clone()).await;
if let Some(estimated_tokens) = sess
.clone_history()
.await
.estimate_token_count(turn_context.as_ref())
{
sess.override_last_token_usage_estimate(turn_context.as_ref(), estimated_tokens)
.await;
}
sess.recompute_token_usage(turn_context).await;
let compacted_item = CompactedItem {
message: String::new(),