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