klieo_flows/

graph.rs

//! DAG of flows with conditional edges. Walked iteratively (no recursion)
//! to avoid stack overflow on deep graphs.

use crate::error::FlowError;
use crate::flow::Flow;
use async_trait::async_trait;
use klieo_core::agent::AgentContext;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;

const DEFAULT_MAX_STEPS: u32 = 32;

/// Outgoing edge from a graph node.
pub enum Edge {
    /// Always route to `to` regardless of node output.
    Always {
        /// Target node name.
        to: String,
    },
    /// Branch on a predicate over the source node's output.
    Conditional {
        /// Predicate. Receives the source node's output.
        predicate: Arc<dyn Fn(&Value) -> bool + Send + Sync>,
        /// Target when predicate returns `true`.
        on_true: String,
        /// Target when predicate returns `false`.
        on_false: String,
    },
}

/// DAG of flows. Each node is itself a flow; edges route based on node
/// output. A node without an outgoing edge is terminal — its output is
/// the graph's output.
pub struct GraphFlow {
    name: String,
    entry: String,
    nodes: HashMap<String, Arc<dyn Flow>>,
    edges: HashMap<String, Edge>,
    max_steps: u32,
}

impl GraphFlow {
    /// New empty graph rooted at `entry`. The entry node must be added via
    /// [`GraphFlow::node`] before [`GraphFlow::run`] is called.
    pub fn new(name: impl Into<String>, entry: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            entry: entry.into(),
            nodes: HashMap::new(),
            edges: HashMap::new(),
            max_steps: DEFAULT_MAX_STEPS,
        }
    }

    /// Insert a node.
    pub fn node(mut self, name: impl Into<String>, flow: Arc<dyn Flow>) -> Self {
        self.nodes.insert(name.into(), flow);
        self
    }

    /// Convenience builder that wraps a concrete [`klieo_core::agent::Agent`]
    /// in [`crate::flow::AgentFlow`] + upcasts to `Arc<dyn Flow>` in one
    /// call.
    pub fn node_agent<A>(self, name: impl Into<String>, agent: A) -> Self
    where
        A: klieo_core::agent::Agent + 'static,
    {
        self.node(
            name,
            std::sync::Arc::new(crate::flow::AgentFlow::new(agent))
                as std::sync::Arc<dyn crate::flow::Flow>,
        )
    }

    /// Add an unconditional edge `from -> to`.
    pub fn edge_to(mut self, from: impl Into<String>, to: impl Into<String>) -> Self {
        self.edges
            .insert(from.into(), Edge::Always { to: to.into() });
        self
    }

    /// Add a conditional edge from `from` that branches on `predicate`.
    pub fn edge_branch<F>(
        mut self,
        from: impl Into<String>,
        predicate: F,
        on_true: impl Into<String>,
        on_false: impl Into<String>,
    ) -> Self
    where
        F: Fn(&Value) -> bool + Send + Sync + 'static,
    {
        self.edges.insert(
            from.into(),
            Edge::Conditional {
                predicate: Arc::new(predicate),
                on_true: on_true.into(),
                on_false: on_false.into(),
            },
        );
        self
    }

    /// Override the default step cap (32). Hitting it returns
    /// `FlowError::Graph("max_steps exceeded; possible cycle")`.
    pub fn with_max_steps(mut self, n: u32) -> Self {
        self.max_steps = n;
        self
    }

    /// Validate the graph at run-start: entry exists, every edge target
    /// exists. Cheap O(edges) check.
    fn validate(&self) -> Result<(), FlowError> {
        if !self.nodes.contains_key(&self.entry) {
            return Err(FlowError::Graph(format!("missing node: {}", self.entry)));
        }
        for (from, edge) in &self.edges {
            if !self.nodes.contains_key(from) {
                return Err(FlowError::Graph(format!("missing node: {from}")));
            }
            match edge {
                Edge::Always { to } => {
                    if !self.nodes.contains_key(to) {
                        return Err(FlowError::Graph(format!("missing node: {to}")));
                    }
                }
                Edge::Conditional {
                    on_true, on_false, ..
                } => {
                    if !self.nodes.contains_key(on_true) {
                        return Err(FlowError::Graph(format!("missing node: {on_true}")));
                    }
                    if !self.nodes.contains_key(on_false) {
                        return Err(FlowError::Graph(format!("missing node: {on_false}")));
                    }
                }
            }
        }
        Ok(())
    }
}

#[async_trait]
impl Flow for GraphFlow {
    fn name(&self) -> &str {
        &self.name
    }

    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        let span = tracing::info_span!(
            "graph_flow.run",
            flow = %self.name,
            entry = %self.entry
        );
        let _guard = span.enter();

        self.validate()?;

        let mut current_node = self.entry.clone();
        let mut current_input = input;

        // note: walked in a `for` loop (NOT async-recursion) so deep DAGs
        // can't overflow the stack. Hard requirement #6.
        for _ in 0..self.max_steps {
            crate::flow::bail_if_cancelled(&ctx)?;
            let flow = self
                .nodes
                .get(&current_node)
                .ok_or_else(|| FlowError::Graph(format!("missing node: {current_node}")))?;
            let output = flow.run(ctx.clone(), current_input).await?;

            let next = match self.edges.get(&current_node) {
                None => return Ok(output), // terminal
                Some(Edge::Always { to }) => to.clone(),
                Some(Edge::Conditional {
                    predicate,
                    on_true,
                    on_false,
                }) => {
                    if predicate(&output) {
                        on_true.clone()
                    } else {
                        on_false.clone()
                    }
                }
            };
            current_node = next;
            current_input = output;
        }

