assistant2: Focus prompt editor after dismissing context picker (#22786)

https://github.com/user-attachments/assets/6d0ac75e-fbc2-4bc2-be13-2d109f61361b




Release Notes:

- N/A

crates/assistant2/src/context_strip.rs

@@ -1,7 +1,10 @@
 use std::rc::Rc;
 
 use editor::Editor;
-use gpui::{AppContext, FocusHandle, Model, View, WeakModel, WeakView};
+use gpui::{
+    AppContext, DismissEvent, EventEmitter, FocusHandle, Model, Subscription, View, WeakModel,
+    WeakView,
+};
 use language::Buffer;
 use ui::{prelude::*, KeyBinding, PopoverMenu, PopoverMenuHandle, Tooltip};
 use workspace::Workspace;
@@ -21,6 +24,7 @@ pub struct ContextStrip {
     focus_handle: FocusHandle,
     suggest_context_kind: SuggestContextKind,
     workspace: WeakView<Workspace>,
+    _context_picker_subscription: Subscription,
 }
 
 impl ContextStrip {
@@ -33,21 +37,27 @@ impl ContextStrip {
         suggest_context_kind: SuggestContextKind,
         cx: &mut ViewContext<Self>,
     ) -> Self {
+        let context_picker = cx.new_view(|cx| {
+            ContextPicker::new(
+                workspace.clone(),
+                thread_store.clone(),
+                context_store.downgrade(),
+                ConfirmBehavior::KeepOpen,
+                cx,
+            )
+        });
+
+        let context_picker_subscription =
+            cx.subscribe(&context_picker, Self::handle_context_picker_event);
+
         Self {
             context_store: context_store.clone(),
-            context_picker: cx.new_view(|cx| {
+            context_picker,
-                ContextPicker::new(
-                    workspace.clone(),
-                    thread_store.clone(),
-                    context_store.downgrade(),
-                    ConfirmBehavior::KeepOpen,
-                    cx,
-                )
-            }),
             context_picker_menu_handle,
             focus_handle,
             suggest_context_kind,
             workspace,
+            _context_picker_subscription: context_picker_subscription,
         }
     }
 
@@ -107,6 +117,15 @@ impl ContextStrip {
             thread: weak_active_thread,
         })
     }
+
+    fn handle_context_picker_event(
+        &mut self,
+        _picker: View<ContextPicker>,
+        _event: &DismissEvent,
+        cx: &mut ViewContext<Self>,
+    ) {
+        cx.emit(ContextStripEvent::PickerDismissed);
+    }
 }
 
 impl Render for ContextStrip {
@@ -221,6 +240,12 @@ impl Render for ContextStrip {
     }
 }
 
+pub enum ContextStripEvent {
+    PickerDismissed,
+}
+
+impl EventEmitter<ContextStripEvent> for ContextStrip {}
+
 pub enum SuggestContextKind {
     File,
     Thread,

crates/assistant2/src/inline_prompt_editor.rs

@@ -2,7 +2,7 @@ use crate::assistant_model_selector::AssistantModelSelector;
 use crate::buffer_codegen::BufferCodegen;
 use crate::context_picker::ContextPicker;
 use crate::context_store::ContextStore;
-use crate::context_strip::{ContextStrip, SuggestContextKind};
+use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
 use crate::terminal_codegen::TerminalCodegen;
 use crate::thread_store::ThreadStore;
 use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist};
@@ -47,6 +47,7 @@ pub struct PromptEditor<T> {
     pending_prompt: String,
     _codegen_subscription: Subscription,
     editor_subscriptions: Vec<Subscription>,
+    _context_strip_subscription: Subscription,
     show_rate_limit_notice: bool,
     _phantom: std::marker::PhantomData<T>,
 }
@@ -726,6 +727,16 @@ impl<T: 'static> PromptEditor<T> {
             })
             .into_any_element()
     }
