use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use crate::error::{Result, ZeptoError};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentTemplate {
pub name: String,
pub description: String,
pub system_prompt: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub allowed_tools: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub blocked_tools: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_tool_iterations: Option<u32>,
#[serde(default)]
pub tags: Vec<String>,
}
fn builtin_coder() -> AgentTemplate {
AgentTemplate {
name: "coder".to_string(),
description: "Expert software engineer that writes clean, well-tested code".to_string(),
system_prompt: concat!(
"You are an expert software engineer. You write clean, idiomatic, ",
"well-structured code following best practices for the language and ",
"framework in use. You explain your reasoning, consider edge cases, ",
"handle errors properly, and write tests when appropriate. You prefer ",
"simple, readable solutions over clever ones."
)
.to_string(),
model: None,
max_tokens: None,
temperature: None,
allowed_tools: None, blocked_tools: None,
max_tool_iterations: None,
tags: vec!["development".to_string(), "coding".to_string()],
}
}
fn builtin_researcher() -> AgentTemplate {
AgentTemplate {
name: "researcher".to_string(),
description: "Research assistant that finds, analyzes, and summarizes information"
.to_string(),
system_prompt: concat!(
"You are a thorough research assistant. You find relevant information, ",
"analyze it critically, summarize your findings clearly, and cite your ",
"sources. You distinguish between facts and opinions, flag uncertainties, ",
"and present multiple perspectives when they exist."
)
.to_string(),
model: None,
max_tokens: None,
temperature: None,
allowed_tools: Some(vec![
"web_search".to_string(),
"web_fetch".to_string(),
"memory_search".to_string(),
"memory_get".to_string(),
"longterm_memory".to_string(),
]),
blocked_tools: None,
max_tool_iterations: None,
tags: vec!["research".to_string(), "information".to_string()],
}
}
fn builtin_writer() -> AgentTemplate {
AgentTemplate {
name: "writer".to_string(),
description: "Professional writer that produces clear, concise, well-structured content"
.to_string(),
system_prompt: concat!(
"You are a professional writer. You produce clear, concise, and ",
"well-structured content. You adapt your tone and style to the context ",
"and audience. You pay attention to grammar, flow, and readability. ",
"You organize ideas logically and use examples to illustrate points."
)
.to_string(),
model: None,
max_tokens: None,
temperature: None,
allowed_tools: Some(vec![
"read_file".to_string(),
"write_file".to_string(),
"edit_file".to_string(),
"memory_search".to_string(),
"memory_get".to_string(),
]),
blocked_tools: None,
max_tool_iterations: None,
tags: vec!["writing".to_string(), "content".to_string()],
}
}
fn builtin_assistant() -> AgentTemplate {
AgentTemplate {
name: "assistant".to_string(),
description: "Helpful general assistant for everyday tasks and questions".to_string(),
system_prompt: concat!(
"You are a helpful general assistant. You answer questions accurately, ",
"perform tasks efficiently, and communicate clearly. You ask for ",
"clarification when a request is ambiguous and provide concise but ",
"complete responses."
)
.to_string(),
model: None,
max_tokens: None,
temperature: None,
allowed_tools: None, blocked_tools: None,
max_tool_iterations: None,
tags: vec!["general".to_string()],
}
}
fn builtin_templates() -> Vec<AgentTemplate> {
vec![
builtin_coder(),
builtin_researcher(),
builtin_writer(),
builtin_assistant(),
]
}
#[derive(Debug, Clone)]
pub struct TemplateRegistry {
templates: HashMap<String, AgentTemplate>,
}
impl TemplateRegistry {
pub fn new() -> Self {
let mut templates = HashMap::new();
for tpl in builtin_templates() {
templates.insert(tpl.name.clone(), tpl);
}
Self { templates }
}
pub fn register(&mut self, template: AgentTemplate) {
self.templates.insert(template.name.clone(), template);
}
pub fn get(&self, name: &str) -> Option<&AgentTemplate> {
self.templates.get(name)
}
pub fn list(&self) -> Vec<&AgentTemplate> {
self.templates.values().collect()
}
pub fn list_by_tag(&self, tag: &str) -> Vec<&AgentTemplate> {
self.templates
.values()
.filter(|t| t.tags.iter().any(|t_tag| t_tag == tag))
.collect()
}
pub fn names(&self) -> Vec<&str> {
self.templates.keys().map(|k| k.as_str()).collect()
}
pub fn load_from_dir(dir: &Path) -> Result<Vec<AgentTemplate>> {
if !dir.exists() {
return Ok(Vec::new());
}
if !dir.is_dir() {
return Err(ZeptoError::Config(format!(
"Template path is not a directory: {}",
dir.display()
)));
}
let entries = std::fs::read_dir(dir)?;
let mut templates = Vec::new();
for entry in entries {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("json") {
continue;
}
match std::fs::read_to_string(&path) {
Ok(content) => match serde_json::from_str::<AgentTemplate>(&content) {
Ok(template) => {
templates.push(template);
}
Err(e) => {
tracing::warn!(
path = %path.display(),
error = %e,
"Skipping invalid template file"
);
}
},
Err(e) => {
tracing::warn!(
path = %path.display(),
error = %e,
"Failed to read template file"
);
}
}
}
Ok(templates)
}
pub fn merge_from_dir(&mut self, dir: &Path) -> Result<usize> {
let templates = Self::load_from_dir(dir)?;
let count = templates.len();
for template in templates {
self.register(template);
}
Ok(count)
}
}
impl Default for TemplateRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
#[test]
fn test_builtin_coder_exists() {
let registry = TemplateRegistry::new();
let coder = registry.get("coder");
assert!(coder.is_some());
let coder = coder.unwrap();
assert_eq!(coder.name, "coder");
assert!(coder.allowed_tools.is_none()); assert!(coder.tags.contains(&"development".to_string()));
assert!(coder.tags.contains(&"coding".to_string()));
}
#[test]
fn test_builtin_researcher_exists() {
let registry = TemplateRegistry::new();
let researcher = registry.get("researcher");
assert!(researcher.is_some());
let researcher = researcher.unwrap();
assert_eq!(researcher.name, "researcher");
assert!(researcher.allowed_tools.is_some());
let tools = researcher.allowed_tools.as_ref().unwrap();
assert!(tools.contains(&"web_search".to_string()));
assert!(tools.contains(&"web_fetch".to_string()));
assert!(tools.contains(&"memory_search".to_string()));
assert!(researcher.tags.contains(&"research".to_string()));
}
#[test]
fn test_builtin_writer_exists() {
let registry = TemplateRegistry::new();
let writer = registry.get("writer");
assert!(writer.is_some());
let writer = writer.unwrap();
assert_eq!(writer.name, "writer");
assert!(writer.allowed_tools.is_some());
let tools = writer.allowed_tools.as_ref().unwrap();
assert!(tools.contains(&"read_file".to_string()));
assert!(tools.contains(&"write_file".to_string()));
assert!(tools.contains(&"edit_file".to_string()));
assert!(writer.tags.contains(&"writing".to_string()));
}
#[test]
fn test_builtin_assistant_exists() {
let registry = TemplateRegistry::new();
let assistant = registry.get("assistant");
assert!(assistant.is_some());
let assistant = assistant.unwrap();
assert_eq!(assistant.name, "assistant");
assert!(assistant.allowed_tools.is_none()); assert!(assistant.tags.contains(&"general".to_string()));
}
#[test]
fn test_lookup_returns_none_for_unknown() {
let registry = TemplateRegistry::new();
assert!(registry.get("nonexistent-template").is_none());
assert!(registry.get("").is_none());
assert!(registry.get("CODER").is_none()); }
#[test]
fn test_list_all_templates() {
let registry = TemplateRegistry::new();
let all = registry.list();
assert_eq!(all.len(), 4);
let names: Vec<&str> = all.iter().map(|t| t.name.as_str()).collect();
assert!(names.contains(&"coder"));
assert!(names.contains(&"researcher"));
assert!(names.contains(&"writer"));
assert!(names.contains(&"assistant"));
}
#[test]
fn test_list_by_tag() {
let registry = TemplateRegistry::new();
let dev = registry.list_by_tag("development");
assert_eq!(dev.len(), 1);
assert_eq!(dev[0].name, "coder");
let research = registry.list_by_tag("research");
assert_eq!(research.len(), 1);
assert_eq!(research[0].name, "researcher");
let general = registry.list_by_tag("general");
assert_eq!(general.len(), 1);
assert_eq!(general[0].name, "assistant");
let empty = registry.list_by_tag("nonexistent-tag");
assert!(empty.is_empty());
}
#[test]
fn test_register_custom_template() {
let mut registry = TemplateRegistry::new();
let custom = AgentTemplate {
name: "devops".to_string(),
description: "DevOps engineer".to_string(),
system_prompt: "You are a DevOps engineer.".to_string(),
model: Some("gpt-5.1".to_string()),
max_tokens: Some(4096),
temperature: Some(0.3),
allowed_tools: Some(vec!["shell".to_string()]),
blocked_tools: None,
max_tool_iterations: Some(10),
tags: vec!["devops".to_string(), "infrastructure".to_string()],
};
registry.register(custom);
assert_eq!(registry.list().len(), 5);
let devops = registry.get("devops").unwrap();
assert_eq!(devops.model, Some("gpt-5.1".to_string()));
assert_eq!(devops.max_tokens, Some(4096));
assert_eq!(devops.temperature, Some(0.3));
assert_eq!(devops.max_tool_iterations, Some(10));
}
#[test]
fn test_custom_template_overrides_builtin() {
let mut registry = TemplateRegistry::new();
let original = registry.get("coder").unwrap();
assert!(original.model.is_none());
let custom_coder = AgentTemplate {
name: "coder".to_string(),
description: "Custom coder".to_string(),
system_prompt: "You are a Rust expert.".to_string(),
model: Some("claude-sonnet-4-5-20250929".to_string()),
max_tokens: None,
temperature: None,
allowed_tools: Some(vec!["shell".to_string(), "read_file".to_string()]),
blocked_tools: None,
max_tool_iterations: None,
tags: vec!["development".to_string(), "rust".to_string()],
};
registry.register(custom_coder);
let overridden = registry.get("coder").unwrap();
assert_eq!(overridden.description, "Custom coder");
assert_eq!(
overridden.model,
Some("claude-sonnet-4-5-20250929".to_string())
);
assert!(overridden.tags.contains(&"rust".to_string()));
assert_eq!(registry.list().len(), 4);
}
#[test]
fn test_names_list() {
let registry = TemplateRegistry::new();
let names = registry.names();
assert_eq!(names.len(), 4);
assert!(names.contains(&"coder"));
assert!(names.contains(&"researcher"));
assert!(names.contains(&"writer"));
assert!(names.contains(&"assistant"));
}
#[test]
fn test_load_from_directory_with_json_files() {
let temp_dir = std::env::temp_dir().join("zeptoclaw_tpl_test_load");
let _ = fs::remove_dir_all(&temp_dir);
fs::create_dir_all(&temp_dir).unwrap();
let tpl1 = r#"{
"name": "ops-agent",
"description": "Operations agent",
"system_prompt": "You handle operations.",
"model": "gpt-5.1",
"tags": ["ops"]
}"#;
fs::write(temp_dir.join("ops-agent.json"), tpl1).unwrap();
let tpl2 = r#"{
"name": "data-analyst",
"description": "Data analysis agent",
"system_prompt": "You analyze data.",
"allowed_tools": ["read_file", "shell"],
"tags": ["data", "analysis"]
}"#;
fs::write(temp_dir.join("data-analyst.json"), tpl2).unwrap();
fs::write(temp_dir.join("README.md"), "# Templates").unwrap();
let templates = TemplateRegistry::load_from_dir(&temp_dir).unwrap();
assert_eq!(templates.len(), 2);
let names: Vec<&str> = templates.iter().map(|t| t.name.as_str()).collect();
assert!(names.contains(&"ops-agent"));
assert!(names.contains(&"data-analyst"));
let ops = templates.iter().find(|t| t.name == "ops-agent").unwrap();
assert_eq!(ops.model, Some("gpt-5.1".to_string()));
fs::remove_dir_all(&temp_dir).ok();
}
#[test]
fn test_load_from_empty_directory() {
let temp_dir = std::env::temp_dir().join("zeptoclaw_tpl_test_empty");
let _ = fs::remove_dir_all(&temp_dir);
fs::create_dir_all(&temp_dir).unwrap();
let templates = TemplateRegistry::load_from_dir(&temp_dir).unwrap();
assert!(templates.is_empty());
fs::remove_dir_all(&temp_dir).ok();
}
#[test]
fn test_load_from_nonexistent_directory() {
let path = Path::new("/tmp/zeptoclaw_tpl_nonexistent_98765");
let templates = TemplateRegistry::load_from_dir(path).unwrap();
assert!(templates.is_empty());
}
#[test]
fn test_load_invalid_json_file() {
let temp_dir = std::env::temp_dir().join("zeptoclaw_tpl_test_invalid");
let _ = fs::remove_dir_all(&temp_dir);
fs::create_dir_all(&temp_dir).unwrap();
fs::write(temp_dir.join("broken.json"), "{ not valid json }}}").unwrap();
let valid = r#"{
"name": "valid-one",
"description": "Valid template",
"system_prompt": "Hello.",
"tags": []
}"#;
fs::write(temp_dir.join("valid.json"), valid).unwrap();
let templates = TemplateRegistry::load_from_dir(&temp_dir).unwrap();
assert_eq!(templates.len(), 1);
assert_eq!(templates[0].name, "valid-one");
fs::remove_dir_all(&temp_dir).ok();
}
#[test]
fn test_template_serialization_roundtrip() {
let template = AgentTemplate {
name: "roundtrip".to_string(),
description: "Roundtrip test".to_string(),
system_prompt: "Test prompt.".to_string(),
model: Some("gpt-5.1".to_string()),
max_tokens: Some(2048),
temperature: Some(0.5),
allowed_tools: Some(vec!["shell".to_string(), "read_file".to_string()]),
blocked_tools: Some(vec!["web_search".to_string()]),
max_tool_iterations: Some(15),
tags: vec!["test".to_string()],
};
let json = serde_json::to_string_pretty(&template).unwrap();
let deserialized: AgentTemplate = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.name, "roundtrip");
assert_eq!(deserialized.description, "Roundtrip test");
assert_eq!(deserialized.system_prompt, "Test prompt.");
assert_eq!(deserialized.model, Some("gpt-5.1".to_string()));
assert_eq!(deserialized.max_tokens, Some(2048));
assert_eq!(deserialized.temperature, Some(0.5));
assert_eq!(
deserialized.allowed_tools,
Some(vec!["shell".to_string(), "read_file".to_string()])
);
assert_eq!(
deserialized.blocked_tools,
Some(vec!["web_search".to_string()])
);
assert_eq!(deserialized.max_tool_iterations, Some(15));
assert_eq!(deserialized.tags, vec!["test".to_string()]);
}
#[test]
fn test_blocked_tools_field() {
let json = r#"{
"name": "restricted",
"description": "Restricted agent",
"system_prompt": "You have restrictions.",
"blocked_tools": ["shell", "write_file"],
"tags": ["restricted"]
}"#;
let template: AgentTemplate = serde_json::from_str(json).unwrap();
assert!(template.allowed_tools.is_none());
assert!(template.blocked_tools.is_some());
let blocked = template.blocked_tools.unwrap();
assert_eq!(blocked.len(), 2);
assert!(blocked.contains(&"shell".to_string()));
assert!(blocked.contains(&"write_file".to_string()));
}
#[test]
fn test_allowed_tools_none_means_all_tools() {
let json = r#"{
"name": "unrestricted",
"description": "Unrestricted agent",
"system_prompt": "You have all tools.",
"tags": []
}"#;
let template: AgentTemplate = serde_json::from_str(json).unwrap();
assert!(template.allowed_tools.is_none());
assert!(template.blocked_tools.is_none());
let registry = TemplateRegistry::new();
assert!(registry.get("coder").unwrap().allowed_tools.is_none());
assert!(registry.get("assistant").unwrap().allowed_tools.is_none());
}
#[test]
fn test_merge_from_dir() {
let temp_dir = std::env::temp_dir().join("zeptoclaw_tpl_test_merge");
let _ = fs::remove_dir_all(&temp_dir);
fs::create_dir_all(&temp_dir).unwrap();
let custom_coder = r#"{
"name": "coder",
"description": "Custom coder from file",
"system_prompt": "You are a Go expert.",
"model": "gpt-5.1",
"tags": ["development", "go"]
}"#;
fs::write(temp_dir.join("coder.json"), custom_coder).unwrap();
let new_tpl = r#"{
"name": "translator",
"description": "Language translator",
"system_prompt": "You translate text between languages.",
"tags": ["translation"]
}"#;
fs::write(temp_dir.join("translator.json"), new_tpl).unwrap();
let mut registry = TemplateRegistry::new();
assert_eq!(registry.list().len(), 4);
let count = registry.merge_from_dir(&temp_dir).unwrap();
assert_eq!(count, 2);
assert_eq!(registry.list().len(), 5);
let coder = registry.get("coder").unwrap();
assert_eq!(coder.description, "Custom coder from file");
assert_eq!(coder.model, Some("gpt-5.1".to_string()));
let translator = registry.get("translator").unwrap();
assert_eq!(translator.description, "Language translator");
fs::remove_dir_all(&temp_dir).ok();
}
#[test]
fn test_default_impl_matches_new() {
let from_new = TemplateRegistry::new();
let from_default = TemplateRegistry::default();
assert_eq!(from_new.names().len(), from_default.names().len());
}
#[test]
fn test_template_optional_fields_default_to_none() {
let minimal_json = r#"{
"name": "minimal",
"description": "Minimal template",
"system_prompt": "Hello."
}"#;
let template: AgentTemplate = serde_json::from_str(minimal_json).unwrap();
assert_eq!(template.name, "minimal");
assert!(template.model.is_none());
assert!(template.max_tokens.is_none());
assert!(template.temperature.is_none());
assert!(template.allowed_tools.is_none());
assert!(template.blocked_tools.is_none());
assert!(template.max_tool_iterations.is_none());
assert!(template.tags.is_empty());
}
#[test]
fn test_serialization_skips_none_fields() {
let template = AgentTemplate {
name: "sparse".to_string(),
description: "Sparse template".to_string(),
system_prompt: "Hello.".to_string(),
model: None,
max_tokens: None,
temperature: None,
allowed_tools: None,
blocked_tools: None,
max_tool_iterations: None,
tags: vec![],
};
let json = serde_json::to_string(&template).unwrap();
assert!(!json.contains("model"));
assert!(!json.contains("max_tokens"));
assert!(!json.contains("temperature"));
assert!(!json.contains("allowed_tools"));
assert!(!json.contains("blocked_tools"));
assert!(!json.contains("max_tool_iterations"));
}
#[test]
fn test_load_from_dir_not_a_directory() {
let temp_file = std::env::temp_dir().join("zeptoclaw_tpl_test_notdir.txt");
fs::write(&temp_file, "not a directory").unwrap();
let result = TemplateRegistry::load_from_dir(&temp_file);
assert!(result.is_err());
let err = result.unwrap_err();
assert!(matches!(err, ZeptoError::Config(msg) if msg.contains("not a directory")));
fs::remove_file(&temp_file).ok();
}
}