        Err(FlowError::Graph(
            "max_steps exceeded; possible cycle".into(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::{cancelled_ctx, ctx};

    struct Boom;
    #[async_trait]
    impl Flow for Boom {
        fn name(&self) -> &str {
            "boom"
        }
        async fn run(&self, _c: AgentContext, _i: Value) -> Result<Value, FlowError> {
            panic!("node must not run under a cancelled context")
        }
    }

    #[tokio::test]
    async fn cancelled_ctx_returns_cancelled_before_any_node() {
        let f = GraphFlow::new("g", "start").node("start", Arc::new(Boom));
        let err = f
            .run(cancelled_ctx(), serde_json::json!({}))
            .await
            .unwrap_err();
        assert!(matches!(err, FlowError::Cancelled), "got {err:?}");
    }

    struct ConstFlow {
        name: String,
        value: Value,
    }

    #[async_trait]
    impl Flow for ConstFlow {
        fn name(&self) -> &str {
            &self.name
        }
        async fn run(&self, _ctx: AgentContext, _input: Value) -> Result<Value, FlowError> {
            Ok(self.value.clone())
        }
    }

    struct IncrementFlow {
        name: String,
    }

    #[async_trait]
    impl Flow for IncrementFlow {
        fn name(&self) -> &str {
            &self.name
        }
        async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
            let n = input.get("n").and_then(|v| v.as_u64()).unwrap_or(0);
            Ok(serde_json::json!({ "n": n + 1 }))
        }
    }

    #[tokio::test]
    async fn linear_three_node_graph() {
        let g = GraphFlow::new("linear", "a")
            .node("a", Arc::new(IncrementFlow { name: "a".into() }))
            .node("b", Arc::new(IncrementFlow { name: "b".into() }))
            .node("c", Arc::new(IncrementFlow { name: "c".into() }))
            .edge_to("a", "b")
            .edge_to("b", "c");
        let out = g.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 3}));
    }

    #[tokio::test]
    async fn conditional_edge_routes_on_true() {
        let g = GraphFlow::new("cond_t", "start")
            .node(
                "start",
                Arc::new(ConstFlow {
                    name: "s".into(),
                    value: serde_json::json!({"big": true}),
                }),
            )
            .node(
                "yes",
                Arc::new(ConstFlow {
                    name: "y".into(),
                    value: serde_json::json!("yes"),
                }),
            )
            .node(
                "no",
                Arc::new(ConstFlow {
                    name: "n".into(),
                    value: serde_json::json!("no"),
                }),
            )
            .edge_branch(
                "start",
                |v| v.get("big").and_then(|b| b.as_bool()).unwrap_or(false),
                "yes",
                "no",
            );
        let out = g.run(ctx(), serde_json::json!(null)).await.unwrap();
        assert_eq!(out, serde_json::json!("yes"));
    }

    #[tokio::test]
    async fn conditional_edge_routes_on_false() {
        let g = GraphFlow::new("cond_f", "start")
            .node(
                "start",
                Arc::new(ConstFlow {
                    name: "s".into(),
                    value: serde_json::json!({"big": false}),
                }),
            )
            .node(
                "yes",
                Arc::new(ConstFlow {
                    name: "y".into(),
                    value: serde_json::json!("yes"),
                }),
            )
            .node(
                "no",
                Arc::new(ConstFlow {
                    name: "n".into(),
                    value: serde_json::json!("no"),
                }),
            )
            .edge_branch(
                "start",
                |v| v.get("big").and_then(|b| b.as_bool()).unwrap_or(false),
                "yes",
                "no",
            );
        let out = g.run(ctx(), serde_json::json!(null)).await.unwrap();
        assert_eq!(out, serde_json::json!("no"));
    }

    #[tokio::test]
    async fn cycle_detection_via_max_steps() {
        let g = GraphFlow::new("cycle", "a")
            .node("a", Arc::new(IncrementFlow { name: "a".into() }))
            .node("b", Arc::new(IncrementFlow { name: "b".into() }))
            .edge_to("a", "b")
            .edge_to("b", "a")
            .with_max_steps(4);
        let err = g.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Graph(s) => assert!(s.contains("max_steps")),
            other => panic!("expected Graph, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn missing_entry_node() {
        let g = GraphFlow::new("missing_entry", "absent")
            .node("a", Arc::new(IncrementFlow { name: "a".into() }));
        let err = g.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Graph(s) => assert!(s.contains("absent")),
            other => panic!("expected Graph, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn missing_edge_target() {
        let g = GraphFlow::new("dangling", "a")
            .node("a", Arc::new(IncrementFlow { name: "a".into() }))
            .edge_to("a", "ghost");
        let err = g.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Graph(s) => assert!(s.contains("ghost")),
            other => panic!("expected Graph, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn terminal_node_returns_its_output() {
        let g = GraphFlow::new("solo", "only").node(
            "only",
            Arc::new(ConstFlow {
                name: "only".into(),
                value: serde_json::json!(42),
            }),
        );
        let out = g.run(ctx(), serde_json::json!(null)).await.unwrap();
        assert_eq!(out, serde_json::json!(42));
    }
}