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 available_skills: String,
 43}
 44
 45impl Template for SystemPromptTemplate<'_> {
 46    const TEMPLATE_NAME: &'static str = "system_prompt.hbs";
 47}
 48
 49/// Context for a single skill in the skills prompt template.
 50#[derive(Serialize)]
 51pub struct SkillContext {
 52    pub name: String,
 53    pub description: String,
 54    pub path: String,
 55}
 56
 57/// Template for rendering the available skills section of the system prompt.
 58#[derive(Serialize)]
 59pub struct SkillsPromptTemplate {
 60    pub skills: Vec<SkillContext>,
 61    pub has_skills: bool,
 62}
 63
 64impl Template for SkillsPromptTemplate {
 65    const TEMPLATE_NAME: &'static str = "skills_prompt.hbs";
 66}
 67
 68/// Handlebars helper for checking if an item is in a list
 69fn contains(
 70    h: &handlebars::Helper,
 71    _: &handlebars::Handlebars,
 72    _: &handlebars::Context,
 73    _: &mut handlebars::RenderContext,
 74    out: &mut dyn handlebars::Output,
 75) -> handlebars::HelperResult {
 76    let list = h
 77        .param(0)
 78        .and_then(|v| v.value().as_array())
 79        .ok_or_else(|| {
 80            handlebars::RenderError::new("contains: missing or invalid list parameter")
 81        })?;
 82    let query = h.param(1).map(|v| v.value()).ok_or_else(|| {
 83        handlebars::RenderError::new("contains: missing or invalid query parameter")
 84    })?;
 85
 86    if list.contains(query) {
 87        out.write("true")?;
 88    }
 89
 90    Ok(())
 91}
 92
 93#[cfg(test)]
 94mod tests {
 95    use super::*;
 96
 97    #[test]
 98    fn test_system_prompt_template() {
 99        let project = prompt_store::ProjectContext::default();
100        let template = SystemPromptTemplate {
101            project: &project,
102            available_tools: vec!["echo".into()],
103            available_skills: String::new(),
104            model_name: Some("test-model".to_string()),
105        };
106        let templates = Templates::new();
107        let rendered = template.render(&templates).unwrap();
108        assert!(rendered.contains("## Fixing Diagnostics"));
109        assert!(!rendered.contains("## Planning"));
110        assert!(rendered.contains("test-model"));
111    }
112
113    #[test]
114    fn test_skills_prompt_template() {
115        let templates = Templates::new();
116        let template = SkillsPromptTemplate {
117            skills: vec![SkillContext {
118                name: "test-skill".to_string(),
119                description: "A test skill description".to_string(),
120                path: "/home/user/.config/zed/skills/test-skill".to_string(),
121            }],
122            has_skills: true,
123        };
124        let rendered = template.render(&templates).unwrap();
125        assert!(rendered.contains("## Available Agent Skills"));
126        assert!(rendered.contains(
127            "| test-skill | A test skill description | /home/user/.config/zed/skills/test-skill |"
128        ));
129    }
130
131    #[test]
132    fn test_skills_prompt_template_empty() {
133        let templates = Templates::new();
134        let template = SkillsPromptTemplate {
135            skills: vec![],
136            has_skills: false,
137        };
138        let rendered = template.render(&templates).unwrap();
139        assert!(rendered.is_empty());
140    }
141}