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

Michael Sloan created

Release Notes:

- N/A

Change summary

crates/agent/src/active_thread.rs  |  95 +++++++++++++++-------------
crates/agent/src/message_editor.rs | 105 +++++++++++++++++--------------
2 files changed, 107 insertions(+), 93 deletions(-)

Detailed changes

crates/agent/src/active_thread.rs 🔗

@@ -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,58 +1268,65 @@ impl ActiveThread {
                     .await;
             }
 
-            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;
-                };
-                let message_text = editor.read(cx).text(cx);
-
-                if message_text.is_empty() && message.loaded_context.is_empty() {
-                    return None;
-                }
+            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;
+                    };
+                    let message_text = editor.read(cx).text(cx);
 
-                let mut request_message = LanguageModelRequestMessage {
-                    role: language_model::Role::User,
-                    content: Vec::new(),
-                    cache: false,
-                };
+                    if message_text.is_empty() && message.loaded_context.is_empty() {
+                        return None;
+                    }
 
-                message
-                    .loaded_context
-                    .add_to_request_message(&mut request_message);
+                    let mut request_message = LanguageModelRequestMessage {
+                        role: language_model::Role::User,
+                        content: Vec::new(),
+                        cache: false,
+                    };
 
-                if !message_text.is_empty() {
-                    request_message
-                        .content
-                        .push(MessageContent::Text(message_text));
-                }
+                    message
+                        .loaded_context
+                        .add_to_request_message(&mut request_message);
 
-                let request = language_model::LanguageModelRequest {
-                    thread_id: None,
-                    prompt_id: None,
-                    mode: None,
-                    messages: vec![request_message],
-                    tools: vec![],
-                    stop: vec![],
-                    temperature: None,
-                };
+                    if !message_text.is_empty() {
+                        request_message
+                            .content
+                            .push(MessageContent::Text(message_text));
+                    }
 
-                Some(configured_model.model.count_tokens(request, cx))
-            })? {
-                task.await?
+                    let request = language_model::LanguageModelRequest {
+                        thread_id: None,
+                        prompt_id: None,
+                        mode: None,
+                        messages: vec![request_message],
+                        tools: vec![],
+                        stop: vec![],
+                        temperature: None,
+                    };
+
+                    Some(configured_model.model.count_tokens(request, cx))
+                })
+                .ok()
+                .flatten()
+            {
+                task.await.log_err()
             } else {
-                0
+                Some(0)
             };
 
-            this.update(cx, |this, cx| {
-                let Some((_message_id, state)) = this.editing_message.as_mut() else {
-                    return;
-                };
+            if let Some(token_count) = token_count {
+                this.update(cx, |this, cx| {
+                    let Some((_message_id, state)) = this.editing_message.as_mut() else {
+                        return;
+                    };
 
-                state.last_estimated_token_count = Some(token_count);
-                cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged);
-            })
+                    state.last_estimated_token_count = Some(token_count);
+                    cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged);
+                })
+                .ok();
+            };
         }));
     }
 

crates/agent/src/message_editor.rs 🔗

@@ -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,57 +1088,64 @@ impl MessageEditor {
                     .await;
             }
 
-            let token_count = if let Some(task) = this.update(cx, |this, cx| {
-                let loaded_context = this
-                    .last_loaded_context
-                    .as_ref()
-                    .map(|context_load_result| &context_load_result.loaded_context);
-                let message_text = editor.read(cx).text(cx);
-
-                if message_text.is_empty()
-                    && loaded_context.map_or(true, |loaded_context| loaded_context.is_empty())
-                {
-                    return None;
-                }
-
-                let mut request_message = LanguageModelRequestMessage {
-                    role: language_model::Role::User,
-                    content: Vec::new(),
-                    cache: false,
-                };
-
-                if let Some(loaded_context) = loaded_context {
-                    loaded_context.add_to_request_message(&mut request_message);
-                }
-
-                if !message_text.is_empty() {
-                    request_message
-                        .content
-                        .push(MessageContent::Text(message_text));
-                }
-
-                let request = language_model::LanguageModelRequest {
-                    thread_id: None,
-                    prompt_id: None,
-                    mode: None,
-                    messages: vec![request_message],
-                    tools: vec![],
-                    stop: vec![],
-                    temperature: None,
-                };
-
-                Some(model.model.count_tokens(request, cx))
-            })? {
-                task.await?
+            let token_count = if let Some(task) = this
+                .update(cx, |this, cx| {
+                    let loaded_context = this
+                        .last_loaded_context
+                        .as_ref()
+                        .map(|context_load_result| &context_load_result.loaded_context);
+                    let message_text = editor.read(cx).text(cx);
+
+                    if message_text.is_empty()
+                        && loaded_context.map_or(true, |loaded_context| loaded_context.is_empty())
+                    {
+                        return None;
+                    }
+
+                    let mut request_message = LanguageModelRequestMessage {
+                        role: language_model::Role::User,
+                        content: Vec::new(),
+                        cache: false,
+                    };
+
+                    if let Some(loaded_context) = loaded_context {
+                        loaded_context.add_to_request_message(&mut request_message);
+                    }
+
+                    if !message_text.is_empty() {
+                        request_message
+                            .content
+                            .push(MessageContent::Text(message_text));
+                    }
+
+                    let request = language_model::LanguageModelRequest {
+                        thread_id: None,
+                        prompt_id: None,
+                        mode: None,
+                        messages: vec![request_message],
+                        tools: vec![],
+                        stop: vec![],
+                        temperature: None,
+                    };
+
+                    Some(model.model.count_tokens(request, cx))
+                })
+                .ok()
+                .flatten()
+            {
+                task.await.log_err()
             } else {
-                0
+                Some(0)
             };
 
-            this.update(cx, |this, cx| {
-                this.last_estimated_token_count = Some(token_count);
-                cx.emit(MessageEditorEvent::EstimatedTokenCount);
-                this.update_token_count_task.take();
-            })
+            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();
+            }
         }));
     }
 }