Record token usage telemetry (#26962)

Richard Feldman created 9 months ago

<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

Change summary

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

Detailed changes

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();
         });