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