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
120///
121pub trait AttachmentTool {
122 type Output: 'static;
123 type View: Render;
124
125 fn run(&self, cx: &mut WindowContext) -> Task<Result<Self::Output>>;
126
127 fn format(output: &Result<Self::Output>) -> Option<String>;
128
129 fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View>;
130}
131
132pub struct ActiveEditorAttachment {
133 filename: Arc<str>,
134 language: Arc<str>,
135 text: Arc<str>,
136}
137
138pub struct FileAttachmentView {
139 output: Result<ActiveEditorAttachment>,
140}
141
142impl Render for FileAttachmentView {
143 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
144 match &self.output {
145 Ok(attachment) => {
146 let filename = attachment.filename.clone();
147
148 // todo!(): make the button link to the actual file to open
149 ButtonLike::new("file-attachment")
150 .child(
151 h_flex()
152 .gap_1()
153 .bg(cx.theme().colors().editor_background)
154 .rounded_md()
155 .child(ui::Icon::new(IconName::File))
156 .child(filename.to_string()),
157 )
158 .tooltip({
159 move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx)
160 })
161 .into_any_element()
162 }
163 // todo!(): show a better error view when the file attaching didn't work
164 Err(err) => div().child(err.to_string()).into_any_element(),
165 }
166 }
167}
168
169pub struct ActiveEditorAttachmentTool {
170 workspace: WeakView<Workspace>,
171}
172
173impl ActiveEditorAttachmentTool {
174 pub fn new(workspace: WeakView<Workspace>, _cx: &mut WindowContext) -> Self {
175 Self { workspace }
176 }
177}
178
179impl AttachmentTool for ActiveEditorAttachmentTool {
180 type Output = ActiveEditorAttachment;
181 type View = FileAttachmentView;
182
183 fn run(&self, cx: &mut WindowContext) -> Task<Result<ActiveEditorAttachment>> {
184 Task::ready(maybe!({
185 let active_buffer = self
186 .workspace
187 .update(cx, |workspace, cx| {
188 workspace
189 .active_item(cx)
190 .and_then(|item| Some(item.act_as::<Editor>(cx)?.read(cx).buffer().clone()))
191 })?
192 .ok_or_else(|| anyhow!("no active buffer"))?;
193
194 let buffer = active_buffer.read(cx);
195
196 if let Some(singleton) = buffer.as_singleton() {
197 let singleton = singleton.read(cx);
198
199 let filename = singleton
200 .file()
201 .map(|file| file.path().to_string_lossy())
202 .unwrap_or("Untitled".into());
203
204 let text = singleton.text();
205
206 let language = singleton
207 .language()
208 .map(|l| {
209 let name = l.code_fence_block_name();
210 name.to_string()
211 })
212 .unwrap_or_default();
213
214 return Ok(ActiveEditorAttachment {
215 filename: filename.into(),
216 language: language.into(),
217 text: text.into(),
218 });
219 }
220
221 Err(anyhow!("no active buffer"))
222 }))
223 }
224
225 fn format(output: &Result<Self::Output>) -> Option<String> {
226 let output = output.as_ref().ok()?;
227
228 let filename = &output.filename;
229 let language = &output.language;
230 let text = &output.text;
231
232 Some(format!(
233 "User's active file `{filename}`:\n\n```{language}\n{text}```\n\n"
234 ))
235 }
236
237 fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View> {
238 cx.new_view(|_cx| FileAttachmentView { output })
239 }
240}