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