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