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

pub(crate) mod artifact;
pub(crate) mod circuit_breaker;
pub(crate) mod event;
pub(crate) mod fidelity;
pub(crate) mod git;
pub(crate) mod hook;

use std::collections::HashMap;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Instant;

use async_trait::async_trait;
use fabro_core::error::{Error as CoreError, 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_graphviz::graph::types::Graph as GvGraph;
use fabro_hooks::HookRunner;
use fabro_sandbox::RunSandbox;
use fabro_types::RunId;

use self::artifact::ArtifactLifecycle;
use self::circuit_breaker::CircuitBreakerLifecycle;
use self::event::EventLifecycle;
use self::fidelity::FidelityLifecycle;
use self::git::{GitCheckpointResult, GitLifecycle};
use self::hook::HookLifecycle;
use crate::artifact_upload::ArtifactSink;
use crate::context;
use crate::error::FailureSignature;
use crate::event::Emitter;
use crate::graph::{WorkflowGraph, WorkflowNode};
use crate::outcome::{BilledModelUsage, Outcome};
use crate::run_control::RunControlState;
use crate::run_metadata::{RunMetadataRuntime, RunMetadataWriterHandle};
use crate::run_options::RunOptions;
use crate::runtime_store::RunStoreHandle;
use crate::sandbox_git_runtime::SandboxGitRuntime;
use crate::services::RunLocations;
use crate::stage_execution::StageExecutionTracker;

type WfRunState = ExecutionState<Option<BilledModelUsage>>;
type WfNodeResult = NodeResult<Option<BilledModelUsage>>;
type WfNodeDecision = NodeDecision<Option<BilledModelUsage>>;

/// Orchestrates all sub-lifecycles with explicit per-callback ordering.
/// Implements `RunLifecycle<WorkflowGraph>` by delegating to focused structs.
pub(crate) struct WorkflowLifecycle {
    event:                 EventLifecycle,
    hook:                  HookLifecycle,
    fidelity:              FidelityLifecycle,
    circuit_breaker:       Arc<CircuitBreakerLifecycle>,
    git:                   GitLifecycle,
    artifact:              ArtifactLifecycle,
    sandbox:               Arc<RunSandbox>,
    on_node:               crate::OnNodeCallback,
    emitter:               Arc<Emitter>,
    run_control:           Option<Arc<RunControlState>>,
    /// Set in on_edge_selected when loop_restart approved; read+cleared by
    /// EventLifecycle::on_run_start
    restarted_from:        Arc<Mutex<Option<(String, String)>>>,
    /// Shared git checkpoint result (written by git, read by event)
    checkpoint_git_result: Arc<Mutex<Option<GitCheckpointResult>>>,
    /// True when constructed with a checkpoint; cleared after first
    /// on_run_start. Gates context seeding on initial resume.
    is_initial_resume:     AtomicBool,
    /// Run-scoped stage execution allocator shared with `RunServices`.
    stage_executions:      StageExecutionTracker,
    // Config needed for context seeding
    graph:                 Arc<GvGraph>,
    run_id:                RunId,
    sandbox_work_dir:      Option<String>,
}

impl WorkflowLifecycle {
    #[allow(
        clippy::too_many_arguments,
        reason = "Workflow startup wires many run-scoped collaborators at once."
    )]
    pub(crate) fn new(
        emitter: &Arc<Emitter>,
        hook_runner: Option<Arc<HookRunner>>,
        sandbox: &Arc<RunSandbox>,
        graph: Arc<GvGraph>,
        run_dir: &Path,
        run_store: &RunStoreHandle,
        artifact_sink: Option<ArtifactSink>,
        locations: &RunLocations,
        run_options: &Arc<RunOptions>,
        sandbox_git: Arc<SandboxGitRuntime>,
        metadata_runtime: Arc<RunMetadataRuntime>,
        metadata_writer: Option<RunMetadataWriterHandle>,
        is_resume: bool,
        on_node: crate::OnNodeCallback,
        run_control: Option<Arc<RunControlState>>,
        stage_executions: StageExecutionTracker,
    ) -> Self {
        let restarted_from: Arc<Mutex<Option<(String, String)>>> = Arc::new(Mutex::new(None));
        let loop_restart_signature_limit = graph.loop_restart_signature_limit();
        let checkpoint_git_result: Arc<Mutex<Option<GitCheckpointResult>>> =
            Arc::new(Mutex::new(None));
        let last_git_sha: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));

        let circuit_breaker = Arc::new(CircuitBreakerLifecycle::new(loop_restart_signature_limit));

        let has_run_branch = run_options
            .git
            .as_ref()
            .and_then(|g| g.run_branch.as_ref())
            .is_some();
        let run_branch_sandbox_work_dir = if has_run_branch {
            locations
                .sandbox_work_dir
                .as_ref()
                .map(|path| path.display().to_string())
        } else {
            None
        };

        let event = EventLifecycle {
            emitter:               Arc::clone(emitter),
            graph_name:            graph.name.clone(),
            run_id:                run_options.run_id,
            run_start:             Mutex::new(Instant::now()),
            restarted_from:        Arc::clone(&restarted_from),
            base_branch:           run_options.base_branch.clone(),
            base_sha:              run_options.git.as_ref().and_then(|g| g.base_sha.clone()),
            run_branch:            run_options.git.as_ref().and_then(|g| g.run_branch.clone()),
            worktree_dir:          run_branch_sandbox_work_dir.clone(),
            goal:                  (!graph.goal().is_empty()).then(|| graph.goal().to_string()),
            checkpoint_git_result: Arc::clone(&checkpoint_git_result),
            circuit_breaker:       Arc::clone(&circuit_breaker),
            stage_executions:      stage_executions.clone(),
        };

        let hook = HookLifecycle {
            hook_runner,
            sandbox: Arc::clone(sandbox),
            hook_execution_context: locations.hook_execution_context(),
            run_id: run_options.run_id,
            graph_name: graph.name.clone(),
        };

        let fidelity = FidelityLifecycle::new(
            Arc::clone(&graph),
            Arc::clone(sandbox),
            run_store.clone(),
            run_dir.to_path_buf(),
        );

        let start_node_id = graph.find_start_node().map(|n| n.id.clone());

        let git = GitLifecycle {
            sandbox: Arc::clone(sandbox),
            emitter: Arc::clone(emitter),
            run_id: run_options.run_id,
            run_store: run_store.clone(),
            run_options: Arc::clone(run_options),
            sandbox_git,
            metadata_runtime,
            metadata_writer,
            start_node_id,
            checkpoint_git_result: Arc::clone(&checkpoint_git_result),
            last_git_sha,
            stage_executions: stage_executions.clone(),
        };

        let artifact = ArtifactLifecycle::new(
            Arc::clone(sandbox),
            run_store.clone(),
            Arc::clone(emitter),
            run_options.run_id,
            run_options.artifact_glob_patterns(),
            artifact_sink,
            stage_executions.clone(),
        );

        Self {
            event,
            hook,
            fidelity,
            circuit_breaker,
            git,
            artifact,
            sandbox: Arc::clone(sandbox),
            on_node,
            emitter: Arc::clone(emitter),
            run_control,
            restarted_from,
            checkpoint_git_result,
            is_initial_resume: AtomicBool::new(is_resume),
            stage_executions,
            graph,
            run_id: run_options.run_id,
            sandbox_work_dir: run_branch_sandbox_work_dir,
        }
    }

    /// Restore circuit breaker state from a checkpoint (for resume).
    pub(crate) fn restore_circuit_breaker(
        &self,
        loop_sigs: HashMap<FailureSignature, usize>,
        restart_sigs: HashMap<FailureSignature, usize>,
    ) {
        self.circuit_breaker.restore(loop_sigs, restart_sigs);
    }

    /// Set the fidelity degradation flag for checkpoint resume.
    pub(crate) fn set_degrade_fidelity_on_resume(&self, flag: bool) {
        self.fidelity.set_degrade_fidelity_on_resume(flag);
    }
}

