klieo_flows/
flow.rs

1//! Type-erased [`Flow`] trait + [`AgentFlow<A>`] adapter for [`klieo_core::Agent`].
2
3use crate::error::FlowError;
4use async_trait::async_trait;
5use klieo_core::agent::{Agent, AgentContext};
6use serde_json::Value;
7use std::sync::Arc;
8
9/// Stop a flow that has been cancelled before it starts the next unit of work.
10/// Every composition shape calls this at each iteration / step boundary so a
11/// cancelled run does no further LLM / IO / bus work (cooperative cancellation
12/// — the in-flight body, if any, is awaited to completion first).
13pub(crate) fn bail_if_cancelled(ctx: &AgentContext) -> Result<(), FlowError> {
14    if ctx.cancel.is_cancelled() {
15        return Err(FlowError::Cancelled);
16    }
17    Ok(())
18}
19
20/// Type-erased composition primitive. All composition shapes
21/// ([`crate::sequential::SequentialFlow`], [`crate::parallel::ParallelFlow`],
22/// [`crate::loop_flow::LoopFlow`], [`crate::graph::GraphFlow`]) implement this
23/// trait.
24///
25/// Unlike [`klieo_core::Agent`], `Flow` has no associated types — I/O is
26/// always `serde_json::Value`. This trades typing precision for the ability
27/// to compose flows behind `Arc<dyn Flow>`.
28#[async_trait]
29pub trait Flow: Send + Sync {
30    /// Stable flow name. Used in tracing spans.
31    fn name(&self) -> &str;
32
33    /// Execute the flow once.
34    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError>;
35}
36
37/// Adapter wrapping a concrete [`Agent`] into a [`Flow`] by JSON-marshalling
38/// at the I/O boundary.
39///
40/// This is parameterised on a concrete `A: Agent + 'static`, *not*
41/// `dyn Agent` — the latter is impossible because `Agent` has associated
42/// types. Each `AgentFlow<A>` is a distinct type, but each implements
43/// `Flow` so callers can hold them as `Arc<dyn Flow>`.
44pub struct AgentFlow<A: Agent> {
45    agent: Arc<A>,
46}
47
48impl<A: Agent + 'static> AgentFlow<A> {
49    /// Wrap an agent so it can participate in flows.
50    pub fn new(agent: A) -> Self {
51        Self {
52            agent: Arc::new(agent),
53        }
54    }
55
56    /// Wrap an `Arc<A>` so the same agent can be referenced from multiple
57    /// flows without cloning the underlying state.
58    pub fn from_arc(agent: Arc<A>) -> Self {
59        Self { agent }
60    }
61}
62
63#[async_trait]
64impl<A: Agent + 'static> Flow for AgentFlow<A> {
65    fn name(&self) -> &str {
66        self.agent.name()
67    }
68
69    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
70        let span = tracing::info_span!("agent_flow.run", agent = %self.agent.name());
71        let _guard = span.enter();
72        let typed_input: A::Input = serde_json::from_value(input)?;
73        let typed_output = self
74            .agent
75            .run(ctx, typed_input)
76            .await
77            .map_err(|e| FlowError::Agent(e.to_string()))?;
78        Ok(serde_json::to_value(typed_output)?)
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use crate::test_helpers::ctx;
86    use async_trait::async_trait;
87    use klieo_core::agent::{Agent, AgentContext};
88    use klieo_core::error::Error as CoreError;
89    use klieo_core::llm::ToolDef;
90    use serde::{Deserialize, Serialize};
91
92    #[derive(Deserialize, Serialize, PartialEq, Debug)]
93    struct EchoIn {
94        msg: String,
95    }
96    #[derive(Deserialize, Serialize, PartialEq, Debug)]
97    struct EchoOut {
98        echoed: String,
99    }
100
101    struct EchoAgent;
102
103    #[async_trait]
104    impl Agent for EchoAgent {
105        type Input = EchoIn;
106        type Output = EchoOut;
107        type Error = CoreError;
108        fn name(&self) -> &str {
109            "echo"
110        }
111        fn system_prompt(&self) -> &str {
112            ""
113        }
114        fn tools(&self) -> &[ToolDef] {
115            &[]
116        }
117        async fn run(&self, _ctx: AgentContext, input: EchoIn) -> Result<EchoOut, CoreError> {
118            Ok(EchoOut { echoed: input.msg })
119        }
120    }
121
122    struct FailingAgent;
123
124    #[async_trait]
125    impl Agent for FailingAgent {
126        type Input = EchoIn;
127        type Output = EchoOut;
128        type Error = CoreError;
129        fn name(&self) -> &str {
130            "failing"
131        }
132        fn system_prompt(&self) -> &str {
133            ""
134        }
135        fn tools(&self) -> &[ToolDef] {
136            &[]
137        }
138        async fn run(&self, _ctx: AgentContext, _input: EchoIn) -> Result<EchoOut, CoreError> {
139            Err(CoreError::Cancelled)
140        }
141    }
142
143    #[tokio::test]
144    async fn name_passes_through() {
145        let f = AgentFlow::new(EchoAgent);
146        assert_eq!(f.name(), "echo");
147    }
148
149    #[tokio::test]
150    async fn input_output_round_trip() {
151        let f = AgentFlow::new(EchoAgent);
152        let out = f
153            .run(ctx(), serde_json::json!({"msg": "hi"}))
154            .await
155            .unwrap();
156        assert_eq!(out, serde_json::json!({"echoed": "hi"}));
157    }
158
159    #[tokio::test]
160    async fn agent_error_is_stringified() {
161        let f = AgentFlow::new(FailingAgent);
162        let err = f
163            .run(ctx(), serde_json::json!({"msg": "x"}))
164            .await
165            .unwrap_err();
166        match err {
167            FlowError::Agent(s) => assert!(s.contains("cancelled")),
168            other => panic!("expected FlowError::Agent, got {other:?}"),
169        }
170    }
171}
klieo_flows/flow.rs

klieo_flows/
flow.rs
