pe_graph/

compiled.rs

//! Compiled graph — the user-facing execution handle.
//!
//! `CompiledGraph` wraps a validated `StateGraph` and provides the
//! execution API: `invoke()`, `resume()`, `get_state()`, `get_state_history()`.

use crate::checkpoint_data::CheckpointData;
use crate::checkpointer::Checkpointer;
use crate::command::Command;
use crate::config::GraphConfig;
use crate::graph::StateGraph;
use crate::pregel::PregelEngine;
use crate::snapshot::StateSnapshot;
use pe_core::error::PeError;
use pe_core::lobe::LobeRuntimeServiceFactory;
use pe_core::node::InterruptRequest;
use pe_core::state::State;
use std::any::Any;
use std::sync::Arc;

/// Outcome of a graph execution.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ExecutionOutcome<S: State> {
    /// Graph ran to END — contains final state.
    Completed(S),

    /// Graph hit an interrupt — contains state at pause point.
    Interrupted {
        /// State at the moment of interruption.
        state: S,
        /// The interrupt request from the node.
        request: InterruptRequest<S::Update>,
    },
}

/// A validated, executable graph. Produced by `StateGraph::compile()`.
///
/// This is the primary execution handle. Use `invoke()` to run a graph
/// from initial state, `resume()` to continue after an interrupt, and
/// `get_state()` to inspect the current checkpoint.
///
/// # Example
///
/// ```ignore
/// let outcome = graph.invoke(initial_state, GraphConfig::default()).await?;
/// match outcome {
///     ExecutionOutcome::Completed(state) => println!("Done: {:?}", state),
///     ExecutionOutcome::Interrupted { state, request } => {
///         println!("Paused: {}", request.reason);
///     }
/// }
/// ```
pub struct CompiledGraph<S: State> {
    pub(crate) graph: Arc<StateGraph<S>>,
    checkpointer: Option<Arc<dyn Checkpointer>>,
    /// Optional matrix layer hook for convergence tracking and learned routing.
    matrix_hook: Option<crate::matrix_hook::MatrixHookHandle>,
    /// Optional agent bound to this graph (set by GraphBuilder).
    agent: Option<pe_core::agent::Agent>,
}

impl<S: State> std::fmt::Debug for CompiledGraph<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CompiledGraph")
            .field("nodes", &self.graph.nodes.keys().collect::<Vec<_>>())
            .field("has_checkpointer", &self.checkpointer.is_some())
            .field("has_matrix_hook", &self.matrix_hook.is_some())
            .field("has_agent", &self.agent.is_some())
            .finish()
    }
}

impl<S: State> CompiledGraph<S> {
    /// Create a new compiled graph (called by `StateGraph::compile()`).
    pub(crate) fn new(graph: Arc<StateGraph<S>>) -> Self {
        Self {
            graph,
            checkpointer: None,
            agent: None,
            matrix_hook: None,
        }
    }

    /// Attach a checkpointer for durable state persistence.
    ///
    /// Without a checkpointer, `resume()` and `get_state()` will return errors.
    pub fn with_checkpointer(mut self, cp: impl Checkpointer + 'static) -> Self {
        self.checkpointer = Some(Arc::new(cp));
        self
    }

    /// Bind an agent to this compiled graph.
    ///
    /// The agent's identity, system prompt, and boundaries are preserved
    /// on the compiled graph for runtime inspection and enforcement.
    pub fn with_agent(mut self, agent: pe_core::agent::Agent) -> Self {
        self.agent = Some(agent);
        self
    }

    /// Get the bound agent, if any.
    pub fn agent(&self) -> Option<&pe_core::agent::Agent> {
        self.agent.as_ref()
    }

    /// Attach a shared checkpointer (already behind `Arc`).
    pub fn with_checkpointer_arc(mut self, cp: Arc<dyn Checkpointer>) -> Self {
        self.checkpointer = Some(cp);
        self
    }

