fabro/lib/components/fabro-workflow/src/lifecycle/event.rs

use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};
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,
};
use fabro_core::outcome::NodeResult;
use fabro_core::state::ExecutionState;
use fabro_types::{Principal, RunId, StageTiming};

use super::circuit_breaker::CircuitBreakerLifecycle;
use super::git::GitCheckpointResult;
use crate::context::{Context, WorkflowContext};
use crate::event::{Emitter, Event, StageScope};
use crate::graph::{WorkflowGraph, WorkflowNode};
use crate::outcome::{FailureCategory, FailureDetail, ModelUsage, Outcome, StageOutcome};
use crate::stage_execution::{StageExecution, StageExecutionTracker};
use crate::{artifact, context};

type WfRunState = ExecutionState<Option<ModelUsage>>;
type WfNodeResult = NodeResult<Option<ModelUsage>>;
type FailureSignatureSnapshot = (
    Option<BTreeMap<String, usize>>,
    Option<BTreeMap<String, usize>>,
);

/// Sub-lifecycle responsible for emitting workflow run events.
pub(crate) struct EventLifecycle {
    pub emitter:               Arc<Emitter>,
    pub graph_name:            String,
    pub run_id:                RunId,
    pub run_start:             Mutex<Instant>,
    /// Set in on_edge_selected when loop_restart approved; emitted+cleared in
    /// on_run_start.
    pub restarted_from:        Arc<Mutex<Option<(String, String)>>>,
    // Config for WorkflowRunStarted payload
    pub base_branch:           Option<String>,
    pub base_sha:              Option<String>,
    pub run_branch:            Option<String>,
    pub worktree_dir:          Option<String>,
    pub goal:                  Option<String>,
    /// Shared git checkpoint result (written by GitLifecycle, read by
    /// EventLifecycle when emitting CheckpointCompleted).
    pub checkpoint_git_result: Arc<Mutex<Option<GitCheckpointResult>>>,
    pub circuit_breaker:       Arc<CircuitBreakerLifecycle>,
    /// Run-scoped stage execution allocator shared with `RunServices`.
    pub stage_executions:      StageExecutionTracker,
}

fn snapshot_failure_signatures(
    circuit_breaker: &CircuitBreakerLifecycle,
) -> FailureSignatureSnapshot {
    let (loop_sigs, restart_sigs) = circuit_breaker.snapshot();
    let loop_sigs = (!loop_sigs.is_empty()).then(|| {
        loop_sigs
            .into_iter()
            .map(|(sig, count)| (sig.to_string(), count))
            .collect::<BTreeMap<_, _>>()
    });
    let restart_sigs = (!restart_sigs.is_empty()).then(|| {
        restart_sigs
            .into_iter()
            .map(|(sig, count)| (sig.to_string(), count))
            .collect::<BTreeMap<_, _>>()
    });
    (loop_sigs, restart_sigs)
}

fn actor_for_stage_failure(failure: &FailureDetail) -> Option<Principal> {
    failure
        .system_actor
        .map(|system_kind| Principal::System { system_kind })
}

/// Build a [`StageTiming`] from a [`WfNodeResult`]. Inference and tool time
/// flow from the executor's `NodeResult` fields, which are populated from
/// `outcome.timing` by [`fabro_core`]. Handlers without an active-time
/// breakdown produce a wall-only timing.
fn node_result_timing(result: &WfNodeResult) -> StageTiming {
    StageTiming::new(
        crate::millis_u64(result.wall_time),
        crate::millis_u64(result.inference_time),
        crate::millis_u64(result.tool_time),
    )
}

fn response_from_outcome(node_id: &str, outcome: &Outcome) -> Option<String> {
    outcome
        .context_updates
        .get(&context::keys::response_key(node_id))
        .and_then(|value| value.as_str().map(ToOwned::to_owned))
}

