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