Attachment store for assistant2 (#11327)

This sets up a way for the user (or Zed) to _push_ context instead of
having the model retrieve it with a function. Our first use is the
contents of the current file.

<img width="399" alt="image"
src="https://github.com/zed-industries/zed/assets/836375/198429a5-82af-4b82-86f6-cb961f10de5c">

<img width="393" alt="image"
src="https://github.com/zed-industries/zed/assets/836375/cfb52444-723b-4fc1-bddc-57e1810c512b">

I heard the asst2 example was deleted in another branch so I deleted
that here too since we wanted the workspace access.

Release Notes:

- N/A

---------

Co-authored-by: Marshall <marshall@zed.dev>

crates/assistant2/src/ui/active_file_button.rs

@@ -0,0 +1,133 @@
+use crate::attachments::{ActiveEditorAttachmentTool, UserAttachmentStore};
+use editor::Editor;
+use gpui::{prelude::*, Subscription, View};
+use std::sync::Arc;
+use ui::{prelude::*, ButtonLike, Color, Icon, IconName, Tooltip};
+use workspace::Workspace;
+
+#[derive(Clone)]
+enum Status {
+    ActiveFile(String),
+    #[allow(dead_code)]
+    NoFile,
+}
+
+pub struct ActiveFileButton {
+    attachment_store: Arc<UserAttachmentStore>,
+    status: Status,
+    #[allow(dead_code)]
+    workspace_subscription: Subscription,
+}
+
+impl ActiveFileButton {
+    pub fn new(
+        attachment_store: Arc<UserAttachmentStore>,
+        workspace: View<Workspace>,
+        cx: &mut ViewContext<Self>,
+    ) -> Self {
+        let workspace_subscription = cx.subscribe(&workspace, Self::handle_workspace_event);
+
+        cx.defer(move |this, cx| this.update_active_buffer(workspace.clone(), cx));
+
+        Self {
+            attachment_store,
+            status: Status::NoFile,
+            workspace_subscription,
+        }
+    }
+
+    pub fn set_enabled(&mut self, enabled: bool) {
+        self.attachment_store
+            .set_attachment_tool_enabled::<ActiveEditorAttachmentTool>(enabled);
+    }
+
+    pub fn update_active_buffer(&mut self, workspace: View<Workspace>, cx: &mut ViewContext<Self>) {
+        let active_buffer = workspace
+            .read(cx)
+            .active_item(cx)
+            .and_then(|item| Some(item.act_as::<Editor>(cx)?.read(cx).buffer().clone()));
+
+        if let Some(buffer) = active_buffer {
+            let buffer = buffer.read(cx);
+
+            if let Some(singleton) = buffer.as_singleton() {
+                let singleton = singleton.read(cx);
+
+                let filename: String = singleton
+                    .file()
+                    .map(|file| file.path().to_string_lossy())
+                    .unwrap_or("Untitled".into())
+                    .into();
+
+                self.status = Status::ActiveFile(filename);
+            }
+        }
+    }
+
+    fn handle_workspace_event(
+        &mut self,
+        workspace: View<Workspace>,
+        event: &workspace::Event,
+        cx: &mut ViewContext<Self>,
+    ) {
+        if let workspace::Event::ActiveItemChanged = event {
+            self.update_active_buffer(workspace, cx);
+        }
+    }
+}
+
+impl Render for ActiveFileButton {
+    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
+        let is_enabled = self
+            .attachment_store
+            .is_attachment_tool_enabled::<ActiveEditorAttachmentTool>();
+
+        let icon = if is_enabled {
+            Icon::new(IconName::File)
+                .size(IconSize::XSmall)
+                .color(Color::Default)
+        } else {
+            Icon::new(IconName::File)
+                .size(IconSize::XSmall)
+                .color(Color::Disabled)
+        };
+
+        let indicator = None;
+
+        let status = self.status.clone();
+
+        ButtonLike::new("active-file-button")
+            .child(
+                ui::IconWithIndicator::new(icon, indicator)
+                    .indicator_border_color(Some(gpui::transparent_black())),
+            )
+            .tooltip({
+                move |cx| {
+                    let status = status.clone();
+                    let (tooltip, meta) = match (is_enabled, status) {
+                        (false, _) => (
+                            "Active file disabled".to_string(),
+                            Some("Click to enable".to_string()),
+                        ),
+                        (true, Status::ActiveFile(filename)) => (
+                            format!("Active file {filename} enabled"),
+                            Some("Click to disable".to_string()),
+                        ),
+                        (true, Status::NoFile) => {
+                            ("No file active for conversation".to_string(), None)
+                        }
+                    };
+
+                    if let Some(meta) = meta {
+                        Tooltip::with_meta(tooltip, None, meta, cx)
+                    } else {
+                        Tooltip::text(tooltip, cx)
+                    }
+                }
+            })
+            .on_click(cx.listener(move |this, _, cx| {
+                this.set_enabled(!is_enabled);
+                cx.notify();
+            }))
+    }
+}

crates/assistant2/src/ui/chat_message.rs

@@ -16,6 +16,7 @@ pub struct ChatMessage {
     id: MessageId,
     player: UserOrAssistant,
     message: Option<AnyElement>,
+    tools_used: Option<AnyElement>,
     collapsed: bool,
     on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
 }
