assistant2: Suggest current file as context (#22526)

Suggest adding the current file as context in the new assistant panel.



https://github.com/user-attachments/assets/62bc267b-3dfe-4a3b-a6af-c89af2c779a8


Note: This doesn't include suggesting the current thread in the inline
assistant.

Release Notes:

- N/A

crates/assistant2/src/context.rs

@@ -1,5 +1,6 @@
 use gpui::SharedString;
 use language_model::{LanguageModelRequestMessage, MessageContent};
+use project::ProjectEntryId;
 use serde::{Deserialize, Serialize};
 use util::post_inc;
 
@@ -23,7 +24,7 @@ pub struct Context {
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum ContextKind {
-    File,
+    File(ProjectEntryId),
     Directory,
     FetchedUrl,
     Thread,
@@ -40,7 +41,7 @@ pub fn attach_context_to_message(
 
     for context in context.into_iter() {
         match context.kind {
-            ContextKind::File => {
+            ContextKind::File(_) => {
                 file_context.push_str(&context.text);
                 file_context.push('\n');
             }

crates/assistant2/src/context_picker.rs

@@ -15,7 +15,6 @@ use ui::{prelude::*, ListItem, ListItemSpacing};
 use util::ResultExt;
 use workspace::Workspace;
 
-use crate::context::ContextKind;
 use crate::context_picker::directory_context_picker::DirectoryContextPicker;
 use crate::context_picker::fetch_context_picker::FetchContextPicker;
 use crate::context_picker::file_context_picker::FileContextPicker;
@@ -54,7 +53,7 @@ impl ContextPicker {
         let mut entries = Vec::new();
         entries.push(ContextPickerEntry {
             name: "File".into(),
-            kind: ContextKind::File,
+            kind: ContextPickerEntryKind::File,
             icon: IconName::File,
         });
         let release_channel = ReleaseChannel::global(cx);
@@ -63,20 +62,20 @@ impl ContextPicker {
         if release_channel == ReleaseChannel::Dev {
             entries.push(ContextPickerEntry {
                 name: "Folder".into(),
-                kind: ContextKind::Directory,
+                kind: ContextPickerEntryKind::Directory,
                 icon: IconName::Folder,
             });
         }
         entries.push(ContextPickerEntry {
             name: "Fetch".into(),
-            kind: ContextKind::FetchedUrl,
+            kind: ContextPickerEntryKind::FetchedUrl,
             icon: IconName::Globe,
         });
 
         if thread_store.is_some() {
             entries.push(ContextPickerEntry {
                 name: "Thread".into(),
-                kind: ContextKind::Thread,
+                kind: ContextPickerEntryKind::Thread,
                 icon: IconName::MessageCircle,
             });
         }
@@ -140,10 +139,18 @@ impl Render for ContextPicker {
 #[derive(Clone)]
 struct ContextPickerEntry {
     name: SharedString,
-    kind: ContextKind,
+    kind: ContextPickerEntryKind,
     icon: IconName,
 }
 
+#[derive(Debug, Clone)]
+enum ContextPickerEntryKind {
+    File,
+    Directory,
+    FetchedUrl,
+    Thread,
+}
+
 pub(crate) struct ContextPickerDelegate {
     context_picker: WeakView<ContextPicker>,
     workspace: WeakView<Workspace>,
@@ -183,7 +190,7 @@ impl PickerDelegate for ContextPickerDelegate {
             self.context_picker
                 .update(cx, |this, cx| {
                     match entry.kind {
-                        ContextKind::File => {
+                        ContextPickerEntryKind::File => {
                             this.mode = ContextPickerMode::File(cx.new_view(|cx| {
                                 FileContextPicker::new(
                                     self.context_picker.clone(),
@@ -194,7 +201,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                 )
                             }));
                         }
-                        ContextKind::Directory => {
+                        ContextPickerEntryKind::Directory => {
                             this.mode = ContextPickerMode::Directory(cx.new_view(|cx| {
                                 DirectoryContextPicker::new(
                                     self.context_picker.clone(),
@@ -205,7 +212,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                 )
                             }));
                         }
-                        ContextKind::FetchedUrl => {
+                        ContextPickerEntryKind::FetchedUrl => {
                             this.mode = ContextPickerMode::Fetch(cx.new_view(|cx| {
                                 FetchContextPicker::new(
                                     self.context_picker.clone(),
@@ -216,7 +223,7 @@ impl PickerDelegate for ContextPickerDelegate {
                                 )
                             }));
                         }
-                        ContextKind::Thread => {
+                        ContextPickerEntryKind::Thread => {
                             if let Some(thread_store) = self.thread_store.as_ref() {
                                 this.mode = ContextPickerMode::Thread(cx.new_view(|cx| {
                                     ThreadContextPicker::new(

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 use fuzzy::PathMatch;
 use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakModel, WeakView};
 use picker::{Picker, PickerDelegate};
-use project::{PathMatchCandidateSet, WorktreeId};
+use project::{PathMatchCandidateSet, ProjectPath, WorktreeId};
 use ui::{prelude::*, ListItem};
 use util::ResultExt as _;
 use workspace::Workspace;
@@ -207,11 +207,20 @@ impl PickerDelegate for FileContextPickerDelegate {
         let worktree_id = WorktreeId::from_usize(mat.worktree_id);
         let confirm_behavior = self.confirm_behavior;
         cx.spawn(|this, mut cx| async move {
-            let Some(open_buffer_task) = project
+            let Some((entry_id, open_buffer_task)) = project
                 .update(&mut cx, |project, cx| {
-                    project.open_buffer((worktree_id, path.clone()), cx)
+                    let project_path = ProjectPath {
+                        worktree_id,
+                        path: path.clone(),
+                    };
+
+                    let entry_id = project.entry_for_path(&project_path, cx)?.id;
+                    let task = project.open_buffer(project_path, cx);
+
+                    Some((entry_id, task))
                 })
                 .ok()
+                .flatten()
             else {
                 return anyhow::Ok(());
             };
@@ -232,7 +241,7 @@ impl PickerDelegate for FileContextPickerDelegate {
                         text.push_str("```\n");
 
                         context_store.insert_context(
-                            ContextKind::File,
+                            ContextKind::File(entry_id),
                             path.to_string_lossy().to_string(),
                             text,
                         );

crates/assistant2/src/context_store.rs

@@ -1,4 +1,5 @@
 use gpui::SharedString;
+use project::ProjectEntryId;
 
 use crate::context::{Context, ContextId, ContextKind};
 
@@ -44,4 +45,13 @@ impl ContextStore {
     pub fn remove_context(&mut self, id: &ContextId) {
         self.context.retain(|context| context.id != *id);
     }
+
+    pub fn contains_project_entry(&self, entry_id: ProjectEntryId) -> bool {
+        self.context.iter().any(|probe| match probe.kind {
+            ContextKind::File(probe_entry_id) => probe_entry_id == entry_id,
+            ContextKind::Directory => false,
+            ContextKind::FetchedUrl => false,
+            ContextKind::Thread => false,
+        })
+    }
 }

crates/assistant2/src/context_strip.rs

@@ -1,9 +1,13 @@
 use std::rc::Rc;
 
-use gpui::{FocusHandle, Model, View, WeakModel, WeakView};
+use editor::Editor;
+use gpui::{EntityId, FocusHandle, Model, Subscription, View, WeakModel, WeakView};
+use language::Buffer;
+use project::ProjectEntryId;
 use ui::{prelude::*, PopoverMenu, PopoverMenuHandle, Tooltip};
-use workspace::Workspace;
+use workspace::{ItemHandle, Workspace};
 
+use crate::context::ContextKind;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 use crate::thread_store::ThreadStore;
@@ -16,6 +20,21 @@ pub struct ContextStrip {
     context_picker: View<ContextPicker>,
     context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
     focus_handle: FocusHandle,
+    workspace_active_pane_id: Option<EntityId>,
+    suggested_context: Option<SuggestedContext>,
+    _subscription: Option<Subscription>,
+}
+
+pub enum SuggestContextKind {
+    File,
+    Thread,
+}
+
+#[derive(Clone)]
+pub struct SuggestedContext {
+    entry_id: ProjectEntryId,
+    title: SharedString,
+    buffer: WeakModel<Buffer>,
 }
 
 impl ContextStrip {
@@ -25,8 +44,23 @@ impl ContextStrip {
         thread_store: Option<WeakModel<ThreadStore>>,
         focus_handle: FocusHandle,
         context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
+        suggest_context_kind: SuggestContextKind,
         cx: &mut ViewContext<Self>,
     ) -> Self {
+        let subscription = match suggest_context_kind {
+            SuggestContextKind::File => {
+                if let Some(workspace) = workspace.upgrade() {
+                    Some(cx.subscribe(&workspace, Self::handle_workspace_event))
+                } else {
+                    None
+                }
+            }
+            SuggestContextKind::Thread => {
+                // TODO: Suggest current thread
+                None
+            }
+        };
+
         Self {
             context_store: context_store.clone(),
             context_picker: cx.new_view(|cx| {
@@ -40,16 +74,73 @@ impl ContextStrip {
             }),
             context_picker_menu_handle,
             focus_handle,
+            workspace_active_pane_id: None,
+            suggested_context: None,
+            _subscription: subscription,
         }
     }
+
+    fn handle_workspace_event(
+        &mut self,
+        workspace: View<Workspace>,
+        event: &workspace::Event,
+        cx: &mut ViewContext<Self>,
+    ) {
+        match event {
+            workspace::Event::WorkspaceCreated(_) | workspace::Event::ActiveItemChanged => {
+                let workspace = workspace.read(cx);
+
+                if let Some(active_item) = workspace.active_item(cx) {
+                    let new_active_item_id = Some(active_item.item_id());
+
+                    if self.workspace_active_pane_id != new_active_item_id {
+                        self.suggested_context = Self::suggested_file(active_item, cx);
+                        self.workspace_active_pane_id = new_active_item_id;
+                    }
+                } else {
+                    self.suggested_context = None;
+                    self.workspace_active_pane_id = None;
+                }
+            }
+            _ => {}
+        }
+    }
+
+    fn suggested_file(
+        active_item: Box<dyn ItemHandle>,
+        cx: &WindowContext,
+    ) -> Option<SuggestedContext> {
+        let entry_id = *active_item.project_entry_ids(cx).first()?;
+
+        let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
+        let active_buffer = editor.buffer().read(cx).as_singleton()?;
+
+        let file = active_buffer.read(cx).file()?;
+        let title = file.path().to_string_lossy().into_owned().into();
+
+        Some(SuggestedContext {
+            entry_id,
+            title,
+            buffer: active_buffer.downgrade(),
+        })
+    }
 }
 
 impl Render for ContextStrip {
     fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
-        let context = self.context_store.read(cx).context().clone();
+        let context_store = self.context_store.read(cx);
+        let context = context_store.context().clone();
         let context_picker = self.context_picker.clone();
         let focus_handle = self.focus_handle.clone();
 
+        let suggested_context = self.suggested_context.as_ref().and_then(|suggested| {
+            if context_store.contains_project_entry(suggested.entry_id) {
+                None
+            } else {
+                Some(suggested.clone())
+            }
+        });
+
         h_flex()
             .flex_wrap()
             .gap_1()
@@ -60,13 +151,17 @@ impl Render for ContextStrip {
                         IconButton::new("add-context", IconName::Plus)
                             .icon_size(IconSize::Small)
                             .style(ui::ButtonStyle::Filled)
-                            .tooltip(move |cx| {
+                            .tooltip({
+                                let focus_handle = focus_handle.clone();
+
+                                move |cx| {
                                     Tooltip::for_action_in(
                                         "Add Context",
                                         &ToggleContextPicker,
                                         &focus_handle,
                                         cx,
                                     )
+                                }
                             }),
                     )
                     .attach(gpui::Corner::TopLeft)
@@ -77,7 +172,7 @@ impl Render for ContextStrip {
                     })
                     .with_handle(self.context_picker_menu_handle.clone()),
             )
-            .when(context.is_empty(), {
+            .when(context.is_empty() && self.suggested_context.is_none(), {
                 |parent| {
                     parent.child(
                         h_flex()
@@ -91,7 +186,7 @@ impl Render for ContextStrip {
                             .children(
                                 ui::KeyBinding::for_action_in(
                                     &ToggleContextPicker,
-                                    &self.focus_handle,
+                                    &focus_handle,
                                     cx,
                                 )
                                 .map(|binding| binding.into_any_element()),
@@ -112,6 +207,41 @@ impl Render for ContextStrip {
                     }))
                 })
             }))
+            .when_some(suggested_context, |el, suggested| {
+                el.child(
+                    Button::new("add-suggested-context", suggested.title.clone())
+                        .on_click({
+                            let context_store = self.context_store.clone();
+
+                            cx.listener(move |_this, _event, cx| {
+                                let Some(buffer) = suggested.buffer.upgrade() else {
+                                    return;
+                                };
+
+                                let title = suggested.title.clone();
+                                let text = buffer.read(cx).text();
+
+                                context_store.update(cx, move |context_store, _cx| {
+                                    context_store.insert_context(
+                                        ContextKind::File(suggested.entry_id),
+                                        title,
+                                        text,
+                                    );
+                                });
+                                cx.notify();
+                            })
+                        })
+                        .icon(IconName::Plus)
+                        .icon_position(IconPosition::Start)
+                        .icon_size(IconSize::XSmall)
+                        .icon_color(Color::Muted)
+                        .label_size(LabelSize::Small)
+                        .style(ButtonStyle::Filled)
+                        .tooltip(|cx| {
+                            Tooltip::with_meta("Suggested Context", None, "Click to add it", cx)
+                        }),
+                )
+            })
             .when(!context.is_empty(), {
                 move |parent| {
                     parent.child(

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;
+use crate::context_strip::{ContextStrip, SuggestContextKind};
 use crate::terminal_codegen::TerminalCodegen;
 use crate::thread_store::ThreadStore;
 use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist};
@@ -793,6 +793,7 @@ impl PromptEditor<BufferCodegen> {
                     thread_store.clone(),
                     prompt_editor.focus_handle(cx),
                     context_picker_menu_handle.clone(),
+                    SuggestContextKind::Thread,
                     cx,
                 )
             }),
@@ -932,6 +933,7 @@ impl PromptEditor<TerminalCodegen> {
                     thread_store.clone(),
                     prompt_editor.focus_handle(cx),
                     context_picker_menu_handle.clone(),
+                    SuggestContextKind::Thread,
                     cx,
                 )
             }),

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;
+use crate::context_strip::{ContextStrip, SuggestContextKind};
 use crate::thread::{RequestKind, Thread};
 use crate::thread_store::ThreadStore;
 use crate::{Chat, ToggleContextPicker, ToggleModelSelector};
@@ -87,6 +87,7 @@ impl MessageEditor {
                     Some(thread_store.clone()),
                     editor.focus_handle(cx),
                     context_picker_menu_handle.clone(),
+                    SuggestContextKind::File,
                     cx,
                 )
             }),

crates/assistant2/src/ui/context_pill.rs

@@ -33,7 +33,7 @@ impl RenderOnce for ContextPill {
             px(4.)
         };
         let icon = match self.context.kind {
-            ContextKind::File => IconName::File,
+            ContextKind::File(_) => IconName::File,
             ContextKind::Directory => IconName::Folder,
             ContextKind::FetchedUrl => IconName::Globe,
             ContextKind::Thread => IconName::MessageCircle,