Simplify node failure policy resolution

lib/components/fabro-validate/src/rules/on_failure_valid.rs

@@ -9,6 +9,35 @@ pub(super) fn rule() -> Box<dyn LintRule> {
 
 struct Rule;
 
+fn invalid_value_diagnostic(
+    rule: &str,
+    node_id: Option<&str>,
+    value: &AttrValue,
+) -> Option<Diagnostic> {
+    let message = match value {
+        AttrValue::String(value) if value.parse::<OnFailure>().is_ok() => return None,
+        AttrValue::String(value) => match node_id {
+            Some(node_id) => format!("Node '{node_id}' has invalid on_failure value '{value}'"),
+            None => format!("Graph has invalid on_failure value '{value}'"),
+        },
+        _ => match node_id {
+            Some(node_id) => {
+                format!("Node '{node_id}' attribute 'on_failure' must be a string")
+            }
+            None => "Graph attribute 'on_failure' must be a string".to_string(),
+        },
+    };
+
+    Some(Diagnostic {
+        rule: rule.to_string(),
+        severity: Severity::Error,
+        message,
+        node_id: node_id.map(str::to_string),
+        fix: Some(format!("Use one of: {}", OnFailure::expected_values())),
+        ..Diagnostic::default()
+    })
+}
+
 impl LintRule for Rule {
     fn name(&self) -> &'static str {
         "on_failure_valid"
@@ -17,39 +46,22 @@ impl LintRule for Rule {
     fn apply(&self, graph: &Graph) -> Vec<Diagnostic> {
         let mut diagnostics = Vec::new();
 
-        let invalid_value = |subject: &str, value: &AttrValue| -> Option<Diagnostic> {
-            let message = match value {
-                AttrValue::String(value) if value.parse::<OnFailure>().is_ok() => return None,
-                AttrValue::String(value) => {
-                    format!("{subject} has invalid on_failure value '{value}'")
-                }
-                _ => format!("{subject} attribute 'on_failure' must be a string"),
-            };
-            Some(Diagnostic {
-                rule: self.name().to_string(),
-                severity: Severity::Error,
-                message,
-                fix: Some(format!("Use one of: {}", OnFailure::expected_values())),
-                ..Diagnostic::default()
-            })
-        };
-
         if let Some(value) = graph.attrs.get("on_failure") {
-            diagnostics.extend(invalid_value("Graph", value));
+            diagnostics.extend(invalid_value_diagnostic(self.name(), None, value));
         }
 
-        let mut node_ids: Vec<&String> = graph.nodes.keys().collect();
-        node_ids.sort();
-        for node_id in node_ids {
-            let node = &graph.nodes[node_id];
-            if let Some(value) = node.attrs.get("on_failure") {
-                if let Some(diagnostic) = invalid_value(&format!("Node '{node_id}'"), value) {
-                    diagnostics.push(Diagnostic {
-                        node_id: Some(node_id.clone()),
-                        ..diagnostic
-                    });
-                }
-            }
+        let mut node_values: Vec<_> = graph
+            .nodes
+            .iter()
+            .filter_map(|(node_id, node)| {
+                node.attrs
+                    .get("on_failure")
+                    .map(|value| (node_id.as_str(), value))
+            })
+            .collect();
+        node_values.sort_unstable_by_key(|(node_id, _)| *node_id);
+        for (node_id, value) in node_values {
+            diagnostics.extend(invalid_value_diagnostic(self.name(), Some(node_id), value));
         }
 
         for edge in &graph.edges {

lib/components/fabro-workflow/src/graph.rs

@@ -131,7 +131,7 @@ impl Graph for WorkflowGraph {
         routing::get_retry_target(failed_node_id, self.inner())
     }
 
-    fn resolve_on_failure(&self, node_id: &str) -> ResolvedOnFailure {
-        self.inner().resolve_on_failure(node_id)
+    fn resolve_on_failure(&self, node: &Self::Node) -> ResolvedOnFailure {
+        self.inner().resolve_on_failure(node.inner())
     }
 }

lib/components/fabro-workflow/src/graph/routing.rs

@@ -75,7 +75,7 @@ pub(crate) fn select_edge<'a>(
         }
     }
 
-    if outcome.status.is_failure() && graph.resolve_on_failure(&node.id).policy == OnFailure::Exit {
+    if outcome.status.is_failure() && graph.resolve_on_failure(node).policy() == OnFailure::Exit {
         return None;
     }
 

lib/foundation/fabro-core/src/executor.rs

@@ -283,15 +283,15 @@ impl<G: Graph + 'static> Executor<G> {
                 NextStep::End => {
                     let mut outcome = last_outcome.clone();
                     if outcome.status.is_failure() {
-                        let resolved = graph.resolve_on_failure(node.id());
-                        let message = match resolved.policy {
+                        let resolved = graph.resolve_on_failure(&node);
+                        let message = match resolved.policy() {
                             OnFailure::Route => {
                                 format!("stage {} failed with no outgoing fail edge", node.id())
                             }
                             OnFailure::Exit => format!(
                                 "stage {} failed and {} on_failure=exit stopped routing",
                                 node.id(),
-                                resolved.scope
+                                resolved.scope()
                             ),
                         };
                         outcome = Outcome::fail(&message);
@@ -2266,7 +2266,7 @@ mod tests {
     async fn executor_node_exit_policy_overrides_graph_route_and_names_node_scope() {
         let graph = TestGraph::new(
             vec![
-                TestNode::new("work"),
+                TestNode::new("work").with_on_failure(OnFailure::Exit),
                 TestNode::new("downstream"),
                 TestNode::terminal("end"),
             ],
@@ -2275,8 +2275,7 @@ mod tests {
                 TestEdge::new("downstream", "end"),
             ],
             "work",
-        )
-        .with_node_on_failure("work", OnFailure::Exit);
+        );
         let state = ExecutionState::new(&graph).unwrap();
         let executor = ExecutorBuilder::new(
             Arc::new(AlwaysFailHandler::new("boom")) as Arc<dyn NodeHandler<TestGraph>>
@@ -2302,7 +2301,7 @@ mod tests {
     async fn executor_node_route_policy_overrides_graph_exit() {
         let graph = TestGraph::new(
             vec![
-                TestNode::new("work"),
+                TestNode::new("work").with_on_failure(OnFailure::Route),
                 TestNode::new("downstream"),
                 TestNode::terminal("end"),
             ],
@@ -2312,8 +2311,7 @@ mod tests {
             ],
             "work",
         )
-        .with_on_failure(OnFailure::Exit)
-        .with_node_on_failure("work", OnFailure::Route);
+        .with_on_failure(OnFailure::Exit);
         let state = ExecutionState::new(&graph).unwrap();
         let handler = DispatchHandler::new(Arc::new(AlwaysSucceedHandler))
             .with_handler("work", Arc::new(AlwaysFailHandler::new("boom")));

lib/foundation/fabro-core/src/graph.rs

@@ -44,5 +44,5 @@ pub trait Graph: Send + Sync {
     fn get_retry_target(&self, failed_node_id: &str) -> Option<String>;
     /// Effective failure routing policy for a node: node-level `on_failure`
     /// overrides the graph level, and an absent node attribute inherits it.
-    fn resolve_on_failure(&self, node_id: &str) -> ResolvedOnFailure;
+    fn resolve_on_failure(&self, node: &Self::Node) -> ResolvedOnFailure;
 }

lib/foundation/fabro-core/src/test_fixtures.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicU32, Ordering};
 
 use async_trait::async_trait;
-use fabro_types::{OnFailure, OnFailureScope, ResolvedOnFailure};
+use fabro_types::{OnFailure, ResolvedOnFailure};
 
 use crate::context::Context;
 use crate::error::{Error, HandlerErrorDetail, Result};
@@ -20,6 +20,7 @@ pub struct TestNode {
     pub terminal:   bool,
     pub max_visits: Option<usize>,
     pub goal_gate:  Option<(String, StageOutcome)>,
+    pub on_failure: Option<OnFailure>,
 }
 
 impl TestNode {
@@ -29,6 +30,7 @@ impl TestNode {
             terminal:   false,
             max_visits: None,
             goal_gate:  None,
+            on_failure: None,
         }
     }
 
@@ -38,6 +40,7 @@ impl TestNode {
             terminal:   true,
             max_visits: None,
             goal_gate:  None,
+            on_failure: None,
         }
     }
 
@@ -52,6 +55,12 @@ impl TestNode {
         self.goal_gate = Some((node_id.to_string(), required_status));
         self
     }
+
+    #[must_use]
+    pub fn with_on_failure(mut self, on_failure: OnFailure) -> Self {
+        self.on_failure = Some(on_failure);
+        self
+    }
 }
 
 impl NodeSpec for TestNode {
@@ -119,12 +128,11 @@ impl EdgeSpec for TestEdge {
 
 #[derive(Debug, Clone)]
 pub struct TestGraph {
-    pub nodes:           Vec<TestNode>,
-    pub edges:           Vec<TestEdge>,
-    pub start_node_id:   String,
-    pub retry_targets:   HashMap<String, String>,
-    pub on_failure:      OnFailure,
-    pub node_on_failure: HashMap<String, OnFailure>,
+    pub nodes:         Vec<TestNode>,
+    pub edges:         Vec<TestEdge>,
+    pub start_node_id: String,
+    pub retry_targets: HashMap<String, String>,
+    pub on_failure:    OnFailure,
 }
 
 impl TestGraph {
@@ -135,7 +143,6 @@ impl TestGraph {
             start_node_id: start.to_string(),
             retry_targets: HashMap::new(),
             on_failure: OnFailure::Route,
-            node_on_failure: HashMap::new(),
         }
     }
 
@@ -150,12 +157,6 @@ impl TestGraph {
         self.on_failure = on_failure;
         self
     }
-
-    #[must_use]
-    pub fn with_node_on_failure(mut self, node_id: &str, on_failure: OnFailure) -> Self {
-        self.node_on_failure.insert(node_id.to_string(), on_failure);
-        self
-    }
 }
 
 impl Graph for TestGraph {
@@ -225,8 +226,7 @@ impl Graph for TestGraph {
             }
         }
 
-        if outcome.status.is_failure()
-            && self.resolve_on_failure(node.id()).policy == OnFailure::Exit
+        if outcome.status.is_failure() && self.resolve_on_failure(node).policy() == OnFailure::Exit
         {
             return None;
         }
@@ -267,16 +267,10 @@ impl Graph for TestGraph {
         self.retry_targets.get(failed_node_id).cloned()
     }
 
-    fn resolve_on_failure(&self, node_id: &str) -> ResolvedOnFailure {
-        match self.node_on_failure.get(node_id) {
-            Some(policy) => ResolvedOnFailure {
-                policy: *policy,
-                scope:  OnFailureScope::Node,
-            },
-            None => ResolvedOnFailure {
-                policy: self.on_failure,
-                scope:  OnFailureScope::Graph,
-            },
+    fn resolve_on_failure(&self, node: &Self::Node) -> ResolvedOnFailure {
+        match node.on_failure {
+            Some(policy) => ResolvedOnFailure::node(policy),
+            None => ResolvedOnFailure::graph(self.on_failure),
         }
     }
 }

lib/foundation/fabro-types/src/graph.rs

@@ -36,20 +36,40 @@ impl OnFailure {
     }
 }
 
-/// The scope whose `on_failure` attribute supplied a resolved policy.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::Display, strum::IntoStaticStr)]
-#[strum(serialize_all = "snake_case")]
-pub enum OnFailureScope {
-    Node,
-    Graph,
-}
-
 /// A failure routing policy together with the scope that supplied it, so
 /// failure messages can name the attribute that stopped routing.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct ResolvedOnFailure {
-    pub policy: OnFailure,
-    pub scope:  OnFailureScope,
+    policy: OnFailure,
+    scope:  AttributeScope,
+}
+
+impl ResolvedOnFailure {
+    #[must_use]
+    pub const fn node(policy: OnFailure) -> Self {
+        Self {
+            policy,
+            scope: AttributeScope::Node,
+        }
+    }
+
+    #[must_use]
+    pub const fn graph(policy: OnFailure) -> Self {
+        Self {
+            policy,
+            scope: AttributeScope::Graph,
+        }
+    }
+
+    #[must_use]
+    pub const fn policy(self) -> OnFailure {
+        self.policy
+    }
+
+    #[must_use]
+    pub const fn scope(self) -> AttributeScope {
+        self.scope
+    }
 }
 
 /// Typed attribute values for nodes, edges, and graph-level attributes.
@@ -610,17 +630,10 @@ impl Graph {
     /// attribute overrides the graph level; an absent (or invalid, hence
     /// validation-rejected) node attribute inherits the graph policy.
     #[must_use]
-    pub fn resolve_on_failure(&self, node_id: &str) -> ResolvedOnFailure {
-        let node_policy = self.nodes.get(node_id).and_then(Node::on_failure);
-        match node_policy {
-            Some(policy) => ResolvedOnFailure {
-                policy,
-                scope: OnFailureScope::Node,
-            },
-            None => ResolvedOnFailure {
-                policy: self.on_failure(),
-                scope:  OnFailureScope::Graph,
-            },
+    pub fn resolve_on_failure(&self, node: &Node) -> ResolvedOnFailure {
+        match node.on_failure() {
+            Some(policy) => ResolvedOnFailure::node(policy),
+            None => ResolvedOnFailure::graph(self.on_failure()),
         }
     }
 
@@ -677,7 +690,8 @@ impl Graph {
 }
 
 /// Where an attribute appears in a workflow graph.
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+#[derive(Clone, Copy, Debug, Eq, PartialEq, strum::Display, strum::IntoStaticStr)]
+#[strum(serialize_all = "snake_case")]
 pub enum AttributeScope {
     Graph,
     Node,
@@ -823,16 +837,16 @@ mod tests {
         graph.nodes.insert("route".to_string(), route);
 
         // Node attribute wins over the graph policy.
-        assert_eq!(graph.resolve_on_failure("route"), ResolvedOnFailure {
-            policy: OnFailure::Route,
-            scope:  OnFailureScope::Node,
-        });
-        // Absent, invalid, and unknown nodes inherit the graph policy.
-        for node_id in ["bare", "invalid", "missing"] {
-            assert_eq!(graph.resolve_on_failure(node_id), ResolvedOnFailure {
-                policy: OnFailure::Exit,
-                scope:  OnFailureScope::Graph,
-            });
+        assert_eq!(
+            graph.resolve_on_failure(&graph.nodes["route"]),
+            ResolvedOnFailure::node(OnFailure::Route)
+        );
+        // Absent and invalid node attributes inherit the graph policy.
+        for node_id in ["bare", "invalid"] {
+            assert_eq!(
+                graph.resolve_on_failure(&graph.nodes[node_id]),
+                ResolvedOnFailure::graph(OnFailure::Exit)
+            );
         }
     }
 

lib/foundation/fabro-types/src/lib.rs

@@ -78,7 +78,7 @@ pub use event_envelope::EventEnvelope;
 pub use fabro_model::ReasoningEffort;
 pub use failure_signature::FailureSignature;
 pub use graph::{
-    AttrValue, ContextKeyAttr, Edge, Graph, KNOWN_HANDLER_TYPES, Node, OnFailure, OnFailureScope,
+    AttrValue, AttributeScope, ContextKeyAttr, Edge, Graph, KNOWN_HANDLER_TYPES, Node, OnFailure,
     ResolvedOnFailure, is_known_handler_type, is_llm_handler_type, shape_to_handler_type,
 };
 pub use input_scalar::{JsonScalarToTomlError, json_scalar_to_toml_value};