1use std::path::Path;
2use std::rc::Rc;
3
4use collections::HashSet;
5use editor::Editor;
6use file_icons::FileIcons;
7use gpui::{
8 App, Bounds, ClickEvent, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable,
9 Subscription, WeakEntity,
10};
11use itertools::Itertools;
12use language::Buffer;
13use project::ProjectItem;
14use ui::{KeyBinding, PopoverMenu, PopoverMenuHandle, Tooltip, prelude::*};
15use workspace::Workspace;
16
17use crate::context::{AgentContext, ContextKind};
18use crate::context_picker::ContextPicker;
19use crate::context_store::ContextStore;
20use crate::thread::Thread;
21use crate::thread_store::ThreadStore;
22use crate::ui::{AddedContext, ContextPill};
23use crate::{
24 AcceptSuggestedContext, AssistantPanel, FocusDown, FocusLeft, FocusRight, FocusUp,
25 RemoveAllContext, RemoveFocusedContext, ToggleContextPicker,
26};
27
28pub struct ContextStrip {
29 context_store: Entity<ContextStore>,
30 context_picker: Entity<ContextPicker>,
31 context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
32 focus_handle: FocusHandle,
33 suggest_context_kind: SuggestContextKind,
34 workspace: WeakEntity<Workspace>,
35 thread_store: Option<WeakEntity<ThreadStore>>,
36 _subscriptions: Vec<Subscription>,
37 focused_index: Option<usize>,
38 children_bounds: Option<Vec<Bounds<Pixels>>>,
39}
40
41impl ContextStrip {
42 pub fn new(
43 context_store: Entity<ContextStore>,
44 workspace: WeakEntity<Workspace>,
45 thread_store: Option<WeakEntity<ThreadStore>>,
46 context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
47 suggest_context_kind: SuggestContextKind,
48 window: &mut Window,
49 cx: &mut Context<Self>,
50 ) -> Self {
51 let context_picker = cx.new(|cx| {
52 ContextPicker::new(
53 workspace.clone(),
54 thread_store.clone(),
55 context_store.downgrade(),
56 window,
57 cx,
58 )
59 });
60
61 let focus_handle = cx.focus_handle();
62
63 let subscriptions = vec![
64 cx.observe(&context_store, |_, _, cx| cx.notify()),
65 cx.subscribe_in(&context_picker, window, Self::handle_context_picker_event),
66 cx.on_focus(&focus_handle, window, Self::handle_focus),
67 cx.on_blur(&focus_handle, window, Self::handle_blur),
68 ];
69
70 Self {
71 context_store: context_store.clone(),
72 context_picker,
73 context_picker_menu_handle,
74 focus_handle,
75 suggest_context_kind,
76 workspace,
77 thread_store,
78 _subscriptions: subscriptions,
79 focused_index: None,
80 children_bounds: None,
81 }
82 }
83
84 fn added_contexts(&self, cx: &App) -> Vec<AddedContext> {
85 if let Some(workspace) = self.workspace.upgrade() {
86 let project = workspace.read(cx).project().read(cx);
87 let prompt_store = self
88 .thread_store
89 .as_ref()
90 .and_then(|thread_store| thread_store.upgrade())
91 .and_then(|thread_store| thread_store.read(cx).prompt_store().as_ref());
92 self.context_store
93 .read(cx)
94 .context()
95 .flat_map(|context| AddedContext::new(context.clone(), prompt_store, project, cx))
96 .collect::<Vec<_>>()
97 } else {
98 Vec::new()
99 }
100 }
101
102 fn suggested_context(&self, cx: &Context<Self>) -> Option<SuggestedContext> {
103 match self.suggest_context_kind {
104 SuggestContextKind::File => self.suggested_file(cx),
105 SuggestContextKind::Thread => self.suggested_thread(cx),
106 }
107 }
108
109 fn suggested_file(&self, cx: &Context<Self>) -> Option<SuggestedContext> {
110 let workspace = self.workspace.upgrade()?;
111 let active_item = workspace.read(cx).active_item(cx)?;
112
113 let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
114 let active_buffer_entity = editor.buffer().read(cx).as_singleton()?;
115 let active_buffer = active_buffer_entity.read(cx);
116 let project_path = active_buffer.project_path(cx)?;
117
118 if self
119 .context_store
120 .read(cx)
121 .file_path_included(&project_path, cx)
122 .is_some()
123 {
124 return None;
125 }
126
127 let file_name = active_buffer.file()?.file_name(cx);
128 let icon_path = FileIcons::get_icon(&Path::new(&file_name), cx);
129 Some(SuggestedContext::File {
130 name: file_name.to_string_lossy().into_owned().into(),
131 buffer: active_buffer_entity.downgrade(),
132 icon_path,
133 })
134 }
135
136 fn suggested_thread(&self, cx: &Context<Self>) -> Option<SuggestedContext> {
137 if !self.context_picker.read(cx).allow_threads() {
138 return None;
139 }
140
141 let workspace = self.workspace.upgrade()?;
142 let active_thread = workspace
143 .read(cx)
144 .panel::<AssistantPanel>(cx)?
145 .read(cx)
146 .active_thread(cx);
147 let weak_active_thread = active_thread.downgrade();
148
149 let active_thread = active_thread.read(cx);
150
151 if self
152 .context_store
153 .read(cx)
154 .includes_thread(active_thread.id())
155 {
156 return None;
157 }
158
159 Some(SuggestedContext::Thread {
160 name: active_thread.summary_or_default(),
161 thread: weak_active_thread,
162 })
163 }
164
165 fn handle_context_picker_event(
166 &mut self,
167 _picker: &Entity<ContextPicker>,
168 _event: &DismissEvent,
169 _window: &mut Window,
170 cx: &mut Context<Self>,
171 ) {
172 cx.emit(ContextStripEvent::PickerDismissed);
173 }
174
175 fn handle_focus(&mut self, _window: &mut Window, cx: &mut Context<Self>) {
176 self.focused_index = self.last_pill_index();
177 cx.notify();
178 }
179
180 fn handle_blur(&mut self, _window: &mut Window, cx: &mut Context<Self>) {
181 self.focused_index = None;
182 cx.notify();
183 }
184
185 fn focus_left(&mut self, _: &FocusLeft, _window: &mut Window, cx: &mut Context<Self>) {
186 self.focused_index = match self.focused_index {
187 Some(index) if index > 0 => Some(index - 1),
188 _ => self.last_pill_index(),
189 };
190
191 cx.notify();
192 }
193
194 fn focus_right(&mut self, _: &FocusRight, _window: &mut Window, cx: &mut Context<Self>) {
195 let Some(last_index) = self.last_pill_index() else {
196 return;
197 };
198
199 self.focused_index = match self.focused_index {
200 Some(index) if index < last_index => Some(index + 1),
201 _ => Some(0),
202 };
203
204 cx.notify();
205 }
206
207 fn focus_up(&mut self, _: &FocusUp, _window: &mut Window, cx: &mut Context<Self>) {
208 let Some(focused_index) = self.focused_index else {
209 return;
210 };
211
212 if focused_index == 0 {
213 return cx.emit(ContextStripEvent::BlurredUp);
214 }
215
216 let Some((focused, pills)) = self.focused_bounds(focused_index) else {
217 return;
218 };
219
220 let iter = pills[..focused_index].iter().enumerate().rev();
221 self.focused_index = Self::find_best_horizontal_match(focused, iter).or(Some(0));
222 cx.notify();
223 }
224
225 fn focus_down(&mut self, _: &FocusDown, _window: &mut Window, cx: &mut Context<Self>) {
226 let Some(focused_index) = self.focused_index else {
227 return;
228 };
229
230 let last_index = self.last_pill_index();
231
232 if self.focused_index == last_index {
233 return cx.emit(ContextStripEvent::BlurredDown);
234 }
235
236 let Some((focused, pills)) = self.focused_bounds(focused_index) else {
237 return;
238 };
239
240 let iter = pills.iter().enumerate().skip(focused_index + 1);
241 self.focused_index = Self::find_best_horizontal_match(focused, iter).or(last_index);
242 cx.notify();
243 }
244
245 fn focused_bounds(&self, focused: usize) -> Option<(&Bounds<Pixels>, &[Bounds<Pixels>])> {
246 let pill_bounds = self.pill_bounds()?;
247 let focused = pill_bounds.get(focused)?;
248
249 Some((focused, pill_bounds))
250 }
251
252 fn pill_bounds(&self) -> Option<&[Bounds<Pixels>]> {
253 let bounds = self.children_bounds.as_ref()?;
254 let eraser = if bounds.len() < 3 { 0 } else { 1 };
255 let pills = &bounds[1..bounds.len() - eraser];
256
257 if pills.is_empty() { None } else { Some(pills) }
258 }
259
260 fn last_pill_index(&self) -> Option<usize> {
261 Some(self.pill_bounds()?.len() - 1)
262 }
263
264 fn find_best_horizontal_match<'a>(
265 focused: &'a Bounds<Pixels>,
266 iter: impl Iterator<Item = (usize, &'a Bounds<Pixels>)>,
267 ) -> Option<usize> {
268 let mut best = None;
269
270 let focused_left = focused.left();
271 let focused_right = focused.right();
272
273 for (index, probe) in iter {
274 if probe.origin.y == focused.origin.y {
275 continue;
276 }
277
278 let overlap = probe.right().min(focused_right) - probe.left().max(focused_left);
279
280 best = match best {
281 Some((_, prev_overlap, y)) if probe.origin.y != y || prev_overlap > overlap => {
282 break;
283 }
284 Some(_) | None => Some((index, overlap, probe.origin.y)),
285 };
286 }
287
288 best.map(|(index, _, _)| index)
289 }
290
291 fn open_context(&mut self, context: &AgentContext, window: &mut Window, cx: &mut App) {
292 let Some(workspace) = self.workspace.upgrade() else {
293 return;
294 };
295
296 crate::active_thread::open_context(context, workspace, window, cx);
297 }
298
299 fn remove_focused_context(
300 &mut self,
301 _: &RemoveFocusedContext,
302 _window: &mut Window,
303 cx: &mut Context<Self>,
304 ) {
305 if let Some(index) = self.focused_index {
306 let added_contexts = self.added_contexts(cx);
307 let Some(context) = added_contexts.get(index) else {
308 return;
309 };
310
311 self.context_store.update(cx, |this, cx| {
312 this.remove_context(&context.context, cx);
313 });
314
315 let is_now_empty = added_contexts.len() == 1;
316 if is_now_empty {
317 cx.emit(ContextStripEvent::BlurredEmpty);
318 } else {
319 self.focused_index = Some(index.saturating_sub(1));
320 cx.notify();
321 }
322 }
323 }
324
325 fn is_suggested_focused(&self, added_contexts: &Vec<AddedContext>) -> bool {
326 // We only suggest one item after the actual context
327 self.focused_index == Some(added_contexts.len())
328 }
329
330 fn accept_suggested_context(
331 &mut self,
332 _: &AcceptSuggestedContext,
333 _window: &mut Window,
334 cx: &mut Context<Self>,
335 ) {
336 if let Some(suggested) = self.suggested_context(cx) {
337 if self.is_suggested_focused(&self.added_contexts(cx)) {
338 self.add_suggested_context(&suggested, cx);
339 }
340 }
341 }
342
343 fn add_suggested_context(&mut self, suggested: &SuggestedContext, cx: &mut Context<Self>) {
344 self.context_store.update(cx, |context_store, cx| {
345 context_store.add_suggested_context(&suggested, cx)
346 });
347 cx.notify();
348 }
349}
350
351impl Focusable for ContextStrip {
352 fn focus_handle(&self, _cx: &App) -> FocusHandle {
353 self.focus_handle.clone()
354 }
355}
356
357impl Render for ContextStrip {
358 fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
359 let context_picker = self.context_picker.clone();
360 let focus_handle = self.focus_handle.clone();
361
362 let added_contexts = self.added_contexts(cx);
363 let dupe_names = added_contexts
364 .iter()
365 .map(|c| c.name.clone())
366 .sorted()
367 .tuple_windows()
368 .filter(|(a, b)| a == b)
369 .map(|(a, _)| a)
370 .collect::<HashSet<SharedString>>();
371 let no_added_context = added_contexts.is_empty();
372
373 let suggested_context = self.suggested_context(cx).map(|suggested_context| {
374 (
375 suggested_context,
376 self.is_suggested_focused(&added_contexts),
377 )
378 });
379
380 h_flex()
381 .flex_wrap()
382 .gap_1()
383 .track_focus(&focus_handle)
384 .key_context("ContextStrip")
385 .on_action(cx.listener(Self::focus_up))
386 .on_action(cx.listener(Self::focus_right))
387 .on_action(cx.listener(Self::focus_down))
388 .on_action(cx.listener(Self::focus_left))
389 .on_action(cx.listener(Self::remove_focused_context))
390 .on_action(cx.listener(Self::accept_suggested_context))
391 .on_children_prepainted({
392 let entity = cx.entity().downgrade();
393 move |children_bounds, _window, cx| {
394 entity
395 .update(cx, |this, _| {
396 this.children_bounds = Some(children_bounds);
397 })
398 .ok();
399 }
400 })
401 .child(
402 PopoverMenu::new("context-picker")
403 .menu(move |window, cx| {
404 context_picker.update(cx, |this, cx| {
405 this.init(window, cx);
406 });
407
408 Some(context_picker.clone())
409 })
410 .trigger_with_tooltip(
411 IconButton::new("add-context", IconName::Plus)
412 .icon_size(IconSize::Small)
413 .style(ui::ButtonStyle::Filled),
414 {
415 let focus_handle = focus_handle.clone();
416 move |window, cx| {
417 Tooltip::for_action_in(
418 "Add Context",
419 &ToggleContextPicker,
420 &focus_handle,
421 window,
422 cx,
423 )
424 }
425 },
426 )
427 .attach(gpui::Corner::TopLeft)
428 .anchor(gpui::Corner::BottomLeft)
429 .offset(gpui::Point {
430 x: px(0.0),
431 y: px(-2.0),
432 })
433 .with_handle(self.context_picker_menu_handle.clone()),
434 )
435 .when(no_added_context && suggested_context.is_none(), {
436 |parent| {
437 parent.child(
438 h_flex()
439 .ml_1p5()
440 .gap_2()
441 .child(
442 Label::new("Add Context")
443 .size(LabelSize::Small)
444 .color(Color::Muted),
445 )
446 .opacity(0.5)
447 .children(
448 KeyBinding::for_action_in(
449 &ToggleContextPicker,
450 &focus_handle,
451 window,
452 cx,
453 )
454 .map(|binding| binding.into_any_element()),
455 ),
456 )
457 }
458 })
459 .children(
460 added_contexts
461 .into_iter()
462 .enumerate()
463 .map(|(i, added_context)| {
464 let name = added_context.name.clone();
465 let context = added_context.context.clone();
466 ContextPill::added(
467 added_context,
468 dupe_names.contains(&name),
469 self.focused_index == Some(i),
470 Some({
471 let context = context.clone();
472 let context_store = self.context_store.clone();
473 Rc::new(cx.listener(move |_this, _event, _window, cx| {
474 context_store.update(cx, |this, cx| {
475 this.remove_context(&context, cx);
476 });
477 cx.notify();
478 }))
479 }),
480 )
481 .on_click({
482 Rc::new(cx.listener(move |this, event: &ClickEvent, window, cx| {
483 if event.down.click_count > 1 {
484 this.open_context(&context, window, cx);
485 } else {
486 this.focused_index = Some(i);
487 }
488 cx.notify();
489 }))
490 })
491 }),
492 )
493 .when_some(suggested_context, |el, (suggested, focused)| {
494 el.child(
495 ContextPill::suggested(
496 suggested.name().clone(),
497 suggested.icon_path(),
498 suggested.kind(),
499 focused,
500 )
501 .on_click(Rc::new(cx.listener(
502 move |this, _event, _window, cx| {
503 this.add_suggested_context(&suggested, cx);
504 },
505 ))),
506 )
507 })
508 .when(!no_added_context, {
509 move |parent| {
510 parent.child(
511 IconButton::new("remove-all-context", IconName::Eraser)
512 .icon_size(IconSize::Small)
513 .tooltip({
514 let focus_handle = focus_handle.clone();
515 move |window, cx| {
516 Tooltip::for_action_in(
517 "Remove All Context",
518 &RemoveAllContext,
519 &focus_handle,
520 window,
521 cx,
522 )
523 }
524 })
525 .on_click(cx.listener({
526 let focus_handle = focus_handle.clone();
527 move |_this, _event, window, cx| {
528 focus_handle.dispatch_action(&RemoveAllContext, window, cx);
529 }
530 })),
531 )
532 }
533 })
534 .into_any()
535 }
536}
537
538pub enum ContextStripEvent {
539 PickerDismissed,
540 BlurredEmpty,
541 BlurredDown,
542 BlurredUp,
543}
544
545impl EventEmitter<ContextStripEvent> for ContextStrip {}
546
547pub enum SuggestContextKind {
548 File,
549 Thread,
550}
551
552#[derive(Clone)]
553pub enum SuggestedContext {
554 File {
555 name: SharedString,
556 icon_path: Option<SharedString>,
557 buffer: WeakEntity<Buffer>,
558 },
559 Thread {
560 name: SharedString,
561 thread: WeakEntity<Thread>,
562 },
563}
564
565impl SuggestedContext {
566 pub fn name(&self) -> &SharedString {
567 match self {
568 Self::File { name, .. } => name,
569 Self::Thread { name, .. } => name,
570 }
571 }
572
573 pub fn icon_path(&self) -> Option<SharedString> {
574 match self {
575 Self::File { icon_path, .. } => icon_path.clone(),
576 Self::Thread { .. } => None,
577 }
578 }
579
580 pub fn kind(&self) -> ContextKind {
581 match self {
582 Self::File { .. } => ContextKind::File,
583 Self::Thread { .. } => ContextKind::Thread,
584 }
585 }
586}