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