1use anyhow::Result;
2use assets::Assets;
3use fs::Fs;
4use futures::StreamExt;
5use gpui::AssetSource;
6use handlebars::{Handlebars, RenderError};
7use language::BufferSnapshot;
8use parking_lot::Mutex;
9use serde::Serialize;
10use std::{ops::Range, path::PathBuf, sync::Arc, time::Duration};
11use util::ResultExt;
12
13#[derive(Serialize)]
14pub struct ContentPromptContext {
15 pub content_type: String,
16 pub language_name: Option<String>,
17 pub is_insert: bool,
18 pub is_truncated: bool,
19 pub document_content: String,
20 pub user_prompt: String,
21 pub rewrite_section: Option<String>,
22}
23
24#[derive(Serialize)]
25pub struct TerminalAssistantPromptContext {
26 pub os: String,
27 pub arch: String,
28 pub shell: Option<String>,
29 pub working_directory: Option<String>,
30 pub latest_output: Vec<String>,
31 pub user_prompt: String,
32}
33
34/// Context required to generate a workflow step resolution prompt.
35#[derive(Debug, Serialize)]
36pub struct StepResolutionContext {
37 /// The full context, including <step>...</step> tags
38 pub workflow_context: String,
39 /// The text of the specific step from the context to resolve
40 pub step_to_resolve: String,
41}
42
43pub struct PromptLoadingParams<'a> {
44 pub fs: Arc<dyn Fs>,
45 pub repo_path: Option<PathBuf>,
46 pub cx: &'a gpui::AppContext,
47}
48
49pub struct PromptBuilder {
50 handlebars: Arc<Mutex<Handlebars<'static>>>,
51}
52
53impl PromptBuilder {
54 pub fn new(loading_params: Option<PromptLoadingParams>) -> Result<Self> {
55 let mut handlebars = Handlebars::new();
56 Self::register_built_in_templates(&mut handlebars)?;
57
58 let handlebars = Arc::new(Mutex::new(handlebars));
59
60 if let Some(params) = loading_params {
61 Self::watch_fs_for_template_overrides(params, handlebars.clone());
62 }
63
64 Ok(Self { handlebars })
65 }
66
67 /// Watches the filesystem for changes to prompt template overrides.
68 ///
69 /// This function sets up a file watcher on the prompt templates directory. It performs
70 /// an initial scan of the directory and registers any existing template overrides.
71 /// Then it continuously monitors for changes, reloading templates as they are
72 /// modified or added.
73 ///
74 /// If the templates directory doesn't exist initially, it waits for it to be created.
75 /// If the directory is removed, it restores the built-in templates and waits for the
76 /// directory to be recreated.
77 ///
78 /// # Arguments
79 ///
80 /// * `params` - A `PromptLoadingParams` struct containing the filesystem, repository path,
81 /// and application context.
82 /// * `handlebars` - An `Arc<Mutex<Handlebars>>` for registering and updating templates.
83 fn watch_fs_for_template_overrides(
84 mut params: PromptLoadingParams,
85 handlebars: Arc<Mutex<Handlebars<'static>>>,
86 ) {
87 params.repo_path = None;
88 let templates_dir = paths::prompt_overrides_dir(params.repo_path.as_deref());
89 params.cx.background_executor()
90 .spawn(async move {
91 let Some(parent_dir) = templates_dir.parent() else {
92 return;
93 };
94
95 let mut found_dir_once = false;
96 loop {
97 // Check if the templates directory exists and handle its status
98 // If it exists, log its presence and check if it's a symlink
99 // If it doesn't exist:
100 // - Log that we're using built-in prompts
101 // - Check if it's a broken symlink and log if so
102 // - Set up a watcher to detect when it's created
103 // After the first check, set the `found_dir_once` flag
104 // This allows us to avoid logging when looping back around after deleting the prompt overrides directory.
105 let dir_status = params.fs.is_dir(&templates_dir).await;
106 let symlink_status = params.fs.read_link(&templates_dir).await.ok();
107 if dir_status {
108 let mut log_message = format!("Prompt template overrides directory found at {}", templates_dir.display());
109 if let Some(target) = symlink_status {
110 log_message.push_str(" -> ");
111 log_message.push_str(&target.display().to_string());
112 }
113 log::info!("{}.", log_message);
114 } else {
115 if !found_dir_once {
116 log::info!("No prompt template overrides directory found at {}. Using built-in prompts.", templates_dir.display());
117 if let Some(target) = symlink_status {
118 log::info!("Symlink found pointing to {}, but target is invalid.", target.display());
119 }
120 }
121
122 if params.fs.is_dir(parent_dir).await {
123 let (mut changes, _watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
124 while let Some(changed_paths) = changes.next().await {
125 if changed_paths.iter().any(|p| p == &templates_dir) {
126 let mut log_message = format!("Prompt template overrides directory detected at {}", templates_dir.display());
127 if let Ok(target) = params.fs.read_link(&templates_dir).await {
128 log_message.push_str(" -> ");
129 log_message.push_str(&target.display().to_string());
130 }
131 log::info!("{}.", log_message);
132 break;
133 }
134 }
135 } else {
136 return;
137 }
138 }
139
140 found_dir_once = true;
141
142 // Initial scan of the prompt overrides directory
143 if let Ok(mut entries) = params.fs.read_dir(&templates_dir).await {
144 while let Some(Ok(file_path)) = entries.next().await {
145 if file_path.to_string_lossy().ends_with(".hbs") {
146 if let Ok(content) = params.fs.load(&file_path).await {
147 let file_name = file_path.file_stem().unwrap().to_string_lossy();
148 log::info!("Registering prompt template override: {}", file_name);
149 handlebars.lock().register_template_string(&file_name, content).log_err();
150 }
151 }
152 }
153 }
154
155 // Watch both the parent directory and the template overrides directory:
156 // - Monitor the parent directory to detect if the template overrides directory is deleted.
157 // - Monitor the template overrides directory to re-register templates when they change.
158 // Combine both watch streams into a single stream.
159 let (parent_changes, parent_watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
160 let (changes, watcher) = params.fs.watch(&templates_dir, Duration::from_secs(1)).await;
161 let mut combined_changes = futures::stream::select(changes, parent_changes);
162
163 while let Some(changed_paths) = combined_changes.next().await {
164 if changed_paths.iter().any(|p| p == &templates_dir) {
165 if !params.fs.is_dir(&templates_dir).await {
166 log::info!("Prompt template overrides directory removed. Restoring built-in prompt templates.");
167 Self::register_built_in_templates(&mut handlebars.lock()).log_err();
168 break;
169 }
170 }
171 for changed_path in changed_paths {
172 if changed_path.starts_with(&templates_dir) && changed_path.extension().map_or(false, |ext| ext == "hbs") {
173 log::info!("Reloading prompt template override: {}", changed_path.display());
174 if let Some(content) = params.fs.load(&changed_path).await.log_err() {
175 let file_name = changed_path.file_stem().unwrap().to_string_lossy();
176 handlebars.lock().register_template_string(&file_name, content).log_err();
177 }
178 }
179 }
180 }
181
182 drop(watcher);
183 drop(parent_watcher);
184 }
185 })
186 .detach();
187 }
188
189 fn register_built_in_templates(handlebars: &mut Handlebars) -> Result<()> {
190 for path in Assets.list("prompts")? {
191 if let Some(id) = path.split('/').last().and_then(|s| s.strip_suffix(".hbs")) {
192 if let Some(prompt) = Assets.load(path.as_ref()).log_err().flatten() {
193 log::info!("Registering built-in prompt template: {}", id);
194 handlebars
195 .register_template_string(id, String::from_utf8_lossy(prompt.as_ref()))?
196 }
197 }
198 }
199
200 Ok(())
201 }
202
203 pub fn generate_content_prompt(
204 &self,
205 user_prompt: String,
206 language_name: Option<&str>,
207 buffer: BufferSnapshot,
208 range: Range<usize>,
209 ) -> Result<String, RenderError> {
210 let content_type = match language_name {
211 None | Some("Markdown" | "Plain Text") => "text",
212 Some(_) => "code",
213 };
214
215 const MAX_CTX: usize = 50000;
216 let is_insert = range.is_empty();
217 let mut is_truncated = false;
218
219 let before_range = 0..range.start;
220 let truncated_before = if before_range.len() > MAX_CTX {
221 is_truncated = true;
222 range.start - MAX_CTX..range.start
223 } else {
224 before_range
225 };
226
227 let after_range = range.end..buffer.len();
228 let truncated_after = if after_range.len() > MAX_CTX {
229 is_truncated = true;
230 range.end..range.end + MAX_CTX
231 } else {
232 after_range
233 };
234
235 let mut document_content = String::new();
236 for chunk in buffer.text_for_range(truncated_before) {
237 document_content.push_str(chunk);
238 }
239 if is_insert {
240 document_content.push_str("<insert_here></insert_here>");
241 } else {
242 document_content.push_str("<rewrite_this>\n");
243 for chunk in buffer.text_for_range(range.clone()) {
244 document_content.push_str(chunk);
245 }
246 document_content.push_str("\n</rewrite_this>");
247 }
248 for chunk in buffer.text_for_range(truncated_after) {
249 document_content.push_str(chunk);
250 }
251
252 let rewrite_section = if !is_insert {
253 let mut section = String::new();
254 for chunk in buffer.text_for_range(range.clone()) {
255 section.push_str(chunk);
256 }
257 Some(section)
258 } else {
259 None
260 };
261
262 let context = ContentPromptContext {
263 content_type: content_type.to_string(),
264 language_name: language_name.map(|s| s.to_string()),
265 is_insert,
266 is_truncated,
267 document_content,
268 user_prompt,
269 rewrite_section,
270 };
271
272 self.handlebars.lock().render("content_prompt", &context)
273 }
274
275 pub fn generate_terminal_assistant_prompt(
276 &self,
277 user_prompt: &str,
278 shell: Option<&str>,
279 working_directory: Option<&str>,
280 latest_output: &[String],
281 ) -> Result<String, RenderError> {
282 let context = TerminalAssistantPromptContext {
283 os: std::env::consts::OS.to_string(),
284 arch: std::env::consts::ARCH.to_string(),
285 shell: shell.map(|s| s.to_string()),
286 working_directory: working_directory.map(|s| s.to_string()),
287 latest_output: latest_output.to_vec(),
288 user_prompt: user_prompt.to_string(),
289 };
290
291 self.handlebars
292 .lock()
293 .render("terminal_assistant_prompt", &context)
294 }
295
296 pub fn generate_workflow_prompt(&self) -> Result<String, RenderError> {
297 self.handlebars.lock().render("edit_workflow", &())
298 }
299
300 pub fn generate_step_resolution_prompt(
301 &self,
302 context: &StepResolutionContext,
303 ) -> Result<String, RenderError> {
304 self.handlebars.lock().render("step_resolution", context)
305 }
306}