tenant.openagents/omega

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