klieo_flows/

loop_flow.rs

//! Iterative flow. Repeats `body` until predicate returns `Done` or `max_iters` is hit.

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

/// Predicate verdict for [`LoopFlow`].
#[derive(Debug, Clone)]
pub enum LoopVerdict {
    /// Run another iteration.
    Continue,
    /// Terminate; return the most recent output.
    Done,
    /// Terminate with a [`FlowError::Loop`] containing this reason.
    Failed(String),
}

const DEFAULT_MAX_ITERS: u32 = 16;

/// Loop a single body flow until the predicate returns `Done`.
pub struct LoopFlow {
    name: String,
    body: Arc<dyn Flow>,
    predicate: Arc<dyn Fn(&Value) -> LoopVerdict + Send + Sync>,
    max_iters: u32,
}

impl LoopFlow {
    /// New loop. Default predicate returns `Done` immediately (single iter)
    /// so the caller is forced to call [`LoopFlow::until`].
    pub fn new(name: impl Into<String>, body: Arc<dyn Flow>) -> Self {
        Self {
            name: name.into(),
            body,
            predicate: Arc::new(|_| LoopVerdict::Done),
            max_iters: DEFAULT_MAX_ITERS,
        }
    }

    /// Replace the termination predicate.
    pub fn until<F>(mut self, predicate: F) -> Self
    where
        F: Fn(&Value) -> LoopVerdict + Send + Sync + 'static,
    {
        self.predicate = Arc::new(predicate);
        self
    }

    /// Override the default cap (16). The cap bounds the number of body
    /// invocations — the body runs at most `max_iters` times.
    pub fn with_max_iters(mut self, n: u32) -> Self {
        self.max_iters = n;
        self
    }

