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