aether_core/testing/
utils.rs

1use std::collections::BTreeMap;
2use std::error::Error;
3use std::sync::{Arc, Mutex};
4use std::time::Duration;
5use tokio::sync::{Notify, mpsc};
6
7use crate::context::CompactionConfig;
8use crate::core::{Prompt, RetryConfig, agent};
9use crate::events::{
10    AgentCommand, AgentEvent, AgentObserver, Command, ContextEvent, ToolEvent, TurnEvent, UserCommand,
11};
12use crate::mcp::mcp;
13use crate::testing::fake_mcp::fake_mcp;
14use crate::testing::{AgentTrace, FakeAgentObserver, FakeMcpServer};
15use llm::{ChatMessage, Context, LlmError, LlmModel, LlmResponse, ModelSettings, StreamingModelProvider};
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/// An ordered interaction with a test agent.
57pub enum TestAgentStep {
58    Send(Command),
59    WaitFor(Box<dyn Fn(&AgentEvent) -> bool + Send>),
60    /// Run an arbitrary side effect (e.g. releasing a paused LLM stream) between steps.
61    Perform(Box<dyn FnOnce() + Send>),
62}
63
64impl TestAgentStep {
65    pub fn send(command: Command) -> Self {
66        Self::Send(command)
67    }
68
69    pub fn user_text(text: impl Into<String>) -> Self {
70        Self::send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text(text.into())] }))
71    }
72
73    pub fn cancel() -> Self {
74        Self::send(Command::UserCommand(UserCommand::Cancel))
75    }
76
77    pub fn switch_model(provider: impl StreamingModelProvider + 'static) -> Self {
78        Self::send(Command::AgentCommand(AgentCommand::SwitchModel(Box::new(provider))))
79    }
80
81    pub fn replace_conversation(messages: Vec<ChatMessage>) -> Self {
82        Self::send(Command::AgentCommand(AgentCommand::ReplaceConversation(messages)))
83    }
84
85    pub fn perform(action: impl FnOnce() + Send + 'static) -> Self {
86        Self::Perform(Box::new(action))
87    }
88
89    pub fn wait_for(predicate: impl Fn(&AgentEvent) -> bool + Send + 'static) -> Self {
90        Self::WaitFor(Box::new(predicate))
91    }
92
93    pub fn wait_for_turn_end() -> Self {
94        Self::wait_for(|event| matches!(event, AgentEvent::Turn(TurnEvent::Ended { .. })))
95    }
96
97    pub fn wait_for_compaction_start() -> Self {
98        Self::wait_for(|event| matches!(event, AgentEvent::Context(ContextEvent::CompactionStarted { .. })))
99    }
100
101    pub fn wait_for_retry(attempt: u32) -> Self {
102        Self::wait_for(
103            move |event| matches!(event, AgentEvent::Turn(TurnEvent::RetryScheduled { attempt: actual, .. }) if *actual == attempt),
104        )
105    }
106}
107
108/// A fluent sequence of commands and synchronization points for a test agent.
109#[derive(Default)]
110pub struct TestScenario {
111    steps: Vec<TestAgentStep>,
112}
113
114impl TestScenario {
115    pub fn new() -> Self {
116        Self::default()
117    }
118
119    pub fn send(mut self, command: Command) -> Self {
120        self.steps.push(TestAgentStep::send(command));
121        self
122    }
123
124    pub fn user_text(mut self, text: impl Into<String>) -> Self {
125        self.steps.push(TestAgentStep::user_text(text));
126        self
127    }
128
129    pub fn cancel(mut self) -> Self {
130        self.steps.push(TestAgentStep::cancel());
131        self
132    }
133
134    pub fn switch_model(mut self, provider: impl StreamingModelProvider + 'static) -> Self {
135        self.steps.push(TestAgentStep::switch_model(provider));
136        self
137    }
138
139    pub fn replace_conversation(mut self, messages: Vec<ChatMessage>) -> Self {
140        self.steps.push(TestAgentStep::replace_conversation(messages));
141        self
142    }
143
144    pub fn wait_for(mut self, predicate: impl Fn(&AgentEvent) -> bool + Send + 'static) -> Self {
145        self.steps.push(TestAgentStep::wait_for(predicate));
146        self
147    }
148
149    pub fn wait_for_turn_end(mut self) -> Self {
150        self.steps.push(TestAgentStep::wait_for_turn_end());
151        self
152    }
153
154    pub fn wait_for_compaction_start(mut self) -> Self {
155        self.steps.push(TestAgentStep::wait_for_compaction_start());
156        self
157    }
158
159    pub fn wait_for_retry(mut self, attempt: u32) -> Self {
160        self.steps.push(TestAgentStep::wait_for_retry(attempt));
161        self
162    }
163
164    /// Run an arbitrary side effect between scenario steps, e.g. releasing a
165    /// paused LLM stream so a queued message can be injected mid-turn.
166    pub fn perform(mut self, action: impl FnOnce() + Send + 'static) -> Self {
167        self.steps.push(TestAgentStep::perform(action));
168        self
169    }
170}
171
172impl From<Vec<TestAgentStep>> for TestScenario {
173    fn from(steps: Vec<TestAgentStep>) -> Self {
174        Self { steps }
175    }
176}
177
178/// Result of running a test agent, including messages and captured contexts.
179pub struct TestAgentResult {
180    pub messages: Vec<AgentEvent>,
181    pub captured_contexts: Arc<Mutex<Vec<Context>>>,
182}
183
184struct ProviderTestConfig {
185    responses: Vec<Vec<Result<LlmResponse, LlmError>>>,
186    model: Option<LlmModel>,
187    context_window: Option<u32>,
188    pause: Option<(usize, usize, Arc<Notify>)>,
189}
190
191struct AgentTestConfig {
192    context_window_override: Option<u32>,
193    timeout: Option<Duration>,
194    max_auto_continues: Option<u32>,
195    retry_config: Option<RetryConfig>,
196    observers: Vec<Box<dyn AgentObserver>>,
197    include_fake_mcp: bool,
198    initial_messages: Vec<ChatMessage>,
199    system_prompt: Option<Prompt>,
200    compaction: Option<CompactionConfig>,
201    model_settings: Option<ModelSettings>,
202}
203
204enum TestExecution {
205    CommandsUntilTurnEnd(Vec<Command>),
206    Scenario(TestScenario),
207}
208
209pub struct TestAgentBuilder {
210    provider: ProviderTestConfig,
211    agent: AgentTestConfig,
212    execution: Option<TestExecution>,
213}
214
215impl Default for TestAgentBuilder {
216    fn default() -> Self {
217        Self::new()
218    }
219}
220
221impl TestAgentBuilder {
222    pub fn new() -> Self {
223        Self {
224            provider: ProviderTestConfig { responses: Vec::new(), model: None, context_window: None, pause: None },
225            agent: AgentTestConfig {
226                context_window_override: None,
227                timeout: None,
228                max_auto_continues: None,
229                retry_config: None,
230                observers: Vec::new(),
231                include_fake_mcp: true,
232                initial_messages: Vec::new(),
233                system_prompt: None,
234                compaction: None,
235                model_settings: None,
236            },
237            execution: None,
238        }
239    }
240
241    pub fn commands(self, commands: Vec<Command>) -> Self {
242        self.with_execution(TestExecution::CommandsUntilTurnEnd(commands))
243    }
244
245    pub fn scenario(self, scenario: impl Into<TestScenario>) -> Self {
246        self.with_execution(TestExecution::Scenario(scenario.into()))
247    }
248
249    pub fn user_text(self, text: &str) -> Self {
250        self.commands(vec![Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text(text)] })])
251    }
252
253    pub fn llm_responses(mut self, llm_responses: &[Vec<LlmResponse>]) -> Self {
254        self.provider.responses = llm_responses.iter().map(|turn| turn.iter().cloned().map(Ok).collect()).collect();
255        self
256    }
257
258    pub fn llm_result_responses(mut self, llm_responses: &[Vec<Result<LlmResponse, LlmError>>]) -> Self {
259        self.provider.responses = Vec::from(llm_responses);
260        self
261    }
262
263    pub fn model(mut self, model: LlmModel) -> Self {
264        self.provider.model = Some(model);
265        self
266    }
267
268    pub fn provider_context_window(mut self, window: Option<u32>) -> Self {
269        self.provider.context_window = window;
270        self
271    }
272
273    pub fn context_window_override(mut self, window: u32) -> Self {
274        self.agent.context_window_override = Some(window);
275        self
276    }
277
278    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
279        self.agent.timeout = Some(timeout);
280        self
281    }
282
283    pub fn max_auto_continues(mut self, max: u32) -> Self {
284        self.agent.max_auto_continues = Some(max);
285        self
286    }
287
288    pub fn retry_config(mut self, config: RetryConfig) -> Self {
289        self.agent.retry_config = Some(config);
290        self
291    }
292
293    /// Run without the default fake MCP server when the scenario does not exercise tools.
294    pub fn without_mcp(mut self) -> Self {
295        self.agent.include_fake_mcp = false;
296        self
297    }
298
299    /// Pre-populate the context with conversation history.
300    pub fn messages(mut self, messages: Vec<ChatMessage>) -> Self {
301        self.agent.initial_messages = messages;
302        self
303    }
304
305    /// Set the system prompt.
306    pub fn system_prompt(mut self, prompt: Prompt) -> Self {
307        self.agent.system_prompt = Some(prompt);
308        self
309    }
310
311    /// Configure context compaction settings.
312    pub fn compaction_config(mut self, config: CompactionConfig) -> Self {
313        self.agent.compaction = Some(config);
314        self
315    }
316
317    /// Set the model settings applied to every LLM call.
318    pub fn model_settings(mut self, settings: ModelSettings) -> Self {
319        self.agent.model_settings = Some(settings);
320        self
321    }
322
323    /// Pause the fake LLM stream at `turn_index` / `chunk_index` until
324    /// `release.notify_one()` is called. Used for deterministic timing tests.
325    pub fn pause_turn_after(mut self, turn_index: usize, chunk_index: usize, release: Arc<Notify>) -> Self {
326        self.provider.pause = Some((turn_index, chunk_index, release));
327        self
328    }
329
330    /// Attach an observer of the test agent's event stream.
331    pub fn observer(mut self, observer: Box<dyn AgentObserver>) -> Self {
332        self.agent.observers.push(observer);
333        self
334    }
335
336    pub async fn run(self) -> Result<Vec<AgentEvent>, Box<dyn Error>> {
337        let result = self.run_with_context().await?;
338        Ok(result.messages)
339    }
340
341    /// Runs the test agent with a recording observer attached and returns the
342    /// full event trace, including internal events.
343    pub async fn run_trace(self) -> Result<AgentTrace, Box<dyn Error>> {
344        let observer = FakeAgentObserver::new();
345        let events = observer.events();
346        self.observer(Box::new(observer)).run().await?;
347        Ok(AgentTrace::from_observer_events(&events))
348    }
349
350    /// Runs the test agent and returns both messages and captured contexts.
351    ///
352    /// Use this when you need to verify what context was passed to the LLM,
353    /// for example when testing that file attachments are properly formatted.
354    pub async fn run_with_context(self) -> Result<TestAgentResult, Box<dyn Error>> {
355        let Self { provider, agent: config, execution } = self;
356        let mut llm = FakeLlmProvider::from_results(provider.responses).with_context_window(provider.context_window);
357        if let Some(model) = provider.model {
358            llm = llm.with_model(model);
359        }
360        if let Some((turn_index, chunk_index, release)) = provider.pause {
361            llm = llm.pause_turn_after(turn_index, chunk_index, release);
362        }
363        let captured_contexts = llm.captured_contexts();
364
365        let mut mcp_spawn = if config.include_fake_mcp {
366            Some(mcp("/workspace").with_servers(vec![fake_mcp("test", FakeMcpServer::new())]).spawn().await?)
367        } else {
368            None
369        };
370
371        let mut builder = agent(llm);
372        if let Some(spawn) = &mut mcp_spawn {
373            let snapshot = spawn.block_until_ready().await.expect("bootstrap completes");
374            builder = builder.tools(spawn.command_tx.clone(), snapshot.tool_definitions);
375        }
376        if let Some(timeout) = config.timeout {
377            builder = builder.tool_timeout(timeout);
378        }
379        if let Some(max) = config.max_auto_continues {
380            builder = builder.max_auto_continues(max);
381        }
382        if let Some(retry) = config.retry_config {
383            builder = builder.retry(retry);
384        } else {
385            builder = builder.retry(RetryConfig::disabled());
386        }
387        if let Some(prompt) = config.system_prompt {
388            builder = builder.system_prompt(prompt);
389        }
390        if let Some(compaction) = config.compaction {
391            builder = builder.compaction(compaction);
392        }
393        if let Some(settings) = config.model_settings {
394            builder = builder.model_settings(settings);
395        }
396        builder = builder.context_window(config.context_window_override);
397        if !config.initial_messages.is_empty() {
398            builder = builder.messages(config.initial_messages);
399        }
400        for observer in config.observers {
401            builder = builder.observer(observer);
402        }
403
404        let steps = match execution.expect("test agent requires commands(), user_text(), or scenario()") {
405            TestExecution::CommandsUntilTurnEnd(commands) => {
406                assert!(!commands.is_empty(), "commands() requires at least one command");
407                let mut steps = commands.into_iter().map(TestAgentStep::send).collect::<Vec<_>>();
408                steps.push(TestAgentStep::wait_for_turn_end());
409                steps
410            }
411            TestExecution::Scenario(scenario) => {
412                assert!(!scenario.steps.is_empty(), "scenario() requires at least one step");
413                scenario.steps
414            }
415        };
416        let (tx, mut rx, handle) = builder.spawn().await?;
417        let mut messages = Vec::new();
418
419        for step in steps {
420            match step {
421                TestAgentStep::Send(command) => tx.send(command).await?,
422                TestAgentStep::WaitFor(predicate) => loop {
423                    let message = rx.recv().await.expect("agent event channel closed before scenario step matched");
424                    let matched = predicate(&message);
425                    messages.push(message);
426                    if matched {
427                        break;
428                    }
429                },
430                TestAgentStep::Perform(action) => action(),
431            }
432        }
433        drop(tx);
434
435        handle.await_completion().await;
436
437        Ok(TestAgentResult { messages, captured_contexts })
438    }
439
440    fn with_execution(mut self, execution: TestExecution) -> Self {
441        assert!(self.execution.is_none(), "commands(), user_text(), and scenario() are mutually exclusive");
442        self.execution = Some(execution);
443        self
444    }
445}
aether_core/testing/utils.rs

aether_core/testing/
utils.rs
