pe-graph 0.1.0

//! Pregel BSP execution engine — the heart of the graph runtime.
//!
//! Implements the Bulk Synchronous Parallel model:
//! 1. **Plan** — determine which nodes are active and why
//! 2. **Execute** — run active nodes in parallel on cloned state snapshots
//! 3. **Update** — collect all results, track pending writes, apply to state
//! 4. **Check** — halt if no nodes active, error if recursion limit hit
//!
//! ## Extension points
//!
//! The engine is fully functional without any optional layers. Extension points
//! allow higher layers to enhance behavior without modifying the core:
//!
//! - **Activation resolution** (`resolve_activations`): determines which nodes
//!   run next. By default, follows fixed + conditional edges. The optional
//!   matrix layer can replace conditional edge routing with learned routing.
//! - **Node context** (`NodeContext`): runtime metadata injected into each node.
//!   pe-runtime populates the metadata map with agent_id, thread_id, etc.
//! - **Pending writes**: fault tolerance tracking for optional RetryPolicy.

use crate::activation::Activation;
use crate::checkpoint_data::CheckpointData;
use crate::checkpointer::{CheckpointMeta, Checkpointer};
use crate::compiled::ExecutionOutcome;
use crate::config::GraphConfig;
use crate::graph::StateGraph;
use crate::matrix_hook::MatrixHookHandle;
use crate::retry::RetryPolicy;
use futures::StreamExt;
use pe_core::error::PeError;
use pe_core::lobe::LobeRuntimeServiceFactory;
use pe_core::node::{ActivationReason, NodeContext, NodeObserver, NodeResult};
use pe_core::state::State;
use pe_core::types::END;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

/// The BSP execution engine. Created internally by `CompiledGraph`.
pub(crate) struct PregelEngine<S: State> {
    graph: Arc<StateGraph<S>>,
    config: GraphConfig,
    checkpointer: Option<Arc<dyn Checkpointer>>,
    /// Type-erased stream sender for the streaming layer (Plan 011).
    /// Injected into every `NodeContext` so nodes can push events.
    stream_sender: Option<Arc<dyn std::any::Any + Send + Sync>>,
    /// Optional observer for node lifecycle events (Plan 011).
    /// pe-runtime provides a streaming implementation.
    observer: Option<Arc<dyn NodeObserver>>,
    /// Optional observer for tool call lifecycle events (Plan 011).
    /// Passed through to NodeContext for pe-tools to use.
    tool_observer: Option<Arc<dyn pe_core::node::ToolObserver>>,
    /// Optional factory for runtime-owned services exposed to lobes.
    lobe_runtime_service_factory: Option<Arc<dyn LobeRuntimeServiceFactory>>,
    /// Optional matrix layer hook for convergence tracking and learned routing.
    matrix_hook: Option<MatrixHookHandle>,
    /// Optional retry policy for failed nodes.
    retry_policy: Option<RetryPolicy>,
}

impl<S: State> PregelEngine<S> {
    pub fn new(
        graph: Arc<StateGraph<S>>,
        config: GraphConfig,
        checkpointer: Option<Arc<dyn Checkpointer>>,
    ) -> Self {
        let retry_policy = config.retry_policy.clone();
        Self {
            graph,
            config,
            checkpointer,
            stream_sender: None,
            observer: None,
            tool_observer: None,
            lobe_runtime_service_factory: None,
            matrix_hook: None,
            retry_policy,
        }
    }

    /// Create an engine with a stream sender for event emission.
    pub fn with_stream_sender(mut self, sender: Arc<dyn std::any::Any + Send + Sync>) -> Self {
        self.stream_sender = Some(sender);
        self
    }

    /// Attach a node lifecycle observer for phase-aware streaming.
    pub fn with_observer(mut self, observer: Arc<dyn NodeObserver>) -> Self {
        self.observer = Some(observer);
        self
    }

    /// Attach a tool call lifecycle observer for streaming.
    pub fn with_tool_observer(mut self, observer: Arc<dyn pe_core::node::ToolObserver>) -> Self {
        self.tool_observer = Some(observer);
        self
    }

    /// Attach runtime-owned services for lobes.
    pub fn with_lobe_runtime_service_factory(
        mut self,
        factory: Arc<dyn LobeRuntimeServiceFactory>,
    ) -> Self {
        self.lobe_runtime_service_factory = Some(factory);
        self
    }