@@ -25,6 +26,7 @@ impl ChatMessage {
         id: MessageId,
         player: UserOrAssistant,
         message: Option<AnyElement>,
+        tools_used: Option<AnyElement>,
         collapsed: bool,
         on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
     ) -> Self {
@@ -32,6 +34,7 @@ impl ChatMessage {
             id,
             player,
             message,
+            tools_used,
             collapsed,
             on_collapse_handle_click,
         }
@@ -66,6 +69,10 @@ impl RenderOnce for ChatMessage {
         // Clamp the message height to exactly 1.5 lines when collapsed.
         let collapsed_height = content_padding.to_pixels(cx.rem_size()) + cx.line_height() * 1.5;
 
+        let tools_used = self
+            .tools_used
+            .map(|attachment| div().mt_3().child(attachment));
+
         let content = self.message.map(|message| {
             div()
                 .overflow_hidden()
@@ -75,6 +82,7 @@ impl RenderOnce for ChatMessage {
                 .when(self.collapsed, |this| this.h(collapsed_height))
                 .bg(cx.theme().colors().surface_background)
                 .child(message)
+                .children(tools_used)
         });
 
         v_flex()

crates/assistant2/src/ui/composer.rs

@@ -1,14 +1,18 @@
-use crate::{ui::ProjectIndexButton, AssistantChat, CompletionProvider};
+use crate::{
+    ui::{ActiveFileButton, ProjectIndexButton},
+    AssistantChat, CompletionProvider,
+};
 use editor::{Editor, EditorElement, EditorStyle};
 use gpui::{AnyElement, FontStyle, FontWeight, TextStyle, View, WeakView, WhiteSpace};
 use settings::Settings;
 use theme::ThemeSettings;
-use ui::{popover_menu, prelude::*, ButtonLike, ContextMenu, Tooltip};
+use ui::{popover_menu, prelude::*, ButtonLike, ContextMenu, Divider, Tooltip};
 
 #[derive(IntoElement)]
 pub struct Composer {
     editor: View<Editor>,
     project_index_button: Option<View<ProjectIndexButton>>,
+    active_file_button: Option<View<ActiveFileButton>>,
     model_selector: AnyElement,
 }
 
@@ -16,11 +20,13 @@ impl Composer {
     pub fn new(
         editor: View<Editor>,
         project_index_button: Option<View<ProjectIndexButton>>,
+        active_file_button: Option<View<ActiveFileButton>>,
         model_selector: AnyElement,
     ) -> Self {
         Self {
             editor,
             project_index_button,
+            active_file_button,
             model_selector,
         }
     }
@@ -32,6 +38,14 @@ impl Composer {
                 .map(|view| view.into_any_element()),
         )
     }
+
+    fn render_attachment_tools(&mut self, _cx: &mut WindowContext) -> impl IntoElement {
+        h_flex().children(
+            self.active_file_button
+                .clone()
+                .map(|view| view.into_any_element()),
+        )
+    }
 }
 
 impl RenderOnce for Composer {
@@ -83,7 +97,15 @@ impl RenderOnce for Composer {
                                     .gap_2()
                                     .justify_between()
                                     .w_full()
-                                    .child(h_flex().gap_1().child(self.render_tools(cx)))
+                                    .child(
+                                        h_flex().gap_1().child(
+                                            h_flex()
+                                                .gap_2()
+                                                .child(self.render_tools(cx))
+                                                .child(Divider::vertical())
+                                                .child(self.render_attachment_tools(cx)),
+                                        ),
+                                    )
                                     .child(h_flex().gap_1().child(self.model_selector)),
                             ),
                     ),

crates/assistant2/src/ui/stories/chat_message.rs

@@ -29,6 +29,7 @@ impl Render for ChatMessageStory {
                         MessageId(0),
                         UserOrAssistant::User(Some(user_1.clone())),
                         Some(div().child("What can I do here?").into_any_element()),
+                        None,
                         false,
                         Box::new(|_, _| {}),
                     ),
@@ -39,6 +40,7 @@ impl Render for ChatMessageStory {
                         MessageId(0),
                         UserOrAssistant::User(Some(user_1.clone())),
                         Some(div().child("What can I do here?").into_any_element()),
+                        None,
                         true,
                         Box::new(|_, _| {}),
                     ),
@@ -52,6 +54,7 @@ impl Render for ChatMessageStory {
                         MessageId(0),
                         UserOrAssistant::Assistant,
                         Some(div().child("You can talk to me!").into_any_element()),
+                        None,
                         false,
                         Box::new(|_, _| {}),
                     ),
@@ -62,6 +65,7 @@ impl Render for ChatMessageStory {
                         MessageId(0),
                         UserOrAssistant::Assistant,
                         Some(div().child(MULTI_LINE_MESSAGE).into_any_element()),
+                        None,
                         true,
                         Box::new(|_, _| {}),
                     ),
@@ -76,6 +80,7 @@ impl Render for ChatMessageStory {
                         MessageId(0),
                         UserOrAssistant::User(Some(user_1.clone())),
                         Some(div().child("What is Rust??").into_any_element()),
+                        None,
                         false,
                         Box::new(|_, _| {}),
                     ))
@@ -83,6 +88,7 @@ impl Render for ChatMessageStory {
                         MessageId(0),
                         UserOrAssistant::Assistant,
                         Some(div().child("Rust is a multi-paradigm programming language focused on performance and safety").into_any_element()),
+                        None,
                         false,
                         Box::new(|_, _| {}),
                     ))
@@ -90,6 +96,7 @@ impl Render for ChatMessageStory {
                         MessageId(0),
                         UserOrAssistant::User(Some(user_1)),
                         Some(div().child("Sounds pretty cool!").into_any_element()),
+                        None,
                         false,
                         Box::new(|_, _| {}),
                     )),