klieo-flows 0.40.0

//! Fan-out / fan-in flow. All branches see the same input concurrently.

use crate::error::FlowError;
use crate::flow::Flow;
use async_trait::async_trait;
use futures::future::try_join_all;
use klieo_core::agent::AgentContext;
use serde_json::{Map, Value};
use std::sync::Arc;

/// Fan-out the input to N branches, run them concurrently, return a JSON
/// object keyed by branch name. First error fails the whole flow (other
/// branches are dropped).
pub struct ParallelFlow {
    name: String,
    branches: Vec<(String, Arc<dyn Flow>)>,
}

impl ParallelFlow {
    /// New empty parallel flow.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            branches: Vec::new(),
        }
    }

    /// Append a branch. The `key` becomes the JSON key in the output object.
    /// Duplicate keys overwrite — caller must ensure uniqueness.
    pub fn branch(mut self, key: impl Into<String>, flow: Arc<dyn Flow>) -> Self {
        self.branches.push((key.into(), flow));
        self
    }
}

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

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

        let futs = self.branches.iter().map(|(key, flow)| {
            let ctx = ctx.clone();
            let input = input.clone();
            let key = key.clone();
            let flow = flow.clone();
            async move {
                let out = flow.run(ctx, input).await?;
                Ok::<(String, Value), FlowError>((key, out))
            }
        });

        let results = try_join_all(futs).await?;
        let mut obj = Map::with_capacity(results.len());
        for (k, v) in results {
            obj.insert(k, v);
        }
        Ok(Value::Object(obj))
    }
}

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

    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 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 three_branches_collated_by_key() {
        let f = ParallelFlow::new("fan")
            .branch(
                "a",
                Arc::new(ConstFlow {
                    name: "a".into(),
                    value: serde_json::json!(1),
                }),
            )
            .branch(
                "b",
                Arc::new(ConstFlow {
                    name: "b".into(),
                    value: serde_json::json!("two"),
                }),
            )
            .branch(
                "c",
                Arc::new(ConstFlow {
                    name: "c".into(),
                    value: serde_json::json!(true),
                }),
            );
        let out = f.run(ctx(), serde_json::json!(null)).await.unwrap();
        assert_eq!(out, serde_json::json!({"a": 1, "b": "two", "c": true}));
    }

    #[tokio::test]
    async fn one_branch_error_fails_whole_flow() {
        let f = ParallelFlow::new("p")
            .branch(
                "ok",
                Arc::new(ConstFlow {
                    name: "ok".into(),
                    value: serde_json::json!(1),
                }),
            )
            .branch("bad", Arc::new(FailFlow));
        let err = f.run(ctx(), serde_json::json!(null)).await.unwrap_err();
        assert!(matches!(err, FlowError::Agent(s) if s == "boom"));
    }

    #[tokio::test]
    async fn empty_branch_list_returns_empty_object() {
        let f = ParallelFlow::new("empty");
        let out = f.run(ctx(), serde_json::json!({"x": 1})).await.unwrap();
        assert_eq!(out, serde_json::json!({}));
    }

    #[tokio::test]
    async fn weird_keys_preserved_verbatim() {
        let f = ParallelFlow::new("weird")
            .branch(
                "dot.key",
                Arc::new(ConstFlow {
                    name: "n1".into(),
                    value: serde_json::json!(1),
                }),
            )
            .branch(
                "dash-key",
                Arc::new(ConstFlow {
                    name: "n2".into(),
                    value: serde_json::json!(2),
                }),
            );
        let out = f.run(ctx(), serde_json::json!(null)).await.unwrap();
        assert_eq!(out, serde_json::json!({"dot.key": 1, "dash-key": 2}));
    }
}