/// Context values for `StageCompleted` events. Runtime-only keys are stripped.
fn stage_context_values(workflow_context: &Context) -> Option<BTreeMap<String, serde_json::Value>> {
    let mut snapshot = workflow_context.snapshot();
    artifact::strip_transient_keys(&mut snapshot);
    (!snapshot.is_empty()).then(|| snapshot.into_iter().collect())
}

pub(super) fn stage_visit(state: &WfRunState, node_id: &str) -> u32 {
    let visits = state.node_visits.get(node_id).copied().unwrap_or(1);
    u32::try_from(visits).unwrap_or(u32::MAX)
}

fn stage_scope_from_execution(
    execution: Option<&StageExecution>,
    state: &WfRunState,
    node_id: &str,
) -> StageScope {
    let (node_id, visit) = execution.map_or_else(
        || (node_id.to_owned(), stage_visit(state, node_id)),
        |execution| {
            (
                execution.stage_id.node_id().to_owned(),
                execution.stage_id.visit(),
            )
        },
    );
    StageScope {
        node_id,
        visit,
        parallel_group_id: state.context.parallel_group_id(),
        parallel_branch_id: state.context.parallel_branch_id(),
    }
}

/// Build the emission scope for a node from its active stage execution.
/// Falls back to the graph visit for direct unit-test call sites that emit
/// without a reservation; the two are equal for a first execution.
pub(crate) fn stage_scope_for(
    stage_executions: &StageExecutionTracker,
    state: &WfRunState,
    node_id: &str,
) -> StageScope {
    let execution = stage_executions.active(node_id);
    stage_scope_from_execution(execution.as_deref(), state, node_id)
}

#[async_trait]
impl RunLifecycle<WorkflowGraph> for EventLifecycle {
    async fn on_run_start(&self, _graph: &WorkflowGraph, _state: &WfRunState) -> CoreResult<()> {
        // If restarted_from is Some, emit LoopRestart and clear it
        {
            let mut restarted = self.restarted_from.lock()
                .expect("event lifecycle mutex should not be poisoned: no code panics while holding this lock");
            if let Some((from_node, to_node)) = restarted.take() {
                self.emitter
                    .emit(&Event::LoopRestart { from_node, to_node });
            }
        }

        // Reset run_start for duration measurement
        *self.run_start.lock().expect(
            "event lifecycle mutex should not be poisoned: no code panics while holding this lock",
        ) = Instant::now();

        // Emit RunStarted
        self.emitter.emit(&Event::WorkflowRunStarted {
            name:         self.graph_name.clone(),
            run_id:       self.run_id,
            base_branch:  self.base_branch.clone(),
            base_sha:     self.base_sha.clone(),
            run_branch:   self.run_branch.clone(),
            worktree_dir: self.worktree_dir.clone(),
            goal:         self.goal.clone(),
        });
        self.emitter.emit(&Event::RunRunning);

        Ok(())
    }

    async fn on_terminal_reached(
        &self,
        node: &WorkflowNode,
        goal_gates_passed: bool,
        state: &WfRunState,
    ) {
        if !goal_gates_passed {
            return;
        }
        let gv = node.inner();
        let stage_index = state.stage_index;
        // Terminal nodes bypass `before_node`/`before_attempt`, so their
        // synthetic paired events reserve an execution here.
        let execution = self
            .stage_executions
            .reserve(&gv.id, stage_visit(state, &gv.id));
        let scope = stage_scope_from_execution(Some(&execution), state, &gv.id);
        let (loop_failure_signatures, restart_failure_signatures) =
            snapshot_failure_signatures(&self.circuit_breaker);
        self.emitter.emit_scoped(
            &Event::StageStarted {
                node_id:               gv.id.clone(),
                name:                  gv.label().to_string(),
                index:                 stage_index,
                handler_type:          gv.handler_type().unwrap_or_default().to_string(),
                attempt:               1,
                max_attempts:          1,
                graph_visit:           Some(execution.graph_visit),
                resumed_from_stage_id: execution.resumed_from.clone(),
            },
            &scope,
        );
        self.emitter.emit_scoped(
            &Event::StageCompleted {
                node_id: gv.id.clone(),
                name: gv.label().to_string(),
                index: stage_index,
                timing: StageTiming::wall_only(0),
                status: StageOutcome::Succeeded.to_string(),
                preferred_label: None,
                suggested_next_ids: Vec::new(),
                usage_by_model: Vec::new(),
                usage: None,
                failure: None,
                notes: None,
                files_touched: Vec::new(),
                context_updates: None,
                jump_to_node: None,
                context_values: None,
                node_visits: None,
                loop_failure_signatures,
                restart_failure_signatures,
                response: state
                    .context
                    .get(&context::keys::response_key(&gv.id))
                    .and_then(|value| value.as_str().map(ToOwned::to_owned)),
                attempt: 1,
                max_attempts: 1,
            },
            &scope,
        );
    }

