Add fabro-core after_record lifecycle hook

This commit is contained in:
Bryan Helmkamp 2026-03-24 16:18:49 -04:00
parent b536564d5e
commit f3f15967e2
6 changed files with 186 additions and 9 deletions

View file

@ -195,6 +195,9 @@ impl<G: Graph + 'static> Executor<G> {
};
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<TestGraph> for ContextWriter {
async fn execute(
&self,
node: &TestNode,
_context: &Context,
_g: &TestGraph,
) -> Result<Outcome> {
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::<String>::new()));
struct RecordTracker(Arc<Mutex<Vec<String>>>);
#[async_trait]
impl RunLifecycle<TestGraph> 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<EdgeDecision> {
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<dyn NodeHandler<TestGraph>>)
.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()));

View file

@ -91,6 +91,15 @@ pub trait RunLifecycle<G: Graph>: Send + Sync {
Ok(())
}
async fn after_record(
&self,
_node: &G::Node,
_result: &NodeResult<G::Meta>,
_state: &RunState<G::Meta>,
) -> Result<()> {
Ok(())
}
async fn on_edge_selected(
&self,
_ctx: &EdgeContext<'_, G>,
@ -203,6 +212,18 @@ impl<G: Graph + 'static> RunLifecycle<G> for CompositeLifecycle<G> {
Ok(())
}
async fn after_record(
&self,
node: &G::Node,
result: &NodeResult<G::Meta>,
state: &RunState<G::Meta>,
) -> 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()));

View file

@ -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<Option<StageUsage>>;
type WfNodeResult = NodeResult<Option<StageUsage>>;
@ -293,6 +295,56 @@ impl RunLifecycle<WorkflowGraph> 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>,

View file

@ -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,

View file

@ -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())

View file

@ -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]