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