aether_core/testing/
utils.rs

1use crate::events::{ToolEvent, TurnEvent};
2use std::collections::BTreeMap;
3use std::error::Error;
4use std::sync::{Arc, Mutex};
5use std::time::Duration;
6use tokio::sync::mpsc;
7
8use futures::future::join_all;
9
10use crate::core::{RetryConfig, agent};
11use crate::events::{AgentEvent, AgentObserver, Command, UserCommand};
12use crate::mcp::mcp;
13use crate::testing::fake_mcp::fake_mcp;
14use crate::testing::{AgentTrace, FakeAgentObserver, FakeMcpServer};
15use llm::{Context, LlmError, LlmModel, LlmResponse};
16
17use llm::testing::FakeLlmProvider;
18
19pub async fn drain_until(
20    receiver: &mut mpsc::Receiver<AgentEvent>,
21    predicate: impl Fn(&AgentEvent) -> bool,
22) -> Vec<AgentEvent> {
23    let mut events = Vec::new();
24    while let Some(event) = receiver.recv().await {
25        let matched = predicate(&event);
26        events.push(event);
27        if matched {
28            return events;
29        }
30    }
31    panic!("agent event channel closed before predicate matched");
32}
33
34pub fn content_events(events: Vec<AgentEvent>) -> Vec<AgentEvent> {
35    events
36        .into_iter()
37        .filter(|event| {
38            !matches!(
39                event,
40                AgentEvent::Turn(
41                    TurnEvent::Started { .. } | TurnEvent::LlmCallStarted { .. } | TurnEvent::LlmCallEnded { .. }
42                ) | AgentEvent::Tool(ToolEvent::ExecutionStarted { .. } | ToolEvent::DefinitionsUpdated { .. })
43            )
44        })
45        .collect()
46}
47
48pub fn mcp_instructions(entries: &[(&str, &str)]) -> BTreeMap<String, String> {
49    entries.iter().map(|(k, v)| ((*k).to_string(), (*v).to_string())).collect()
50}
51
52pub fn test_agent() -> TestAgentBuilder {
53    TestAgentBuilder::new()
54}
55
56/// Result of running a test agent, including messages and captured contexts.
57pub struct TestAgentResult {
58    pub messages: Vec<AgentEvent>,
59    pub captured_contexts: Arc<Mutex<Vec<Context>>>,
60}
61
62type CancelPredicate = Box<dyn Fn(&AgentEvent) -> bool + Send>;
63
64pub struct TestAgentBuilder {
65    messages: Vec<Command>,
66    responses: Vec<Vec<Result<LlmResponse, LlmError>>>,
67    model: Option<LlmModel>,
68    context_window: Option<u32>,
69    timeout: Option<Duration>,
70    max_auto_continues: Option<u32>,
71    retry_config: Option<RetryConfig>,
72    observers: Vec<Box<dyn AgentObserver>>,
73    cancel_when: Option<CancelPredicate>,
74}
75
76impl Default for TestAgentBuilder {
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82impl TestAgentBuilder {
83    pub fn new() -> Self {
84        Self {
85            messages: Vec::new(),
86            responses: Vec::new(),
87            model: None,
88            context_window: None,
89            timeout: None,
90            max_auto_continues: None,
91            retry_config: None,
92            observers: Vec::new(),
93            cancel_when: None,
94        }
95    }
96
97    pub fn user_messages(mut self, user_messages: Vec<Command>) -> Self {
98        self.messages = user_messages;
99        self
100    }
101
102    pub fn user_text(self, text: &str) -> Self {
103        self.user_messages(vec![Command::UserCommand(UserCommand::Text {
104            content: vec![llm::ContentBlock::text(text)],
105        })])
106    }
107
108    pub fn llm_responses(mut self, llm_responses: &[Vec<LlmResponse>]) -> Self {
109        self.responses = llm_responses.iter().map(|turn| turn.iter().cloned().map(Ok).collect()).collect();
110        self
111    }
112
113    pub fn llm_result_responses(mut self, llm_responses: &[Vec<Result<LlmResponse, LlmError>>]) -> Self {
114        self.responses = Vec::from(llm_responses);
115        self
116    }
117
118    pub fn model(mut self, model: LlmModel) -> Self {
119        self.model = Some(model);
120        self
121    }
122
123    pub fn context_window(mut self, window: u32) -> Self {
124        self.context_window = Some(window);
125        self
126    }
127
128    /// Sends `UserCommand::Cancel` as soon as a received event matches
129    /// `predicate`, enabling deterministic mid-turn cancellation tests.
130    pub fn cancel_when(mut self, predicate: impl Fn(&AgentEvent) -> bool + Send + 'static) -> Self {
131        self.cancel_when = Some(Box::new(predicate));
132        self
133    }
134
135    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
136        self.timeout = Some(timeout);
137        self
138    }
139
140    pub fn max_auto_continues(mut self, max: u32) -> Self {
141        self.max_auto_continues = Some(max);
142        self
143    }
144
145    pub fn retry_config(mut self, config: RetryConfig) -> Self {
146        self.retry_config = Some(config);
147        self
148    }
149
150    /// Attach an observer of the test agent's event stream.
151    pub fn observer(mut self, observer: Box<dyn AgentObserver>) -> Self {
152        self.observers.push(observer);
153        self
154    }
155
156    pub async fn run(self) -> Result<Vec<AgentEvent>, Box<dyn Error>> {
157        let result = self.run_with_context().await?;
158        Ok(result.messages)
159    }
160
161    /// Runs the test agent with a recording observer attached and returns the
162    /// full event trace, including internal events.
163    pub async fn run_trace(self) -> Result<AgentTrace, Box<dyn Error>> {
164        let observer = FakeAgentObserver::new();
165        let events = observer.events();
166        self.observer(Box::new(observer)).run().await?;
167        Ok(AgentTrace::from_observer_events(&events))
168    }
169
170    /// Runs the test agent and returns both messages and captured contexts.
171    ///
172    /// Use this when you need to verify what context was passed to the LLM,
173    /// for example when testing that file attachments are properly formatted.
174    pub async fn run_with_context(self) -> Result<TestAgentResult, Box<dyn Error>> {
175        let mut llm = FakeLlmProvider::from_results(self.responses).with_context_window(self.context_window);
176        if let Some(model) = self.model {
177            llm = llm.with_model(model);
178        }
179        let captured_contexts = llm.captured_contexts();
180
181        let mut spawn = mcp("/workspace").with_servers(vec![fake_mcp("test", FakeMcpServer::new())]).spawn().await?;
182        let snapshot = spawn.block_until_ready().await.expect("bootstrap completes");
183
184        let mut builder = agent(llm).tools(spawn.command_tx, snapshot.tool_definitions);
185        if let Some(timeout) = self.timeout {
186            builder = builder.tool_timeout(timeout);
187        }
188        if let Some(max) = self.max_auto_continues {
189            builder = builder.max_auto_continues(max);
190        }
191        if let Some(retry) = self.retry_config {
192            builder = builder.retry(retry);
193        } else {
194            builder = builder.retry(RetryConfig::disabled());
195        }
196        for observer in self.observers {
197            builder = builder.observer(observer);
198        }
199
200        let (tx, mut rx, handle) = builder.spawn().await?;
201        let futures: Vec<_> = self.messages.into_iter().map(|m| tx.send(m)).collect();
202
203        join_all(futures).await;
204
205        let mut command_tx = if self.cancel_when.is_some() {
206            Some(tx)
207        } else {
208            drop(tx);
209            None
210        };
211        let mut messages = Vec::new();
212        while let Some(message) = rx.recv().await {
213            messages.push(message.clone());
214            if self.cancel_when.as_ref().is_some_and(|predicate| predicate(&message))
215                && let Some(tx) = command_tx.take()
216            {
217                tx.send(Command::UserCommand(UserCommand::Cancel)).await?;
218            }
219            if matches!(message, AgentEvent::Turn(TurnEvent::Ended { .. })) {
220                break;
221            }
222        }
223        drop(command_tx);
224
225        handle.await_completion().await;
226
227        Ok(TestAgentResult { messages, captured_contexts })
228    }
229}
aether_core/testing/utils.rs

aether_core/testing/
utils.rs