#[async_trait]
impl RunLifecycle<WorkflowGraph> for WorkflowLifecycle {
    async fn on_run_start(&self, graph: &WorkflowGraph, state: &WfRunState) -> CoreResult<()> {
        // Re-seed context keys (fires on initial start AND after every loop restart).
        // Skip on initial checkpoint resume (context already has them).
        if self.is_initial_resume.swap(false, Ordering::Relaxed) {
            // First on_run_start after checkpoint resume — skip context seeding
        } else {
            // Mirror graph-level attributes into the core context
            if !self.graph.goal().is_empty() {
                state.context.set(
                    context::keys::GRAPH_GOAL,
                    serde_json::json!(self.graph.goal()),
                );
            }
            for (key, val) in &self.graph.attrs {
                state.context.set(
                    context::keys::graph_attr_key(key),
                    serde_json::json!(val.to_string_value()),
                );
            }
        }
        // Always set run_id and work_dir (idempotent)
        state.context.set(
            context::keys::INTERNAL_RUN_ID,
            serde_json::json!(self.run_id),
        );
        if let Some(ref wd) = self.sandbox_work_dir {
            state
                .context
                .set(context::keys::INTERNAL_WORK_DIR, serde_json::json!(wd));
        }

        // Reset restart-scoped state
        self.fidelity.on_run_start(graph, state).await?;
        self.artifact.on_run_start(graph, state).await?;
        // Observable callbacks
        self.event.on_run_start(graph, state).await?;
        self.hook.on_run_start(graph, state).await?;
        self.git.on_run_start(graph, state).await?;
        Ok(())
    }

