From 8bc59ab075284f83c5e0f50ecd0580df601dda5f Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Wed, 7 May 2025 10:10:30 +0200 Subject: [PATCH 1/4] WIP --- crates/agent/src/active_thread.rs | 849 ++++++++++++------------------ 1 file changed, 335 insertions(+), 514 deletions(-) diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index e22726f56a70f5..88fd3ad4512071 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -1,4 +1,3 @@ -use crate::AssistantPanel; use crate::context::{AgentContextHandle, RULES_ICON}; use crate::context_picker::{ContextPicker, MentionLink}; use crate::context_store::ContextStore; @@ -12,10 +11,11 @@ use crate::tool_use::{PendingToolUseStatus, ToolUse}; use crate::ui::{ AddedContext, AgentNotification, AgentNotificationEvent, AnimatedLabel, ContextPill, }; +use crate::{AssistantPanel, ToggleContextPicker}; use anyhow::Context as _; use assistant_settings::{AssistantSettings, NotifyWhenAgentWaiting}; use assistant_tool::ToolUseStatus; -use collections::{HashMap, HashSet}; +use collections::{HashMap, HashSet, hash_map::Entry}; use editor::actions::{MoveUp, Paste}; use editor::scroll::Autoscroll; use editor::{Editor, EditorElement, EditorEvent, EditorStyle, MultiBuffer}; @@ -67,7 +67,8 @@ pub struct ActiveThread { hide_scrollbar_task: Option>, rendered_messages_by_id: HashMap, rendered_tool_uses: HashMap, - editing_message: Option<(MessageId, EditingMessageState)>, + user_message_views: HashMap, + focused_user_message: Option, expanded_tool_uses: HashMap, expanded_thinking_segments: HashMap<(MessageId, usize), bool>, expanded_code_blocks: HashMap<(MessageId, usize), bool>, @@ -740,7 +741,7 @@ fn open_markdown_link( } } -struct EditingMessageState { +struct UserMessageView { editor: Entity, context_strip: Entity, context_picker_menu_handle: PopoverMenuHandle, @@ -792,7 +793,8 @@ impl ActiveThread { scrollbar_state: ScrollbarState::new(list_state), show_scrollbar: false, hide_scrollbar_task: None, - editing_message: None, + user_message_views: HashMap::default(), + focused_user_message: None, last_error: None, copied_code_block_ids: HashSet::default(), notifications: Vec::new(), @@ -852,9 +854,9 @@ impl ActiveThread { /// Returns the editing message id and the estimated token count in the content pub fn editing_message_id(&self) -> Option<(MessageId, usize)> { - self.editing_message - .as_ref() - .map(|(id, state)| (*id, state.last_estimated_token_count.unwrap_or(0))) + let message_id = self.focused_user_message?; + let user_message_editor = self.user_message_views.get(&message_id)?; + Some((message_id, user_message_editor.last_estimated_token_count?)) } pub fn context_store(&self) -> &Entity { @@ -960,7 +962,7 @@ impl ActiveThread { ) { match event { ThreadEvent::CancelEditing => { - if self.editing_message.is_some() { + if self.focused_user_message.is_some() { self.cancel_editing_message(&menu::Cancel, window, cx); } } @@ -1269,195 +1271,94 @@ impl ActiveThread { })); } - fn start_editing_message( - &mut self, - message_id: MessageId, - message_segments: &[MessageSegment], - window: &mut Window, - cx: &mut Context, - ) { - // User message should always consist of a single text segment, - // therefore we can skip returning early if it's not a text segment. - let Some(MessageSegment::Text(message_text)) = message_segments.first() else { - return; - }; - - // Cancel any ongoing streaming when user starts editing a previous message - self.cancel_last_completion(window, cx); - - let editor = crate::message_editor::create_editor( - self.workspace.clone(), - self.context_store.downgrade(), - self.thread_store.downgrade(), - self.text_thread_store.downgrade(), - window, - cx, - ); - editor.update(cx, |editor, cx| { - editor.set_text(message_text.clone(), window, cx); - editor.focus_handle(cx).focus(window); - editor.move_to_end(&editor::actions::MoveToEnd, window, cx); - }); - let buffer_edited_subscription = cx.subscribe(&editor, |this, _, event, cx| match event { - EditorEvent::BufferEdited => { - this.update_editing_message_token_count(true, cx); - } - _ => {} - }); - - let context_picker_menu_handle = PopoverMenuHandle::default(); - let context_strip = cx.new(|cx| { - ContextStrip::new( - self.context_store.clone(), - self.workspace.clone(), - Some(self.thread_store.downgrade()), - Some(self.text_thread_store.downgrade()), - context_picker_menu_handle.clone(), - SuggestContextKind::File, - window, - cx, - ) - }); - - let context_strip_subscription = - cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event); - - self.editing_message = Some(( - message_id, - EditingMessageState { - editor: editor.clone(), - context_strip, - context_picker_menu_handle, - last_estimated_token_count: None, - _subscriptions: [buffer_edited_subscription, context_strip_subscription], - _update_token_count_task: None, - }, - )); - self.update_editing_message_token_count(false, cx); - cx.notify(); - } - - fn handle_context_strip_event( - &mut self, - _context_strip: &Entity, - event: &ContextStripEvent, - window: &mut Window, - cx: &mut Context, - ) { - if let Some((_, state)) = self.editing_message.as_ref() { - match event { - ContextStripEvent::PickerDismissed - | ContextStripEvent::BlurredEmpty - | ContextStripEvent::BlurredDown => { - let editor_focus_handle = state.editor.focus_handle(cx); - window.focus(&editor_focus_handle); - } - ContextStripEvent::BlurredUp => {} - } - } - } - fn update_editing_message_token_count(&mut self, debounce: bool, cx: &mut Context) { - let Some((message_id, state)) = self.editing_message.as_mut() else { - return; - }; - - cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); - state._update_token_count_task.take(); - - let Some(configured_model) = self.thread.read(cx).configured_model() else { - state.last_estimated_token_count.take(); - return; - }; - - let editor = state.editor.clone(); - let thread = self.thread.clone(); - let message_id = *message_id; - - state._update_token_count_task = Some(cx.spawn(async move |this, cx| { - if debounce { - cx.background_executor() - .timer(Duration::from_millis(200)) - .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 mut request_message = LanguageModelRequestMessage { - role: language_model::Role::User, - content: Vec::new(), - cache: false, - }; - - message - .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: AssistantSettings::temperature_for_model( - &configured_model.model, - cx, - ), - }; - - Some(configured_model.model.count_tokens(request, cx)) - }) - .ok() - .flatten() - { - task.await.log_err() - } else { - 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; - }; - - state.last_estimated_token_count = Some(token_count); - cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); - }) - .ok(); - }; - })); - } - - fn toggle_context_picker( - &mut self, - _: &crate::ToggleContextPicker, - window: &mut Window, - cx: &mut Context, - ) { - if let Some((_, state)) = self.editing_message.as_mut() { - let handle = state.context_picker_menu_handle.clone(); - window.defer(cx, move |window, cx| { - handle.toggle(window, cx); - }); - } + // todo!() + // let Some((message_id, state)) = self.user_message_views.as_mut() else { + // return; + // }; + + // cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); + // state._update_token_count_task.take(); + + // let Some(configured_model) = self.thread.read(cx).configured_model() else { + // state.last_estimated_token_count.take(); + // return; + // }; + + // let editor = state.editor.clone(); + // let thread = self.thread.clone(); + // let message_id = *message_id; + + // state._update_token_count_task = Some(cx.spawn(async move |this, cx| { + // if debounce { + // cx.background_executor() + // .timer(Duration::from_millis(200)) + // .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 mut request_message = LanguageModelRequestMessage { + // role: language_model::Role::User, + // content: Vec::new(), + // cache: false, + // }; + + // message + // .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: AssistantSettings::temperature_for_model( + // &configured_model.model, + // cx, + // ), + // }; + + // Some(configured_model.model.count_tokens(request, cx)) + // }) + // .ok() + // .flatten() + // { + // task.await.log_err() + // } else { + // Some(0) + // }; + + // if let Some(token_count) = token_count { + // this.update(cx, |this, cx| { + // let Some((_message_id, state)) = this.user_message_views.as_mut() else { + // return; + // }; + + // state.last_estimated_token_count = Some(token_count); + // cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); + // }) + // .ok(); + // }; + // })); } fn remove_all_context( @@ -1470,14 +1371,15 @@ impl ActiveThread { cx.notify(); } - fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context) { - if let Some((_, state)) = self.editing_message.as_mut() { - if state.context_picker_menu_handle.is_deployed() { - cx.propagate(); - } else { - state.context_strip.focus_handle(cx).focus(window); - } - } + fn move_up(&mut self, _: &MoveUp, _window: &mut Window, _cx: &mut Context) { + // todo!() + // if let Some((_, state)) = self.user_message_views.as_mut() { + // if state.context_picker_menu_handle.is_deployed() { + // cx.propagate(); + // } else { + // state.context_strip.focus_handle(cx).focus(window); + // } + // } } fn paste(&mut self, _: &Paste, _window: &mut Window, cx: &mut Context) { @@ -1509,7 +1411,8 @@ impl ActiveThread { } fn cancel_editing_message(&mut self, _: &menu::Cancel, _: &mut Window, cx: &mut Context) { - self.editing_message.take(); + self.focused_user_message.take(); + // todo!("restore focus to new message editor in parent panel") cx.notify(); } @@ -1519,7 +1422,10 @@ impl ActiveThread { window: &mut Window, cx: &mut Context, ) { - let Some((message_id, state)) = self.editing_message.take() else { + let Some(message_id) = self.focused_user_message else { + return; + }; + let Some(user_message_editor) = self.user_message_views.get(&message_id) else { return; }; @@ -1535,7 +1441,7 @@ impl ActiveThread { return; } - let edited_text = state.editor.read(cx).text(cx); + let edited_text = user_message_editor.editor.read(cx).text(cx); let new_context = self .context_store @@ -1694,12 +1600,87 @@ impl ActiveThread { } } - fn render_edit_message_editor( - &self, - state: &EditingMessageState, - _window: &mut Window, - cx: &Context, - ) -> impl IntoElement { + fn render_user_message( + &mut self, + message_id: MessageId, + window: &mut Window, + cx: &mut Context, + ) -> Option { + // Cancel any ongoing streaming when user starts editing a previous message + // todo!("do this when submitting an edited message") + // self.cancel_last_completion(window, cx); + + let view = match self.user_message_views.entry(message_id) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => { + let message = self.thread.read(cx).message(message_id)?; + + // User message should always consist of a single text segment, + // therefore we can skip returning early if it's not a text segment. + let Some(MessageSegment::Text(message_text)) = message.segments.first() else { + return None; + }; + let message_text = message_text.clone(); + + let editor = crate::message_editor::create_editor( + self.workspace.clone(), + self.context_store.downgrade(), + self.thread_store.downgrade(), + self.text_thread_store.downgrade(), + window, + cx, + ); + editor.update(cx, |editor, cx| { + editor.set_text(message_text.clone(), window, cx); + }); + let buffer_edited_subscription = + cx.subscribe(&editor, |this, _, event, cx| match event { + EditorEvent::BufferEdited => { + this.update_editing_message_token_count(true, cx); + } + _ => {} + }); + + let context_picker_menu_handle = PopoverMenuHandle::default(); + let context_strip = cx.new(|cx| { + ContextStrip::new( + self.context_store.clone(), + self.workspace.clone(), + Some(self.thread_store.downgrade()), + Some(self.text_thread_store.downgrade()), + context_picker_menu_handle.clone(), + SuggestContextKind::File, + window, + cx, + ) + }); + + let editor_focus_handle = editor.focus_handle(cx); + let context_strip_subscription = cx.subscribe_in( + &context_strip, + window, + move |_this, _context_strip, event, window, _cx| match event { + ContextStripEvent::PickerDismissed + | ContextStripEvent::BlurredEmpty + | ContextStripEvent::BlurredDown => window.focus(&editor_focus_handle), + ContextStripEvent::BlurredUp => {} + }, + ); + + entry.insert(UserMessageView { + editor: editor.clone(), + context_strip, + context_picker_menu_handle, + last_estimated_token_count: None, + _subscriptions: [buffer_edited_subscription, context_strip_subscription], + _update_token_count_task: None, + }) + } + }; + + // todo!("on focus call this") + // self.update_editing_message_token_count(false, cx); + let settings = ThemeSettings::get_global(cx); let font_size = TextSize::Small .rems(cx) @@ -1717,33 +1698,45 @@ impl ActiveThread { line_height: line_height.into(), ..Default::default() }; + let context_picker_menu_handle = view.context_picker_menu_handle.clone(); - v_flex() - .key_context("EditMessageEditor") - .on_action(cx.listener(Self::toggle_context_picker)) - .on_action(cx.listener(Self::remove_all_context)) - .on_action(cx.listener(Self::move_up)) - .on_action(cx.listener(Self::cancel_editing_message)) - .on_action(cx.listener(Self::confirm_editing_message)) - .capture_action(cx.listener(Self::paste)) - .min_h_6() - .flex_grow() - .w_full() - .gap_2() - .child(EditorElement::new( - &state.editor, - EditorStyle { - background: colors.editor_background, - local_player: cx.theme().players().local(), - text: text_style, - syntax: cx.theme().syntax().clone(), - ..Default::default() - }, - )) - .child(state.context_strip.clone()) + Some( + v_flex() + .key_context("EditMessageEditor") + .on_action(move |_: &ToggleContextPicker, window, cx| { + let handle = context_picker_menu_handle.clone(); + window.defer(cx, move |window, cx| handle.toggle(window, cx)); + }) + .on_action(cx.listener(Self::remove_all_context)) + .on_action(cx.listener(Self::move_up)) + .on_action(cx.listener(Self::cancel_editing_message)) + .on_action(cx.listener(Self::confirm_editing_message)) + .capture_action(cx.listener(Self::paste)) + .min_h_6() + .flex_grow() + .w_full() + .gap_2() + .child(EditorElement::new( + &view.editor, + EditorStyle { + background: colors.editor_background, + local_player: cx.theme().players().local(), + text: text_style, + syntax: cx.theme().syntax().clone(), + ..Default::default() + }, + )) + // todo!("only show this if it's not empty") + .child(view.context_strip.clone()), + ) } - fn render_message(&self, ix: usize, window: &mut Window, cx: &mut Context) -> AnyElement { + fn render_message( + &mut self, + ix: usize, + window: &mut Window, + cx: &mut Context, + ) -> AnyElement { let message_id = self.messages[ix]; let Some(message) = self.thread.read(cx).message(message_id) else { return Empty.into_any(); @@ -1762,6 +1755,9 @@ impl ActiveThread { .context_for_message(message_id) .map(|context| AddedContext::new_attached(context, cx)) .collect::>(); + let message_is_empty = message.should_display_content(); + let has_content = !message_is_empty || !added_context.is_empty(); + let role = message.role; let tool_uses = thread.tool_uses_for_message(message_id, cx); let has_tool_uses = !tool_uses.is_empty(); @@ -1774,15 +1770,6 @@ impl ActiveThread { let loading_dots = (is_generating_stale && is_last_message) .then(|| AnimatedLabel::new("").size(LabelSize::Small)); - let editing_message_state = self - .editing_message - .as_ref() - .filter(|(id, _)| *id == message_id) - .map(|(_, state)| state); - - let colors = cx.theme().colors(); - let editor_bg_color = colors.editor_background; - let open_as_markdown = IconButton::new(("open-as-markdown", ix), IconName::FileCode) .shape(ui::IconButtonShape::Square) .icon_size(IconSize::XSmall) @@ -1918,50 +1905,53 @@ impl ActiveThread { .into_any_element(), }; - let message_is_empty = message.should_display_content(); - let has_content = !message_is_empty || !added_context.is_empty(); - - let message_content = has_content.then(|| { - if let Some(state) = editing_message_state.as_ref() { - self.render_edit_message_editor(state, window, cx) - .into_any_element() + let message_content = if has_content { + if message.role == Role::User { + self.render_user_message(message_id, window, cx) + .map(|m| m.into_any_element()) } else { - v_flex() - .w_full() - .gap_1() - .when(!message_is_empty, |parent| { - parent.child(div().min_h_6().child(self.render_message_content( - message_id, - rendered_message, - has_tool_uses, - workspace.clone(), - window, - cx, - ))) - }) - .when(!added_context.is_empty(), |parent| { - parent.child(h_flex().flex_wrap().gap_1().children( - added_context.into_iter().map(|added_context| { - let context = added_context.handle.clone(); - ContextPill::added(added_context, false, false, None).on_click( - Rc::new(cx.listener({ - let workspace = workspace.clone(); - move |_, _, window, cx| { - if let Some(workspace) = workspace.upgrade() { - open_context(&context, workspace, window, cx); - cx.notify(); + Some( + v_flex() + .w_full() + .gap_1() + .when(!message_is_empty, |parent| { + parent.child(div().min_h_6().child(self.render_message_content( + message_id, + rendered_message, + has_tool_uses, + workspace.clone(), + window, + cx, + ))) + }) + .when(!added_context.is_empty(), |parent| { + parent.child(h_flex().flex_wrap().gap_1().children( + added_context.into_iter().map(|added_context| { + let context = added_context.handle.clone(); + ContextPill::added(added_context, false, false, None).on_click( + Rc::new(cx.listener({ + let workspace = workspace.clone(); + move |_, _, window, cx| { + if let Some(workspace) = workspace.upgrade() { + open_context(&context, workspace, window, cx); + cx.notify(); + } } - } - })), - ) - }), - )) - }) - .into_any_element() + })), + ) + }), + )) + }) + .into_any_element(), + ) } - }); + } else { + None + }; - let styled_message = match message.role { + let colors = cx.theme().colors(); + let editor_bg_color = colors.editor_background; + let styled_message = match role { Role::User => v_flex() .id(("message-container", ix)) .pt_2() @@ -1979,77 +1969,60 @@ impl ActiveThread { .hover(|hover| hover.border_color(colors.text_accent.opacity(0.5))) .cursor_pointer() .child( - h_flex() - .p_2p5() - .gap_1() - .children(message_content) - .when_some(editing_message_state, |this, state| { - let focus_handle = state.editor.focus_handle(cx).clone(); - this.w_full().justify_between().child( - h_flex() - .gap_0p5() - .child( - IconButton::new( - "cancel-edit-message", - IconName::Close, - ) - .shape(ui::IconButtonShape::Square) - .icon_color(Color::Error) - .tooltip({ - let focus_handle = focus_handle.clone(); - move |window, cx| { - Tooltip::for_action_in( - "Cancel Edit", - &menu::Cancel, - &focus_handle, - window, - cx, - ) - } - }) - .on_click(cx.listener(Self::handle_cancel_click)), - ) - .child( - IconButton::new( - "confirm-edit-message", - IconName::Check, - ) - .disabled(state.editor.read(cx).is_empty(cx)) - .shape(ui::IconButtonShape::Square) - .icon_color(Color::Success) - .tooltip({ - let focus_handle = focus_handle.clone(); - move |window, cx| { - Tooltip::for_action_in( - "Regenerate", - &menu::Confirm, - &focus_handle, - window, - cx, - ) - } - }) - .on_click( - cx.listener(Self::handle_regenerate_click), - ), - ), - ) - }), - ) - .when(editing_message_state.is_none(), |this| { - this.tooltip(Tooltip::text("Click To Edit")) - }) - .on_click(cx.listener({ - let message_segments = message.segments.clone(); - move |this, _, window, cx| { - this.start_editing_message( - message_id, - &message_segments, - window, - cx, - ); - } - })), + h_flex().p_2p5().gap_1().children(message_content), // todo!("handle this in render_user_message") + // .when_some(editing_message_state, |this, state| { + // let focus_handle = state.editor.focus_handle(cx).clone(); + // this.w_full().justify_between().child( + // h_flex() + // .gap_0p5() + // .child( + // IconButton::new( + // "cancel-edit-message", + // IconName::Close, + // ) + // .shape(ui::IconButtonShape::Square) + // .icon_color(Color::Error) + // .tooltip({ + // let focus_handle = focus_handle.clone(); + // move |window, cx| { + // Tooltip::for_action_in( + // "Cancel Edit", + // &menu::Cancel, + // &focus_handle, + // window, + // cx, + // ) + // } + // }) + // .on_click(cx.listener(Self::handle_cancel_click)), + // ) + // .child( + // IconButton::new( + // "confirm-edit-message", + // IconName::Check, + // ) + // .disabled(state.editor.read(cx).is_empty(cx)) + // .shape(ui::IconButtonShape::Square) + // .icon_color(Color::Success) + // .tooltip({ + // let focus_handle = focus_handle.clone(); + // move |window, cx| { + // Tooltip::for_action_in( + // "Regenerate", + // &menu::Confirm, + // &focus_handle, + // window, + // cx, + // ) + // } + // }) + // .on_click( + // cx.listener(Self::handle_regenerate_click), + // ), + // ), + // ) + // }), + ), ), Role::Assistant => v_flex() .id(("message-container", ix)) @@ -2069,12 +2042,9 @@ impl ActiveThread { ), }; - let after_editing_message = self - .editing_message - .as_ref() - .map_or(false, |(editing_message_id, _)| { - message_id > *editing_message_id - }); + let below_focused_user_message = self + .focused_user_message + .map_or(false, |focused_message_id| message_id > focused_message_id); let panel_background = cx.theme().colors().panel_background; @@ -2245,7 +2215,7 @@ impl ActiveThread { }, ) }) - .when(after_editing_message, |parent| { + .when(below_focused_user_message, |parent| { // Backdrop to dim out the whole thread below the editing user message parent.relative().child( div() @@ -3583,152 +3553,3 @@ fn open_editor_at_position( } }) } - -#[cfg(test)] -mod tests { - use assistant_tool::{ToolRegistry, ToolWorkingSet}; - use editor::EditorSettings; - use fs::FakeFs; - use gpui::{TestAppContext, VisualTestContext}; - use language_model::{LanguageModel, fake_provider::FakeLanguageModel}; - use project::Project; - use prompt_store::PromptBuilder; - use serde_json::json; - use settings::SettingsStore; - use util::path; - - use crate::{ContextLoadResult, thread_store}; - - use super::*; - - #[gpui::test] - async fn test_current_completion_cancelled_when_message_edited(cx: &mut TestAppContext) { - init_test_settings(cx); - - let project = create_test_project( - cx, - json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}), - ) - .await; - - let (cx, active_thread, thread, model) = setup_test_environment(cx, project.clone()).await; - - // Insert user message without any context (empty context vector) - let message = thread.update(cx, |thread, cx| { - let message_id = thread.insert_user_message( - "What is the best way to learn Rust?", - ContextLoadResult::default(), - None, - vec![], - cx, - ); - thread - .message(message_id) - .expect("message should exist") - .clone() - }); - - // Stream response to user message - thread.update(cx, |thread, cx| { - let request = thread.to_completion_request(model.clone(), cx); - thread.stream_completion(request, model, cx.active_window(), cx) - }); - let generating = thread.update(cx, |thread, _cx| thread.is_generating()); - assert!(generating, "There should be one pending completion"); - - // Edit the previous message - active_thread.update_in(cx, |active_thread, window, cx| { - active_thread.start_editing_message(message.id, &message.segments, window, cx); - }); - - // Check that the stream was cancelled - let generating = thread.update(cx, |thread, _cx| thread.is_generating()); - assert!(!generating, "The completion should have been cancelled"); - } - - fn init_test_settings(cx: &mut TestAppContext) { - cx.update(|cx| { - let settings_store = SettingsStore::test(cx); - cx.set_global(settings_store); - language::init(cx); - Project::init_settings(cx); - AssistantSettings::register(cx); - prompt_store::init(cx); - thread_store::init(cx); - workspace::init_settings(cx); - language_model::init_settings(cx); - ThemeSettings::register(cx); - EditorSettings::register(cx); - ToolRegistry::default_global(cx); - }); - } - - // Helper to create a test project with test files - async fn create_test_project( - cx: &mut TestAppContext, - files: serde_json::Value, - ) -> Entity { - let fs = FakeFs::new(cx.executor()); - fs.insert_tree(path!("/test"), files).await; - Project::test(fs, [path!("/test").as_ref()], cx).await - } - - async fn setup_test_environment( - cx: &mut TestAppContext, - project: Entity, - ) -> ( - &mut VisualTestContext, - Entity, - Entity, - Arc, - ) { - let (workspace, cx) = - cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx)); - - let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap()); - let thread_store = cx - .update(|_, cx| { - ThreadStore::load( - project.clone(), - cx.new(|_| ToolWorkingSet::default()), - None, - prompt_builder.clone(), - cx, - ) - }) - .await - .unwrap(); - let text_thread_store = cx - .update(|_, cx| { - TextThreadStore::new(project.clone(), prompt_builder, Default::default(), cx) - }) - .await - .unwrap(); - - let thread = thread_store.update(cx, |store, cx| store.create_thread(cx)); - let context_store = cx.new(|_cx| ContextStore::new(project.downgrade(), None)); - - let model = FakeLanguageModel::default(); - let model: Arc = Arc::new(model); - - let language_registry = LanguageRegistry::new(cx.executor()); - let language_registry = Arc::new(language_registry); - - let active_thread = cx.update(|window, cx| { - cx.new(|cx| { - ActiveThread::new( - thread.clone(), - thread_store.clone(), - text_thread_store.clone(), - context_store.clone(), - language_registry.clone(), - workspace.downgrade(), - window, - cx, - ) - }) - }); - - (cx, active_thread, thread, model) - } -} From 06d5e2a0255c66e1b9ab2de7cc8df8b61edf3643 Mon Sep 17 00:00:00 2001 From: Cole Miller Date: Thu, 8 May 2025 19:50:51 -0400 Subject: [PATCH 2/4] Checkpoint --- crates/agent/src/active_thread.rs | 242 +++++++++++++++++------------- crates/agent/src/context_strip.rs | 2 +- 2 files changed, 142 insertions(+), 102 deletions(-) diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index 88fd3ad4512071..9ded515ca1aa20 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -746,7 +746,7 @@ struct UserMessageView { context_strip: Entity, context_picker_menu_handle: PopoverMenuHandle, last_estimated_token_count: Option, - _subscriptions: [Subscription; 2], + _subscriptions: [Subscription; 4], _update_token_count_task: Option>, } @@ -1410,9 +1410,26 @@ impl ActiveThread { }); } - fn cancel_editing_message(&mut self, _: &menu::Cancel, _: &mut Window, cx: &mut Context) { + fn focus_new_message_editor(&mut self, window: &mut Window, cx: &mut Context) { + self.workspace + .update(cx, |workspace, cx| { + if let Some(panel) = workspace.panel::(cx) { + // FIXME + panel.focus_handle(cx).focus(window); + } + }) + .ok(); + } + + // FIXME still needed? + fn cancel_editing_message( + &mut self, + _: &menu::Cancel, + window: &mut Window, + cx: &mut Context, + ) { self.focused_user_message.take(); - // todo!("restore focus to new message editor in parent panel") + self.focus_new_message_editor(window, cx); cx.notify(); } @@ -1422,9 +1439,12 @@ impl ActiveThread { window: &mut Window, cx: &mut Context, ) { - let Some(message_id) = self.focused_user_message else { + self.cancel_last_completion(window, cx); + + let Some(message_id) = self.focused_user_message.take() else { return; }; + self.focus_new_message_editor(window, cx); let Some(user_message_editor) = self.user_message_views.get(&message_id) else { return; }; @@ -1490,9 +1510,9 @@ impl ActiveThread { .unwrap_or(&[]) } - fn handle_cancel_click(&mut self, _: &ClickEvent, window: &mut Window, cx: &mut Context) { - self.cancel_editing_message(&menu::Cancel, window, cx); - } + // fn handle_cancel_click(&mut self, _: &ClickEvent, window: &mut Window, cx: &mut Context) { + // self.cancel_editing_message(&menu::Cancel, window, cx); + // } fn handle_regenerate_click( &mut self, @@ -1605,20 +1625,18 @@ impl ActiveThread { message_id: MessageId, window: &mut Window, cx: &mut Context, - ) -> Option { - // Cancel any ongoing streaming when user starts editing a previous message - // todo!("do this when submitting an edited message") - // self.cancel_last_completion(window, cx); - + ) -> Vec { let view = match self.user_message_views.entry(message_id) { Entry::Occupied(entry) => entry.into_mut(), Entry::Vacant(entry) => { - let message = self.thread.read(cx).message(message_id)?; + let Some(message) = self.thread.read(cx).message(message_id) else { + return vec![]; + }; // User message should always consist of a single text segment, // therefore we can skip returning early if it's not a text segment. let Some(MessageSegment::Text(message_text)) = message.segments.first() else { - return None; + return vec![]; }; let message_text = message_text.clone(); @@ -1667,12 +1685,30 @@ impl ActiveThread { }, ); + // FIXME focus out too? + let editor_focus_handle = editor.focus_handle(cx); + let editor_focus_in_subscription = + cx.on_focus_in(&editor_focus_handle, window, move |this, window, cx| { + this.focused_user_message = Some(message_id); + }); + let editor_focus_out_subscription = + cx.on_focus_out(&editor_focus_handle, window, move |this, _, window, cx| { + if this.focused_user_message == Some(message_id) { + this.focused_user_message = None; + } + }); + entry.insert(UserMessageView { editor: editor.clone(), context_strip, context_picker_menu_handle, last_estimated_token_count: None, - _subscriptions: [buffer_edited_subscription, context_strip_subscription], + _subscriptions: [ + buffer_edited_subscription, + context_strip_subscription, + editor_focus_in_subscription, + editor_focus_out_subscription, + ], _update_token_count_task: None, }) } @@ -1700,35 +1736,87 @@ impl ActiveThread { }; let context_picker_menu_handle = view.context_picker_menu_handle.clone(); - Some( - v_flex() - .key_context("EditMessageEditor") - .on_action(move |_: &ToggleContextPicker, window, cx| { - let handle = context_picker_menu_handle.clone(); - window.defer(cx, move |window, cx| handle.toggle(window, cx)); - }) - .on_action(cx.listener(Self::remove_all_context)) - .on_action(cx.listener(Self::move_up)) - .on_action(cx.listener(Self::cancel_editing_message)) - .on_action(cx.listener(Self::confirm_editing_message)) - .capture_action(cx.listener(Self::paste)) - .min_h_6() - .flex_grow() - .w_full() - .gap_2() - .child(EditorElement::new( - &view.editor, - EditorStyle { - background: colors.editor_background, - local_player: cx.theme().players().local(), - text: text_style, - syntax: cx.theme().syntax().clone(), - ..Default::default() - }, - )) - // todo!("only show this if it's not empty") - .child(view.context_strip.clone()), - ) + let context_strip_is_empty = view.context_strip.read(cx).added_contexts(cx).is_empty(); + + let editor = v_flex() + .key_context("EditMessageEditor") + .on_action(move |_: &ToggleContextPicker, window, cx| { + let handle = context_picker_menu_handle.clone(); + window.defer(cx, move |window, cx| handle.toggle(window, cx)); + }) + .on_action(cx.listener(Self::remove_all_context)) + .on_action(cx.listener(Self::move_up)) + .on_action(cx.listener(Self::cancel_editing_message)) + .on_action(cx.listener(Self::confirm_editing_message)) + .capture_action(cx.listener(Self::paste)) + .min_h_6() + .flex_grow() + .w_full() + .gap_2() + .child(EditorElement::new( + &view.editor, + EditorStyle { + background: colors.editor_background, + local_player: cx.theme().players().local(), + text: text_style, + syntax: cx.theme().syntax().clone(), + ..Default::default() + }, + )) + .when(!context_strip_is_empty, |el| { + el.child(view.context_strip.clone()) + }) + .into_any_element(); + + let focus_handle = view.editor.focus_handle(cx).clone(); + let is_focused = focus_handle.is_focused(window); + let is_empty = view.editor.read(cx).is_empty(cx); + + let buttons = h_flex() + .gap_0p5() + // FIXME seems like we don't need this anymore + // .child( + // IconButton::new("cancel-edit-message", IconName::Close) + // .shape(ui::IconButtonShape::Square) + // .icon_color(Color::Error) + // .tooltip({ + // let focus_handle = focus_handle.clone(); + // move |window, cx| { + // Tooltip::for_action_in( + // "Cancel Edit", + // &menu::Cancel, + // &focus_handle, + // window, + // cx, + // ) + // } + // }) + // .on_click(cx.listener(Self::handle_cancel_click)), + // ) + .when(is_focused, |el| { + el.child( + IconButton::new("confirm-edit-message", IconName::Check) + .disabled(is_empty) + .shape(ui::IconButtonShape::Square) + .icon_color(Color::Success) + .tooltip({ + let focus_handle = focus_handle.clone(); + move |window, cx| { + Tooltip::for_action_in( + "Regenerate", + &menu::Confirm, + &focus_handle, + window, + cx, + ) + } + }) + .on_click(cx.listener(Self::handle_regenerate_click)), + ) + }) + .into_any_element(); + + vec![editor, buttons] } fn render_message( @@ -1908,9 +1996,8 @@ impl ActiveThread { let message_content = if has_content { if message.role == Role::User { self.render_user_message(message_id, window, cx) - .map(|m| m.into_any_element()) } else { - Some( + vec![ v_flex() .w_full() .gap_1() @@ -1943,10 +2030,10 @@ impl ActiveThread { )) }) .into_any_element(), - ) + ] } } else { - None + vec![] }; let colors = cx.theme().colors(); @@ -1969,59 +2056,12 @@ impl ActiveThread { .hover(|hover| hover.border_color(colors.text_accent.opacity(0.5))) .cursor_pointer() .child( - h_flex().p_2p5().gap_1().children(message_content), // todo!("handle this in render_user_message") - // .when_some(editing_message_state, |this, state| { - // let focus_handle = state.editor.focus_handle(cx).clone(); - // this.w_full().justify_between().child( - // h_flex() - // .gap_0p5() - // .child( - // IconButton::new( - // "cancel-edit-message", - // IconName::Close, - // ) - // .shape(ui::IconButtonShape::Square) - // .icon_color(Color::Error) - // .tooltip({ - // let focus_handle = focus_handle.clone(); - // move |window, cx| { - // Tooltip::for_action_in( - // "Cancel Edit", - // &menu::Cancel, - // &focus_handle, - // window, - // cx, - // ) - // } - // }) - // .on_click(cx.listener(Self::handle_cancel_click)), - // ) - // .child( - // IconButton::new( - // "confirm-edit-message", - // IconName::Check, - // ) - // .disabled(state.editor.read(cx).is_empty(cx)) - // .shape(ui::IconButtonShape::Square) - // .icon_color(Color::Success) - // .tooltip({ - // let focus_handle = focus_handle.clone(); - // move |window, cx| { - // Tooltip::for_action_in( - // "Regenerate", - // &menu::Confirm, - // &focus_handle, - // window, - // cx, - // ) - // } - // }) - // .on_click( - // cx.listener(Self::handle_regenerate_click), - // ), - // ), - // ) - // }), + h_flex() + .p_2p5() + .gap_1() + .w_full() + .justify_between() + .children(message_content), ), ), Role::Assistant => v_flex() diff --git a/crates/agent/src/context_strip.rs b/crates/agent/src/context_strip.rs index a0c22de66b4b58..16e88c3bc377bf 100644 --- a/crates/agent/src/context_strip.rs +++ b/crates/agent/src/context_strip.rs @@ -84,7 +84,7 @@ impl ContextStrip { } } - fn added_contexts(&self, cx: &App) -> Vec { + pub fn added_contexts(&self, cx: &App) -> Vec { if let Some(workspace) = self.workspace.upgrade() { let project = workspace.read(cx).project().read(cx); let prompt_store = self From 4ea851b90fe2bb19e01ccd3112460f806338d42f Mon Sep 17 00:00:00 2001 From: Cole Miller Date: Thu, 8 May 2025 20:00:34 -0400 Subject: [PATCH 3/4] Don't limit height of past message editor --- crates/agent/src/active_thread.rs | 3 ++- crates/agent/src/message_editor.rs | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index 9ded515ca1aa20..d6fc16f98314c8 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -18,7 +18,7 @@ use assistant_tool::ToolUseStatus; use collections::{HashMap, HashSet, hash_map::Entry}; use editor::actions::{MoveUp, Paste}; use editor::scroll::Autoscroll; -use editor::{Editor, EditorElement, EditorEvent, EditorStyle, MultiBuffer}; +use editor::{Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, MultiBuffer}; use gpui::{ AbsoluteLength, Animation, AnimationExt, AnyElement, App, ClickEvent, ClipboardEntry, ClipboardItem, DefiniteLength, EdgesRefinement, Empty, Entity, EventEmitter, Focusable, Hsla, @@ -1641,6 +1641,7 @@ impl ActiveThread { let message_text = message_text.clone(); let editor = crate::message_editor::create_editor( + usize::MAX, // XXX self.workspace.clone(), self.context_store.downgrade(), self.thread_store.downgrade(), diff --git a/crates/agent/src/message_editor.rs b/crates/agent/src/message_editor.rs index cb096231eb014b..456ab1c575e85d 100644 --- a/crates/agent/src/message_editor.rs +++ b/crates/agent/src/message_editor.rs @@ -80,6 +80,7 @@ pub struct MessageEditor { const MAX_EDITOR_LINES: usize = 8; pub(crate) fn create_editor( + max_lines: usize, workspace: WeakEntity, context_store: WeakEntity, thread_store: WeakEntity, @@ -99,9 +100,7 @@ pub(crate) fn create_editor( let buffer = cx.new(|cx| Buffer::local("", cx).with_language(Arc::new(language), cx)); let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx)); let mut editor = Editor::new( - editor::EditorMode::AutoHeight { - max_lines: MAX_EDITOR_LINES, - }, + EditorMode::AutoHeight { max_lines }, buffer, None, window, @@ -159,6 +158,7 @@ impl MessageEditor { let model_selector_menu_handle = PopoverMenuHandle::default(); let editor = create_editor( + MAX_EDITOR_LINES, workspace.clone(), context_store.downgrade(), thread_store.clone(), From bef68db74755194278999114e10056bc481b08b2 Mon Sep 17 00:00:00 2001 From: Cole Miller Date: Thu, 8 May 2025 20:19:34 -0400 Subject: [PATCH 4/4] Restore token counting --- crates/agent/src/active_thread.rs | 216 +++++++++++++++--------------- 1 file changed, 109 insertions(+), 107 deletions(-) diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index d6fc16f98314c8..3c9bc004ad3d1d 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -16,9 +16,9 @@ use anyhow::Context as _; use assistant_settings::{AssistantSettings, NotifyWhenAgentWaiting}; use assistant_tool::ToolUseStatus; use collections::{HashMap, HashSet, hash_map::Entry}; -use editor::actions::{MoveUp, Paste}; +use editor::actions::Paste; use editor::scroll::Autoscroll; -use editor::{Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, MultiBuffer}; +use editor::{Editor, EditorElement, EditorEvent, EditorStyle, MultiBuffer}; use gpui::{ AbsoluteLength, Animation, AnimationExt, AnyElement, App, ClickEvent, ClipboardEntry, ClipboardItem, DefiniteLength, EdgesRefinement, Empty, Entity, EventEmitter, Focusable, Hsla, @@ -1272,93 +1272,95 @@ impl ActiveThread { } fn update_editing_message_token_count(&mut self, debounce: bool, cx: &mut Context) { - // todo!() - // let Some((message_id, state)) = self.user_message_views.as_mut() else { - // return; - // }; - - // cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); - // state._update_token_count_task.take(); - - // let Some(configured_model) = self.thread.read(cx).configured_model() else { - // state.last_estimated_token_count.take(); - // return; - // }; - - // let editor = state.editor.clone(); - // let thread = self.thread.clone(); - // let message_id = *message_id; - - // state._update_token_count_task = Some(cx.spawn(async move |this, cx| { - // if debounce { - // cx.background_executor() - // .timer(Duration::from_millis(200)) - // .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 mut request_message = LanguageModelRequestMessage { - // role: language_model::Role::User, - // content: Vec::new(), - // cache: false, - // }; - - // message - // .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: AssistantSettings::temperature_for_model( - // &configured_model.model, - // cx, - // ), - // }; - - // Some(configured_model.model.count_tokens(request, cx)) - // }) - // .ok() - // .flatten() - // { - // task.await.log_err() - // } else { - // Some(0) - // }; - - // if let Some(token_count) = token_count { - // this.update(cx, |this, cx| { - // let Some((_message_id, state)) = this.user_message_views.as_mut() else { - // return; - // }; - - // state.last_estimated_token_count = Some(token_count); - // cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); - // }) - // .ok(); - // }; - // })); + let Some(message_id) = self.focused_user_message else { + return; + }; + + let Some(state) = self.user_message_views.get_mut(&message_id) else { + return; + }; + + cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); + state._update_token_count_task.take(); + + let Some(configured_model) = self.thread.read(cx).configured_model() else { + state.last_estimated_token_count.take(); + return; + }; + + let editor = state.editor.clone(); + let thread = self.thread.clone(); + + state._update_token_count_task = Some(cx.spawn(async move |this, cx| { + if debounce { + cx.background_executor() + .timer(Duration::from_millis(200)) + .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 mut request_message = LanguageModelRequestMessage { + role: language_model::Role::User, + content: Vec::new(), + cache: false, + }; + + message + .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: AssistantSettings::temperature_for_model( + &configured_model.model, + cx, + ), + }; + + Some(configured_model.model.count_tokens(request, cx)) + }) + .ok() + .flatten() + { + task.await.log_err() + } else { + Some(0) + }; + + if let Some(token_count) = token_count { + this.update(cx, |this, cx| { + let Some(state) = this.user_message_views.get_mut(&message_id) else { + return; + }; + + state.last_estimated_token_count = Some(token_count); + cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged); + }) + .ok(); + }; + })); } fn remove_all_context( @@ -1371,16 +1373,15 @@ impl ActiveThread { cx.notify(); } - fn move_up(&mut self, _: &MoveUp, _window: &mut Window, _cx: &mut Context) { - // todo!() - // if let Some((_, state)) = self.user_message_views.as_mut() { - // if state.context_picker_menu_handle.is_deployed() { - // cx.propagate(); - // } else { - // state.context_strip.focus_handle(cx).focus(window); - // } - // } - } + // fn move_up(&mut self, _: &MoveUp, _window: &mut Window, _cx: &mut Context) { + // if let Some((_, state)) = self.user_message_views.as_mut() { + // if state.context_picker_menu_handle.is_deployed() { + // cx.propagate(); + // } else { + // state.context_strip.focus_handle(cx).focus(window); + // } + // } + // } fn paste(&mut self, _: &Paste, _window: &mut Window, cx: &mut Context) { let images = cx @@ -1689,15 +1690,19 @@ impl ActiveThread { // FIXME focus out too? let editor_focus_handle = editor.focus_handle(cx); let editor_focus_in_subscription = - cx.on_focus_in(&editor_focus_handle, window, move |this, window, cx| { + cx.on_focus_in(&editor_focus_handle, window, move |this, _window, cx| { this.focused_user_message = Some(message_id); + this.update_editing_message_token_count(false, cx); }); - let editor_focus_out_subscription = - cx.on_focus_out(&editor_focus_handle, window, move |this, _, window, cx| { + let editor_focus_out_subscription = cx.on_focus_out( + &editor_focus_handle, + window, + move |this, _, _window, _cx| { if this.focused_user_message == Some(message_id) { this.focused_user_message = None; } - }); + }, + ); entry.insert(UserMessageView { editor: editor.clone(), @@ -1715,9 +1720,6 @@ impl ActiveThread { } }; - // todo!("on focus call this") - // self.update_editing_message_token_count(false, cx); - let settings = ThemeSettings::get_global(cx); let font_size = TextSize::Small .rems(cx) @@ -1746,7 +1748,7 @@ impl ActiveThread { window.defer(cx, move |window, cx| handle.toggle(window, cx)); }) .on_action(cx.listener(Self::remove_all_context)) - .on_action(cx.listener(Self::move_up)) + // .on_action(cx.listener(Self::move_up)) .on_action(cx.listener(Self::cancel_editing_message)) .on_action(cx.listener(Self::confirm_editing_message)) .capture_action(cx.listener(Self::paste))