1use crate::manifest::Edge;
2use crate::result::BrickResult;
3
4#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct RoutedEdge {
7 pub edge_id: String,
8 pub target_node: String,
9}
10
11pub fn route(
18 outbound_edges: &[&Edge],
19 result: &BrickResult,
20 output_confidence: Option<f64>,
21) -> Vec<RoutedEdge> {
22 match result {
23 BrickResult::Failure { error } | BrickResult::LowConfidence { error, .. } => {
25 route_on_error(outbound_edges, &error.error_class)
26 }
27 BrickResult::Success { .. } => {
28 if let Some(c) = output_confidence {
31 if !c.is_finite() || !(0.0..=1.0).contains(&c) {
32 return route_on_error(outbound_edges, "INVALID_INPUT");
33 }
34 }
35 route_on_success(outbound_edges, output_confidence)
36 }
37 }
38}
39
40fn route_on_error(edges: &[&Edge], error_class: &str) -> Vec<RoutedEdge> {
43 let mut candidates: Vec<&Edge> = edges
44 .iter()
45 .filter(|e| {
46 e.on_error
49 .as_ref()
50 .is_some_and(|m| m.contains_key(error_class))
51 })
52 .copied()
53 .collect();
54
55 candidates.sort_by(|a, b| {
57 let pa = a.priority.unwrap_or(0);
58 let pb = b.priority.unwrap_or(0);
59 pb.cmp(&pa).then_with(|| a.edge_id.cmp(&b.edge_id))
60 });
61
62 candidates
64 .first()
65 .map(|e| {
66 vec![RoutedEdge {
67 edge_id: e.edge_id.clone(),
68 target_node: e.target_node.clone(),
69 }]
70 })
71 .unwrap_or_default()
72}
73
74fn route_on_success(edges: &[&Edge], output_confidence: Option<f64>) -> Vec<RoutedEdge> {
77 let mut candidates: Vec<&Edge> = edges
78 .iter()
79 .filter(|e| e.on_success.is_some())
80 .filter(|e| {
81 let on_success = e.on_success.as_ref().unwrap();
82 match (on_success.threshold, output_confidence) {
83 (Some(threshold), Some(confidence)) => confidence >= threshold,
85 (Some(_), None) => true,
87 (None, _) => true,
89 }
90 })
91 .copied()
92 .collect();
93
94 candidates.sort_by(|a, b| {
96 let mut wa = a.on_success.as_ref().and_then(|s| s.weight).unwrap_or(0.0);
97 let mut wb = b.on_success.as_ref().and_then(|s| s.weight).unwrap_or(0.0);
98
99 if !wa.is_finite() {
101 wa = 0.0;
102 }
103 if !wb.is_finite() {
104 wb = 0.0;
105 }
106
107 wb.total_cmp(&wa).then_with(|| a.edge_id.cmp(&b.edge_id))
108 });
109
110 candidates
111 .iter()
112 .map(|e| RoutedEdge {
113 edge_id: e.edge_id.clone(),
114 target_node: e.target_node.clone(),
115 })
116 .collect()
117}
118
119#[cfg(test)]
120mod tests {
121 use super::*;
122 use crate::manifest::{Edge, OnSuccess};
123 use crate::result::{BrickResult, ErrorObject};
124 use std::collections::HashMap;
125
126 fn failure(error_class: &str) -> BrickResult {
127 BrickResult::Failure {
128 error: ErrorObject {
129 error_class: error_class.to_string(),
130 message: "test".to_string(),
131 retry_advice: None,
132 severity: None,
133 },
134 }
135 }
136
137 fn low_confidence() -> BrickResult {
138 BrickResult::LowConfidence {
139 output: crate::result::CborValue::Null,
140 error: ErrorObject {
141 error_class: "LOW_CONFIDENCE".to_string(),
142 message: "test".to_string(),
143 retry_advice: None,
144 severity: None,
145 },
146 }
147 }
148
149 fn success() -> BrickResult {
150 BrickResult::Success {
151 output: crate::result::CborValue::Null,
152 }
153 }
154
155 fn edge(id: &str, target: &str) -> Edge {
156 Edge {
157 edge_id: id.to_string(),
158 source_node: "src".to_string(),
159 target_node: target.to_string(),
160 mapping: vec![],
161 on_success: None,
162 on_error: None,
163 priority: None,
164 }
165 }
166
167 fn ids(routed: &[RoutedEdge]) -> Vec<&str> {
168 routed.iter().map(|r| r.edge_id.as_str()).collect()
169 }
170
171 #[test]
174 fn v1_on_error_single_match() {
175 let mut e1 = edge("e1", "recovery_node");
176 e1.on_error = Some(HashMap::from([(
177 "COMPUTATION_ERROR".into(),
178 vec!["recovery_node".into()],
179 )]));
180 let mut e2 = edge("e2", "fallback_node");
181 e2.on_error = Some(HashMap::from([(
182 "INVALID_INPUT".into(),
183 vec!["fallback_node".into()],
184 )]));
185
186 let result = failure("COMPUTATION_ERROR");
187 let edges: Vec<&Edge> = vec![&e1, &e2];
188 assert_eq!(ids(&route(&edges, &result, None)), vec!["e1"]);
189 }
190
191 #[test]
192 fn v2_on_error_no_match_terminal() {
193 let mut e1 = edge("e1", "fallback_node");
194 e1.on_error = Some(HashMap::from([(
195 "INVALID_INPUT".into(),
196 vec!["fallback_node".into()],
197 )]));
198
199 let result = failure("COMPUTATION_ERROR");
200 let edges: Vec<&Edge> = vec![&e1];
201 assert!(route(&edges, &result, None).is_empty());
202 }
203
204 #[test]
205 fn v3_on_error_priority_ordering() {
206 let mut e1 = edge("e1", "node_a");
207 e1.priority = Some(5);
208 e1.on_error = Some(HashMap::from([(
209 "COMPUTATION_ERROR".into(),
210 vec!["node_a".into()],
211 )]));
212 let mut e2 = edge("e2", "node_b");
213 e2.priority = Some(10);
214 e2.on_error = Some(HashMap::from([(
215 "COMPUTATION_ERROR".into(),
216 vec!["node_b".into()],
217 )]));
218
219 let result = failure("COMPUTATION_ERROR");
220 let edges: Vec<&Edge> = vec![&e1, &e2];
221 assert_eq!(ids(&route(&edges, &result, None)), vec!["e2"]);
222 }
223
224 #[test]
225 fn v4_on_error_tiebreak_edge_id() {
226 let mut eb = edge("edge_b", "node_a");
227 eb.priority = Some(10);
228 eb.on_error = Some(HashMap::from([(
229 "COMPUTATION_ERROR".into(),
230 vec!["node_a".into()],
231 )]));
232 let mut ea = edge("edge_a", "node_b");
233 ea.priority = Some(10);
234 ea.on_error = Some(HashMap::from([(
235 "COMPUTATION_ERROR".into(),
236 vec!["node_b".into()],
237 )]));
238
239 let result = failure("COMPUTATION_ERROR");
240 let edges: Vec<&Edge> = vec![&eb, &ea];
241 assert_eq!(ids(&route(&edges, &result, None)), vec!["edge_a"]);
242 }
243
244 #[test]
245 fn v5_low_confidence_routes_via_on_error() {
246 let mut e1 = edge("e1", "llm_node");
247 e1.on_error = Some(HashMap::from([(
248 "LOW_CONFIDENCE".into(),
249 vec!["llm_node".into()],
250 )]));
251 let mut e2 = edge("e2", "next_node");
252 e2.on_success = Some(OnSuccess {
253 weight: Some(1.0),
254 threshold: None,
255 });
256
257 let result = low_confidence();
258 let edges: Vec<&Edge> = vec![&e1, &e2];
259 assert_eq!(ids(&route(&edges, &result, None)), vec!["e1"]);
260 }
261
262 #[test]
263 fn v6_on_success_above_threshold() {
264 let mut e1 = edge("e1", "next_node");
265 e1.on_success = Some(OnSuccess {
266 threshold: Some(0.5),
267 weight: Some(1.0),
268 });
269
270 let result = success();
271 let edges: Vec<&Edge> = vec![&e1];
272 assert_eq!(ids(&route(&edges, &result, Some(0.8))), vec!["e1"]);
273 }
274
275 #[test]
276 fn v7_on_success_below_threshold_terminal() {
277 let mut e1 = edge("e1", "next_node");
278 e1.on_success = Some(OnSuccess {
279 threshold: Some(0.9),
280 weight: Some(1.0),
281 });
282
283 let result = success();
284 let edges: Vec<&Edge> = vec![&e1];
285 assert!(route(&edges, &result, Some(0.7)).is_empty());
286 }
287
288 #[test]
289 fn v8_on_success_no_confidence_eligible() {
290 let mut e1 = edge("e1", "next_node");
291 e1.on_success = Some(OnSuccess {
292 threshold: Some(0.9),
293 weight: Some(1.0),
294 });
295
296 let result = success();
297 let edges: Vec<&Edge> = vec![&e1];
298 assert_eq!(ids(&route(&edges, &result, None)), vec!["e1"]);
299 }
300
301 #[test]
302 fn v9_on_success_fanout_weight_ordering() {
303 let mut e1 = edge("e1", "node_a");
304 e1.on_success = Some(OnSuccess {
305 weight: Some(0.5),
306 threshold: None,
307 });
308 let mut e2 = edge("e2", "node_b");
309 e2.on_success = Some(OnSuccess {
310 weight: Some(0.9),
311 threshold: None,
312 });
313 let mut e3 = edge("e3", "node_c");
314 e3.on_success = Some(OnSuccess {
315 weight: Some(0.9),
316 threshold: None,
317 });
318
319 let result = success();
320 let edges: Vec<&Edge> = vec![&e1, &e2, &e3];
321 assert_eq!(ids(&route(&edges, &result, None)), vec!["e2", "e3", "e1"]);
322 }
323
324 #[test]
325 fn v10_on_success_weight_tiebreak_edge_id() {
326 let mut ezz = edge("zz_edge", "node_a");
327 ezz.on_success = Some(OnSuccess {
328 weight: Some(1.0),
329 threshold: None,
330 });
331 let mut eaa = edge("aa_edge", "node_b");
332 eaa.on_success = Some(OnSuccess {
333 weight: Some(1.0),
334 threshold: None,
335 });
336
337 let result = success();
338 let edges: Vec<&Edge> = vec![&ezz, &eaa];
339 assert_eq!(
340 ids(&route(&edges, &result, None)),
341 vec!["aa_edge", "zz_edge"]
342 );
343 }
344
345 #[test]
346 fn v11_invalid_confidence_routes_as_invalid_input_via_on_error() {
347 let mut on_err = edge("e_err", "err_node");
348 on_err.on_error = Some(HashMap::from([(
349 "INVALID_INPUT".into(),
350 vec!["err_node".into()],
351 )]));
352
353 let mut on_ok = edge("e_ok", "next");
354 on_ok.on_success = Some(OnSuccess {
355 threshold: Some(0.5),
356 weight: Some(1.0),
357 });
358
359 let result = success();
360 let edges: Vec<&Edge> = vec![&on_ok, &on_err];
361
362 assert_eq!(ids(&route(&edges, &result, Some(1.2))), vec!["e_err"]);
364 }
365
366 #[test]
369 fn v12_nan_weight_sorted_after_finite() {
370 let mut ea = edge("ea", "node_a");
371 ea.on_success = Some(OnSuccess {
372 weight: Some(f64::NAN),
373 threshold: None,
374 });
375 let mut eb = edge("eb", "node_b");
376 eb.on_success = Some(OnSuccess {
377 weight: Some(0.1),
378 threshold: None,
379 });
380
381 let result = success();
382 let edges: Vec<&Edge> = vec![&ea, &eb];
383 assert_eq!(ids(&route(&edges, &result, None)), vec!["eb", "ea"]);
385 }
386
387 #[test]
388 fn v13_both_nan_weights_tiebreak_by_edge_id() {
389 let mut eb = edge("eb", "node_b");
390 eb.on_success = Some(OnSuccess {
391 weight: Some(f64::NAN),
392 threshold: None,
393 });
394 let mut ea = edge("ea", "node_a");
395 ea.on_success = Some(OnSuccess {
396 weight: Some(f64::NAN),
397 threshold: None,
398 });
399
400 let result = success();
401 let edges: Vec<&Edge> = vec![&eb, &ea];
402 assert_eq!(ids(&route(&edges, &result, None)), vec!["ea", "eb"]);
404 }
405}