Skip to main content

minion_engine/engine/
context.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::{Arc, Mutex};
4
5use crate::prompts::detector::StackInfo;
6use crate::steps::{ParsedValue, StepOutput};
7
8/// A single chat message (user or assistant turn)
9#[derive(Debug, Clone)]
10pub struct ChatMessage {
11    pub role: String,
12    pub content: String,
13}
14
15/// Ordered history of messages for a named chat session
16#[derive(Debug, Clone, Default)]
17pub struct ChatHistory {
18    pub messages: Vec<ChatMessage>,
19}
20
21/// Shared chat session store — Arc so child contexts inherit the same store
22type ChatSessionStore = Arc<Mutex<HashMap<String, ChatHistory>>>;
23
24/// Tree-structured context that stores step outputs
25pub struct Context {
26    steps: HashMap<String, StepOutput>,
27    parsed_outputs: HashMap<String, ParsedValue>,
28    variables: HashMap<String, serde_json::Value>,
29    parent: Option<Arc<Context>>,
30    pub scope_value: Option<serde_json::Value>,
31    pub scope_index: usize,
32    pub session_id: Option<String>,
33    /// Shared chat session store — inherited by child contexts via Arc clone
34    chat_sessions: ChatSessionStore,
35    /// Detected stack info for prompt resolution (Story 11.5/11.6)
36    pub stack_info: Option<StackInfo>,
37    /// Directory where prompt template files live (defaults to "prompts")
38    pub prompts_dir: PathBuf,
39}
40
41impl Context {
42    pub fn new(target: String, vars: HashMap<String, serde_json::Value>) -> Self {
43        let mut variables = vars;
44        variables.insert("target".to_string(), serde_json::Value::String(target));
45
46        Self {
47            steps: HashMap::new(),
48            parsed_outputs: HashMap::new(),
49            variables,
50            parent: None,
51            scope_value: None,
52            scope_index: 0,
53            session_id: None,
54            chat_sessions: Arc::new(Mutex::new(HashMap::new())),
55            stack_info: None,
56            prompts_dir: PathBuf::from("prompts"),
57        }
58    }
59
60    /// Store a step output
61    pub fn store(&mut self, name: &str, output: StepOutput) {
62        if let StepOutput::Agent(ref agent) = output {
63            if let Some(ref sid) = agent.session_id {
64                self.session_id = Some(sid.clone());
65            }
66        }
67        self.steps.insert(name.to_string(), output);
68    }
69
70    /// Get a step output (looks in parent if not found locally)
71    pub fn get_step(&self, name: &str) -> Option<&StepOutput> {
72        self.steps
73            .get(name)
74            .or_else(|| self.parent.as_ref().and_then(|p| p.get_step(name)))
75    }
76
77    /// Insert a variable into this context
78    pub fn insert_var(&mut self, name: impl Into<String>, value: serde_json::Value) {
79        self.variables.insert(name.into(), value);
80    }
81
82    /// Get a variable
83    pub fn get_var(&self, name: &str) -> Option<&serde_json::Value> {
84        self.variables
85            .get(name)
86            .or_else(|| self.parent.as_ref().and_then(|p| p.get_var(name)))
87    }
88
89    /// Get session ID (searches parent chain)
90    pub fn get_session(&self) -> Option<&str> {
91        self.session_id
92            .as_deref()
93            .or_else(|| self.parent.as_ref().and_then(|p| p.get_session()))
94    }
95
96    /// Store a parsed value for a step
97    pub fn store_parsed(&mut self, name: &str, parsed: ParsedValue) {
98        self.parsed_outputs.insert(name.to_string(), parsed);
99    }
100
101    /// Get a parsed value for a step (looks in parent if not found locally)
102    pub fn get_parsed(&self, name: &str) -> Option<&ParsedValue> {
103        self.parsed_outputs
104            .get(name)
105            .or_else(|| self.parent.as_ref().and_then(|p| p.get_parsed(name)))
106    }
107
108    /// Create a child context for a scope
109    pub fn child(
110        parent: Arc<Context>,
111        scope_value: Option<serde_json::Value>,
112        index: usize,
113    ) -> Self {
114        let stack_info = parent.stack_info.clone();
115        let prompts_dir = parent.prompts_dir.clone();
116        Self {
117            steps: HashMap::new(),
118            parsed_outputs: HashMap::new(),
119            variables: HashMap::new(),
120            parent: Some(parent.clone()),
121            scope_value,
122            scope_index: index,
123            session_id: parent.session_id.clone(),
124            chat_sessions: Arc::clone(&parent.chat_sessions),
125            stack_info,
126            prompts_dir,
127        }
128    }
129
130    /// Get all variables (local + parent chain) merged into a flat HashMap
131    pub fn all_variables(&self) -> HashMap<String, serde_json::Value> {
132        let mut result = HashMap::new();
133        if let Some(ref parent) = self.parent {
134            result = parent.all_variables();
135        }
136        result.extend(self.variables.clone());
137        result
138    }
139
140    /// Get the stack_info from this context or any parent
141    pub fn get_stack_info(&self) -> Option<&StackInfo> {
142        self.stack_info
143            .as_ref()
144            .or_else(|| self.parent.as_ref().and_then(|p| p.get_stack_info()))
145    }
146
147    /// Get the Tera-ready value for a step by name (used by from() preprocessing).
148    /// Returns None if the step doesn't exist in this context or any parent.
149    pub fn get_from_value(&self, name: &str) -> Option<serde_json::Value> {
150        let step = self.get_step(name)?;
151        let parsed = self.get_parsed(name);
152        Some(step_output_to_value_with_parsed(step, parsed))
153    }
154
155    /// Check if a dotted path variable exists in this context (used for strict accessor)
156    pub fn var_exists(&self, path: &str) -> bool {
157        let parts: Vec<&str> = path.split('.').collect();
158        if parts.is_empty() {
159            return false;
160        }
161        let root = parts[0];
162        if let Some(step) = self.get_step(root) {
163            if parts.len() == 1 {
164                return true;
165            }
166            let val = step_output_to_value_with_parsed(step, self.get_parsed(root));
167            return check_json_path(&val, &parts[1..]);
168        }
169        if let Some(var) = self.get_var(root) {
170            if parts.len() == 1 {
171                return true;
172            }
173            return check_json_path(var, &parts[1..]);
174        }
175        false
176    }
177
178    /// Return all stored messages for a chat session (empty vec if session doesn't exist)
179    pub fn get_chat_messages(&self, session: &str) -> Vec<ChatMessage> {
180        let guard = self
181            .chat_sessions
182            .lock()
183            .expect("chat_sessions lock poisoned");
184        guard
185            .get(session)
186            .map(|h| h.messages.clone())
187            .unwrap_or_default()
188    }
189
190    /// Append messages to a chat session, creating the session if it doesn't exist
191    pub fn append_chat_messages(&self, session: &str, messages: Vec<ChatMessage>) {
192        let mut guard = self
193            .chat_sessions
194            .lock()
195            .expect("chat_sessions lock poisoned");
196        let history = guard
197            .entry(session.to_string())
198            .or_insert_with(ChatHistory::default);
199        history.messages.extend(messages);
200    }
201
202    /// Convert to Tera template context
203    pub fn to_tera_context(&self) -> tera::Context {
204        let mut ctx = tera::Context::new();
205
206        // Add variables (parent first, then override with local)
207        if let Some(parent) = &self.parent {
208            ctx = parent.to_tera_context();
209        }
210        for (k, v) in &self.variables {
211            ctx.insert(k, v);
212        }
213
214        // Build full steps map (parent + local)
215        let mut steps_map: HashMap<String, serde_json::Value> = HashMap::new();
216        if let Some(parent) = &self.parent {
217            collect_steps_with_parsed(parent, &mut steps_map);
218        }
219        for (name, output) in &self.steps {
220            let parsed = self.parsed_outputs.get(name);
221            let val = step_output_to_value_with_parsed(output, parsed);
222            steps_map.insert(name.clone(), val);
223        }
224
225        // Insert steps both under "steps" and directly by name for flexible access
226        for (name, val) in &steps_map {
227            ctx.insert(name.as_str(), val);
228        }
229        ctx.insert("steps", &steps_map);
230
231        // Add scope
232        if let Some(sv) = &self.scope_value {
233            let mut scope_map = HashMap::new();
234            scope_map.insert("value".to_string(), sv.clone());
235            scope_map.insert("index".to_string(), serde_json::json!(self.scope_index));
236            ctx.insert("scope", &scope_map);
237        }
238
239        ctx
240    }
241
242    /// Render a template string with this context
243    pub fn render_template(&self, template: &str) -> Result<String, crate::error::StepError> {
244        // Pre-process template: handle ?, !, and from("name") → __from_name__ substitution
245        let pre = crate::engine::template::preprocess_template(template, self)?;
246
247        // Build base Tera context (steps, vars, scope)
248        let mut tera_ctx = self.to_tera_context();
249
250        // Inject from() lookup variables into the Tera context
251        for (k, v) in &pre.injected {
252            tera_ctx.insert(k.as_str(), v);
253        }
254
255        // Build Tera instance and render
256        let mut tera = tera::Tera::default();
257        tera.add_raw_template("__tmpl__", &pre.template)
258            .map_err(|e| crate::error::StepError::Template(format!("{e}")))?;
259
260        tera.render("__tmpl__", &tera_ctx)
261            .map_err(|e| crate::error::StepError::Template(format!("{e}")))
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use crate::steps::{CmdOutput, StepOutput};
269    use std::time::Duration;
270
271    fn cmd_output(stdout: &str, exit_code: i32) -> StepOutput {
272        StepOutput::Cmd(CmdOutput {
273            stdout: stdout.to_string(),
274            stderr: String::new(),
275            exit_code,
276            duration: Duration::ZERO,
277        })
278    }
279
280    #[test]
281    fn store_and_retrieve() {
282        let mut ctx = Context::new("123".to_string(), HashMap::new());
283        ctx.store("step1", cmd_output("hello", 0));
284        let out = ctx.get_step("step1").unwrap();
285        assert_eq!(out.text(), "hello");
286        assert_eq!(out.exit_code(), 0);
287    }
288
289    #[test]
290    fn parent_context_inheritance() {
291        let mut parent = Context::new("456".to_string(), HashMap::new());
292        parent.store("parent_step", cmd_output("from parent", 0));
293        let child = Context::child(Arc::new(parent), None, 0);
294        let out = child.get_step("parent_step").unwrap();
295        assert_eq!(out.text(), "from parent");
296    }
297
298    #[test]
299    fn target_variable_resolves() {
300        let ctx = Context::new("42".to_string(), HashMap::new());
301        let result = ctx.render_template("{{ target }}").unwrap();
302        assert_eq!(result, "42");
303    }
304
305    #[test]
306    fn render_template_with_step_stdout() {
307        let mut ctx = Context::new("".to_string(), HashMap::new());
308        ctx.store("fetch", cmd_output("some output", 0));
309        let result = ctx.render_template("{{ steps.fetch.stdout }}").unwrap();
310        assert_eq!(result, "some output");
311    }
312
313    #[test]
314    fn render_scope_value() {
315        let parent = Context::new("".to_string(), HashMap::new());
316        let child = Context::child(Arc::new(parent), Some(serde_json::json!("my_value")), 0);
317        let result = child.render_template("{{ scope.value }}").unwrap();
318        assert_eq!(result, "my_value");
319    }
320
321    #[test]
322    fn render_template_with_step_exit_code() {
323        let mut ctx = Context::new("".to_string(), HashMap::new());
324        ctx.store("prev", cmd_output("output", 0));
325        let result = ctx.render_template("{{ steps.prev.exit_code }}").unwrap();
326        assert_eq!(result, "0");
327    }
328
329    #[test]
330    fn agent_session_id_accessible_in_template() {
331        use crate::steps::{AgentOutput, AgentStats, StepOutput};
332        let mut ctx = Context::new("".to_string(), HashMap::new());
333        ctx.store(
334            "scan",
335            StepOutput::Agent(AgentOutput {
336                response: "done".to_string(),
337                session_id: Some("sess-abc".to_string()),
338                stats: AgentStats::default(),
339            }),
340        );
341        let result = ctx.render_template("{{ steps.scan.session_id }}").unwrap();
342        assert_eq!(result, "sess-abc");
343    }
344
345    #[test]
346    fn cmd_step_session_id_is_empty_string() {
347        let mut ctx = Context::new("".to_string(), HashMap::new());
348        ctx.store("build", cmd_output("output", 0));
349        let result = ctx.render_template("{{ steps.build.session_id }}").unwrap();
350        assert_eq!(result, "");
351    }
352
353    #[test]
354    fn agent_session_id_none_renders_empty_string() {
355        use crate::steps::{AgentOutput, AgentStats, StepOutput};
356        let mut ctx = Context::new("".to_string(), HashMap::new());
357        ctx.store(
358            "scan",
359            StepOutput::Agent(AgentOutput {
360                response: "done".to_string(),
361                session_id: None,
362                stats: AgentStats::default(),
363            }),
364        );
365        let result = ctx.render_template("{{ steps.scan.session_id }}").unwrap();
366        assert_eq!(result, "");
367    }
368
369    #[test]
370    fn child_inherits_parent_steps() {
371        let mut parent = Context::new("test".to_string(), HashMap::new());
372        parent.store("a", cmd_output("alpha", 0));
373        let mut child = Context::child(Arc::new(parent), None, 0);
374        child.store("b", cmd_output("beta", 0));
375        // Child can see parent step
376        assert!(child.get_step("a").is_some());
377        // Child can see own step
378        assert!(child.get_step("b").is_some());
379    }
380
381    #[test]
382    fn output_key_defaults_to_text() {
383        let mut ctx = Context::new("".to_string(), HashMap::new());
384        ctx.store("fetch", cmd_output("hello world", 0));
385        // Without parsed value, {{fetch.output}} returns the raw text
386        let result = ctx.render_template("{{ fetch.output }}").unwrap();
387        assert_eq!(result, "hello world");
388    }
389
390    #[test]
391    fn output_key_with_json_parsed_value() {
392        use crate::steps::ParsedValue;
393        let mut ctx = Context::new("".to_string(), HashMap::new());
394        ctx.store("scan", cmd_output(r#"{"count": 5}"#, 0));
395        ctx.store_parsed("scan", ParsedValue::Json(serde_json::json!({"count": 5})));
396        // JSON parsed value allows dot-path access
397        let result = ctx.render_template("{{ scan.output.count }}").unwrap();
398        assert_eq!(result, "5");
399    }
400
401    #[test]
402    fn output_key_with_lines_parsed_value() {
403        use crate::steps::ParsedValue;
404        let mut ctx = Context::new("".to_string(), HashMap::new());
405        ctx.store("files", cmd_output("a.rs\nb.rs\nc.rs", 0));
406        ctx.store_parsed(
407            "files",
408            ParsedValue::Lines(vec!["a.rs".into(), "b.rs".into(), "c.rs".into()]),
409        );
410        // Lines parsed value renders with | length filter
411        let result = ctx.render_template("{{ files.output | length }}").unwrap();
412        assert_eq!(result, "3");
413    }
414
415    #[test]
416    fn step_accessible_directly_by_name() {
417        let mut ctx = Context::new("".to_string(), HashMap::new());
418        ctx.store("greet", cmd_output("hi", 0));
419        // Steps are also accessible directly by name (not just via steps.)
420        let result = ctx.render_template("{{ greet.output }}").unwrap();
421        assert_eq!(result, "hi");
422    }
423
424    #[test]
425    fn from_accesses_step_by_name() {
426        let mut ctx = Context::new("".to_string(), HashMap::new());
427        ctx.store("global-config", cmd_output("prod", 0));
428        // from("name") syntax allows accessing any step by name
429        let result = ctx
430            .render_template(r#"{{ from("global-config").output }}"#)
431            .unwrap();
432        assert_eq!(result, "prod");
433    }
434
435    #[test]
436    fn from_fails_for_nonexistent_step() {
437        let ctx = Context::new("".to_string(), HashMap::new());
438        let err = ctx
439            .render_template(r#"{{ from("nonexistent").output }}"#)
440            .unwrap_err();
441        assert!(
442            err.to_string().contains("not found"),
443            "expected 'not found' error, got: {err}"
444        );
445    }
446
447    #[test]
448    fn from_with_json_dot_access() {
449        use crate::steps::ParsedValue;
450        let mut ctx = Context::new("".to_string(), HashMap::new());
451        ctx.store("scan", cmd_output(r#"{"issues": [1, 2]}"#, 0));
452        ctx.store_parsed(
453            "scan",
454            ParsedValue::Json(serde_json::json!({"issues": [1, 2]})),
455        );
456        // from() with JSON output allows deep dot-path access
457        let result = ctx
458            .render_template(r#"{{ from("scan").output.issues | length }}"#)
459            .unwrap();
460        assert_eq!(result, "2");
461    }
462
463    #[test]
464    fn from_traverses_parent_scope() {
465        let mut parent = Context::new("".to_string(), HashMap::new());
466        parent.store("root-step", cmd_output("root-value", 0));
467        let child = Context::child(Arc::new(parent), None, 0);
468        // from() inside child scope can access parent scope steps
469        let result = child
470            .render_template(r#"{{ from("root-step").output }}"#)
471            .unwrap();
472        assert_eq!(result, "root-value");
473    }
474
475    #[test]
476    fn from_safe_accessor_returns_empty_when_step_missing() {
477        let ctx = Context::new("".to_string(), HashMap::new());
478        // from("nonexistent").output? should return "" not fail
479        let result = ctx
480            .render_template(r#"{{ from("nonexistent").output? }}"#)
481            .unwrap();
482        assert_eq!(result, "");
483    }
484}
485
486fn collect_steps_with_parsed(ctx: &Context, map: &mut HashMap<String, serde_json::Value>) {
487    if let Some(parent) = &ctx.parent {
488        collect_steps_with_parsed(parent, map);
489    }
490    for (name, output) in &ctx.steps {
491        let parsed = ctx.parsed_outputs.get(name);
492        map.insert(
493            name.clone(),
494            step_output_to_value_with_parsed(output, parsed),
495        );
496    }
497}
498
499fn step_output_to_value_with_parsed(
500    output: &StepOutput,
501    parsed: Option<&ParsedValue>,
502) -> serde_json::Value {
503    let mut val = serde_json::to_value(output).unwrap_or(serde_json::Value::Null);
504
505    if let serde_json::Value::Object(ref mut map) = val {
506        // Add "output" key for template access (typed output parsing)
507        let output_val = match parsed {
508            Some(ParsedValue::Json(j)) => j.clone(),
509            Some(ParsedValue::Lines(lines)) => serde_json::json!(lines),
510            Some(ParsedValue::Integer(n)) => serde_json::json!(n),
511            Some(ParsedValue::Boolean(b)) => serde_json::json!(b),
512            Some(ParsedValue::Text(t)) => serde_json::Value::String(t.clone()),
513            None => serde_json::Value::String(output.text().to_string()),
514        };
515        map.insert("output".to_string(), output_val);
516
517        // Story 2.3: ensure session_id is always a string (empty if not an agent output)
518        let sid = match map.get("session_id") {
519            Some(serde_json::Value::String(s)) => serde_json::Value::String(s.clone()),
520            _ => serde_json::Value::String(String::new()),
521        };
522        map.insert("session_id".to_string(), sid);
523    }
524
525    val
526}
527
528fn check_json_path(val: &serde_json::Value, path: &[&str]) -> bool {
529    if path.is_empty() {
530        return true;
531    }
532    match val {
533        serde_json::Value::Object(map) => {
534            if let Some(next) = map.get(path[0]) {
535                check_json_path(next, &path[1..])
536            } else {
537                false
538            }
539        }
540        _ => false,
541    }
542}