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