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#[derive(Debug, Clone)]
10pub struct ChatMessage {
11 pub role: String,
12 pub content: String,
13}
14
15#[derive(Debug, Clone, Default)]
17pub struct ChatHistory {
18 pub messages: Vec<ChatMessage>,
19}
20
21type ChatSessionStore = Arc<Mutex<HashMap<String, ChatHistory>>>;
23
24pub 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 chat_sessions: ChatSessionStore,
35 pub stack_info: Option<StackInfo>,
37 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 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 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 pub fn insert_var(&mut self, name: impl Into<String>, value: serde_json::Value) {
79 self.variables.insert(name.into(), value);
80 }
81
82 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 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 pub fn store_parsed(&mut self, name: &str, parsed: ParsedValue) {
98 self.parsed_outputs.insert(name.to_string(), parsed);
99 }
100
101 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 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 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 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 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 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 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 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 pub fn to_tera_context(&self) -> tera::Context {
204 let mut ctx = tera::Context::new();
205
206 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 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 for (name, val) in &steps_map {
227 ctx.insert(name.as_str(), val);
228 }
229 ctx.insert("steps", &steps_map);
230
231 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 pub fn render_template(&self, template: &str) -> Result<String, crate::error::StepError> {
244 let pre = crate::engine::template::preprocess_template(template, self)?;
246
247 let mut tera_ctx = self.to_tera_context();
249
250 for (k, v) in &pre.injected {
252 tera_ctx.insert(k.as_str(), v);
253 }
254
255 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 assert!(child.get_step("a").is_some());
377 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 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 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 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 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 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 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 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 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 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 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}