1use std::rc::Rc;
2
3use anyhow::Result;
4use collections::HashSet;
5use editor::Editor;
6use file_icons::FileIcons;
7use gpui::{
8 DismissEvent, EventEmitter, FocusHandle, Model, ModelContext, Subscription, Task, View,
9 WeakModel, WeakView,
10};
11use itertools::Itertools;
12use language::Buffer;
13use ui::{prelude::*, KeyBinding, PopoverMenu, PopoverMenuHandle, Tooltip};
14use workspace::Workspace;
15
16use crate::context::ContextKind;
17use crate::context_picker::{ConfirmBehavior, ContextPicker};
18use crate::context_store::ContextStore;
19use crate::thread::Thread;
20use crate::thread_store::ThreadStore;
21use crate::ui::ContextPill;
22use crate::{AssistantPanel, RemoveAllContext, ToggleContextPicker};
23
24pub struct ContextStrip {
25 context_store: Model<ContextStore>,
26 pub context_picker: View<ContextPicker>,
27 context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
28 focus_handle: FocusHandle,
29 suggest_context_kind: SuggestContextKind,
30 workspace: WeakView<Workspace>,
31 _context_picker_subscription: Subscription,
32}
33
34impl ContextStrip {
35 pub fn new(
36 context_store: Model<ContextStore>,
37 workspace: WeakView<Workspace>,
38 thread_store: Option<WeakModel<ThreadStore>>,
39 focus_handle: FocusHandle,
40 context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
41 suggest_context_kind: SuggestContextKind,
42 cx: &mut ViewContext<Self>,
43 ) -> Self {
44 let context_picker = cx.new_view(|cx| {
45 ContextPicker::new(
46 workspace.clone(),
47 thread_store.clone(),
48 context_store.downgrade(),
49 ConfirmBehavior::KeepOpen,
50 cx,
51 )
52 });
53
54 let context_picker_subscription =
55 cx.subscribe(&context_picker, Self::handle_context_picker_event);
56
57 Self {
58 context_store: context_store.clone(),
59 context_picker,
60 context_picker_menu_handle,
61 focus_handle,
62 suggest_context_kind,
63 workspace,
64 _context_picker_subscription: context_picker_subscription,
65 }
66 }
67
68 fn suggested_context(&self, cx: &ViewContext<Self>) -> Option<SuggestedContext> {
69 match self.suggest_context_kind {
70 SuggestContextKind::File => self.suggested_file(cx),
71 SuggestContextKind::Thread => self.suggested_thread(cx),
72 }
73 }
74
75 fn suggested_file(&self, cx: &ViewContext<Self>) -> Option<SuggestedContext> {
76 let workspace = self.workspace.upgrade()?;
77 let active_item = workspace.read(cx).active_item(cx)?;
78
79 let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
80 let active_buffer_model = editor.buffer().read(cx).as_singleton()?;
81 let active_buffer = active_buffer_model.read(cx);
82
83 let path = active_buffer.file()?.path();
84
85 if self
86 .context_store
87 .read(cx)
88 .will_include_buffer(active_buffer.remote_id(), path)
89 .is_some()
90 {
91 return None;
92 }
93
94 let name = match path.file_name() {
95 Some(name) => name.to_string_lossy().into_owned().into(),
96 None => path.to_string_lossy().into_owned().into(),
97 };
98
99 let icon_path = FileIcons::get_icon(path, cx);
100
101 Some(SuggestedContext::File {
102 name,
103 buffer: active_buffer_model.downgrade(),
104 icon_path,
105 })
106 }
107
108 fn suggested_thread(&self, cx: &ViewContext<Self>) -> Option<SuggestedContext> {
109 let workspace = self.workspace.upgrade()?;
110 let active_thread = workspace
111 .read(cx)
112 .panel::<AssistantPanel>(cx)?
113 .read(cx)
114 .active_thread(cx);
115 let weak_active_thread = active_thread.downgrade();
116
117 let active_thread = active_thread.read(cx);
118
119 if self
120 .context_store
121 .read(cx)
122 .includes_thread(active_thread.id())
123 .is_some()
124 {
125 return None;
126 }
127
128 Some(SuggestedContext::Thread {
129 name: active_thread.summary_or_default(),
130 thread: weak_active_thread,
131 })
132 }
133
134 fn handle_context_picker_event(
135 &mut self,
136 _picker: View<ContextPicker>,
137 _event: &DismissEvent,
138 cx: &mut ViewContext<Self>,
139 ) {
140 cx.emit(ContextStripEvent::PickerDismissed);
141 }
142}
143
144impl Render for ContextStrip {
145 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
146 let context_store = self.context_store.read(cx);
147 let context = context_store
148 .context()
149 .iter()
150 .flat_map(|context| context.snapshot(cx))
151 .collect::<Vec<_>>();
152 let context_picker = self.context_picker.clone();
153 let focus_handle = self.focus_handle.clone();
154
155 let suggested_context = self.suggested_context(cx);
156
157 let dupe_names = context
158 .iter()
159 .map(|context| context.name.clone())
160 .sorted()
161 .tuple_windows()
162 .filter(|(a, b)| a == b)
163 .map(|(a, _)| a)
164 .collect::<HashSet<SharedString>>();
165
166 h_flex()
167 .flex_wrap()
168 .gap_1()
169 .child(
170 PopoverMenu::new("context-picker")
171 .menu(move |cx| {
172 context_picker.update(cx, |this, cx| {
173 this.reset_mode(cx);
174 });
175
176 Some(context_picker.clone())
177 })
178 .trigger(
179 IconButton::new("add-context", IconName::Plus)
180 .icon_size(IconSize::Small)
181 .style(ui::ButtonStyle::Filled)
182 .tooltip({
183 let focus_handle = focus_handle.clone();
184
185 move |cx| {
186 Tooltip::for_action_in(
187 "Add Context",
188 &ToggleContextPicker,
189 &focus_handle,
190 cx,
191 )
192 }
193 }),
194 )
195 .attach(gpui::Corner::TopLeft)
196 .anchor(gpui::Corner::BottomLeft)
197 .offset(gpui::Point {
198 x: px(0.0),
199 y: px(-2.0),
200 })
201 .with_handle(self.context_picker_menu_handle.clone()),
202 )
203 .when(context.is_empty() && suggested_context.is_none(), {
204 |parent| {
205 parent.child(
206 h_flex()
207 .ml_1p5()
208 .gap_2()
209 .child(
210 Label::new("Add Context")
211 .size(LabelSize::Small)
212 .color(Color::Muted),
213 )
214 .opacity(0.5)
215 .children(
216 KeyBinding::for_action_in(&ToggleContextPicker, &focus_handle, cx)
217 .map(|binding| binding.into_any_element()),
218 ),
219 )
220 }
221 })
222 .children(context.iter().map(|context| {
223 ContextPill::new_added(
224 context.clone(),
225 dupe_names.contains(&context.name),
226 Some({
227 let id = context.id;
228 let context_store = self.context_store.clone();
229 Rc::new(cx.listener(move |_this, _event, cx| {
230 context_store.update(cx, |this, _cx| {
231 this.remove_context(id);
232 });
233 cx.notify();
234 }))
235 }),
236 )
237 }))
238 .when_some(suggested_context, |el, suggested| {
239 el.child(ContextPill::new_suggested(
240 suggested.name().clone(),
241 suggested.icon_path(),
242 suggested.kind(),
243 {
244 let context_store = self.context_store.clone();
245 Rc::new(cx.listener(move |this, _event, cx| {
246 let task = context_store.update(cx, |context_store, cx| {
247 suggested.accept(context_store, cx)
248 });
249
250 let workspace = this.workspace.clone();
251 cx.spawn(|this, mut cx| async move {
252 match task.await {
253 Ok(()) => {
254 if let Some(this) = this.upgrade() {
255 this.update(&mut cx, |_, cx| cx.notify())?;
256 }
257 }
258 Err(err) => {
259 let Some(workspace) = workspace.upgrade() else {
260 return anyhow::Ok(());
261 };
262
263 workspace.update(&mut cx, |workspace, cx| {
264 workspace.show_error(&err, cx);
265 })?;
266 }
267 }
268 anyhow::Ok(())
269 })
270 .detach_and_log_err(cx);
271 }))
272 },
273 ))
274 })
275 .when(!context.is_empty(), {
276 move |parent| {
277 parent.child(
278 IconButton::new("remove-all-context", IconName::Eraser)
279 .icon_size(IconSize::Small)
280 .tooltip({
281 let focus_handle = focus_handle.clone();
282 move |cx| {
283 Tooltip::for_action_in(
284 "Remove All Context",
285 &RemoveAllContext,
286 &focus_handle,
287 cx,
288 )
289 }
290 })
291 .on_click(cx.listener({
292 let focus_handle = focus_handle.clone();
293 move |_this, _event, cx| {
294 focus_handle.dispatch_action(&RemoveAllContext, cx);
295 }
296 })),
297 )
298 }
299 })
300 }
301}
302
303pub enum ContextStripEvent {
304 PickerDismissed,
305}
306
307impl EventEmitter<ContextStripEvent> for ContextStrip {}
308
309pub enum SuggestContextKind {
310 File,
311 Thread,
312}
313
314#[derive(Clone)]
315pub enum SuggestedContext {
316 File {
317 name: SharedString,
318 icon_path: Option<SharedString>,
319 buffer: WeakModel<Buffer>,
320 },
321 Thread {
322 name: SharedString,
323 thread: WeakModel<Thread>,
324 },
325}
326
327impl SuggestedContext {
328 pub fn name(&self) -> &SharedString {
329 match self {
330 Self::File { name, .. } => name,
331 Self::Thread { name, .. } => name,
332 }
333 }
334
335 pub fn icon_path(&self) -> Option<SharedString> {
336 match self {
337 Self::File { icon_path, .. } => icon_path.clone(),
338 Self::Thread { .. } => None,
339 }
340 }
341
342 pub fn accept(
343 &self,
344 context_store: &mut ContextStore,
345 cx: &mut ModelContext<ContextStore>,
346 ) -> Task<Result<()>> {
347 match self {
348 Self::File {
349 buffer,
350 icon_path: _,
351 name: _,
352 } => {
353 if let Some(buffer) = buffer.upgrade() {
354 return context_store.add_file_from_buffer(buffer, cx);
355 };
356 }
357 Self::Thread { thread, name: _ } => {
358 if let Some(thread) = thread.upgrade() {
359 context_store.insert_thread(thread, cx);
360 };
361 }
362 }
363 Task::ready(Ok(()))
364 }
365
366 pub fn kind(&self) -> ContextKind {
367 match self {
368 Self::File { .. } => ContextKind::File,
369 Self::Thread { .. } => ContextKind::Thread,
370 }
371 }
372}