tenant.openagents/omega

1use anyhow::Result;
2use gpui::SharedString;
3use handlebars::Handlebars;
4use rust_embed::RustEmbed;
5use serde::Serialize;
6use std::sync::Arc;
7
8#[derive(RustEmbed)]
9#[folder = "src/templates"]
10#[include = "*.hbs"]
11struct Assets;
12
13pub struct Templates(Handlebars<'static>);
14
15impl Templates {
16    pub fn new() -> Arc<Self> {
17        let mut handlebars = Handlebars::new();
18        handlebars.set_strict_mode(true);
19        handlebars.register_helper("contains", Box::new(contains));
20        handlebars.register_embed_templates::<Assets>().unwrap();
21        Arc::new(Self(handlebars))
22    }
23}
24
25pub trait Template: Sized {
26    const TEMPLATE_NAME: &'static str;
27
28    fn render(&self, templates: &Templates) -> Result<String>
29    where
30        Self: Serialize + Sized,
31    {
32        Ok(templates.0.render(Self::TEMPLATE_NAME, self)?)
33    }
34}
35
36#[derive(Serialize)]
37pub struct SystemPromptTemplate<'a> {
38    #[serde(flatten)]
39    pub project: &'a prompt_store::ProjectContext,
40    pub available_tools: Vec<SharedString>,
41    pub model_name: Option<String>,
42    pub date: String,
43    /// Contents of the user-global `~/.config/zed/AGENTS.md` file (or the
44    /// platform equivalent), if present and non-empty.
45    pub user_agents_md: Option<SharedString>,
46    /// Whether agent-run terminal commands are wrapped in an OS-level
47    /// sandbox for this thread. When `true`, the rendered prompt
48    /// describes the sandbox's read/write/network rules and the
49    /// per-command flags the model can request to relax them. When
50    /// `false`, the prompt omits the sandbox section entirely.
51    pub sandboxing: bool,
52    /// Whether the host is Linux. The writable-temp story differs by
53    /// platform (Linux exposes an ephemeral `tmpfs` over `/tmp`; other
54    /// platforms provide a persistent per-thread `$TMPDIR`), so the sandbox
55    /// section describes the right one rather than advertising a `$TMPDIR`
56    /// that doesn't behave as stated.
57    pub is_linux: bool,
58    /// Whether sandboxed terminal commands run through WSL on Windows.
59    pub is_windows: bool,
60}
61
62impl Template for SystemPromptTemplate<'_> {
63    const TEMPLATE_NAME: &'static str = "system_prompt.hbs";
64}
65
66/// Handlebars helper for checking if an item is in a list
67fn contains(
68    h: &handlebars::Helper,
69    _: &handlebars::Handlebars,
70    _: &handlebars::Context,
71    _: &mut handlebars::RenderContext,
72    out: &mut dyn handlebars::Output,
73) -> handlebars::HelperResult {
74    let list = h
75        .param(0)
76        .and_then(|v| v.value().as_array())
77        .ok_or_else(|| {
78            handlebars::RenderError::new("contains: missing or invalid list parameter")
79        })?;
80    let query = h.param(1).map(|v| v.value()).ok_or_else(|| {
81        handlebars::RenderError::new("contains: missing or invalid query parameter")
82    })?;
83
84    if list.contains(query) {
85        out.write("true")?;
86    }
87
88    Ok(())
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn test_system_prompt_template() {
97        let project = prompt_store::ProjectContext::default();
98        let template = SystemPromptTemplate {
99            project: &project,
100            available_tools: vec!["echo".into()],
101            model_name: Some("test-model".to_string()),
102            date: "2026-01-01".to_string(),
103            user_agents_md: None,
104            sandboxing: false,
105            is_linux: false,
106            is_windows: false,
107        };
108        let templates = Templates::new();
109        let rendered = template.render(&templates).unwrap();
110        assert!(rendered.contains("You are the Omega coding agent"));
111        assert!(rendered.contains("Today's Date: 2026-01-01"));
112        assert!(rendered.contains("## Fixing Diagnostics"));
113        assert!(rendered.contains("test-model"));
114    }
115
116    #[test]
117    fn test_system_prompt_renders_user_agents_md_before_project_rules() {
118        use prompt_store::{ProjectContext, RulesFileContext, WorktreeContext};
119        use util::rel_path::RelPath;
120
121        let worktrees = vec![WorktreeContext {
122            root_name: "my-project".to_string(),
123            abs_path: std::path::Path::new("/tmp/my-project").into(),
124            rules_file: Some(RulesFileContext {
125                path_in_worktree: RelPath::from_unix_str("AGENTS.md").unwrap().into(),
126                text: "project-specific guidance".to_string(),
127                project_entry_id: 1,
128            }),
129        }];
130        let project = ProjectContext::new(worktrees);
131        let template = SystemPromptTemplate {
132            project: &project,
133            available_tools: vec!["echo".into()],
134            model_name: Some("test-model".to_string()),
135            date: "2026-01-01".to_string(),
136            user_agents_md: Some("always be concise".into()),
137            sandboxing: false,
138            is_linux: false,
139            is_windows: false,
140        };
141        let templates = Templates::new();
142        let rendered = template.render(&templates).unwrap();
143
144        assert!(rendered.contains("### Personal `AGENTS.md`"));
145        assert!(rendered.contains("always be concise"));
146        assert!(rendered.contains("### Project Rules"));
147        assert!(rendered.contains("project-specific guidance"));
148
149        let personal_idx = rendered.find("### Personal `AGENTS.md`").unwrap();
150        let project_idx = rendered.find("### Project Rules").unwrap();
151        assert!(
152            personal_idx < project_idx,
153            "personal AGENTS.md should render before project rules so project rules can override it"
154        );
155    }
156
157    #[test]
158    fn test_system_prompt_omits_sandbox_section_when_sandboxing_disabled() {
159        let project = prompt_store::ProjectContext::default();
160        let template = SystemPromptTemplate {
161            project: &project,
162            available_tools: vec!["echo".into()],
163            model_name: Some("test-model".to_string()),
164            date: "2026-01-01".to_string(),
165            user_agents_md: None,
166            sandboxing: false,
167            is_linux: false,
168            is_windows: false,
169        };
170        let templates = Templates::new();
171        let rendered = template.render(&templates).unwrap();
172        assert!(!rendered.contains("## Terminal sandbox"));
173        assert!(!rendered.contains("allow_hosts"));
174    }
175
176    #[test]
177    fn test_system_prompt_renders_sandbox_section_with_worktrees_when_enabled() {
178        use prompt_store::{ProjectContext, WorktreeContext};
179
180        let worktrees = vec![
181            WorktreeContext {
182                root_name: "alpha".to_string(),
183                abs_path: std::path::Path::new("/tmp/alpha").into(),
184                rules_file: None,
185            },
186            WorktreeContext {
187                root_name: "beta".to_string(),
188                abs_path: std::path::Path::new("/tmp/beta").into(),
189                rules_file: None,
190            },
191        ];
192        let project = ProjectContext::new(worktrees);
193        let template = SystemPromptTemplate {
194            project: &project,
195            available_tools: vec!["echo".into()],
196            model_name: Some("test-model".to_string()),
197            date: "2026-01-01".to_string(),
198            user_agents_md: None,
199            sandboxing: true,
200            is_linux: false,
201            is_windows: false,
202        };
203        let templates = Templates::new();
204        let rendered = template.render(&templates).unwrap();
205
206        assert!(rendered.contains("## Terminal sandbox"));
207        assert!(rendered.contains("`/tmp/alpha`"));
208        assert!(rendered.contains("`/tmp/beta`"));
209        assert!(rendered.contains("allow_hosts"));
210        assert!(rendered.contains("allow_all_hosts: true"));
211        assert!(rendered.contains("fs_write_paths"));
212        assert!(rendered.contains("allow_fs_write_all: true"));
213        assert!(rendered.contains("unsandboxed: true"));
214        assert!(rendered.contains("`.git` directories remain protected"));
215        assert!(rendered.contains("Git metadata writes are never grantable inside the sandbox"));
216        assert!(rendered.contains("request `unsandboxed: true` with a reason"));
217        assert!(rendered.contains("git --no-optional-locks status"));
218        assert!(rendered.contains("for the rest of the thread"));
219    }
220
221    #[test]
222    fn test_system_prompt_linux_sandbox_section_omits_tmpdir() {
223        use prompt_store::{ProjectContext, WorktreeContext};
224
225        let worktrees = vec![WorktreeContext {
226            root_name: "alpha".to_string(),
227            abs_path: std::path::Path::new("/tmp/alpha").into(),
228            rules_file: None,
229        }];
230        let project = ProjectContext::new(worktrees);
231        let template = SystemPromptTemplate {
232            project: &project,
233            available_tools: vec!["echo".into()],
234            model_name: Some("test-model".to_string()),
235            date: "2026-01-01".to_string(),
236            user_agents_md: None,
237            sandboxing: true,
238            is_linux: true,
239            is_windows: false,
240        };
241        let templates = Templates::new();
242        let rendered = template.render(&templates).unwrap();
243
244        assert!(rendered.contains("## Terminal sandbox"));
245        // On Linux we must not advertise the special persistent `$TMPDIR`.
246        assert!(!rendered.contains("$TMPDIR"));
247        assert!(rendered.contains("`/tmp` is writable"));
248        assert!(rendered.contains("`/tmp/alpha`"));
249    }
250
251    #[test]
252    fn test_system_prompt_windows_sandbox_section_rejects_host_specific_network() {
253        use prompt_store::{ProjectContext, WorktreeContext};
254
255        let worktrees = vec![WorktreeContext {
256            root_name: "alpha".to_string(),
257            abs_path: std::path::Path::new("C:/Users/me/project").into(),
258            rules_file: None,
259        }];
260        let project = ProjectContext::new(worktrees);
261        let template = SystemPromptTemplate {
262            project: &project,
263            available_tools: vec!["echo".into()],
264            model_name: Some("test-model".to_string()),
265            date: "2026-01-01".to_string(),
266            user_agents_md: None,
267            sandboxing: true,
268            is_linux: false,
269            is_windows: true,
270        };
271        let templates = Templates::new();
272        let rendered = template.render(&templates).unwrap();
273
274        assert!(rendered.contains("commands run inside WSL under Bubblewrap"));
275        assert!(rendered.contains("Protected Git metadata remains read-only"));
276        assert!(rendered.contains("do not use this on Windows"));
277        assert!(rendered.contains("such requests are rejected"));
278        assert!(rendered.contains("allow_all_hosts: true"));
279        assert!(rendered.contains("git --no-optional-locks status"));
280    }
281
282    #[test]
283    fn test_system_prompt_sandbox_section_handles_zero_worktrees() {
284        let project = prompt_store::ProjectContext::default();
285        let template = SystemPromptTemplate {
286            project: &project,
287            available_tools: vec!["echo".into()],
288            model_name: Some("test-model".to_string()),
289            date: "2026-01-01".to_string(),
290            user_agents_md: None,
291            sandboxing: true,
292            is_linux: false,
293            is_windows: false,
294        };
295        let templates = Templates::new();
296        let rendered = template.render(&templates).unwrap();
297
298        assert!(rendered.contains("## Terminal sandbox"));
299        assert!(rendered.contains("No project directories are currently writable"));
300    }
301
302    #[test]
303    fn test_system_prompt_omits_user_agents_md_section_when_absent() {
304        let project = prompt_store::ProjectContext::default();
305        let template = SystemPromptTemplate {
306            project: &project,
307            available_tools: vec!["echo".into()],
308            model_name: Some("test-model".to_string()),
309            date: "2026-01-01".to_string(),
310            user_agents_md: None,
311            sandboxing: false,
312            is_linux: false,
313            is_windows: false,
314        };
315        let templates = Templates::new();
316        let rendered = template.render(&templates).unwrap();
317        assert!(!rendered.contains("### Personal `AGENTS.md`"));
318    }
319
320    #[test]
321    fn test_system_prompt_does_not_render_legacy_zed_rules_section() {
322        let project = prompt_store::ProjectContext::default();
323        let template = SystemPromptTemplate {
324            project: &project,
325            available_tools: vec!["echo".into()],
326            model_name: Some("test-model".to_string()),
327            date: "2026-01-01".to_string(),
328            user_agents_md: None,
329            sandboxing: false,
330            is_linux: false,
331            is_windows: false,
332        };
333        let templates = Templates::new();
334        let rendered = template.render(&templates).unwrap();
335
336        assert!(!rendered.contains("The user has specified the following rules"));
337        assert!(!rendered.contains("Rules title:"));
338    }
339}
340