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