autoagents_core/agent/executor/
turn_engine.rs

1use crate::agent::Context;
2use crate::agent::executor::event_helper::EventHelper;
3use crate::agent::executor::memory_policy::{MemoryAdapter, MemoryPolicy};
4use crate::agent::executor::tool_processor::ToolProcessor;
5use crate::agent::hooks::AgentHooks;
6use crate::agent::task::Task;
7use crate::channel::{Sender, channel};
8use crate::tool::{ToolCallResult, ToolT, to_llm_tool};
9use crate::utils::stream_from_producer;
10use autoagents_llm::ToolCall;
11use autoagents_llm::chat::{ChatMessage, ChatRole, MessageType, StreamChunk, StreamResponse};
12use autoagents_llm::error::LLMError;
13use autoagents_protocol::{Event, SubmissionId};
14#[cfg(target_arch = "wasm32")]
15use futures::SinkExt;
16use futures::{Stream, StreamExt};
17use serde_json::Value;
18use std::collections::HashSet;
19use std::pin::Pin;
20use std::sync::Arc;
21use thiserror::Error;
22
23#[cfg(not(target_arch = "wasm32"))]
24use tokio::sync::mpsc;
25
26#[cfg(target_arch = "wasm32")]
27use futures::channel::mpsc;
28
29/// Defines if tools are enabled for a given execution plan.
30#[derive(Debug, Clone, Copy)]
31pub enum ToolMode {
32    Enabled,
33    Disabled,
34}
35
36/// Defines which streaming primitive to use.
37#[derive(Debug, Clone, Copy)]
38pub enum StreamMode {
39    Structured,
40    Tool,
41}
42
43/// Configuration for the shared executor engine.
44#[derive(Debug, Clone)]
45pub struct TurnEngineConfig {
46    pub max_turns: usize,
47    pub tool_mode: ToolMode,
48    pub stream_mode: StreamMode,
49    pub memory_policy: MemoryPolicy,
50}
51
52impl TurnEngineConfig {
53    pub fn basic(max_turns: usize) -> Self {
54        Self {
55            max_turns,
56            tool_mode: ToolMode::Disabled,
57            stream_mode: StreamMode::Structured,
58            memory_policy: MemoryPolicy::basic(),
59        }
60    }
61
62    pub fn react(max_turns: usize) -> Self {
63        Self {
64            max_turns,
65            tool_mode: ToolMode::Enabled,
66            stream_mode: StreamMode::Tool,
67            memory_policy: MemoryPolicy::react(),
68        }
69    }
70}
71
72/// Normalized output emitted by the engine for a single turn.
73#[derive(Debug, Clone)]
74pub struct TurnEngineOutput {
75    pub response: String,
76    pub reasoning_content: String,
77    pub tool_calls: Vec<ToolCallResult>,
78}
79
80/// Streaming deltas emitted per turn.
81#[derive(Debug)]
82pub enum TurnDelta {
83    Text(String),
84    ReasoningContent(String),
85    ToolResults(Vec<ToolCallResult>),
86    Done(crate::agent::executor::TurnResult<TurnEngineOutput>),
87}
88
89#[derive(Error, Debug)]
90pub enum TurnEngineError {
91    #[error("LLM error: {0}")]
92    LLMError(
93        #[from]
94        #[source]
95        LLMError,
96    ),
97
98    #[error("Run aborted by hook")]
99    Aborted,
100
101    #[error("Other error: {0}")]
102    Other(String),
103}
104
105/// Per-run state for the turn engine.
106#[derive(Clone)]
107pub struct TurnState {
108    memory: MemoryAdapter,
109    stored_user: bool,
110}
111
112impl TurnState {
113    pub fn new(context: &Context, policy: MemoryPolicy) -> Self {
114        Self {
115            memory: MemoryAdapter::new(context.memory(), policy),
116            stored_user: false,
117        }
118    }
119
120    pub fn memory(&self) -> &MemoryAdapter {
121        &self.memory
122    }
123
124    pub fn stored_user(&self) -> bool {
125        self.stored_user
126    }
127
128    fn mark_user_stored(&mut self) {
129        self.stored_user = true;
130    }
131}
132
133/// Shared turn engine that handles memory, tools, and events consistently.
134#[derive(Debug, Clone)]
135pub struct TurnEngine {
136    config: TurnEngineConfig,
137}
138
139impl TurnEngine {
140    pub fn new(config: TurnEngineConfig) -> Self {
141        Self { config }
142    }
143
144    pub fn turn_state(&self, context: &Context) -> TurnState {
145        TurnState::new(context, self.config.memory_policy.clone())
146    }
147
148    pub async fn run_turn<H: AgentHooks>(
149        &self,
150        hooks: &H,
151        task: &Task,
152        context: &Context,
153        turn_state: &mut TurnState,
154        turn_index: usize,
155        max_turns: usize,
156    ) -> Result<crate::agent::executor::TurnResult<TurnEngineOutput>, TurnEngineError> {
157        let max_turns = normalize_max_turns(max_turns, self.config.max_turns);
158        let tx_event = context.tx().ok();
159        EventHelper::send_turn_started(
160            &tx_event,
161            task.submission_id,
162            context.config().id,
163            turn_index,
164            max_turns,
165        )
166        .await;
167
168        hooks.on_turn_start(turn_index, context).await;
169
170        let include_user_prompt =
171            should_include_user_prompt(turn_state.memory(), turn_state.stored_user());
172        let messages = self
173            .build_messages(context, task, turn_state.memory(), include_user_prompt)
174            .await;
175        let store_user = should_store_user(turn_state);
176
177        let tools = context.tools();
178        let response = self.get_llm_response(context, &messages, tools).await?;
179        let response_text = response.text().unwrap_or_default();
180        let reasoning_content = response.thinking().unwrap_or_default();
181        if store_user {
182            turn_state.memory.store_user(task).await?;
183            turn_state.mark_user_stored();
184        }
185
186        let tool_calls = if matches!(self.config.tool_mode, ToolMode::Enabled) {
187            response.tool_calls().unwrap_or_default()
188        } else {
189            Vec::new()
190        };
191
192        if !tool_calls.is_empty() {
193            let tool_results = process_tool_calls_with_hooks(
194                hooks,
195                context,
196                task.submission_id,
197                tools,
198                &tool_calls,
199                &tx_event,
200            )
201            .await;
202
203            turn_state
204                .memory
205                .store_tool_interaction(&tool_calls, &tool_results, &response_text)
206                .await?;
207            record_tool_calls_state(context, &tool_results);
208
209            EventHelper::send_turn_completed(
210                &tx_event,
211                task.submission_id,
212                context.config().id,
213                turn_index,
214                false,
215            )
216            .await;
217            hooks.on_turn_complete(turn_index, context).await;
218
219            return Ok(crate::agent::executor::TurnResult::Continue(Some(
220                TurnEngineOutput {
221                    response: response_text,
222                    reasoning_content,
223                    tool_calls: tool_results,
224                },
225            )));
226        }
227
228        if !response_text.is_empty() {
229            turn_state.memory.store_assistant(&response_text).await?;
230        }
231
232        EventHelper::send_turn_completed(
233            &tx_event,
234            task.submission_id,
235            context.config().id,
236            turn_index,
237            true,
238        )
239        .await;
240        hooks.on_turn_complete(turn_index, context).await;
241
242        Ok(crate::agent::executor::TurnResult::Complete(
243            TurnEngineOutput {
244                response: response_text,
245                reasoning_content,
246                tool_calls: Vec::new(),
247            },
248        ))
249    }
250
251    pub async fn run_turn_stream<H>(
252        &self,
253        hooks: H,
254        task: &Task,
255        context: Arc<Context>,
256        turn_state: &mut TurnState,
257        turn_index: usize,
258        max_turns: usize,
259    ) -> Result<crate::utils::BoxRuntimeStream<Result<TurnDelta, TurnEngineError>>, TurnEngineError>
260    where
261        H: AgentHooks + Clone + Send + Sync + 'static,
262    {
263        let max_turns = normalize_max_turns(max_turns, self.config.max_turns);
264        let include_user_prompt =
265            should_include_user_prompt(turn_state.memory(), turn_state.stored_user());
266        let messages = self
267            .build_messages(&context, task, turn_state.memory(), include_user_prompt)
268            .await;
269        let store_user = should_store_user(turn_state);
270        if store_user {
271            turn_state.mark_user_stored();
272        }
273
274        let (mut tx, rx) = channel::<Result<TurnDelta, TurnEngineError>>(100);
275        let engine = self.clone();
276        let context_clone = context.clone();
277        let task = task.clone();
278        let hooks = hooks.clone();
279        let memory = turn_state.memory.clone();
280        let messages = messages.clone();
281
282        let producer = async move {
283            let tx_event = context_clone.tx().ok();
284            EventHelper::send_turn_started(
285                &tx_event,
286                task.submission_id,
287                context_clone.config().id,
288                turn_index,
289                max_turns,
290            )
291            .await;
292            hooks.on_turn_start(turn_index, &context_clone).await;
293
294            let result = match engine.config.stream_mode {
295                StreamMode::Structured => {
296                    engine
297                        .stream_structured(
298                            &context_clone,
299                            &task,
300                            &memory,
301                            &mut tx,
302                            &messages,
303                            store_user,
304                        )
305                        .await
306                }
307                StreamMode::Tool => {
308                    engine
309                        .stream_with_tools(
310                            &hooks,
311                            &context_clone,
312                            &task,
313                            context_clone.tools(),
314                            &memory,
315                            &mut tx,
316                            &messages,
317                            store_user,
318                        )
319                        .await
320                }
321            };
322
323            match result {
324                Ok(turn_result) => {
325                    let final_turn =
326                        matches!(turn_result, crate::agent::executor::TurnResult::Complete(_));
327                    EventHelper::send_turn_completed(
328                        &tx_event,
329                        task.submission_id,
330                        context_clone.config().id,
331                        turn_index,
332                        final_turn,
333                    )
334                    .await;
335                    hooks.on_turn_complete(turn_index, &context_clone).await;
336                    let _ = tx.send(Ok(TurnDelta::Done(turn_result))).await;
337                }
338                Err(err) => {
339                    let _ = tx.send(Err(err)).await;
340                }
341            }
342        };
343
344        Ok(stream_from_producer(rx, producer))
345    }
346
347    async fn stream_structured(
348        &self,
349        context: &Context,
350        task: &Task,
351        memory: &MemoryAdapter,
352        tx: &mut Sender<Result<TurnDelta, TurnEngineError>>,
353        messages: &[ChatMessage],
354        store_user: bool,
355    ) -> Result<crate::agent::executor::TurnResult<TurnEngineOutput>, TurnEngineError> {
356        let mut stream = self.get_structured_stream(context, messages).await?;
357        if store_user {
358            memory.store_user(task).await?;
359        }
360        let mut response_text = String::default();
361        let mut reasoning_content = String::default();
362
363        while let Some(chunk_result) = stream.next().await {
364            let chunk = chunk_result.map_err(TurnEngineError::LLMError)?;
365            let delta = chunk.choices.first().map(|choice| &choice.delta);
366            let content = delta
367                .and_then(|d| d.content.as_ref())
368                .map(String::as_str)
369                .unwrap_or("")
370                .to_string();
371            let reasoning = delta
372                .and_then(|d| d.reasoning_content.as_ref())
373                .map(String::as_str)
374                .unwrap_or("")
375                .to_string();
376
377            let tx_event = context.tx().ok();
378            if !content.is_empty() {
379                response_text.push_str(&content);
380                let _ = tx.send(Ok(TurnDelta::Text(content.clone()))).await;
381                EventHelper::send_stream_chunk(
382                    &tx_event,
383                    task.submission_id,
384                    StreamChunk::Text(content),
385                )
386                .await;
387            }
388            if !reasoning.is_empty() {
389                reasoning_content.push_str(&reasoning);
390                let _ = tx
391                    .send(Ok(TurnDelta::ReasoningContent(reasoning.clone())))
392                    .await;
393                EventHelper::send_stream_chunk(
394                    &tx_event,
395                    task.submission_id,
396                    StreamChunk::ReasoningContent(reasoning),
397                )
398                .await;
399            }
400        }
401
402        if !response_text.is_empty() {
403            memory.store_assistant(&response_text).await?;
404        }
405
406        Ok(crate::agent::executor::TurnResult::Complete(
407            TurnEngineOutput {
408                response: response_text,
409                reasoning_content,
410                tool_calls: Vec::default(),
411            },
412        ))
413    }
414
415    #[allow(clippy::too_many_arguments)]
416    async fn stream_with_tools<H: AgentHooks>(
417        &self,
418        hooks: &H,
419        context: &Context,
420        task: &Task,
421        tools: &[Box<dyn ToolT>],
422        memory: &MemoryAdapter,
423        tx: &mut Sender<Result<TurnDelta, TurnEngineError>>,
424        messages: &[ChatMessage],
425        store_user: bool,
426    ) -> Result<crate::agent::executor::TurnResult<TurnEngineOutput>, TurnEngineError> {
427        let mut stream = self.get_tool_stream(context, messages, tools).await?;
428        if store_user {
429            memory.store_user(task).await?;
430        }
431        let mut response_text = String::default();
432        let mut reasoning_content = String::default();
433        let mut tool_calls = Vec::default();
434        let mut tool_call_ids: HashSet<String> = HashSet::default();
435
436        while let Some(chunk_result) = stream.next().await {
437            let chunk = chunk_result.map_err(TurnEngineError::LLMError)?;
438            let chunk_clone = chunk.clone();
439
440            match chunk {
441                StreamChunk::Text(content) => {
442                    response_text.push_str(&content);
443                    let _ = tx.send(Ok(TurnDelta::Text(content.clone()))).await;
444                }
445                StreamChunk::ReasoningContent(content) => {
446                    reasoning_content.push_str(&content);
447                    let _ = tx.send(Ok(TurnDelta::ReasoningContent(content))).await;
448                }
449                StreamChunk::ToolUseComplete { tool_call, .. }
450                    if tool_call_ids.insert(tool_call.id.clone()) =>
451                {
452                    tool_calls.push(tool_call.clone());
453                    let tx_event = context.tx().ok();
454                    EventHelper::send_stream_tool_call(
455                        &tx_event,
456                        task.submission_id,
457                        serde_json::to_value(tool_call).unwrap_or(Value::Null),
458                    )
459                    .await;
460                }
461                StreamChunk::Usage(_) => {}
462                _ => {}
463            }
464
465            let tx_event = context.tx().ok();
466            EventHelper::send_stream_chunk(&tx_event, task.submission_id, chunk_clone).await;
467        }
468
469        if tool_calls.is_empty() {
470            if !response_text.is_empty() {
471                memory.store_assistant(&response_text).await?;
472            }
473            return Ok(crate::agent::executor::TurnResult::Complete(
474                TurnEngineOutput {
475                    response: response_text,
476                    reasoning_content,
477                    tool_calls: Vec::new(),
478                },
479            ));
480        }
481
482        let tx_event = context.tx().ok();
483        let tool_results = process_tool_calls_with_hooks(
484            hooks,
485            context,
486            task.submission_id,
487            tools,
488            &tool_calls,
489            &tx_event,
490        )
491        .await;
492
493        memory
494            .store_tool_interaction(&tool_calls, &tool_results, &response_text)
495            .await?;
496        record_tool_calls_state(context, &tool_results);
497
498        let _ = tx
499            .send(Ok(TurnDelta::ToolResults(tool_results.clone())))
500            .await;
501
502        Ok(crate::agent::executor::TurnResult::Continue(Some(
503            TurnEngineOutput {
504                response: response_text,
505                reasoning_content,
506                tool_calls: tool_results,
507            },
508        )))
509    }
510
511    async fn get_llm_response(
512        &self,
513        context: &Context,
514        messages: &[ChatMessage],
515        tools: &[Box<dyn ToolT>],
516    ) -> Result<Box<dyn autoagents_llm::chat::ChatResponse>, TurnEngineError> {
517        let llm = context.llm();
518        let output_schema = context.config().output_schema.clone();
519
520        if matches!(self.config.tool_mode, ToolMode::Enabled) && !tools.is_empty() {
521            let cached = context.serialized_tools();
522            let tools_serialized = if let Some(cached) = cached {
523                cached
524            } else {
525                Arc::new(tools.iter().map(to_llm_tool).collect::<Vec<_>>())
526            };
527            llm.chat_with_tools(messages, Some(&tools_serialized), output_schema)
528                .await
529                .map_err(TurnEngineError::LLMError)
530        } else {
531            llm.chat(messages, output_schema)
532                .await
533                .map_err(TurnEngineError::LLMError)
534        }
535    }
536
537    async fn get_structured_stream(
538        &self,
539        context: &Context,
540        messages: &[ChatMessage],
541    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamResponse, LLMError>> + Send>>, TurnEngineError>
542    {
543        context
544            .llm()
545            .chat_stream_struct(messages, None, context.config().output_schema.clone())
546            .await
547            .map_err(TurnEngineError::LLMError)
548    }
549
550    async fn get_tool_stream(
551        &self,
552        context: &Context,
553        messages: &[ChatMessage],
554        tools: &[Box<dyn ToolT>],
555    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, LLMError>> + Send>>, TurnEngineError>
556    {
557        let cached = context.serialized_tools();
558        let tools_serialized = if let Some(cached) = cached {
559            cached
560        } else {
561            Arc::new(tools.iter().map(to_llm_tool).collect::<Vec<_>>())
562        };
563        context
564            .llm()
565            .chat_stream_with_tools(
566                messages,
567                if tools_serialized.is_empty() {
568                    None
569                } else {
570                    Some(&tools_serialized)
571                },
572                context.config().output_schema.clone(),
573            )
574            .await
575            .map_err(TurnEngineError::LLMError)
576    }
577
578    async fn build_messages(
579        &self,
580        context: &Context,
581        task: &Task,
582        memory: &MemoryAdapter,
583        include_user_prompt: bool,
584    ) -> Vec<ChatMessage> {
585        let system_prompt = task
586            .system_prompt
587            .as_deref()
588            .unwrap_or_else(|| &context.config().description);
589        let mut messages = vec![ChatMessage {
590            role: ChatRole::System,
591            message_type: MessageType::Text,
592            content: system_prompt.to_string(),
593        }];
594
595        let recalled = memory.recall_messages(task).await;
596        messages.extend(recalled);
597
598        if include_user_prompt {
599            messages.push(user_message(task));
600        }
601
602        messages
603    }
604}
605
606pub fn record_task_state(context: &Context, task: &Task) {
607    let state = context.state();
608    #[cfg(not(target_arch = "wasm32"))]
609    if let Ok(mut guard) = state.try_lock() {
610        guard.record_task(task.clone());
611    };
612    #[cfg(target_arch = "wasm32")]
613    if let Some(mut guard) = state.try_lock() {
614        guard.record_task(task.clone());
615    };
616}
617
618fn user_message(task: &Task) -> ChatMessage {
619    if let Some((mime, image_data)) = &task.image {
620        ChatMessage {
621            role: ChatRole::User,
622            message_type: MessageType::Image(((*mime).into(), image_data.clone())),
623            content: task.prompt.clone(),
624        }
625    } else {
626        ChatMessage {
627            role: ChatRole::User,
628            message_type: MessageType::Text,
629            content: task.prompt.clone(),
630        }
631    }
632}
633
634fn should_include_user_prompt(memory: &MemoryAdapter, stored_user: bool) -> bool {
635    if !memory.is_enabled() {
636        return true;
637    }
638    if !memory.policy().recall {
639        return true;
640    }
641    if !memory.policy().store_user {
642        return true;
643    }
644    !stored_user
645}
646
647fn should_store_user(turn_state: &TurnState) -> bool {
648    if !turn_state.memory.is_enabled() {
649        return false;
650    }
651    if !turn_state.memory.policy().store_user {
652        return false;
653    }
654    !turn_state.stored_user
655}
656
657fn normalize_max_turns(max_turns: usize, fallback: usize) -> usize {
658    if max_turns == 0 {
659        return fallback.max(1);
660    }
661    max_turns
662}
663
664fn record_tool_calls_state(context: &Context, tool_results: &[ToolCallResult]) {
665    if tool_results.is_empty() {
666        return;
667    }
668    let state = context.state();
669    #[cfg(not(target_arch = "wasm32"))]
670    if let Ok(mut guard) = state.try_lock() {
671        for result in tool_results {
672            guard.record_tool_call(result.clone());
673        }
674    };
675    #[cfg(target_arch = "wasm32")]
676    if let Some(mut guard) = state.try_lock() {
677        for result in tool_results {
678            guard.record_tool_call(result.clone());
679        }
680    };
681}
682
683async fn process_tool_calls_with_hooks<H: AgentHooks>(
684    hooks: &H,
685    context: &Context,
686    submission_id: SubmissionId,
687    tools: &[Box<dyn ToolT>],
688    tool_calls: &[ToolCall],
689    tx_event: &Option<mpsc::Sender<Event>>,
690) -> Vec<ToolCallResult> {
691    let mut results = Vec::new();
692    for call in tool_calls {
693        if let Some(result) = ToolProcessor::process_single_tool_call_with_hooks(
694            hooks,
695            context,
696            submission_id,
697            tools,
698            call,
699            tx_event,
700        )
701        .await
702        {
703            results.push(result);
704        }
705    }
706    results
707}
708
709#[cfg(test)]
710mod tests {
711    use super::*;
712    use crate::agent::memory::{MemoryProvider, MemoryType, SlidingWindowMemory};
713    use crate::agent::task::Task;
714    use crate::agent::{AgentConfig, Context};
715    use crate::tests::{ConfigurableLLMProvider, StaticChatResponse};
716    use async_trait::async_trait;
717    use autoagents_llm::LLMProvider;
718    use autoagents_llm::ToolCall;
719    use autoagents_llm::chat::{StreamChoice, StreamChunk, StreamDelta, StreamResponse};
720    use autoagents_llm::error::GuardrailPhase;
721    use autoagents_protocol::ActorID;
722    use futures::StreamExt;
723
724    #[derive(Debug)]
725    struct LocalTool {
726        name: String,
727        output: serde_json::Value,
728    }
729
730    impl LocalTool {
731        fn new(name: &str, output: serde_json::Value) -> Self {
732            Self {
733                name: name.to_string(),
734                output,
735            }
736        }
737    }
738
739    impl crate::tool::ToolT for LocalTool {
740        fn name(&self) -> &str {
741            &self.name
742        }
743
744        fn description(&self) -> &str {
745            "local tool"
746        }
747
748        fn args_schema(&self) -> serde_json::Value {
749            serde_json::json!({"type": "object"})
750        }
751    }
752
753    #[async_trait]
754    impl crate::tool::ToolRuntime for LocalTool {
755        async fn execute(
756            &self,
757            _args: serde_json::Value,
758        ) -> Result<serde_json::Value, crate::tool::ToolCallError> {
759            Ok(self.output.clone())
760        }
761    }
762
763    #[derive(Debug)]
764    struct GuardrailRejectLLMProvider;
765
766    #[derive(Clone)]
767    struct FailingMemoryProvider;
768
769    #[async_trait]
770    impl MemoryProvider for FailingMemoryProvider {
771        async fn remember(&mut self, _message: &ChatMessage) -> Result<(), LLMError> {
772            Err(LLMError::ProviderError("memory write failed".to_string()))
773        }
774
775        async fn recall(
776            &self,
777            _query: &str,
778            _limit: Option<usize>,
779        ) -> Result<Vec<ChatMessage>, LLMError> {
780            Ok(Vec::new())
781        }
782
783        async fn clear(&mut self) -> Result<(), LLMError> {
784            Ok(())
785        }
786
787        fn memory_type(&self) -> MemoryType {
788            MemoryType::Custom
789        }
790
791        fn size(&self) -> usize {
792            0
793        }
794
795        fn clone_box(&self) -> Box<dyn MemoryProvider> {
796            Box::new(self.clone())
797        }
798    }
799
800    fn guardrail_block_error() -> LLMError {
801        LLMError::GuardrailBlocked {
802            phase: GuardrailPhase::Input,
803            guard: "prompt-injection".to_string().into(),
804            rule_id: "prompt_injection_detected".to_string().into(),
805            category: "prompt_injection".to_string().into(),
806            severity: "high".to_string().into(),
807            message: "detected suspicious instruction pattern: jailbreak"
808                .to_string()
809                .into(),
810        }
811    }
812
813    #[async_trait]
814    impl autoagents_llm::chat::ChatProvider for GuardrailRejectLLMProvider {
815        async fn chat(
816            &self,
817            _messages: &[ChatMessage],
818            _json_schema: Option<autoagents_llm::chat::StructuredOutputFormat>,
819        ) -> Result<Box<dyn autoagents_llm::chat::ChatResponse>, LLMError> {
820            Err(guardrail_block_error())
821        }
822
823        async fn chat_with_tools(
824            &self,
825            _messages: &[ChatMessage],
826            _tools: Option<&[autoagents_llm::chat::Tool]>,
827            _json_schema: Option<autoagents_llm::chat::StructuredOutputFormat>,
828        ) -> Result<Box<dyn autoagents_llm::chat::ChatResponse>, LLMError> {
829            Err(guardrail_block_error())
830        }
831
832        async fn chat_stream_struct(
833            &self,
834            _messages: &[ChatMessage],
835            _tools: Option<&[autoagents_llm::chat::Tool]>,
836            _json_schema: Option<autoagents_llm::chat::StructuredOutputFormat>,
837        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamResponse, LLMError>> + Send>>, LLMError>
838        {
839            Err(guardrail_block_error())
840        }
841
842        async fn chat_stream_with_tools(
843            &self,
844            _messages: &[ChatMessage],
845            _tools: Option<&[autoagents_llm::chat::Tool]>,
846            _json_schema: Option<autoagents_llm::chat::StructuredOutputFormat>,
847        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, LLMError>> + Send>>, LLMError>
848        {
849            Err(guardrail_block_error())
850        }
851    }
852
853    #[async_trait]
854    impl autoagents_llm::completion::CompletionProvider for GuardrailRejectLLMProvider {
855        async fn complete(
856            &self,
857            _req: &autoagents_llm::completion::CompletionRequest,
858            _json_schema: Option<autoagents_llm::chat::StructuredOutputFormat>,
859        ) -> Result<autoagents_llm::completion::CompletionResponse, LLMError> {
860            Ok(autoagents_llm::completion::CompletionResponse {
861                text: String::default(),
862            })
863        }
864    }
865
866    #[async_trait]
867    impl autoagents_llm::embedding::EmbeddingProvider for GuardrailRejectLLMProvider {
868        async fn embed(&self, _input: Vec<String>) -> Result<Vec<Vec<f32>>, LLMError> {
869            Ok(Vec::new())
870        }
871    }
872
873    #[async_trait]
874    impl autoagents_llm::models::ModelsProvider for GuardrailRejectLLMProvider {}
875
876    impl LLMProvider for GuardrailRejectLLMProvider {}
877
878    fn context_with_memory(llm: Arc<dyn LLMProvider>) -> Context {
879        let config = AgentConfig {
880            id: ActorID::new_v4(),
881            name: "memory_agent".to_string(),
882            description: "desc".to_string(),
883            output_schema: None,
884        };
885        let memory: Box<dyn MemoryProvider> = Box::new(SlidingWindowMemory::new(20));
886        Context::new(llm, None)
887            .with_config(config)
888            .with_memory(Some(Arc::new(tokio::sync::Mutex::new(memory))))
889    }
890
891    fn context_with_failing_memory(llm: Arc<dyn LLMProvider>) -> Context {
892        let config = AgentConfig {
893            id: ActorID::new_v4(),
894            name: "memory_agent".to_string(),
895            description: "desc".to_string(),
896            output_schema: None,
897        };
898        let memory: Box<dyn MemoryProvider> = Box::new(FailingMemoryProvider);
899        Context::new(llm, None)
900            .with_config(config)
901            .with_memory(Some(Arc::new(tokio::sync::Mutex::new(memory))))
902    }
903
904    fn assert_turn_memory_error(error: TurnEngineError) {
905        match error {
906            TurnEngineError::LLMError(LLMError::ProviderError(message)) => {
907                assert_eq!(message, "memory write failed");
908            }
909            other => panic!("expected memory provider error, got {other:?}"),
910        }
911    }
912
913    async fn recalled_messages(context: &Context) -> Vec<ChatMessage> {
914        let memory = context.memory().expect("memory should exist");
915        memory
916            .lock()
917            .await
918            .recall("", None)
919            .await
920            .expect("memory recall should succeed")
921    }
922
923    #[test]
924    fn test_turn_engine_config_basic() {
925        let config = TurnEngineConfig::basic(5);
926        assert_eq!(config.max_turns, 5);
927        assert!(matches!(config.tool_mode, ToolMode::Disabled));
928        assert!(matches!(config.stream_mode, StreamMode::Structured));
929        assert!(config.memory_policy.recall);
930    }
931
932    #[test]
933    fn test_turn_engine_config_react() {
934        let config = TurnEngineConfig::react(10);
935        assert_eq!(config.max_turns, 10);
936        assert!(matches!(config.tool_mode, ToolMode::Enabled));
937        assert!(matches!(config.stream_mode, StreamMode::Tool));
938        assert!(config.memory_policy.recall);
939    }
940
941    #[tokio::test]
942    async fn test_run_turn_llm_error_does_not_store_user_message() {
943        use crate::tests::MockAgentImpl;
944
945        let llm: Arc<dyn LLMProvider> = Arc::new(GuardrailRejectLLMProvider);
946        let context = context_with_memory(llm);
947        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
948        let mut turn_state = engine.turn_state(&context);
949        let task = Task::new("jailbreak");
950        let hooks = MockAgentImpl::new("test", "test");
951
952        let result = engine
953            .run_turn(&hooks, &task, &context, &mut turn_state, 0, 1)
954            .await;
955        assert!(matches!(
956            result,
957            Err(TurnEngineError::LLMError(LLMError::GuardrailBlocked { .. }))
958        ));
959
960        let stored = recalled_messages(&context).await;
961        assert!(stored.is_empty());
962    }
963
964    #[tokio::test]
965    async fn test_run_turn_success_stores_user_once_in_memory() {
966        use crate::tests::MockAgentImpl;
967
968        let llm: Arc<dyn LLMProvider> = Arc::new(ConfigurableLLMProvider::default());
969        let context = context_with_memory(llm);
970        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
971        let mut turn_state = engine.turn_state(&context);
972        let task = Task::new("hello");
973        let hooks = MockAgentImpl::new("test", "test");
974
975        let result = engine
976            .run_turn(&hooks, &task, &context, &mut turn_state, 0, 1)
977            .await;
978        assert!(matches!(
979            result,
980            Ok(crate::agent::executor::TurnResult::Complete(_))
981        ));
982
983        let stored = recalled_messages(&context).await;
984        let user_count = stored
985            .iter()
986            .filter(|m| m.role == ChatRole::User && m.content == "hello")
987            .count();
988        let assistant_count = stored
989            .iter()
990            .filter(|m| m.role == ChatRole::Assistant)
991            .count();
992
993        assert_eq!(user_count, 1);
994        assert_eq!(assistant_count, 1);
995    }
996
997    #[tokio::test]
998    async fn test_run_turn_returns_memory_write_failure() {
999        use crate::tests::MockAgentImpl;
1000
1001        let llm: Arc<dyn LLMProvider> = Arc::new(ConfigurableLLMProvider::default());
1002        let context = context_with_failing_memory(llm);
1003        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1004        let mut turn_state = engine.turn_state(&context);
1005        let task = Task::new("hello");
1006        let hooks = MockAgentImpl::new("test", "test");
1007
1008        let result = engine
1009            .run_turn(&hooks, &task, &context, &mut turn_state, 0, 1)
1010            .await;
1011
1012        assert_turn_memory_error(result.expect_err("memory write should fail"));
1013    }
1014
1015    #[test]
1016    fn test_normalize_max_turns_nonzero() {
1017        assert_eq!(normalize_max_turns(5, 10), 5);
1018    }
1019
1020    #[test]
1021    fn test_normalize_max_turns_zero_uses_fallback() {
1022        assert_eq!(normalize_max_turns(0, 10), 10);
1023    }
1024
1025    #[test]
1026    fn test_normalize_max_turns_zero_fallback_zero() {
1027        assert_eq!(normalize_max_turns(0, 0), 1);
1028    }
1029
1030    #[test]
1031    fn test_should_include_user_prompt_no_memory() {
1032        let adapter = MemoryAdapter::new(None, MemoryPolicy::basic());
1033        assert!(should_include_user_prompt(&adapter, false));
1034    }
1035
1036    #[test]
1037    fn test_should_include_user_prompt_recall_disabled() {
1038        let mut policy = MemoryPolicy::basic();
1039        policy.recall = false;
1040        let mem: Box<dyn crate::agent::memory::MemoryProvider> =
1041            Box::new(crate::agent::memory::SlidingWindowMemory::new(10));
1042        let adapter = MemoryAdapter::new(
1043            Some(std::sync::Arc::new(tokio::sync::Mutex::new(mem))),
1044            policy,
1045        );
1046        assert!(should_include_user_prompt(&adapter, false));
1047    }
1048
1049    #[test]
1050    fn test_should_include_user_prompt_store_user_disabled() {
1051        let mut policy = MemoryPolicy::basic();
1052        policy.store_user = false;
1053        let mem: Box<dyn crate::agent::memory::MemoryProvider> =
1054            Box::new(crate::agent::memory::SlidingWindowMemory::new(10));
1055        let adapter = MemoryAdapter::new(
1056            Some(std::sync::Arc::new(tokio::sync::Mutex::new(mem))),
1057            policy,
1058        );
1059        assert!(should_include_user_prompt(&adapter, false));
1060    }
1061
1062    #[test]
1063    fn test_should_include_user_prompt_already_stored() {
1064        let mem: Box<dyn crate::agent::memory::MemoryProvider> =
1065            Box::new(crate::agent::memory::SlidingWindowMemory::new(10));
1066        let adapter = MemoryAdapter::new(
1067            Some(std::sync::Arc::new(tokio::sync::Mutex::new(mem))),
1068            MemoryPolicy::basic(),
1069        );
1070        // stored_user = true => should not include
1071        assert!(!should_include_user_prompt(&adapter, true));
1072    }
1073
1074    #[test]
1075    fn test_should_store_user_no_memory() {
1076        let state = TurnState {
1077            memory: MemoryAdapter::new(None, MemoryPolicy::basic()),
1078            stored_user: false,
1079        };
1080        assert!(!should_store_user(&state));
1081    }
1082
1083    #[test]
1084    fn test_should_store_user_already_stored() {
1085        let mem: Box<dyn crate::agent::memory::MemoryProvider> =
1086            Box::new(crate::agent::memory::SlidingWindowMemory::new(10));
1087        let state = TurnState {
1088            memory: MemoryAdapter::new(
1089                Some(std::sync::Arc::new(tokio::sync::Mutex::new(mem))),
1090                MemoryPolicy::basic(),
1091            ),
1092            stored_user: true,
1093        };
1094        assert!(!should_store_user(&state));
1095    }
1096
1097    #[test]
1098    fn test_user_message_text() {
1099        let task = Task::new("hello");
1100        let msg = user_message(&task);
1101        assert!(matches!(msg.role, ChatRole::User));
1102        assert!(matches!(msg.message_type, MessageType::Text));
1103        assert_eq!(msg.content, "hello");
1104    }
1105
1106    #[test]
1107    fn test_user_message_image() {
1108        let mut task = Task::new("describe");
1109        task.image = Some((autoagents_protocol::ImageMime::PNG, vec![1, 2, 3]));
1110        let msg = user_message(&task);
1111        assert!(matches!(msg.role, ChatRole::User));
1112        assert!(matches!(msg.message_type, MessageType::Image(_)));
1113    }
1114
1115    #[test]
1116    fn test_turn_state_new_and_mark_user_stored() {
1117        let config = AgentConfig {
1118            id: ActorID::new_v4(),
1119            name: "test".to_string(),
1120            description: "test".to_string(),
1121            output_schema: None,
1122        };
1123        let llm = std::sync::Arc::new(crate::tests::MockLLMProvider {});
1124        let context = Context::new(llm, None).with_config(config);
1125
1126        let mut state = TurnState::new(&context, MemoryPolicy::basic());
1127        assert!(!state.stored_user());
1128        state.mark_user_stored();
1129        assert!(state.stored_user());
1130    }
1131
1132    #[tokio::test]
1133    async fn test_build_messages_with_system_prompt() {
1134        let config = AgentConfig {
1135            id: ActorID::new_v4(),
1136            name: "test".to_string(),
1137            description: "default desc".to_string(),
1138            output_schema: None,
1139        };
1140        let llm = std::sync::Arc::new(crate::tests::MockLLMProvider {});
1141        let context = Context::new(llm, None).with_config(config);
1142
1143        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1144        let adapter = MemoryAdapter::new(None, MemoryPolicy::basic());
1145        let mut task = Task::new("user input");
1146        task.system_prompt = Some("custom system".to_string());
1147
1148        let messages = engine.build_messages(&context, &task, &adapter, true).await;
1149        // System + user
1150        assert_eq!(messages.len(), 2);
1151        assert_eq!(messages[0].content, "custom system");
1152        assert_eq!(messages[0].role, ChatRole::System);
1153        assert_eq!(messages[1].content, "user input");
1154    }
1155
1156    #[tokio::test]
1157    async fn test_build_messages_without_user_prompt() {
1158        let config = AgentConfig {
1159            id: ActorID::new_v4(),
1160            name: "test".to_string(),
1161            description: "desc".to_string(),
1162            output_schema: None,
1163        };
1164        let llm = std::sync::Arc::new(crate::tests::MockLLMProvider {});
1165        let context = Context::new(llm, None).with_config(config);
1166
1167        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1168        let adapter = MemoryAdapter::new(None, MemoryPolicy::basic());
1169        let task = Task::new("user input");
1170
1171        let messages = engine
1172            .build_messages(&context, &task, &adapter, false)
1173            .await;
1174        // Only system prompt
1175        assert_eq!(messages.len(), 1);
1176        assert_eq!(messages[0].role, ChatRole::System);
1177    }
1178
1179    #[tokio::test]
1180    async fn test_run_turn_no_tools_single_turn() {
1181        use crate::tests::MockAgentImpl;
1182        let config = AgentConfig {
1183            id: ActorID::new_v4(),
1184            name: "test".to_string(),
1185            description: "test desc".to_string(),
1186            output_schema: None,
1187        };
1188        let llm = std::sync::Arc::new(crate::tests::MockLLMProvider {});
1189        let context = Context::new(llm, None).with_config(config);
1190
1191        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1192        let mut turn_state = engine.turn_state(&context);
1193        let task = Task::new("test prompt");
1194        let hooks = MockAgentImpl::new("test", "test");
1195
1196        let result = engine
1197            .run_turn(&hooks, &task, &context, &mut turn_state, 0, 1)
1198            .await;
1199        assert!(result.is_ok());
1200        let turn_result = result.unwrap();
1201        assert!(matches!(
1202            turn_result,
1203            crate::agent::executor::TurnResult::Complete(_)
1204        ));
1205        if let crate::agent::executor::TurnResult::Complete(output) = turn_result {
1206            assert_eq!(output.response, "Mock response");
1207        }
1208    }
1209
1210    #[tokio::test]
1211    async fn test_run_turn_with_tool_calls_continues() {
1212        use crate::tests::MockAgentImpl;
1213        let tool_call = ToolCall {
1214            id: "call_1".to_string(),
1215            call_type: "function".to_string(),
1216            function: autoagents_llm::FunctionCall {
1217                name: "tool_a".to_string(),
1218                arguments: r#"{"value":1}"#.to_string(),
1219            },
1220        };
1221
1222        let llm = Arc::new(ConfigurableLLMProvider {
1223            chat_response: StaticChatResponse {
1224                text: Some("Use tool".to_string()),
1225                tool_calls: Some(vec![tool_call.clone()]),
1226                usage: None,
1227                thinking: None,
1228            },
1229            ..ConfigurableLLMProvider::default()
1230        });
1231
1232        let config = AgentConfig {
1233            id: ActorID::new_v4(),
1234            name: "tool_agent".to_string(),
1235            description: "desc".to_string(),
1236            output_schema: None,
1237        };
1238        let tool = LocalTool::new("tool_a", serde_json::json!({"ok": true}));
1239        let context = Context::new(llm, None)
1240            .with_config(config)
1241            .with_tools(vec![Box::new(tool)]);
1242
1243        let engine = TurnEngine::new(TurnEngineConfig {
1244            max_turns: 2,
1245            tool_mode: ToolMode::Enabled,
1246            stream_mode: StreamMode::Structured,
1247            memory_policy: MemoryPolicy::basic(),
1248        });
1249        let mut turn_state = engine.turn_state(&context);
1250        let task = Task::new("prompt");
1251        let hooks = MockAgentImpl::new("test", "test");
1252
1253        let result = engine
1254            .run_turn(&hooks, &task, &context, &mut turn_state, 0, 2)
1255            .await
1256            .unwrap();
1257
1258        match result {
1259            crate::agent::executor::TurnResult::Continue(Some(output)) => {
1260                assert_eq!(output.response, "Use tool");
1261                assert_eq!(output.tool_calls.len(), 1);
1262                assert!(output.tool_calls[0].success);
1263            }
1264            _ => panic!("expected Continue(Some)"),
1265        }
1266
1267        #[cfg(not(target_arch = "wasm32"))]
1268        if let Ok(state) = context.state().try_lock() {
1269            assert_eq!(state.tool_calls.len(), 1);
1270        }
1271    }
1272
1273    #[tokio::test]
1274    async fn test_run_turn_tool_mode_disabled_ignores_tool_calls() {
1275        use crate::tests::MockAgentImpl;
1276        let tool_call = ToolCall {
1277            id: "call_1".to_string(),
1278            call_type: "function".to_string(),
1279            function: autoagents_llm::FunctionCall {
1280                name: "tool_a".to_string(),
1281                arguments: r#"{"value":1}"#.to_string(),
1282            },
1283        };
1284
1285        let llm = Arc::new(ConfigurableLLMProvider {
1286            chat_response: StaticChatResponse {
1287                text: Some("No tools".to_string()),
1288                tool_calls: Some(vec![tool_call]),
1289                usage: None,
1290                thinking: None,
1291            },
1292            ..ConfigurableLLMProvider::default()
1293        });
1294
1295        let config = AgentConfig {
1296            id: ActorID::new_v4(),
1297            name: "tool_agent".to_string(),
1298            description: "desc".to_string(),
1299            output_schema: None,
1300        };
1301        let context = Context::new(llm, None).with_config(config);
1302
1303        let engine = TurnEngine::new(TurnEngineConfig {
1304            max_turns: 1,
1305            tool_mode: ToolMode::Disabled,
1306            stream_mode: StreamMode::Structured,
1307            memory_policy: MemoryPolicy::basic(),
1308        });
1309        let mut turn_state = engine.turn_state(&context);
1310        let task = Task::new("prompt");
1311        let hooks = MockAgentImpl::new("test", "test");
1312
1313        let result = engine
1314            .run_turn(&hooks, &task, &context, &mut turn_state, 0, 1)
1315            .await
1316            .unwrap();
1317
1318        match result {
1319            crate::agent::executor::TurnResult::Complete(output) => {
1320                assert_eq!(output.response, "No tools");
1321                assert!(output.tool_calls.is_empty());
1322            }
1323            _ => panic!("expected Complete"),
1324        }
1325    }
1326
1327    #[tokio::test]
1328    async fn test_run_turn_propagates_reasoning_content() {
1329        use crate::tests::MockAgentImpl;
1330
1331        let llm = Arc::new(ConfigurableLLMProvider {
1332            chat_response: StaticChatResponse {
1333                text: Some("answer".to_string()),
1334                tool_calls: None,
1335                usage: None,
1336                thinking: Some("reasoning".to_string()),
1337            },
1338            ..ConfigurableLLMProvider::default()
1339        });
1340
1341        let config = AgentConfig {
1342            id: ActorID::new_v4(),
1343            name: "reasoning_agent".to_string(),
1344            description: "desc".to_string(),
1345            output_schema: None,
1346        };
1347        let context = Context::new(llm, None).with_config(config);
1348        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1349        let mut turn_state = engine.turn_state(&context);
1350        let task = Task::new("prompt");
1351        let hooks = MockAgentImpl::new("test", "test");
1352
1353        let result = engine
1354            .run_turn(&hooks, &task, &context, &mut turn_state, 0, 1)
1355            .await
1356            .unwrap();
1357
1358        match result {
1359            crate::agent::executor::TurnResult::Complete(output) => {
1360                assert_eq!(output.response, "answer");
1361                assert_eq!(output.reasoning_content, "reasoning");
1362            }
1363            _ => panic!("expected Complete"),
1364        }
1365    }
1366
1367    #[tokio::test]
1368    async fn test_run_turn_stream_structured_aggregates_text() {
1369        use crate::tests::MockAgentImpl;
1370        let llm = Arc::new(ConfigurableLLMProvider {
1371            structured_stream: vec![
1372                StreamResponse {
1373                    choices: vec![StreamChoice {
1374                        delta: StreamDelta {
1375                            content: Some("Hello ".to_string()),
1376                            reasoning_content: None,
1377                            tool_calls: None,
1378                        },
1379                    }],
1380                    usage: None,
1381                },
1382                StreamResponse {
1383                    choices: vec![StreamChoice {
1384                        delta: StreamDelta {
1385                            content: Some("world".to_string()),
1386                            reasoning_content: None,
1387                            tool_calls: None,
1388                        },
1389                    }],
1390                    usage: None,
1391                },
1392            ],
1393            ..ConfigurableLLMProvider::default()
1394        });
1395
1396        let config = AgentConfig {
1397            id: ActorID::new_v4(),
1398            name: "stream_agent".to_string(),
1399            description: "desc".to_string(),
1400            output_schema: None,
1401        };
1402        let context = Arc::new(Context::new(llm, None).with_config(config));
1403        let engine = TurnEngine::new(TurnEngineConfig {
1404            max_turns: 1,
1405            tool_mode: ToolMode::Disabled,
1406            stream_mode: StreamMode::Structured,
1407            memory_policy: MemoryPolicy::basic(),
1408        });
1409        let mut turn_state = engine.turn_state(&context);
1410        let task = Task::new("prompt");
1411        let hooks = MockAgentImpl::new("test", "test");
1412
1413        let mut stream = engine
1414            .run_turn_stream(hooks, &task, context, &mut turn_state, 0, 1)
1415            .await
1416            .unwrap();
1417
1418        let mut final_text = String::default();
1419        while let Some(delta) = stream.next().await {
1420            if let Ok(TurnDelta::Done(result)) = delta {
1421                final_text = match result {
1422                    crate::agent::executor::TurnResult::Complete(output) => output.response,
1423                    crate::agent::executor::TurnResult::Continue(Some(output)) => output.response,
1424                    crate::agent::executor::TurnResult::Continue(None) => String::default(),
1425                };
1426                break;
1427            }
1428        }
1429
1430        assert_eq!(final_text, "Hello world");
1431    }
1432
1433    #[tokio::test]
1434    async fn test_run_turn_stream_structured_emits_reasoning_content() {
1435        use crate::tests::MockAgentImpl;
1436        let llm = Arc::new(ConfigurableLLMProvider {
1437            structured_stream: vec![StreamResponse {
1438                choices: vec![StreamChoice {
1439                    delta: StreamDelta {
1440                        content: None,
1441                        reasoning_content: Some("think".to_string()),
1442                        tool_calls: None,
1443                    },
1444                }],
1445                usage: None,
1446            }],
1447            ..ConfigurableLLMProvider::default()
1448        });
1449
1450        let config = AgentConfig {
1451            id: ActorID::new_v4(),
1452            name: "stream_reasoning_agent".to_string(),
1453            description: "desc".to_string(),
1454            output_schema: None,
1455        };
1456        let context = Arc::new(Context::new(llm, None).with_config(config));
1457        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1458        let mut turn_state = engine.turn_state(&context);
1459        let task = Task::new("prompt");
1460        let hooks = MockAgentImpl::new("test", "test");
1461
1462        let mut stream = engine
1463            .run_turn_stream(hooks, &task, context, &mut turn_state, 0, 1)
1464            .await
1465            .unwrap();
1466
1467        let mut saw_delta = false;
1468        let mut final_reasoning = String::default();
1469        while let Some(delta) = stream.next().await {
1470            match delta {
1471                Ok(TurnDelta::ReasoningContent(text)) => {
1472                    saw_delta = true;
1473                    assert_eq!(text, "think");
1474                }
1475                Ok(TurnDelta::Done(result)) => {
1476                    final_reasoning = match result {
1477                        crate::agent::executor::TurnResult::Complete(output) => {
1478                            output.reasoning_content
1479                        }
1480                        crate::agent::executor::TurnResult::Continue(Some(output)) => {
1481                            output.reasoning_content
1482                        }
1483                        crate::agent::executor::TurnResult::Continue(None) => String::default(),
1484                    };
1485                    break;
1486                }
1487                _ => {}
1488            }
1489        }
1490
1491        assert!(saw_delta);
1492        assert_eq!(final_reasoning, "think");
1493    }
1494
1495    #[tokio::test]
1496    async fn test_run_turn_stream_with_tools_executes_tools() {
1497        use crate::tests::MockAgentImpl;
1498        let tool_call = ToolCall {
1499            id: "call_1".to_string(),
1500            call_type: "function".to_string(),
1501            function: autoagents_llm::FunctionCall {
1502                name: "tool_a".to_string(),
1503                arguments: r#"{"value":1}"#.to_string(),
1504            },
1505        };
1506
1507        let llm = Arc::new(ConfigurableLLMProvider {
1508            stream_chunks: vec![
1509                StreamChunk::Text("thinking".to_string()),
1510                StreamChunk::ToolUseComplete {
1511                    index: 0,
1512                    tool_call: tool_call.clone(),
1513                },
1514                StreamChunk::Done {
1515                    stop_reason: "tool_use".to_string(),
1516                },
1517            ],
1518            ..ConfigurableLLMProvider::default()
1519        });
1520
1521        let config = AgentConfig {
1522            id: ActorID::new_v4(),
1523            name: "tool_stream_agent".to_string(),
1524            description: "desc".to_string(),
1525            output_schema: None,
1526        };
1527        let tool = LocalTool::new("tool_a", serde_json::json!({"ok": true}));
1528        let context = Arc::new(
1529            Context::new(llm, None)
1530                .with_config(config)
1531                .with_tools(vec![Box::new(tool)]),
1532        );
1533        let engine = TurnEngine::new(TurnEngineConfig {
1534            max_turns: 1,
1535            tool_mode: ToolMode::Enabled,
1536            stream_mode: StreamMode::Tool,
1537            memory_policy: MemoryPolicy::basic(),
1538        });
1539        let mut turn_state = engine.turn_state(&context);
1540        let task = Task::new("prompt");
1541        let hooks = MockAgentImpl::new("test", "test");
1542
1543        let mut stream = engine
1544            .run_turn_stream(hooks, &task, context, &mut turn_state, 0, 1)
1545            .await
1546            .unwrap();
1547
1548        let mut final_result = None;
1549        while let Some(delta) = stream.next().await {
1550            if let Ok(TurnDelta::Done(result)) = delta {
1551                final_result = Some(result);
1552                break;
1553            }
1554        }
1555
1556        match final_result.expect("done") {
1557            crate::agent::executor::TurnResult::Continue(Some(output)) => {
1558                assert_eq!(output.tool_calls.len(), 1);
1559                assert!(output.tool_calls[0].success);
1560            }
1561            _ => panic!("expected Continue(Some)"),
1562        }
1563    }
1564
1565    #[tokio::test]
1566    async fn test_run_turn_stream_llm_error_does_not_store_user_message() {
1567        use crate::tests::MockAgentImpl;
1568
1569        let llm: Arc<dyn LLMProvider> = Arc::new(GuardrailRejectLLMProvider);
1570        let context = Arc::new(context_with_memory(llm));
1571        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1572        let mut turn_state = engine.turn_state(&context);
1573        let task = Task::new("jailbreak");
1574        let hooks = MockAgentImpl::new("test", "test");
1575
1576        let mut stream = engine
1577            .run_turn_stream(hooks, &task, context.clone(), &mut turn_state, 0, 1)
1578            .await
1579            .expect("stream should initialize");
1580
1581        let first = stream
1582            .next()
1583            .await
1584            .expect("stream should emit an error event");
1585        assert!(matches!(
1586            first,
1587            Err(TurnEngineError::LLMError(LLMError::GuardrailBlocked { .. }))
1588        ));
1589
1590        let stored = recalled_messages(&context).await;
1591        assert!(stored.is_empty());
1592    }
1593
1594    #[tokio::test]
1595    async fn test_run_turn_stream_emits_memory_write_failure() {
1596        use crate::tests::MockAgentImpl;
1597
1598        let llm: Arc<dyn LLMProvider> = Arc::new(ConfigurableLLMProvider {
1599            structured_stream: vec![StreamResponse {
1600                choices: vec![StreamChoice {
1601                    delta: StreamDelta {
1602                        content: Some("hello".to_string()),
1603                        reasoning_content: None,
1604                        tool_calls: None,
1605                    },
1606                }],
1607                usage: None,
1608            }],
1609            ..ConfigurableLLMProvider::default()
1610        });
1611        let context = Arc::new(context_with_failing_memory(llm));
1612        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1613        let mut turn_state = engine.turn_state(&context);
1614        let task = Task::new("hello");
1615        let hooks = MockAgentImpl::new("test", "test");
1616
1617        let mut stream = engine
1618            .run_turn_stream(hooks, &task, context, &mut turn_state, 0, 1)
1619            .await
1620            .expect("stream should initialize");
1621
1622        let first = stream
1623            .next()
1624            .await
1625            .expect("stream should emit memory error");
1626        match first {
1627            Err(error) => assert_turn_memory_error(error),
1628            other => panic!("expected memory error, got {other:?}"),
1629        }
1630    }
1631
1632    #[tokio::test]
1633    async fn test_run_turn_stream_success_stores_user_once_in_memory() {
1634        use crate::tests::MockAgentImpl;
1635
1636        let llm: Arc<dyn LLMProvider> = Arc::new(ConfigurableLLMProvider {
1637            structured_stream: vec![StreamResponse {
1638                choices: vec![StreamChoice {
1639                    delta: StreamDelta {
1640                        content: Some("hello".to_string()),
1641                        reasoning_content: None,
1642                        tool_calls: None,
1643                    },
1644                }],
1645                usage: None,
1646            }],
1647            ..ConfigurableLLMProvider::default()
1648        });
1649        let context = Arc::new(context_with_memory(llm));
1650        let engine = TurnEngine::new(TurnEngineConfig::basic(1));
1651        let mut turn_state = engine.turn_state(&context);
1652        let task = Task::new("hello");
1653        let hooks = MockAgentImpl::new("test", "test");
1654
1655        let mut stream = engine
1656            .run_turn_stream(hooks, &task, context.clone(), &mut turn_state, 0, 1)
1657            .await
1658            .expect("stream should initialize");
1659
1660        while let Some(delta) = stream.next().await {
1661            if matches!(delta, Ok(TurnDelta::Done(_))) {
1662                break;
1663            }
1664        }
1665
1666        let stored = recalled_messages(&context).await;
1667        let user_count = stored
1668            .iter()
1669            .filter(|m| m.role == ChatRole::User && m.content == "hello")
1670            .count();
1671        let assistant_count = stored
1672            .iter()
1673            .filter(|m| m.role == ChatRole::Assistant)
1674            .count();
1675
1676        assert_eq!(user_count, 1);
1677        assert_eq!(assistant_count, 1);
1678    }
1679}
autoagents_core/agent/executor/turn_engine.rs

autoagents_core/agent/executor/
turn_engine.rs