    /// Attach a matrix layer hook for convergence tracking and learned routing.
    pub fn with_matrix_hook(mut self, hook: MatrixHookHandle) -> Self {
        self.matrix_hook = Some(hook);
        self
    }

    /// Run the BSP loop from initial state.
    pub async fn run(&self, state: S) -> Result<ExecutionOutcome<S>, PeError> {
        let activations = self.initial_activations();
        let phase_store = pe_core::phase_store::PhaseStateStore::new();
        self.execute_loop(state, 0, activations, phase_store).await
    }

    /// Run the BSP loop from a resumed checkpoint.
    pub async fn run_from_checkpoint(
        &self,
        data: CheckpointData<S>,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        let activations = data
            .next_nodes
            .into_iter()
            .map(|name| Activation {
                node_name: name,
                reason: ActivationReason::Resume,
            })
            .collect();
        self.execute_loop(data.state, data.step, activations, data.phase_state)
            .await
    }

    /// The single BSP loop — both `run` and `run_from_checkpoint` delegate here.
    async fn execute_loop(
        &self,
        mut state: S,
        mut step: u32,
        mut activations: Vec<Activation>,
        phase_store: pe_core::phase_store::PhaseStateStore,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        // Note: tracing spans use drop-before-await pattern to stay Send-safe.
        // Event-level tracing (info!, debug!) is used instead of span guards
        // that would be held across await points.
        tracing::info!(thread_id = %self.config.thread_id, "pregel_run_start");

        let start_time = Instant::now();

        loop {
            // HALT: no active nodes means graph completed
            if activations.is_empty() {
                return Ok(ExecutionOutcome::Completed(state));
            }

            // EXECUTION TIMEOUT: checked before running nodes so we never
            // cancel a node mid-execution. Uses wall-clock elapsed time.
            if let Some(max_duration) = self.config.max_execution_time {
                let elapsed = start_time.elapsed();
                if elapsed > max_duration {
                    return Err(PeError::Timeout {
                        seconds: elapsed.as_secs_f64(),
                    });
                }
            }

            // RECURSION LIMIT
            step += 1;
            if step > self.config.recursion_limit {
                tracing::warn!(
                    limit = self.config.recursion_limit,
                    step,
                    "recursion_limit_exceeded"
                );
                return Err(PeError::GraphRecursion {
                    limit: self.config.recursion_limit,
                });
            }

            tracing::debug!(step, "superstep_start");

            // Save the pre-step activation set so resume/goto policies can
            // recover even if an active node fails before post-step checkpointing.
            self.save_checkpoint(&state, &activations, step, &phase_store)
                .await?;

            // EXECUTE: run all active nodes in parallel on cloned snapshots
            let results = self
                .execute_parallel(&state, &activations, step, &phase_store)
                .await?;

            // NOTE(R01): PendingWrites tracking removed — was populated but never consumed.
            // RetryPolicy is now wired (retries individual failed nodes inline).
            // PendingWrites may be useful for skip-already-succeeded optimization later.
            // See: Checkpointer::put_writes (exists but uncalled), plans/R01-global-review/PLAN.md C2

            // CHECK for interrupts — first interrupt wins.
            // If multiple parallel nodes interrupt, only the first is returned.
            // Others' partial updates are NOT applied (they ran on snapshots).
            // This is the defined policy: parallel interrupts are rare, and the
            // human should resolve one at a time. Log if we drop any.
            let mut interrupt_found = false;
            let mut interrupt_outcome = None;
            for (node_name, result) in &results {
                if let NodeResult::Interrupt(request) = result {
                    if !interrupt_found {
                        interrupt_found = true;
                        tracing::info!(node = %node_name, "interrupt_requested");
                        if let Some(ref partial) = request.partial_update {
                            state.apply(partial.clone());
                        }
                        // Save checkpoint with the interrupted node as the
                        // next node to re-activate on resume. We store the
                        // interrupted node itself (not its successors) because
                        // on resume it needs to re-run from its current phase.
                        let resume_activations = vec![Activation {
                            node_name: node_name.clone(),
                            reason: ActivationReason::Resume,
                        }];
                        self.save_interrupt_checkpoint(
                            &state,
                            &resume_activations,
                            step,
                            node_name,
                            &phase_store,
                        )
                        .await?;
                        interrupt_outcome = Some(ExecutionOutcome::Interrupted {
                            state: state.clone(),
                            request: request.clone(),
                        });
                    } else {
                        tracing::warn!(node = %node_name, "dropped_parallel_interrupt");
                    }
                }
            }
            if let Some(outcome) = interrupt_outcome {
                return Ok(outcome);
            }

            // UPDATE: apply all results to canonical state
            for (node_name, result) in results {
                match result {
                    NodeResult::Update(update) => {
                        state.apply(update);
                    }
                    NodeResult::Converge(signal) => {
                        // With the matrix layer: record signal for convergence
                        // tracking and routing learning. Without it: degrade to
                        // Update. Either way, the partial_update is always applied.
                        if let Some(ref hook) = self.matrix_hook {
                            hook.on_converge(
                                &node_name,
                                signal.actual_contribution,
                                signal.surprise,
                                signal.quality,
                            );
                        }
                        state.apply(signal.partial_update);
                    }
                    NodeResult::Error(e) => {
                        // Retry the individual failed node if a retry policy
                        // is configured and the error is retryable.
                        if let Some(ref policy) = self.retry_policy {
                            if e.is_retryable() {
                                let retried = self
                                    .retry_single_node(
                                        &node_name,
                                        &state,
                                        step,
                                        &phase_store,
                                        policy,
                                    )
                                    .await;
                                match retried {
                                    Ok(update) => {
                                        state.apply(update);
                                        continue;
                                    }
                                    Err(final_err) => return Err(final_err),
                                }
                            }
                        }
                        return Err(e);
                    }
                    NodeResult::Interrupt(_) => unreachable!("handled above"),
                    // NOTE: Forward compat — new NodeResult variants must
                    // NOT be silently ignored. Return an error so the developer
                    // knows they need to handle it. Silent `_ => {}` would mask bugs.
                    #[allow(unreachable_patterns)]
                    other => {
                        return Err(PeError::Internal {
                            details: format!(
                                "Node '{}' returned unhandled NodeResult variant: {:?}",
                                node_name, other
                            ),
                        });
                    }
                }
            }

            // RESOLVE: determine next activations via edge traversal.
            // The optional matrix layer can replace conditional routing here.
            let completed_names: Vec<String> =
                activations.iter().map(|a| a.node_name.clone()).collect();
            let next = self.resolve_activations(&state, &completed_names);

            // Record transitions for matrix learning (if active)
            if let Some(ref hook) = self.matrix_hook {
                for from in &completed_names {
                    for to_activation in &next {
                        // Quality 1.0 for successful transitions — the convergence
                        // signals provide finer-grained quality data
                        hook.record_transition(from, &to_activation.node_name, 1.0);
                    }
                }
            }

            // CHECKPOINT the resolved successor set after a successful superstep
            self.save_checkpoint(&state, &next, step, &phase_store)
                .await?;

            activations = next;
        }
    }

