1use std::{
2 any::TypeId,
3 sync::{
4 atomic::{AtomicBool, Ordering::SeqCst},
5 Arc,
6 },
7};
8
9use anyhow::{anyhow, Result};
10use collections::HashMap;
11use editor::Editor;
12use futures::future::join_all;
13use gpui::{AnyView, Render, Task, View, WeakView};
14use ui::{prelude::*, ButtonLike, Tooltip, WindowContext};
15use util::{maybe, ResultExt};
16use workspace::Workspace;
17
18/// A collected attachment from running an attachment tool
19pub struct UserAttachment {
20 pub message: Option<String>,
21 pub view: AnyView,
22}
23
24pub struct UserAttachmentStore {
25 attachment_tools: HashMap<TypeId, DynamicAttachment>,
26}
27
28/// Internal representation of an attachment tool to allow us to treat them dynamically
29struct DynamicAttachment {
30 enabled: AtomicBool,
31 call: Box<dyn Fn(&mut WindowContext) -> Task<Result<UserAttachment>>>,
32}
33
34impl UserAttachmentStore {
35 pub fn new() -> Self {
36 Self {
37 attachment_tools: HashMap::default(),
38 }
39 }
40
41 pub fn register<A: AttachmentTool + 'static>(&mut self, attachment: A) {
42 let call = Box::new(move |cx: &mut WindowContext| {
43 let result = attachment.run(cx);
44
45 cx.spawn(move |mut cx| async move {
46 let result: Result<A::Output> = result.await;
47 let message = A::format(&result);
48 let view = cx.update(|cx| A::view(result, cx))?;
49
50 Ok(UserAttachment {
51 message,
52 view: view.into(),
53 })
54 })
55 });
56
57 self.attachment_tools.insert(
58 TypeId::of::<A>(),
59 DynamicAttachment {
60 call,
61 enabled: AtomicBool::new(true),
62 },
63 );
64 }
65
66 pub fn set_attachment_tool_enabled<A: AttachmentTool + 'static>(&self, is_enabled: bool) {
67 if let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) {
68 attachment.enabled.store(is_enabled, SeqCst);
69 }
70 }
71
72 pub fn is_attachment_tool_enabled<A: AttachmentTool + 'static>(&self) -> bool {
73 if let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) {
74 attachment.enabled.load(SeqCst)
75 } else {
76 false
77 }
78 }
79
80 pub fn call<A: AttachmentTool + 'static>(
81 &self,
82 cx: &mut WindowContext,
83 ) -> Task<Result<UserAttachment>> {
84 let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) else {
85 return Task::ready(Err(anyhow!("no attachment tool")));
86 };
87
88 (attachment.call)(cx)
89 }
90
91 pub fn call_all_attachment_tools(
92 self: Arc<Self>,
93 cx: &mut WindowContext<'_>,
94 ) -> Task<Result<Vec<UserAttachment>>> {
95 let this = self.clone();
96 cx.spawn(|mut cx| async move {
97 let attachment_tasks = cx.update(|cx| {
98 let mut tasks = Vec::new();
99 for attachment in this
100 .attachment_tools
101 .values()
102 .filter(|attachment| attachment.enabled.load(SeqCst))
103 {
104 tasks.push((attachment.call)(cx))
105 }
106
107 tasks
108 })?;
109
110 let attachments = join_all(attachment_tasks.into_iter()).await;
111
112 Ok(attachments
113 .into_iter()
114 .filter_map(|attachment| attachment.log_err())
115 .collect())
116 })
117 }
118}
119
120pub trait AttachmentTool {
121 type Output: 'static;
122 type View: Render;
123
124 fn run(&self, cx: &mut WindowContext) -> Task<Result<Self::Output>>;
125
126 fn format(output: &Result<Self::Output>) -> Option<String>;
127
128 fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View>;
129}
130
131pub struct ActiveEditorAttachment {
132 filename: Arc<str>,
133 language: Arc<str>,
134 text: Arc<str>,
135}
136
137pub struct FileAttachmentView {
138 output: Result<ActiveEditorAttachment>,
139}
140
141impl Render for FileAttachmentView {
142 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
143 match &self.output {
144 Ok(attachment) => {
145 let filename = attachment.filename.clone();
146
147 // todo!(): make the button link to the actual file to open
148 ButtonLike::new("file-attachment")
149 .child(
150 h_flex()
151 .gap_1()
152 .bg(cx.theme().colors().editor_background)
153 .rounded_md()
154 .child(ui::Icon::new(IconName::File))
155 .child(filename.to_string()),
156 )
157 .tooltip({
158 move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx)
159 })
160 .into_any_element()
161 }
162 Err(err) => div().child(err.to_string()).into_any_element(),
163 }
164 }
165}
166
167pub struct ActiveEditorAttachmentTool {
168 workspace: WeakView<Workspace>,
169}
170
171impl ActiveEditorAttachmentTool {
172 pub fn new(workspace: WeakView<Workspace>, _cx: &mut WindowContext) -> Self {
173 Self { workspace }
174 }
175}
176
177impl AttachmentTool for ActiveEditorAttachmentTool {
178 type Output = ActiveEditorAttachment;
179 type View = FileAttachmentView;
180
181 fn run(&self, cx: &mut WindowContext) -> Task<Result<ActiveEditorAttachment>> {
182 Task::ready(maybe!({
183 let active_buffer = self
184 .workspace
185 .update(cx, |workspace, cx| {
186 workspace
187 .active_item(cx)
188 .and_then(|item| Some(item.act_as::<Editor>(cx)?.read(cx).buffer().clone()))
189 })?
190 .ok_or_else(|| anyhow!("no active buffer"))?;
191
192 let buffer = active_buffer.read(cx);
193
194 if let Some(singleton) = buffer.as_singleton() {
195 let singleton = singleton.read(cx);
196
197 let filename = singleton
198 .file()
199 .map(|file| file.path().to_string_lossy())
200 .unwrap_or("Untitled".into());
201
202 let text = singleton.text();
203
204 let language = singleton
205 .language()
206 .map(|l| {
207 let name = l.code_fence_block_name();
208 name.to_string()
209 })
210 .unwrap_or_default();
211
212 return Ok(ActiveEditorAttachment {
213 filename: filename.into(),
214 language: language.into(),
215 text: text.into(),
216 });
217 }
218
219 Err(anyhow!("no active buffer"))
220 }))
221 }
222
223 fn format(output: &Result<Self::Output>) -> Option<String> {
224 let output = output.as_ref().ok()?;
225
226 let filename = &output.filename;
227 let language = &output.language;
228 let text = &output.text;
229
230 Some(format!(
231 "User's active file `{filename}`:\n\n```{language}\n{text}```\n\n"
232 ))
233 }
234
235 fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View> {
236 cx.new_view(|_cx| FileAttachmentView { output })
237 }
238}