    async fn before_attempt(
        &self,
        ctx: &AttemptContext<'_, WorkflowGraph>,
        state: &WfRunState,
    ) -> CoreResult<NodeDecision<Option<ModelUsage>>> {
        let gv = ctx.node.inner();
        let execution = self.stage_executions.active(&gv.id);
        let scope = stage_scope_from_execution(execution.as_deref(), state, &gv.id);
        let graph_visit = execution
            .as_ref()
            .map_or_else(|| stage_visit(state, &gv.id), |e| e.graph_visit);
        self.emitter.emit_scoped(
            &Event::StageStarted {
                node_id:               gv.id.clone(),
                name:                  gv.label().to_string(),
                index:                 state.stage_index,
                handler_type:          gv.handler_type().unwrap_or_default().to_string(),
                attempt:               ctx.attempt as usize,
                max_attempts:          ctx.max_attempts as usize,
                graph_visit:           Some(graph_visit),
                resumed_from_stage_id: execution
                    .as_ref()
                    .and_then(|execution| execution.resumed_from.clone()),
            },
            &scope,
        );
        Ok(NodeDecision::Continue)
    }

    async fn after_attempt(
        &self,
        ctx: &AttemptResultContext<'_, WorkflowGraph>,
        state: &WfRunState,
    ) -> CoreResult<()> {
        if ctx.will_retry {
            let gv = ctx.node.inner();
            let outcome = &ctx.result.outcome;
            let stage_index = state.stage_index;
            let scope = stage_scope_for(&self.stage_executions, state, &gv.id);

            let timing = node_result_timing(ctx.result);
            let failure = outcome.failure.clone().unwrap_or_else(|| {
                FailureDetail::new("handler failed", FailureCategory::TransientInfra)
            });
            let actor = actor_for_stage_failure(&failure);
            self.emitter.emit_scoped(
                &Event::StageFailed {
                    node_id: gv.id.clone(),
                    name: gv.label().to_string(),
                    index: stage_index,
                    failure,
                    will_retry: true,
                    timing,
                    usage: outcome.usage.clone(),
                    usage_by_model: outcome.usage_by_model.clone(),
                    actor,
                },
                &scope,
            );

            self.emitter.emit_scoped(
                &Event::StageRetrying {
                    node_id:      gv.id.clone(),
                    name:         gv.label().to_string(),
                    index:        stage_index,
                    attempt:      ctx.attempt as usize,
                    max_attempts: ctx.result.max_attempts as usize,
                    delay_ms:     ctx.backoff_delay.map_or(0, crate::millis_u64),
                },
                &scope,
            );
        }
        Ok(())
    }

