heartbit_core/template/
variables.rs

1//! Prompt variable substitution — replaces `{variable}` placeholders in system prompts.
2
3use std::collections::HashMap;
4
5/// Substitute `{var_name}` placeholders in text.
6///
7/// Single-pass, non-recursive. Unknown variables are left as-is
8/// (allows literal `{braces}` in prompts).
9pub fn substitute(text: &str, variables: &HashMap<String, String>) -> String {
10    if variables.is_empty() || !text.contains('{') {
11        return text.to_string();
12    }
13
14    let mut result = String::with_capacity(text.len());
15    let mut chars = text.char_indices().peekable();
16
17    while let Some((i, ch)) = chars.next() {
18        if ch == '{' {
19            // Look for closing brace
20            let rest = &text[i + 1..];
21            if let Some(close) = rest.find('}') {
22                let var_name = &rest[..close];
23                // Variable names: alphanumeric + underscores only
24                if !var_name.is_empty()
25                    && var_name.chars().all(|c| c.is_alphanumeric() || c == '_')
26                    && let Some(value) = variables.get(var_name)
27                {
28                    result.push_str(value);
29                    // Skip past the closing brace
30                    for _ in 0..close + 1 {
31                        chars.next();
32                    }
33                    continue;
34                }
35            }
36            result.push(ch);
37        } else {
38            result.push(ch);
39        }
40    }
41
42    result
43}
44
45/// Build the full variable context from built-in + custom variables.
46pub fn build_variables(
47    agent_name: &str,
48    agent_description: &str,
49    workspace: Option<&str>,
50    custom: &HashMap<String, String>,
51) -> HashMap<String, String> {
52    let mut vars = HashMap::with_capacity(custom.len() + 5);
53    vars.extend(custom.iter().map(|(k, v)| (k.clone(), v.clone())));
54
55    // Built-ins override custom (intentional — prevents spoofing)
56    vars.insert("agent_name".into(), agent_name.into());
57    vars.insert("agent_description".into(), agent_description.into());
58    vars.insert(
59        "workspace".into(),
60        workspace.unwrap_or("not configured").into(),
61    );
62    vars.insert(
63        "date".into(),
64        chrono::Utc::now().format("%Y-%m-%d").to_string(),
65    );
66    vars.insert(
67        "datetime".into(),
68        chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(),
69    );
70
71    vars
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn substitute_known_variable() {
80        let mut vars = HashMap::new();
81        vars.insert("name".into(), "Alice".into());
82        assert_eq!(substitute("Hello {name}!", &vars), "Hello Alice!");
83    }
84
85    #[test]
86    fn substitute_unknown_variable_preserved() {
87        let vars = HashMap::new();
88        assert_eq!(substitute("Hello {name}!", &vars), "Hello {name}!");
89    }
90
91    #[test]
92    fn substitute_multiple_variables() {
93        let mut vars = HashMap::new();
94        vars.insert("a".into(), "1".into());
95        vars.insert("b".into(), "2".into());
96        assert_eq!(substitute("{a} + {b} = 3", &vars), "1 + 2 = 3");
97    }
98
99    #[test]
100    fn substitute_empty_map_noop() {
101        let vars = HashMap::new();
102        assert_eq!(substitute("no vars here", &vars), "no vars here");
103    }
104
105    #[test]
106    fn substitute_no_braces_noop() {
107        let mut vars = HashMap::new();
108        vars.insert("x".into(), "y".into());
109        assert_eq!(substitute("no braces", &vars), "no braces");
110    }
111
112    #[test]
113    fn substitute_literal_braces_preserved() {
114        let vars = HashMap::new();
115        // JSON-like braces without valid var names
116        assert_eq!(
117            substitute("{\"key\": \"value\"}", &vars),
118            "{\"key\": \"value\"}"
119        );
120    }
121
122    #[test]
123    fn substitute_empty_braces_preserved() {
124        let vars = HashMap::new();
125        assert_eq!(substitute("empty {} here", &vars), "empty {} here");
126    }
127
128    #[test]
129    fn substitute_nested_braces_preserved() {
130        let mut vars = HashMap::new();
131        vars.insert("x".into(), "val".into());
132        // Only the inner valid var is replaced (outer brace is literal)
133        assert_eq!(substitute("{{x}}", &vars), "{val}");
134    }
135
136    #[test]
137    fn substitute_underscore_in_var_name() {
138        let mut vars = HashMap::new();
139        vars.insert("agent_name".into(), "coder".into());
140        assert_eq!(substitute("I am {agent_name}", &vars), "I am coder");
141    }
142
143    #[test]
144    fn build_variables_includes_builtins() {
145        let vars = build_variables(
146            "test-agent",
147            "A test agent",
148            Some("/workspace"),
149            &HashMap::new(),
150        );
151        assert_eq!(vars.get("agent_name").unwrap(), "test-agent");
152        assert_eq!(vars.get("agent_description").unwrap(), "A test agent");
153        assert_eq!(vars.get("workspace").unwrap(), "/workspace");
154        assert!(vars.contains_key("date"));
155        assert!(vars.contains_key("datetime"));
156    }
157
158    #[test]
159    fn build_variables_custom_merged() {
160        let mut custom = HashMap::new();
161        custom.insert("project".into(), "heartbit".into());
162        let vars = build_variables("a", "d", None, &custom);
163        assert_eq!(vars.get("project").unwrap(), "heartbit");
164        assert_eq!(vars.get("workspace").unwrap(), "not configured");
165    }
166
167    #[test]
168    fn build_variables_builtins_override_custom() {
169        let mut custom = HashMap::new();
170        custom.insert("agent_name".into(), "spoofed".into());
171        let vars = build_variables("real", "d", None, &custom);
172        assert_eq!(vars.get("agent_name").unwrap(), "real");
173    }
174}
heartbit_core/template/variables.rs

heartbit_core/template/
variables.rs
