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