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::context_picker::ContextPicker;
11use agent::context::RULES_ICON;
12use agent::context_store::{self, ContextStore};
13
14pub struct RulesContextPicker {
15 picker: Entity<Picker<RulesContextPickerDelegate>>,
16}
17
18impl RulesContextPicker {
19 pub fn new(
20 prompt_store: Entity<PromptStore>,
21 context_picker: WeakEntity<ContextPicker>,
22 context_store: WeakEntity<context_store::ContextStore>,
23 window: &mut Window,
24 cx: &mut Context<Self>,
25 ) -> Self {
26 let delegate = RulesContextPickerDelegate::new(prompt_store, context_picker, context_store);
27 let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
28
29 RulesContextPicker { picker }
30 }
31}
32
33impl Focusable for RulesContextPicker {
34 fn focus_handle(&self, cx: &App) -> FocusHandle {
35 self.picker.focus_handle(cx)
36 }
37}
38
39impl Render for RulesContextPicker {
40 fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
41 self.picker.clone()
42 }
43}
44
45#[derive(Debug, Clone)]
46pub struct RulesContextEntry {
47 pub prompt_id: UserPromptId,
48 pub title: SharedString,
49}
50
51pub struct RulesContextPickerDelegate {
52 prompt_store: Entity<PromptStore>,
53 context_picker: WeakEntity<ContextPicker>,
54 context_store: WeakEntity<context_store::ContextStore>,
55 matches: Vec<RulesContextEntry>,
56 selected_index: usize,
57}
58
59impl RulesContextPickerDelegate {
60 pub fn new(
61 prompt_store: Entity<PromptStore>,
62 context_picker: WeakEntity<ContextPicker>,
63 context_store: WeakEntity<context_store::ContextStore>,
64 ) -> Self {
65 RulesContextPickerDelegate {
66 prompt_store,
67 context_picker,
68 context_store,
69 matches: Vec::new(),
70 selected_index: 0,
71 }
72 }
73}
74
75impl PickerDelegate for RulesContextPickerDelegate {
76 type ListItem = ListItem;
77
78 fn match_count(&self) -> usize {
79 self.matches.len()
80 }
81
82 fn selected_index(&self) -> usize {
83 self.selected_index
84 }
85
86 fn set_selected_index(
87 &mut self,
88 ix: usize,
89 _window: &mut Window,
90 _cx: &mut Context<Picker<Self>>,
91 ) {
92 self.selected_index = ix;
93 }
94
95 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
96 "Search available rules…".into()
97 }
98
99 fn update_matches(
100 &mut self,
101 query: String,
102 window: &mut Window,
103 cx: &mut Context<Picker<Self>>,
104 ) -> Task<()> {
105 let search_task = search_rules(
106 query,
107 Arc::new(AtomicBool::default()),
108 &self.prompt_store,
109 cx,
110 );
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}