From 3d4e0780c440343e14269530fab6cd8db2349b2b Mon Sep 17 00:00:00 2001
From: Danilo Leal <67129314+danilo-leal@users.noreply.github.com>
Date: Tue, 24 Dec 2024 01:26:09 -0300
Subject: [PATCH] assistant2: Add icons to the context pill (#22385)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR adds an icon field to the `ContextKind` enum, which means that
icons will now display on context pills, both on the message editor and
on the active thread.

<img width="800" alt="Screenshot 2024-12-24 at 12 23 17 AM"
src="https://github.com/user-attachments/assets/f00e540b-30fe-49ac-b3df-7c7a5dc86d65"
/>

Release Notes:

- N/A
---
 crates/assistant2/src/context.rs         | 13 +++++++++++++
 crates/assistant2/src/context_store.rs   |  3 ++-
 crates/assistant2/src/ui/context_pill.rs | 15 +++++++++++++--
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/crates/assistant2/src/context.rs b/crates/assistant2/src/context.rs
index 1aff6c982d..096dbcab9c 100644
--- a/crates/assistant2/src/context.rs
+++ b/crates/assistant2/src/context.rs
@@ -1,6 +1,7 @@
 use gpui::SharedString;
 use language_model::{LanguageModelRequestMessage, MessageContent};
 use serde::{Deserialize, Serialize};
+use ui::prelude::*;
 use util::post_inc;
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
@@ -19,6 +20,7 @@ pub struct Context {
     pub name: SharedString,
     pub kind: ContextKind,
     pub text: SharedString,
+    pub icon: IconName,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -29,6 +31,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,
+        }
+    }
+}
+
 pub fn attach_context_to_message(
     message: &mut LanguageModelRequestMessage,
     context: impl IntoIterator<Item = Context>,
diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs
index febd1f597d..0a34e65881 100644
--- a/crates/assistant2/src/context_store.rs
+++ b/crates/assistant2/src/context_store.rs
@@ -36,8 +36,9 @@ impl ContextStore {
         self.context.push(Context {
             id: self.next_context_id.post_inc(),
             name: name.into(),
-            kind,
+            kind: kind.clone(),
             text: text.into(),
+            icon: kind.icon(),
         });
     }
 
diff --git a/crates/assistant2/src/ui/context_pill.rs b/crates/assistant2/src/ui/context_pill.rs
index d2f020a582..768eddfc3f 100644
--- a/crates/assistant2/src/ui/context_pill.rs
+++ b/crates/assistant2/src/ui/context_pill.rs
@@ -27,15 +27,26 @@ impl ContextPill {
 
 impl RenderOnce for ContextPill {
     fn render(self, cx: &mut WindowContext) -> impl IntoElement {
+        let padding_right = if self.on_remove.is_some() {
+            px(2.)
+        } else {
+            px(4.)
+        };
+
         h_flex()
             .gap_1()
-            .pl_1p5()
-            .pr_0p5()
+            .pl_1()
+            .pr(padding_right)
             .pb(px(1.))
             .border_1()
             .border_color(cx.theme().colors().border.opacity(0.5))
             .bg(cx.theme().colors().element_background)
             .rounded_md()
+            .child(
+                Icon::new(self.context.icon)
+                    .size(IconSize::XSmall)
+                    .color(Color::Muted),
+            )
             .child(Label::new(self.context.name.clone()).size(LabelSize::Small))
             .when_some(self.on_remove, |parent, on_remove| {
                 parent.child(