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