From 699f58aeba56b10e99f789f9fb492c76fbeea81b Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Wed, 20 Aug 2025 18:04:32 +0200
Subject: [PATCH] Capture telemetry when requesting completions in agent2
 (#36600)

Release Notes:

- N/A
---
 Cargo.lock                  |  1 +
 crates/agent2/Cargo.toml    |  1 +
 crates/agent2/src/thread.rs | 26 ++++++++++++++++++++++++--
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fdc858ef50..342bb1058f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -239,6 +239,7 @@ dependencies = [
  "smol",
  "sqlez",
  "task",
+ "telemetry",
  "tempfile",
  "terminal",
  "text",
diff --git a/crates/agent2/Cargo.toml b/crates/agent2/Cargo.toml
index bc32a79622..2a5d879e9e 100644
--- a/crates/agent2/Cargo.toml
+++ b/crates/agent2/Cargo.toml
@@ -58,6 +58,7 @@ settings.workspace = true
 smol.workspace = true
 sqlez.workspace = true
 task.workspace = true
+telemetry.workspace = true
 terminal.workspace = true
 text.workspace = true
 ui.workspace = true
diff --git a/crates/agent2/src/thread.rs b/crates/agent2/src/thread.rs
index c7b1a08b92..f407ee7de5 100644
--- a/crates/agent2/src/thread.rs
+++ b/crates/agent2/src/thread.rs
@@ -1195,6 +1195,15 @@ impl Thread {
 
         let mut attempt = None;
         'retry: loop {
+            telemetry::event!(
+                "Agent Thread Completion",
+                thread_id = this.read_with(cx, |this, _| this.id.to_string())?,
+                prompt_id = this.read_with(cx, |this, _| this.prompt_id.to_string())?,
+                model = model.telemetry_id(),
+                model_provider = model.provider_id().to_string(),
+                attempt
+            );
+
             let mut events = model.stream_completion(request.clone(), cx).await?;
             let mut tool_uses = FuturesUnordered::new();
             while let Some(event) = events.next().await {
@@ -1211,8 +1220,21 @@ impl Thread {
                             this.update_model_request_usage(amount, limit, cx)
                         })?;
                     }
-                    Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
-                        this.update(cx, |this, cx| this.update_token_usage(token_usage, cx))?;
+                    Ok(LanguageModelCompletionEvent::UsageUpdate(usage)) => {
+                        telemetry::event!(
+                            "Agent Thread Completion Usage Updated",
+                            thread_id = this.read_with(cx, |this, _| this.id.to_string())?,
+                            prompt_id = this.read_with(cx, |this, _| this.prompt_id.to_string())?,
+                            model = model.telemetry_id(),
+                            model_provider = model.provider_id().to_string(),
+                            attempt,
+                            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,
+                        );
+
+                        this.update(cx, |this, cx| this.update_token_usage(usage, cx))?;
                     }
                     Ok(LanguageModelCompletionEvent::Stop(StopReason::Refusal)) => {
                         *refusal = true;