Track cumulative token usage in assistant2 when using anthropic API (#26738)

Release Notes:

- N/A

crates/assistant2/src/thread.rs

@@ -11,7 +11,7 @@ use language_model::{
     LanguageModel, LanguageModelCompletionEvent, LanguageModelRegistry, LanguageModelRequest,
     LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
     LanguageModelToolUseId, MaxMonthlySpendReachedError, MessageContent, PaymentRequiredError,
-    Role, StopReason,
+    Role, StopReason, TokenUsage,
 };
 use project::Project;
 use prompt_store::{AssistantSystemPromptWorktree, PromptBuilder};
@@ -81,6 +81,7 @@ pub struct Thread {
     tool_use: ToolUseState,
     scripting_session: Entity<ScriptingSession>,
     scripting_tool_use: ToolUseState,
+    cumulative_token_usage: TokenUsage,
 }
 
 impl Thread {
@@ -109,6 +110,7 @@ impl Thread {
             tool_use: ToolUseState::new(),
             scripting_session,
             scripting_tool_use: ToolUseState::new(),
+            cumulative_token_usage: TokenUsage::default(),
         }
     }
 
@@ -158,6 +160,8 @@ impl Thread {
             tool_use,
             scripting_session,
             scripting_tool_use,
+            // TODO: persist token usage?
+            cumulative_token_usage: TokenUsage::default(),
         }
     }
 
@@ -490,6 +494,7 @@ impl Thread {
             let stream_completion = async {
                 let mut events = stream.await?;
                 let mut stop_reason = StopReason::EndTurn;
+                let mut current_token_usage = TokenUsage::default();
 
                 while let Some(event) = events.next().await {
                     let event = event?;
@@ -502,6 +507,12 @@ impl Thread {
                             LanguageModelCompletionEvent::Stop(reason) => {
                                 stop_reason = reason;
                             }
+                            LanguageModelCompletionEvent::UsageUpdate(token_usage) => {
+                                thread.cumulative_token_usage =
+                                    thread.cumulative_token_usage.clone() + token_usage.clone()
+                                        - current_token_usage.clone();
+                                current_token_usage = token_usage;
+                            }
                             LanguageModelCompletionEvent::Text(chunk) => {
                                 if let Some(last_message) = thread.messages.last_mut() {
                                     if last_message.role == Role::Assistant {
@@ -843,6 +854,10 @@ impl Thread {
 
         Ok(String::from_utf8_lossy(&markdown).to_string())
     }
+
+    pub fn cumulative_token_usage(&self) -> TokenUsage {
+        self.cumulative_token_usage.clone()
+    }
 }
 
 #[derive(Debug, Clone)]