1use std::rc::Rc;
2
3use anyhow::Result;
4use collections::HashSet;
5use editor::Editor;
6use file_icons::FileIcons;
7use gpui::{
8 DismissEvent, EventEmitter, FocusHandle, Model, ModelContext, Subscription, Task, View,
9 WeakModel, WeakView,
10};
11use itertools::Itertools;
12use language::Buffer;
13use ui::{prelude::*, KeyBinding, PopoverMenu, PopoverMenuHandle, Tooltip};
14use workspace::Workspace;
15
16use crate::context::ContextKind;
17use crate::context_picker::{ConfirmBehavior, ContextPicker};
18use crate::context_store::ContextStore;
19use crate::thread::Thread;
20use crate::thread_store::ThreadStore;
21use crate::ui::ContextPill;
22use crate::{AssistantPanel, RemoveAllContext, ToggleContextPicker};
23
24pub struct ContextStrip {
25 context_store: Model<ContextStore>,
26 context_picker: View<ContextPicker>,
27 context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
28 focus_handle: FocusHandle,
29 suggest_context_kind: SuggestContextKind,
30 workspace: WeakView<Workspace>,
31 _context_picker_subscription: Subscription,
32}
33
34impl ContextStrip {
35 pub fn new(
36 context_store: Model<ContextStore>,
37 workspace: WeakView<Workspace>,
38 thread_store: Option<WeakModel<ThreadStore>>,
39 focus_handle: FocusHandle,
40 context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
41 suggest_context_kind: SuggestContextKind,
42 cx: &mut ViewContext<Self>,
43 ) -> Self {
44 let context_picker = cx.new_view(|cx| {
45 ContextPicker::new(
46 workspace.clone(),
47 thread_store.clone(),
48 context_store.downgrade(),
49 ConfirmBehavior::KeepOpen,
50 cx,
51 )
52 });
53
54 let context_picker_subscription =
55 cx.subscribe(&context_picker, Self::handle_context_picker_event);
56
57 Self {
58 context_store: context_store.clone(),
59 context_picker,
60 context_picker_menu_handle,
61 focus_handle,
62 suggest_context_kind,
63 workspace,
64 _context_picker_subscription: context_picker_subscription,
65 }
66 }
67
68 fn suggested_context(&self, cx: &ViewContext<Self>) -> Option<SuggestedContext> {
69 match self.suggest_context_kind {
70 SuggestContextKind::File => self.suggested_file(cx),
71 SuggestContextKind::Thread => self.suggested_thread(cx),
72 }
73 }
74
75 fn suggested_file(&self, cx: &ViewContext<Self>) -> Option<SuggestedContext> {
76 let workspace = self.workspace.upgrade()?;
77 let active_item = workspace.read(cx).active_item(cx)?;
78
79 let editor = active_item.to_any().downcast::<Editor>().ok()?.read(cx);
80 let active_buffer_model = editor.buffer().read(cx).as_singleton()?;
81 let active_buffer = active_buffer_model.read(cx);
82
83 let path = active_buffer.file()?.path();
84
85 if self
86 .context_store
87 .read(cx)
88 .will_include_buffer(active_buffer.remote_id(), path)
89 .is_some()
90 {
91 return None;
92 }
93
94 let name = match path.file_name() {
95 Some(name) => name.to_string_lossy().into_owned().into(),
96 None => path.to_string_lossy().into_owned().into(),
97 };
98
99 let icon_path = FileIcons::get_icon(path, cx);
100
101 Some(SuggestedContext::File {
102 name,
103 buffer: active_buffer_model.downgrade(),
104 icon_path,
105 })
106 }
107
108 fn suggested_thread(&self, cx: &ViewContext<Self>) -> Option<SuggestedContext> {
109 let workspace = self.workspace.upgrade()?;
110 let active_thread = workspace
111 .read(cx)
112 .panel::<AssistantPanel>(cx)?
113 .read(cx)
114 .active_thread(cx);
115 let weak_active_thread = active_thread.downgrade();
116
117 let active_thread = active_thread.read(cx);
118
119 if self
120 .context_store
121 .read(cx)
122 .includes_thread(active_thread.id())
123 .is_some()
124 {
125 return None;
126 }
127
128 Some(SuggestedContext::Thread {
129 name: active_thread.summary().unwrap_or("New Thread".into()),
130 thread: weak_active_thread,
131 })
132 }
133
134 fn handle_context_picker_event(
135 &mut self,
136 _picker: View<ContextPicker>,
137 _event: &DismissEvent,
138 cx: &mut ViewContext<Self>,
139 ) {
140 cx.emit(ContextStripEvent::PickerDismissed);
141 }
142}
143
144impl Render for ContextStrip {
145 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
146 let context_store = self.context_store.read(cx);
147 let context = context_store
148 .context()
149 .iter()
150 .flat_map(|context| context.snapshot(cx))
151 .collect::<Vec<_>>();
152 let context_picker = self.context_picker.clone();
153 let focus_handle = self.focus_handle.clone();
154
155 let suggested_context = self.suggested_context(cx);
156
157 let dupe_names = context
158 .iter()
159 .map(|context| context.name.clone())
160 .sorted()
161 .tuple_windows()
162 .filter(|(a, b)| a == b)
163 .map(|(a, _)| a)
164 .collect::<HashSet<SharedString>>();
165
166 h_flex()
167 .flex_wrap()
168 .gap_1()
169 .child(
170 PopoverMenu::new("context-picker")
171 .menu(move |_cx| Some(context_picker.clone()))
172 .trigger(
173 IconButton::new("add-context", IconName::Plus)
174 .icon_size(IconSize::Small)
175 .style(ui::ButtonStyle::Filled)
176 .tooltip({
177 let focus_handle = focus_handle.clone();
178
179 move |cx| {
180 Tooltip::for_action_in(
181 "Add Context",
182 &ToggleContextPicker,
183 &focus_handle,
184 cx,
185 )
186 }
187 }),
188 )
189 .attach(gpui::Corner::TopLeft)
190 .anchor(gpui::Corner::BottomLeft)
191 .offset(gpui::Point {
192 x: px(0.0),
193 y: px(-16.0),
194 })
195 .with_handle(self.context_picker_menu_handle.clone()),
196 )
197 .when(context.is_empty() && suggested_context.is_none(), {
198 |parent| {
199 parent.child(
200 h_flex()
201 .ml_1p5()
202 .gap_2()
203 .child(
204 Label::new("Add Context")
205 .size(LabelSize::Small)
206 .color(Color::Muted),
207 )
208 .opacity(0.5)
209 .children(
210 KeyBinding::for_action_in(&ToggleContextPicker, &focus_handle, cx)
211 .map(|binding| binding.into_any_element()),
212 ),
213 )
214 }
215 })
216 .children(context.iter().map(|context| {
217 ContextPill::new_added(
218 context.clone(),
219 dupe_names.contains(&context.name),
220 Some({
221 let id = context.id;
222 let context_store = self.context_store.clone();
223 Rc::new(cx.listener(move |_this, _event, cx| {
224 context_store.update(cx, |this, _cx| {
225 this.remove_context(id);
226 });
227 cx.notify();
228 }))
229 }),
230 )
231 }))
232 .when_some(suggested_context, |el, suggested| {
233 el.child(ContextPill::new_suggested(
234 suggested.name().clone(),
235 suggested.icon_path(),
236 suggested.kind(),
237 {
238 let context_store = self.context_store.clone();
239 Rc::new(cx.listener(move |this, _event, cx| {
240 let task = context_store.update(cx, |context_store, cx| {
241 suggested.accept(context_store, cx)
242 });
243
244 let workspace = this.workspace.clone();
245 cx.spawn(|this, mut cx| async move {
246 match task.await {
247 Ok(()) => {
248 if let Some(this) = this.upgrade() {
249 this.update(&mut cx, |_, cx| cx.notify())?;
250 }
251 }
252 Err(err) => {
253 let Some(workspace) = workspace.upgrade() else {
254 return anyhow::Ok(());
255 };
256
257 workspace.update(&mut cx, |workspace, cx| {
258 workspace.show_error(&err, cx);
259 })?;
260 }
261 }
262 anyhow::Ok(())
263 })
264 .detach_and_log_err(cx);
265 }))
266 },
267 ))
268 })
269 .when(!context.is_empty(), {
270 move |parent| {
271 parent.child(
272 IconButton::new("remove-all-context", IconName::Eraser)
273 .icon_size(IconSize::Small)
274 .tooltip({
275 let focus_handle = focus_handle.clone();
276 move |cx| {
277 Tooltip::for_action_in(
278 "Remove All Context",
279 &RemoveAllContext,
280 &focus_handle,
281 cx,
282 )
283 }
284 })
285 .on_click(cx.listener({
286 let focus_handle = focus_handle.clone();
287 move |_this, _event, cx| {
288 focus_handle.dispatch_action(&RemoveAllContext, cx);
289 }
290 })),
291 )
292 }
293 })
294 }
295}
296
297pub enum ContextStripEvent {
298 PickerDismissed,
299}
300
301impl EventEmitter<ContextStripEvent> for ContextStrip {}
302
303pub enum SuggestContextKind {
304 File,
305 Thread,
306}
307
308#[derive(Clone)]
309pub enum SuggestedContext {
310 File {
311 name: SharedString,
312 icon_path: Option<SharedString>,
313 buffer: WeakModel<Buffer>,
314 },
315 Thread {
316 name: SharedString,
317 thread: WeakModel<Thread>,
318 },
319}
320
321impl SuggestedContext {
322 pub fn name(&self) -> &SharedString {
323 match self {
324 Self::File { name, .. } => name,
325 Self::Thread { name, .. } => name,
326 }
327 }
328
329 pub fn icon_path(&self) -> Option<SharedString> {
330 match self {
331 Self::File { icon_path, .. } => icon_path.clone(),
332 Self::Thread { .. } => None,
333 }
334 }
335
336 pub fn accept(
337 &self,
338 context_store: &mut ContextStore,
339 cx: &mut ModelContext<ContextStore>,
340 ) -> Task<Result<()>> {
341 match self {
342 Self::File {
343 buffer,
344 icon_path: _,
345 name: _,
346 } => {
347 if let Some(buffer) = buffer.upgrade() {
348 return context_store.add_file_from_buffer(buffer, cx);
349 };
350 }
351 Self::Thread { thread, name: _ } => {
352 if let Some(thread) = thread.upgrade() {
353 context_store.insert_thread(thread, cx);
354 };
355 }
356 }
357 Task::ready(Ok(()))
358 }
359
360 pub fn kind(&self) -> ContextKind {
361 match self {
362 Self::File { .. } => ContextKind::File,
363 Self::Thread { .. } => ContextKind::Thread,
364 }
365 }
366}