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