    /// Run active nodes in parallel. Each gets a cloned state snapshot
    /// and a [`NodeContext`] with step info and activation reason.
    ///
    /// Decision 3: snapshot isolation — no locks, no contention.
    /// Respects `max_concurrency` if configured.
    async fn execute_parallel(
        &self,
        state: &S,
        activations: &[Activation],
        step: u32,
        phase_store: &pe_core::phase_store::PhaseStateStore,
    ) -> Result<Vec<(String, NodeResult<S::Update>)>, PeError> {
        let observer = self.observer.clone();
        // Pre-resolve activated nodes to avoid borrowing self.graph inside
        // the iterator chain. This sidesteps the HRTB lifetime issue that
        // arises when dyn NodeFn trait objects are captured in closures
        // called through Pin<Box<dyn Future>>.
        // Pre-resolve activated nodes and build futures in a for loop.
        // Using a for loop instead of .map() avoids HRTB lifetime issues
        // with dyn NodeFn trait objects captured in closures.
        let mut futures_vec = Vec::new();
        for activation in activations {
            let node = match self.graph.nodes.get(&activation.node_name) {
                Some(n) => Arc::clone(n),
                None => continue,
            };
            let activation = activation.clone();
            let snapshot = state.clone();
            let ctx = NodeContext {
                step,
                recursion_limit: self.config.recursion_limit,
                node_name: activation.node_name.clone(),
                activation: activation.reason.clone(),
                metadata: HashMap::new(), // pe-runtime populates (Plan 007)
                phase_store: phase_store.clone(),
                stream_sender: self.stream_sender.clone(),
                tool_observer: self.tool_observer.clone(),
                lobe_runtime_service_factory: self.lobe_runtime_service_factory.clone(),
            };
            let obs = observer.clone();
            futures_vec.push(async move {
                tracing::info!(
                    node = %activation.node_name,
                    step,
                    reason = ?activation.reason,
                    "node_execute_start"
                );

                // Notify observer before node starts
                if let Some(ref o) = obs {
                    o.on_node_start(&activation.node_name, step).await;
                }

                let start = std::time::Instant::now();
                let result = tokio::spawn(async move { node.call(&snapshot, &ctx).await })
                    .await
                    .map_err(|e| PeError::Internal {
                        details: format!("Node '{}' task panicked: {e}", activation.node_name),
                    });
                let elapsed = start.elapsed();

                // Notify observer after node completes
                if let Some(ref o) = obs {
                    match &result {
                        Ok(pe_core::node::NodeResult::Error(e)) => {
                            o.on_node_error(&activation.node_name, step, &e.to_string())
                                .await;
                        }
                        Err(e) => {
                            o.on_node_error(&activation.node_name, step, &e.to_string())
                                .await;
                        }
                        _ => {
                            o.on_node_complete(&activation.node_name, step, elapsed)
                                .await;
                        }
                    }
                }

                tracing::debug!(node = %activation.node_name, "node_completed");
                (activation.node_name, result)
            });
        }
        let futures_iter = futures_vec.into_iter();

        // Honor max_concurrency: limit parallel node executions
        let max = self.config.max_concurrency.unwrap_or(usize::MAX);
        let results: Vec<_> = futures::stream::iter(futures_iter)
            .buffer_unordered(max)
            .collect()
            .await;

        results
            .into_iter()
            .map(|(name, r): (String, Result<_, _>)| r.map(|result| (name, result)))
            .collect()
    }

