tirea_agentos/runtime/
behavior.rs

1use crate::contracts::runtime::behavior::{AgentBehavior, ReadOnlyContext};
2use crate::contracts::runtime::phase::{
3    ActionSet, AfterInferenceAction, AfterToolExecuteAction, BeforeInferenceAction,
4    BeforeToolExecuteAction, LifecycleAction,
5};
6use async_trait::async_trait;
7use std::sync::Arc;
8
9/// Compose multiple behaviors into a single [`AgentBehavior`].
10///
11/// If the list contains a single behavior, returns it directly.
12/// If it contains multiple, wraps them in a composite that concatenates
13/// their action lists in order.
14///
15/// This is the public API for behavior composition — callers never need
16/// to know about the concrete composite type.
17pub fn compose_behaviors(
18    id: impl Into<String>,
19    behaviors: Vec<Arc<dyn AgentBehavior>>,
20) -> Arc<dyn AgentBehavior> {
21    match behaviors.len() {
22        0 => Arc::new(crate::contracts::runtime::behavior::NoOpBehavior),
23        1 => behaviors.into_iter().next().unwrap(),
24        _ => Arc::new(CompositeBehavior::new(id, behaviors)),
25    }
26}
27
28/// An [`AgentBehavior`] that composes multiple sub-behaviors.
29///
30/// Each phase hook executes all sub-behaviors concurrently, merging their
31/// action sets in registration order. All sub-behaviors receive the same
32/// [`ReadOnlyContext`] snapshot — they do not see each other's effects
33/// within the same phase. The loop validates and applies all collected
34/// actions after the composite hook returns.
35pub(crate) struct CompositeBehavior {
36    id: String,
37    behaviors: Vec<Arc<dyn AgentBehavior>>,
38}
39
40impl CompositeBehavior {
41    pub(crate) fn new(id: impl Into<String>, behaviors: Vec<Arc<dyn AgentBehavior>>) -> Self {
42        Self {
43            id: id.into(),
44            behaviors,
45        }
46    }
47}
48
49#[async_trait]
50impl AgentBehavior for CompositeBehavior {
51    fn id(&self) -> &str {
52        &self.id
53    }
54
55    fn behavior_ids(&self) -> Vec<&str> {
56        self.behaviors
57            .iter()
58            .flat_map(|b| b.behavior_ids())
59            .collect()
60    }
61
62    async fn run_start(&self, ctx: &ReadOnlyContext<'_>) -> ActionSet<LifecycleAction> {
63        let futs: Vec<_> = self.behaviors.iter().map(|b| b.run_start(ctx)).collect();
64        futures::future::join_all(futs)
65            .await
66            .into_iter()
67            .fold(ActionSet::empty(), |acc, a| acc.and(a))
68    }
69
70    async fn step_start(&self, ctx: &ReadOnlyContext<'_>) -> ActionSet<LifecycleAction> {
71        let futs: Vec<_> = self.behaviors.iter().map(|b| b.step_start(ctx)).collect();
72        futures::future::join_all(futs)
73            .await
74            .into_iter()
75            .fold(ActionSet::empty(), |acc, a| acc.and(a))
76    }
77
78    async fn before_inference(
79        &self,
80        ctx: &ReadOnlyContext<'_>,
81    ) -> ActionSet<BeforeInferenceAction> {
82        let futs: Vec<_> = self
83            .behaviors
84            .iter()
85            .map(|b| b.before_inference(ctx))
86            .collect();
87        futures::future::join_all(futs)
88            .await
89            .into_iter()
90            .fold(ActionSet::empty(), |acc, a| acc.and(a))
91    }
92
93    async fn after_inference(&self, ctx: &ReadOnlyContext<'_>) -> ActionSet<AfterInferenceAction> {
94        let futs: Vec<_> = self
95            .behaviors
96            .iter()
97            .map(|b| b.after_inference(ctx))
98            .collect();
99        futures::future::join_all(futs)
100            .await
101            .into_iter()
102            .fold(ActionSet::empty(), |acc, a| acc.and(a))
103    }
104
105    async fn before_tool_execute(
106        &self,
107        ctx: &ReadOnlyContext<'_>,
108    ) -> ActionSet<BeforeToolExecuteAction> {
109        let futs: Vec<_> = self
110            .behaviors
111            .iter()
112            .map(|b| b.before_tool_execute(ctx))
113            .collect();
114        futures::future::join_all(futs)
115            .await
116            .into_iter()
117            .fold(ActionSet::empty(), |acc, a| acc.and(a))
118    }
119
120    async fn after_tool_execute(
121        &self,
122        ctx: &ReadOnlyContext<'_>,
123    ) -> ActionSet<AfterToolExecuteAction> {
124        let futs: Vec<_> = self
125            .behaviors
126            .iter()
127            .map(|b| b.after_tool_execute(ctx))
128            .collect();
129        futures::future::join_all(futs)
130            .await
131            .into_iter()
132            .fold(ActionSet::empty(), |acc, a| acc.and(a))
133    }
134
135    async fn step_end(&self, ctx: &ReadOnlyContext<'_>) -> ActionSet<LifecycleAction> {
136        let futs: Vec<_> = self.behaviors.iter().map(|b| b.step_end(ctx)).collect();
137        futures::future::join_all(futs)
138            .await
139            .into_iter()
140            .fold(ActionSet::empty(), |acc, a| acc.and(a))
141    }
142
143    async fn run_end(&self, ctx: &ReadOnlyContext<'_>) -> ActionSet<LifecycleAction> {
144        let futs: Vec<_> = self.behaviors.iter().map(|b| b.run_end(ctx)).collect();
145        futures::future::join_all(futs)
146            .await
147            .into_iter()
148            .fold(ActionSet::empty(), |acc, a| acc.and(a))
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::contracts::runtime::phase::BeforeInferenceAction;
156    use crate::contracts::runtime::phase::Phase;
157    use crate::contracts::RunPolicy;
158    use serde_json::json;
159    use tirea_state::DocCell;
160
161    struct ContextBehavior {
162        id: String,
163        text: String,
164    }
165
166    #[async_trait]
167    impl AgentBehavior for ContextBehavior {
168        fn id(&self) -> &str {
169            &self.id
170        }
171
172        async fn before_inference(
173            &self,
174            _ctx: &ReadOnlyContext<'_>,
175        ) -> ActionSet<BeforeInferenceAction> {
176            ActionSet::single(BeforeInferenceAction::AddSystemContext(self.text.clone()))
177        }
178    }
179
180    struct BlockBehavior;
181
182    #[async_trait]
183    impl AgentBehavior for BlockBehavior {
184        fn id(&self) -> &str {
185            "blocker"
186        }
187
188        async fn before_tool_execute(
189            &self,
190            ctx: &ReadOnlyContext<'_>,
191        ) -> ActionSet<BeforeToolExecuteAction> {
192            if ctx.tool_name() == Some("dangerous") {
193                ActionSet::single(BeforeToolExecuteAction::Block("denied".into()))
194            } else {
195                ActionSet::empty()
196            }
197        }
198    }
199
200    fn make_ctx<'a>(
201        doc: &'a DocCell,
202        run_policy: &'a RunPolicy,
203        phase: Phase,
204    ) -> ReadOnlyContext<'a> {
205        ReadOnlyContext::new(phase, "thread_1", &[], run_policy, doc)
206    }
207
208    #[tokio::test]
209    async fn composite_merges_actions() {
210        let behaviors: Vec<Arc<dyn AgentBehavior>> = vec![
211            Arc::new(ContextBehavior {
212                id: "a".into(),
213                text: "ctx_a".into(),
214            }),
215            Arc::new(ContextBehavior {
216                id: "b".into(),
217                text: "ctx_b".into(),
218            }),
219        ];
220        let composite = CompositeBehavior::new("test", behaviors);
221
222        let doc = DocCell::new(json!({}));
223        let run_policy = RunPolicy::new();
224        let ctx = make_ctx(&doc, &run_policy, Phase::BeforeInference);
225        let actions = composite.before_inference(&ctx).await;
226
227        assert_eq!(actions.len(), 2);
228        let v = actions.into_vec();
229        assert!(matches!(v[0], BeforeInferenceAction::AddSystemContext(_)));
230        assert!(matches!(v[1], BeforeInferenceAction::AddSystemContext(_)));
231    }
232
233    #[tokio::test]
234    async fn composite_empty_behaviors_returns_empty() {
235        let composite = CompositeBehavior::new("empty", vec![]);
236        let doc = DocCell::new(json!({}));
237        let run_policy = RunPolicy::new();
238        let ctx = make_ctx(&doc, &run_policy, Phase::BeforeInference);
239
240        let actions = composite.before_inference(&ctx).await;
241        assert!(actions.is_empty());
242    }
243
244    #[tokio::test]
245    async fn composite_preserves_action_order() {
246        let behaviors: Vec<Arc<dyn AgentBehavior>> = vec![
247            Arc::new(ContextBehavior {
248                id: "first".into(),
249                text: "1".into(),
250            }),
251            Arc::new(BlockBehavior),
252            Arc::new(ContextBehavior {
253                id: "last".into(),
254                text: "2".into(),
255            }),
256        ];
257        let composite = CompositeBehavior::new("order_test", behaviors);
258
259        let doc = DocCell::new(json!({}));
260        let run_policy = RunPolicy::new();
261        let ctx = make_ctx(&doc, &run_policy, Phase::BeforeInference);
262        let actions = composite.before_inference(&ctx).await;
263
264        // BlockBehavior returns empty for BeforeInference, so 2 actions
265        assert_eq!(actions.len(), 2);
266        let v = actions.into_vec();
267        assert!(matches!(v[0], BeforeInferenceAction::AddSystemContext(_)));
268        assert!(matches!(v[1], BeforeInferenceAction::AddSystemContext(_)));
269    }
270
271    #[test]
272    fn compose_behaviors_empty_returns_noop() {
273        let behavior = compose_behaviors("test", Vec::new());
274
275        assert_eq!(behavior.id(), "noop");
276        assert_eq!(behavior.behavior_ids(), vec!["noop"]);
277    }
278
279    #[test]
280    fn compose_behaviors_single_passthrough() {
281        let input = Arc::new(ContextBehavior {
282            id: "single".into(),
283            text: "ctx".into(),
284        }) as Arc<dyn AgentBehavior>;
285        let behavior = compose_behaviors("ignored", vec![input.clone()]);
286
287        assert!(Arc::ptr_eq(&behavior, &input));
288        assert_eq!(behavior.id(), "single");
289        assert_eq!(behavior.behavior_ids(), vec!["single"]);
290    }
291
292    #[test]
293    fn compose_behaviors_multiple_keeps_leaf_behavior_ids_order() {
294        let behavior = compose_behaviors(
295            "composed",
296            vec![
297                Arc::new(ContextBehavior {
298                    id: "a".into(),
299                    text: "ctx_a".into(),
300                }),
301                Arc::new(ContextBehavior {
302                    id: "b".into(),
303                    text: "ctx_b".into(),
304                }),
305            ],
306        );
307
308        assert_eq!(behavior.id(), "composed");
309        assert_eq!(behavior.behavior_ids(), vec!["a", "b"]);
310    }
311}
tirea_agentos/runtime/behavior.rs

tirea_agentos/runtime/
behavior.rs
