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