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::{
18 ResultExt, get_default_system_shell_preferring_bash, rel_path::RelPath, shell::ShellKind,
19};
20
21use crate::UserPromptId;
22
23#[derive(Default, Debug, Clone, Serialize)]
24pub struct ProjectContext {
25 pub worktrees: Vec<WorktreeContext>,
26 /// Whether any worktree has a rules_file. Provided as a field because handlebars can't do this.
27 pub has_rules: bool,
28 pub user_rules: Vec<UserRulesContext>,
29 /// `!user_rules.is_empty()` - provided as a field because handlebars can't do this.
30 pub has_user_rules: bool,
31 pub os: String,
32 pub arch: String,
33 pub shell: String,
34}
35
36impl ProjectContext {
37 pub fn new(worktrees: Vec<WorktreeContext>, default_user_rules: Vec<UserRulesContext>) -> Self {
38 let has_rules = worktrees
39 .iter()
40 .any(|worktree| worktree.rules_file.is_some());
41 Self {
42 worktrees,
43 has_rules,
44 has_user_rules: !default_user_rules.is_empty(),
45 user_rules: default_user_rules,
46 os: std::env::consts::OS.to_string(),
47 arch: std::env::consts::ARCH.to_string(),
48 shell: ShellKind::new(&get_default_system_shell_preferring_bash(), cfg!(windows))
49 .to_string(),
50 }
51 }
52}
53
54#[derive(Debug, Clone, Serialize)]
55pub struct UserRulesContext {
56 pub uuid: UserPromptId,
57 pub title: Option<String>,
58 pub contents: String,
59}
60
61#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
62pub struct WorktreeContext {
63 pub root_name: String,
64 pub abs_path: Arc<Path>,
65 pub rules_file: Option<RulesFileContext>,
66}
67
68#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
69pub struct RulesFileContext {
70 pub path_in_worktree: Arc<RelPath>,
71 pub text: String,
72 // This used for opening rules files. TODO: Since it isn't related to prompt templating, this
73 // should be moved elsewhere.
74 #[serde(skip)]
75 pub project_entry_id: usize,
76}
77
78#[derive(Serialize)]
79pub struct ContentPromptDiagnosticContext {
80 pub line_number: usize,
81 pub error_message: String,
82 pub code_content: String,
83}
84
85#[derive(Serialize)]
86pub struct ContentPromptContext {
87 pub content_type: String,
88 pub language_name: Option<String>,
89 pub is_insert: bool,
90 pub is_truncated: bool,
91 pub document_content: String,
92 pub user_prompt: String,
93 pub rewrite_section: Option<String>,
94 pub diagnostic_errors: Vec<ContentPromptDiagnosticContext>,
95}
96
97#[derive(Serialize)]
98pub struct ContentPromptContextV2 {
99 pub content_type: String,
100 pub language_name: Option<String>,
101 pub is_truncated: bool,
102 pub document_content: String,
103 pub rewrite_section: Option<String>,
104 pub diagnostic_errors: Vec<ContentPromptDiagnosticContext>,
105}
106
107#[derive(Serialize)]
108pub struct TerminalAssistantPromptContext {
109 pub os: String,
110 pub arch: String,
111 pub shell: Option<String>,
112 pub working_directory: Option<String>,
113 pub latest_output: Vec<String>,
114 pub user_prompt: String,
115}
116
117pub struct PromptLoadingParams<'a> {
118 pub fs: Arc<dyn Fs>,
119 pub repo_path: Option<PathBuf>,
120 pub cx: &'a gpui::App,
121}
122
123pub struct PromptBuilder {
124 handlebars: Arc<Mutex<Handlebars<'static>>>,
125}
126
127impl PromptBuilder {
128 pub fn load(fs: Arc<dyn Fs>, stdout_is_a_pty: bool, cx: &mut App) -> Arc<Self> {
129 Self::new(Some(PromptLoadingParams {
130 fs: fs.clone(),
131 repo_path: stdout_is_a_pty
132 .then(|| std::env::current_dir().log_err())
133 .flatten(),
134 cx,
135 }))
136 .log_err()
137 .map(Arc::new)
138 .unwrap_or_else(|| Arc::new(Self::new(None).unwrap()))
139 }
140
141 pub fn new(loading_params: Option<PromptLoadingParams>) -> Result<Self> {
142 let mut handlebars = Handlebars::new();
143 Self::register_built_in_templates(&mut handlebars)?;
144
145 let handlebars = Arc::new(Mutex::new(handlebars));
146
147 if let Some(params) = loading_params {
148 Self::watch_fs_for_template_overrides(params, handlebars.clone());
149 }
150
151 Ok(Self { handlebars })
152 }
153
154 /// Watches the filesystem for changes to prompt template overrides.
155 ///
156 /// This function sets up a file watcher on the prompt templates directory. It performs
157 /// an initial scan of the directory and registers any existing template overrides.
158 /// Then it continuously monitors for changes, reloading templates as they are
159 /// modified or added.
160 ///
161 /// If the templates directory doesn't exist initially, it waits for it to be created.
162 /// If the directory is removed, it restores the built-in templates and waits for the
163 /// directory to be recreated.
164 ///
165 /// # Arguments
166 ///
167 /// * `params` - A `PromptLoadingParams` struct containing the filesystem, repository path,
168 /// and application context.
169 /// * `handlebars` - An `Arc<Mutex<Handlebars>>` for registering and updating templates.
170 fn watch_fs_for_template_overrides(
171 params: PromptLoadingParams,
172 handlebars: Arc<Mutex<Handlebars<'static>>>,
173 ) {
174 let templates_dir = paths::prompt_overrides_dir(params.repo_path.as_deref());
175 params.cx.background_spawn(async move {
176 let Some(parent_dir) = templates_dir.parent() else {
177 return;
178 };
179
180 let mut found_dir_once = false;
181 loop {
182 // Check if the templates directory exists and handle its status
183 // If it exists, log its presence and check if it's a symlink
184 // If it doesn't exist:
185 // - Log that we're using built-in prompts
186 // - Check if it's a broken symlink and log if so
187 // - Set up a watcher to detect when it's created
188 // After the first check, set the `found_dir_once` flag
189 // This allows us to avoid logging when looping back around after deleting the prompt overrides directory.
190 let dir_status = params.fs.is_dir(&templates_dir).await;
191 let symlink_status = params.fs.read_link(&templates_dir).await.ok();
192 if dir_status {
193 let mut log_message = format!("Prompt template overrides directory found at {}", templates_dir.display());
194 if let Some(target) = symlink_status {
195 log_message.push_str(" -> ");
196 log_message.push_str(&target.display().to_string());
197 }
198 log::trace!("{}.", log_message);
199 } else {
200 if !found_dir_once {
201 log::trace!("No prompt template overrides directory found at {}. Using built-in prompts.", templates_dir.display());
202 if let Some(target) = symlink_status {
203 log::trace!("Symlink found pointing to {}, but target is invalid.", target.display());
204 }
205 }
206
207 if params.fs.is_dir(parent_dir).await {
208 let (mut changes, _watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
209 while let Some(changed_paths) = changes.next().await {
210 if changed_paths.iter().any(|p| &p.path == &templates_dir) {
211 let mut log_message = format!("Prompt template overrides directory detected at {}", templates_dir.display());
212 if let Ok(target) = params.fs.read_link(&templates_dir).await {
213 log_message.push_str(" -> ");
214 log_message.push_str(&target.display().to_string());
215 }
216 log::trace!("{}.", log_message);
217 break;
218 }
219 }
220 } else {
221 return;
222 }
223 }
224
225 found_dir_once = true;
226
227 // Initial scan of the prompt overrides directory
228 if let Ok(mut entries) = params.fs.read_dir(&templates_dir).await {
229 while let Some(Ok(file_path)) = entries.next().await {
230 if file_path.to_string_lossy().ends_with(".hbs")
231 && let Ok(content) = params.fs.load(&file_path).await {
232 let file_name = file_path.file_stem().unwrap().to_string_lossy();
233 log::debug!("Registering prompt template override: {}", file_name);
234 handlebars.lock().register_template_string(&file_name, content).log_err();
235 }
236 }
237 }
238
239 // Watch both the parent directory and the template overrides directory:
240 // - Monitor the parent directory to detect if the template overrides directory is deleted.
241 // - Monitor the template overrides directory to re-register templates when they change.
242 // Combine both watch streams into a single stream.
243 let (parent_changes, parent_watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
244 let (changes, watcher) = params.fs.watch(&templates_dir, Duration::from_secs(1)).await;
245 let mut combined_changes = futures::stream::select(changes, parent_changes);
246
247 while let Some(changed_paths) = combined_changes.next().await {
248 if changed_paths.iter().any(|p| &p.path == &templates_dir)
249 && !params.fs.is_dir(&templates_dir).await {
250 log::info!("Prompt template overrides directory removed. Restoring built-in prompt templates.");
251 Self::register_built_in_templates(&mut handlebars.lock()).log_err();
252 break;
253 }
254 for event in changed_paths {
255 if event.path.starts_with(&templates_dir) && event.path.extension().is_some_and(|ext| ext == "hbs") {
256 log::info!("Reloading prompt template override: {}", event.path.display());
257 if let Some(content) = params.fs.load(&event.path).await.log_err() {
258 let file_name = event.path.file_stem().unwrap().to_string_lossy();
259 handlebars.lock().register_template_string(&file_name, content).log_err();
260 }
261 }
262 }
263 }
264
265 drop(watcher);
266 drop(parent_watcher);
267 }
268 })
269 .detach();
270 }
271
272 fn register_built_in_templates(handlebars: &mut Handlebars) -> Result<()> {
273 for path in Assets.list("prompts")? {
274 if let Some(id) = path
275 .split('/')
276 .next_back()
277 .and_then(|s| s.strip_suffix(".hbs"))
278 && let Some(prompt) = Assets.load(path.as_ref()).log_err().flatten()
279 {
280 log::debug!("Registering built-in prompt template: {}", id);
281 let prompt = String::from_utf8_lossy(prompt.as_ref());
282 handlebars.register_template_string(id, LineEnding::normalize_cow(prompt))?
283 }
284 }
285
286 Ok(())
287 }
288
289 pub fn generate_inline_transformation_prompt_tools(
290 &self,
291 language_name: Option<&LanguageName>,
292 buffer: BufferSnapshot,
293 range: Range<usize>,
294 ) -> Result<String, RenderError> {
295 let content_type = match language_name.as_ref().map(|l| l.as_ref()) {
296 None | Some("Markdown" | "Plain Text") => "text",
297 Some(_) => "code",
298 };
299
300 const MAX_CTX: usize = 50000;
301 let is_insert = range.is_empty();
302 let mut is_truncated = false;
303
304 let before_range = 0..range.start;
305 let truncated_before = if before_range.len() > MAX_CTX {
306 is_truncated = true;
307 let start = buffer.clip_offset(range.start - MAX_CTX, text::Bias::Right);
308 start..range.start
309 } else {
310 before_range
311 };
312
313 let after_range = range.end..buffer.len();
314 let truncated_after = if after_range.len() > MAX_CTX {
315 is_truncated = true;
316 let end = buffer.clip_offset(range.end + MAX_CTX, text::Bias::Left);
317 range.end..end
318 } else {
319 after_range
320 };
321
322 let mut document_content = String::new();
323 for chunk in buffer.text_for_range(truncated_before) {
324 document_content.push_str(chunk);
325 }
326 if is_insert {
327 document_content.push_str("<insert_here></insert_here>");
328 } else {
329 document_content.push_str("<rewrite_this>\n");
330 for chunk in buffer.text_for_range(range.clone()) {
331 document_content.push_str(chunk);
332 }
333 document_content.push_str("\n</rewrite_this>");
334 }
335 for chunk in buffer.text_for_range(truncated_after) {
336 document_content.push_str(chunk);
337 }
338
339 let rewrite_section = if !is_insert {
340 let mut section = String::new();
341 for chunk in buffer.text_for_range(range.clone()) {
342 section.push_str(chunk);
343 }
344 Some(section)
345 } else {
346 None
347 };
348 let diagnostics = buffer.diagnostics_in_range::<_, Point>(range, false);
349 let diagnostic_errors: Vec<ContentPromptDiagnosticContext> = diagnostics
350 .map(|entry| {
351 let start = entry.range.start;
352 ContentPromptDiagnosticContext {
353 line_number: (start.row + 1) as usize,
354 error_message: entry.diagnostic.message.clone(),
355 code_content: buffer.text_for_range(entry.range).collect(),
356 }
357 })
358 .collect();
359
360 let context = ContentPromptContextV2 {
361 content_type: content_type.to_string(),
362 language_name: language_name.map(|s| s.to_string()),
363 is_truncated,
364 document_content,
365 rewrite_section,
366 diagnostic_errors,
367 };
368 self.handlebars.lock().render("content_prompt_v2", &context)
369 }
370
371 pub fn generate_inline_transformation_prompt(
372 &self,
373 user_prompt: String,
374 language_name: Option<&LanguageName>,
375 buffer: BufferSnapshot,
376 range: Range<usize>,
377 ) -> Result<String, RenderError> {
378 let content_type = match language_name.as_ref().map(|l| l.as_ref()) {
379 None | Some("Markdown" | "Plain Text") => "text",
380 Some(_) => "code",
381 };
382
383 const MAX_CTX: usize = 50000;
384 let is_insert = range.is_empty();
385 let mut is_truncated = false;
386
387 let before_range = 0..range.start;
388 let truncated_before = if before_range.len() > MAX_CTX {
389 is_truncated = true;
390 let start = buffer.clip_offset(range.start - MAX_CTX, text::Bias::Right);
391 start..range.start
392 } else {
393 before_range
394 };
395
396 let after_range = range.end..buffer.len();
397 let truncated_after = if after_range.len() > MAX_CTX {
398 is_truncated = true;
399 let end = buffer.clip_offset(range.end + MAX_CTX, text::Bias::Left);
400 range.end..end
401 } else {
402 after_range
403 };
404
405 let mut document_content = String::new();
406 for chunk in buffer.text_for_range(truncated_before) {
407 document_content.push_str(chunk);
408 }
409 if is_insert {
410 document_content.push_str("<insert_here></insert_here>");
411 } else {
412 document_content.push_str("<rewrite_this>\n");
413 for chunk in buffer.text_for_range(range.clone()) {
414 document_content.push_str(chunk);
415 }
416 document_content.push_str("\n</rewrite_this>");
417 }
418 for chunk in buffer.text_for_range(truncated_after) {
419 document_content.push_str(chunk);
420 }
421
422 let rewrite_section = if !is_insert {
423 let mut section = String::new();
424 for chunk in buffer.text_for_range(range.clone()) {
425 section.push_str(chunk);
426 }
427 Some(section)
428 } else {
429 None
430 };
431 let diagnostics = buffer.diagnostics_in_range::<_, Point>(range, false);
432 let diagnostic_errors: Vec<ContentPromptDiagnosticContext> = diagnostics
433 .map(|entry| {
434 let start = entry.range.start;
435 ContentPromptDiagnosticContext {
436 line_number: (start.row + 1) as usize,
437 error_message: entry.diagnostic.message.clone(),
438 code_content: buffer.text_for_range(entry.range).collect(),
439 }
440 })
441 .collect();
442
443 let context = ContentPromptContext {
444 content_type: content_type.to_string(),
445 language_name: language_name.map(|s| s.to_string()),
446 is_insert,
447 is_truncated,
448 document_content,
449 user_prompt,
450 rewrite_section,
451 diagnostic_errors,
452 };
453 self.handlebars.lock().render("content_prompt", &context)
454 }
455
456 pub fn generate_terminal_assistant_prompt(
457 &self,
458 user_prompt: &str,
459 shell: Option<&str>,
460 working_directory: Option<&str>,
461 latest_output: &[String],
462 ) -> Result<String, RenderError> {
463 let context = TerminalAssistantPromptContext {
464 os: std::env::consts::OS.to_string(),
465 arch: std::env::consts::ARCH.to_string(),
466 shell: shell.map(|s| s.to_string()),
467 working_directory: working_directory.map(|s| s.to_string()),
468 latest_output: latest_output.to_vec(),
469 user_prompt: user_prompt.to_string(),
470 };
471
472 self.handlebars
473 .lock()
474 .render("terminal_assistant_prompt", &context)
475 }
476}