From f3f15967e29983fcbe21d85292ad7fbf8535e3b5 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Tue, 24 Mar 2026 16:18:49 -0400 Subject: [PATCH] Add fabro-core after_record lifecycle hook --- lib/crates/fabro-core/src/executor.rs | 79 +++++++++++++++++++ lib/crates/fabro-core/src/lifecycle.rs | 50 ++++++++++++ .../src/core_adapter/lifecycle/mod.rs | 52 ++++++++++++ lib/crates/fabro-workflows/src/engine.rs | 5 +- lib/crates/fabro-workflows/src/lib.rs | 7 +- .../fabro-workflows/tests/integration.rs | 2 +- 6 files changed, 186 insertions(+), 9 deletions(-) diff --git a/lib/crates/fabro-core/src/executor.rs b/lib/crates/fabro-core/src/executor.rs index a63bc03fa..835ff8609 100644 --- a/lib/crates/fabro-core/src/executor.rs +++ b/lib/crates/fabro-core/src/executor.rs @@ -195,6 +195,9 @@ impl Executor { }; state.record(node.id(), &node_result); + self.lifecycle + .after_record(&node, &node_result, &state) + .await?; // Determine next step let last_outcome = state.node_outcomes.get(node.id()).unwrap(); @@ -1539,6 +1542,82 @@ mod tests { ); } + #[tokio::test] + async fn executor_after_record_runs_after_record_and_before_edge_selection() { + use serde_json::json; + + struct ContextWriter; + #[async_trait] + impl NodeHandler for ContextWriter { + async fn execute( + &self, + node: &TestNode, + _context: &Context, + _g: &TestGraph, + ) -> Result { + let mut outcome = Outcome::success(); + if node.id() == "start" { + outcome + .context_updates + .insert("shared".into(), json!("hello")); + } + Ok(outcome) + } + } + + let log = Arc::new(Mutex::new(Vec::::new())); + struct RecordTracker(Arc>>); + #[async_trait] + impl RunLifecycle for RecordTracker { + async fn after_record( + &self, + node: &TestNode, + _result: &NodeResult, + state: &RunState, + ) -> Result<()> { + let shared = state.context.get_string("shared", "missing"); + let completed = state.completed_nodes.join(","); + self.0.lock().unwrap().push(format!( + "after_record:{}:{}:{}", + node.id(), + completed, + shared + )); + Ok(()) + } + + async fn on_edge_selected( + &self, + ctx: &EdgeContext<'_, TestGraph>, + state: &RunState, + ) -> Result { + let shared = state.context.get_string("shared", "missing"); + self.0 + .lock() + .unwrap() + .push(format!("on_edge_selected:{}:{}", ctx.from, shared)); + Ok(EdgeDecision::Continue) + } + } + + let g = linear_graph(&["start", "end"]); + let state = RunState::new(&g).unwrap(); + let executor = + ExecutorBuilder::new(Arc::new(ContextWriter) as Arc>) + .lifecycle(Box::new(RecordTracker(log.clone()))) + .build(); + + executor.run(&g, state).await.unwrap(); + + assert_eq!( + *log.lock().unwrap(), + vec![ + "after_record:start:start:hello".to_string(), + "on_edge_selected:start:hello".to_string(), + ] + ); + } + #[tokio::test] async fn executor_terminal_reached_receives_goal_gate_result() { let log = Arc::new(Mutex::new(Vec::<(String, bool)>::new())); diff --git a/lib/crates/fabro-core/src/lifecycle.rs b/lib/crates/fabro-core/src/lifecycle.rs index 8ee859e6e..67ba92296 100644 --- a/lib/crates/fabro-core/src/lifecycle.rs +++ b/lib/crates/fabro-core/src/lifecycle.rs @@ -91,6 +91,15 @@ pub trait RunLifecycle: Send + Sync { Ok(()) } + async fn after_record( + &self, + _node: &G::Node, + _result: &NodeResult, + _state: &RunState, + ) -> Result<()> { + Ok(()) + } + async fn on_edge_selected( &self, _ctx: &EdgeContext<'_, G>, @@ -203,6 +212,18 @@ impl RunLifecycle for CompositeLifecycle { Ok(()) } + async fn after_record( + &self, + node: &G::Node, + result: &NodeResult, + state: &RunState, + ) -> Result<()> { + for child in &self.children { + child.after_record(node, result, state).await?; + } + Ok(()) + } + async fn on_edge_selected( &self, ctx: &EdgeContext<'_, G>, @@ -360,6 +381,19 @@ mod tests { Ok(()) } + async fn after_record( + &self, + _node: &TestNode, + _result: &NodeResult, + _state: &RunState, + ) -> Result<()> { + self.log + .lock() + .unwrap() + .push(format!("{}:after_record", self.name)); + Ok(()) + } + async fn on_edge_selected( &self, _ctx: &EdgeContext<'_, TestGraph>, @@ -630,6 +664,22 @@ mod tests { assert_eq!(calls, vec!["a:after_node", "b:after_node"]); } + #[tokio::test] + async fn composite_after_record_calls_all() { + let log = Arc::new(Mutex::new(Vec::new())); + let lc = CompositeLifecycle::new(vec![ + Box::new(RecordingLifecycle::new("a", log.clone())), + Box::new(RecordingLifecycle::new("b", log.clone())), + ]); + let g = linear_graph(&["start", "end"]); + let state = RunState::new(&g).unwrap(); + let node = g.get_node("start").unwrap(); + let result = NodeResult::new(Outcome::success(), Duration::ZERO, 1, 1); + lc.after_record(&node, &result, &state).await.unwrap(); + let calls = log.lock().unwrap().clone(); + assert_eq!(calls, vec!["a:after_record", "b:after_record"]); + } + #[tokio::test] async fn composite_ordering_is_preserved() { let log = Arc::new(Mutex::new(Vec::new())); diff --git a/lib/crates/fabro-workflows/src/core_adapter/lifecycle/mod.rs b/lib/crates/fabro-workflows/src/core_adapter/lifecycle/mod.rs index 56c8b2ad4..5893ff98f 100644 --- a/lib/crates/fabro-workflows/src/core_adapter/lifecycle/mod.rs +++ b/lib/crates/fabro-workflows/src/core_adapter/lifecycle/mod.rs @@ -16,6 +16,7 @@ use std::time::Instant; use async_trait::async_trait; use fabro_core::error::Result as CoreResult; +use fabro_core::graph::NodeSpec; use fabro_core::lifecycle::{ AttemptContext, AttemptResultContext, EdgeContext, EdgeDecision, NodeDecision, RunLifecycle, }; @@ -40,6 +41,7 @@ use self::event::EventLifecycle; use self::fidelity::FidelityLifecycle; use self::git::{GitCheckpointResult, GitLifecycle}; use self::hook::HookLifecycle; +use crate::outcome::OutcomeExt; type WfRunState = RunState>; type WfNodeResult = NodeResult>; @@ -293,6 +295,56 @@ impl RunLifecycle for WorkflowLifecycle { Ok(()) } + async fn after_record( + &self, + node: &WorkflowNode, + result: &WfNodeResult, + state: &WfRunState, + ) -> CoreResult<()> { + let outcome = &result.outcome; + let retry_count = state.node_retries.get(node.id()).copied().unwrap_or(0); + let failure_class = crate::engine::classify_outcome(outcome); + let failure_signature = failure_class + .map(|category| { + let signature_hint = outcome + .failure + .as_ref() + .and_then(|f| f.signature.as_deref()); + crate::error::FailureSignature::new( + node.id(), + category, + signature_hint, + outcome.failure_reason(), + ) + .to_string() + }) + .unwrap_or_default(); + + state.context.set( + context::keys::retry_count_key(node.id()), + serde_json::json!(retry_count), + ); + state.context.set( + context::keys::OUTCOME, + serde_json::json!(outcome.status.to_string()), + ); + state.context.set( + context::keys::FAILURE_CLASS, + serde_json::json!(failure_class.map_or(String::new(), |fc| fc.to_string())), + ); + state.context.set( + context::keys::FAILURE_SIGNATURE, + serde_json::json!(failure_signature), + ); + if let Some(ref preferred_label) = outcome.preferred_label { + state.context.set( + context::keys::PREFERRED_LABEL, + serde_json::json!(preferred_label), + ); + } + Ok(()) + } + async fn on_edge_selected( &self, ctx: &EdgeContext<'_, WorkflowGraph>, diff --git a/lib/crates/fabro-workflows/src/engine.rs b/lib/crates/fabro-workflows/src/engine.rs index bbccb336b..2c86cde03 100644 --- a/lib/crates/fabro-workflows/src/engine.rs +++ b/lib/crates/fabro-workflows/src/engine.rs @@ -2109,10 +2109,11 @@ impl WorkflowRunEngine { .await? }; // Gap #5: Track retry count per node - node_retries.insert(node.id.clone(), attempts_used); + let retry_count = attempts_used.saturating_sub(1); + node_retries.insert(node.id.clone(), retry_count); context.set( context::keys::retry_count_key(&node.id), - serde_json::json!(attempts_used), + serde_json::json!(retry_count), ); // Gap #1: Auto status -- when auto_status=true and outcome is non-success, diff --git a/lib/crates/fabro-workflows/src/lib.rs b/lib/crates/fabro-workflows/src/lib.rs index 4a59c9d2e..5d6aeb082 100644 --- a/lib/crates/fabro-workflows/src/lib.rs +++ b/lib/crates/fabro-workflows/src/lib.rs @@ -38,12 +38,7 @@ pub fn build_completed_stages( for node_id in &cp.completed_nodes { let outcome = cp.node_outcomes.get(node_id); - let retries = cp - .node_retries - .get(node_id) - .copied() - .unwrap_or(1) - .saturating_sub(1); + let retries = cp.node_retries.get(node_id).copied().unwrap_or(0); let status = outcome .map(|o| o.status.to_string()) diff --git a/lib/crates/fabro-workflows/tests/integration.rs b/lib/crates/fabro-workflows/tests/integration.rs index fda42d981..b6b7ebfc2 100644 --- a/lib/crates/fabro-workflows/tests/integration.rs +++ b/lib/crates/fabro-workflows/tests/integration.rs @@ -2804,7 +2804,7 @@ async fn scenario_node_retries_on_retry_status() { .node_retries .get("flaky") .expect("flaky should have retries"); - assert_eq!(*retry_count, 2, "should have been called 2x"); + assert_eq!(*retry_count, 1, "should have retried once"); } #[tokio::test]