    async fn on_terminal_reached(
        &self,
        node: &WorkflowNode,
        goal_gates_passed: bool,
        state: &WfRunState,
    ) {
        self.event
            .on_terminal_reached(node, goal_gates_passed, state)
            .await;
    }

    async fn before_node(
        &self,
        node: &WorkflowNode,
        state: &WfRunState,
    ) -> CoreResult<WfNodeDecision> {
        if let Some(run_control) = &self.run_control {
            run_control.wait_if_paused(self.emitter.as_ref()).await;
        }
        // A provider may auto-stop while the run is paused between nodes.
        self.sandbox.activate().await.map_err(|err| {
            CoreError::context(
                format!("failed to activate sandbox before node {}", node.id()),
                err,
            )
        })?;
        if let Some(on_node) = &self.on_node {
            on_node(node.id());
        }
        // Node boundary: clear the prior execution scope so the next
        // observable attempt reserves a fresh ordinal. No reservation happens
        // here — a hook block or process exit before any stage-scoped event
        // must not consume an ordinal.
        self.stage_executions.begin_node(node.id());
        state.context.set(
            context::keys::INTERNAL_STAGE_EXECUTION_ORDINAL,
            serde_json::Value::Null,
        );
        self.fidelity.before_node(node, state).await
    }

    async fn before_attempt(
        &self,
        ctx: &AttemptContext<'_, WorkflowGraph>,
        state: &WfRunState,
    ) -> CoreResult<WfNodeDecision> {
        // Hook first (can skip/block)
        match self.hook.before_attempt(ctx, state).await? {
            NodeDecision::Continue => {}
            decision => return Ok(decision),
        }
        // Reserve the stage execution once per handler invocation: the first
        // attempt allocates the ordinal and automatic retries reuse it.
        let node_id = ctx.node.id();
        let execution = self
            .stage_executions
            .ensure(node_id, event::stage_visit(state, node_id));
        state.context.set(
            context::keys::INTERNAL_STAGE_EXECUTION_ORDINAL,
            serde_json::json!(execution.stage_id.visit()),
        );
        // Event emission
        self.event.before_attempt(ctx, state).await?;
        // Record epoch AFTER hook+event (engine.rs:968→1006)
        self.artifact.before_attempt(ctx, state).await?;
        Ok(NodeDecision::Continue)
    }

