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 #[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 pub fn store_parsed(&mut self, name: &str, parsed: ParsedValue) {
99 self.parsed_outputs.insert(name.to_string(), parsed);
100 }
101
102 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 #[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 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 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 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 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 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 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 pub fn to_tera_context(&self) -> tera::Context {
206 let mut ctx = tera::Context::new();
207
208 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 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 for (name, val) in &steps_map {
229 ctx.insert(name.as_str(), val);
230 }
231 ctx.insert("steps", &steps_map);
232
233 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 pub fn render_template(&self, template: &str) -> Result<String, crate::error::StepError> {
246 let pre = crate::engine::template::preprocess_template(template, self)?;
248
249 let mut tera_ctx = self.to_tera_context();
251
252 for (k, v) in &pre.injected {
254 tera_ctx.insert(k.as_str(), v);
255 }
256
257 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 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 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 assert!(child.get_step("a").is_some());
437 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 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 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 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 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 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 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 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 let result = ctx
540 .render_template(r#"{{ from("nonexistent").output? }}"#)
541 .unwrap();
542 assert_eq!(result, "");
543 }
544}