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 {
374 self.add_suggested_context(&suggested, cx);
375 }
376 }
377
378 fn add_suggested_context(&mut self, suggested: &SuggestedContext, cx: &mut Context<Self>) {
379 self.context_store.update(cx, |context_store, cx| {
380 context_store.add_suggested_context(suggested, cx)
381 });
382 cx.notify();
383 }
384}
385
386impl Focusable for ContextStrip {
387 fn focus_handle(&self, _cx: &App) -> FocusHandle {
388 self.focus_handle.clone()
389 }
390}
391
392impl Render for ContextStrip {
393 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
394 let context_picker = self.context_picker.clone();
395 let focus_handle = self.focus_handle.clone();
396
397 let added_contexts = self.added_contexts(cx);
398 let dupe_names = added_contexts
399 .iter()
400 .map(|c| c.name.clone())
401 .sorted()
402 .tuple_windows()
403 .filter(|(a, b)| a == b)
404 .map(|(a, _)| a)
405 .collect::<HashSet<SharedString>>();
406 let no_added_context = added_contexts.is_empty();
407
408 let suggested_context = self.suggested_context(cx).map(|suggested_context| {
409 (
410 suggested_context,
411 self.is_suggested_focused(&added_contexts),
412 )
413 });
414
415 h_flex()
416 .flex_wrap()
417 .gap_1()
418 .track_focus(&focus_handle)
419 .key_context("ContextStrip")
420 .on_action(cx.listener(Self::focus_up))
421 .on_action(cx.listener(Self::focus_right))
422 .on_action(cx.listener(Self::focus_down))
423 .on_action(cx.listener(Self::focus_left))
424 .on_action(cx.listener(Self::remove_focused_context))
425 .on_action(cx.listener(Self::accept_suggested_context))
426 .on_children_prepainted({
427 let entity = cx.entity().downgrade();
428 move |children_bounds, _window, cx| {
429 entity
430 .update(cx, |this, _| {
431 this.children_bounds = Some(children_bounds);
432 })
433 .ok();
434 }
435 })
436 .child(
437 PopoverMenu::new("context-picker")
438 .menu({
439 let context_picker = context_picker.clone();
440 move |window, cx| {
441 context_picker.update(cx, |this, cx| {
442 this.init(window, cx);
443 });
444
445 Some(context_picker.clone())
446 }
447 })
448 .on_open({
449 let context_picker = context_picker.downgrade();
450 Rc::new(move |window, cx| {
451 context_picker
452 .update(cx, |context_picker, cx| {
453 context_picker.select_first(window, cx);
454 })
455 .ok();
456 })
457 })
458 .trigger_with_tooltip(
459 IconButton::new("add-context", IconName::Plus)
460 .icon_size(IconSize::Small)
461 .style(ui::ButtonStyle::Filled),
462 {
463 let focus_handle = focus_handle.clone();
464 move |window, cx| {
465 Tooltip::for_action_in(
466 "Add Context",
467 &ToggleContextPicker,
468 &focus_handle,
469 window,
470 cx,
471 )
472 }
473 },
474 )
475 .attach(gpui::Corner::TopLeft)
476 .anchor(gpui::Corner::BottomLeft)
477 .offset(gpui::Point {
478 x: px(0.0),
479 y: px(-2.0),
480 })
481 .with_handle(self.context_picker_menu_handle.clone()),
482 )
483 .children(
484 added_contexts
485 .into_iter()
486 .enumerate()
487 .map(|(i, added_context)| {
488 let name = added_context.name.clone();
489 let context = added_context.handle.clone();
490 ContextPill::added(
491 added_context,
492 dupe_names.contains(&name),
493 self.focused_index == Some(i),
494 Some({
495 let context = context.clone();
496 let context_store = self.context_store.clone();
497 Rc::new(cx.listener(move |_this, _event, _window, cx| {
498 context_store.update(cx, |this, cx| {
499 this.remove_context(&context, cx);
500 });
501 cx.notify();
502 }))
503 }),
504 )
505 .on_click({
506 Rc::new(cx.listener(move |this, event: &ClickEvent, window, cx| {
507 if event.click_count() > 1 {
508 this.open_context(&context, window, cx);
509 } else {
510 this.focused_index = Some(i);
511 }
512 cx.notify();
513 }))
514 })
515 }),
516 )
517 .when_some(suggested_context, |el, (suggested, focused)| {
518 el.child(
519 ContextPill::suggested(
520 suggested.name().clone(),
521 suggested.icon_path(),
522 suggested.kind(),
523 focused,
524 )
525 .on_click(Rc::new(cx.listener(
526 move |this, _event, _window, cx| {
527 this.add_suggested_context(&suggested, cx);
528 },
529 ))),
530 )
531 })
532 .when(!no_added_context, {
533 move |parent| {
534 parent.child(
535 IconButton::new("remove-all-context", IconName::Eraser)
536 .icon_size(IconSize::Small)
537 .tooltip({
538 let focus_handle = focus_handle.clone();
539 move |window, cx| {
540 Tooltip::for_action_in(
541 "Remove All Context",
542 &RemoveAllContext,
543 &focus_handle,
544 window,
545 cx,
546 )
547 }
548 })
549 .on_click(cx.listener({
550 let focus_handle = focus_handle.clone();
551 move |_this, _event, window, cx| {
552 focus_handle.dispatch_action(&RemoveAllContext, window, cx);
553 }
554 })),
555 )
556 }
557 })
558 .into_any()
559 }
560}
561
562pub enum ContextStripEvent {
563 PickerDismissed,
564 BlurredEmpty,
565 BlurredDown,
566 BlurredUp,
567}
568
569impl EventEmitter<ContextStripEvent> for ContextStrip {}
570
571pub enum SuggestContextKind {
572 File,
573 Thread,
574}