From 33abf1ee7cac2f0a91b92375cf40142954ea32a5 Mon Sep 17 00:00:00 2001 From: Michael Sloan Date: Tue, 29 Apr 2025 15:55:54 -0600 Subject: [PATCH] agent: Log errors from token counting APIs (#29626) Release Notes: - N/A --- crates/agent/src/active_thread.rs | 95 ++++++++++++++------------ crates/agent/src/message_editor.rs | 105 +++++++++++++++-------------- 2 files changed, 107 insertions(+), 93 deletions(-) diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index 8aab7dd0f7849cf2a4df5c9a5a57948c0b305a2f..7d112169315b0bb996e7f99612002e7386bdc4e9 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -702,7 +702,7 @@ struct EditMessageState { editor: Entity, last_estimated_token_count: Option, _subscription: Subscription, - _update_token_count_task: Option>>, + _update_token_count_task: Option>, } 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(); + }; })); } diff --git a/crates/agent/src/message_editor.rs b/crates/agent/src/message_editor.rs index 6ab4e899bc32267ea921363c4e5dc1b786c1f1cb..e223b5821f652b5e7f9932da87e3fccd51249115 100644 --- a/crates/agent/src/message_editor.rs +++ b/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, - update_token_count_task: Option>>, + update_token_count_task: Option>, _subscriptions: Vec, } @@ -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(); + } })); } }