    /// Retry a single failed node up to `policy.max_attempts` times with backoff.
    ///
    /// Returns `Ok(update)` if a retry succeeds, or `Err(last_error)` if all
    /// retries are exhausted or a non-retryable error occurs.
    async fn retry_single_node(
        &self,
        node_name: &str,
        state: &S,
        step: u32,
        phase_store: &pe_core::phase_store::PhaseStateStore,
        policy: &RetryPolicy,
    ) -> Result<S::Update, PeError> {
        let node = self.graph.nodes.get(node_name).ok_or(PeError::Internal {
            details: format!("Retry target node '{}' not found", node_name),
        })?;
        let node = Arc::clone(node);

        let mut delay = policy.initial_interval;
        let mut last_err = PeError::Internal {
            details: "retry_single_node called but never executed".into(),
        };

        for attempt in 1..=policy.max_attempts {
            tracing::info!(
                node = %node_name,
                attempt,
                max = policy.max_attempts,
                "node_retry"
            );

            if let Some(ref observer) = self.observer {
                observer
                    .on_node_retry(node_name, step, attempt, policy.max_attempts)
                    .await;
            }

            // Backoff before each retry attempt
            let sleep_dur = if policy.jitter {
                crate::retry::apply_jitter(delay)
            } else {
                delay
            };
            tokio::time::sleep(sleep_dur).await;
            delay = crate::retry::next_delay(delay, policy.backoff_factor, policy.max_interval);

            let snapshot = state.clone();
            let ctx = NodeContext {
                step,
                recursion_limit: self.config.recursion_limit,
                node_name: node_name.to_string(),
                activation: ActivationReason::Retry { attempt },
                metadata: HashMap::new(),
                phase_store: phase_store.clone(),
                stream_sender: self.stream_sender.clone(),
                tool_observer: self.tool_observer.clone(),
                lobe_runtime_service_factory: self.lobe_runtime_service_factory.clone(),
            };

            let result = node.call(&snapshot, &ctx).await;

            match result {
                NodeResult::Update(update) => return Ok(update),
                NodeResult::Error(e) if e.is_retryable() => {
                    last_err = e;
                    // continue to next retry
                }
                NodeResult::Error(e) => return Err(e),
                NodeResult::Interrupt(req) => {
                    // Interrupts during retry are not retried — propagate them.
                    // This is unusual but should not be swallowed.
                    return Err(PeError::Internal {
                        details: format!(
                            "Node '{}' interrupted during retry: {}",
                            node_name, req.reason
                        ),
                    });
                }
                NodeResult::Converge(signal) => return Ok(signal.partial_update),
                #[allow(unreachable_patterns)]
                _ => {
                    return Err(PeError::Internal {
                        details: format!(
                            "Node '{}' returned unhandled variant during retry",
                            node_name
                        ),
                    });
                }
            }
        }

        Err(last_err)
    }

