From 63037513258977e5da68ca8f63f64db3f16abc27 Mon Sep 17 00:00:00 2001
From: Richard Feldman <oss@rtfeldman.com>
Date: Wed, 19 Mar 2025 10:47:46 -0400
Subject: [PATCH] Record token usage telemetry (#26962)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

<img width="1103" alt="Screenshot 2025-03-17 at 9 47 32 PM"
src="https://github.com/user-attachments/assets/947cf33d-4464-4305-8ff0-3630529d2f81"
/>


Release Notes:

- N/A
---
 crates/assistant2/src/thread.rs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs
index a89dba5660fba2336093eb2cb91f1517c0f035dd..f241b4a571b279222db23d88c4c7f0de910b6700 100644
--- a/crates/assistant2/src/thread.rs
+++ b/crates/assistant2/src/thread.rs
@@ -604,6 +604,8 @@ impl Thread {
 
         let task = cx.spawn(async move |thread, cx| {
             let stream = model.stream_completion(request, &cx);
+            let initial_token_usage =
+                thread.read_with(&cx, |thread, _cx| thread.cumulative_token_usage.clone());
             let stream_completion = async {
                 let mut events = stream.await?;
                 let mut stop_reason = StopReason::EndTurn;
@@ -718,6 +720,21 @@ impl Thread {
                         }
                     }
                     cx.emit(ThreadEvent::DoneStreaming);
+
+                    if let Ok(initial_usage) = initial_token_usage {
+                        let usage = thread.cumulative_token_usage.clone() - initial_usage;
+
+                        telemetry::event!(
+                            "Assistant Thread Completion",
+                            thread_id = thread.id().to_string(),
+                            model = model.telemetry_id(),
+                            model_provider = model.provider_id().to_string(),
+                            input_tokens = usage.input_tokens,
+                            output_tokens = usage.output_tokens,
+                            cache_creation_input_tokens = usage.cache_creation_input_tokens,
+                            cache_read_input_tokens = usage.cache_read_input_tokens,
+                        );
+                    }
                 })
                 .ok();
         });