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(Default, 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 ModelContext {
53 pub available_tools: Vec<String>,
54}
55
56#[derive(Serialize)]
57struct PromptTemplateContext {
58 #[serde(flatten)]
59 project: ProjectContext,
60
61 #[serde(flatten)]
62 model: ModelContext,
63
64 has_tools: bool,
65}
66
67#[derive(Debug, Clone, Serialize)]
68pub struct UserRulesContext {
69 pub uuid: UserPromptId,
70 pub title: Option<String>,
71 pub contents: String,
72}
73
74#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
75pub struct WorktreeContext {
76 pub root_name: String,
77 pub abs_path: Arc<Path>,
78 pub rules_file: Option<RulesFileContext>,
79}
80
81#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
82pub struct RulesFileContext {
83 pub path_in_worktree: Arc<Path>,
84 pub text: String,
85 // This used for opening rules files. TODO: Since it isn't related to prompt templating, this
86 // should be moved elsewhere.
87 #[serde(skip)]
88 pub project_entry_id: usize,
89}
90
91#[derive(Serialize)]
92pub struct ContentPromptDiagnosticContext {
93 pub line_number: usize,
94 pub error_message: String,
95 pub code_content: String,
96}
97
98#[derive(Serialize)]
99pub struct ContentPromptContext {
100 pub content_type: String,
101 pub language_name: Option<String>,
102 pub is_insert: bool,
103 pub is_truncated: bool,
104 pub document_content: String,
105 pub user_prompt: String,
106 pub rewrite_section: Option<String>,
107 pub diagnostic_errors: Vec<ContentPromptDiagnosticContext>,
108}
109
110#[derive(Serialize)]
111pub struct TerminalAssistantPromptContext {
112 pub os: String,
113 pub arch: String,
114 pub shell: Option<String>,
115 pub working_directory: Option<String>,
116 pub latest_output: Vec<String>,
117 pub user_prompt: String,
118}
119
120pub struct PromptLoadingParams<'a> {
121 pub fs: Arc<dyn Fs>,
122 pub repo_path: Option<PathBuf>,
123 pub cx: &'a gpui::App,
124}
125
126pub struct PromptBuilder {
127 handlebars: Arc<Mutex<Handlebars<'static>>>,
128}
129
130impl PromptBuilder {
131 pub fn load(fs: Arc<dyn Fs>, stdout_is_a_pty: bool, cx: &mut App) -> Arc<Self> {
132 Self::new(Some(PromptLoadingParams {
133 fs: fs.clone(),
134 repo_path: stdout_is_a_pty
135 .then(|| std::env::current_dir().log_err())
136 .flatten(),
137 cx,
138 }))
139 .log_err()
140 .map(Arc::new)
141 .unwrap_or_else(|| Arc::new(Self::new(None).unwrap()))
142 }
143
144 /// Helper function for handlebars templates to check if a specific tool is enabled
145 fn has_tool_helper(
146 h: &handlebars::Helper,
147 _: &Handlebars,
148 ctx: &handlebars::Context,
149 _: &mut handlebars::RenderContext,
150 out: &mut dyn handlebars::Output,
151 ) -> handlebars::HelperResult {
152 let tool_name = h.param(0).and_then(|v| v.value().as_str()).ok_or_else(|| {
153 handlebars::RenderError::new("has_tool helper: missing or invalid tool name parameter")
154 })?;
155
156 let enabled_tools = ctx
157 .data()
158 .get("available_tools")
159 .and_then(|v| v.as_array())
160 .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect::<Vec<&str>>())
161 .ok_or_else(|| {
162 handlebars::RenderError::new(
163 "has_tool handlebars helper: available_tools not found or not an array",
164 )
165 })?;
166
167 if enabled_tools.contains(&tool_name) {
168 out.write("true")?;
169 }
170
171 Ok(())
172 }
173
174 pub fn new(loading_params: Option<PromptLoadingParams>) -> Result<Self> {
175 let mut handlebars = Handlebars::new();
176 Self::register_built_in_templates(&mut handlebars)?;
177 handlebars.register_helper("has_tool", Box::new(Self::has_tool_helper));
178
179 let handlebars = Arc::new(Mutex::new(handlebars));
180
181 if let Some(params) = loading_params {
182 Self::watch_fs_for_template_overrides(params, handlebars.clone());
183 }
184
185 Ok(Self { handlebars })
186 }
187
188 /// Watches the filesystem for changes to prompt template overrides.
189 ///
190 /// This function sets up a file watcher on the prompt templates directory. It performs
191 /// an initial scan of the directory and registers any existing template overrides.
192 /// Then it continuously monitors for changes, reloading templates as they are
193 /// modified or added.
194 ///
195 /// If the templates directory doesn't exist initially, it waits for it to be created.
196 /// If the directory is removed, it restores the built-in templates and waits for the
197 /// directory to be recreated.
198 ///
199 /// # Arguments
200 ///
201 /// * `params` - A `PromptLoadingParams` struct containing the filesystem, repository path,
202 /// and application context.
203 /// * `handlebars` - An `Arc<Mutex<Handlebars>>` for registering and updating templates.
204 fn watch_fs_for_template_overrides(
205 params: PromptLoadingParams,
206 handlebars: Arc<Mutex<Handlebars<'static>>>,
207 ) {
208 let templates_dir = paths::prompt_overrides_dir(params.repo_path.as_deref());
209 params.cx.background_spawn(async move {
210 let Some(parent_dir) = templates_dir.parent() else {
211 return;
212 };
213
214 let mut found_dir_once = false;
215 loop {
216 // Check if the templates directory exists and handle its status
217 // If it exists, log its presence and check if it's a symlink
218 // If it doesn't exist:
219 // - Log that we're using built-in prompts
220 // - Check if it's a broken symlink and log if so
221 // - Set up a watcher to detect when it's created
222 // After the first check, set the `found_dir_once` flag
223 // This allows us to avoid logging when looping back around after deleting the prompt overrides directory.
224 let dir_status = params.fs.is_dir(&templates_dir).await;
225 let symlink_status = params.fs.read_link(&templates_dir).await.ok();
226 if dir_status {
227 let mut log_message = format!("Prompt template overrides directory found at {}", templates_dir.display());
228 if let Some(target) = symlink_status {
229 log_message.push_str(" -> ");
230 log_message.push_str(&target.display().to_string());
231 }
232 log::trace!("{}.", log_message);
233 } else {
234 if !found_dir_once {
235 log::trace!("No prompt template overrides directory found at {}. Using built-in prompts.", templates_dir.display());
236 if let Some(target) = symlink_status {
237 log::trace!("Symlink found pointing to {}, but target is invalid.", target.display());
238 }
239 }
240
241 if params.fs.is_dir(parent_dir).await {
242 let (mut changes, _watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
243 while let Some(changed_paths) = changes.next().await {
244 if changed_paths.iter().any(|p| &p.path == &templates_dir) {
245 let mut log_message = format!("Prompt template overrides directory detected at {}", templates_dir.display());
246 if let Ok(target) = params.fs.read_link(&templates_dir).await {
247 log_message.push_str(" -> ");
248 log_message.push_str(&target.display().to_string());
249 }
250 log::trace!("{}.", log_message);
251 break;
252 }
253 }
254 } else {
255 return;
256 }
257 }
258
259 found_dir_once = true;
260
261 // Initial scan of the prompt overrides directory
262 if let Ok(mut entries) = params.fs.read_dir(&templates_dir).await {
263 while let Some(Ok(file_path)) = entries.next().await {
264 if file_path.to_string_lossy().ends_with(".hbs")
265 && let Ok(content) = params.fs.load(&file_path).await {
266 let file_name = file_path.file_stem().unwrap().to_string_lossy();
267 log::debug!("Registering prompt template override: {}", file_name);
268 handlebars.lock().register_template_string(&file_name, content).log_err();
269 }
270 }
271 }
272
273 // Watch both the parent directory and the template overrides directory:
274 // - Monitor the parent directory to detect if the template overrides directory is deleted.
275 // - Monitor the template overrides directory to re-register templates when they change.
276 // Combine both watch streams into a single stream.
277 let (parent_changes, parent_watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
278 let (changes, watcher) = params.fs.watch(&templates_dir, Duration::from_secs(1)).await;
279 let mut combined_changes = futures::stream::select(changes, parent_changes);
280
281 while let Some(changed_paths) = combined_changes.next().await {
282 if changed_paths.iter().any(|p| &p.path == &templates_dir)
283 && !params.fs.is_dir(&templates_dir).await {
284 log::info!("Prompt template overrides directory removed. Restoring built-in prompt templates.");
285 Self::register_built_in_templates(&mut handlebars.lock()).log_err();
286 break;
287 }
288 for event in changed_paths {
289 if event.path.starts_with(&templates_dir) && event.path.extension().is_some_and(|ext| ext == "hbs") {
290 log::info!("Reloading prompt template override: {}", event.path.display());
291 if let Some(content) = params.fs.load(&event.path).await.log_err() {
292 let file_name = event.path.file_stem().unwrap().to_string_lossy();
293 handlebars.lock().register_template_string(&file_name, content).log_err();
294 }
295 }
296 }
297 }
298
299 drop(watcher);
300 drop(parent_watcher);
301 }
302 })
303 .detach();
304 }
305
306 fn register_built_in_templates(handlebars: &mut Handlebars) -> Result<()> {
307 for path in Assets.list("prompts")? {
308 if let Some(id) = path
309 .split('/')
310 .next_back()
311 .and_then(|s| s.strip_suffix(".hbs"))
312 && let Some(prompt) = Assets.load(path.as_ref()).log_err().flatten()
313 {
314 log::debug!("Registering built-in prompt template: {}", id);
315 let prompt = String::from_utf8_lossy(prompt.as_ref());
316 handlebars.register_template_string(id, LineEnding::normalize_cow(prompt))?
317 }
318 }
319
320 Ok(())
321 }
322
323 pub fn generate_assistant_system_prompt(
324 &self,
325 context: &ProjectContext,
326 model_context: &ModelContext,
327 ) -> Result<String, RenderError> {
328 let template_context = PromptTemplateContext {
329 project: context.clone(),
330 model: model_context.clone(),
331 has_tools: !model_context.available_tools.is_empty(),
332 };
333
334 self.handlebars
335 .lock()
336 .render("assistant_system_prompt", &template_context)
337 }
338
339 pub fn generate_inline_transformation_prompt(
340 &self,
341 user_prompt: String,
342 language_name: Option<&LanguageName>,
343 buffer: BufferSnapshot,
344 range: Range<usize>,
345 ) -> Result<String, RenderError> {
346 let content_type = match language_name.as_ref().map(|l| l.as_ref()) {
347 None | Some("Markdown" | "Plain Text") => "text",
348 Some(_) => "code",
349 };
350
351 const MAX_CTX: usize = 50000;
352 let is_insert = range.is_empty();
353 let mut is_truncated = false;
354
355 let before_range = 0..range.start;
356 let truncated_before = if before_range.len() > MAX_CTX {
357 is_truncated = true;
358 let start = buffer.clip_offset(range.start - MAX_CTX, text::Bias::Right);
359 start..range.start
360 } else {
361 before_range
362 };
363
364 let after_range = range.end..buffer.len();
365 let truncated_after = if after_range.len() > MAX_CTX {
366 is_truncated = true;
367 let end = buffer.clip_offset(range.end + MAX_CTX, text::Bias::Left);
368 range.end..end
369 } else {
370 after_range
371 };
372
373 let mut document_content = String::new();
374 for chunk in buffer.text_for_range(truncated_before) {
375 document_content.push_str(chunk);
376 }
377 if is_insert {
378 document_content.push_str("<insert_here></insert_here>");
379 } else {
380 document_content.push_str("<rewrite_this>\n");
381 for chunk in buffer.text_for_range(range.clone()) {
382 document_content.push_str(chunk);
383 }
384 document_content.push_str("\n</rewrite_this>");
385 }
386 for chunk in buffer.text_for_range(truncated_after) {
387 document_content.push_str(chunk);
388 }
389
390 let rewrite_section = if !is_insert {
391 let mut section = String::new();
392 for chunk in buffer.text_for_range(range.clone()) {
393 section.push_str(chunk);
394 }
395 Some(section)
396 } else {
397 None
398 };
399 let diagnostics = buffer.diagnostics_in_range::<_, Point>(range, false);
400 let diagnostic_errors: Vec<ContentPromptDiagnosticContext> = diagnostics
401 .map(|entry| {
402 let start = entry.range.start;
403 ContentPromptDiagnosticContext {
404 line_number: (start.row + 1) as usize,
405 error_message: entry.diagnostic.message.clone(),
406 code_content: buffer.text_for_range(entry.range).collect(),
407 }
408 })
409 .collect();
410
411 let context = ContentPromptContext {
412 content_type: content_type.to_string(),
413 language_name: language_name.map(|s| s.to_string()),
414 is_insert,
415 is_truncated,
416 document_content,
417 user_prompt,
418 rewrite_section,
419 diagnostic_errors,
420 };
421 self.handlebars.lock().render("content_prompt", &context)
422 }
423
424 pub fn generate_terminal_assistant_prompt(
425 &self,
426 user_prompt: &str,
427 shell: Option<&str>,
428 working_directory: Option<&str>,
429 latest_output: &[String],
430 ) -> Result<String, RenderError> {
431 let context = TerminalAssistantPromptContext {
432 os: std::env::consts::OS.to_string(),
433 arch: std::env::consts::ARCH.to_string(),
434 shell: shell.map(|s| s.to_string()),
435 working_directory: working_directory.map(|s| s.to_string()),
436 latest_output: latest_output.to_vec(),
437 user_prompt: user_prompt.to_string(),
438 };
439
440 self.handlebars
441 .lock()
442 .render("terminal_assistant_prompt", &context)
443 }
444}
445
446#[cfg(test)]
447mod test {
448 use super::*;
449 use serde_json;
450 use uuid::Uuid;
451
452 #[test]
453 fn test_assistant_system_prompt_renders() {
454 let worktrees = vec![WorktreeContext {
455 root_name: "path".into(),
456 abs_path: Path::new("/path/to/root").into(),
457 rules_file: Some(RulesFileContext {
458 path_in_worktree: Path::new(".rules").into(),
459 text: "".into(),
460 project_entry_id: 0,
461 }),
462 }];
463 let default_user_rules = vec![UserRulesContext {
464 uuid: UserPromptId(Uuid::nil()),
465 title: Some("Rules title".into()),
466 contents: "Rules contents".into(),
467 }];
468 let project_context = ProjectContext::new(worktrees, default_user_rules);
469 let model_context = ModelContext {
470 available_tools: ["grep".into()].to_vec(),
471 };
472 let prompt = PromptBuilder::new(None)
473 .unwrap()
474 .generate_assistant_system_prompt(&project_context, &model_context)
475 .unwrap();
476 assert!(
477 prompt.contains("Rules contents"),
478 "Expected default user rules to be in rendered prompt"
479 );
480 }
481
482 #[test]
483 fn test_assistant_system_prompt_depends_on_enabled_tools() {
484 let worktrees = vec![WorktreeContext {
485 root_name: "path".into(),
486 abs_path: Path::new("/path/to/root").into(),
487 rules_file: None,
488 }];
489 let default_user_rules = vec![];
490 let project_context = ProjectContext::new(worktrees, default_user_rules);
491 let prompt_builder = PromptBuilder::new(None).unwrap();
492
493 // When the `grep` tool is enabled, it should be mentioned in the prompt
494 let model_context = ModelContext {
495 available_tools: ["grep".into()].to_vec(),
496 };
497 let prompt_with_grep = prompt_builder
498 .generate_assistant_system_prompt(&project_context, &model_context)
499 .unwrap();
500 assert!(
501 prompt_with_grep.contains("grep"),
502 "`grep` tool should be mentioned in prompt when the tool is enabled"
503 );
504
505 // When the `grep` tool is disabled, it should not be mentioned in the prompt
506 let model_context = ModelContext {
507 available_tools: [].to_vec(),
508 };
509 let prompt_without_grep = prompt_builder
510 .generate_assistant_system_prompt(&project_context, &model_context)
511 .unwrap();
512 assert!(
513 !prompt_without_grep.contains("grep"),
514 "`grep` tool should not be mentioned in prompt when the tool is disabled"
515 );
516 }
517
518 #[test]
519 fn test_has_tool_helper() {
520 let mut handlebars = Handlebars::new();
521 handlebars.register_helper("has_tool", Box::new(PromptBuilder::has_tool_helper));
522 handlebars
523 .register_template_string(
524 "test_template",
525 "{{#if (has_tool 'grep')}}grep is enabled{{else}}grep is disabled{{/if}}",
526 )
527 .unwrap();
528
529 // grep available
530 let data = serde_json::json!({"available_tools": ["grep", "fetch"]});
531 let result = handlebars.render("test_template", &data).unwrap();
532 assert_eq!(result, "grep is enabled");
533
534 // grep not available
535 let data = serde_json::json!({"available_tools": ["terminal", "fetch"]});
536 let result = handlebars.render("test_template", &data).unwrap();
537 assert_eq!(result, "grep is disabled");
538 }
539}