    async fn after_node(
        &self,
        node: &WorkflowNode,
        result: &mut WfNodeResult,
        state: &WfRunState,
    ) -> CoreResult<()> {
        let outcome = &result.outcome;
        // Skipped nodes had no StageStarted, so skip completion events (engine.rs:2080)
        if outcome.status == StageOutcome::Skipped {
            return Ok(());
        }
        let gv = node.inner();
        let stage_index = state.stage_index;
        let scope = stage_scope_for(&self.stage_executions, state, &gv.id);
        let timing = node_result_timing(result);
        let (loop_failure_signatures, restart_failure_signatures) =
            snapshot_failure_signatures(&self.circuit_breaker);

        if outcome.status.is_failure() {
            let failure = outcome.failure.clone().unwrap_or_else(|| {
                FailureDetail::new("handler failed", FailureCategory::Deterministic)
            });
            let actor = actor_for_stage_failure(&failure);
            self.emitter.emit_scoped(
                &Event::StageFailed {
                    node_id: gv.id.clone(),
                    name: gv.label().to_string(),
                    index: stage_index,
                    failure,
                    will_retry: false,
                    timing,
                    usage: outcome.usage.clone(),
                    usage_by_model: outcome.usage_by_model.clone(),
                    actor,
                },
                &scope,
            );
        } else {
            self.emitter.emit_scoped(
                &Event::StageCompleted {
                    node_id: gv.id.clone(),
                    name: gv.label().to_string(),
                    index: stage_index,
                    timing,
                    status: outcome.status.to_string(),
                    preferred_label: outcome.preferred_label.clone(),
                    suggested_next_ids: outcome.suggested_next_ids.clone(),
                    usage: outcome.usage.clone(),
                    usage_by_model: outcome.usage_by_model.clone(),
                    failure: outcome.failure.clone(),
                    notes: outcome.notes.clone(),
                    files_touched: outcome.files_touched.clone(),
                    context_updates: (!outcome.context_updates.is_empty()).then(|| {
                        outcome
                            .context_updates
                            .clone()
                            .into_iter()
                            .collect::<BTreeMap<_, _>>()
                    }),
                    jump_to_node: outcome.jump_to_node.clone(),
                    context_values: stage_context_values(&state.context),
                    node_visits: (!state.node_visits.is_empty()).then(|| {
                        state
                            .node_visits
                            .clone()
                            .into_iter()
                            .collect::<BTreeMap<_, _>>()
                    }),
                    loop_failure_signatures,
                    restart_failure_signatures,
                    response: response_from_outcome(&gv.id, outcome),
                    attempt: result.attempts as usize,
                    max_attempts: result.max_attempts as usize,
                },
                &scope,
            );
        }
        Ok(())
    }

    async fn on_edge_selected(
        &self,
        ctx: &EdgeContext<'_, WorkflowGraph>,
        _state: &WfRunState,
    ) -> CoreResult<EdgeDecision> {
        let outcome = ctx.outcome;
        let label = ctx
            .edge
            .as_ref()
            .and_then(|e| e.inner().label().map(String::from));
        let condition = ctx
            .edge
            .as_ref()
            .and_then(|e| e.inner().condition().map(String::from));
        self.emitter.emit(&Event::EdgeSelected {
            from_node: ctx.from.to_string(),
            to_node: ctx.to.to_string(),
            label,
            condition,
            reason: ctx.reason.to_string(),
            preferred_label: outcome.preferred_label.clone(),
            suggested_next_ids: outcome.suggested_next_ids.clone(),
            stage_status: outcome.status.to_string(),
            is_jump: ctx.is_jump,
        });
        Ok(EdgeDecision::Continue)
    }