+
+    fn handle_context_strip_event(
+        &mut self,
+        _context_strip: View<ContextStrip>,
+        ContextStripEvent::PickerDismissed: &ContextStripEvent,
+        cx: &mut ViewContext<Self>,
+    ) {
+        let editor_focus_handle = self.editor.focus_handle(cx);
+        cx.focus(&editor_focus_handle);
+    }
 }
 
 pub enum PromptEditorMode {
@@ -803,19 +814,24 @@ impl PromptEditor<BufferCodegen> {
         let context_picker_menu_handle = PopoverMenuHandle::default();
         let model_selector_menu_handle = PopoverMenuHandle::default();
 
+        let context_strip = cx.new_view(|cx| {
+            ContextStrip::new(
+                context_store,
+                workspace.clone(),
+                thread_store.clone(),
+                prompt_editor.focus_handle(cx),
+                context_picker_menu_handle.clone(),
+                SuggestContextKind::Thread,
+                cx,
+            )
+        });
+
+        let context_strip_subscription =
+            cx.subscribe(&context_strip, Self::handle_context_strip_event);
+
         let mut this: PromptEditor<BufferCodegen> = PromptEditor {
             editor: prompt_editor.clone(),
-            context_strip: cx.new_view(|cx| {
+            context_strip,
-                ContextStrip::new(
-                    context_store,
-                    workspace.clone(),
-                    thread_store.clone(),
-                    prompt_editor.focus_handle(cx),
-                    context_picker_menu_handle.clone(),
-                    SuggestContextKind::Thread,
-                    cx,
-                )
-            }),
             context_picker_menu_handle,
             model_selector: cx.new_view(|cx| {
                 AssistantModelSelector::new(fs, model_selector_menu_handle.clone(), cx)
@@ -827,6 +843,7 @@ impl PromptEditor<BufferCodegen> {
             pending_prompt: String::new(),
             _codegen_subscription: codegen_subscription,
             editor_subscriptions: Vec::new(),
+            _context_strip_subscription: context_strip_subscription,
             show_rate_limit_notice: false,
             mode,
             _phantom: Default::default(),
@@ -943,19 +960,24 @@ impl PromptEditor<TerminalCodegen> {
         let context_picker_menu_handle = PopoverMenuHandle::default();
         let model_selector_menu_handle = PopoverMenuHandle::default();
 
+        let context_strip = cx.new_view(|cx| {
+            ContextStrip::new(
+                context_store,
+                workspace.clone(),
+                thread_store.clone(),
+                prompt_editor.focus_handle(cx),
+                context_picker_menu_handle.clone(),
+                SuggestContextKind::Thread,
+                cx,
+            )
+        });
+
+        let context_strip_subscription =
+            cx.subscribe(&context_strip, Self::handle_context_strip_event);
+
         let mut this = Self {
             editor: prompt_editor.clone(),
-            context_strip: cx.new_view(|cx| {
+            context_strip,
-                ContextStrip::new(
-                    context_store,
-                    workspace.clone(),
-                    thread_store.clone(),
-                    prompt_editor.focus_handle(cx),
-                    context_picker_menu_handle.clone(),
-                    SuggestContextKind::Thread,
-                    cx,
-                )
-            }),
             context_picker_menu_handle,
             model_selector: cx.new_view(|cx| {
                 AssistantModelSelector::new(fs, model_selector_menu_handle.clone(), cx)
@@ -967,6 +989,7 @@ impl PromptEditor<TerminalCodegen> {
             pending_prompt: String::new(),
             _codegen_subscription: codegen_subscription,
             editor_subscriptions: Vec::new(),
+            _context_strip_subscription: context_strip_subscription,
             mode,
             show_rate_limit_notice: false,
             _phantom: Default::default(),

crates/assistant2/src/message_editor.rs

@@ -20,7 +20,7 @@ use workspace::Workspace;
 use crate::assistant_model_selector::AssistantModelSelector;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
-use crate::context_strip::{ContextStrip, SuggestContextKind};
+use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
 use crate::thread::{RequestKind, Thread};
 use crate::thread_store::ThreadStore;
 use crate::{Chat, ToggleContextPicker, ToggleModelSelector};
@@ -59,6 +59,7 @@ impl MessageEditor {
 
             editor
         });
+
         let inline_context_picker = cx.new_view(|cx| {
             ContextPicker::new(
                 workspace.clone(),
@@ -68,29 +69,33 @@ impl MessageEditor {
                 cx,
             )
         });
+
+        let context_strip = cx.new_view(|cx| {
+            ContextStrip::new(
+                context_store.clone(),
+                workspace.clone(),
+                Some(thread_store.clone()),
+                editor.focus_handle(cx),
+                context_picker_menu_handle.clone(),
+                SuggestContextKind::File,
+                cx,
+            )
+        });
+
         let subscriptions = vec![
             cx.subscribe(&editor, Self::handle_editor_event),
             cx.subscribe(
                 &inline_context_picker,
                 Self::handle_inline_context_picker_event,
             ),
+            cx.subscribe(&context_strip, Self::handle_context_strip_event),
         ];
 
         Self {
             thread,
             editor: editor.clone(),
-            context_store: context_store.clone(),
+            context_store,
-            context_strip: cx.new_view(|cx| {
+            context_strip,
-                ContextStrip::new(
-                    context_store,
-                    workspace.clone(),
-                    Some(thread_store.clone()),
-                    editor.focus_handle(cx),
-                    context_picker_menu_handle.clone(),
-                    SuggestContextKind::File,
-                    cx,
-                )
-            }),
             context_picker_menu_handle,
             inline_context_picker,
             inline_context_picker_menu_handle,
@@ -195,6 +200,16 @@ impl MessageEditor {
         let editor_focus_handle = self.editor.focus_handle(cx);
         cx.focus(&editor_focus_handle);
     }
+
+    fn handle_context_strip_event(
+        &mut self,
+        _context_strip: View<ContextStrip>,
+        ContextStripEvent::PickerDismissed: &ContextStripEvent,
+        cx: &mut ViewContext<Self>,
+    ) {
+        let editor_focus_handle = self.editor.focus_handle(cx);
+        cx.focus(&editor_focus_handle);
+    }
 }
 
 impl FocusableView for MessageEditor {