    /// Attach a matrix layer hook for convergence tracking and learned routing.
    ///
    /// When attached, the Pregel engine will:
    /// - Record `ConvergenceSignal` metadata via the hook
    /// - Consult the hook for conditional edge routing decisions
    /// - Record transitions for learning
    ///
    /// Without a hook, `NodeResult::Converge` degrades to `Update` and
    /// conditional edges use the user's router function directly.
    pub fn with_matrix_hook(mut self, hook: crate::matrix_hook::MatrixHookHandle) -> Self {
        self.matrix_hook = Some(hook);
        self
    }

    /// Get a reference to the matrix hook (if attached).
    pub fn matrix_hook(&self) -> Option<&crate::matrix_hook::MatrixHookHandle> {
        self.matrix_hook.as_ref()
    }

    /// Run graph from START with initial state.
    ///
    /// If `config.checkpoint_id` is set and a checkpointer is attached,
    /// the graph resumes from that specific checkpoint (time travel)
    /// instead of running from the provided `state`.
    ///
    /// Executes the BSP loop until END, interrupt, or recursion limit.
    #[must_use = "the execution outcome contains the final state"]
    pub async fn invoke(
        &self,
        state: S,
        config: GraphConfig,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        self.invoke_with_lobe_runtime_services(state, config, None)
            .await
    }

    /// Run graph from START with optional runtime-owned services for lobes.
    #[must_use = "the execution outcome contains the final state"]
    pub async fn invoke_with_lobe_runtime_services(
        &self,
        state: S,
        config: GraphConfig,
        lobe_runtime_service_factory: Option<Arc<dyn LobeRuntimeServiceFactory>>,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        self.invoke_with_observer_and_lobe_runtime_services(
            state,
            config,
            None,
            None,
            lobe_runtime_service_factory,
        )
        .await
    }

    /// Run graph from START with optional observer, tool observer, and runtime-owned services.
    #[must_use = "the execution outcome contains the final state"]
    pub async fn invoke_with_observer_and_lobe_runtime_services(
        &self,
        state: S,
        config: GraphConfig,
        observer: Option<Arc<dyn pe_core::node::NodeObserver>>,
        tool_observer: Option<Arc<dyn pe_core::node::ToolObserver>>,
        lobe_runtime_service_factory: Option<Arc<dyn LobeRuntimeServiceFactory>>,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        // Time travel: if checkpoint_id is specified, resume from that checkpoint
        if let Some(ref cp_id) = config.checkpoint_id {
            let data = self.load_checkpoint_data(&config.thread_id, cp_id).await?;
            let mut engine = self.make_engine(config);
            if let Some(obs) = observer.clone() {
                engine = engine.with_observer(obs);
            }
            if let Some(tobs) = tool_observer.clone() {
                engine = engine.with_tool_observer(tobs);
            }
            if let Some(factory) = lobe_runtime_service_factory {
                engine = engine.with_lobe_runtime_service_factory(factory);
            }
            return engine.run_from_checkpoint(data).await;
        }

        let mut engine = self.make_engine(config);
        if let Some(obs) = observer {
            engine = engine.with_observer(obs);
        }
        if let Some(tobs) = tool_observer {
            engine = engine.with_tool_observer(tobs);
        }
        if let Some(factory) = lobe_runtime_service_factory {
            engine = engine.with_lobe_runtime_service_factory(factory);
        }
        engine.run(state).await
    }

    /// Run graph from START with streaming support.
    ///
    /// Like [`invoke`](Self::invoke), but injects a type-erased stream
    /// sender into every [`NodeContext`](pe_core::node::NodeContext)
    /// and an optional [`NodeObserver`](pe_core::node::NodeObserver) for phase lifecycle events.
    ///
    /// Supports time-travel: if `config.checkpoint_id` is set, loads and
    /// resumes from that checkpoint (mirroring [`invoke`](Self::invoke)).
    ///
    /// Called by pe-runtime's streaming layer — not typically used directly.
    #[must_use = "the execution outcome contains the final state"]
    pub async fn invoke_with_stream(
        &self,
        state: S,
        config: GraphConfig,
        stream_sender: Arc<dyn Any + Send + Sync>,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        self.invoke_with_stream_and_observer(state, config, stream_sender, None, None)
            .await
    }

