Skip to main content

pe_graph/
lib.rs

1//! # pe-graph — Graph execution engine for Potential Expectations
2//!
3//! Implements the core graph primitives that agent topologies are built on:
4//!
5//! - [`StateGraph`] — declarative graph definition with typed nodes and edges
6//! - [`CompiledGraph`] — validated, executable graph with `invoke()` / `resume()`
7//! - [`GraphConfig`] — execution configuration (thread ID, recursion limit, etc.)
8//! - [`Checkpointer`] — trait for durable state persistence
9//! - [`GraphRegistry`] — named storage for compiled graphs
10//!
11//! The execution engine uses the **Pregel BSP model** (Bulk Synchronous Parallel):
12//! nodes execute in parallel supersteps with snapshot isolation, writes are
13//! collected and applied atomically between steps.
14//!
15//! Depends only on `pe-core` plus `tokio` and `futures` for async execution.
16
17mod activation;
18mod checkpoint_data;
19pub mod checkpointer;
20pub mod command;
21pub mod compiled;
22pub mod config;
23pub mod graph;
24pub mod matrix_hook;
25pub mod pending_writes;
26pub mod phase_store;
27mod pregel;
28pub mod registry;
29pub mod retry;
30pub mod snapshot;
31
32// Primary re-exports
33pub use checkpoint_data::CheckpointData;
34pub use checkpointer::{CheckpointMeta, Checkpointer, InMemoryCheckpointer, PendingWrite};
35pub use command::Command;
36pub use compiled::{CompiledGraph, ExecutionOutcome};
37pub use config::GraphConfig;
38pub use graph::StateGraph;
39pub use matrix_hook::{
40    ConvergenceRecorder, DefaultMatrixHook, MatrixHook, MatrixHookHandle, RoutingResolver,
41};
42pub use pending_writes::PendingWrites;
43pub use phase_store::{PhaseStateStore, PhaseStoreError};
44pub use registry::GraphRegistry;
45pub use retry::{RetryPolicy, with_retry};
46pub use snapshot::StateSnapshot;
47
48// Re-export START/END from pe-core for convenience
49pub use pe_core::types::{END, START};
50
51// ── Test support ──────────────────────────────────────────────────────
52// Shared test types used across multiple test modules in this crate.
53
54#[cfg(test)]
55#[allow(dead_code)]
56pub(crate) mod tests {
57    use pe_core::node::{NodeContext, NodeFn, NodeFuture, NodeResult};
58    use pe_core::state::{State, StateUpdate};
59    use serde::{Deserialize, Serialize};
60
61    // ── Test State ────────────────────────────────────────────────────
62
63    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
64    pub struct TestState {
65        pub messages: Vec<String>,
66        pub counter: u32,
67        pub thread_id: String,
68    }
69
70    #[derive(Debug, Clone, Serialize, Deserialize, Default)]
71    pub struct TestUpdate {
72        pub messages: Option<Vec<String>>,
73        pub counter: Option<u32>,
74    }
75
76    impl StateUpdate for TestUpdate {}
77
78    impl State for TestState {
79        type Update = TestUpdate;
80
81        fn apply(&mut self, update: TestUpdate) {
82            // messages: Appender semantics
83            if let Some(msgs) = update.messages {
84                self.messages.extend(msgs);
85            }
86            // counter: LastValue semantics
87            if let Some(c) = update.counter {
88                self.counter = c;
89            }
90        }
91    }
92
93    impl TestState {
94        pub fn new() -> Self {
95            Self {
96                messages: Vec::new(),
97                counter: 0,
98                thread_id: "test-thread".into(),
99            }
100        }
101    }
102
103    impl TestUpdate {
104        pub fn with_message(msg: impl Into<String>) -> Self {
105            Self {
106                messages: Some(vec![msg.into()]),
107                counter: None,
108            }
109        }
110
111        pub fn with_counter(n: u32) -> Self {
112            Self {
113                messages: None,
114                counter: Some(n),
115            }
116        }
117    }
118
119    // ── Test Nodes ────────────────────────────────────────────────────
120
121    /// Node that appends a message to state.
122    pub struct AppendNode {
123        node_name: &'static str,
124        message: &'static str,
125    }
126
127    impl AppendNode {
128        pub fn new(name: &'static str, message: &'static str) -> Self {
129            Self {
130                node_name: name,
131                message,
132            }
133        }
134    }
135
136    impl NodeFn<TestState> for AppendNode {
137        fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
138            let msg = self.message.to_string();
139            Box::pin(async move { NodeResult::Update(TestUpdate::with_message(msg)) })
140        }
141
142        fn name(&self) -> &str {
143            self.node_name
144        }
145    }
146}