zen_engine/decision_graph/
tracer.rs

1use crate::nodes::NodeResult;
2use ahash::HashMap;
3use serde::{Deserialize, Serialize};
4use std::sync::Arc;
5use zen_expression::variable::ToVariable;
6use zen_types::decision::{DecisionNode, DecisionNodeKind};
7use zen_types::variable::Variable;
8
9pub(crate) struct NodeTracer(Option<HashMap<Arc<str>, DecisionGraphTrace>>);
10
11impl NodeTracer {
12    pub fn new(enabled: bool) -> Self {
13        Self(enabled.then(|| HashMap::default()))
14    }
15
16    pub fn record_execution(
17        &mut self,
18        node: &DecisionNode,
19        input_trace: Variable,
20        result: &NodeResult,
21        duration: std::time::Duration,
22    ) {
23        let Some(traces) = &mut self.0 else {
24            return;
25        };
26
27        if matches!(node.kind, DecisionNodeKind::SwitchNode { .. }) {
28            return;
29        }
30
31        let input = match &node.kind {
32            DecisionNodeKind::InputNode { .. } => Variable::Null,
33            _ => input_trace,
34        };
35
36        let mut trace = DecisionGraphTrace {
37            id: node.id.clone(),
38            name: node.name.clone(),
39            input,
40            order: traces.len() as u32,
41            output: Variable::Null,
42            trace_data: None,
43            performance: Some(Arc::from(format!("{:.1?}", duration))),
44        };
45
46        match &result {
47            Ok(ok) => {
48                trace.trace_data = ok.trace_data.clone();
49                if !matches!(node.kind, DecisionNodeKind::OutputNode { .. }) {
50                    trace.output = ok.output.clone();
51                }
52            }
53            Err(err) => {
54                trace.trace_data = err.trace.clone();
55            }
56        };
57
58        traces.insert(node.id.clone(), trace);
59    }
60
61    pub fn trace_callback(&mut self) -> Option<impl FnMut(DecisionGraphTrace) + '_> {
62        let Some(traces) = &mut self.0 else {
63            return None;
64        };
65
66        Some(|mut trace: DecisionGraphTrace| {
67            trace.order = traces.len() as u32;
68            traces.insert(trace.id.clone(), trace);
69        })
70    }
71
72    pub fn into_traces(self) -> Option<HashMap<Arc<str>, DecisionGraphTrace>> {
73        self.0
74    }
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize, ToVariable)]
78#[serde(rename_all = "camelCase")]
79pub struct DecisionGraphTrace {
80    pub input: Variable,
81    pub output: Variable,
82    pub name: Arc<str>,
83    pub id: Arc<str>,
84    pub performance: Option<Arc<str>>,
85    pub trace_data: Option<Variable>,
86    pub order: u32,
87}
zen_engine/decision_graph/tracer.rs

zen_engine/decision_graph/
tracer.rs
