Capture telemetry when requesting completions in agent2 (#36600)

Release Notes:

- N/A
This commit is contained in:
Antonio Scandurra 2025-08-20 18:04:32 +02:00 committed by GitHub
parent 1e6cefaa56
commit 699f58aeba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 26 additions and 2 deletions

1
Cargo.lock generated
View file

@ -239,6 +239,7 @@ dependencies = [
"smol", "smol",
"sqlez", "sqlez",
"task", "task",
"telemetry",
"tempfile", "tempfile",
"terminal", "terminal",
"text", "text",

View file

@ -58,6 +58,7 @@ settings.workspace = true
smol.workspace = true smol.workspace = true
sqlez.workspace = true sqlez.workspace = true
task.workspace = true task.workspace = true
telemetry.workspace = true
terminal.workspace = true terminal.workspace = true
text.workspace = true text.workspace = true
ui.workspace = true ui.workspace = true

View file

@ -1195,6 +1195,15 @@ impl Thread {
let mut attempt = None; let mut attempt = None;
'retry: loop { '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 events = model.stream_completion(request.clone(), cx).await?;
let mut tool_uses = FuturesUnordered::new(); let mut tool_uses = FuturesUnordered::new();
while let Some(event) = events.next().await { while let Some(event) = events.next().await {
@ -1211,8 +1220,21 @@ impl Thread {
this.update_model_request_usage(amount, limit, cx) this.update_model_request_usage(amount, limit, cx)
})?; })?;
} }
Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => { Ok(LanguageModelCompletionEvent::UsageUpdate(usage)) => {
this.update(cx, |this, cx| this.update_token_usage(token_usage, cx))?; 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)) => { Ok(LanguageModelCompletionEvent::Stop(StopReason::Refusal)) => {
*refusal = true; *refusal = true;