yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

agent: Accumulate and persist cumulative token usage for threads (#58378)

Browse Source 
`Thread::cumulative_token_usage` has been dead code since the agent2
rewrite: the field was added as a persistence placeholder in #36483 with
`#[allow(unused)]`, and the accumulation logic was deleted along with
agent1 in #40495. Since then, every saved thread persists
`cumulative_token_usage` as `{}` in `threads.db`, breaking offline
token-usage analysis and silently degrading eval reporting (`eval_cli`
prefers this value and was always falling back to the last request's
snapshot, undercounting multi-step turns).

This PR restores the accumulation, mirroring the original agent1
implementation:

- Providers emit multiple cumulative-per-request `UsageUpdate` snapshots
during a single request, so naive `+=` would double-count. We track the
previous snapshot for the in-flight request
(`current_request_token_usage`) and accumulate only the delta, resetting
the tracker each time a new completion request is built in
`run_turn_internal`. This counts each request exactly once, including
tool-call rounds and retries within a turn.
- Unlike agent1, the delta uses per-field `saturating_sub` instead of
the panicking `Sub` impl, so a provider reporting a smaller snapshot
can't crash Zed.
- Removed the `#[allow(unused)]`; the value now round-trips through
`to_db`/`from_db`, surviving restarts and continuing to grow when a
thread is reopened.
- Added a `cumulative_token_usage()` accessor following the
`latest_request_token_usage` convention.

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- N/A

---------

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>
This commit is contained in:
Anant GoelandBen Brandt authored and GitHub committed 2026-06-04 15:24:02 +00:00
1 parent 05f6f164a3
commit 5f2eba257d
2 files changed
+295 -2
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/agent/src/tests/mod.rs
+140
View File
@@ -3206,6 +3206,146 @@ async fn test_latest_token_usage_counts_cached_input_tokens(cx: &mut TestAppCont
});
}
#[gpui::test]
async fn test_cumulative_token_usage(cx: &mut TestAppContext) {
let ThreadTest {
model,
thread,
project_context,
..
} = setup(cx, TestModel::Fake).await;
let fake_model = model.as_fake();
thread
.update(cx, |thread, cx| {
thread.add_tool(EchoTool);
thread.send(UserMessageId::new(), ["Use the echo tool"], cx)
})
.unwrap();
cx.run_until_parked();
// The first request emits two cumulative snapshots; only the final values
// must be counted, exactly once.
fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 100,
output_tokens: 10,
..Default::default()
},
));
fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 100,
output_tokens: 50,
..Default::default()
},
));
fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
LanguageModelToolUse {
id: "tool_1".into(),
name: EchoTool::NAME.into(),
raw_input: json!({"text": "hello"}).to_string(),
input: json!({"text": "hello"}),
is_input_complete: true,
thought_signature: None,
},
));
fake_model.end_last_completion_stream();
cx.run_until_parked();
// The second request (after the tool call) is counted in addition to the first.
fake_model.send_last_completion_stream_text_chunk("Done");
fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 200,
output_tokens: 30,
..Default::default()
},
));
fake_model.end_last_completion_stream();
cx.run_until_parked();
let expected = TokenUsage {
input_tokens: 300,
output_tokens: 80,
..Default::default()
};
thread.read_with(cx, |thread, _| {
assert_eq!(thread.cumulative_token_usage(), expected);
});
let db_thread = thread.read_with(cx, |thread, cx| thread.to_db(cx)).await;
assert_eq!(db_thread.cumulative_token_usage, expected);
cx.update(|cx| {
LanguageModelRegistry::test(cx);
});
let restored = cx.update(|cx| {
let thread = thread.read(cx);
let project = thread.project.clone();
let context_server_registry = thread.context_server_registry.clone();
let templates = thread.templates.clone();
cx.new(|cx| {
Thread::from_db(
acp::SessionId::new("restored"),
db_thread,
project,
project_context.clone(),
context_server_registry,
templates,
cx,
)
})
});
restored.read_with(cx, |thread, _| {
assert_eq!(thread.cumulative_token_usage(), expected);
});
}
#[gpui::test]
async fn test_cumulative_token_usage_keeps_accounted_usage_monotonic(cx: &mut TestAppContext) {
let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
let fake_model = model.as_fake();
thread
.update(cx, |thread, cx| {
thread.send(UserMessageId::new(), ["hello"], cx)
})
.unwrap();
cx.run_until_parked();
fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 100,
output_tokens: 10,
..Default::default()
},
));
fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage::default(),
));
fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 100,
output_tokens: 50,
..Default::default()
},
));
fake_model.end_last_completion_stream();
cx.run_until_parked();
thread.read_with(cx, |thread, _| {
assert_eq!(
thread.cumulative_token_usage(),
TokenUsage {
input_tokens: 100,
output_tokens: 50,
..Default::default()
}
);
});
}
#[gpui::test]
async fn test_truncate_second_message(cx: &mut TestAppContext) {
let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
crates/agent/src/thread.rs
+155 -2
View File
@@ -1111,8 +1111,11 @@ pub struct Thread {
pending_message: Option<AgentMessage>,
pub(crate) tools: BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,
request_token_usage: HashMap<UserMessageId, language_model::TokenUsage>,
#[allow(unused)]
cumulative_token_usage: TokenUsage,
/// The per-field maximum usage snapshot already added to
/// `cumulative_token_usage` for the in-flight completion request. Reset at
/// the start of each request.
current_request_token_usage: TokenUsage,
#[allow(unused)]
initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
pub(crate) context_server_registry: Entity<ContextServerRegistry>,
@@ -1245,6 +1248,7 @@ impl Thread {
tools: BTreeMap::default(),
request_token_usage: HashMap::default(),
cumulative_token_usage: TokenUsage::default(),
current_request_token_usage: TokenUsage::default(),
initial_project_snapshot: {
let project_snapshot = Self::project_snapshot(project.clone(), cx);
cx.foreground_executor()
@@ -1611,6 +1615,7 @@ impl Thread {
tools: BTreeMap::default(),
request_token_usage: db_thread.request_token_usage.clone(),
cumulative_token_usage: db_thread.cumulative_token_usage,
current_request_token_usage: TokenUsage::default(),
initial_project_snapshot: Task::ready(db_thread.initial_project_snapshot).shared(),
context_server_registry,
profile_id,
@@ -1991,7 +1996,43 @@ impl Thread {
self.has_queued_message
}
fn accumulate_token_usage(&mut self, update: language_model::TokenUsage) {
let previous_accounted_usage = self.current_request_token_usage;
let current_accounted_usage = TokenUsage {
input_tokens: previous_accounted_usage
.input_tokens
.max(update.input_tokens),
output_tokens: previous_accounted_usage
.output_tokens
.max(update.output_tokens),
cache_creation_input_tokens: previous_accounted_usage
.cache_creation_input_tokens
.max(update.cache_creation_input_tokens),
cache_read_input_tokens: previous_accounted_usage
.cache_read_input_tokens
.max(update.cache_read_input_tokens),
};
self.current_request_token_usage = current_accounted_usage;
self.cumulative_token_usage = self.cumulative_token_usage
+ TokenUsage {
input_tokens: current_accounted_usage
.input_tokens
.saturating_sub(previous_accounted_usage.input_tokens),
output_tokens: current_accounted_usage
.output_tokens
.saturating_sub(previous_accounted_usage.output_tokens),
cache_creation_input_tokens: current_accounted_usage
.cache_creation_input_tokens
.saturating_sub(previous_accounted_usage.cache_creation_input_tokens),
cache_read_input_tokens: current_accounted_usage
.cache_read_input_tokens
.saturating_sub(previous_accounted_usage.cache_read_input_tokens),
};
}
fn update_token_usage(&mut self, update: language_model::TokenUsage, cx: &mut Context<Self>) {
self.accumulate_token_usage(update);
let Some(last_user_message) = self.last_user_message() else {
return;
};
@@ -2032,6 +2073,10 @@ impl Thread {
Some(*tokens)
}
pub fn cumulative_token_usage(&self) -> language_model::TokenUsage {
self.cumulative_token_usage
}
pub fn latest_token_usage(&self) -> Option<acp_thread::TokenUsage> {
let usage = self.latest_request_token_usage()?;
let model = self.model.clone()?;
@@ -2301,6 +2346,7 @@ impl Thread {
.ok_or_else(|| anyhow!(NoModelConfiguredError))?;
this.refresh_turn_tools(cx);
let request = this.build_completion_request(intent, cx)?;
this.current_request_token_usage = TokenUsage::default();
anyhow::Ok((model, request))
})??;
@@ -2532,6 +2578,7 @@ impl Thread {
};
let model = this.model.clone()?;
let request = this.build_compaction_request(insertion_ix, &model, cx);
this.current_request_token_usage = TokenUsage::default();
Some((model, request, insertion_ix))
})?
else {
@@ -2575,10 +2622,14 @@ impl Thread {
summary.push_str(&text);
event_stream.send_context_compaction_update(compaction_id.clone(), &text);
}
LanguageModelCompletionEvent::UsageUpdate(usage) => {
this.update(cx, |this, _cx| {
this.accumulate_token_usage(usage);
})?;
}
LanguageModelCompletionEvent::Stop(_)
| LanguageModelCompletionEvent::Started
| LanguageModelCompletionEvent::Queued { .. }
| LanguageModelCompletionEvent::UsageUpdate(_)
| LanguageModelCompletionEvent::Thinking { .. }
| LanguageModelCompletionEvent::RedactedThinking { .. }
| LanguageModelCompletionEvent::ReasoningDetails(_)
@@ -5646,6 +5697,108 @@ mod tests {
});
}
#[gpui::test]
async fn test_compaction_usage_counts_toward_cumulative_usage(cx: &mut TestAppContext) {
let (thread, _event_stream) = setup_thread_for_test(cx).await;
let model = Arc::new(FakeLanguageModel::default());
let old_user_message_id = UserMessageId::new();
let new_user_message_id = UserMessageId::new();
let prior_usage = TokenUsage {
input_tokens: 960_000,
output_tokens: 25,
..Default::default()
};
let compaction_usage = TokenUsage {
input_tokens: 40,
output_tokens: 9,
cache_creation_input_tokens: 2,
cache_read_input_tokens: 3,
};
let final_usage = TokenUsage {
input_tokens: 500,
output_tokens: 50,
..Default::default()
};
cx.update(|cx| {
cx.update_flags(true, vec!["handoff".to_string()]);
thread.update(cx, |thread, cx| {
thread.set_model(model.clone(), cx);
thread
.messages
.push(user_text_message(old_user_message_id.clone(), "old user"));
thread.messages.push(agent_text_message("old assistant"));
thread
.request_token_usage
.insert(old_user_message_id.clone(), prior_usage);
thread.cumulative_token_usage = prior_usage;
thread.current_request_token_usage = prior_usage;
});
});
let _events = cx
.update(|cx| {
thread.update(cx, |thread, cx| {
thread.send(new_user_message_id.clone(), vec!["new prompt"], cx)
})
})
.unwrap();
cx.run_until_parked();
let compaction_request = model.pending_completions().pop().unwrap();
assert_eq!(
compaction_request.intent,
Some(CompletionIntent::ThreadContextSummarization)
);
model.send_completion_stream_event(
&compaction_request,
LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
input_tokens: 40,
output_tokens: 4,
..Default::default()
}),
);
model.send_completion_stream_event(
&compaction_request,
LanguageModelCompletionEvent::UsageUpdate(compaction_usage),
);
model.send_completion_stream_text_chunk(&compaction_request, "compacted old context");
model.end_completion_stream(&compaction_request);
cx.run_until_parked();
let expected_after_compaction = prior_usage + compaction_usage;
thread.read_with(cx, |thread, _cx| {
assert_eq!(thread.cumulative_token_usage(), expected_after_compaction);
assert!(
!thread
.request_token_usage
.contains_key(&new_user_message_id)
);
});
let final_request = model.pending_completions().pop().unwrap();
assert_eq!(final_request.intent, Some(CompletionIntent::UserPrompt));
model.send_completion_stream_event(
&final_request,
LanguageModelCompletionEvent::UsageUpdate(final_usage),
);
model.end_completion_stream(&final_request);
cx.run_until_parked();
thread.read_with(cx, |thread, _cx| {
assert_eq!(
thread.cumulative_token_usage(),
expected_after_compaction + final_usage
);
assert_eq!(
thread.request_token_usage.get(&new_user_message_id),
Some(&final_usage)
);
});
}
#[gpui::test]
async fn test_replay_emits_context_compaction(cx: &mut TestAppContext) {
let (thread, _event_stream) = setup_thread_for_test(cx).await;