    /// Convenience constructor that wraps a concrete
    /// [`klieo_core::agent::Agent`] in [`crate::flow::AgentFlow`] and
    /// constructs a [`LoopFlow`] with the same defaults as [`LoopFlow::new`]
    /// (single-iteration predicate, max 16 iters).
    pub fn with_body_agent<A>(name: impl Into<String>, agent: A) -> Self
    where
        A: klieo_core::agent::Agent + 'static,
    {
        Self::new(
            name,
            std::sync::Arc::new(crate::flow::AgentFlow::new(agent))
                as std::sync::Arc<dyn crate::flow::Flow>,
        )
    }
}

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

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

        let mut current = input;
        for iter in 0..self.max_iters {
            crate::flow::bail_if_cancelled(&ctx)?;
            current = self.body.run(ctx.clone(), current).await?;
            match (self.predicate)(&current) {
                LoopVerdict::Continue => continue,
                LoopVerdict::Done => return Ok(current),
                LoopVerdict::Failed(reason) => {
                    return Err(FlowError::Loop {
                        reason,
                        iter: iter + 1,
                    });
                }
            }
        }
        Err(FlowError::Loop {
            reason: "max_iters exceeded".into(),
            iter: self.max_iters,
        })
    }
}

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

    struct IncrementFlow;

    #[async_trait]
    impl Flow for IncrementFlow {
        fn name(&self) -> &str {
            "inc"
        }
        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 predicate_done_immediately_runs_once() {
        let f = LoopFlow::new("once", Arc::new(IncrementFlow));
        let out = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 1}));
    }

    struct CountingFlow(std::sync::Arc<std::sync::atomic::AtomicU32>);

    #[async_trait]
    impl Flow for CountingFlow {
        fn name(&self) -> &str {
            "counting"
        }
        async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
            self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            Ok(input)
        }
    }

    #[tokio::test]
    async fn cancelled_ctx_returns_cancelled_without_running_body() {
        use crate::test_helpers::cancelled_ctx;
        let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let f = LoopFlow::new("c", Arc::new(CountingFlow(calls.clone())))
            .until(|_| LoopVerdict::Continue);
        let err = f
            .run(cancelled_ctx(), serde_json::json!({"n": 0}))
            .await
            .unwrap_err();
        assert!(matches!(err, FlowError::Cancelled), "got {err:?}");
        assert_eq!(
            calls.load(std::sync::atomic::Ordering::SeqCst),
            0,
            "body must not run under a cancelled context"
        );
    }

    #[tokio::test]
    async fn continues_then_done_after_three() {
        let f = LoopFlow::new("3x", Arc::new(IncrementFlow)).until(|v| {
            let n = v.get("n").and_then(|v| v.as_u64()).unwrap_or(0);
            if n >= 4 {
                LoopVerdict::Done
            } else {
                LoopVerdict::Continue
            }
        });
        let out = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 4}));
    }

    #[tokio::test]
    async fn predicate_failed_returns_loop_error() {
        let f = LoopFlow::new("fail", Arc::new(IncrementFlow))
            .until(|_| LoopVerdict::Failed("nope".into()));
        let err = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Loop { reason, iter } => {
                assert_eq!(reason, "nope");
                assert_eq!(iter, 1);
            }
            other => panic!("expected Loop, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn max_iters_exceeded() {
        let f = LoopFlow::new("inf", Arc::new(IncrementFlow))
            .until(|_| LoopVerdict::Continue)
            .with_max_iters(2);
        let err = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Loop { reason, iter } => {
                assert_eq!(reason, "max_iters exceeded");
                assert_eq!(iter, 2);
            }
            other => panic!("expected Loop, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn with_body_agent_runs_agent_body_once_by_default() {
        use async_trait::async_trait;
        use klieo_core::agent::{Agent, AgentContext as Ctx};
        use klieo_core::error::Error as CoreError;
        use klieo_core::llm::ToolDef;
        use serde::{Deserialize, Serialize};

        #[derive(Deserialize, Serialize)]
        struct In {
            n: u64,
        }
        #[derive(Deserialize, Serialize, PartialEq, Debug)]
        struct Out {
            n: u64,
        }

        struct IncrementAgent;

        #[async_trait]
        impl Agent for IncrementAgent {
            type Input = In;
            type Output = Out;
            type Error = CoreError;
            fn name(&self) -> &str {
                "inc"
            }
            fn system_prompt(&self) -> &str {
                ""
            }
            fn tools(&self) -> &[ToolDef] {
                &[]
            }
            async fn run(&self, _ctx: Ctx, input: In) -> Result<Out, CoreError> {
                Ok(Out { n: input.n + 1 })
            }
        }

        // Default predicate is Done-immediately → body runs exactly once.
        let f = LoopFlow::with_body_agent("inc-loop", IncrementAgent);
        let out = f.run(ctx(), serde_json::json!({"n": 5})).await.unwrap();
        assert_eq!(out["n"], 6, "agent body must execute once");
    }

    #[tokio::test]
    async fn with_body_agent_error_propagates_from_loop_run() {
        use async_trait::async_trait;
        use klieo_core::agent::{Agent, AgentContext as Ctx};
        use klieo_core::error::Error as CoreError;
        use klieo_core::llm::ToolDef;
        use serde_json::Value;

        struct AlwaysFailAgent;

        #[async_trait]
        impl Agent for AlwaysFailAgent {
            type Input = Value;
            type Output = Value;
            type Error = CoreError;
            fn name(&self) -> &str {
                "always_fail"
            }
            fn system_prompt(&self) -> &str {
                ""
            }
            fn tools(&self) -> &[ToolDef] {
                &[]
            }
            async fn run(&self, _ctx: Ctx, _input: Value) -> Result<Value, CoreError> {
                Err(CoreError::Other {
                    message: "injected agent failure".into(),
                    source: None,
                })
            }
        }

        let f = LoopFlow::with_body_agent("fail-loop", AlwaysFailAgent);
        let result = f.run(ctx(), serde_json::json!({})).await;
        assert!(
            result.is_err(),
            "with_body_agent error must propagate to LoopFlow::run"
        );
        assert!(
            matches!(result.unwrap_err(), FlowError::Agent(_)),
            "error must surface as FlowError::Agent"
        );
    }

    #[tokio::test]
    async fn output_threads_to_next_iteration() {
        let f = LoopFlow::new("thread", Arc::new(IncrementFlow)).until(|v| {
            let n = v.get("n").and_then(|v| v.as_u64()).unwrap_or(0);
            if n >= 5 {
                LoopVerdict::Done
            } else {
                LoopVerdict::Continue
            }
        });
        let out = f.run(ctx(), serde_json::json!({"n": 2})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 5}));
    }
}