autoagents_core/agent/
direct.rs

1use crate::agent::base::AgentType;
2use crate::agent::context::Context;
3use crate::agent::error::{AgentBuildError, RunnableAgentError};
4use crate::agent::executor::event_helper::EventHelper;
5use crate::agent::task::Task;
6use crate::agent::{AgentBuilder, AgentDeriveT, AgentExecutor, AgentHooks, BaseAgent, HookOutcome};
7use crate::error::Error;
8use autoagents_protocol::Event;
9use serde_json::Value;
10use std::sync::Arc;
11
12use crate::agent::constants::DEFAULT_CHANNEL_BUFFER;
13
14use crate::channel::{Receiver, Sender, channel};
15
16#[cfg(not(target_arch = "wasm32"))]
17use crate::event_fanout::EventFanout;
18use crate::utils::{BoxEventStream, receiver_into_stream};
19#[cfg(not(target_arch = "wasm32"))]
20use futures_util::stream;
21
22/// Marker type for direct (non-actor) agents.
23///
24/// Direct agents execute immediately within the caller's task without
25/// requiring a runtime or event wiring. Use this for simple one-shot
26/// invocations and unit tests.
27#[derive(Clone, Copy)]
28pub struct DirectAgent {}
29
30impl AgentType for DirectAgent {
31    fn type_name() -> &'static str {
32        "direct_agent"
33    }
34}
35
36/// Handle for a direct agent containing the agent instance and an event stream
37/// receiver. Use `agent.run(...)` for one-shot calls or `agent.run_stream(...)`
38/// to receive streaming outputs.
39///
40/// Terminal outcomes emit protocol events on [`Self::rx`]: `TaskComplete` on success
41/// and `TaskError` on failure (hook abort, executor error, stream setup error,
42/// in-stream item errors, and empty streams). For `run_stream()`, `TaskComplete`
43/// is emitted when the returned output stream is fully drained; the last successful
44/// item is used for the event payload.
45pub struct DirectAgentHandle<T: AgentDeriveT + AgentExecutor + AgentHooks + Send + Sync> {
46    pub agent: BaseAgent<T, DirectAgent>,
47    pub rx: BoxEventStream<Event>,
48    #[cfg(not(target_arch = "wasm32"))]
49    fanout: Option<EventFanout>,
50}
51
52impl<T: AgentDeriveT + AgentExecutor + AgentHooks> DirectAgentHandle<T> {
53    pub fn new(agent: BaseAgent<T, DirectAgent>, rx: BoxEventStream<Event>) -> Self {
54        Self {
55            agent,
56            rx,
57            #[cfg(not(target_arch = "wasm32"))]
58            fanout: None,
59        }
60    }
61
62    #[cfg(not(target_arch = "wasm32"))]
63    pub fn subscribe_events(&mut self) -> BoxEventStream<Event> {
64        if let Some(fanout) = &self.fanout {
65            return fanout.subscribe();
66        }
67
68        let stream = std::mem::replace(&mut self.rx, Box::pin(stream::empty::<Event>()));
69        let fanout = EventFanout::new(stream, DEFAULT_CHANNEL_BUFFER);
70        self.rx = fanout.subscribe();
71        let stream = fanout.subscribe();
72        self.fanout = Some(fanout);
73        stream
74    }
75}
76
77impl<T: AgentDeriveT + AgentExecutor + AgentHooks> AgentBuilder<T, DirectAgent> {
78    /// Build the BaseAgent and return a wrapper
79    #[allow(clippy::result_large_err)]
80    pub async fn build(self) -> Result<DirectAgentHandle<T>, Error> {
81        let llm = self.llm.ok_or(AgentBuildError::BuildFailure(
82            "LLM provider is required".to_string(),
83        ))?;
84        let (tx, rx): (Sender<Event>, Receiver<Event>) = channel(DEFAULT_CHANNEL_BUFFER);
85        let agent: BaseAgent<T, DirectAgent> =
86            BaseAgent::<T, DirectAgent>::new(self.inner, llm, self.memory, tx, self.stream).await?;
87        let stream = receiver_into_stream(rx);
88        Ok(DirectAgentHandle::new(agent, stream))
89    }
90}
91
92fn wrap_direct_stream_with_terminal_events<T>(
93    agent: BaseAgent<T, DirectAgent>,
94    stream: crate::utils::BoxRuntimeStream<
95        Result<<T as AgentExecutor>::Output, <T as AgentExecutor>::Error>,
96    >,
97    task: Task,
98    context: Arc<Context>,
99    tx_event: Option<crate::channel::Sender<Event>>,
100    submission_id: autoagents_protocol::SubmissionId,
101    actor_id: autoagents_protocol::ActorID,
102) -> crate::utils::BoxRuntimeStream<Result<<T as AgentDeriveT>::Output, Error>>
103where
104    T: AgentDeriveT + AgentExecutor + AgentHooks + Send + Sync,
105    Value: From<<T as AgentExecutor>::Output>,
106    <T as AgentDeriveT>::Output: From<<T as AgentExecutor>::Output>,
107    <T as AgentExecutor>::Output: Clone,
108    <T as AgentExecutor>::Error: Into<RunnableAgentError>,
109{
110    use futures::StreamExt;
111
112    Box::pin(futures::stream::unfold(
113        (
114            stream,
115            false,
116            None::<<T as AgentExecutor>::Output>,
117            task,
118            context,
119        ),
120        move |(mut stream, mut saw_error, last, task, context)| {
121            let agent = agent.clone_shallow();
122            let tx_event = tx_event.clone();
123            async move {
124                match stream.next().await {
125                    Some(result) => {
126                        match EventHelper::map_executor_stream_item(
127                            &tx_event,
128                            submission_id,
129                            actor_id,
130                            result,
131                        )
132                        .await
133                        {
134                            Ok(output) => {
135                                let agent_out: <T as AgentDeriveT>::Output = output.clone().into();
136                                Some((
137                                    Ok(agent_out),
138                                    (stream, saw_error, Some(output), task, context),
139                                ))
140                            }
141                            Err(err) => {
142                                saw_error = true;
143                                Some((
144                                    Err(Error::from(err)),
145                                    (stream, saw_error, last, task, context),
146                                ))
147                            }
148                        }
149                    }
150                    None => {
151                        if !saw_error {
152                            if let Some(executor_out) = last {
153                                let _ = agent
154                                    .finish_executor_run(
155                                        &task,
156                                        context.as_ref(),
157                                        submission_id,
158                                        executor_out,
159                                    )
160                                    .await;
161                            } else {
162                                let err = RunnableAgentError::ExecutorError(
163                                    "Stream completed without output".to_string(),
164                                );
165                                #[cfg(not(target_arch = "wasm32"))]
166                                EventHelper::send_task_error(
167                                    &tx_event,
168                                    submission_id,
169                                    actor_id,
170                                    err.to_string(),
171                                )
172                                .await;
173                            }
174                        }
175                        None
176                    }
177                }
178            }
179        },
180    ))
181}
182
183impl<T: AgentDeriveT + AgentExecutor + AgentHooks> BaseAgent<T, DirectAgent> {
184    /// Execute the agent for a single task and return the final agent output.
185    pub async fn run(&self, task: Task) -> Result<<T as AgentDeriveT>::Output, RunnableAgentError>
186    where
187        Value: From<<T as AgentExecutor>::Output>,
188        <T as AgentDeriveT>::Output: From<<T as AgentExecutor>::Output>,
189        <T as AgentExecutor>::Output: Clone,
190        <T as AgentExecutor>::Error: Into<RunnableAgentError>,
191    {
192        let submission_id = task.submission_id;
193        let tx_event = self.tx.clone();
194        let context = self.create_context();
195
196        //Run Hook
197        let hook_outcome = self.inner.on_run_start(&task, &context).await;
198        match hook_outcome {
199            HookOutcome::Abort => {
200                return Err(
201                    EventHelper::abort_run_from_hook(&tx_event, submission_id, self.id).await,
202                );
203            }
204            HookOutcome::Continue => {}
205        }
206
207        // Execute the agent's logic using the executor
208        match self.inner().execute(&task, context.clone()).await {
209            Ok(output) => {
210                self.finish_executor_run(&task, &context, submission_id, output)
211                    .await
212            }
213            Err(e) => {
214                let err: RunnableAgentError = e.into();
215                #[cfg(not(target_arch = "wasm32"))]
216                EventHelper::send_task_error(&tx_event, submission_id, self.id, err.to_string())
217                    .await;
218                Err(err)
219            }
220        }
221    }
222
223    /// Execute the agent with streaming enabled and receive a stream of
224    /// partial outputs which culminate in a final chunk with `done=true`.
225    pub async fn run_stream(
226        &self,
227        task: Task,
228    ) -> Result<
229        crate::utils::BoxRuntimeStream<Result<<T as AgentDeriveT>::Output, Error>>,
230        RunnableAgentError,
231    >
232    where
233        Value: From<<T as AgentExecutor>::Output>,
234        <T as AgentDeriveT>::Output: From<<T as AgentExecutor>::Output>,
235        <T as AgentExecutor>::Output: Clone,
236        <T as AgentExecutor>::Error: Into<RunnableAgentError>,
237    {
238        let submission_id = task.submission_id;
239        let tx_event = self.tx.clone();
240        let context = self.create_context();
241
242        //Run Hook
243        let hook_outcome = self.inner.on_run_start(&task, &context).await;
244        match hook_outcome {
245            HookOutcome::Abort => {
246                return Err(
247                    EventHelper::abort_run_from_hook(&tx_event, submission_id, self.id).await,
248                );
249            }
250            HookOutcome::Continue => {}
251        }
252
253        // Execute the agent's streaming logic using the executor
254        match self.inner().execute_stream(&task, context.clone()).await {
255            Ok(stream) => Ok(wrap_direct_stream_with_terminal_events(
256                self.clone_shallow(),
257                stream,
258                task,
259                context,
260                tx_event,
261                submission_id,
262                self.id,
263            )),
264            Err(e) => {
265                let err: RunnableAgentError = e.into();
266                #[cfg(not(target_arch = "wasm32"))]
267                EventHelper::send_task_error(&tx_event, submission_id, self.id, err.to_string())
268                    .await;
269                Err(err)
270            }
271        }
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278    use crate::agent::hooks::HookOutcome;
279    use crate::agent::output::AgentOutputT;
280    use crate::agent::prebuilt::executor::{
281        BasicAgent as StableBasicAgent, BasicAgentOutput, ReActAgent as StableReActAgent,
282        ReActAgentOutput,
283    };
284    use crate::agent::task::Task;
285    use crate::agent::{Context, ExecutorConfig};
286    use crate::tests::{
287        ConfigurableLLMProvider, MockAgentImpl, MultiItemStreamAgent, TestAgentOutput, TestError,
288    };
289    use crate::tool::ToolT;
290    use async_trait::async_trait;
291    use futures::StreamExt;
292    use serde::{Deserialize, Serialize};
293    use serde_json::Value;
294    use std::sync::{
295        Arc,
296        atomic::{AtomicBool, AtomicUsize, Ordering},
297    };
298
299    #[tokio::test]
300    async fn test_direct_agent_build_requires_llm() {
301        let mock_agent = MockAgentImpl::new("direct", "direct agent");
302        let err = match AgentBuilder::<_, DirectAgent>::new(mock_agent)
303            .build()
304            .await
305        {
306            Ok(_) => panic!("expected missing llm error"),
307            Err(err) => err,
308        };
309
310        assert!(matches!(err, crate::error::Error::AgentBuildError(_)));
311    }
312
313    #[tokio::test]
314    async fn test_direct_agent_run_success() {
315        let mock_agent = MockAgentImpl::new("direct", "direct agent");
316        let llm = Arc::new(ConfigurableLLMProvider::default());
317        let mut handle = AgentBuilder::<_, DirectAgent>::new(mock_agent)
318            .llm(llm)
319            .build()
320            .await
321            .expect("build should succeed");
322
323        let task = Task::new("hello");
324        let result = handle.agent.run(task).await.expect("run should succeed");
325        assert_eq!(result.result, "Processed: hello");
326
327        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
328            .await
329            .expect("timed out waiting for TaskComplete event")
330            .expect("stream ended without event");
331        match event {
332            Event::TaskComplete { result, .. } => {
333                let parsed: Value =
334                    serde_json::from_str(&result).expect("TaskComplete result should be JSON");
335                assert_eq!(parsed["result"], "Processed: hello");
336            }
337            other => panic!("expected TaskComplete, got {other:?}"),
338        }
339    }
340
341    #[tokio::test]
342    async fn test_direct_agent_run_executor_error() {
343        let mock_agent = MockAgentImpl::new("direct", "direct agent").with_failure(true);
344        let llm = Arc::new(ConfigurableLLMProvider::default());
345        let mut handle = AgentBuilder::<_, DirectAgent>::new(mock_agent)
346            .llm(llm)
347            .build()
348            .await
349            .expect("build should succeed");
350
351        let task = Task::new("fail");
352        let err = handle.agent.run(task).await.expect_err("expected error");
353        assert!(matches!(err, RunnableAgentError::ExecutorError(_)));
354
355        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
356            .await
357            .expect("timed out waiting for TaskError event")
358            .expect("stream ended without event");
359        match event {
360            Event::TaskError { error, .. } => {
361                assert!(error.contains("Mock execution failed"));
362            }
363            other => panic!("expected TaskError, got {other:?}"),
364        }
365    }
366
367    #[derive(Debug, Clone, Serialize, Deserialize)]
368    struct HookCountOutput {
369        result: String,
370    }
371
372    impl AgentOutputT for HookCountOutput {
373        fn output_schema() -> &'static str {
374            r#"{"type":"object","properties":{"result":{"type":"string"}},"required":["result"]}"#
375        }
376
377        fn structured_output_format() -> Value {
378            serde_json::json!({
379                "name": "HookCountOutput",
380                "description": "Hook count output",
381                "schema": {
382                    "type": "object",
383                    "properties": {
384                        "result": {"type": "string"}
385                    },
386                    "required": ["result"]
387                },
388                "strict": true
389            })
390        }
391    }
392
393    impl From<BasicAgentOutput> for HookCountOutput {
394        fn from(output: BasicAgentOutput) -> Self {
395            Self {
396                result: output.response,
397            }
398        }
399    }
400
401    impl From<ReActAgentOutput> for HookCountOutput {
402        fn from(output: ReActAgentOutput) -> Self {
403            Self {
404                result: output.response,
405            }
406        }
407    }
408
409    #[derive(Debug, Clone)]
410    struct CountingHookAgent {
411        on_run_start_calls: Arc<AtomicUsize>,
412    }
413
414    #[async_trait]
415    impl AgentDeriveT for CountingHookAgent {
416        type Output = HookCountOutput;
417
418        fn description(&self) -> &'static str {
419            "counting hook agent"
420        }
421
422        fn output_schema(&self) -> Option<Value> {
423            Some(serde_json::json!({
424                "type": "object",
425                "properties": {"result": {"type": "string"}},
426                "required": ["result"]
427            }))
428        }
429
430        fn name(&self) -> &'static str {
431            "counting_hook_agent"
432        }
433
434        fn tools(&self) -> Vec<Box<dyn ToolT>> {
435            vec![]
436        }
437    }
438
439    #[async_trait]
440    impl AgentHooks for CountingHookAgent {
441        async fn on_run_start(&self, _task: &Task, _ctx: &Context) -> HookOutcome {
442            self.on_run_start_calls.fetch_add(1, Ordering::SeqCst);
443            HookOutcome::Continue
444        }
445    }
446
447    #[tokio::test]
448    async fn test_direct_basic_agent_run_calls_on_run_start_once() {
449        let calls = Arc::new(AtomicUsize::new(0));
450        let llm = Arc::new(ConfigurableLLMProvider::default());
451        let handle =
452            AgentBuilder::<_, DirectAgent>::new(StableBasicAgent::new(CountingHookAgent {
453                on_run_start_calls: Arc::clone(&calls),
454            }))
455            .llm(llm)
456            .build()
457            .await
458            .expect("build should succeed");
459
460        let task = Task::new("hello");
461        let result = handle.agent.run(task).await.expect("run should succeed");
462
463        assert_eq!(result.result, "Mock response");
464        assert_eq!(calls.load(Ordering::SeqCst), 1);
465    }
466
467    #[tokio::test]
468    async fn test_direct_react_agent_run_calls_on_run_start_once() {
469        let calls = Arc::new(AtomicUsize::new(0));
470        let llm = Arc::new(ConfigurableLLMProvider::default());
471        let handle =
472            AgentBuilder::<_, DirectAgent>::new(StableReActAgent::new(CountingHookAgent {
473                on_run_start_calls: Arc::clone(&calls),
474            }))
475            .llm(llm)
476            .build()
477            .await
478            .expect("build should succeed");
479
480        let task = Task::new("hello");
481        let result = handle.agent.run(task).await.expect("run should succeed");
482
483        assert_eq!(result.result, "Mock response");
484        assert_eq!(calls.load(Ordering::SeqCst), 1);
485    }
486
487    #[derive(Clone, Debug)]
488    struct StreamAgent;
489
490    #[async_trait]
491    impl AgentDeriveT for StreamAgent {
492        type Output = TestAgentOutput;
493
494        fn description(&self) -> &'static str {
495            "stream agent"
496        }
497
498        fn output_schema(&self) -> Option<Value> {
499            Some(TestAgentOutput::structured_output_format())
500        }
501
502        fn name(&self) -> &'static str {
503            "stream_agent"
504        }
505
506        fn tools(&self) -> Vec<Box<dyn ToolT>> {
507            vec![]
508        }
509    }
510
511    #[async_trait]
512    impl AgentExecutor for StreamAgent {
513        type Output = TestAgentOutput;
514        type Error = TestError;
515
516        fn config(&self) -> ExecutorConfig {
517            ExecutorConfig::default()
518        }
519
520        async fn execute(
521            &self,
522            task: &Task,
523            _context: Arc<Context>,
524        ) -> Result<Self::Output, Self::Error> {
525            Ok(TestAgentOutput {
526                result: format!("Streamed: {}", task.prompt),
527            })
528        }
529    }
530
531    impl AgentHooks for StreamAgent {}
532
533    #[tokio::test]
534    async fn test_direct_agent_run_stream_default_executes_once() {
535        let llm = Arc::new(ConfigurableLLMProvider::default());
536        let mut handle = AgentBuilder::<_, DirectAgent>::new(StreamAgent)
537            .llm(llm)
538            .build()
539            .await
540            .expect("build should succeed");
541
542        let task = Task::new("stream");
543        let stream = handle
544            .agent
545            .run_stream(task)
546            .await
547            .expect("stream should succeed");
548        let outputs: Vec<_> = stream.collect().await;
549        assert_eq!(outputs.len(), 1);
550        let output = outputs[0].as_ref().expect("expected Ok output");
551        assert_eq!(output.result, "Streamed: stream");
552
553        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
554            .await
555            .expect("timed out waiting for TaskComplete event")
556            .expect("stream ended without event");
557        match event {
558            Event::TaskComplete { result, .. } => {
559                let parsed: Value =
560                    serde_json::from_str(&result).expect("TaskComplete result should be JSON");
561                assert_eq!(parsed["result"], "Streamed: stream");
562            }
563            other => panic!("expected TaskComplete, got {other:?}"),
564        }
565    }
566
567    #[derive(Debug)]
568    struct AbortAgent {
569        executed: Arc<AtomicBool>,
570    }
571
572    #[async_trait]
573    impl AgentDeriveT for AbortAgent {
574        type Output = TestAgentOutput;
575
576        fn description(&self) -> &'static str {
577            "abort agent"
578        }
579
580        fn output_schema(&self) -> Option<Value> {
581            Some(TestAgentOutput::structured_output_format())
582        }
583
584        fn name(&self) -> &'static str {
585            "abort_agent"
586        }
587
588        fn tools(&self) -> Vec<Box<dyn ToolT>> {
589            vec![]
590        }
591    }
592
593    #[async_trait]
594    impl AgentExecutor for AbortAgent {
595        type Output = TestAgentOutput;
596        type Error = TestError;
597
598        fn config(&self) -> ExecutorConfig {
599            ExecutorConfig::default()
600        }
601
602        async fn execute(
603            &self,
604            _task: &Task,
605            _context: Arc<Context>,
606        ) -> Result<Self::Output, Self::Error> {
607            self.executed.store(true, Ordering::SeqCst);
608            Ok(TestAgentOutput {
609                result: "should-not-run".to_string(),
610            })
611        }
612    }
613
614    #[async_trait]
615    impl AgentHooks for AbortAgent {
616        async fn on_run_start(&self, _task: &Task, _ctx: &Context) -> HookOutcome {
617            HookOutcome::Abort
618        }
619    }
620
621    #[tokio::test]
622    async fn test_direct_agent_run_aborts_before_execute() {
623        let executed = Arc::new(AtomicBool::new(false));
624        let agent = AbortAgent {
625            executed: Arc::clone(&executed),
626        };
627        let llm = Arc::new(ConfigurableLLMProvider::default());
628        let mut handle = AgentBuilder::<_, DirectAgent>::new(agent)
629            .llm(llm)
630            .build()
631            .await
632            .expect("build should succeed");
633
634        let task = Task::new("abort");
635        let err = handle.agent.run(task).await.expect_err("expected abort");
636        assert!(matches!(err, RunnableAgentError::Abort));
637        assert!(!executed.load(Ordering::SeqCst));
638
639        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
640            .await
641            .expect("timed out waiting for TaskError event")
642            .expect("stream ended without event");
643        match event {
644            Event::TaskError { error, .. } => {
645                assert!(error.contains("Abort"));
646            }
647            other => panic!("expected TaskError, got {other:?}"),
648        }
649    }
650
651    #[tokio::test]
652    async fn test_direct_agent_run_stream_aborts_before_execute_stream() {
653        let executed = Arc::new(AtomicBool::new(false));
654        let agent = AbortAgent {
655            executed: Arc::clone(&executed),
656        };
657        let llm = Arc::new(ConfigurableLLMProvider::default());
658        let mut handle = AgentBuilder::<_, DirectAgent>::new(agent)
659            .llm(llm)
660            .stream(true)
661            .build()
662            .await
663            .expect("build should succeed");
664
665        let task = Task::new("abort");
666        let err = match handle.agent.run_stream(task).await {
667            Err(err) => err,
668            Ok(_) => panic!("expected abort"),
669        };
670        assert!(matches!(err, RunnableAgentError::Abort));
671        assert!(!executed.load(Ordering::SeqCst));
672
673        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
674            .await
675            .expect("timed out waiting for TaskError event")
676            .expect("stream ended without event");
677        match event {
678            Event::TaskError { error, .. } => {
679                assert!(error.contains("Abort"));
680            }
681            other => panic!("expected TaskError, got {other:?}"),
682        }
683    }
684
685    #[derive(Debug, Clone)]
686    struct FailingStreamSetupAgent;
687
688    #[async_trait]
689    impl AgentDeriveT for FailingStreamSetupAgent {
690        type Output = TestAgentOutput;
691
692        fn description(&self) -> &'static str {
693            "failing stream setup"
694        }
695
696        fn output_schema(&self) -> Option<Value> {
697            Some(TestAgentOutput::structured_output_format())
698        }
699
700        fn name(&self) -> &'static str {
701            "failing_stream_setup"
702        }
703
704        fn tools(&self) -> Vec<Box<dyn ToolT>> {
705            vec![]
706        }
707    }
708
709    #[async_trait]
710    impl AgentExecutor for FailingStreamSetupAgent {
711        type Output = TestAgentOutput;
712        type Error = TestError;
713
714        fn config(&self) -> ExecutorConfig {
715            ExecutorConfig::default()
716        }
717
718        async fn execute(
719            &self,
720            _task: &Task,
721            _context: Arc<Context>,
722        ) -> Result<Self::Output, Self::Error> {
723            Ok(TestAgentOutput {
724                result: "unused".to_string(),
725            })
726        }
727
728        async fn execute_stream(
729            &self,
730            _task: &Task,
731            _context: Arc<Context>,
732        ) -> Result<crate::utils::BoxRuntimeStream<Result<Self::Output, Self::Error>>, Self::Error>
733        {
734            Err(TestError::ExecutionFailed(
735                "stream setup failed".to_string(),
736            ))
737        }
738    }
739
740    impl AgentHooks for FailingStreamSetupAgent {}
741
742    #[tokio::test]
743    async fn test_direct_agent_run_stream_setup_error_emits_task_error() {
744        let llm = Arc::new(ConfigurableLLMProvider::default());
745        let mut handle = AgentBuilder::<_, DirectAgent>::new(FailingStreamSetupAgent)
746            .llm(llm)
747            .stream(true)
748            .build()
749            .await
750            .expect("build should succeed");
751
752        let err = match handle.agent.run_stream(Task::new("fail setup")).await {
753            Err(err) => err,
754            Ok(_) => panic!("expected stream setup failure"),
755        };
756        assert!(matches!(err, RunnableAgentError::ExecutorError(_)));
757
758        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
759            .await
760            .expect("timed out waiting for TaskError event")
761            .expect("stream ended without event");
762        match event {
763            Event::TaskError { error, .. } => {
764                assert!(error.contains("stream setup failed"));
765            }
766            other => panic!("expected TaskError, got {other:?}"),
767        }
768    }
769
770    #[tokio::test]
771    async fn test_direct_agent_run_stream_item_error_emits_task_error() {
772        let mock_agent = MockAgentImpl::new("direct", "direct agent").with_failure(true);
773        let llm = Arc::new(ConfigurableLLMProvider::default());
774        let mut handle = AgentBuilder::<_, DirectAgent>::new(mock_agent)
775            .llm(llm)
776            .stream(true)
777            .build()
778            .await
779            .expect("build should succeed");
780
781        let mut stream = handle
782            .agent
783            .run_stream(Task::new("fail"))
784            .await
785            .expect("default execute_stream should return Ok stream");
786
787        let err = stream
788            .next()
789            .await
790            .expect("stream should yield one item")
791            .expect_err("expected stream item error");
792        assert!(matches!(err, Error::RunnableAgentError(_)));
793
794        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
795            .await
796            .expect("timed out waiting for TaskError event")
797            .expect("stream ended without event");
798        match event {
799            Event::TaskError { error, .. } => {
800                assert!(error.contains("Mock execution failed"));
801            }
802            other => panic!("expected TaskError, got {other:?}"),
803        }
804    }
805
806    #[derive(Debug, Clone)]
807    struct StreamItemErrorAgent;
808
809    #[async_trait]
810    impl AgentDeriveT for StreamItemErrorAgent {
811        type Output = TestAgentOutput;
812
813        fn description(&self) -> &'static str {
814            "stream item error agent"
815        }
816
817        fn output_schema(&self) -> Option<Value> {
818            Some(TestAgentOutput::structured_output_format())
819        }
820
821        fn name(&self) -> &'static str {
822            "stream_item_error"
823        }
824
825        fn tools(&self) -> Vec<Box<dyn ToolT>> {
826            vec![]
827        }
828    }
829
830    #[async_trait]
831    impl AgentExecutor for StreamItemErrorAgent {
832        type Output = TestAgentOutput;
833        type Error = TestError;
834
835        fn config(&self) -> ExecutorConfig {
836            ExecutorConfig::default()
837        }
838
839        async fn execute(
840            &self,
841            _task: &Task,
842            _context: Arc<Context>,
843        ) -> Result<Self::Output, Self::Error> {
844            Ok(TestAgentOutput {
845                result: "unused".to_string(),
846            })
847        }
848
849        async fn execute_stream(
850            &self,
851            _task: &Task,
852            _context: Arc<Context>,
853        ) -> Result<crate::utils::BoxRuntimeStream<Result<Self::Output, Self::Error>>, Self::Error>
854        {
855            Ok(Box::pin(futures::stream::iter([Err(
856                TestError::ExecutionFailed("stream item failed".to_string()),
857            )])))
858        }
859    }
860
861    impl AgentHooks for StreamItemErrorAgent {}
862
863    #[tokio::test]
864    async fn test_direct_agent_run_stream_custom_item_error_emits_task_error() {
865        let llm = Arc::new(ConfigurableLLMProvider::default());
866        let mut handle = AgentBuilder::<_, DirectAgent>::new(StreamItemErrorAgent)
867            .llm(llm)
868            .stream(true)
869            .build()
870            .await
871            .expect("build should succeed");
872
873        let mut stream = handle
874            .agent
875            .run_stream(Task::new("stream error"))
876            .await
877            .expect("stream should start");
878
879        let err = stream
880            .next()
881            .await
882            .expect("stream should yield one item")
883            .expect_err("expected stream item failure");
884        assert!(matches!(err, Error::RunnableAgentError(_)));
885
886        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
887            .await
888            .expect("timed out waiting for TaskError event")
889            .expect("stream ended without event");
890        match event {
891            Event::TaskError { error, .. } => {
892                assert!(error.contains("stream item failed"));
893            }
894            other => panic!("expected TaskError, got {other:?}"),
895        }
896    }
897
898    #[derive(Debug, Clone)]
899    struct EmptyStreamAgent;
900
901    #[async_trait]
902    impl AgentDeriveT for EmptyStreamAgent {
903        type Output = TestAgentOutput;
904
905        fn description(&self) -> &'static str {
906            "empty stream agent"
907        }
908
909        fn output_schema(&self) -> Option<Value> {
910            Some(TestAgentOutput::structured_output_format())
911        }
912
913        fn name(&self) -> &'static str {
914            "empty_stream_agent"
915        }
916
917        fn tools(&self) -> Vec<Box<dyn ToolT>> {
918            vec![]
919        }
920    }
921
922    #[async_trait]
923    impl AgentExecutor for EmptyStreamAgent {
924        type Output = TestAgentOutput;
925        type Error = TestError;
926
927        fn config(&self) -> ExecutorConfig {
928            ExecutorConfig::default()
929        }
930
931        async fn execute(
932            &self,
933            _task: &Task,
934            _context: Arc<Context>,
935        ) -> Result<Self::Output, Self::Error> {
936            Ok(TestAgentOutput {
937                result: "unused".to_string(),
938            })
939        }
940
941        async fn execute_stream(
942            &self,
943            _task: &Task,
944            _context: Arc<Context>,
945        ) -> Result<crate::utils::BoxRuntimeStream<Result<Self::Output, Self::Error>>, Self::Error>
946        {
947            Ok(Box::pin(futures::stream::empty()))
948        }
949    }
950
951    impl AgentHooks for EmptyStreamAgent {}
952
953    #[tokio::test]
954    async fn test_direct_agent_run_stream_empty_stream_emits_task_error() {
955        let llm = Arc::new(ConfigurableLLMProvider::default());
956        let mut handle = AgentBuilder::<_, DirectAgent>::new(EmptyStreamAgent)
957            .llm(llm)
958            .stream(true)
959            .build()
960            .await
961            .expect("build should succeed");
962
963        let stream = handle
964            .agent
965            .run_stream(Task::new("empty"))
966            .await
967            .expect("stream should start");
968        let outputs: Vec<_> = stream.collect().await;
969        assert!(outputs.is_empty());
970
971        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
972            .await
973            .expect("timed out waiting for TaskError event")
974            .expect("stream ended without event");
975        match event {
976            Event::TaskError { error, .. } => {
977                assert!(error.contains("without output"));
978            }
979            other => panic!("expected TaskError, got {other:?}"),
980        }
981    }
982
983    #[tokio::test]
984    async fn test_direct_agent_run_stream_uses_last_item_for_task_complete() {
985        let llm = Arc::new(ConfigurableLLMProvider::default());
986        let mut handle = AgentBuilder::<_, DirectAgent>::new(MultiItemStreamAgent)
987            .llm(llm)
988            .stream(true)
989            .build()
990            .await
991            .expect("build should succeed");
992
993        let stream = handle
994            .agent
995            .run_stream(Task::new("chunked"))
996            .await
997            .expect("stream should start");
998        let outputs: Vec<_> = stream.collect().await;
999        assert_eq!(outputs.len(), 3);
1000        assert_eq!(outputs[2].as_ref().expect("third item").result, "chunked-3");
1001
1002        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
1003            .await
1004            .expect("timed out waiting for TaskComplete event")
1005            .expect("stream ended without event");
1006        match event {
1007            Event::TaskComplete { result, .. } => {
1008                let parsed: Value =
1009                    serde_json::from_str(&result).expect("TaskComplete result should be JSON");
1010                assert_eq!(parsed["sequence"], 3);
1011                assert_eq!(parsed["response"], "chunked-3");
1012            }
1013            other => panic!("expected TaskComplete, got {other:?}"),
1014        }
1015    }
1016
1017    #[derive(Debug, Clone)]
1018    struct OkThenErrStreamAgent;
1019
1020    #[async_trait]
1021    impl AgentDeriveT for OkThenErrStreamAgent {
1022        type Output = TestAgentOutput;
1023
1024        fn description(&self) -> &'static str {
1025            "ok then err stream agent"
1026        }
1027
1028        fn output_schema(&self) -> Option<Value> {
1029            Some(TestAgentOutput::structured_output_format())
1030        }
1031
1032        fn name(&self) -> &'static str {
1033            "ok_then_err_stream_agent"
1034        }
1035
1036        fn tools(&self) -> Vec<Box<dyn ToolT>> {
1037            vec![]
1038        }
1039    }
1040
1041    #[async_trait]
1042    impl AgentExecutor for OkThenErrStreamAgent {
1043        type Output = TestAgentOutput;
1044        type Error = TestError;
1045
1046        fn config(&self) -> ExecutorConfig {
1047            ExecutorConfig::default()
1048        }
1049
1050        async fn execute(
1051            &self,
1052            _task: &Task,
1053            _context: Arc<Context>,
1054        ) -> Result<Self::Output, Self::Error> {
1055            Ok(TestAgentOutput {
1056                result: "unused".to_string(),
1057            })
1058        }
1059
1060        async fn execute_stream(
1061            &self,
1062            _task: &Task,
1063            _context: Arc<Context>,
1064        ) -> Result<crate::utils::BoxRuntimeStream<Result<Self::Output, Self::Error>>, Self::Error>
1065        {
1066            Ok(Box::pin(futures::stream::iter([
1067                Ok(TestAgentOutput {
1068                    result: "partial".to_string(),
1069                }),
1070                Err(TestError::ExecutionFailed("stream item failed".to_string())),
1071            ])))
1072        }
1073    }
1074
1075    impl AgentHooks for OkThenErrStreamAgent {}
1076
1077    #[tokio::test]
1078    async fn test_direct_agent_run_stream_ok_then_err_emits_task_error_not_task_complete() {
1079        let llm = Arc::new(ConfigurableLLMProvider::default());
1080        let mut handle = AgentBuilder::<_, DirectAgent>::new(OkThenErrStreamAgent)
1081            .llm(llm)
1082            .stream(true)
1083            .build()
1084            .await
1085            .expect("build should succeed");
1086
1087        let mut stream = handle
1088            .agent
1089            .run_stream(Task::new("partial failure"))
1090            .await
1091            .expect("stream should start");
1092
1093        let first = stream
1094            .next()
1095            .await
1096            .expect("stream should yield ok item")
1097            .expect("expected ok item");
1098        assert_eq!(first.result, "partial");
1099
1100        let second = stream
1101            .next()
1102            .await
1103            .expect("stream should yield err item")
1104            .expect_err("expected err item");
1105        assert!(matches!(second, Error::RunnableAgentError(_)));
1106
1107        assert!(
1108            stream.next().await.is_none(),
1109            "stream should end after error item"
1110        );
1111
1112        let event = tokio::time::timeout(std::time::Duration::from_secs(1), handle.rx.next())
1113            .await
1114            .expect("timed out waiting for TaskError event")
1115            .expect("stream ended without event");
1116        match event {
1117            Event::TaskError { error, .. } => {
1118                assert!(error.contains("stream item failed"));
1119            }
1120            other => panic!("expected TaskError, got {other:?}"),
1121        }
1122
1123        let no_terminal_success =
1124            tokio::time::timeout(std::time::Duration::from_millis(100), handle.rx.next()).await;
1125        assert!(
1126            no_terminal_success.is_err(),
1127            "Ok-then-Err stream should not emit TaskComplete after TaskError"
1128        );
1129    }
1130}
autoagents_core/agent/direct.rs

autoagents_core/agent/
direct.rs