    /// Run graph with streaming and a [`NodeObserver`](pe_core::node::NodeObserver) for lifecycle events.
    ///
    /// The observer receives `on_node_start` / `on_node_complete` /
    /// `on_node_error` callbacks. pe-runtime provides `StreamingObserver`
    /// which converts these to `StreamEvent`.
    #[must_use = "the execution outcome contains the final state"]
    pub async fn invoke_with_stream_and_observer(
        &self,
        state: S,
        config: GraphConfig,
        stream_sender: Arc<dyn Any + Send + Sync>,
        observer: Option<Arc<dyn pe_core::node::NodeObserver>>,
        tool_observer: Option<Arc<dyn pe_core::node::ToolObserver>>,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        self.invoke_with_stream_observer_and_lobe_runtime_services(
            state,
            config,
            stream_sender,
            observer,
            tool_observer,
            None,
        )
        .await
    }

    /// Run graph with streaming, observers, and optional runtime-owned lobe services.
    #[must_use = "the execution outcome contains the final state"]
    pub async fn invoke_with_stream_observer_and_lobe_runtime_services(
        &self,
        state: S,
        config: GraphConfig,
        stream_sender: Arc<dyn Any + Send + Sync>,
        observer: Option<Arc<dyn pe_core::node::NodeObserver>>,
        tool_observer: Option<Arc<dyn pe_core::node::ToolObserver>>,
        lobe_runtime_service_factory: Option<Arc<dyn LobeRuntimeServiceFactory>>,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        // Time-travel: mirror the checkpoint-based resume from invoke()
        if let Some(ref cp_id) = config.checkpoint_id {
            let data = self.load_checkpoint_data(&config.thread_id, cp_id).await?;
            let mut engine = self.make_engine(config).with_stream_sender(stream_sender);
            if let Some(obs) = observer {
                engine = engine.with_observer(obs);
            }
            if let Some(tobs) = tool_observer {
                engine = engine.with_tool_observer(tobs);
            }
            if let Some(factory) = lobe_runtime_service_factory {
                engine = engine.with_lobe_runtime_service_factory(factory);
            }
            return engine.run_from_checkpoint(data).await;
        }

        let mut engine = self.make_engine(config).with_stream_sender(stream_sender);
        if let Some(obs) = observer {
            engine = engine.with_observer(obs);
        }
        if let Some(tobs) = tool_observer {
            engine = engine.with_tool_observer(tobs);
        }
        if let Some(factory) = lobe_runtime_service_factory {
            engine = engine.with_lobe_runtime_service_factory(factory);
        }
        engine.run(state).await
    }

    /// Resume a previously interrupted graph with human input.
    ///
    /// Loads the latest checkpoint for the thread, applies the input update,
    /// and continues execution from where it paused.
    ///
    #[must_use = "the execution outcome contains the final state"]
    pub async fn resume(
        &self,
        thread_id: &str,
        input: S::Update,
        config: GraphConfig,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        if thread_id != config.thread_id {
            return Err(PeError::GraphValue {
                details: format!(
                    "resume() thread_id '{}' does not match config.thread_id '{}'",
                    thread_id, config.thread_id
                ),
            });
        }

        let cp = self.checkpointer.as_ref().ok_or(PeError::Storage {
            details: "Cannot resume without a checkpointer".into(),
        })?;

        let (bytes, meta) =
            cp.load_latest(thread_id)
                .await?
                .ok_or(PeError::CheckpointNotFound {
                    thread_id: thread_id.to_string(),
                })?;

        let mut data: CheckpointData<S> =
            serde_json::from_slice(&bytes).map_err(|e| PeError::Storage {
                details: format!("Checkpoint deserialization failed: {e}"),
            })?;
        // Seed lineage — checkpoint_id is #[serde(skip)] so we inject it here
        data.checkpoint_id = Some(meta.id.clone());

        // Apply human input to the checkpointed state
        data.state.apply(input);

        // The old resume() API continues from successors, not re-running the
        // interrupted node. Resolve fixed-edge successors.
        if let Some(ref interrupted) = data.interrupted_node {
            let successors = self.graph.fixed_successors(interrupted);
            if !successors.is_empty() {
                data.next_nodes = successors;
            }
        }

        self.make_engine(config).run_from_checkpoint(data).await
    }