    async fn after_attempt(
        &self,
        ctx: &AttemptResultContext<'_, WorkflowGraph>,
        state: &WfRunState,
    ) -> CoreResult<()> {
        if let Some(run_control) = &self.run_control {
            run_control.wait_if_paused(self.emitter.as_ref()).await;
        }
        // Human, wait, and paused stages can return after a long period with
        // no sandbox traffic. Reactivate before artifact and checkpoint work.
        self.sandbox.activate().await.map_err(|err| {
            CoreError::context(
                format!(
                    "failed to activate sandbox after node attempt {}",
                    ctx.node.id()
                ),
                err,
            )
        })?;
        self.artifact.after_attempt(ctx, state).await?;
        self.event.after_attempt(ctx, state).await?;
        Ok(())
    }

    async fn after_node(
        &self,
        node: &WorkflowNode,
        result: &mut WfNodeResult,
        state: &WfRunState,
    ) -> CoreResult<()> {
        self.circuit_breaker.after_node(node, result, state).await?;
        self.artifact.after_node(node, result, state).await?;
        self.event.after_node(node, result, state).await?;
        self.hook.after_node(node, result, state).await?;
        Ok(())
    }

    async fn after_record(
        &self,
        node: &WorkflowNode,
        result: &WfNodeResult,
        state: &WfRunState,
    ) -> CoreResult<()> {
        let retry_count = state.node_retries.get(node.id()).copied().unwrap_or(0);
        context::apply_recorded_outcome_context(
            &state.context,
            node.id(),
            &result.outcome,
            retry_count,
        );
        Ok(())
    }

    async fn on_edge_selected(
        &self,
        ctx: &EdgeContext<'_, WorkflowGraph>,
        state: &WfRunState,
    ) -> CoreResult<EdgeDecision> {
        // Fidelity captures edge data
        self.fidelity.on_edge_selected(ctx, state).await?;
        // Event always fires first
        self.event.on_edge_selected(ctx, state).await?;
        // Hook can override/block
        match self.hook.on_edge_selected(ctx, state).await? {
            EdgeDecision::Continue => {
                // Edge unchanged — check circuit breaker for loop_restart
                let decision = self.circuit_breaker.on_edge_selected(ctx, state).await?;
                // If loop_restart edge approved by both hook and circuit breaker, mark for
                // LoopRestart emission
                if matches!(decision, EdgeDecision::Continue) {
                    if let Some(ref edge) = ctx.edge {
                        if edge.inner().loop_restart() {
                            *self.restarted_from.lock()
                                .expect("lifecycle mutex should not be poisoned: no code panics while holding this lock") =
                                Some((ctx.from.to_string(), ctx.to.to_string()));
                        }
                    }
                }
                Ok(decision)
            }
            decision => Ok(decision), // Override/Block — skip circuit breaker
        }
    }

    async fn on_checkpoint(
        &self,
        node: &WorkflowNode,
        result: &WfNodeResult,
        next_node_id: Option<&str>,
        state: &WfRunState,
    ) -> CoreResult<()> {
        // A StageStart hook can skip before any attempt reserved an execution
        // scope. Ensure one exists so Git metadata-snapshot events and the
        // `checkpoint.completed` envelope attach to a concrete execution;
        // an existing reservation from the attempt path is reused as-is.
        let execution = self
            .stage_executions
            .ensure(node.id(), event::stage_visit(state, node.id()));
        state.context.set(
            context::keys::INTERNAL_STAGE_EXECUTION_ORDINAL,
            serde_json::json!(execution.stage_id.visit()),
        );
        self.git
            .on_checkpoint(node, result, next_node_id, state)
            .await?;
        self.event
            .on_checkpoint(node, result, next_node_id, state)
            .await?;
        self.hook
            .on_checkpoint(node, result, next_node_id, state)
            .await?;
        // Clear checkpoint result for next checkpoint
        *self.checkpoint_git_result.lock().expect(
            "lifecycle mutex should not be poisoned: no code panics while holding this lock",
        ) = None;
        Ok(())
    }

    async fn on_run_end(&self, outcome: &Outcome, state: &WfRunState) {
        self.hook.on_run_end(outcome, state).await;
    }
}