pe_graph/

config.rs

//! Graph execution configuration.
//!
//! `GraphConfig` controls how a compiled graph executes: thread identity,
//! recursion limits, concurrency bounds, and user-defined configurable values.

use crate::retry::RetryPolicy;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration;

/// Configuration for a single graph execution.
///
/// # Example
///
/// ```
/// use pe_graph::GraphConfig;
///
/// let config = GraphConfig::default()
///     .with_thread_id("my-thread")
///     .with_recursion_limit(50);
/// ```
/// NOTE: `#[non_exhaustive]` — will grow with timeout, retry policy, etc.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct GraphConfig {
    /// Conversation thread identifier. Used as the checkpoint key.
    pub thread_id: String,

    /// Resume from a specific past checkpoint (time travel).
    pub checkpoint_id: Option<String>,

    /// Maximum supersteps before `PeError::GraphRecursion`. Default: 25.
    ///
    /// Kept low to catch infinite loops early in agentic systems.
    /// Users can raise this for legitimately deep graphs.
    pub recursion_limit: u32,

    /// Maximum parallel nodes per superstep. Default: unbounded.
    ///
    /// Tokio's thread pool naturally bounds actual parallelism.
    /// This provides an additional application-level limit.
    pub max_concurrency: Option<usize>,

    /// User-defined values readable inside nodes.
    pub configurable: HashMap<String, serde_json::Value>,

    /// Maximum wall-clock time for the entire graph execution.
    ///
    /// The Pregel engine checks elapsed time at the start of each superstep.
    /// If exceeded, returns `PeError::Timeout` before running nodes.
    /// Default: `None` (no timeout).
    #[serde(skip)]
    pub max_execution_time: Option<Duration>,

    /// Retry policy for failed nodes.
    ///
    /// When set, nodes that return `NodeResult::Error(e)` where
    /// `e.is_retryable()` are retried up to `max_attempts` times
    /// with backoff delays. Only the individual failed node is retried,
    /// not the entire superstep.
    /// Default: `None` (no retries — errors are immediate).
    #[serde(skip)]
    pub retry_policy: Option<RetryPolicy>,
}

impl Default for GraphConfig {
    fn default() -> Self {
        Self {
            thread_id: uuid::Uuid::new_v4().to_string(),
            checkpoint_id: None,
            recursion_limit: 25,
            max_concurrency: None,
            configurable: HashMap::new(),
            max_execution_time: None,
            retry_policy: None,
        }
    }
}

impl GraphConfig {
    /// Set the thread identifier for this execution.
    #[must_use = "builder methods return the modified config"]
    pub fn with_thread_id(mut self, id: impl Into<String>) -> Self {
        self.thread_id = id.into();
        self
    }

    /// Set the maximum number of supersteps before recursion error.
    #[must_use = "builder methods return the modified config"]
    pub fn with_recursion_limit(mut self, limit: u32) -> Self {
        self.recursion_limit = limit;
        self
    }

    /// Limit parallel node executions per superstep.
    #[must_use = "builder methods return the modified config"]
    pub fn with_max_concurrency(mut self, n: usize) -> Self {
        self.max_concurrency = Some(n);
        self
    }

    /// Resume from a specific checkpoint.
    #[must_use = "builder methods return the modified config"]
    pub fn with_checkpoint_id(mut self, id: impl Into<String>) -> Self {
        self.checkpoint_id = Some(id.into());
        self
    }

    /// Set the maximum wall-clock time for the entire execution.
    ///
    /// The engine checks this at the start of each superstep, before running
    /// any nodes. Running nodes are never cancelled mid-execution.
    #[must_use = "builder methods return the modified config"]
    pub fn with_max_execution_time(mut self, duration: Duration) -> Self {
        self.max_execution_time = Some(duration);
        self
    }

    /// Set a retry policy for failed nodes.
    ///
    /// When set, nodes returning retryable errors are retried individually
    /// with exponential backoff. Non-retryable errors fail immediately.
    ///
    /// # Example
    ///
    /// ```
    /// use pe_graph::{GraphConfig, RetryPolicy};
    ///
    /// let config = GraphConfig::default()
    ///     .with_retry_policy(RetryPolicy::default());
    /// ```
    #[must_use = "builder methods return the modified config"]
    pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
        self.retry_policy = Some(policy);
        self
    }
}

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

    #[test]
    fn test_default_values() {
        let config = GraphConfig::default();
        assert_eq!(config.recursion_limit, 25);
        assert!(config.max_concurrency.is_none());
        assert!(config.checkpoint_id.is_none());
        assert!(!config.thread_id.is_empty());
        assert!(config.configurable.is_empty());
        assert!(config.max_execution_time.is_none());
        assert!(config.retry_policy.is_none());
    }

    #[test]
    fn test_builder_methods() {
        let config = GraphConfig::default()
            .with_thread_id("test-thread")
            .with_recursion_limit(100)
            .with_max_concurrency(4)
            .with_checkpoint_id("cp-123");

        assert_eq!(config.thread_id, "test-thread");
        assert_eq!(config.recursion_limit, 100);
        assert_eq!(config.max_concurrency, Some(4));
        assert_eq!(config.checkpoint_id.as_deref(), Some("cp-123"));
    }

    #[test]
    fn test_with_retry_policy() {
        let policy = RetryPolicy {
            max_attempts: 5,
            ..Default::default()
        };
        let config = GraphConfig::default().with_retry_policy(policy);
        assert!(config.retry_policy.is_some());
        assert_eq!(config.retry_policy.unwrap().max_attempts, 5);
    }

    #[test]
    fn test_unique_thread_ids() {
        let a = GraphConfig::default();
        let b = GraphConfig::default();
        assert_ne!(a.thread_id, b.thread_id);
    }
}