    /// Resume a previously interrupted graph using a [`Command`].
    ///
    /// This is the preferred resume API. It loads the latest checkpoint,
    /// applies the command (human input, goto, or state update), and
    /// continues execution.
    ///
    /// For `Command::Resume`, the human input is stored in the phase state
    /// so nodes can access it via `PhaseStateStore::get::<HumanInput>()`.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let cmd = Command::resume(HumanInput { approved: true, feedback: None, data: None });
    /// let outcome = graph.resume_with("thread-1", cmd, config).await?;
    /// ```
    #[must_use = "the execution outcome contains the final state"]
    pub async fn resume_with(
        &self,
        thread_id: &str,
        command: Command,
        config: GraphConfig,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        if thread_id != config.thread_id {
            return Err(PeError::GraphValue {
                details: format!(
                    "resume() thread_id '{}' does not match config.thread_id '{}'",
                    thread_id, config.thread_id
                ),
            });
        }

        let data = self.load_and_apply_command(thread_id, command).await?;
        self.make_engine(config).run_from_checkpoint(data).await
    }

    /// Resume with streaming and observer support.
    ///
    /// Like [`resume_with`](Self::resume_with), but injects a stream sender
    /// and optional observer for lifecycle events.
    #[must_use = "the execution outcome contains the final state"]
    pub async fn resume_with_stream(
        &self,
        thread_id: &str,
        command: Command,
        config: GraphConfig,
        stream_sender: Arc<dyn Any + Send + Sync>,
        observer: Option<Arc<dyn pe_core::node::NodeObserver>>,
        tool_observer: Option<Arc<dyn pe_core::node::ToolObserver>>,
    ) -> Result<ExecutionOutcome<S>, PeError> {
        if thread_id != config.thread_id {
            return Err(PeError::GraphValue {
                details: format!(
                    "resume() thread_id '{}' does not match config.thread_id '{}'",
                    thread_id, config.thread_id
                ),
            });
        }

        let data = self.load_and_apply_command(thread_id, command).await?;
        let mut engine = self.make_engine(config).with_stream_sender(stream_sender);
        if let Some(obs) = observer {
            engine = engine.with_observer(obs);
        }
        if let Some(tobs) = tool_observer {
            engine = engine.with_tool_observer(tobs);
        }
        engine.run_from_checkpoint(data).await
    }

    /// Create a base Pregel engine with the matrix hook applied (if attached).
    fn make_engine(&self, config: GraphConfig) -> PregelEngine<S> {
        let mut engine =
            PregelEngine::new(Arc::clone(&self.graph), config, self.checkpointer.clone());
        if let Some(ref hook) = self.matrix_hook {
            engine = engine.with_matrix_hook(hook.clone());
        }
        engine
    }

    /// Load a specific checkpoint by ID (for time-travel).
    ///
    /// Shared between `invoke` and `invoke_with_stream_and_observer` to
    /// avoid duplicating the checkpoint loading + deserialization logic.
    async fn load_checkpoint_data(
        &self,
        thread_id: &str,
        checkpoint_id: &str,
    ) -> Result<CheckpointData<S>, PeError> {
        let cp = self.checkpointer.as_ref().ok_or(PeError::Storage {
            details: "Cannot time-travel without a checkpointer".into(),
        })?;
        let bytes =
            cp.load_by_id(thread_id, checkpoint_id)
                .await?
                .ok_or(PeError::CheckpointNotFound {
                    thread_id: format!("{}@{}", thread_id, checkpoint_id),
                })?;
        serde_json::from_slice(&bytes).map_err(|e| PeError::Storage {
            details: format!("Checkpoint deserialization failed: {e}"),
        })
    }

