assistant2: Show file icons for context entries

crates/assistant2/Cargo.toml

@@ -73,6 +73,7 @@ util.workspace = true
 uuid.workspace = true
 workspace.workspace = true
 zed_actions.workspace = true
+file_icons.workspace = true
 
 [dev-dependencies]
 rand.workspace = true

crates/assistant2/src/context.rs

@@ -3,11 +3,13 @@ use std::rc::Rc;
 use std::sync::Arc;
 
 use collections::BTreeMap;
+use file_icons::FileIcons;
 use gpui::{AppContext, Model, SharedString};
 use language::Buffer;
 use language_model::{LanguageModelRequestMessage, MessageContent};
 use serde::{Deserialize, Serialize};
 use text::BufferId;
+use ui::IconName;
 use util::post_inc;
 
 use crate::thread::Thread;
@@ -28,6 +30,7 @@ pub struct ContextSnapshot {
     pub name: SharedString,
     pub parent: Option<SharedString>,
     pub tooltip: Option<SharedString>,
+    pub icon_path: Option<SharedString>,
     pub kind: ContextKind,
     /// Text to send to the model. This is not refreshed by `snapshot`.
     pub text: SharedString,
@@ -41,6 +44,17 @@ pub enum ContextKind {
     Thread,
 }
 
+impl ContextKind {
+    pub fn icon(&self) -> IconName {
+        match self {
+            ContextKind::File => IconName::File,
+            ContextKind::Directory => IconName::Folder,
+            ContextKind::FetchedUrl => IconName::Globe,
+            ContextKind::Thread => IconName::MessageCircle,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub enum Context {
     File(FileContext),
@@ -135,11 +149,14 @@ impl FileContext {
             .and_then(|p| p.file_name())
             .map(|p| p.to_string_lossy().into_owned().into());
 
+        let icon_path = FileIcons::get_icon(&path, cx);
+
         Some(ContextSnapshot {
             id: self.id,
             name,
             parent,
             tooltip: Some(full_path),
+            icon_path,
             kind: ContextKind::File,
             text: self.text.clone(),
         })
@@ -159,6 +176,7 @@ impl FetchedUrlContext {
             name: self.url.clone(),
             parent: None,
             tooltip: None,
+            icon_path: None,
             kind: ContextKind::FetchedUrl,
             text: self.text.clone(),
         }
@@ -173,6 +191,7 @@ impl ThreadContext {
             name: thread.summary().unwrap_or("New thread".into()),
             parent: None,
             tooltip: None,
+            icon_path: None,
             kind: ContextKind::Thread,
             text: self.text.clone(),
         }

crates/assistant2/src/context_picker/file_context_picker.rs

@@ -2,6 +2,7 @@ use std::path::Path;
 use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
+use file_icons::FileIcons;
 use fuzzy::PathMatch;
 use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakModel, WeakView};
 use picker::{Picker, PickerDelegate};
@@ -281,13 +282,18 @@ impl PickerDelegate for FileContextPickerDelegate {
                 .will_include_file_path(&path_match.path, cx)
         });
 
+        let file_icon = FileIcons::get_icon(&path_match.path.clone(), cx)
+            .map(|icon_path| Icon::from_path(icon_path))
+            .unwrap_or_else(|| Icon::new(IconName::File));
+
         Some(
             ListItem::new(ix)
                 .inset(true)
                 .toggle_state(selected)
                 .child(
                     h_flex()
                         .gap_2()
+                        .child(file_icon.size(IconSize::Small))
                         .child(Label::new(file_name))
                         .children(directory.map(|directory| {
                             Label::new(directory)

crates/assistant2/src/context_store.rs

@@ -234,6 +234,7 @@ impl ContextStore {
                 name,
                 parent,
                 tooltip: Some(full_path),
+                icon_path: None,
                 kind: ContextKind::Directory,
                 text: text.into(),
             },

crates/assistant2/src/context_strip.rs

@@ -2,6 +2,7 @@ use std::rc::Rc;
 
 use collections::HashSet;
 use editor::Editor;
+use file_icons::FileIcons;
 use gpui::{
     AppContext, DismissEvent, EventEmitter, FocusHandle, Model, Subscription, View, WeakModel,
     WeakView,
@@ -94,9 +95,12 @@ impl ContextStrip {
             None => path.to_string_lossy().into_owned().into(),
         };
 
+        let icon_path = FileIcons::get_icon(path, cx);
+
         Some(SuggestedContext::File {
             name,
             buffer: active_buffer_model.downgrade(),
+            icon_path,
         })
     }
 
@@ -227,6 +231,7 @@ impl Render for ContextStrip {
             .when_some(suggested_context, |el, suggested| {
                 el.child(ContextPill::new_suggested(
                     suggested.name().clone(),
+                    suggested.icon_path(),
                     suggested.kind(),
                     {
                         let context_store = self.context_store.clone();
@@ -283,6 +288,7 @@ pub enum SuggestContextKind {
 pub enum SuggestedContext {
     File {
         name: SharedString,
+        icon_path: Option<SharedString>,
         buffer: WeakModel<Buffer>,
     },
     Thread {
@@ -299,9 +305,20 @@ impl SuggestedContext {
         }
     }
 
+    pub fn icon_path(&self) -> Option<SharedString> {
+        match self {
+            Self::File { icon_path, .. } => icon_path.clone(),
+            Self::Thread { .. } => None,
+        }
+    }
+
     pub fn accept(&self, context_store: &mut ContextStore, cx: &mut AppContext) {
         match self {
-            Self::File { buffer, name: _ } => {
+            Self::File {
+                buffer,
+                icon_path: _,
+                name: _,
+            } => {
                 if let Some(buffer) = buffer.upgrade() {
                     context_store.insert_file(buffer, cx);
                 };

crates/assistant2/src/ui/context_pill.rs

@@ -14,6 +14,7 @@ pub enum ContextPill {
     },
     Suggested {
         name: SharedString,
+        icon_path: Option<SharedString>,
         kind: ContextKind,
         on_add: Rc<dyn Fn(&ClickEvent, &mut WindowContext)>,
     },
@@ -34,10 +35,16 @@ impl ContextPill {
 
     pub fn new_suggested(
         name: SharedString,
+        icon_path: Option<SharedString>,
         kind: ContextKind,
         on_add: Rc<dyn Fn(&ClickEvent, &mut WindowContext)>,
     ) -> Self {
-        Self::Suggested { name, kind, on_add }
+        Self::Suggested {
+            name,
+            icon_path,
+            kind,
+            on_add,
+        }
     }
 
     pub fn id(&self) -> ElementId {
@@ -49,23 +56,27 @@ impl ContextPill {
         }
     }
 
-    pub fn kind(&self) -> ContextKind {
+    pub fn icon(&self) -> Icon {
         match self {
-            Self::Added { context, .. } => context.kind,
-            Self::Suggested { kind, .. } => *kind,
+            Self::Added { context, .. } => match &context.icon_path {
+                Some(icon_path) => Icon::from_path(icon_path),
+                None => Icon::new(context.kind.icon()),
+            },
+            Self::Suggested {
+                icon_path: Some(icon_path),
+                ..
+            } => Icon::from_path(icon_path),
+            Self::Suggested {
+                kind,
+                icon_path: None,
+                ..
+            } => Icon::new(kind.icon()),
         }
     }
 }
 
 impl RenderOnce for ContextPill {
     fn render(self, cx: &mut WindowContext) -> impl IntoElement {
-        let icon = match &self.kind() {
-            ContextKind::File => IconName::File,
-            ContextKind::Directory => IconName::Folder,
-            ContextKind::FetchedUrl => IconName::Globe,
-            ContextKind::Thread => IconName::MessageCircle,
-        };
-
         let color = cx.theme().colors();
 
         let base_pill = h_flex()
@@ -75,7 +86,7 @@ impl RenderOnce for ContextPill {
             .border_1()
             .rounded_md()
             .gap_1()
-            .child(Icon::new(icon).size(IconSize::XSmall).color(Color::Muted));
+            .child(self.icon().size(IconSize::XSmall).color(Color::Muted));
 
         match &self {
             ContextPill::Added {
@@ -118,7 +129,12 @@ impl RenderOnce for ContextPill {
                             }),
                     )
                 }),
-            ContextPill::Suggested { name, kind, on_add } => base_pill
+            ContextPill::Suggested {
+                name,
+                icon_path: _,
+                kind,
+                on_add,
+            } => base_pill
                 .cursor_pointer()
                 .pr_1()
                 .border_color(color.border_variant.opacity(0.5))