    async fn on_checkpoint(
        &self,
        node: &WorkflowNode,
        result: &WfNodeResult,
        next_node_id: Option<&str>,
        state: &WfRunState,
    ) -> CoreResult<()> {
        let status = result.outcome.status.to_string();

        // Read git checkpoint result (set by GitLifecycle)
        let git_result = self.checkpoint_git_result.lock()
            .expect("event lifecycle mutex should not be poisoned: no code panics while holding this lock")
            .clone();

        let git_sha = git_result.as_ref().and_then(|r| r.commit_sha.clone());
        let diff = git_result.as_ref().and_then(|r| r.diff.clone());
        let diff_summary = git_result.as_ref().and_then(|r| r.diff_summary);
        let (loop_failure_signatures, restart_failure_signatures) =
            snapshot_failure_signatures(&self.circuit_breaker);
        let context_values = artifact::durable_context_snapshot(&state.context);
        let mut node_outcomes = state.node_outcomes.clone();
        node_outcomes.insert(node.id().to_string(), result.outcome.clone());
        artifact::normalize_durable_outcomes(&mut node_outcomes);

        let execution = self.stage_executions.active(node.id());
        let scope = stage_scope_from_execution(execution.as_deref(), state, node.id());
        let graph_visit = execution
            .as_ref()
            .map_or_else(|| stage_visit(state, node.id()), |e| e.graph_visit);
        self.emitter.emit_scoped(
            &Event::CheckpointCompleted {
                node_id: node.id().to_string(),
                status,
                current_node: node.id().to_string(),
                completed_nodes: state.completed_nodes.clone(),
                node_retries: state
                    .node_retries
                    .clone()
                    .into_iter()
                    .collect::<BTreeMap<_, _>>(),
                context_values: context_values.into_iter().collect::<BTreeMap<_, _>>(),
                node_outcomes: node_outcomes.into_iter().collect::<BTreeMap<_, _>>(),
                next_node_id: next_node_id.map(ToOwned::to_owned),
                git_commit_sha: git_sha.clone(),
                loop_failure_signatures: loop_failure_signatures.unwrap_or_default(),
                restart_failure_signatures: restart_failure_signatures.unwrap_or_default(),
                node_visits: state
                    .node_visits
                    .clone()
                    .into_iter()
                    .collect::<BTreeMap<_, _>>(),
                diff,
                diff_summary,
                graph_visit: Some(graph_visit),
                resumed_from_stage_id: execution
                    .as_ref()
                    .and_then(|execution| execution.resumed_from.clone()),
            },
            &scope,
        );

        // Emit GitCommit + GitPush events if git produced results
        if let Some(ref result) = git_result {
            if let Some(ref sha) = result.commit_sha {
                self.emitter.emit_scoped(
                    &Event::GitCommit {
                        node_id: Some(node.id().to_string()),
                        sha:     sha.clone(),
                    },
                    &scope,
                );
            }
            for push in &result.push_results {
                self.emitter.emit(&Event::GitPush {
                    branch:           push.branch.clone(),
                    success:          push.success,
                    exec_output_tail: push.exec_output_tail.clone(),
                    attempts:         push.attempts.clone(),
                });
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn stage_context_values_drops_runtime_keys_including_current_preamble() {
        let workflow_context = Context::new();
        workflow_context.set(
            context::keys::INTERNAL_PARALLEL_BRANCH_PREAMBLES,
            serde_json::json!([{"fidelity": "summary:high", "preamble": "runtime only"}]),
        );
        workflow_context.set(
            context::keys::INTERNAL_STAGE_EXECUTION_ORDINAL,
            serde_json::json!(2),
        );
        workflow_context.set(
            context::keys::CURRENT_PREAMBLE,
            serde_json::json!("active preamble"),
        );
        workflow_context.set("response.work", serde_json::json!("durable"));

        let values = stage_context_values(&workflow_context).expect("snapshot should not be empty");

        assert!(!values.contains_key(context::keys::INTERNAL_PARALLEL_BRANCH_PREAMBLES));
        assert!(!values.contains_key(context::keys::INTERNAL_STAGE_EXECUTION_ORDINAL));
        assert!(!values.contains_key(context::keys::CURRENT_PREAMBLE));
        assert_eq!(
            values.get("response.work"),
            Some(&serde_json::json!("durable"))
        );
    }
}