klieo_flows/

parallel.rs

//! 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
    }

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

#[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();

        crate::flow::bail_if_cancelled(&ctx)?;
        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::{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!("branch must not run under a cancelled context")
        }
    }

    #[tokio::test]
    async fn cancelled_ctx_returns_cancelled_before_fan_out() {
        let f = ParallelFlow::new("p").branch("b", 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 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 branch_agent_output_appears_under_key() {
        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 {
            msg: String,
        }
        #[derive(Deserialize, Serialize, PartialEq, Debug)]
        struct Out {
            echoed: String,
        }

        struct EchoAgent;

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

        let f = ParallelFlow::new("p").branch_agent("echo_key", EchoAgent);
        let out = f
            .run(ctx(), serde_json::json!({"msg": "hello"}))
            .await
            .unwrap();
        assert_eq!(out["echo_key"]["echoed"], "hello");
    }

    #[tokio::test]
    async fn branch_agent_error_propagates_from_parallel_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 = ParallelFlow::new("p").branch_agent("failing_key", AlwaysFailAgent);
        let result = f.run(ctx(), serde_json::json!({})).await;
        assert!(
            result.is_err(),
            "branch_agent error must propagate to ParallelFlow::run"
        );
        assert!(
            matches!(result.unwrap_err(), FlowError::Agent(_)),
            "error must surface as FlowError::Agent"
        );
    }

    #[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}));
    }
}