From 5f2eba257d7cfa4db56282003930592ac2e122ff Mon Sep 17 00:00:00 2001 From: Anant Goel Date: Thu, 4 Jun 2026 11:24:02 -0400 Subject: [PATCH] agent: Accumulate and persist cumulative token usage for threads (#58378) `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 --- crates/agent/src/tests/mod.rs | 140 ++++++++++++++++++++++++++++++ crates/agent/src/thread.rs | 157 +++++++++++++++++++++++++++++++++- 2 files changed, 295 insertions(+), 2 deletions(-) diff --git a/crates/agent/src/tests/mod.rs b/crates/agent/src/tests/mod.rs index f82c4c7ddd..758a4f51cb 100644 --- a/crates/agent/src/tests/mod.rs +++ b/crates/agent/src/tests/mod.rs @@ -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; diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index d742750b22..68b6bf03ab 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -1111,8 +1111,11 @@ pub struct Thread { pending_message: Option, pub(crate) tools: BTreeMap>, request_token_usage: HashMap, - #[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>>>, pub(crate) context_server_registry: Entity, @@ -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.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 { 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;