    /// Collect initial activations from START edges.
    fn initial_activations(&self) -> Vec<Activation> {
        self.graph
            .edges
            .iter()
            .filter(|e| e.from == pe_core::types::START && e.to != END)
            .map(|e| Activation {
                node_name: e.to.clone(),
                reason: ActivationReason::EntryPoint,
            })
            .collect()
    }

    /// Traverse edges from completed nodes to determine what runs next.
    /// Returns activations with reasons (which edge/router triggered them).
    ///
    /// This is the extension point for learned routing: the optional matrix
    /// layer can replace conditional edge resolution with probabilistic
    /// routing while fixed edges continue to fire normally.
    fn resolve_activations(&self, state: &S, completed: &[String]) -> Vec<Activation> {
        let mut result = Vec::new();
        let mut seen = std::collections::HashSet::new();

        for node_name in completed {
            // Fixed edges
            for edge in &self.graph.edges {
                if edge.from == *node_name && edge.to != END && seen.insert(edge.to.clone()) {
                    tracing::debug!(from = %node_name, to = %edge.to, "transition");
                    result.push(Activation {
                        node_name: edge.to.clone(),
                        reason: ActivationReason::Edge {
                            from: node_name.clone(),
                        },
                    });
                }
            }
            // Conditional edges — matrix layer can override routing
            for ce in &self.graph.conditional_edges {
                if ce.from == *node_name {
                    // Always evaluate the user's router to get candidates
                    let candidates = (ce.router)(state);

                    // If matrix hook is active, let it re-route among candidates
                    let selected = if let Some(ref hook) = self.matrix_hook {
                        hook.route(node_name, &candidates).unwrap_or(candidates)
                    } else {
                        candidates
                    };

                    for target in selected {
                        if target != END && seen.insert(target.clone()) {
                            tracing::debug!(from = %node_name, to = %target, "conditional_transition");
                            result.push(Activation {
                                node_name: target,
                                reason: ActivationReason::ConditionalEdge {
                                    from: node_name.clone(),
                                },
                            });
                        }
                    }
                }
            }
        }

        result
    }

    /// Save checkpoint if a checkpointer is configured.
    async fn save_checkpoint(
        &self,
        state: &S,
        next: &[Activation],
        step: u32,
        phase_store: &pe_core::phase_store::PhaseStateStore,
    ) -> Result<(), PeError> {
        let Some(ref cp) = self.checkpointer else {
            return Ok(());
        };

        let data = CheckpointData {
            state: state.clone(),
            next_nodes: next.iter().map(|a| a.node_name.clone()).collect(),
            step,
            interrupted_node: None,
            phase_state: phase_store.clone(),
            checkpoint_id: None,
        };
        let bytes = serde_json::to_vec(&data).map_err(|e| PeError::Storage {
            details: format!("Checkpoint serialization failed: {e}"),
        })?;

        let checkpoint_id = uuid::Uuid::new_v4().to_string();
        let meta = CheckpointMeta::new(&checkpoint_id, &self.config.thread_id, step);

        cp.save(&self.config.thread_id, &checkpoint_id, &bytes, &meta)
            .await?;
        tracing::info!(checkpoint_id = %checkpoint_id, step, "checkpoint_saved");
        Ok(())
    }

    /// Save an interrupt checkpoint -- includes the interrupted node name
    /// and current phase state so `resume()` can re-activate the correct
    /// node with the correct phase.
    async fn save_interrupt_checkpoint(
        &self,
        state: &S,
        next: &[Activation],
        step: u32,
        interrupted_node: &str,
        phase_store: &pe_core::phase_store::PhaseStateStore,
    ) -> Result<(), PeError> {
        let Some(ref cp) = self.checkpointer else {
            return Ok(());
        };

        let data = CheckpointData {
            state: state.clone(),
            next_nodes: next.iter().map(|a| a.node_name.clone()).collect(),
            step,
            interrupted_node: Some(interrupted_node.to_string()),
            phase_state: phase_store.clone(),
            checkpoint_id: None,
        };
        let bytes = serde_json::to_vec(&data).map_err(|e| PeError::Storage {
            details: format!("Checkpoint serialization failed: {e}"),
        })?;

        let checkpoint_id = uuid::Uuid::new_v4().to_string();
        let meta = CheckpointMeta::new(&checkpoint_id, &self.config.thread_id, step);

        cp.save(&self.config.thread_id, &checkpoint_id, &bytes, &meta)
            .await
    }
}