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