klieo-flows 0.6.0

//! Ordered pipeline of flows; output of step N feeds input of step N+1.

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;

/// Ordered pipeline. Each step's output is threaded as the next step's
/// input. Empty pipelines return their input unchanged (identity).
/// First error short-circuits.
pub struct SequentialFlow {
    name: String,
    steps: Vec<Arc<dyn Flow>>,
}

impl SequentialFlow {
    /// New empty pipeline with the given name.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            steps: Vec::new(),
        }
    }

    /// Append a step. Builder-style: returns self.
    pub fn step(mut self, flow: Arc<dyn Flow>) -> Self {
        self.steps.push(flow);
        self
    }
}

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

    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        let span = tracing::info_span!("sequential_flow.run", flow = %self.name);
        let _guard = span.enter();
        let mut current = input;
        for (i, step) in self.steps.iter().enumerate() {
            tracing::debug!(step_index = i, step_name = step.name(), "sequential step");
            current = step.run(ctx.clone(), current).await?;
        }
        Ok(current)
    }
}

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

    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 }))
        }
    }

    struct FailFlow;

    #[async_trait]
    impl Flow for FailFlow {
        fn name(&self) -> &str {
            "fail"
        }
        async fn run(&self, _ctx: AgentContext, _input: Value) -> Result<Value, FlowError> {
            Err(FlowError::Agent("boom".into()))
        }
    }

    #[tokio::test]
    async fn empty_pipeline_returns_input_unchanged() {
        let f = SequentialFlow::new("empty");
        let out = f.run(ctx(), serde_json::json!({"n": 7})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 7}));
    }

    #[tokio::test]
    async fn three_steps_increment_in_order() {
        let f = SequentialFlow::new("incs")
            .step(Arc::new(IncrementFlow { name: "a".into() }))
            .step(Arc::new(IncrementFlow { name: "b".into() }))
            .step(Arc::new(IncrementFlow { name: "c".into() }));
        let out = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 3}));
    }

    #[tokio::test]
    async fn step_2_error_short_circuits() {
        let f = SequentialFlow::new("with_fail")
            .step(Arc::new(IncrementFlow { name: "a".into() }))
            .step(Arc::new(FailFlow))
            .step(Arc::new(IncrementFlow { name: "c".into() }));
        let err = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Agent(s) => assert_eq!(s, "boom"),
            other => panic!("expected Agent, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn name_returns_configured_name() {
        let f = SequentialFlow::new("my_pipe");
        assert_eq!(f.name(), "my_pipe");
    }
}