Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

agent: Log errors from token counting APIs (#29626)

Browse source 
Release Notes:

- N/A
This commit is contained in:
Michael Sloan 2025-04-29 15:55:54 -06:00 • committed by GitHub
parent f7a3e00074
commit 33abf1ee7c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 100 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

17
crates/agent/src/active_thread.rs
View file

@ -702,7 +702,7 @@ struct EditMessageState {
editor: Entity<Editor>,
last_estimated_token_count: Option<usize>,
_subscription: Subscription,
_update_token_count_task: Option<Task<anyhow::Result<()>>>,
_update_token_count_task: Option<Task<()>>,
}
impl ActiveThread {
@ -1268,7 +1268,8 @@ impl ActiveThread {
.await;
}
let token_count = if let Some(task) = cx.update(|cx| {
let token_count = if let Some(task) = cx
.update(|cx| {
let Some(message) = thread.read(cx).message(message_id) else {
log::error!("Message that was being edited no longer exists");
return None;
@ -1306,12 +1307,16 @@ impl ActiveThread {
};
Some(configured_model.model.count_tokens(request, cx))
})? {
task.await?
})
.ok()
.flatten()
{
task.await.log_err()
} else {
0
Some(0)
};
if let Some(token_count) = token_count {
this.update(cx, |this, cx| {
let Some((_message_id, state)) = this.editing_message.as_mut() else {
return;
@ -1320,6 +1325,8 @@ impl ActiveThread {
state.last_estimated_token_count = Some(token_count);
cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged);
})
.ok();
};
}));
}

17
crates/agent/src/message_editor.rs
View file

@ -63,7 +63,7 @@ pub struct MessageEditor {
edits_expanded: bool,
editor_is_expanded: bool,
last_estimated_token_count: Option<usize>,
update_token_count_task: Option<Task<anyhow::Result<()>>>,
update_token_count_task: Option<Task<()>>,
_subscriptions: Vec<Subscription>,
}
@ -1088,7 +1088,8 @@ impl MessageEditor {
.await;
}
let token_count = if let Some(task) = this.update(cx, |this, cx| {
let token_count = if let Some(task) = this
.update(cx, |this, cx| {
let loaded_context = this
.last_loaded_context
.as_ref()
@ -1128,17 +1129,23 @@ impl MessageEditor {
};
Some(model.model.count_tokens(request, cx))
})? {
task.await?
})
.ok()
.flatten()
{
task.await.log_err()
} else {
0
Some(0)
};
if let Some(token_count) = token_count {
this.update(cx, |this, cx| {
this.last_estimated_token_count = Some(token_count);
cx.emit(MessageEditorEvent::EstimatedTokenCount);
this.update_token_count_task.take();
})
.ok();
}
}));
}
}