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