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