klieo-flows 3.4.0

//! Type-erased [`Flow`] trait + [`AgentFlow<A>`] adapter for [`klieo_core::Agent`].

use crate::error::FlowError;
use async_trait::async_trait;
use klieo_core::agent::{Agent, AgentContext};
use serde_json::Value;
use std::sync::Arc;

/// Stop a flow that has been cancelled before it starts the next unit of work.
/// Every composition shape calls this at each iteration / step boundary so a
/// cancelled run does no further LLM / IO / bus work (cooperative cancellation
/// — the in-flight body, if any, is awaited to completion first).
pub(crate) fn bail_if_cancelled(ctx: &AgentContext) -> Result<(), FlowError> {
    if ctx.cancel.is_cancelled() {
        return Err(FlowError::Cancelled);
    }
    Ok(())
}

/// Type-erased composition primitive. All composition shapes
/// ([`crate::sequential::SequentialFlow`], [`crate::parallel::ParallelFlow`],
/// [`crate::loop_flow::LoopFlow`], [`crate::graph::GraphFlow`]) implement this
/// trait.
///
/// Unlike [`klieo_core::Agent`], `Flow` has no associated types — I/O is
/// always `serde_json::Value`. This trades typing precision for the ability
/// to compose flows behind `Arc<dyn Flow>`.
#[async_trait]
pub trait Flow: Send + Sync {
    /// Stable flow name. Used in tracing spans.
    fn name(&self) -> &str;

    /// Execute the flow once.
    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError>;
}

/// Adapter wrapping a concrete [`Agent`] into a [`Flow`] by JSON-marshalling
/// at the I/O boundary.
///
/// This is parameterised on a concrete `A: Agent + 'static`, *not*
/// `dyn Agent` — the latter is impossible because `Agent` has associated
/// types. Each `AgentFlow<A>` is a distinct type, but each implements
/// `Flow` so callers can hold them as `Arc<dyn Flow>`.
pub struct AgentFlow<A: Agent> {
    agent: Arc<A>,
}

impl<A: Agent + 'static> AgentFlow<A> {
    /// Wrap an agent so it can participate in flows.
    pub fn new(agent: A) -> Self {
        Self {
            agent: Arc::new(agent),
        }
    }

    /// Wrap an `Arc<A>` so the same agent can be referenced from multiple
    /// flows without cloning the underlying state.
    pub fn from_arc(agent: Arc<A>) -> Self {
        Self { agent }
    }
}

#[async_trait]
impl<A: Agent + 'static> Flow for AgentFlow<A> {
    fn name(&self) -> &str {
        self.agent.name()
    }

    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        let span = tracing::info_span!("agent_flow.run", agent = %self.agent.name());
        let _guard = span.enter();
        let typed_input: A::Input = serde_json::from_value(input)?;
        let typed_output = self
            .agent
            .run(ctx, typed_input)
            .await
            .map_err(|e| FlowError::Agent(e.to_string()))?;
        Ok(serde_json::to_value(typed_output)?)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::ctx;
    use async_trait::async_trait;
    use klieo_core::agent::{Agent, AgentContext};
    use klieo_core::error::Error as CoreError;
    use klieo_core::llm::ToolDef;
    use serde::{Deserialize, Serialize};

    #[derive(Deserialize, Serialize, PartialEq, Debug)]
    struct EchoIn {
        msg: String,
    }
    #[derive(Deserialize, Serialize, PartialEq, Debug)]
    struct EchoOut {
        echoed: String,
    }

    struct EchoAgent;

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

    struct FailingAgent;

    #[async_trait]
    impl Agent for FailingAgent {
        type Input = EchoIn;
        type Output = EchoOut;
        type Error = CoreError;
        fn name(&self) -> &str {
            "failing"
        }
        fn system_prompt(&self) -> &str {
            ""
        }
        fn tools(&self) -> &[ToolDef] {
            &[]
        }
        async fn run(&self, _ctx: AgentContext, _input: EchoIn) -> Result<EchoOut, CoreError> {
            Err(CoreError::Cancelled)
        }
    }

    #[tokio::test]
    async fn name_passes_through() {
        let f = AgentFlow::new(EchoAgent);
        assert_eq!(f.name(), "echo");
    }

    #[tokio::test]
    async fn input_output_round_trip() {
        let f = AgentFlow::new(EchoAgent);
        let out = f
            .run(ctx(), serde_json::json!({"msg": "hi"}))
            .await
            .unwrap();
        assert_eq!(out, serde_json::json!({"echoed": "hi"}));
    }

    #[tokio::test]
    async fn agent_error_is_stringified() {
        let f = AgentFlow::new(FailingAgent);
        let err = f
            .run(ctx(), serde_json::json!({"msg": "x"}))
            .await
            .unwrap_err();
        match err {
            FlowError::Agent(s) => assert!(s.contains("cancelled")),
            other => panic!("expected FlowError::Agent, got {other:?}"),
        }
    }
}