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