    /// Load checkpoint and apply a command to it.
    ///
    /// Shared logic between `resume_with` and `resume_with_stream`.
    /// Handles all three command variants: Resume (stores HumanInput in
    /// phase state), Goto (overrides next_nodes), Update (deserializes
    /// JSON and applies to state, then resolves successors).
    async fn load_and_apply_command(
        &self,
        thread_id: &str,
        command: Command,
    ) -> Result<CheckpointData<S>, PeError> {
        let cp = self.checkpointer.as_ref().ok_or(PeError::Storage {
            details: "Cannot resume without a checkpointer".into(),
        })?;

        let (bytes, meta) =
            cp.load_latest(thread_id)
                .await?
                .ok_or(PeError::CheckpointNotFound {
                    thread_id: thread_id.to_string(),
                })?;

        let mut data: CheckpointData<S> =
            serde_json::from_slice(&bytes).map_err(|e| PeError::Storage {
                details: format!("Checkpoint deserialization failed: {e}"),
            })?;
        // Seed lineage for the engine
        data.checkpoint_id = Some(meta.id.clone());

        match command {
            Command::Resume { human_input } => {
                // Store human input in the phase state so resumed nodes can
                // access it via PhaseStateStore::get::<HumanInput>().
                data.phase_state
                    .set(&human_input)
                    .map_err(|e| PeError::Storage {
                        details: format!("Failed to store human input: {e}"),
                    })?;
            }
            Command::Goto { node } => {
                if !self.graph.nodes.contains_key(&node) {
                    return Err(PeError::GraphValue {
                        details: format!("Goto target node '{}' does not exist", node),
                    });
                }
                data.next_nodes = vec![node];
            }
            Command::Update { update } => {
                let typed_update: S::Update =
                    serde_json::from_value(update).map_err(|e| PeError::InvalidUpdate {
                        details: format!("Command::Update deserialization failed: {e}"),
                    })?;
                data.state.apply(typed_update);
                // After applying the update, skip re-running the interrupted node.
                // Resolve its successors using fixed edges so execution continues.
                if let Some(ref interrupted) = data.interrupted_node {
                    let successors = self.graph.fixed_successors(interrupted);
                    if !successors.is_empty() {
                        data.next_nodes = successors;
                    }
                }
            }
        }

        Ok(data)
    }

    /// Get the current state snapshot for a thread.
    ///
    /// Returns `None` if no checkpoints exist for this thread.
    #[must_use = "the snapshot contains the state — inspect it"]
    pub async fn get_state(&self, thread_id: &str) -> Result<Option<StateSnapshot<S>>, PeError> {
        let Some(ref cp) = self.checkpointer else {
            return Ok(None);
        };

        let Some((bytes, meta)) = cp.load_latest(thread_id).await? else {
            return Ok(None);
        };

        deserialize_snapshot(bytes, meta)
    }

    /// Get full history of all checkpoints for a thread (time travel).
    ///
    /// Returns snapshots oldest-first. Each snapshot contains the full
    /// state at that point, which nodes were scheduled next, and metadata.
    #[must_use = "the history contains all past states"]
    pub async fn get_state_history(
        &self,
        thread_id: &str,
    ) -> Result<Vec<StateSnapshot<S>>, PeError> {
        let Some(ref cp) = self.checkpointer else {
            return Ok(Vec::new());
        };

        let metas = cp.list(thread_id).await?;
        let mut snapshots = Vec::with_capacity(metas.len());

        for meta in metas {
            let Some(bytes) = cp.load_by_id(thread_id, &meta.id).await? else {
                continue;
            };
            if let Ok(Some(snapshot)) = deserialize_snapshot(bytes, meta) {
                snapshots.push(snapshot);
            }
        }

        Ok(snapshots)
    }
}

/// Deserialize checkpoint bytes into a StateSnapshot.
fn deserialize_snapshot<S: State>(
    bytes: Vec<u8>,
    meta: crate::checkpointer::CheckpointMeta,
) -> Result<Option<StateSnapshot<S>>, PeError> {
    let data: CheckpointData<S> = serde_json::from_slice(&bytes).map_err(|e| PeError::Storage {
        details: format!("Checkpoint deserialization failed: {e}"),
    })?;

    Ok(Some(StateSnapshot {
        state: data.state,
        checkpoint_id: meta.id.clone(),
        step: data.step,
        thread_id: meta.thread_id,
        parent_checkpoint_id: meta.parent_id.clone(),
        created_at: meta.created_at, // Use checkpoint time, not current time
        next_nodes: data.next_nodes,
    }))
}