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