1use std::sync::Arc;
2use std::sync::atomic::AtomicBool;
3
4use anyhow::anyhow;
5use gpui::{App, DismissEvent, Entity, FocusHandle, Focusable, Task, WeakEntity};
6use picker::{Picker, PickerDelegate};
7use prompt_store::{PromptId, UserPromptId};
8use ui::{ListItem, prelude::*};
9
10use crate::context::RULES_ICON;
11use crate::context_picker::ContextPicker;
12use crate::context_store::{self, ContextStore};
13use crate::thread_store::ThreadStore;
14
15pub struct RulesContextPicker {
16 picker: Entity<Picker<RulesContextPickerDelegate>>,
17}
18
19impl RulesContextPicker {
20 pub fn new(
21 thread_store: WeakEntity<ThreadStore>,
22 context_picker: WeakEntity<ContextPicker>,
23 context_store: WeakEntity<context_store::ContextStore>,
24 window: &mut Window,
25 cx: &mut Context<Self>,
26 ) -> Self {
27 let delegate = RulesContextPickerDelegate::new(thread_store, context_picker, context_store);
28 let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
29
30 RulesContextPicker { picker }
31 }
32}
33
34impl Focusable for RulesContextPicker {
35 fn focus_handle(&self, cx: &App) -> FocusHandle {
36 self.picker.focus_handle(cx)
37 }
38}
39
40impl Render for RulesContextPicker {
41 fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
42 self.picker.clone()
43 }
44}
45
46#[derive(Debug, Clone)]
47pub struct RulesContextEntry {
48 pub prompt_id: UserPromptId,
49 pub title: SharedString,
50}
51
52pub struct RulesContextPickerDelegate {
53 thread_store: WeakEntity<ThreadStore>,
54 context_picker: WeakEntity<ContextPicker>,
55 context_store: WeakEntity<context_store::ContextStore>,
56 matches: Vec<RulesContextEntry>,
57 selected_index: usize,
58}
59
60impl RulesContextPickerDelegate {
61 pub fn new(
62 thread_store: WeakEntity<ThreadStore>,
63 context_picker: WeakEntity<ContextPicker>,
64 context_store: WeakEntity<context_store::ContextStore>,
65 ) -> Self {
66 RulesContextPickerDelegate {
67 thread_store,
68 context_picker,
69 context_store,
70 matches: Vec::new(),
71 selected_index: 0,
72 }
73 }
74}
75
76impl PickerDelegate for RulesContextPickerDelegate {
77 type ListItem = ListItem;
78
79 fn match_count(&self) -> usize {
80 self.matches.len()
81 }
82
83 fn selected_index(&self) -> usize {
84 self.selected_index
85 }
86
87 fn set_selected_index(
88 &mut self,
89 ix: usize,
90 _window: &mut Window,
91 _cx: &mut Context<Picker<Self>>,
92 ) {
93 self.selected_index = ix;
94 }
95
96 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
97 "Search available rules…".into()
98 }
99
100 fn update_matches(
101 &mut self,
102 query: String,
103 window: &mut Window,
104 cx: &mut Context<Picker<Self>>,
105 ) -> Task<()> {
106 let Some(thread_store) = self.thread_store.upgrade() else {
107 return Task::ready(());
108 };
109
110 let search_task = search_rules(query, Arc::new(AtomicBool::default()), thread_store, cx);
111 cx.spawn_in(window, async move |this, cx| {
112 let matches = search_task.await;
113 this.update(cx, |this, cx| {
114 this.delegate.matches = matches;
115 this.delegate.selected_index = 0;
116 cx.notify();
117 })
118 .ok();
119 })
120 }
121
122 fn confirm(&mut self, _secondary: bool, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
123 let Some(entry) = self.matches.get(self.selected_index) else {
124 return;
125 };
126
127 let Some(thread_store) = self.thread_store.upgrade() else {
128 return;
129 };
130
131 let prompt_id = entry.prompt_id;
132
133 let load_rules_task = thread_store.update(cx, |thread_store, cx| {
134 thread_store.load_rules(prompt_id, cx)
135 });
136
137 cx.spawn(async move |this, cx| {
138 let (metadata, text) = load_rules_task.await?;
139 let Some(title) = metadata.title else {
140 return Err(anyhow!("Encountered user rule with no title when attempting to add it to agent context."));
141 };
142 this.update(cx, |this, cx| {
143 this.delegate
144 .context_store
145 .update(cx, |context_store, cx| {
146 context_store.add_rules(prompt_id, title, text, true, cx)
147 })
148 .ok();
149 })
150 })
151 .detach_and_log_err(cx);
152 }
153
154 fn dismissed(&mut self, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
155 self.context_picker
156 .update(cx, |_, cx| {
157 cx.emit(DismissEvent);
158 })
159 .ok();
160 }
161
162 fn render_match(
163 &self,
164 ix: usize,
165 selected: bool,
166 _window: &mut Window,
167 cx: &mut Context<Picker<Self>>,
168 ) -> Option<Self::ListItem> {
169 let thread = &self.matches[ix];
170
171 Some(ListItem::new(ix).inset(true).toggle_state(selected).child(
172 render_thread_context_entry(thread, self.context_store.clone(), cx),
173 ))
174 }
175}
176
177pub fn render_thread_context_entry(
178 user_rules: &RulesContextEntry,
179 context_store: WeakEntity<ContextStore>,
180 cx: &mut App,
181) -> Div {
182 let added = context_store.upgrade().map_or(false, |ctx_store| {
183 ctx_store
184 .read(cx)
185 .includes_user_rules(&user_rules.prompt_id)
186 .is_some()
187 });
188
189 h_flex()
190 .gap_1p5()
191 .w_full()
192 .justify_between()
193 .child(
194 h_flex()
195 .gap_1p5()
196 .max_w_72()
197 .child(
198 Icon::new(RULES_ICON)
199 .size(IconSize::XSmall)
200 .color(Color::Muted),
201 )
202 .child(Label::new(user_rules.title.clone()).truncate()),
203 )
204 .when(added, |el| {
205 el.child(
206 h_flex()
207 .gap_1()
208 .child(
209 Icon::new(IconName::Check)
210 .size(IconSize::Small)
211 .color(Color::Success),
212 )
213 .child(Label::new("Added").size(LabelSize::Small)),
214 )
215 })
216}
217
218pub(crate) fn search_rules(
219 query: String,
220 cancellation_flag: Arc<AtomicBool>,
221 thread_store: Entity<ThreadStore>,
222 cx: &mut App,
223) -> Task<Vec<RulesContextEntry>> {
224 let Some(prompt_store) = thread_store.read(cx).prompt_store() else {
225 return Task::ready(vec![]);
226 };
227 let search_task = prompt_store.read(cx).search(query, cancellation_flag, cx);
228 cx.background_spawn(async move {
229 search_task
230 .await
231 .into_iter()
232 .flat_map(|metadata| {
233 // Default prompts are filtered out as they are automatically included.
234 if metadata.default {
235 None
236 } else {
237 match metadata.id {
238 PromptId::EditWorkflow => None,
239 PromptId::User { uuid } => Some(RulesContextEntry {
240 prompt_id: uuid,
241 title: metadata.title?,
242 }),
243 }
244 }
245 })
246 .collect::<Vec<_>>()
247 })
248}