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